test_elementwise_unary_activation_alignment.py

"""Behavior tests for the ``elementwise_unary_activation`` family.

The manifest L1 signature contract is enforced by
``scripts/validate_manifest.py`` for every op family; these tests
exercise activation-specific *behavior* — ``inplace=True`` aliasing
identity, ``approximate`` validation, kernel_map override
dispatch, and end-to-end correctness against the PyTorch reference.
"""

import inspect

import pytest
import torch

from tileops.utils import get_backend_name, is_available

DEVICE = get_backend_name()

_INPLACE_PARAM_FREE_OPS = (
    "ReluFwdOp", "SiluFwdOp", "HardswishFwdOp",
    "HardsigmoidFwdOp", "MishFwdOp", "SeluFwdOp",
)
_INPLACE_PARAMETRIC_OPS = (
    "LeakyReluFwdOp", "EluFwdOp", "HardtanhFwdOp",
)

_CLAMP_OPS = ("ClampFwdOp", "ClampScalarFwdOp", "ClampMinFwdOp", "ClampMaxFwdOp")


def _selu_ref(x: torch.Tensor) -> torch.Tensor:
    return (
        1.0507009873554805
        * torch.where(
            x.float() > 0,
            x.float(),
            1.6732632423543772 * torch.expm1(x.float()),
        )
    ).to(x.dtype)


def _torch_reference(op_name: str):
    """Map an activation op class to its ``torch.nn.functional`` reference."""
    refs = {
        "ReluFwdOp": torch.nn.functional.relu,
        "SiluFwdOp": torch.nn.functional.silu,
        "HardswishFwdOp": torch.nn.functional.hardswish,
        "HardsigmoidFwdOp": torch.nn.functional.hardsigmoid,
        "MishFwdOp": torch.nn.functional.mish,
        "SeluFwdOp": _selu_ref,
        "LeakyReluFwdOp": torch.nn.functional.leaky_relu,
        "EluFwdOp": torch.nn.functional.elu,
        "HardtanhFwdOp": torch.nn.functional.hardtanh,
    }
    return refs[op_name]


def _construct_inplace_op(mod, op_name: str, n_total: int, dtype: torch.dtype, inplace: bool):
    """Build an instance with the manifest-spec construction signature."""
    cls = getattr(mod, op_name)
    if op_name in _INPLACE_PARAM_FREE_OPS:
        return cls(N_total=n_total, dtype=dtype, inplace=inplace)
    return cls(n_total, dtype, inplace=inplace)


def _clamp_construct_kwargs(op_name: str) -> tuple[tuple, dict]:
    """Return ``(positional, keyword)`` args needed to construct ``op_name``."""
    shape = (2, 4)
    if op_name == "ClampFwdOp":
        return ((shape, shape, shape, torch.float16), {})
    if op_name == "ClampScalarFwdOp":
        return ((shape,), {"min": -1.0, "max": 1.0, "dtype": torch.float16})
    return ((shape, shape, torch.float16), {})


@pytest.mark.smoke
@pytest.mark.skipif(not is_available(), reason=f"{DEVICE.upper()} required")
@pytest.mark.parametrize("op_name", _CLAMP_OPS)
def test_clamp_family_kernel_map_override_is_dispatched(op_name: str) -> None:
    """A user-supplied ``kernel_map`` value must reach the kernel build.

    Construct each Clamp op with a ``kernel_map`` whose value is a
    *subclass* of the default kernel and assert the constructed
    ``self.kernel`` is an instance of that subclass — the load-bearing
    invariant is that the override class is the one actually used to
    build ``self.kernel``.
    """
    import tileops.ops.elementwise as mod

    cls = getattr(mod, op_name)
    pos, kw = _clamp_construct_kwargs(op_name)
    inst = cls(*pos, **kw)
    (key, default_kernel_cls), = inst.default_kernel_map.items()

    class MarkerKernel(default_kernel_cls):  # type: ignore[misc, valid-type]
        """Subclass marker; identical behavior, distinct identity."""

    override = {key: MarkerKernel}
    inst2 = cls(*pos, **kw, kernel_map=override)
    assert inst2.kernel_map[key] is MarkerKernel, (
        f"{op_name}: kernel_map override entry was not stored on "
        f"self.kernel_map (got {inst2.kernel_map[key]!r})"
    )
    assert isinstance(inst2.kernel, MarkerKernel), (
        f"{op_name}: kernel_map override class was not used to build "
        f"self.kernel (kernel type: {type(inst2.kernel).__name__})"
    )


@pytest.mark.smoke
def test_nan_to_num_canonical_kwarg_names() -> None:
    """NanToNumFwdOp accepts the manifest-aligned names end-to-end."""
    import tileops.ops.elementwise as mod

    op = mod.NanToNumFwdOp(
        N_total=8, dtype=torch.float16, nan=0.0, posinf=1.0, neginf=-1.0,
    )
    assert op.nan == 0.0
    assert op.posinf == 1.0
    assert op.neginf == -1.0


