From 5e6c2b3ffb0e4dcad245e31c8ac4a3e5b484b74d Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Sun, 8 Nov 2026 17:04:22 -0500 Subject: [PATCH] [FEAT] Interpret NKI ISA instructions and more nl operations in the NKI interpreter Implements, with NKI's semantics, the nisa instructions that AWS's NKI samples use: nc_matmul, nc_transpose, activation, activation_reduce, tensor_tensor, tensor_scalar, tensor_scalar_reduce, scalar_tensor_tensor, tensor_reduce, tensor_tensor_scan, tensor_copy, memset, affine_select, reciprocal, iota, dma_copy, dma_transpose, get_nc_version; and nl.ds, num_programs, program_ndim, full, ones, shared_constant, shared_identity_matrix, transpose, floor, mod. nisa.dma_copy and nisa.dma_transpose read and write through load/store, so clients see their HBM accesses; NKITransformer passes their subscripted operands as keys. Also: a 1-D tile (partition only) broadcasts along the free axes; shapes accept par_dim objects imported from NKI; nl.mgrid results have .p/.x; assigning a [P, 1] result to a [P] slot reshapes it; a slice past the tensor keeps the tile's size instead of being truncated by numpy, and an unmasked out-of-bounds load raises IndexError as in NKI; shared-HBM tensors were all named after the op-patching wrapper's frame and so aliased one array; the tracer accepts boolean masks and slices with computed bounds. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm --- tests/nki/test_nki_isa.py | 214 ++++++++++ tilelens/clients/tracer/tracer.py | 11 +- tilelens/core/frontend/nki_transform.py | 48 +++ tilelens/core/masked_load_store.py | 2 + tilelens/core/simulation/nki.py | 526 +++++++++++++++++++++++- 5 files changed, 793 insertions(+), 8 deletions(-) create mode 100644 tests/nki/test_nki_isa.py diff --git a/tests/nki/test_nki_isa.py b/tests/nki/test_nki_isa.py new file mode 100644 index 00000000..402e6316 --- /dev/null +++ b/tests/nki/test_nki_isa.py @@ -0,0 +1,214 @@ +"""NKI ISA instructions and language helpers in the interpreter, checked against NKI's own +CPU simulator (``nki.simulate_kernel``) on small kernels.""" + +import numpy as np +import pytest + +try: + from neuronxcc import nki + import neuronxcc.nki.isa as nisa + 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 matmul_kernel(lhsT, rhs): + # dma_copy between HBM and SBUF, and a Tensor Engine matmul: lhsT.T @ rhs. + out = nl.ndarray( + (lhsT.shape[1], rhs.shape[1]), dtype=nl.float32, buffer=nl.shared_hbm + ) + a = nl.ndarray(lhsT.shape, dtype=lhsT.dtype, buffer=nl.sbuf) + b = nl.ndarray(rhs.shape, dtype=rhs.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=a, src=lhsT) + nisa.dma_copy(dst=b, src=rhs) + res = nisa.nc_matmul(a, b) + nl.store(out, value=nl.copy(res, dtype=nl.float32)) + return out + + +def vector_kernel(x): + # activation, tensor_scalar, tensor_tensor, tensor_reduce, affine_select, + # activation_reduce, tensor_copy, memset, and nc_transpose. + out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm) + sums = nl.ndarray((x.shape[0], 1), dtype=x.dtype, buffer=nl.shared_hbm) + rows = nl.ndarray((x.shape[0], 1), dtype=x.dtype, buffer=nl.shared_hbm) + tr = nl.ndarray((x.shape[1], x.shape[0]), dtype=x.dtype, buffer=nl.shared_hbm) + t = nl.load(x) + bias = nisa.memset((x.shape[0], 1), 0.25, nl.float32) + a = nisa.activation(op=nl.exp, data=t, bias=bias, scale=0.5) + s = nisa.tensor_scalar(a, np.multiply, 2.0, op1=np.add, operand1=1.0) + u = nisa.tensor_tensor(s, t, op=np.subtract) + nl.store(rows, value=nisa.tensor_reduce(np.add, u, axis=[1])) + i = nl.arange(x.shape[0])[:, None] + j = nl.arange(x.shape[1])[None, :] + lower = nisa.affine_select(i - j >= 0, u, 0.0) + acc = nl.ndarray((x.shape[0], 1), dtype=nl.float32, buffer=nl.sbuf) + act = nisa.activation_reduce( + op=nl.copy, data=lower, reduce_op=np.add, reduce_res=acc + ) + nl.store(out, value=nisa.tensor_copy(act)) + nl.store(sums, value=acc) + nl.store(tr, value=nisa.nc_transpose(t)) + return out, sums, rows, tr + + +def grid_kernel(x): + # nl.ds slices from program_id, num_programs, full, and ones on a (2,) grid. + out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm) + p = nl.program_id(0) + n = nl.num_programs(0) + t = nl.load(x[nl.ds(p * 128, 128), :]) + c = nl.full((128, x.shape[1]), 2.0, dtype=nl.float32) + o = nl.ones((128, x.shape[1]), dtype=nl.float32) + nl.store( + out[nl.ds(p * 128, 128), :], value=nl.add(nl.multiply(t, c), nl.multiply(o, n)) + ) + return out + + +def dma_slices_kernel(x): + # dma_copy with subscripted source and destination, HBM to HBM through SBUF. + out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm) + for k in nl.affine_range(x.shape[0] // 128): + buf = nl.ndarray((128, x.shape[1]), dtype=x.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=buf, src=x[nl.ds(k * 128, 128), :]) + nisa.dma_copy(dst=out[nl.ds(k * 128, 128), :], src=buf) + return out + + +def misc_kernel(x, y): + # scalar_tensor_tensor, floor, mod, iota, dma_transpose of an HBM slice, + # shared_identity_matrix, and program_ndim. + out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm) + idx = nl.ndarray(x.shape, dtype=nl.int32, buffer=nl.shared_hbm) + tr = nl.ndarray((y.shape[1], y.shape[0]), dtype=y.dtype, buffer=nl.shared_hbm) + t = nl.load(x) + s = nisa.scalar_tensor_tensor( + data=t, op0=np.multiply, operand0=2.0, op1=np.add, operand1=t + ) + eye = nl.load(nl.shared_identity_matrix(x.shape[1], dtype=nl.float32)) + nl.store( + out, value=nl.add(nl.mod(nl.floor(s), 3.0), nl.program_ndim() + eye[0:1, :]) + ) + i = nl.arange(x.shape[0])[:, None] + j = nl.arange(x.shape[1])[None, :] + nl.store(idx, value=nisa.iota(i * x.shape[1] + j, dtype=nl.int32)) + nl.store(tr, value=nisa.dma_transpose(y[:, :])) + return out, idx, tr + + +from neuronxcc.nki.language import par_dim # noqa: E402 - kernels import it directly + + +def partition_kernel(x, scale): + # par_dim imported directly, nl.mgrid's .p / .x, a 1-D (partition) scale vector, + # tensor_tensor_scan, nisa.reciprocal, and a [P, 1] result assigned to a [P] slot. + out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm) + tot = nl.ndarray((x.shape[0], 2), dtype=x.dtype, buffer=nl.shared_hbm) + i = nl.mgrid[0 : x.shape[0], 0 : x.shape[1]] + t = nl.ndarray((par_dim(x.shape[0]), x.shape[1]), dtype=x.dtype) + t[i.p, i.x] = nl.load(x[i.p, i.x]) + s = nl.load(scale[0 : x.shape[0], 0]) # 1-D: one value per partition + a = nisa.activation(op=nl.exp, data=t, scale=s) + scan = nisa.tensor_tensor_scan(a, t, initial=0, op0=np.multiply, op1=np.add) + rows = nl.ndarray((par_dim(x.shape[0]), 2), dtype=x.dtype) + rows[:, 0] = nisa.tensor_reduce(np.add, scan, axis=[1]) # [P, 1] into a [P] column + rows[:, 1] = nisa.tensor_reduce(np.add, t, axis=[1]) + nl.store( + out, + value=nisa.tensor_scalar( + scan, np.multiply, nisa.reciprocal(nisa.tensor_reduce(np.add, t, axis=[1])) + ), + ) + nl.store(tot, value=rows) + return out, tot + + +rng = np.random.default_rng(0) + + +@pytest.mark.parametrize( + "kernel, grid, args", + [ + ( + matmul_kernel, + (1,), + ( + rng.standard_normal((128, 64)).astype(np.float32), + rng.standard_normal((128, 512)).astype(np.float32), + ), + ), + (vector_kernel, (1,), (rng.standard_normal((128, 64)).astype(np.float32),)), + (grid_kernel, (2,), (rng.standard_normal((256, 32)).astype(np.float32),)), + (dma_slices_kernel, (1,), (rng.standard_normal((256, 32)).astype(np.float32),)), + ( + misc_kernel, + (1,), + ( + (4 * rng.standard_normal((128, 64))).astype(np.float32), + rng.standard_normal((32, 128)).astype(nl.bfloat16), + ), + ), + ( + partition_kernel, + (1,), + ( + rng.random((128, 16)).astype(np.float32) + 0.5, + (0.1 * rng.standard_normal((128, 2))).astype(np.float32), + ), + ), + ], + ids=["matmul", "vector", "grid", "dma_slices", "misc", "partition"], +) +def test_isa_kernels_match_nki_simulator(kernel, grid, args): + tilelens.clear() + got = tilelens.trace(client=Tracer(), frontend="nki")(kernel)[grid](*args) + jitted = nki.jit(kernel) + want = nki.simulate_kernel(jitted[grid], *args) + got = got if isinstance(got, tuple) else (got,) + want = want if isinstance(want, tuple) else (want,) + assert len(got) == len(want) + for g, w in zip(got, want): + np.testing.assert_allclose( + g.astype(np.float32), w.astype(np.float32), rtol=1e-4, atol=1e-4 + ) + + +def test_dma_copy_is_traced_as_a_load_and_a_store(): + tilelens.clear() + from tilelens.core.data import Load, Store + from tilelens.core.trace import launches + + x = rng.standard_normal((256, 32)).astype(np.float32) + tilelens.trace(client=Tracer(), frontend="nki")(dma_slices_kernel)[(1,)](x) + records = launches[-1].records + loads = [r for r in records if isinstance(r, Load)] + stores = [r for r in records if isinstance(r, Store)] + # Two iterations, each a DMA from HBM to SBUF and one from SBUF to HBM. + assert len(loads) == 4 and len(stores) == 4 + + +def oob_slice_kernel(x): + out = nl.ndarray((128, x.shape[1]), dtype=x.dtype, buffer=nl.shared_hbm) + nl.store( + out, value=nl.load(x[0:128, 1 : x.shape[1] + 1]) + ) # one column past the end + return out + + +def test_out_of_range_slice_raises_like_nki(): + # numpy would truncate the slice to the tensor; NKI keeps the tile and reports the access. + tilelens.clear() + x = rng.standard_normal((128, 32)).astype(np.float32) + with pytest.raises(IndexError): + nki.simulate_kernel(nki.jit(oob_slice_kernel)[(1,)], x) + with pytest.raises(IndexError): + tilelens.trace(client=Tracer(), frontend="nki")(oob_slice_kernel)[(1,)](x) diff --git a/tilelens/clients/tracer/tracer.py b/tilelens/clients/tracer/tracer.py index 512d1b50..509334b4 100644 --- a/tilelens/clients/tracer/tracer.py +++ b/tilelens/clients/tracer/tracer.py @@ -21,7 +21,9 @@ 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) + mask_data = tile_mask( + offsets, keys, getattr(mask, "data", mask) + ) # NDArray, bool, or 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 @@ -102,6 +104,13 @@ def _convert_keys_to_numpy(keys): """Convert any NDArrays in keys to numpy arrays.""" if isinstance(keys, (tuple, list)): return tuple(_convert_keys_to_numpy(k) for k in keys) + if isinstance(keys, slice): # e.g. nl.ds(start, size) with a computed start + part = ( + lambda v: v + if v is None or isinstance(v, int) + else int(np.asarray(getattr(v, "data", v))) + ) # noqa: E731 + return slice(part(keys.start), part(keys.stop), part(keys.step)) return keys.data if hasattr(keys, "data") else keys @self.lock_fn diff --git a/tilelens/core/frontend/nki_transform.py b/tilelens/core/frontend/nki_transform.py index cc497fb2..8b2ed703 100644 --- a/tilelens/core/frontend/nki_transform.py +++ b/tilelens/core/frontend/nki_transform.py @@ -10,6 +10,8 @@ class NKITransformer(ast.NodeTransformer): - `nl.(x[...], ...)` -> `nl.(x, slice_obj, ...)` where 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]))`) + - `nisa.dma_copy(dst=o[...], src=x[...])` and `nisa.dma_transpose(x[...])` -> keys passed + as `dst_keys` / `src_keys` - strips out kernel decorators """ @@ -20,8 +22,54 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST: 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) + if new_call is None: + new_call = self._rewrite_dma_copy_call(node) return node if new_call is None else new_call + def _rewrite_dma_copy_call(self, call_node: ast.Call) -> ast.Call | None: + """`nisa.dma_copy(dst=o[...], src=x[...])` -> `nisa.dma_copy(dst=o, dst_keys=..., src=x, src_keys=...)`""" + if not ( + isinstance(call_node.func, ast.Attribute) + and isinstance(call_node.func.value, ast.Name) + and call_node.func.value.id == "nisa" + and call_node.func.attr in ("dma_copy", "dma_transpose") + ): + return None + args = list(call_node.args) + keywords = [] + changed = False + if ( + call_node.func.attr == "dma_transpose" + and args + and isinstance(args[0], ast.Subscript) + ): + # `nisa.dma_transpose(x[...])` -> `nisa.dma_transpose(x, src_keys=...)` + keywords.append( + ast.keyword( + arg="src_keys", + value=cast(ast.expr, self._create_slice_value_node(args[0].slice)), + ) + ) + args[0] = args[0].value + changed = True + for kw in call_node.keywords: + if kw.arg in ("dst", "src") and isinstance(kw.value, ast.Subscript): + keywords.append(ast.keyword(arg=kw.arg, value=kw.value.value)) + keywords.append( + ast.keyword( + arg=f"{kw.arg}_keys", + value=cast( + ast.expr, self._create_slice_value_node(kw.value.slice) + ), + ) + ) + changed = True + else: + keywords.append(kw) + if not changed: + return None + return ast.Call(func=call_node.func, args=args, keywords=keywords) + def _rewrite_nl_slice_call(self, call_node: ast.Call) -> ast.Call | None: if not ( isinstance(call_node.func, ast.Attribute) diff --git a/tilelens/core/masked_load_store.py b/tilelens/core/masked_load_store.py index b3c9c732..c73bc3dd 100644 --- a/tilelens/core/masked_load_store.py +++ b/tilelens/core/masked_load_store.py @@ -132,6 +132,8 @@ def masked_load( padded_array[tuple(copy_slices)] = ndarray[tuple(copy_slices)] result = padded_array[keys] + if mask is None: # every lane is active, so an out-of-bounds one is an error + mask = np.ones(np.expand_dims(result, singleton_dims).shape, dtype=bool) assert np.expand_dims(result, singleton_dims).shape == mask.shape mask = np.squeeze(mask, singleton_dims) diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py index 3a485fa2..30b35376 100644 --- a/tilelens/core/simulation/nki.py +++ b/tilelens/core/simulation/nki.py @@ -3,6 +3,7 @@ try: import neuronxcc.nki.language as nl + import neuronxcc.nki.isa as nisa except ( ModuleNotFoundError ) as exc: # pragma: no cover - only hit when optional deps missing @@ -10,11 +11,18 @@ "NeuronX dependencies are missing. Install tilelens[nki] to enable the NKI interpreter." ) from exc import inspect +import os import textwrap from ..frontend.nki_transform import transform_code from ..masked_load_store import masked_load, masked_store, tile_mask, tile_shape +_PACKAGE_DIR = ( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + os.sep +) + + def _value(x): """The numpy value of an NDArray operand; other operands (scalars) as is.""" return x.data if isinstance(x, NDArray) else x @@ -25,6 +33,55 @@ def _axis(axis): return tuple(axis) if isinstance(axis, list) else axis +_UFUNCS = { + "add": np.add, + "subtract": np.subtract, + "multiply": np.multiply, + "divide": np.divide, + "maximum": np.maximum, + "minimum": np.minimum, + "max": np.maximum, + "sum": np.add, + "prod": np.multiply, + "min": np.minimum, + "power": np.power, + "pow": np.power, + "equal": np.equal, + "not_equal": np.not_equal, + "less": np.less, + "less_equal": np.less_equal, + "greater": np.greater, + "greater_equal": np.greater_equal, + "bitwise_and": np.bitwise_and, + "bitwise_or": np.bitwise_or, + "bitwise_xor": np.bitwise_xor, + "logical_and": np.logical_and, + "logical_or": np.logical_or, +} + + +def _ufunc(op): + """The numpy ufunc of an NKI arithmetic operator: a numpy ufunc itself, or an + `nl` function such as `nl.add` (patched or not), looked up by name.""" + if isinstance(op, np.ufunc): + return op + name = getattr(op, "__name__", None) + if name in _UFUNCS: + return _UFUNCS[name] + raise NotImplementedError( + f"NKI operator {op!r} is not supported by the interpreter" + ) + + +def _int(x): + """A Python int from an index value (an NDArray scalar, a numpy scalar, or an int).""" + return x if x is None or isinstance(x, int) else int(np.asarray(_value(x))) + + +def _free_axes(ndim): + return tuple(range(1, ndim)) + + def _extreme(dtype, lowest): """Identity of a masked max (lowest) or min (highest) reduction over `dtype`.""" if np.issubdtype(dtype, np.integer): @@ -33,6 +90,54 @@ def _extreme(dtype, lowest): return -np.inf if lowest else np.inf +def _shape(shape): + """A shape with `par_dim(n)` entries (imported directly from NKI) as plain ints.""" + if isinstance(shape, (tuple, list)): + return tuple(int(getattr(d, "value", d)) for d in shape) + return int(getattr(shape, "value", shape)) + + +def _partition_aligned(operand, ndim): + """NKI broadcasts a 1-D tile, which has only the partition dimension, along the free + dimensions of an N-D tile: as an array it is [P, 1, ...], not [P] aligned to the last axis.""" + value = _value(operand) + if np.ndim(value) == 1 and ndim >= 2: + return np.reshape(value, (np.shape(value)[0],) + (1,) * (ndim - 1)) + return value + + +def _aligned_pair(a, b): + va, vb = _value(a), _value(b) + return _partition_aligned(va, np.ndim(vb)), _partition_aligned(vb, np.ndim(va)) + + +def _expand_oob_slices(keys, shape): + """NKI tiles keep their size when a slice runs past the tensor; numpy truncates the slice. + If a key of ints and slices has a slice out of range, index with one array per slice + instead (laid out like the slices' outer product), so the tile keeps its shape and the + out-of-bounds lanes reach masked_load's checks.""" + keys = keys if isinstance(keys, tuple) else (keys,) + keys = tuple( + int(k) if isinstance(k, np.ndarray) and k.ndim == 0 else k for k in keys + ) + if not all(isinstance(k, (slice, int, np.integer)) for k in keys): + return keys + ranges = [] + for dim, k in enumerate(keys): + if isinstance(k, slice) and dim < len(shape): + start = 0 if k.start is None else k.start + stop = shape[dim] if k.stop is None else k.stop + ranges.append( + (dim, np.arange(start, stop, 1 if k.step is None else k.step)) + ) + if not any(r.size and (r.min() < 0 or r.max() >= shape[d]) for d, r in ranges): + return keys + out = list(keys) + for i, (dim, r) in enumerate(ranges): + out[dim] = r.reshape((1,) * i + (-1,) + (1,) * (len(ranges) - i - 1)) + return tuple(out) + + def _cast(value, dtype): return np.asarray(value) if dtype is None else np.asarray(value).astype(dtype) @@ -116,15 +221,19 @@ def __setitem__(self, keys, value): keys = (keys,) # Apply the slicing to the underlying numpy array - new_keys = [k.data if isinstance(k, NDArray) else k for k in keys] - self.data[tuple(new_keys)] = value.data + new_keys = tuple(k.data if isinstance(k, NDArray) else k for k in keys) + target_shape = self.data[new_keys].shape + value = np.asarray(value.data if isinstance(value, NDArray) else value) + if value.shape != target_shape and value.size == int(np.prod(target_shape)): + value = value.reshape(target_shape) # e.g. a [P, 1] result into a [P] slot + self.data[new_keys] = value return self def _binary_op(self, other, op_func, op_name, op_symbol): if isinstance(other, NDArray): return NDArray( - value=op_func(self.data, other.data), + value=op_func(*_aligned_pair(self.data, other.data)), name=f"{self.name}_{op_name}_{other.name}", ) elif np.isscalar(other): @@ -138,7 +247,7 @@ def _binary_op(self, other, op_func, op_name, op_symbol): def _rbinary_op(self, other, op_func, op_name, op_symbol): if isinstance(other, NDArray): return NDArray( - value=op_func(other.data, self.data), + value=op_func(*_aligned_pair(other.data, self.data)), name=f"{other.name}_{op_name}_{self.name}", ) elif np.isscalar(other): @@ -205,10 +314,39 @@ def broadcast_to(self, *args, **kwargs): ) +class GridIndices(NDArray): + """The result of `nl.mgrid[...]`: unpacks into one index array per axis, and, as in NKI, + also names them `.p` (the partition axis), `.x`, `.y`, and `.z` (the free axes).""" + + def _axis(self, i): + return NDArray(value=self.data[i], name=f"{self.name}_{i}") + + def __iter__(self): + return (self._axis(i) for i in range(self.data.shape[0])) + + p = property(lambda self: self._axis(0)) + x = property(lambda self: self._axis(1)) + y = property(lambda self: self._axis(2)) + z = property(lambda self: self._axis(3)) + + +class MGrid: + def __getitem__(self, keys): + keys = keys if isinstance(keys, tuple) else (keys,) + keys = tuple( + slice(_int(k.start), _int(k.stop), _int(k.step)) + if isinstance(k, slice) + else k + for k in keys + ) + return GridIndices(value=np.mgrid[keys], name="mgrid") + + class Builder: def __init__(self, grid_dims=None): # TODO: infinite grid dims for NKI self.grid_dims = grid_dims if grid_dims is not None else (1, 1, 1) + self.grid_ndim = 1 self.grid_x = None self.grid_y = None self.grid_z = None @@ -224,10 +362,18 @@ def set_grid_idx(self, x, y, z): self.grid_z = z def ndarray(self, shape, dtype, *, buffer=None, name=None, **kwargs): + shape = _shape(shape) if buffer == nl.shared_hbm: if name is None: - # file name + function name + line number + # file name + function name + line number of the allocation in the kernel, + # so every program of the grid gets the same array. Skip TileLens's own + # frames (the op patching wraps this method), or every shared HBM tensor + # of a kernel would share one name, and one array. frame = inspect.currentframe().f_back + while frame is not None and frame.f_code.co_filename.startswith( + _PACKAGE_DIR + ): + frame = frame.f_back file_name = frame.f_code.co_filename function_name = frame.f_code.co_name line_number = frame.f_lineno @@ -246,6 +392,7 @@ def ndarray(self, shape, dtype, *, buffer=None, name=None, **kwargs): return ret def zeros(self, shape, dtype, *, buffer=None, name=None, **kwargs): + shape = _shape(shape) value = np.zeros(shape, dtype=dtype) return self.ndarray( shape, dtype, buffer=buffer, name=name, value=value, **kwargs @@ -270,6 +417,8 @@ def _convert_keys_to_numpy(self, keys): return tuple(self._convert_keys_to_numpy(k) for k in keys) elif isinstance(keys, NDArray): return keys.data + elif isinstance(keys, slice): # e.g. nl.ds(start, size) with a computed start + return slice(_int(keys.start), _int(keys.stop), _int(keys.step)) else: return keys @@ -285,6 +434,7 @@ def load(self, src: NDArray, keys=None, *, mask=None, **kwargs): numpy_keys = self._convert_keys_to_numpy(keys) # Call the actual masked_load function + numpy_keys = _expand_oob_slices(numpy_keys, ndarray.shape) mask_value = tile_mask(ndarray, numpy_keys, mask_value) result = masked_load(ndarray, numpy_keys, mask=mask_value) @@ -306,6 +456,7 @@ def store( numpy_keys = self._convert_keys_to_numpy(keys) # NKI broadcasts the stored value and the mask to the tile shape + numpy_keys = _expand_oob_slices(numpy_keys, ndarray.shape) tile_shape_ = tile_shape(ndarray, numpy_keys) value_array = np.asarray(_value(value)) same_size = value_array.size == int(np.prod(tile_shape_)) @@ -338,7 +489,7 @@ def _unary_op( def _binary_op(self, x, y, np_func, op_name, *, dtype=None, mask=None, **kwargs): return NDArray( - value=_cast(np_func(_value(x), _value(y)), dtype), + value=_cast(np_func(*_aligned_pair(x, y)), dtype), name=f"{getattr(x, 'name', 'scalar')}_{op_name}_{getattr(y, 'name', 'scalar')}", **kwargs, ) @@ -482,6 +633,332 @@ def rsqrt(self, x: NDArray, **kwargs): def multiply(self, x, y, **kwargs): return self._binary_op(x, y, np.multiply, "multiply", **kwargs) + # ---- nl: tensor creation, slicing, grid ------------------------------ + def full(self, shape, fill_value, dtype, *, buffer=None, name=None, **kwargs): + shape = _shape(shape) + value = np.full(shape, fill_value, dtype=dtype) + return self.ndarray( + shape, dtype, buffer=buffer, name=name, value=value, **kwargs + ) + + def ones(self, shape, dtype, *, buffer=None, name=None, **kwargs): + return self.full(shape, 1, dtype, buffer=buffer, name=name, **kwargs) + + def shared_constant(self, constant, dtype=None, **kwargs): + value = ( + np.asarray(constant) if dtype is None else np.asarray(constant, dtype=dtype) + ) + return NDArray(value=value, buffer=nl.shared_hbm, name="shared_constant") + + def transpose(self, x, *, dtype=None, mask=None, **kwargs): + return NDArray( + value=_cast(np.transpose(x.data), dtype), name=f"{x.name}_transpose" + ) + + def ds(self, start, size): + # A slice that keeps a computed start as is, so clients can see where it came from. + return slice(start, start + size) + + def program_ndim(self): + return self.grid_ndim + + def shared_identity_matrix(self, n, dtype=np.uint8, **kwargs): + return NDArray( + value=np.eye(n, dtype=dtype), + buffer=nl.shared_hbm, + name="shared_identity_matrix", + ) + + def floor(self, x, **kwargs): + return self._unary_op(x, np.floor, "floor", **kwargs) + + def mod(self, x, y, **kwargs): + return self._binary_op(x, y, np.mod, "mod", **kwargs) + + def num_programs(self, axes=None): + dims = [d for d in self.grid_dims] + if axes is None: + return int(np.prod(dims)) + if isinstance(axes, (list, tuple)): + return int(np.prod([dims[_int(a)] for a in axes])) + return dims[_int(axes)] + + # ---- nisa: instructions ------------------------------------------- + def _activate(self, op, value): + if isinstance(op, np.ufunc): + return op(value) + name = getattr(op, "__name__", None) + impl = getattr(self, name, None) if name else None + if impl is None or not callable(impl): + raise NotImplementedError( + f"NKI activation {op!r} is not supported by the interpreter" + ) + return impl(NDArray(value=np.asarray(value))).data + + def _reduce_into(self, reduce_op, value, reduce_res): + # Reduce along the free axes into `reduce_res`, a [P, 1] tile. + if reduce_op is None or reduce_res is None: + return + reduced = _ufunc(reduce_op).reduce( + value, axis=_free_axes(value.ndim), keepdims=True + ) + reduce_res.data[...] = np.broadcast_to( + reduced.reshape( + reduce_res.data.shape[:1] + (1,) * (reduce_res.data.ndim - 1) + ), + reduce_res.data.shape, + ) + + def nc_matmul(self, stationary, moving, *, mask=None, **kwargs): + # stationary.T @ moving on the Tensor Engine, accumulated in fp32 (PSUM). + value = np.matmul( + stationary.data.astype(np.float32).T, moving.data.astype(np.float32) + ) + return NDArray(value=value, name=f"{stationary.name}_{moving.name}_nc_matmul") + + def nc_transpose(self, data, *, mask=None, dtype=None, **kwargs): + return NDArray( + value=_cast(np.transpose(data.data), dtype), + name=f"{data.name}_nc_transpose", + ) + + def activation( + self, + op, + data, + *, + bias=None, + scale=1.0, + reduce_op=None, + reduce_res=None, + mask=None, + dtype=None, + **kwargs, + ): + nd = data.data.ndim + value = data.data.astype(np.float32) * np.asarray( + _partition_aligned(scale, nd), dtype=np.float32 + ) + if bias is not None: + value = value + _partition_aligned(bias, nd) + value = self._activate(op, value) + self._reduce_into(reduce_op, value, reduce_res) + return NDArray( + value=_cast(value, dtype if dtype is not None else data.data.dtype), + name=f"{data.name}_activation", + ) + + def activation_reduce( + self, + op, + data, + *, + reduce_op, + reduce_res, + bias=None, + scale=1.0, + mask=None, + dtype=None, + **kwargs, + ): + return self.activation( + op, + data, + bias=bias, + scale=scale, + reduce_op=reduce_op, + reduce_res=reduce_res, + mask=mask, + dtype=dtype, + ) + + def tensor_tensor(self, data1, data2, op, *, dtype=None, mask=None, **kwargs): + value = _ufunc(op)(*_aligned_pair(data1, data2)) + return NDArray(value=_cast(value, dtype), name="tensor_tensor") + + def tensor_scalar( + self, + data, + op0, + operand0, + reverse0=False, + op1=None, + operand1=None, + reverse1=False, + *, + dtype=None, + mask=None, + **kwargs, + ): + def apply(op, x, operand, reverse): + u = _ufunc(op) + operand = _partition_aligned(operand, np.ndim(x)) + return u(operand, x) if reverse else u(x, operand) + + value = apply(op0, data.data, operand0, reverse0) + if op1 is not None: + value = apply(op1, value, operand1, reverse1) + return NDArray( + value=_cast(value, dtype if dtype is not None else data.data.dtype), + name=f"{data.name}_tensor_scalar", + ) + + def tensor_scalar_reduce( + self, + *, + data, + op0, + operand0, + reduce_op, + reduce_res, + reverse0=False, + dtype=None, + mask=None, + **kwargs, + ): + out = self.tensor_scalar(data, op0, operand0, reverse0, dtype=dtype) + self._reduce_into(reduce_op, out.data, reduce_res) + return out + + def tensor_reduce( + self, + op, + data, + axis, + *, + mask=None, + dtype=None, + negate=False, + keepdims=False, + **kwargs, + ): + value = _ufunc(op).reduce(data.data, axis=_axis(axis), keepdims=keepdims) + if not keepdims and np.ndim(value) == 1 and data.data.ndim >= 2: + value = value[:, None] # NKI keeps a free dimension, as in nl reductions + if negate: + value = -value + return NDArray(value=_cast(value, dtype), name=f"{data.name}_tensor_reduce") + + def tensor_copy(self, src, *, mask=None, dtype=None, **kwargs): + return NDArray( + value=_cast(np.copy(src.data), dtype), name=f"{src.name}_tensor_copy" + ) + + def memset(self, shape, value, dtype, *, mask=None, **kwargs): + return NDArray(value=np.full(_shape(shape), value, dtype=dtype), name="memset") + + def affine_select( + self, pred, on_true_tile, on_false_value, *, mask=None, dtype=None, **kwargs + ): + value = np.where( + np.asarray(_value(pred), dtype=bool), _value(on_true_tile), on_false_value + ) + return NDArray( + value=_cast(value, dtype if dtype is not None else on_true_tile.data.dtype), + name="affine_select", + ) + + def dma_copy( + self, + *, + dst, + src, + dst_keys=None, + src_keys=None, + mask=None, + dst_rmw_op=None, + **kwargs, + ): + # A DMA transfer reads `src[src_keys]` and writes `dst[dst_keys]`; both go through + # load/store, so analysis clients see them as a load and a store. + value = self.load(src, src_keys, mask=mask) + if dst_rmw_op is not None: + current = self.load(dst, dst_keys, mask=mask) + value = NDArray(value=_ufunc(dst_rmw_op)(current.data, value.data)) + self.store(dst, dst_keys, value, mask=mask) + + def scalar_tensor_tensor( + self, + *, + data, + op0, + operand0, + op1, + operand1, + reverse0=False, + reverse1=False, + dtype=None, + mask=None, + **kwargs, + ): + # (data operand0) operand1 + first = self.tensor_scalar(data, op0, operand0, reverse0) + u = _ufunc(op1) + a, b = _aligned_pair(first.data, operand1) + value = u(b, a) if reverse1 else u(a, b) + return NDArray( + value=_cast(value, dtype if dtype is not None else data.data.dtype), + name=f"{data.name}_scalar_tensor_tensor", + ) + + def nisa_reciprocal(self, data, *, dtype=None, mask=None, **kwargs): + return NDArray( + value=_cast( + 1 / np.asarray(data.data, dtype=np.float32), dtype or data.data.dtype + ), + name=f"{data.name}_reciprocal", + ) + + def tensor_tensor_scan( + self, + data0, + data1, + initial, + op0, + op1, + reverse0=False, + reverse1=False, + *, + dtype=None, + mask=None, + **kwargs, + ): + # Per partition, along the free axis: + # result[:, i] = op1(op0(data0[:, i], result[:, i-1]), data1[:, i]), result[:, -1] = initial + u0, u1 = _ufunc(op0), _ufunc(op1) + d0 = data0.data.reshape(data0.data.shape[0], -1).astype(np.float32) + d1 = data1.data.reshape(data1.data.shape[0], -1).astype(np.float32) + prev = np.broadcast_to( + np.asarray(_partition_aligned(initial, 2), dtype=np.float32), + (d0.shape[0], 1), + )[:, 0] + result = np.empty_like(d0) + for i in range(d0.shape[1]): + first = u0(prev, d0[:, i]) if reverse0 else u0(d0[:, i], prev) + prev = u1(d1[:, i], first) if reverse1 else u1(first, d1[:, i]) + result[:, i] = prev + out_dtype = dtype if dtype is not None else data0.data.dtype + return NDArray( + value=_cast(result.reshape(data0.data.shape), out_dtype), + name=f"{data0.name}_scan", + ) + + def iota(self, expr, dtype, *, mask=None, **kwargs): + # A constant tile built from an index expression. + return NDArray(value=_cast(_value(expr), dtype), name="iota") + + def dma_transpose( + self, src, *, src_keys=None, axes=None, mask=None, dtype=None, **kwargs + ): + # A DMA read of `src[src_keys]` (seen by clients as a load), transposed. + value = self.load(src, src_keys, mask=mask).data + value = np.transpose(value) if axes is None else np.transpose(value, axes) + return NDArray(value=_cast(value, dtype), name=f"{src.name}_dma_transpose") + + def get_nc_version(self): + # The interpreter models a Trn1 NeuronCore. + return nisa.nc_version.gen2 + def range(self, *args, **kwargs): # affine_range, sequential_range, and static_range all iterate like range() return range(*args) @@ -514,13 +991,47 @@ def _set_attr(obj, name, value): _set_attr(nl, "static_range", nki_builder.range) _set_attr(nl, "par_dim", nki_builder.par_dim) _set_attr(nl, "zeros", nki_builder.zeros) - _set_attr(nl, "mgrid", NDArray(value=np.mgrid, buffer=nl.sbuf, name="mgrid")) + _set_attr(nl, "mgrid", MGrid()) _set_attr(nl, "matmul", nki_builder.matmul) _set_attr(nl, "copy", nki_builder.copy) _set_attr(nl, "sum", nki_builder.sum) _set_attr(nl, "square", nki_builder.square) _set_attr(nl, "rsqrt", nki_builder.rsqrt) _set_attr(nl, "multiply", nki_builder.multiply) + for name in ( + "full", + "ones", + "shared_constant", + "shared_identity_matrix", + "transpose", + "ds", + "num_programs", + "program_ndim", + "floor", + "mod", + ): + _set_attr(nl, name, getattr(nki_builder, name)) + for name in ( + "nc_matmul", + "nc_transpose", + "activation", + "activation_reduce", + "tensor_tensor", + "tensor_scalar", + "tensor_scalar_reduce", + "tensor_reduce", + "tensor_copy", + "memset", + "affine_select", + "dma_copy", + "dma_transpose", + "scalar_tensor_tensor", + "iota", + "tensor_tensor_scan", + "get_nc_version", + ): + _set_attr(nisa, name, getattr(nki_builder, name)) + _set_attr(nisa, "reciprocal", nki_builder.nisa_reciprocal) for name in ( "max", "min", @@ -584,6 +1095,7 @@ def run(self, *args, **kwargs): grid_dims = kwargs.pop( "grid", (1, 1, 1) ) # Remove grid from kwargs to avoid passing it to the function + nki_builder.grid_ndim = len(grid_dims) # make it 3d if not if len(grid_dims) == 1: grid_dims = (grid_dims[0], 1, 1)