test_nki_sim.py

"""Simulator-backed kernel correctness tests (NKI 0.3.0 Stable).

Run with `TRNSPARSE_USE_SIMULATOR=1` on any x86_64 Linux host that has
`nki>=0.3.0` installed. Bypasses torch_xla + NEFF compile; routes
kernel dispatch through `nki.simulate(kernel)(np_args)`.

Intentionally curated to small shapes — the CPU simulator is slow at
1024³ and above. Correctness parity with hardware at these scales is
what we're verifying, not perf.
"""

from __future__ import annotations

import os

import pytest
import torch

import trnsparse

pytestmark = pytest.mark.nki_simulator


@pytest.fixture(autouse=True)
def _simulator_enabled():
    """Skip the whole module if TRNSPARSE_USE_SIMULATOR isn't set.

    The marker alone isn't sufficient — users may `pytest -m
    nki_simulator` on a host where nki isn't importable or the env
    var hasn't been set. Fail loudly vs silently falling back.
    """
    if os.environ.get("TRNSPARSE_USE_SIMULATOR", "").lower() not in (
        "1",
        "true",
        "yes",
    ):
        pytest.skip("TRNSPARSE_USE_SIMULATOR=1 required")

    from trnsparse.nki.dispatch import HAS_NKI

    if not HAS_NKI:
        pytest.skip("nki package not importable on this host")


ATOL, RTOL = 1e-3, 1e-4


@pytest.fixture
def nki_backend():
    prev = trnsparse.get_backend()
    trnsparse.set_backend("nki")
    yield
    trnsparse.set_backend(prev)


def _local_mask(seq_len: int, block_size: int, window: int) -> torch.Tensor:
    n_blocks = seq_len // block_size
    bm = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
    for i in range(n_blocks):
        lo = max(0, i - window)
        hi = min(n_blocks, i + window + 1)
        bm[i, lo:hi] = True
    return bm.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)


def _dilated_mask(seq_len: int, block_size: int, stride: int) -> torch.Tensor:
    n_blocks = seq_len // block_size
    bm = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
    for i in range(n_blocks):
        for j in range(n_blocks):
            if (i - j) % stride == 0:
                bm[i, j] = True
    return bm.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)


class TestCsrSpmmSimulator:
    """CSR SpMM path through the simulator. Mirrors the hardware tests
    in `test_nki_spmm.py` at smaller shapes that fit the CPU simulator
    comfortably.
    """

    def test_aligned_128(self, nki_backend):
        torch.manual_seed(0)
        A_dense = torch.randn(128, 128)
        A_dense[torch.abs(A_dense) < 0.5] = 0.0
        A = trnsparse.from_dense(A_dense)
        B = torch.randn(128, 64)

        got = trnsparse.spmm(A, B)
        torch.testing.assert_close(got, A_dense @ B, atol=ATOL, rtol=RTOL)

    def test_rectangular(self, nki_backend):
        torch.manual_seed(1)
        A_dense = torch.randn(256, 128)
        A_dense *= (torch.rand(256, 128) < 0.1).float()
        A = trnsparse.from_dense(A_dense)
        B = torch.randn(128, 32)

        got = trnsparse.spmm(A, B)
        torch.testing.assert_close(got, A_dense @ B, atol=ATOL, rtol=RTOL)


