bench_activation.py

"""Benchmarks for unary activation ops covering risk points R2-R7.

Risk points covered:
- R2: Divmod overhead on small tensors (relu x fp16 x 4K)
- R3: JIT compilation cost (relu x 10 different N, cold vs warm)
- R4: DEFAULT_STRATEGY confirmation (relu x 3 strategies x 3 dtypes x 3 sizes)
- R5: Boundary auto-guard tail vectorization (relu x aligned/unaligned sizes)
- R6: threads=256 vs 128 for complex ops (relu/erf/mish x thread configs)
- R7: dtype-aware num_per_thread (relu x fp32/fp16 x npt=4/npt=8)

Profiles all 3 strategies (direct, explicit_parallel, register_copy) and
compares against PyTorch baseline to determine optimal DEFAULT_STRATEGY.
"""

from math import prod
from typing import Optional, Protocol

import pytest
import torch

from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.kernels.elementwise import (
    ErfFwdKernel,
    MishFwdKernel,
    ReluFwdKernel,
    _make_unary_explicit,
)
from tileops.ops.elementwise import (
    EluFwdOp,
    ErfFwdOp,
    GeluFwdOp,
    HardsigmoidFwdOp,
    HardswishFwdOp,
    HardtanhFwdOp,
    LeakyReluFwdOp,
    MishFwdOp,
    ReluFwdOp,
    SeluFwdOp,
    SiluFwdOp,
    SoftplusFwdOp,
)
from workloads.activation import ReluTest
from workloads.workload_base import FixtureBase

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

_SHAPES_2D = [
    (1, 4096),         # 4K  -- single-token small
    (1024, 4096),      # 4M  -- small transformer hidden dim
    (1024, 11008),     # ~11M -- non-pow2 LLaMA-7B intermediate
]
_SIZE_LABELS = ("4K", "4M", "11M")
_SHAPE_BY_LABEL = dict(zip(_SIZE_LABELS, _SHAPES_2D, strict=True))

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


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


class _UnaryWorkload(Protocol):
    """Structural type for unary benchmark workloads."""

    shape: tuple[int, ...]
    n_total: int
    dtype: torch.dtype

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


class UnaryBenchCase:
    """Minimal test harness for unary benchmarks.

    Accepts either a shape tuple or a scalar element count. The tuple form
    is preferred so the original input geometry survives into the report.
    """

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

    def gen_inputs(self) -> tuple[torch.Tensor]:
        return (torch.randn(*self.shape, device="cuda", dtype=self.dtype),)


class UnaryBenchmark(BenchmarkBase[_UnaryWorkload]):
    """Bandwidth-oriented benchmark for unary elementwise ops.

    Constructed with the Op instance whose roofline drives the TFLOPs
    column: the bench reads ``op.eval_roofline()`` (or
    ``FLOPS_PER_ELEM`` as a fallback for ops that have not implemented
    ``eval_roofline``) so the reported TFLOPs reflects the manifest
    ``roofline.flops`` contract end-to-end. The R6 / R7 kernel-direct
    benches build a transient Op of the corresponding family and pass
    it here; there is no separate ``flops_per_elem`` override.
    """

    def __init__(
        self,
        workload: _UnaryWorkload,
        *,
        op: object,
    ):
        super().__init__(workload)
        self._op = op

    def calculate_flops(self) -> Optional[float]:
        eval_fn = getattr(self._op, "eval_roofline", None)
        if eval_fn is not None:
            # The base ``Op.eval_roofline`` raises ``NotImplementedError``
            # so subclasses must opt in. Fall through to the per-elem
            # path when the op inherits the unimplemented stub rather
            # than crashing the bench.
            try:
                flops, _ = eval_fn()
            except NotImplementedError:
                pass
            else:
                return float(flops)
        per_elem = getattr(self._op, "FLOPS_PER_ELEM", None)
        if per_elem is not None:
            return float(per_elem) * self.workload.n_total
        return self.workload.n_total

    def calculate_memory(self) -> Optional[float]:
        t = self.workload
        out_dtype = getattr(t, "output_dtype", t.dtype)
        return t.n_total * (t.dtype.itemsize + out_dtype.itemsize)


