[FIX] Run NKI kernels written like the official samples in the NKI interpreter

tests/nki/test_nki_kernel_idioms.py

@@ -0,0 +1,140 @@
+"""Kernels written in the style of the official NKI samples, run by TileLens and by
+NKI's own CPU simulator (``nki.simulate_kernel``); both must give the same outputs."""
+
+import math
+
+import numpy as np
+import pytest
+
+try:
+    from neuronxcc import nki
+    import neuronxcc.nki.language as nl
+except ModuleNotFoundError:
+    pytest.skip(
+        "NeuronX dependencies are missing. Install tilelens[nki] to run these tests.",
+        allow_module_level=True,
+    )
+
+import tilelens
+from tilelens.clients import Tracer
+
+pytestmark = pytest.mark.nki  # only run at "pytest -m nki"
+
+
+def softmax_kernel(a_tensor):
+    # The output is allocated in shared HBM and returned; rows are tiled by 128 and
+    # the tail is masked with a [128, 1] mask that broadcasts over the tile.
+    out_tensor = nl.ndarray(a_tensor.shape, dtype=a_tensor.dtype, buffer=nl.shared_hbm)
+    ix = nl.arange(128)[:, None]
+    iy = nl.arange(a_tensor.shape[1])[None, :]
+    num_rows = a_tensor.shape[0]
+    for i in nl.affine_range(math.ceil(num_rows / 128)):
+        row_mask = i * 128 + ix < num_rows
+        a_tile = nl.load(a_tensor[i * 128 + ix, iy], mask=row_mask)
+        max_vals = nl.max(a_tile, axis=[1], keepdims=True, mask=row_mask)
+        shifted = nl.subtract(a_tile, max_vals, mask=row_mask)
+        numerator = nl.exp(shifted)
+        denominator = nl.sum(numerator, axis=[1])  # [128, 1], as in NKI
+        nl.store(
+            out_tensor[i * 128 + ix, iy], value=numerator / denominator, mask=row_mask
+        )
+    return out_tensor
+
+
+def copy_kernel(in_tensor):
+    # Whole-tensor load and store, and a load nested in a store.
+    out_tensor = nl.ndarray(
+        in_tensor.shape, dtype=in_tensor.dtype, buffer=nl.shared_hbm
+    )
+    tmp = nl.ndarray(in_tensor.shape, dtype=in_tensor.dtype, buffer=nl.shared_hbm)
+    nl.store(tmp, value=nl.load(in_tensor))
+    ix = nl.arange(in_tensor.shape[0])[:, None]
+    iy = nl.arange(in_tensor.shape[1])[None, :]
+    nl.store(out_tensor[ix, iy], value=nl.copy(nl.load(tmp[ix, iy])))
+    return out_tensor
+
+
+def spmd_add_kernel(a_input, b_input):
+    # One [128, 256] tile per program on a 2-D grid.
+    c_output = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.shared_hbm)
+    ix = nl.program_id(0) * 128 + nl.arange(128)[:, None]
+    iy = nl.program_id(1) * 256 + nl.arange(256)[None, :]
+    nl.store(
+        c_output[ix, iy],
+        value=nl.add(nl.load(a_input[ix, iy]), nl.load(b_input[ix, iy])),
+    )
+    return c_output
+
+
+def loops_and_casts_kernel(x):
+    # static_range / sequential_range loops, par_dim in a shape, static_cast,
+    # maximum / minimum / divide, and a reduction over two free axes.
+    out = nl.ndarray((x.shape[0], 1), dtype=nl.float32, buffer=nl.shared_hbm)
+    acc = nl.zeros((nl.par_dim(x.shape[0]), 1), dtype=nl.float32)
+    for k in nl.static_range(2):
+        t = nl.load(x[:, k])
+        clipped = nl.minimum(nl.maximum(t, -1.0), 1.0)
+        acc[...] = nl.add(acc, nl.sum(clipped, axis=[1, 2]))
+    for _ in nl.sequential_range(1):
+        acc[...] = nl.divide(acc, 2.0)
+    nl.store(out, value=nl.static_cast(acc, nl.float32))
+    return out
+
+
+def _tilelens(kernel, grid, *args):
+    return tilelens.trace(client=Tracer(), frontend="nki")(kernel)[grid](*args)
+
+
+def _reference(kernel, grid, *args):
+    jitted = nki.jit(kernel)
+    return nki.simulate_kernel(jitted[grid] if grid != (1,) else jitted, *args)
+
+
+rng = np.random.default_rng(0)
+
+
+@pytest.mark.parametrize(
+    "kernel, grid, args",
+    [
+        (softmax_kernel, (1,), (rng.standard_normal((200, 64)).astype(np.float32),)),
+        (copy_kernel, (1,), (rng.standard_normal((16, 32)).astype(np.float32),)),
+        (
+            spmd_add_kernel,
+            (2, 2),
+            (
+                rng.standard_normal((256, 512)).astype(np.float32),
+                rng.standard_normal((256, 512)).astype(np.float32),
+            ),
+        ),
+        (
+            loops_and_casts_kernel,
+            (1,),
+            (rng.standard_normal((8, 2, 4, 3)).astype(np.float32),),
+        ),
+    ],
+    ids=["softmax", "copy", "spmd_add", "loops_and_casts"],
+)
+def test_kernel_idioms_match_nki_simulator(kernel, grid, args):
+    tilelens.clear()
+    got = _tilelens(kernel, grid, *args)
+    want = _reference(kernel, grid, *args)
+    assert isinstance(got, np.ndarray)
+    np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-6)
+
+
+class _StopAfterFirstProgram(Tracer):
+    def post_run_callback(self, fn):
+        return False
+
+
+def test_outputs_are_returned_when_a_client_stops_the_grid_early():
+    tilelens.clear()
+    a = rng.standard_normal((256, 512)).astype(np.float32)
+    b = rng.standard_normal((256, 512)).astype(np.float32)
+    traced = tilelens.trace(client=_StopAfterFirstProgram(), frontend="nki")(
+        spmd_add_kernel
+    )
+    got = traced[(2, 2)](a, b)
+    assert isinstance(got, np.ndarray)
+    # Only program (0, 0) ran; its tile holds the sum.
+    np.testing.assert_allclose(got[:128, :256], a[:128, :256] + b[:128, :256])