@pytest.mark.smoke
@pytest.mark.skipif(not is_available(), reason=f"{DEVICE.upper()} required")
@pytest.mark.parametrize(
    "op_name", _INPLACE_PARAM_FREE_OPS + _INPLACE_PARAMETRIC_OPS,
)
def test_unary_activation_inplace_true_aliases_input(op_name: str) -> None:
    """``inplace=True`` must mutate ``input`` and return the same tensor.

    PyTorch's contract for ``functional.relu(x, inplace=True)`` (and the
    other activations declaring ``inplace`` in their manifest entry) is
    that the returned tensor *is* ``x`` and that ``x`` now holds the
    activation output.
    """
    import tileops.ops.elementwise as mod

    n_total = 64
    dtype = torch.float16
    op = _construct_inplace_op(mod, op_name, n_total, dtype, inplace=True)
    x = torch.randn(n_total, dtype=dtype, device=DEVICE)
    expected = _torch_reference(op_name)(x.clone())
    y = op(x)
    assert y is x, (
        f"{op_name}: inplace=True must return the input tensor (identity); "
        f"got id(y)={id(y)} id(x)={id(x)}"
    )
    assert torch.allclose(x, expected, rtol=1e-2, atol=1e-2), (
        f"{op_name}: inplace=True did not mutate input to the activation output"
    )


@pytest.mark.smoke
@pytest.mark.skipif(not is_available(), reason=f"{DEVICE.upper()} required")
@pytest.mark.parametrize(
    "op_name", _INPLACE_PARAM_FREE_OPS + _INPLACE_PARAMETRIC_OPS,
)
def test_unary_activation_inplace_false_returns_fresh_tensor(op_name: str) -> None:
    """Default ``inplace=False`` must not alias or mutate the input."""
    import tileops.ops.elementwise as mod

    n_total = 64
    dtype = torch.float16
    op = _construct_inplace_op(mod, op_name, n_total, dtype, inplace=False)
    x = torch.randn(n_total, dtype=dtype, device=DEVICE)
    x_before = x.clone()
    y = op(x)
    assert y is not x, f"{op_name}: inplace=False must return a fresh tensor"
    assert torch.equal(x, x_before), (
        f"{op_name}: inplace=False must not mutate the input tensor"
    )


@pytest.mark.smoke
def test_gelu_approximate_validation() -> None:
    """GeluFwdOp must reject ``approximate`` values outside the manifest set."""
    import tileops.ops.elementwise as mod

    with pytest.raises(ValueError, match="approximate"):
        mod.GeluFwdOp(N_total=8, dtype=torch.float16, approximate="invalid")


@pytest.mark.smoke
@pytest.mark.skipif(not is_available(), reason=f"{DEVICE.upper()} required")
@pytest.mark.parametrize("approximate", ["none", "tanh"])
def test_gelu_approximate_runs_through_forward(approximate: str) -> None:
    """Both ``approximate='none'`` and ``'tanh'`` must dispatch end-to-end.

    Each mode is checked against ``torch.nn.functional.gelu`` with the
    matching ``approximate`` argument so the kernel selection is
    observable from the op layer.
    """
    import tileops.ops.elementwise as mod

    n_total = 128
    dtype = torch.float16
    op = mod.GeluFwdOp(N_total=n_total, dtype=dtype, approximate=approximate)
    x = torch.randn(n_total, dtype=dtype, device=DEVICE)
    y = op(x)
    expected = torch.nn.functional.gelu(x, approximate=approximate)
    assert y.shape == x.shape
    assert torch.allclose(y, expected, rtol=1e-2, atol=1e-2)


# Frozen ``__init__`` signatures for every unary activation Op. The
# refactor pulling shared ``__init__`` / ``forward`` / ``_eager_forward``
# logic up into a base or mixin must keep these byte-identical, because
# downstream code (tests, benches, codegen) relies on them.
_FROZEN_UNARY_ACTIVATION_SIGNATURES = {
    "ReluFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, inplace: bool = False, *, "
        "strategy: Optional[str] = None, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "SiluFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, inplace: bool = False, *, "
        "strategy: Optional[str] = None, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "HardswishFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, inplace: bool = False, *, "
        "strategy: Optional[str] = None, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "HardsigmoidFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, inplace: bool = False, *, "
        "strategy: Optional[str] = None, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "MishFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, inplace: bool = False, *, "
        "strategy: Optional[str] = None, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "SeluFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, inplace: bool = False, *, "
        "strategy: Optional[str] = None, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "LeakyReluFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, "
        "negative_slope: float = 0.01, inplace: bool = False, *, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "EluFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, "
        "alpha: float = 1.0, inplace: bool = False, *, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "HardtanhFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, "
        "min_val: float = -1.0, max_val: float = 1.0, inplace: bool = False, *, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
    "SoftplusFwdOp": (
        "(self, N_total: int, dtype: torch.dtype, "
        "beta: float = 1.0, threshold: float = 20.0, *, "
        "kernel_map: Optional[Dict[str, tileops.kernels.kernel_base.Kernel]] = None, "
        "tune: bool = False)"
    ),
}


@pytest.mark.smoke
@pytest.mark.parametrize(
    "op_name", sorted(_FROZEN_UNARY_ACTIVATION_SIGNATURES.keys()),
)
def test_unary_activation_init_signature_is_frozen(op_name: str) -> None:
    """All ten unary activation Ops keep their ``__init__`` signature.

    The shared base/mixin refactor must not alter the constructor
    contract (parameter names, defaults, keyword-only-ness) of any
    leaf Op; downstream code (codegen, tests, benches) depends on
    these exact signatures.
    """
    import tileops.ops.elementwise as mod

    cls = getattr(mod, op_name)
    sig = inspect.signature(cls.__init__)
    assert str(sig) == _FROZEN_UNARY_ACTIVATION_SIGNATURES[op_name], (
        f"{op_name}.__init__ signature drifted: "
        f"got {sig!s}, expected {_FROZEN_UNARY_ACTIVATION_SIGNATURES[op_name]}"
    )


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