# ---------------------------------------------------------------------------
# R2: Divmod overhead on small tensors
# ---------------------------------------------------------------------------


class R2SmallTensorFixture(FixtureBase):
    PARAMS = [
        ("shape, dtype", [
            pytest.param((1, 4096), torch.float16, marks=pytest.mark.smoke),
        ]),
    ]


@R2SmallTensorFixture
def test_r2_small_tensor_unary(shape: tuple[int, ...], dtype: torch.dtype) -> None:
    """R2: Benchmark divmod overhead on small tensors (unary relu, 4K)."""
    test = UnaryBenchCase(shape, dtype)
    inputs = test.gen_inputs()

    n_total = prod(shape)
    op = ReluFwdOp(N_total=n_total, dtype=dtype)
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(op, locals(), result, tag="tileops")

    def baseline_fn(x):
        return torch.relu(x)

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


# ---------------------------------------------------------------------------
# R3: JIT compilation cost
# ---------------------------------------------------------------------------


# R3 uses 1D shapes whose total element count varies per case; the goal is
# to measure JIT compile cost as a function of N, so we keep a 1D layout
# but record the shape tuple so the report stays consistent.
_R3_SHAPES = [
    (1_000,), (2_000,), (4_000,), (8_000,), (16_000,),
    (32_000,), (64_000,), (128_000,), (256_000,), (512_000,),
]


class R3JitFixture(FixtureBase):
    PARAMS = [
        ("shape", [
            pytest.param(s, marks=pytest.mark.full) for s in _R3_SHAPES
        ]),
    ]


@R3JitFixture
def test_r3_jit_compilation_cost(shape: tuple[int, ...]) -> None:
    """R3: Benchmark JIT compilation cost — relu with 10 different N values.

    Each test case creates a new kernel (different N -> different codegen),
    measuring both first-call (cold JIT) and subsequent (warm) latency.
    """
    import time

    dtype = torch.float16
    n_total = prod(shape)
    x = torch.randn(*shape, device="cuda", dtype=dtype)

    # Cold: time the first call including JIT compilation
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    op = ReluFwdOp(N_total=n_total, dtype=dtype)
    _ = op(x)
    torch.cuda.synchronize()
    cold_ms = (time.perf_counter() - t0) * 1000.0

    # Warm: profile subsequent calls
    test = UnaryBenchCase(shape, dtype)
    bm = UnaryBenchmark(test, op=op)
    warm_result = bm.profile(op, x)

    BenchmarkReport.record(
        "r3_jit_cost",
        {"shape": shape, "cold_ms": round(cold_ms, 2)},
        warm_result,
        tag="relu_jit",
    )


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


_R4_PARAMS = []
for size_label, _shape in _SHAPE_BY_LABEL.items():
    for dt in _DTYPES:
        for strategy in _UNARY_STRATEGIES:
            mark = pytest.mark.smoke if (
                size_label == "4M" and dt == torch.float16
                and strategy == "register_copy"
            ) else pytest.mark.full
            _R4_PARAMS.append(
                pytest.param(
                    _shape, size_label, dt, strategy,
                    id=f"{size_label}-{dt}-{strategy}",
                    marks=mark,
                )
            )


class R4StrategyFixture(FixtureBase):
    PARAMS = [
        ("shape, size_label, dtype, strategy", _R4_PARAMS),
    ]


@R4StrategyFixture
def test_r4_default_strategy_unary(
    shape: tuple[int, ...],
    size_label: str,
    dtype: torch.dtype,
    strategy: str,
) -> None:
    """R4: Benchmark all 3 unary strategies to confirm DEFAULT_STRATEGY."""
    test = UnaryBenchCase(shape, dtype)
    inputs = test.gen_inputs()

    n_total = prod(shape)
    op = ReluFwdOp(N_total=n_total, dtype=dtype, strategy=strategy)
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r4_strategy_unary",
        {"shape": shape, "size_label": size_label,
         "dtype": dtype, "strategy": strategy},
        result,
        tag=f"relu_{strategy}",
    )


