bench_binary_arith.py

from tileops.utils import get_backend_name

DEVICE = get_backend_name()
"""Benchmarks for binary arithmetic ops covering risk points R1, R2, R4.

Risk points covered:
- R1: Stride-based load vectorization (add x explicit_parallel x fp16 x
      {1D same-shape, 2D bias-add, 3D interleaved})
- R2: Divmod overhead on small tensors (add same-shape/3D-broadcast x fp16 x 4K)
- R4: DEFAULT_STRATEGY confirmation (add x 2 strategies x 3 dtypes x 3 sizes x
      {same-shape, 2D bias-add, 3D interleaved})

Profiles both binary strategies (direct, explicit_parallel) and compares
against PyTorch baseline.
"""

from math import prod
from typing import Optional, Protocol

import pytest
import torch

from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.ops.elementwise import AddFwdOp, LerpTensorFwdOp, WhereFwdOp
from workloads.binary_arith import AddSameShapeTest
from workloads.workload_base import FixtureBase

# ---------------------------------------------------------------------------
# LLM-realistic shapes (LLaMA-family defaults)
# ---------------------------------------------------------------------------

# Per-strategy/broadcast matrix sizes. Each label maps to a 2D shape that
# both the same-shape and broadcast patterns can derive from. The third
# entry is non-pow2 in the hidden dim to exercise tail handling.
_SHAPE_BY_LABEL: dict[str, tuple[int, int]] = {
    "4K": (1, 4096),
    "1M": (1024, 1024),
    "11M": (1024, 11008),
}

_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
_BINARY_STRATEGIES = ("direct", "explicit_parallel")

def _make_interleaved_3d(n: int) -> tuple[tuple, tuple]:
    """Build (A,1,C) + (1,B,1) -> (A,B,C) with A*B*C == n exactly.

    Uses A=8 (or 1 for very small n). Finds the largest B <= sqrt(n/A)
    that divides n/A evenly, then C = n/(A*B).
    """
    if n < 8:
        return (1, 1, n), (1, n, 1)
    a_dim = 8
    remainder = n // a_dim
    b_dim = int(remainder ** 0.5)
    while b_dim > 1 and remainder % b_dim != 0:
        b_dim -= 1
    c_dim = remainder // b_dim
    return (a_dim, 1, c_dim), (1, b_dim, 1)


# Broadcast patterns for binary ops. Each pattern derives a (a_shape,
# b_shape) pair from a 2D output shape (M, N), preserving model geometry.
_BROADCAST_PATTERNS = {
    "same_shape": lambda mn: (mn, mn),
    "bias_add_2d": lambda mn: (mn, (1, mn[1])),
    "interleaved_3d": lambda mn: _make_interleaved_3d(mn[0] * mn[1]),
}


# ---------------------------------------------------------------------------
# Benchmark harness
# ---------------------------------------------------------------------------


class BinaryWorkload(Protocol):
    """Structural type for binary benchmark workloads.

    Requires ``n_total``, ``dtype``, and ``gen_inputs``.  Attributes
    ``a_shape`` / ``b_shape`` are optional — ``BinaryBenchmark`` falls
    back to ``n_total`` when they are absent.
    """

    n_total: int
    dtype: torch.dtype

    def gen_inputs(self) -> tuple[torch.Tensor, ...]: ...


class BinaryBenchCase:
    """Minimal test harness for binary benchmarks."""

    def __init__(
        self, a_shape: tuple, b_shape: tuple, dtype: torch.dtype,
    ):
        self.a_shape = a_shape
        self.b_shape = b_shape
        self.dtype = dtype
        self.n_total = prod(torch.broadcast_shapes(a_shape, b_shape))

    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor]:
        a = torch.randn(self.a_shape, device=DEVICE, dtype=self.dtype)
        b = torch.randn(self.b_shape, device=DEVICE, dtype=self.dtype)
        return a, b


class BinaryBenchmark(BenchmarkBase[BinaryWorkload]):
    """Bandwidth-oriented benchmark for binary elementwise ops."""

    def calculate_flops(self) -> Optional[float]:
        return self.workload.n_total

    def calculate_memory(self) -> Optional[float]:
        t = self.workload
        elem_bytes = t.dtype.itemsize
        # Read a + read b + write output
        a_elems = prod(getattr(t, "a_shape", (t.n_total,)))
        b_elems = prod(getattr(t, "b_shape", (t.n_total,)))
        return (a_elems + b_elems + t.n_total) * elem_bytes