class TestBsrSpmmSimulator:
    """BSR SpMM through the simulator. Block size 128 matches the Tensor
    Engine tile; we pick a minimal (2, 2) block grid so the simulator
    runs in seconds rather than minutes.
    """

    def test_block_dense_2x2(self, nki_backend):
        torch.manual_seed(2)
        b = 128
        # Fill both blocks in each row — exercises the non-zero-block path.
        A_dense = torch.randn(2 * b, 2 * b)
        bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=b)
        assert bsr.n_blocks == 4

        B = torch.randn(2 * b, 64)
        got = trnsparse.bsr_spmm(bsr, B)
        torch.testing.assert_close(got, A_dense @ B, atol=ATOL, rtol=RTOL)

    def test_block_sparse_2x2(self, nki_backend):
        """Drop one block to exercise the zero-block-padding path."""
        torch.manual_seed(3)
        b = 128
        A_dense = torch.zeros(2 * b, 2 * b)
        A_dense[:b, :b] = torch.randn(b, b)
        A_dense[b:, b:] = torch.randn(b, b)  # block-diagonal
        bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=b)
        assert bsr.n_blocks == 2

        B = torch.randn(2 * b, 32)
        got = trnsparse.bsr_spmm(bsr, B)
        torch.testing.assert_close(got, A_dense @ B, atol=ATOL, rtol=RTOL)


class TestAttnTiledSimulator:
    """NKI two-pass attention kernels through the simulator (#25).

    Shapes are intentionally tiny (seq_len=256, head_dim=32, 2×2 block grid)
    so the CPU simulator completes in seconds rather than minutes.
    """

    def test_stats_kernel_shapes(self, nki_backend):
        """Pass-1 kernel output shapes are (M_tiles, K_max, 128)."""
        import nki as _nki
        import numpy as np

        from trnsparse.nki.dispatch import _attn_gather, _attn_host_reduction
        from trnsparse.nki.kernels import _attn_stats_kernel

        torch.manual_seed(20)
        seq_len, head_dim, block_size = 256, 32, 128
        M_tiles = seq_len // block_size

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)

        mask = _local_mask(seq_len, block_size, window=1)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        scale = head_dim**-0.5
        qs, kg, vg, K_max, _ = _attn_gather(Q, K, V, mask_bsr, scale)

        t_max_np, t_sum_np = _nki.simulate(_attn_stats_kernel)(
            qs.contiguous().numpy(), kg.contiguous().numpy()
        )
        t_max = torch.from_numpy(np.asarray(t_max_np))
        t_sum = torch.from_numpy(np.asarray(t_sum_np))

        # NKI 0.3.0 keepdims: output may be (M_tiles, K_max, block_size, 1)
        t_max = t_max.squeeze(-1) if t_max.dim() == 4 else t_max
        t_sum = t_sum.squeeze(-1) if t_sum.dim() == 4 else t_sum
        assert t_max.shape == (M_tiles, K_max, block_size), f"tile_max shape: {t_max.shape}"
        assert t_sum.shape == (M_tiles, K_max, block_size), f"tile_sumexp shape: {t_sum.shape}"

    def test_local_window_parity(self, nki_backend):
        """NKI tiled attention matches PyTorch reference, local window."""
        torch.manual_seed(21)
        seq_len, head_dim, block_size = 256, 32, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)

        mask = _local_mask(seq_len, block_size, window=1)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        # PyTorch reference (pytorch backend)
        trnsparse.set_backend("pytorch")
        ref = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        # NKI simulator path
        trnsparse.set_backend("nki")
        got = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        torch.testing.assert_close(got, ref, atol=ATOL, rtol=RTOL)

    def test_dilated_parity(self, nki_backend):
        """NKI tiled attention matches PyTorch reference, dilated pattern."""
        torch.manual_seed(22)
        seq_len, head_dim, block_size = 256, 32, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)

        mask = _dilated_mask(seq_len, block_size, stride=2)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        trnsparse.set_backend("pytorch")
        ref = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        trnsparse.set_backend("nki")
        got = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        torch.testing.assert_close(got, ref, atol=ATOL, rtol=RTOL)


def _pytorch_grads(Q, K, V, mask_bsr):
    """Reference dQ, dK, dV via the PyTorch backend."""
    prev = trnsparse.get_backend()
    trnsparse.set_backend("pytorch")
    try:
        Qr = Q.clone().requires_grad_(True)
        Kr = K.clone().requires_grad_(True)
        Vr = V.clone().requires_grad_(True)
        out = trnsparse.block_sparse_attention_tiled(Qr, Kr, Vr, mask_bsr)
        dO = torch.randn_like(out)
        out.backward(dO)
        return Qr.grad.detach(), Kr.grad.detach(), Vr.grad.detach(), dO
    finally:
        trnsparse.set_backend(prev)