# Also benchmark gelu to verify strategy choice holds for transcendental ops.
# Both ``approximate`` modes (``none`` -> erf-based; ``tanh`` -> tanh
# approximation) are exercised so the bench captures kernel selection.
_R4_GELU_PARAMS = []
for size_label, _shape in _SHAPE_BY_LABEL.items():
    for dt in _DTYPES:
        for strategy in _UNARY_STRATEGIES:
            for approximate in ("none", "tanh"):
                _R4_GELU_PARAMS.append(
                    pytest.param(
                        _shape, size_label, dt, strategy, approximate,
                        id=f"gelu-{approximate}-{size_label}-{dt}-{strategy}",
                        marks=pytest.mark.full,
                    )
                )


class R4GeluStrategyFixture(FixtureBase):
    PARAMS = [
        ("shape, size_label, dtype, strategy, approximate", _R4_GELU_PARAMS),
    ]


@R4GeluStrategyFixture
def test_r4_default_strategy_gelu(
    shape: tuple[int, ...],
    size_label: str,
    dtype: torch.dtype,
    strategy: str,
    approximate: str,
) -> None:
    """R4: Benchmark gelu strategies (transcendental op) to confirm DEFAULT_STRATEGY."""
    test = UnaryBenchCase(shape, dtype)
    inputs = test.gen_inputs()

    n_total = prod(shape)
    op = GeluFwdOp(
        N_total=n_total, dtype=dtype, strategy=strategy,
        approximate=approximate,
    )
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r4_strategy_gelu",
        {"shape": shape, "size_label": size_label,
         "dtype": dtype, "strategy": strategy, "approximate": approximate},
        result,
        tag=f"gelu_{approximate}_{strategy}",
    )


# ---------------------------------------------------------------------------
# R5: Boundary auto-guard tail vectorization
# ---------------------------------------------------------------------------


# npt=8, threads=256 -> block_size=2048
_BLOCK_SIZE = 256 * 8
_R5_SHAPES = [
    ((_BLOCK_SIZE * 1000,), "aligned"),                   # perfectly aligned
    ((_BLOCK_SIZE * 1000 + 1,), "unaligned_plus_1"),      # minimal tail
    ((_BLOCK_SIZE * 1000 + 127,), "unaligned_plus_127"),  # large partial tail
]


class R5BoundaryFixture(FixtureBase):
    PARAMS = [
        ("shape, align_label", [
            pytest.param(s, label, marks=pytest.mark.full)
            for s, label in _R5_SHAPES
        ]),
    ]


@R5BoundaryFixture
def test_r5_boundary_guard(shape: tuple[int, ...], align_label: str) -> None:
    """R5: Benchmark boundary auto-guard tail vectorization.

    Compares aligned vs unaligned sizes under explicit_parallel strategy
    to detect performance cliff from boundary guard overhead.
    """
    dtype = torch.float16
    test = UnaryBenchCase(shape, dtype)
    inputs = test.gen_inputs()

    n_total = prod(shape)
    op = ReluFwdOp(N_total=n_total, dtype=dtype, strategy="explicit_parallel")
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(
        "r5_boundary",
        {"shape": shape, "align_label": align_label},
        result,
        tag=f"relu_{align_label}",
    )


# ---------------------------------------------------------------------------
# R6: threads=256 vs 128 for complex ops
# ---------------------------------------------------------------------------


_R6_KERNEL_OPS = [
    ("relu", ReluFwdKernel),
    ("erf", ErfFwdKernel),
    ("mish", MishFwdKernel),
]


