tests/test_nix_bundle_qualification.py

"""Qualification harness: load causal-conv1d kernel from Nix store bundle, assert provenance, run boundary tests."""

from __future__ import annotations

import json
import sys
from datetime import UTC, datetime
from pathlib import Path

import pytest

RESULT_DIR = Path(__file__).resolve().parents[1] / "result"


@pytest.fixture(scope="module")
def cc_bundle():
    """Load the kernel from the Nix store result/ symlink."""
    kernels = pytest.importorskip("kernels")
    assert RESULT_DIR.is_dir(), f"result/ not found at {RESULT_DIR}"
    cc = kernels.get_local_kernel(RESULT_DIR.resolve(), backend="rocm")
    return cc


@pytest.fixture(scope="module")
def causal_conv1d_fn(cc_bundle):
    """Extract the causal_conv1d_fn from the bundle."""
    return cc_bundle.causal_conv1d_fn


@pytest.fixture(scope="module")
def causal_conv1d_update(cc_bundle):
    """Extract the causal_conv1d_update from the bundle."""
    return cc_bundle.causal_conv1d_update


class TestNixStoreProvenance:
    """Fail-hard: kernel must come from the immutable Nix store, not editable source."""

    def test_module_file_is_in_nix_store(self, cc_bundle):
        module_path = Path(cc_bundle.__file__).resolve()
        assert "/nix/store/" in str(module_path), (
            f"Kernel loaded from editable source, not Nix store: {module_path}"
        )

    def test_module_parent_is_torch_rocm(self, cc_bundle):
        module_path = Path(cc_bundle.__file__).resolve()
        assert module_path.parent.name == "torch-rocm", (
            f"Expected torch-rocm directory, got: {module_path.parent.name}"
        )

    def test_module_name_has_rocm_suffix(self, cc_bundle):
        name = cc_bundle.__name__
        assert "rocm" in name.lower(), f"Module name missing rocm suffix: {name}"

    def test_not_from_torch_ext_editable(self, cc_bundle):
        module_path = Path(cc_bundle.__file__).resolve()
        path_str = str(module_path)
        assert "/nix/store/" in path_str or "torch-ext/" not in path_str, (
            f"Kernel loaded from editable torch-ext, not Nix bundle: {module_path}"
        )


class TestCausalConv1dAPIBoundary:
    """Verify the causal_conv1d_fn and causal_conv1d_update API signatures."""

    def test_causal_conv1d_fn_signature(self, causal_conv1d_fn):
        import inspect

        params = tuple(inspect.signature(causal_conv1d_fn).parameters)
        expected = (
            "x", "weight", "bias", "seq_idx", "initial_states",
            "return_final_states", "final_states_out", "activation",
        )
        assert params == expected, f"Signature mismatch: {params}"

    def test_causal_conv1d_update_signature(self, causal_conv1d_update):
        import inspect

        params = tuple(inspect.signature(causal_conv1d_update).parameters)
        expected = (
            "x", "conv_state", "weight", "bias", "activation",
            "cache_seqlens", "conv_state_indices",
        )
        assert params == expected, f"Signature mismatch: {params}"

    def test_causal_conv1d_fn_rejects_cpu(self, causal_conv1d_fn):
        """Triton kernel must reject CPU tensors — no silent fallback in the bundle."""
        torch = pytest.importorskip("torch")
        x = torch.zeros((1, 4, 8), dtype=torch.float32)
        weight = torch.zeros((4, 4), dtype=torch.float32)
        bias = torch.zeros((4,), dtype=torch.float32)
        # The Nix bundle's functional.py routes through Triton directly.
        # CPU tensors must fail — no silent CPU fallback in the qualified bundle.
        with pytest.raises((ValueError, RuntimeError, AssertionError)):
            causal_conv1d_fn(x, weight, bias, activation="silu")