class WhereBenchCase:
    """Test harness for where op benchmarks."""

    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
        self.shape = shape
        self.n_total = prod(shape)
        self.dtype = dtype

    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        cond = torch.randint(0, 2, self.shape, device=DEVICE, dtype=torch.bool)
        x = torch.randn(*self.shape, device=DEVICE, dtype=self.dtype)
        y = torch.randn(*self.shape, device=DEVICE, dtype=self.dtype)
        return cond, x, y


class WhereBenchmark(BenchmarkBase[WhereBenchCase]):
    """Benchmark for where op."""

    def calculate_flops(self) -> Optional[float]:
        return self.workload.n_total

    def calculate_memory(self) -> Optional[float]:
        t = self.workload
        elem_bytes = t.dtype.itemsize
        # Read cond (1 byte) + read x + read y + write output
        return t.n_total * (1 + 3 * elem_bytes)


# ---------------------------------------------------------------------------
# R1: Stride-based load vectorization
# ---------------------------------------------------------------------------

_R1_PATTERNS = [
    ("same_shape_1d", (1_000_000,), (1_000_000,)),
    # bias-add: (1000, 1000) + (1, 1000) -> 1,000,000 output elements
    ("bias_add_2d", (1000, 1000), (1, 1000)),
    # interleaved: (8,1,1024) + (1,128,1) -> (8,128,1024) = 1,048,576 output
    ("interleaved_3d", (8, 1, 1024), (1, 128, 1)),
]


class R1VectorizationFixture(FixtureBase):
    PARAMS = [
        ("pattern_name, a_shape, b_shape", [
            pytest.param(name, a, b, marks=pytest.mark.smoke if name == "same_shape_1d"
                         else pytest.mark.full)
            for name, a, b in _R1_PATTERNS
        ]),
    ]


@R1VectorizationFixture
def test_r1_vectorization(
    pattern_name: str,
    a_shape: tuple,
    b_shape: tuple,
) -> None:
    """R1: Benchmark stride-based load vectorization.

    Binary divmod offset may prevent uint4 vectorized loads.
    Compares same-shape (no divmod) vs broadcast patterns (divmod required).
    """
    dtype = torch.float16
    test = BinaryBenchCase(a_shape, b_shape, dtype)
    bm = BinaryBenchmark(test)
    inputs = test.gen_inputs()

    op = AddFwdOp(
        a_shape=a_shape, b_shape=b_shape, dtype=dtype,
        strategy="explicit_parallel",
    )
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r1_vectorization",
        {"pattern_name": pattern_name, "a_shape": a_shape, "b_shape": b_shape},
        result,
        tag=f"add_{pattern_name}",
    )

    # Baseline: PyTorch add with broadcast
    a, b = inputs

    def baseline_fn(a, b):
        return a + b

    result_bl = bm.profile(baseline_fn, a, b)
    BenchmarkReport.record(
        "r1_vectorization",
        {"pattern_name": pattern_name, "a_shape": a_shape, "b_shape": b_shape},
        result_bl,
        tag=f"torch-{pattern_name}",
    )


# ---------------------------------------------------------------------------
# R2: Divmod overhead on small tensors (binary)
# ---------------------------------------------------------------------------


class R2BinaryFixture(FixtureBase):
    PARAMS = [
        ("pattern_name, a_shape, b_shape", [
            pytest.param("same_shape", (4096,), (4096,), marks=pytest.mark.smoke),
            pytest.param(
                "broadcast_3d", (4, 1, 32), (1, 32, 1),
                marks=pytest.mark.full,
            ),
        ]),
    ]


@R2BinaryFixture
def test_r2_small_tensor_binary(
    pattern_name: str,
    a_shape: tuple,
    b_shape: tuple,
) -> None:
    """R2: Benchmark divmod overhead on small tensors (binary add, 4K)."""
    dtype = torch.float16
    test = BinaryBenchCase(a_shape, b_shape, dtype)
    bm = BinaryBenchmark(test)
    inputs = test.gen_inputs()

    op = AddFwdOp(a_shape=a_shape, b_shape=b_shape, dtype=dtype)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r2_small_tensor_binary",
        {"pattern_name": pattern_name, "a_shape": a_shape, "b_shape": b_shape},
        result,
        tag=f"add_{pattern_name}",
    )

    a, b = inputs

    def baseline_fn(a, b):
        return a + b

    result_bl = bm.profile(baseline_fn, a, b)
    BenchmarkReport.record(
        "r2_small_tensor_binary",
        {"pattern_name": pattern_name, "a_shape": a_shape, "b_shape": b_shape},
        result_bl,
        tag=f"torch-{pattern_name}",
    )