# Op classes corresponding to the kernel-direct benches below. Each
# R6 / R7 case constructs a transient instance and hands it to
# ``UnaryBenchmark`` so the TFLOPs column reads ``op.eval_roofline()``
# (or ``FLOPS_PER_ELEM``) — no manifest expression parsing in the bench.
_R6_OP_BY_NAME = {
    "relu": ReluFwdOp,
    "erf": ErfFwdOp,
    "mish": MishFwdOp,
}

_R6_THREADS = [128, 256]

_R6_PARAMS = []
for size_label, _shape in _SHAPE_BY_LABEL.items():
    for op_name, _ in _R6_KERNEL_OPS:
        for threads in _R6_THREADS:
            mark = pytest.mark.smoke if (
                size_label == "4M" and op_name == "relu" and threads == 256
            ) else pytest.mark.full
            _R6_PARAMS.append(
                pytest.param(
                    _shape, size_label, op_name, threads,
                    id=f"{op_name}-{size_label}-t{threads}",
                    marks=mark,
                )
            )


class R6ThreadsFixture(FixtureBase):
    PARAMS = [
        ("shape, size_label, op_name, threads", _R6_PARAMS),
    ]


_R6_KERNEL_MAP = {name: cls for name, cls in _R6_KERNEL_OPS}


@R6ThreadsFixture
def test_r6_threads_comparison(
    shape: tuple[int, ...],
    size_label: str,
    op_name: str,
    threads: int,
) -> None:
    """R6: Benchmark threads=256 vs 128 for simple and complex ops.

    Complex ops (erf, mish) may benefit from fewer threads (128) due to
    higher register pressure. Simple ops (relu) should prefer 256.

    Builds kernels directly via _make_unary_explicit to ensure block_size
    is baked with the requested threads/npt at build time, not overridden
    after construction.
    """
    dtype = torch.float16
    dtype_str = "float16"
    test = UnaryBenchCase(shape, dtype)
    n_total = prod(shape)
    # Build a transient Op of the same family; ``UnaryBenchmark`` reads
    # ``op.eval_roofline()`` to drive the TFLOPs column so the kernel-
    # direct R6 path stays aligned with the manifest coefficient
    # without parsing roofline expressions in the bench file.
    flops_op = _R6_OP_BY_NAME[op_name](N_total=n_total, dtype=dtype)
    bm = UnaryBenchmark(test, op=flops_op)
    inputs = test.gen_inputs()

    npt = 8  # default for fp16
    kernel_cls = _R6_KERNEL_MAP[op_name]
    # Build kernel directly with the desired threads/npt so block_size is correct
    kernel_fn = _make_unary_explicit(
        n_total, dtype_str, kernel_cls.op_func, threads=threads, num_per_thread=npt,
    )
    # The explicit-parallel kernel expects a 1D contiguous tensor, so flatten
    # the (possibly multi-dim) input here. The shape tuple is still recorded
    # via ``BenchmarkReport.record(...)`` so the report carries the original
    # input geometry.
    flat_inputs = tuple(t.reshape(-1) for t in inputs)
    # Profile: call the JIT kernel with matching runtime args
    compiled = kernel_fn(threads, npt)
    result = bm.profile(compiled, *flat_inputs)
    BenchmarkReport.record(
        "r6_threads",
        {"shape": shape, "size_label": size_label,
         "op_name": op_name, "threads": threads},
        result,
        tag=f"{op_name}_t{threads}",
    )


# ---------------------------------------------------------------------------
# R7: dtype-aware num_per_thread
# ---------------------------------------------------------------------------


_R7_PARAMS = []
for dt, dt_label in [(torch.float32, "fp32"), (torch.float16, "fp16")]:
    for npt in [4, 8]:
        _R7_PARAMS.append(
            pytest.param(
                dt, dt_label, npt,
                id=f"{dt_label}-npt{npt}",
                marks=pytest.mark.full,
            )
        )


class R7NptFixture(FixtureBase):
    PARAMS = [
        ("dtype, dtype_label, num_per_thread", _R7_PARAMS),
    ]


