TileOPs-Metax/tests/ops/test_reduction_defaults.py

"""Regression tests for reduction-op constructor defaults and empty-dim semantics.

Pins two manifest-conformance invariants for the reduction op family:

1. For the ten ops whose manifest declares ``default: null`` on ``dim``
   (Sum/Mean/Amax/Amin/Var/Std/VarMean/All/Any/CountNonzero), constructing
   the op with only ``dtype=`` performs a full reduction (output shape
   equals ``torch.<op>(x).shape``). ``ProdFwdOp`` keeps its documented
   ``dim=-1`` default.

2. ``AllFwdOp`` / ``AnyFwdOp`` honor the spec's ``dim=[]`` / ``dim=()``
   no-op contract: output shape equals the input shape, output dtype is
   ``bool``, and values equal ``x.bool()``.
"""

from __future__ import annotations

import pytest
import torch

pytestmark = pytest.mark.skipif(
    not torch.cuda.is_available(), reason="CUDA required"
)


_FLOAT_SHAPE = (2, 4, 8)
_LOGICAL_SHAPE = (2, 4, 8)


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


def _make_logical(shape: tuple, dtype: torch.dtype) -> torch.Tensor:
    # values in {-1, 0, 1} so .bool() has both T and F.
    return (torch.randint(-1, 2, shape, device="cuda")).to(dtype)


# default dim=None for the ten ops -> full reduction on 3-D input


@pytest.mark.smoke
def test_sum_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import SumFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = SumFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.sum(x).shape


@pytest.mark.smoke
def test_mean_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import MeanFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = MeanFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.mean(x).shape


@pytest.mark.smoke
def test_amax_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import AmaxFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = AmaxFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.amax(x).shape


@pytest.mark.smoke
def test_amin_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import AminFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = AminFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.amin(x).shape


@pytest.mark.smoke
def test_var_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import VarFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = VarFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.var(x).shape


@pytest.mark.smoke
def test_std_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import StdFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = StdFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.std(x).shape


@pytest.mark.smoke
def test_var_mean_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.reduce import VarMeanFwdOp

    x = _make_float(_FLOAT_SHAPE, torch.float16)
    op = VarMeanFwdOp(dtype=torch.float16)
    var_out, mean_out = op(x)
    ref_var, ref_mean = torch.var_mean(x)
    assert var_out.shape == ref_var.shape
    assert mean_out.shape == ref_mean.shape


@pytest.mark.smoke
def test_all_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.logical_reduce import AllFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = AllFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.all(x.bool()).shape
    assert y.dtype == torch.bool


@pytest.mark.smoke
def test_any_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.logical_reduce import AnyFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = AnyFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.any(x.bool()).shape
    assert y.dtype == torch.bool


@pytest.mark.smoke
def test_count_nonzero_default_dim_full_reduction() -> None:
    from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = CountNonzeroFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.count_nonzero(x).shape
    assert y.dtype == torch.int64


# ProdFwdOp keeps documented dim=-1 default


@pytest.mark.smoke
def test_prod_default_dim_last_axis() -> None:
    from tileops.ops.reduction.reduce import ProdFwdOp

    # use a narrow value range so fp16 prod is numerically stable
    x = torch.rand(*_FLOAT_SHAPE, dtype=torch.float16, device="cuda") * 0.01 + 0.99
    op = ProdFwdOp(dtype=torch.float16)
    y = op(x)
    assert y.shape == torch.prod(x, dim=-1).shape


# AllFwdOp/AnyFwdOp dim=[] / dim=() noop contract


@pytest.mark.smoke
@pytest.mark.parametrize("empty_dim", [[], ()])
def test_all_empty_dim_noop(empty_dim) -> None:
    from tileops.ops.reduction.logical_reduce import AllFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = AllFwdOp(dtype=torch.float16, dim=empty_dim)
    y = op(x)
    assert y.shape == x.shape
    assert y.dtype == torch.bool
    assert torch.equal(y, x.bool())


@pytest.mark.smoke
@pytest.mark.parametrize("empty_dim", [[], ()])
def test_any_empty_dim_noop(empty_dim) -> None:
    from tileops.ops.reduction.logical_reduce import AnyFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = AnyFwdOp(dtype=torch.float16, dim=empty_dim)
    y = op(x)
    assert y.shape == x.shape
    assert y.dtype == torch.bool
    assert torch.equal(y, x.bool())


# normalize_dim noop policy returns []


@pytest.mark.smoke
def test_normalize_dim_noop_returns_empty() -> None:
    from tileops.ops.reduction._multidim import normalize_dim

    assert normalize_dim([], ndim=3, empty_dim_policy="noop") == []
    assert normalize_dim((), ndim=3, empty_dim_policy="noop") == []


@pytest.mark.smoke
def test_normalize_dim_reject_raises_on_empty() -> None:
    from tileops.ops.reduction._multidim import normalize_dim

    with pytest.raises(ValueError):
        normalize_dim([], ndim=3, empty_dim_policy="reject")


@pytest.mark.smoke
def test_normalize_dim_full_returns_all() -> None:
    from tileops.ops.reduction._multidim import normalize_dim

    assert normalize_dim([], ndim=3, empty_dim_policy="full") == [0, 1, 2]


@pytest.mark.smoke
def test_empty_dim_policy_class_attrs() -> None:
    """Per-op empty_dim_policy bindings."""
    from tileops.ops.reduction.logical_reduce import AllFwdOp, AnyFwdOp, CountNonzeroFwdOp
    from tileops.ops.reduction.reduce import (
        AmaxFwdOp,
        AminFwdOp,
        MeanFwdOp,
        ProdFwdOp,
        StdFwdOp,
        SumFwdOp,
        VarFwdOp,
        VarMeanFwdOp,
        _ReduceOpBase,
    )

    assert _ReduceOpBase._empty_dim_policy == "reject"
    assert AllFwdOp._empty_dim_policy == "noop"
    assert AnyFwdOp._empty_dim_policy == "noop"
    for cls in (
        SumFwdOp, MeanFwdOp, AmaxFwdOp, AminFwdOp,
        StdFwdOp, VarFwdOp, VarMeanFwdOp, CountNonzeroFwdOp,
    ):
        assert cls._empty_dim_policy == "full", cls.__name__
    # ProdFwdOp inherits default (reject); empty dim is not in its contract
    assert ProdFwdOp._empty_dim_policy == "reject"


# Empty-dim noop must NOT bypass input validation or roofline binding


@pytest.mark.smoke
def test_all_empty_dim_noop_rejects_cpu_tensor() -> None:
    """dim=[] must still validate device; non-CUDA input must raise."""
    from tileops.ops.reduction.logical_reduce import AllFwdOp

    x = (torch.randint(-1, 2, _LOGICAL_SHAPE)).to(torch.float16)  # cpu
    op = AllFwdOp(dtype=torch.float16, dim=[])
    with pytest.raises(ValueError, match="CUDA tensor"):
        op(x)


@pytest.mark.smoke
def test_any_empty_dim_noop_rejects_cpu_tensor() -> None:
    from tileops.ops.reduction.logical_reduce import AnyFwdOp

    x = (torch.randint(-1, 2, _LOGICAL_SHAPE)).to(torch.float16)  # cpu
    op = AnyFwdOp(dtype=torch.float16, dim=[])
    with pytest.raises(ValueError, match="CUDA tensor"):
        op(x)


@pytest.mark.smoke
def test_all_empty_dim_noop_rejects_wrong_dtype() -> None:
    """dim=[] must still validate dtype against the op's declared dtype."""
    from tileops.ops.reduction.logical_reduce import AllFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float32)  # cuda, fp32
    op = AllFwdOp(dtype=torch.float16, dim=[])
    with pytest.raises(ValueError, match="Expected x.dtype"):
        op(x)


@pytest.mark.smoke
def test_any_empty_dim_noop_rejects_wrong_dtype() -> None:
    from tileops.ops.reduction.logical_reduce import AnyFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float32)
    op = AnyFwdOp(dtype=torch.float16, dim=[])
    with pytest.raises(ValueError, match="Expected x.dtype"):
        op(x)


@pytest.mark.smoke
def test_all_empty_dim_noop_binds_roofline() -> None:
    """eval_roofline() must succeed after a dim=[] noop forward and
    report non-zero data-movement (the noop still reads the input and
    writes an equal-shape cast result)."""
    from tileops.ops.reduction.logical_reduce import AllFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = AllFwdOp(dtype=torch.float16, dim=[])
    op(x)
    flops, mem_bytes = op.eval_roofline()
    numel = x.numel()
    elem_bytes = x.element_size()
    # Noop binds (M=numel, N=1); for the "all" op_kind this gives
    # mem_bytes = numel * elem_bytes + numel (input read + bool write).
    expected_lower = numel * elem_bytes
    expected_upper = 2 * numel * elem_bytes + numel
    assert mem_bytes >= expected_lower, (
        f"noop bandwidth {mem_bytes} under-counts input read "
        f"({expected_lower} bytes)"
    )
    assert mem_bytes <= expected_upper
    # flops are degenerate (one op per element); contract is non-negative.
    assert flops >= 0


@pytest.mark.smoke
def test_any_empty_dim_noop_binds_roofline() -> None:
    from tileops.ops.reduction.logical_reduce import AnyFwdOp

    x = _make_logical(_LOGICAL_SHAPE, torch.float16)
    op = AnyFwdOp(dtype=torch.float16, dim=[])
    op(x)
    flops, mem_bytes = op.eval_roofline()
    numel = x.numel()
    elem_bytes = x.element_size()
    expected_lower = numel * elem_bytes
    expected_upper = 2 * numel * elem_bytes + numel
    assert mem_bytes >= expected_lower
    assert mem_bytes <= expected_upper
    assert flops >= 0

@pytest.mark.smoke
def test_validate_dim_rejects_bool_scalar() -> None:
    """`bool` subclasses `int`, but a boolean dim is never a valid axis;
    `_validate_dim` must reject it explicitly."""
    from tileops.ops.reduction.reduce import SumFwdOp

    with pytest.raises(TypeError, match="dim must not be bool"):
        SumFwdOp(dtype=torch.float16, dim=True)


@pytest.mark.smoke
def test_validate_dim_rejects_bool_in_list() -> None:
    """Same guard applies element-wise to `list[int]` / `tuple[int, ...]`."""
    from tileops.ops.reduction.reduce import SumFwdOp

    with pytest.raises(TypeError, match="must be int .not bool"):
        SumFwdOp(dtype=torch.float16, dim=[True, 0])