class TestAttnBwdSimulator:
    """NKI backward kernel pair through the simulator (v0.5.1).

    Correctness-only: NKI dQ/dK/dV must match PyTorch backward at atol=1e-3.
    Shapes are tiny (seq_len=256, head_dim=32) so the CPU simulator runs
    in seconds.
    """

    def test_bwd_dq_shapes(self, nki_backend):
        """dQ output from NKI backward kernel has shape (seq_len, head_dim)."""
        torch.manual_seed(30)
        seq_len, head_dim, block_size = 256, 32, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)
        mask = _local_mask(seq_len, block_size, window=1)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        Qr = Q.clone().requires_grad_(True)
        Kr = K.clone().requires_grad_(True)
        Vr = V.clone().requires_grad_(True)
        out = trnsparse.block_sparse_attention_tiled(Qr, Kr, Vr, mask_bsr)
        out.sum().backward()

        assert Qr.grad is not None and Qr.grad.shape == Q.shape
        assert Kr.grad is not None and Kr.grad.shape == K.shape
        assert Vr.grad is not None and Vr.grad.shape == V.shape

    @pytest.mark.xfail(
        strict=False,
        reason="NKI simulator: dQ backward has ~1.0 systematic error under investigation; "
        "dK/dV correct; hardware path unaffected",
    )
    def test_bwd_dq_parity(self, nki_backend):
        """NKI dQ matches PyTorch dQ at atol=1e-3, local window mask."""
        torch.manual_seed(31)
        seq_len, head_dim, block_size = 256, 32, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)
        mask = _local_mask(seq_len, block_size, window=1)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        dQ_ref, dK_ref, dV_ref, dO = _pytorch_grads(Q, K, V, mask_bsr)

        trnsparse.set_backend("nki")
        Qr = Q.clone().requires_grad_(True)
        Kr = K.clone().requires_grad_(True)
        Vr = V.clone().requires_grad_(True)
        out = trnsparse.block_sparse_attention_tiled(Qr, Kr, Vr, mask_bsr)
        out.backward(dO)

        torch.testing.assert_close(Qr.grad, dQ_ref, atol=ATOL, rtol=RTOL)

    def test_bwd_dkdv_parity(self, nki_backend):
        """NKI dK+dV match PyTorch at atol=1e-3, dilated mask."""
        torch.manual_seed(32)
        seq_len, head_dim, block_size = 256, 32, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)
        mask = _dilated_mask(seq_len, block_size, stride=2)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        dQ_ref, dK_ref, dV_ref, dO = _pytorch_grads(Q, K, V, mask_bsr)

        trnsparse.set_backend("nki")
        Qr = Q.clone().requires_grad_(True)
        Kr = K.clone().requires_grad_(True)
        Vr = V.clone().requires_grad_(True)
        out = trnsparse.block_sparse_attention_tiled(Qr, Kr, Vr, mask_bsr)
        out.backward(dO)

        torch.testing.assert_close(Kr.grad, dK_ref, atol=ATOL, rtol=RTOL)
        torch.testing.assert_close(Vr.grad, dV_ref, atol=ATOL, rtol=RTOL)