@R7NptFixture
def test_r7_dtype_npt(
    dtype: torch.dtype,
    dtype_label: str,
    num_per_thread: int,
) -> None:
    """R7: Benchmark dtype-aware num_per_thread.

    Compares npt=4 vs npt=8 for fp32 and fp16 to validate whether the
    current default (fp32->4, fp16->8) is optimal. If no difference,
    simplify to fixed npt=8.

    Builds kernels directly via _make_unary_explicit to ensure block_size
    is baked with the requested npt at build time.
    """
    shape = (1_000_000,)
    n_total = prod(shape)
    threads = 256
    dtype_str = "float32" if dtype == torch.float32 else "float16"
    test = UnaryBenchCase(shape, dtype)
    # Transient Op carries the manifest roofline; bench reads it.
    flops_op = ReluFwdOp(N_total=n_total, dtype=dtype)
    bm = UnaryBenchmark(test, op=flops_op)
    inputs = test.gen_inputs()

    # Build kernel directly with the desired threads/npt so block_size is correct
    kernel_fn = _make_unary_explicit(
        n_total, dtype_str, ReluFwdKernel.op_func,
        threads=threads, num_per_thread=num_per_thread,
    )
    compiled = kernel_fn(threads, num_per_thread)
    result = bm.profile(compiled, *inputs)
    BenchmarkReport.record(
        "r7_dtype_npt",
        {"shape": shape, "dtype_label": dtype_label,
         "num_per_thread": num_per_thread},
        result,
        tag=f"relu_{dtype_label}_npt{num_per_thread}",
    )


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


_RELU_BENCH_PARAMS = [
    pytest.param(_SHAPES_2D[1], torch.float16, id="throughput-fp16"),
    pytest.param(_SHAPES_2D[1], torch.bfloat16, id="throughput-bf16"),
    pytest.param(_SHAPES_2D[1], torch.float32, id="baseline-fp32"),
]


@pytest.mark.parametrize("shape, dtype", _RELU_BENCH_PARAMS)
def test_relu_bench(shape: tuple[int, ...], dtype: torch.dtype) -> None:
    n_total = prod(shape)
    # ``ReluTest`` (workloads) accepts a flat element count; the bench
    # harness still records the original shape tuple via ``record(...)``.
    test = ReluTest(n_total, dtype)
    inputs = test.gen_inputs()

    op = ReluFwdOp(N_total=n_total, dtype=dtype)
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(op, locals(), result, tag="tileops")

    def baseline_fn(x):
        return torch.relu(x)

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


# ---------------------------------------------------------------------------
# Throughput coverage for the param-free unary activations declared in
# tileops/manifest/elementwise_unary_activation.yaml. Each op gets a single
# fp16 case at the LLaMA hidden-dim shape so the bench file produces a
# number for downstream perf tracking without inflating CI runtime.
# ---------------------------------------------------------------------------


_PARAM_FREE_ACTIVATION_OPS = [
    pytest.param(SiluFwdOp, "silu", torch.nn.functional.silu, id="silu"),
    pytest.param(
        HardswishFwdOp, "hardswish", torch.nn.functional.hardswish, id="hardswish",
    ),
    pytest.param(
        HardsigmoidFwdOp, "hardsigmoid", torch.nn.functional.hardsigmoid,
        id="hardsigmoid",
    ),
    pytest.param(MishFwdOp, "mish", torch.nn.functional.mish, id="mish"),
    pytest.param(SeluFwdOp, "selu", torch.nn.functional.selu, id="selu"),
]