tests/nki/test_nki_tracer.py

@@ -127,7 +127,8 @@ def reduce_sum_kernel(x_ptr, out_ptr):
     assert record.input_shape == (block_m, block_n)
     assert record.index == 1
     assert record.keep_dims is False
-    assert record.output_shape == (block_m,)
+    # As in NKI, reducing the only free axis of a [P, F] tile keeps it: [P, 1].
+    assert record.output_shape == (block_m, 1)
 
 
 def test_tracer_records_dot():

tilelens/clients/tracer/tracer.py

@@ -13,10 +13,20 @@
     Allocate,
 )
 from ...utils.traceback_utils import extract_user_frames
-from tilelens.core.masked_load_store import masked_load
+from tilelens.core.masked_load_store import masked_load, tile_mask, tile_shape
 import numpy as np
 
 
+def _nki_offsets_and_mask(tensor, keys, mask):
+    """Element offsets and mask of an NKI access ``tensor[keys]``. NKI masks broadcast
+    to the tile shape, and an access without a mask touches every lane."""
+    offsets = tensor.get_offsets().data
+    mask_data = tile_mask(offsets, keys, mask.data if mask is not None else None)
+    if mask_data is None:
+        mask_data = np.ones(tile_shape(offsets, keys), dtype=bool)
+    return masked_load(offsets, keys, mask=mask_data), mask_data
+
+
 def _convert_grid_idx(grid_idx) -> tuple[int, int, int] | None:
     if grid_idx is None:
         return grid_idx
@@ -103,12 +113,13 @@ def pre_load_callback(ptr, mask, keys):
                 first_ptr = np.reshape(ptr.data, (-1))[0]
                 tensor = self._get_tensor(first_ptr)
                 offsets = ptr.data - tensor.data_ptr()
+                mask_data = mask.data
             else:
                 keys = _convert_keys_to_numpy(keys)
-                offsets = masked_load(ptr.get_offsets().data, keys, mask=mask.data)
+                offsets, mask_data = _nki_offsets_and_mask(ptr, keys, mask)
                 tensor = ptr
 
-            rec = Load(tensor.data_ptr(), offsets, mask.data)
+            rec = Load(tensor.data_ptr(), offsets, mask_data)
             rec.call_path = extract_user_frames(num_frames=1)
             self.records.append(rec)
 
@@ -124,12 +135,7 @@ def pre_store_callback(ptr, mask, keys):
                 mask_data = mask.data
             else:
                 keys = _convert_keys_to_numpy(keys)
-                if mask is None:
-                    offsets = masked_load(ptr.get_offsets().data, keys)
-                    mask_data = np.ones_like(offsets).astype(bool)
-                else:
-                    mask_data = mask.data
-                    offsets = masked_load(ptr.get_offsets().data, keys, mask=mask_data)
+                offsets, mask_data = _nki_offsets_and_mask(ptr, keys, mask)
                 tensor = ptr
 
             rec = Store(tensor.data_ptr(), offsets, mask_data)