class TestAttnKTilingSimulator:
    """NKI K-tiling (head_dim=256) through the simulator (v0.6.0).

    Verifies that the K-tile loop path produces identical output to the
    PyTorch reference at head_dim=256. Both forward and backward are checked.
    """

    def test_forward_head_dim_256(self, nki_backend):
        """NKI forward at head_dim=256 matches PyTorch reference."""
        torch.manual_seed(60)
        seq_len, head_dim, block_size = 256, 256, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)
        mask = _local_mask(seq_len, block_size, window=1)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        trnsparse.set_backend("pytorch")
        ref = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        trnsparse.set_backend("nki")
        got = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        torch.testing.assert_close(got, ref, atol=ATOL, rtol=RTOL)
        assert got.shape == (seq_len, head_dim)

    @pytest.mark.xfail(
        strict=False,
        reason="NKI simulator: dQ backward systematic error (same issue as test_bwd_dq_parity)",
    )
    def test_backward_head_dim_256(self, nki_backend):
        """NKI dQ/dK/dV match PyTorch at head_dim=256."""
        torch.manual_seed(61)
        seq_len, head_dim, block_size = 256, 256, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)
        mask = _local_mask(seq_len, block_size, window=1)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        dQ_ref, dK_ref, dV_ref, dO = _pytorch_grads(Q, K, V, mask_bsr)

        trnsparse.set_backend("nki")
        Qr = Q.clone().requires_grad_(True)
        Kr = K.clone().requires_grad_(True)
        Vr = V.clone().requires_grad_(True)
        out = trnsparse.block_sparse_attention_tiled(Qr, Kr, Vr, mask_bsr)
        out.backward(dO)

        torch.testing.assert_close(Qr.grad, dQ_ref, atol=ATOL, rtol=RTOL)
        torch.testing.assert_close(Kr.grad, dK_ref, atol=ATOL, rtol=RTOL)
        torch.testing.assert_close(Vr.grad, dV_ref, atol=ATOL, rtol=RTOL)

    def test_forward_dilated_head_dim_256(self, nki_backend):
        """K-tiling forward with dilated pattern matches PyTorch."""
        torch.manual_seed(62)
        seq_len, head_dim, block_size = 256, 256, 128

        Q = torch.randn(seq_len, head_dim)
        K = torch.randn(seq_len, head_dim)
        V = torch.randn(seq_len, head_dim)
        mask = _dilated_mask(seq_len, block_size, stride=2)
        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)

        trnsparse.set_backend("pytorch")
        ref = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        trnsparse.set_backend("nki")
        got = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)

        torch.testing.assert_close(got, ref, atol=ATOL, rtol=RTOL)


class TestScreenedSpmmSimulator:
    """Fused screened SpMM through the simulator (#19).

    The NKI kernel fuses Q outer-product + threshold mask + nc_matmul.
    Small tile-aligned shapes so the simulator runs in seconds.
    """

    def test_threshold_zero_equals_plain_matmul(self, nki_backend):
        """threshold=0 → mask passes all entries → screened_spmm == A @ B."""
        torch.manual_seed(10)
        n = 128
        A = torch.randn(n, n)
        diag = torch.abs(torch.randn(n)) + 0.1
        B = torch.randn(n, 64)

        got = trnsparse.screened_spmm(A, diag, B, threshold=0.0)
        torch.testing.assert_close(got, A @ B, atol=ATOL, rtol=RTOL)

    @pytest.mark.xfail(
        strict=False,
        reason="NKI simulator: boolean mask to float conversion not yet correct",
    )
    def test_non_trivial_threshold_parity(self, nki_backend):
        """Non-trivial threshold drops some entries; NKI kernel must match
        the explicit (A * mask) @ B spec.
        """
        import math

        torch.manual_seed(11)
        n = 128
        A = torch.randn(n, n)
        diag = torch.abs(torch.randn(n)) * 4.0
        B = torch.randn(n, 64)
        threshold = 0.5

        got = trnsparse.screened_spmm(A, diag, B, threshold=threshold)

        Q = torch.sqrt(torch.abs(diag))
        mask = (Q.unsqueeze(-1) * Q.unsqueeze(0)) > math.sqrt(threshold)
        expected = (A * mask.to(A.dtype)) @ B

        torch.testing.assert_close(got, expected, atol=ATOL, rtol=RTOL)