@pytest.mark.parametrize("op_cls, op_label, torch_ref", _PARAM_FREE_ACTIVATION_OPS)
@pytest.mark.parametrize("dtype", [torch.float16])
def test_param_free_unary_bench(
    op_cls, op_label: str, torch_ref, dtype: torch.dtype,
) -> None:
    """Throughput bench for param-free unary activations (manifest-aligned).

    One representative shape x fp16 keeps each op covered without
    expanding the matrix. Records both the TileOps op and the matching
    ``torch.nn.functional`` reference so each row has an external
    baseline per the benchmark contract.
    """
    shape = _SHAPES_2D[1]  # (1024, 4096), LLaMA hidden dim
    n_total = prod(shape)
    test = UnaryBenchCase(shape, dtype)
    inputs = test.gen_inputs()

    op = op_cls(N_total=n_total, dtype=dtype)
    # Pass ``op=`` so UnaryBenchmark reads ``op.eval_roofline()`` /
    # ``FLOPS_PER_ELEM`` and reports manifest-aligned TFLOPs (e.g. SiLU
    # 4*N, Mish 7*N, SELU 5*N) rather than the bandwidth-only 1*N
    # default. The torch baseline is profiled with the same harness so
    # both rows share the same FLOP normalization.
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(op, locals(), result, tag="tileops")

    def baseline_fn(x):
        return torch_ref(x)

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


# ---------------------------------------------------------------------------
# Throughput coverage for the parametric unary activations (leaky_relu, elu,
# hardtanh, softplus). These ops take scalar constructor params and need
# matching kwargs on the torch baseline so the two rows compute the same
# function. Layout mirrors the param-free block: one fp16 case at the
# LLaMA hidden-dim shape per op.
# ---------------------------------------------------------------------------


def _leaky_relu_baseline(x: torch.Tensor) -> torch.Tensor:
    return torch.nn.functional.leaky_relu(x, negative_slope=0.01)


def _elu_baseline(x: torch.Tensor) -> torch.Tensor:
    return torch.nn.functional.elu(x, alpha=1.0)


def _hardtanh_baseline(x: torch.Tensor) -> torch.Tensor:
    return torch.nn.functional.hardtanh(x, min_val=-1.0, max_val=1.0)


def _softplus_baseline(x: torch.Tensor) -> torch.Tensor:
    return torch.nn.functional.softplus(x, beta=1.0, threshold=20.0)


_PARAMETRIC_ACTIVATION_OPS = [
    pytest.param(
        LeakyReluFwdOp, "leaky_relu",
        {"negative_slope": 0.01},
        _leaky_relu_baseline,
        id="leaky_relu",
    ),
    pytest.param(
        EluFwdOp, "elu",
        {"alpha": 1.0},
        _elu_baseline,
        id="elu",
    ),
    pytest.param(
        HardtanhFwdOp, "hardtanh",
        {"min_val": -1.0, "max_val": 1.0},
        _hardtanh_baseline,
        id="hardtanh",
    ),
    pytest.param(
        SoftplusFwdOp, "softplus",
        {"beta": 1.0, "threshold": 20.0},
        _softplus_baseline,
        id="softplus",
    ),
]


@pytest.mark.parametrize(
    "op_cls, op_label, op_kwargs, torch_ref", _PARAMETRIC_ACTIVATION_OPS,
)
@pytest.mark.parametrize("dtype", [torch.float16])
def test_parametric_unary_bench(
    op_cls,
    op_label: str,
    op_kwargs: dict,
    torch_ref,
    dtype: torch.dtype,
) -> None:
    """Throughput bench for parametric unary activations (manifest-aligned).

    Each op is constructed with its default scalar params and the torch
    baseline is wrapped with the same kwargs so both rows compute the
    same function. One representative shape x fp16 keeps the matrix
    small while still producing a number per touched op.
    """
    shape = _SHAPES_2D[1]  # (1024, 4096), LLaMA hidden dim
    n_total = prod(shape)
    test = UnaryBenchCase(shape, dtype)
    inputs = test.gen_inputs()

    op = op_cls(N_total=n_total, dtype=dtype, **op_kwargs)
    bm = UnaryBenchmark(test, op=op)
    result = bm.profile(op, *inputs)
    BenchmarkReport.record(op, locals(), result, tag="tileops")

    def baseline_fn(x):
        return torch_ref(x)

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


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