# ---------------------------------------------------------------------------
# R4: DEFAULT_STRATEGY confirmation (binary full matrix)
# ---------------------------------------------------------------------------


_R4_BINARY_PARAMS = []
for size_label, _shape_2d in _SHAPE_BY_LABEL.items():
    for dt in _DTYPES:
        for strategy in _BINARY_STRATEGIES:
            for pat_name, pat_fn in _BROADCAST_PATTERNS.items():
                a_shape, b_shape = pat_fn(_shape_2d)
                mark = pytest.mark.smoke if (
                    size_label == "1M" and dt == torch.float16
                    and strategy == "explicit_parallel"
                    and pat_name == "same_shape"
                ) else pytest.mark.full
                _R4_BINARY_PARAMS.append(
                    pytest.param(
                        a_shape, b_shape, dt, strategy, size_label, pat_name,
                        id=f"{size_label}-{dt}-{strategy}-{pat_name}",
                        marks=mark,
                    )
                )


class R4BinaryStrategyFixture(FixtureBase):
    PARAMS = [
        ("a_shape, b_shape, dtype, strategy, size_label, pattern_name",
         _R4_BINARY_PARAMS),
    ]


@R4BinaryStrategyFixture
def test_r4_default_strategy_binary(
    a_shape: tuple,
    b_shape: tuple,
    dtype: torch.dtype,
    strategy: str,
    size_label: str,
    pattern_name: str,
) -> None:
    """R4: Benchmark both binary strategies across full matrix.

    Covers: add x {direct, explicit_parallel} x {fp32, fp16, bf16}
            x {4K, 1M, 16M} x {same-shape, bias-add, interleaved-3D}
    """
    test = BinaryBenchCase(a_shape, b_shape, dtype)
    bm = BinaryBenchmark(test)
    inputs = test.gen_inputs()

    op = AddFwdOp(
        a_shape=a_shape, b_shape=b_shape, dtype=dtype, strategy=strategy,
    )
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r4_strategy_binary",
        {
            "size_label": size_label,
            "pattern_name": pattern_name,
            "a_shape": a_shape,
            "b_shape": b_shape,
            "dtype": dtype,
            "strategy": strategy,
        },
        result,
        tag=f"add_{strategy}_{pattern_name}",
    )


# ---------------------------------------------------------------------------
# R4: Where op strategy comparison (3-input op)
# ---------------------------------------------------------------------------


_R4_WHERE_PARAMS = []
for size_label, _shape_2d in _SHAPE_BY_LABEL.items():
    _R4_WHERE_PARAMS.append(
        pytest.param(
            _shape_2d, size_label, torch.float16,
            id=f"where-{size_label}-fp16",
            marks=pytest.mark.full,
        )
    )


class R4WhereFixture(FixtureBase):
    PARAMS = [
        ("shape, size_label, dtype", _R4_WHERE_PARAMS),
    ]


@R4WhereFixture
def test_r4_where_bench(
    shape: tuple[int, ...],
    size_label: str,
    dtype: torch.dtype,
) -> None:
    """R4: Benchmark where op across sizes."""
    test = WhereBenchCase(shape, dtype)
    bm = WhereBenchmark(test)
    inputs = test.gen_inputs()

    op = WhereFwdOp(condition=shape, input=shape, other=shape, dtype=dtype)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r4_where",
        {"shape": shape, "size_label": size_label, "dtype": dtype},
        result,
        tag="tileops-where",
    )

    cond, x, y = inputs

    def baseline_fn(cond, x, y):
        return torch.where(cond, x, y)

    result_bl = bm.profile(baseline_fn, cond, x, y)
    BenchmarkReport.record(
        "r4_where",
        {"shape": shape, "size_label": size_label, "dtype": dtype},
        result_bl,
        tag="torch",
    )


# ---------------------------------------------------------------------------
# Baseline throughput benchmarks (existing, refined with LLaMA shapes)
# ---------------------------------------------------------------------------


_ADD_BENCH_PARAMS = [
    pytest.param((1024, 4096), torch.float16, id="throughput-fp16"),
    pytest.param((1024, 4096), torch.bfloat16, id="throughput-bf16"),
    pytest.param((1024, 4096), torch.float32, id="baseline-fp32"),
]