tilelens/core/frontend/nki.py

@@ -19,7 +19,7 @@
 HAS_NKI = False
 nki_builder = None
 try:
-    from tilelens.core.simulation.nki import NDArray, nki_builder  # type: ignore
+    from tilelens.core.simulation.nki import NDArray, full_keys, nki_builder  # type: ignore
 
     HAS_NKI = True
 except ModuleNotFoundError:
@@ -62,15 +62,21 @@ def _nki_reduce_sum_adapter(
     NKI_ADAPTERS = {
         ProgramId: lambda axis, *_args, **_kwargs: AdapterResult(axis),
         Allocate: lambda *_args, **_kwargs: AdapterResult(),
-        Load: lambda src, keys, *, mask=None, **_kwargs: AdapterResult(
+        # `nl.load(t)` and `nl.store(t, value=v)` access the whole tensor.
+        Load: lambda src, keys=None, *, mask=None, **_kwargs: AdapterResult(
             src,
             mask,
-            keys,
+            full_keys(src) if keys is None else keys,
         ),
-        Store: lambda dst, keys, value, *, mask=None, **_kwargs: AdapterResult(
+        Store: lambda dst,
+        keys=None,
+        value=None,
+        *,
+        mask=None,
+        **_kwargs: AdapterResult(
             dst,
             mask,
-            keys,
+            full_keys(dst) if keys is None else keys,
         ),
         Dot: _nki_dot_adapter,
         ReduceSum: _nki_reduce_sum_adapter,

tilelens/core/frontend/nki_transform.py

@@ -8,35 +8,19 @@ class NKITransformer(ast.NodeTransformer):
 
     Transformations:
     - `nl.<func>(x[...], ...)` -> `nl.<func>(x, slice_obj, ...)`
-       where <func> is one of {load, store, load_transpose2d}
+       where <func> is one of {load, store, load_transpose2d}, wherever the call
+       appears (e.g. nested in another call, as in `nl.store(o[i], nl.load(x[i]))`)
     - strips out kernel decorators
     """
 
-    def visit_Expr(self, node: ast.Expr) -> ast.AST | list[ast.AST]:
-        if not isinstance(node.value, ast.Call):
-            return self.generic_visit(node)
-
-        new_call = self._rewrite_nl_slice_call(node.value)
-        if new_call is None:
-            return self.generic_visit(node)
-        return ast.Expr(value=new_call)
-
     def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
         node.decorator_list = []
         return self.generic_visit(node)
 
-    def visit_Assign(self, node: ast.Assign) -> ast.AST:
-        if not isinstance(node.value, ast.Call):
-            return self.generic_visit(node)
-
-        new_call = self._rewrite_nl_slice_call(node.value)
-        if new_call is None:
-            return self.generic_visit(node)
-        return ast.Assign(
-            targets=node.targets,
-            value=new_call,
-            type_comment=getattr(node, "type_comment", None),
-        )
+    def visit_Call(self, node: ast.Call) -> ast.AST:
+        node = cast(ast.Call, self.generic_visit(node))
+        new_call = self._rewrite_nl_slice_call(node)
+        return node if new_call is None else new_call
 
     def _rewrite_nl_slice_call(self, call_node: ast.Call) -> ast.Call | None:
         if not (

tilelens/core/masked_load_store.py

@@ -67,6 +67,30 @@ def _get_valid_indices(
     return valid_mask
 
 
+def tile_shape(ndarray: np.ndarray, keys) -> tuple:
+    """Shape of ``ndarray[keys]``, also when some integer-array keys are out of bounds."""
+    if not isinstance(keys, tuple):
+        keys = (keys,)
+    clipped = []
+    for dim, key in enumerate(keys):
+        if (
+            isinstance(key, np.ndarray)
+            and key.dtype.kind in "iu"
+            and dim < ndarray.ndim
+        ):
+            key = np.clip(key, 0, max(ndarray.shape[dim] - 1, 0))
+        clipped.append(key)
+    return ndarray[tuple(clipped)].shape
+
+
+def tile_mask(ndarray: np.ndarray, keys, mask) -> np.ndarray | None:
+    """``mask`` broadcast to the shape of ``ndarray[keys]``, as NKI broadcasts masks;
+    None stays None."""
+    if mask is None:
+        return None
+    return np.broadcast_to(np.asarray(mask, dtype=bool), tile_shape(ndarray, keys))
+
+
 def masked_load(
     ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None
 ) -> np.ndarray: