TileOPs-Metax/tests/test_op_base.py

"""Tests for tileops.ops.op_base.

Covers ``Op._cache_key`` default behavior, the runtime warning fired when
a subclass with empty ``_static_axes`` does not override ``_cache_key``,
composite kernel-map overrides, and ``Op.autotune`` kernel discovery.
"""

import warnings

import pytest

from tileops.kernels.kernel_base import Kernel
from tileops.ops import op_base
from tileops.ops.op_base import Op

pytestmark = pytest.mark.smoke


@pytest.fixture(autouse=True)
def _reset_warned_types():
    """Clear the module-level dedup set so each test sees a fresh warn state."""
    op_base._EMPTY_STATIC_DIMS_WARNED.clear()
    yield
    op_base._EMPTY_STATIC_DIMS_WARNED.clear()


def _make_op_subclass(*, static_axes=frozenset(), override_cache_key=False):
    """Build a minimal concrete Op subclass for testing.

    ``static_axes`` populates ``_static_axes``.
    ``override_cache_key=True`` attaches a subclass-level override.
    """
    attrs = {
        "_static_axes": static_axes,
        "default_kernel_map": property(lambda self: {}),
        "forward": lambda self, *a, **kw: None,
    }
    if override_cache_key:
        attrs["_cache_key"] = lambda self, *shapes: ("overridden",)
    return type("TestOp", (Op,), attrs)


class TestCacheKeyDefault:
    def test_static_axes_exclude_single_input(self):
        """_static_axes=[(0,1)] on a 3D input excludes axis 1 from the key."""
        Cls = _make_op_subclass(static_axes=frozenset({(0, 1)}))
        op = Cls()
        key = op._cache_key((2, 4, 8))
        assert key == (2, 8)

    def test_static_axes_across_multiple_inputs(self):
        """_static_axes can reference axes in different input positions."""
        Cls = _make_op_subclass(static_axes=frozenset({(0, 1), (1, 0)}))
        op = Cls()
        key = op._cache_key((2, 4, 8), (16, 32))
        # Input 0: exclude axis 1 -> (2, 8); Input 1: exclude axis 0 -> (32,)
        assert key == (2, 8, 32)

    def test_empty_static_axes_returns_full_shape(self):
        """With no static axes, the key concatenates all input shape values."""
        Cls = _make_op_subclass(static_axes=frozenset())
        op = Cls()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")  # warning tested separately
            key = op._cache_key((2, 4, 8), (3, 5))
        assert key == (2, 4, 8, 3, 5)


class TestCacheKeyWarning:
    def test_empty_static_axes_warns_once_per_type(self):
        """Default path with empty _static_axes warns exactly once per subclass,
        even across multiple instances and repeated calls."""
        Cls = _make_op_subclass(static_axes=frozenset())

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            Cls()._cache_key((2, 4))
            Cls()._cache_key((3, 5))
            Cls()._cache_key((7, 9))
            Cls()._cache_key((11, 13))

        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
        assert len(user_warnings) == 1
        assert "TestOp" in str(user_warnings[0].message)
        assert "_cache_key" in str(user_warnings[0].message)

    def test_override_suppresses_warning(self):
        """When the subclass overrides _cache_key, no warning fires."""
        Cls = _make_op_subclass(
            static_axes=frozenset(), override_cache_key=True
        )

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            result = Cls()._cache_key((2, 4))

        assert result == ("overridden",)
        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
        assert user_warnings == []

    def test_populated_static_axes_suppresses_warning(self):
        """Non-empty _static_axes means the user committed at ctor; no warning
        fires regardless of whether _cache_key was overridden."""
        Cls = _make_op_subclass(static_axes=frozenset({(0, 0)}))

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            Cls()._cache_key((2, 4))

        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
        assert user_warnings == []

    def test_distinct_subclasses_each_warn_once(self):
        """Two different subclasses each warn once; the dedup set is keyed by
        type, not globally suppressed after the first warning."""
        ClsA = _make_op_subclass(static_axes=frozenset())
        ClsB = _make_op_subclass(static_axes=frozenset())

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            ClsA()._cache_key((1,))
            ClsA()._cache_key((2,))  # no re-warn for A
            ClsB()._cache_key((3,))  # fresh warn for B
            ClsB()._cache_key((4,))  # no re-warn for B

        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
        assert len(user_warnings) == 2


class TestCompositeKernelMapOverride:
    """Composite ops (empty ``default_kernel_map``) accept a non-empty override and store it verbatim."""

    def test_empty_default_with_empty_override_yields_empty_map(self):
        Cls = _make_op_subclass()
        op = Cls()
        op.dispatch_kernel(None)
        assert op.kernel_map == {}

    def test_empty_default_with_non_empty_override_stores_override(self):
        Cls = _make_op_subclass()
        op = Cls()
        override = {"permute_nopad_kernel": object(), "unpermute_kernel": object()}
        op.dispatch_kernel(override)
        assert op.kernel_map == override

    def test_empty_default_override_is_copied_not_aliased(self):
        Cls = _make_op_subclass()
        op = Cls()
        override = {"permute_nopad_kernel": object()}
        op.dispatch_kernel(override)
        override["extra"] = object()
        assert "extra" not in op.kernel_map


class TestAutotune:
    """``Op.autotune`` tunes every Kernel-typed attribute and nothing else."""

    def test_autotune_dispatches_to_each_kernel_attribute(self):
        tuned: list[str] = []

        class FakeKernel(Kernel):
            def __init__(self, name):
                super().__init__()
                self.name = name

            def forward(self):
                return None

            def autotune(self, warmup=25, rep=50):
                tuned.append(self.name)

        class FakeOp(Op):
            def __init__(self):
                self.k1 = FakeKernel("k1")
                self.k2 = FakeKernel("k2")
                self.not_a_kernel = object()

            def forward(self, *a, **kw):
                return None

            @property
            def default_kernel_map(self):
                return {}

        FakeOp().autotune()
        assert sorted(tuned) == ["k1", "k2"]