class TestGfx1030BoundaryMatrices:
    """Run boundary shapes against the Nix bundle on GPU."""

    @pytest.fixture(autouse=True)
    def _require_gpu(self):
        torch = pytest.importorskip("torch")
        if not torch.cuda.is_available():
            pytest.skip("No GPU available for gfx1030 boundary test")

    def test_small_fp32(self, causal_conv1d_fn):
        torch = pytest.importorskip("torch")
        b, d, seq_len, width = 1, 4, 8, 4
        x = torch.randn(b, d, seq_len, device="cuda", dtype=torch.float32)
        weight = torch.randn(d, width, device="cuda", dtype=torch.float32)
        bias = torch.randn(d, device="cuda", dtype=torch.float32)

        out = causal_conv1d_fn(x, weight, bias, activation="silu")
        assert out.shape == x.shape
        assert not torch.isnan(out).any(), "NaN in output"
        assert not torch.isinf(out).any(), "Inf in output"

    def test_small_fp16(self, causal_conv1d_fn):
        torch = pytest.importorskip("torch")
        b, d, seq_len, width = 1, 16, 32, 4
        x = torch.randn(b, d, seq_len, device="cuda", dtype=torch.float16)
        weight = torch.randn(d, width, device="cuda", dtype=torch.float16)
        bias = torch.randn(d, device="cuda", dtype=torch.float16)

        out = causal_conv1d_fn(x, weight, bias, activation="silu")
        assert out.shape == x.shape
        assert out.dtype == torch.float16
        assert not torch.isnan(out).any(), "NaN in output"

    def test_initial_states(self, causal_conv1d_fn):
        torch = pytest.importorskip("torch")
        b, d, seq_len, width = 2, 8, 16, 4
        x = torch.randn(b, d, seq_len, device="cuda")
        weight = torch.randn(d, width, device="cuda")
        bias = torch.randn(d, device="cuda")
        init = torch.randn(b, d, width - 1, device="cuda")

        out, final = causal_conv1d_fn(
            x, weight, bias, initial_states=init, return_final_states=True
        )
        # The Triton kernel prepends initial_states to x before convolution,
        # so the output length is seq_len + (width - 1). The final_states
        # are the last (width - 1) elements of the original x (not the
        # prepended version).
        expected_len = seq_len + (width - 1)
        assert out.shape == (b, d, expected_len), f"Expected {(b, d, expected_len)}, got {out.shape}"
        assert not torch.isnan(out).any()
        assert final.shape == (b, d, width - 1)

    def test_update_matches_prefill(self, causal_conv1d_fn, causal_conv1d_update):
        """Decode-step update must equal the prefill output for the last token."""
        torch = pytest.importorskip("torch")
        b, d, seq_len, width = 1, 8, 6, 4
        x = torch.randn(b, d, seq_len, device="cuda")
        weight = torch.randn(d, width, device="cuda")
        bias = torch.randn(d, device="cuda")

        prefill = causal_conv1d_fn(x, weight, bias, activation="silu")
        last_token = prefill[:, :, -1]

        conv_state = x[:, :, : width - 1]
        for t in range(width - 1, seq_len):
            token = x[:, :, t]
            out = causal_conv1d_update(
                token, conv_state, weight, bias, activation="silu"
            )
            torch.testing.assert_close(out, prefill[:, :, t], atol=1e-4, rtol=1e-4)
            conv_state = torch.cat([conv_state[:, :, 1:], token.unsqueeze(-1)], dim=-1)

        torch.testing.assert_close(out, last_token, atol=1e-4, rtol=1e-4)

    def test_backward_finite_gradients(self, causal_conv1d_fn):
        """Autograd backward produces finite gradients."""
        torch = pytest.importorskip("torch")
        b, d, seq_len, width = 1, 8, 16, 4
        x = torch.randn(b, d, seq_len, device="cuda", dtype=torch.float32, requires_grad=True)
        weight = torch.randn(d, width, device="cuda", dtype=torch.float32, requires_grad=True)
        bias = torch.randn(d, device="cuda", dtype=torch.float32, requires_grad=True)

        out = causal_conv1d_fn(x, weight, bias, activation="silu")
        loss = out.float().square().mean()
        loss.backward()
        torch.cuda.synchronize()

        assert torch.isfinite(x.grad).all(), "dx contains non-finite values"
        assert torch.isfinite(weight.grad).all(), "dweight contains non-finite values"
        assert torch.isfinite(bias.grad).all(), "dbias contains non-finite values"


def _record_qualification_artifact(cc_bundle, result_dir: Path) -> Path:
    """Write qualification artifact JSON with provenance metadata."""
    module_path = Path(cc_bundle.__file__).resolve()

    artifact = {
        "kernel_source": "nix-local-bundle",
        "kernel_module": cc_bundle.__name__,
        "kernel_path": str(module_path),
        "backend": "rocm",
        "triton_target": "hip:gfx1030",
        "result_dir": str(result_dir.resolve()),
        "timestamp": datetime.now(UTC).isoformat(),
        "python": sys.version,
    }

    out_path = result_dir / "qualification_artifact.json"
    out_path.write_text(json.dumps(artifact, indent=2) + "\n")
    return out_path