FrRay
/
TileOPs-Metax
forked from ccf-ai-infra/TileOPs-Metax
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
TileOPs-Metax/benchmarks/ops/bench_binary_arith.py

504 lines
16 KiB
Python
Raw Blame History

"""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="cuda", dtype=self.dtype)
b = torch.randn(self.b_shape, device="cuda", 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="cuda", dtype=torch.bool)
x = torch.randn(*self.shape, device="cuda", dtype=self.dtype)
y = torch.randn(*self.shape, device="cuda", 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="cuda", dtype=self.dtype)
b = torch.randn(self.shape, device="cuda", dtype=self.dtype)
# Keep weight in [0, 1] to stay close to typical lerp usage.
w = torch.rand(self.shape, device="cuda", 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"])
Reference in New Issue View Git Blame Copy Permalink