@pytest.mark.parametrize("shape, dtype", _ADD_BENCH_PARAMS)
def test_add_bench(shape: tuple[int, ...], dtype: torch.dtype) -> None:
    n_total = prod(shape)
    # ``AddSameShapeTest`` (workloads) accepts a flat element count; the
    # bench harness still records the original shape tuple via
    # ``record(...)`` so the report carries the input geometry verbatim.
    test = AddSameShapeTest(n_total, dtype)
    bm = BinaryBenchmark(test)
    inputs = test.gen_inputs()

    op = AddFwdOp(a_shape=shape, b_shape=shape, dtype=dtype)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(op, locals(), result, tag="tileops")

    def baseline_fn(a, b):
        return a + b

    result_bl = bm.profile(baseline_fn, *inputs)
    BenchmarkReport.record(op, locals(), result_bl, tag="torch")


# ---------------------------------------------------------------------------
# LerpTensorFwdOp — Tensor-weight torch.lerp benchmark.
#
# Per output element: 3 flops (sub + mul + add); 3 reads + 1 write at
# post-broadcast ``N_total`` (matches
# ``tileops.perf.formulas.lerp_tensor_fwd_roofline``). Same-shape inputs
# only here; the broadcast contract is exercised by the test suite.
# ---------------------------------------------------------------------------


class LerpTensorBenchCase:
    """Same-shape input/end/weight; output broadcast equals the shape."""

    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
        self.shape = shape
        self.n_total = prod(shape)
        self.dtype = dtype

    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        a = torch.randn(self.shape, device=DEVICE, dtype=self.dtype)
        b = torch.randn(self.shape, device=DEVICE, dtype=self.dtype)
        # Keep weight in [0, 1] to stay close to typical lerp usage.
        w = torch.rand(self.shape, device=DEVICE, dtype=self.dtype)
        return a, b, w


class LerpTensorBenchmark(BenchmarkBase[LerpTensorBenchCase]):
    """Bandwidth-oriented benchmark for ``LerpTensorFwdOp``."""

    def calculate_flops(self) -> Optional[float]:
        return 3 * self.workload.n_total

    def calculate_memory(self) -> Optional[float]:
        t = self.workload
        return 4 * t.n_total * t.dtype.itemsize


_LERP_TENSOR_BENCH_PARAMS = [
    pytest.param((1024, 4096), torch.float16,
                 id="lerp-tensor-fp16-1024x4096", marks=pytest.mark.smoke),
    pytest.param((1024, 4096), torch.bfloat16,
                 id="lerp-tensor-bf16-1024x4096", marks=pytest.mark.full),
    pytest.param((1024, 4096), torch.float32,
                 id="lerp-tensor-fp32-1024x4096", marks=pytest.mark.full),
    pytest.param((1024, 10240), torch.float16,
                 id="lerp-tensor-fp16-1024x10240", marks=pytest.mark.full),
    pytest.param((1024, 11008), torch.float16,
                 id="lerp-tensor-fp16-1024x11008", marks=pytest.mark.full),
]


@pytest.mark.parametrize("shape, dtype", _LERP_TENSOR_BENCH_PARAMS)
def test_lerp_tensor_bench(shape: tuple[int, ...], dtype: torch.dtype) -> None:
    from tileops.perf.formulas import lerp_tensor_fwd_roofline

    test = LerpTensorBenchCase(shape, dtype)
    bm = LerpTensorBenchmark(test)
    a, b, w = test.gen_inputs()

    op = LerpTensorFwdOp(
        input=tuple(shape), end=tuple(shape), weight=tuple(shape), dtype=dtype,
    )
    # Cross-check the bench harness' inline flop/byte counts against the
    # manifest-bound roofline formula so a drift in either direction
    # surfaces as a bench failure rather than silent perf misreporting.
    formula_flops, formula_bytes = lerp_tensor_fwd_roofline(op)
    assert formula_flops == bm.calculate_flops(), (
        f"flop mismatch: formula={formula_flops}, bench={bm.calculate_flops()}"
    )
    assert formula_bytes == bm.calculate_memory(), (
        f"byte mismatch: formula={formula_bytes}, bench={bm.calculate_memory()}"
    )

    result = bm.profile(op, a, b, w)
    BenchmarkReport.record(op, locals(), result, tag="tileops")

    result_bl = bm.profile(torch.lerp, a, b, w)
    BenchmarkReport.record(op, locals(), result_bl, tag="torch")


if __name__ == "__main__":
    pytest.main([__file__, "-vvs"])