diff --git a/ktir_cpu/memory.py b/ktir_cpu/memory.py index 09f2ee3..7b9aa79 100644 --- a/ktir_cpu/memory.py +++ b/ktir_cpu/memory.py @@ -220,72 +220,92 @@ def _read_flat( n_elements: int, np_dtype: np.dtype, elem_size: int, + *, + offsets: Optional[np.ndarray] = None, ) -> np.ndarray: - """Read *n_elements* elements starting at byte address *ptr*. + """Read elements starting at byte address *ptr*. - Returns a flat array of length *n_elements*. Elements beyond the end of - the containing allocation are zero-padded. Raises ``ValueError`` if *ptr* - is unmapped. + Two modes: + - **Contiguous** (offsets=None): reads *n_elements* starting at *ptr*. + Elements beyond the allocation end are zero-padded. + - **Sparse** (offsets=array): gathers elements at the given element + offsets relative to the allocation base. OOB offsets (past end) + return zero; *n_elements* and *np_dtype* are derived from offsets. - Example — reading 13 elements from inside a 4×4 f16 allocation:: - - # 4×4 f16 tensor at ptr=0x1000, values 0..15 - memory = {0x1000: np.arange(16, dtype=np.float16)} - # Read 13 elements starting at element 2 (byte offset 4 from base) - flat = _read_flat(memory, ptr=0x1004, n_elements=13, - np_dtype=np.float16, elem_size=2) - # flat == [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] + Raises ``ValueError`` if *ptr* is unmapped. """ alloc = _find_allocation(memory, ptr, elem_size) if alloc is None: raise ValueError(f"Read from unmapped address 0x{ptr:x} (n_elements={n_elements})") _, data, elem_offset = alloc - # ravel (a view when the allocation is contiguous, which it is here) instead of - # flatten (always a full copy): we only slice + astype(copy=True) out of `flat`, - # never mutate it, so copying the whole allocation per read is pure waste — it - # was the dominant cost of a whole-model pass once the offset loop was vectorized. flat = data.ravel() + + if offsets is not None: + indices = elem_offset + offsets + oob = indices >= flat.size + if oob.any(): + indices = np.where(oob, 0, indices) + result = flat[indices] + result[oob] = 0 + return result + return flat[indices] + end = elem_offset + n_elements if end <= flat.size: return flat[elem_offset:end].astype(np_dtype, copy=True) - # Partial allocation — pad remainder with zeros result = np.zeros(n_elements, dtype=np_dtype) avail = flat.size - elem_offset result[:avail] = flat[elem_offset:] return result -def _write_flat(memory: Dict[int, np.ndarray], ptr: int, data: np.ndarray): - """Write *data* (flat ndarray) at byte address *ptr*. - - Patches an existing allocation in-place when *ptr* falls within one. - Creates a new allocation at *ptr* if unmapped. +def _write_flat( + memory: Dict[int, np.ndarray], + ptr: int, + data: np.ndarray, + *, + offsets: Optional[np.ndarray] = None, +): + """Write *data* at byte address *ptr*. + + Two modes: + - **Contiguous** (offsets=None): writes *data* sequentially starting at + *ptr*. Creates a new allocation if unmapped. + - **Sparse** (offsets=array): scatters *data* into the given element + offsets relative to the allocation base. OOB offsets (past end) + are silently dropped. Raises ``ValueError`` if *ptr* is unmapped. + """ + bytes_per_elem_val = data.itemsize + alloc = _find_allocation(memory, ptr, bytes_per_elem_val) - Example — writing a single element into the middle of a 4×4 f16 tensor:: + if offsets is not None: + if alloc is None: + raise ValueError(f"Scatter to unmapped address 0x{ptr:x}") + _, buf, elem_offset = alloc + indices = elem_offset + offsets + inbounds = indices < buf.ravel().size + if not inbounds.all(): + buf.ravel()[indices[inbounds]] = data.ravel()[inbounds] + else: + buf.ravel()[indices] = data.ravel() + return - # 4×4 f16 tensor at ptr=0x1000, all zeros - memory = {0x1000: np.zeros(16, dtype=np.float16)} - # Write value 99 at element [1,2] (flat offset 6, byte offset 12) - _write_flat(memory, ptr=0x100C, data=np.array([99.0], dtype=np.float16)) - # memory[0x1000].reshape(4,4)[1, 2] == 99.0, all other elements unchanged - """ - bytes_per_elem = data.itemsize - alloc = _find_allocation(memory, ptr, bytes_per_elem) if alloc is not None: base_ptr, existing, elem_offset = alloc - flat = existing.flatten() # flatten already returns a fresh (mutable) copy - src = data.ravel() # read-only — a view is fine + flat = existing.flatten() + src = data.ravel() end_elem = elem_offset + src.size if end_elem <= flat.size: flat[elem_offset:end_elem] = src memory[base_ptr] = flat.reshape(existing.shape) return - # src extends past allocation end — write what fits fit = flat.size - elem_offset flat[elem_offset:] = src[:fit] memory[base_ptr] = flat.reshape(existing.shape) return - memory[ptr] = data.flatten() # flatten already copies; the extra .copy() was redundant + memory[ptr] = data.flatten() + + class HBMSimulator: @@ -331,37 +351,37 @@ def allocate(self, size: int) -> int: self.next_ptr = (self.next_ptr + self.STICK_BYTES - 1) & ~(self.STICK_BYTES - 1) return ptr // self.STICK_BYTES - def read(self, stick: int, n_elements: int, dtype: str, *, intra_byte: int = 0) -> np.ndarray: - """Read *n_elements* elements from HBM. + def read(self, stick: int, n_elements: int, dtype: str, *, intra_byte: int = 0, offsets: Optional[np.ndarray] = None) -> np.ndarray: + """Read elements from HBM (contiguous or sparse). Args: - stick: HBM stick index (from \`\`allocate()\`\` or \`\`MemRef.split_addr\`\`). - n_elements: Number of elements to read. + stick: HBM stick index. + n_elements: Number of elements (contiguous mode). dtype: Data type. intra_byte: Byte offset within the stick (default 0). - - Returns: - Flat NumPy array of length n_elements. + offsets: If provided, gather at these element offsets instead. """ assert 0 <= intra_byte < self.STICK_BYTES, ( f"intra_byte {intra_byte} out of range [0, {self.STICK_BYTES})" ) np_dtype = to_np_dtype(dtype) return _read_flat(self.memory, stick * self.STICK_BYTES + intra_byte, - n_elements, np_dtype, bytes_per_elem(dtype)) + n_elements, np_dtype, bytes_per_elem(dtype), offsets=offsets) - def write(self, stick: int, data: np.ndarray, *, intra_byte: int = 0): - """Write *data* (flat ndarray) to HBM. + def write(self, stick: int, data: np.ndarray, *, intra_byte: int = 0, offsets: Optional[np.ndarray] = None): + """Write data to HBM (contiguous or sparse). Args: - stick: HBM stick index (from \`\`allocate()\`\` or \`\`MemRef.split_addr\`\`). + stick: HBM stick index. data: Flat NumPy array to write. intra_byte: Byte offset within the stick (default 0). + offsets: If provided, scatter into these element offsets instead. """ assert 0 <= intra_byte < self.STICK_BYTES, ( f"intra_byte {intra_byte} out of range [0, {self.STICK_BYTES})" ) - _write_flat(self.memory, stick * self.STICK_BYTES + intra_byte, data) + _write_flat(self.memory, stick * self.STICK_BYTES + intra_byte, data, offsets=offsets) + def read_element(self, addr: int, dtype: str = "f16"): """Read a single element by byte address. @@ -391,34 +411,28 @@ def __init__(self, size_mb: int = 2, core_id: int = 0): self.memory: Dict[int, np.ndarray] = _AllocStore() self.next_ptr = 0 # Local address space - def read(self, ptr: int, n_elements: int, dtype: str) -> np.ndarray: - """Read *n_elements* elements starting at byte address *ptr*. - - Returns a flat array of length *n_elements*. Raises ValueError if - *ptr* is unmapped. + def read(self, ptr: int, n_elements: int, dtype: str, *, offsets: Optional[np.ndarray] = None) -> np.ndarray: + """Read elements from LX scratchpad (contiguous or sparse). Args: ptr: Local address (byte offset) - n_elements: Number of elements to read + n_elements: Number of elements (contiguous mode) dtype: Data type - - Returns: - Flat NumPy array of length n_elements + offsets: If provided, gather at these element offsets instead. """ np_dtype = to_np_dtype(dtype) - return _read_flat(self.memory, ptr, n_elements, np_dtype, bytes_per_elem(dtype)) + return _read_flat(self.memory, ptr, n_elements, np_dtype, bytes_per_elem(dtype), offsets=offsets) - def write(self, ptr: int, data: np.ndarray): - """Write *data* (flat ndarray) starting at byte address *ptr*. - - Patches an existing allocation in-place when *ptr* falls within one. - Creates a new allocation at *ptr* if unmapped. + def write(self, ptr: int, data: np.ndarray, *, offsets: Optional[np.ndarray] = None): + """Write data to LX scratchpad (contiguous or sparse). Args: ptr: Local address (byte offset) data: Flat NumPy array to write + offsets: If provided, scatter into these element offsets instead. """ - _write_flat(self.memory, ptr, data) + _write_flat(self.memory, ptr, data, offsets=offsets) + def clear(self): """Clear scratchpad and reset allocation.""" diff --git a/ktir_cpu/ops/memory_ops.py b/ktir_cpu/ops/memory_ops.py index dd0fa29..ba14fb6 100644 --- a/ktir_cpu/ops/memory_ops.py +++ b/ktir_cpu/ops/memory_ops.py @@ -30,6 +30,8 @@ from ..grid import CoreContext from ..memory import HBMSimulator +_MIN_BLOCKING_FACTOR = 16 + class _MemAccessor: """Resolves a (context, memory_space, byte_addr) triple into simulator @@ -91,8 +93,21 @@ def count_sticks( return None return len({a // HBMSimulator.STICK_BYTES for a in byte_addresses}) - def read(self, n: int, dtype: str) -> np.ndarray: - return self._sim.read(*self._args, n, dtype, **self._kwargs) + @classmethod + def count_sticks_array( + cls, memory_space: str, base_ptr: int, offsets: np.ndarray, bpe: int, + ) -> Optional[int]: + """Vectorized stick counting for large offset arrays. + + Same semantics as :meth:`count_sticks` but avoids Python iteration + over element offsets — uses numpy unique on the stick indices directly. + """ + if memory_space != "HBM": + return None + return int(np.unique((base_ptr + offsets * bpe) // HBMSimulator.STICK_BYTES).size) + + def read(self, n: int, dtype: str, *, offsets: Optional[np.ndarray] = None) -> np.ndarray: + return self._sim.read(*self._args, n, dtype, **self._kwargs, offsets=offsets) def read_scattered( self, byte_addresses: List[int], dtype: str, @@ -163,8 +178,9 @@ def read_scattered( ) return values, unique_sticks - def write(self, data: np.ndarray) -> None: - self._sim.write(*self._args, data, **self._kwargs) + def write(self, data: np.ndarray, *, offsets: Optional[np.ndarray] = None) -> None: + self._sim.write(*self._args, data, **self._kwargs, offsets=offsets) + def hbm_read(hbm: "HBMSimulator", byte_addr: int, n_elements: int, dtype: str) -> np.ndarray: @@ -180,59 +196,151 @@ def hbm_write(hbm: "HBMSimulator", byte_addr: int, data: np.ndarray) -> None: hbm.write(stick, data, intra_byte=intra) -def _enumerate_in_vso_order(iat: "IndirectAccessTile") -> List[Tuple[int, ...]]: - """Enumerate variable-space points in ``variables_space_order``-permuted order. +def _expr_dependent_vars(expr: tuple) -> set: + """Return the set of iteration-variable indices that *expr* depends on. - Identity (or absent) ``vso`` returns the natural row-major enumeration; - otherwise points are sorted by ``vso.eval(pt)`` per RFC 0682 §473. - - Both :func:`_resolve_idx_reads` and :func:`_build_indirect_coords` route - through this so their pt iteration stays in lockstep — they consume - ``idx_values`` positionally, so any divergence would silently mismatch - indirect dims to coords. Callers are expected to have already rejected - non-permutation ``vso`` upstream; this function trusts the guard. + Walks the subscript-expression AST produced by ``parse_subscript_expr`` + and collects every ``("dim", i)`` reference. ``("const", ...)`` and + ``("ssa", ...)`` nodes contribute nothing — they are loop-invariant. """ - points = iat.variables_space_set.enumerate(iat.shape) + tag = expr[0] + if tag == "dim": + return {expr[1]} + if tag == "const" or tag == "ssa": + return set() + if tag in ("add", "sub"): + return _expr_dependent_vars(expr[1]) | _expr_dependent_vars(expr[2]) + if tag == "mul": + # ("mul", const_int, sub_expr) + return _expr_dependent_vars(expr[2]) + if tag == "neg": + return _expr_dependent_vars(expr[1]) + if tag in ("floordiv", "mod"): + # ("floordiv", sub_expr, const_int) + return _expr_dependent_vars(expr[1]) + return set() + + +def _analyze_blocked_indirect(iat: "IndirectAccessTile"): + """Analyze an indirect-access expression (IAT), extract access pattern + for the consideration of fast emulation of the datamoves associated with + indirect-accesses. + + The condition for the fast-path is to meet all of the following: + 1.1. the IAT has at least one indirect subscript + 1.2. at least one direct subscript + 2. no direct_expr subscripts and identity VSO (meshgrid-compatible) + 3.1 blocking factor is greater than 16. Here the blocking factor is defined as + the ratio of sizes of the two spaces: resulting data tensor (N) vs distinct + accesses of indirect subscripts (K). Equivalently, the sizes of the two + spaces map to the total number of points in the iteration space and the + number of unique index lookups, respectively. + 3.2 For store op, we use the source data tensor. + + Returns (indirect_subs, dep_vars, dep_var_list, dep_extents) if the IAT + qualifies for the blocked-indirect fast-path, or None otherwise. + + The tuple has 4 ordered fields: + - indirect_subs: subscripts with kind=="indirect" (the index lookups) + - dep_vars: set of variable-space dims the index exprs depend on + - dep_var_list: sorted list form of dep_vars (stable iteration order) + - dep_extents: iteration extent per dependent dim (aka, K, the number of + unique index lookups) + + Example: W[e_idx[e], m, n] from MoE (synthetic) + - indirect_subs has one item: e_idx[e] + - dep_vars = {e} for the dependency of the e_idx[e] expression on 'e', index + or id of an expert, likely associated with a loop induction var. + They are referred to as "indirect variable" (vs "direct") sometimes. + """ + # --- Gate 1: must have both indirect and direct subscript dimensions --- + indirect_subs = [s for s in iat.dim_subscripts if s.get("kind") == "indirect"] + if len(indirect_subs) < 1: + return None + + direct_subs = [s for s in iat.dim_subscripts if s.get("kind") == "direct"] + if len(direct_subs) < 1: + return None + + # --- Gate 2: reject cases that can't use pure meshgrid broadcast --- + has_direct_expr = any(s.get("kind") == "direct_expr" for s in iat.dim_subscripts) vso = iat.variables_space_order - if vso is not None and not vso.is_identity(): - points = sorted(points, key=lambda pt: vso.eval(pt)) - return points + non_identity_vso = vso is not None and not vso.is_identity() + if has_direct_expr or non_identity_vso: + return None + + # --- Collect iteration-space dims the index exprs depend on --- + dep_vars: set = set() + for sub in indirect_subs: + for expr in sub["idx_exprs"]: + dep_vars |= _expr_dependent_vars(expr) + + vss = iat.variables_space_set + if not isinstance(vss, BoxSet): + return None + + # --- Gate 3: blocking factor N/K must be ≥ _MIN_BLOCKING_FACTOR --- + unique_lookups = 1 + for d in dep_vars: + extent = int(vss.hi[d]) - int(vss.lo[d]) + if extent <= 0: + continue + unique_lookups *= extent + total_points = 1 + for d in range(vss.n_dims): + extent = int(vss.hi[d]) - int(vss.lo[d]) + if extent > 0: + total_points *= extent -def _resolve_idx_reads( + if unique_lookups * _MIN_BLOCKING_FACTOR > total_points: + return None + + dep_var_list = sorted(dep_vars) + dep_extents = [int(vss.hi[d]) - int(vss.lo[d]) for d in dep_var_list] + return indirect_subs, dep_vars, dep_var_list, dep_extents + + +def _prepare_dep_var_sub_space( + iat: "IndirectAccessTile", dep_vars: set, dep_var_list: list, +) -> list: + """K sampling coordinates for the dep-var subspace. + + Output: K tuples, each n_dims wide. Dep-var positions sweep their + full range; direct positions are pinned to lo (irrelevant to index + lookups). K = product of dep-var extents. + """ + import itertools + vss = iat.variables_space_set + if dep_vars: + dep_ranges = [range(int(vss.lo[d]), int(vss.hi[d])) for d in dep_var_list] + # Cartesian product over dep dims only; non-dep dims pinned to lo + base = list(vss.lo) + points = [] + for dpt in itertools.product(*dep_ranges): + pt = list(base) + for i, d in enumerate(dep_var_list): + pt[d] = dpt[i] + points.append(tuple(pt)) + return points + return [tuple(vss.lo)] + + +def _runtime_read_and_expand_sub_space( context: CoreContext, iat: "IndirectAccessTile", + points, indirect_subs: list, ) -> Tuple[Dict[int, np.ndarray], int]: - """Read every idx-tensor value the IAT enumeration needs. - - For each indirect dimension, enumerates its address in pt order, then - issues one ``_MemAccessor.read_scattered`` per index view (so all - reads to one view share a single accessor and a single dedup pass). - - Returns ``(per_view_values, total_idx_unique_sticks)``: - - * ``per_view_values[idx_view_idx]`` is an ``np.ndarray`` whose ``i``-th - entry is the idx value resolved for the ``i``-th enumerated point's - use of that view. Indirect dims sharing the same view share the - array (consumed in pt-major, dim-minor order). - * ``total_idx_unique_sticks`` is the sum across HBM views; ``0`` when - every idx view lives in LX (LX has no stick concept). The return - type is always ``int``: callers receiving ``None`` would have to - special-case it, and the LX-only case is a defined "zero HBM - traffic" answer, so the function returns the integer directly. - - Per-view loop-invariants (``bpe``, ``strides``, ``byte_address``) - are hoisted out of the pt loop for million-point scale. - - This is the canonical idx-side resolver: ``indirect_load`` and - ``indirect_store`` both call it so their stick accounting stays in - sync (guard symmetry). - """ - points = _enumerate_in_vso_order(iat) - indirect_subs = [s for s in iat.dim_subscripts if s.get("kind") == "indirect"] + """K index values per indirect subscription, read from HBM. + + For each of the K points, computes byte addresses for each sub's + index expression, then batch-reads via scattered DMA. - # Hoist per-view loop-invariants once before enumerating points. + Output: ``per_sub_values[sub_i]`` — a K-element int array. + These K values are the input to the K→N broadcast step. + Raises ``IndexError`` on negative indices. + """ + # --- Phase 1: cache view constants (dedup across subs sharing a view) --- per_view_consts: Dict[int, Tuple[int, List[int], int]] = {} - per_view_addrs: Dict[int, List[int]] = {} for sub in indirect_subs: iv_idx = sub["index_view_idx"] if iv_idx in per_view_consts: @@ -241,37 +349,181 @@ def _resolve_idx_reads( per_view_consts[iv_idx] = ( _bytes_per_elem(iv.dtype), list(iv.strides), iv.byte_address, ) - per_view_addrs[iv_idx] = [] + # --- Phase 2: compute byte addresses per subscription expression --- + per_sub_addrs: Dict[int, List[int]] = {i: [] for i in range(len(indirect_subs))} for pt in points: - for sub in indirect_subs: + for sub_i, sub in enumerate(indirect_subs): iv_idx = sub["index_view_idx"] bpe, strides, base = per_view_consts[iv_idx] offset = sum( eval_subscript_expr(e, pt) * s for e, s in zip(sub["idx_exprs"], strides) ) - per_view_addrs[iv_idx].append(base + offset * bpe) + per_sub_addrs[sub_i].append(base + offset * bpe) - per_view_values: Dict[int, np.ndarray] = {} + # --- Phase 3: batch-read per sub via its view's accessor --- + per_sub_values: Dict[int, np.ndarray] = {} total_sticks = 0 - for iv_idx, addrs in per_view_addrs.items(): - # Zero-extent enumeration: no points, no addresses, no read. - # _build_indirect_coords iterates the same enumeration, so it - # also produces zero coords and never consumes from this view. + for sub_i, addrs in per_sub_addrs.items(): if not addrs: continue + iv_idx = indirect_subs[sub_i]["index_view_idx"] idx_view = iat.index_views[iv_idx] accessor = _MemAccessor( context, idx_view.memory_space, idx_view.byte_address, idx_view.lx_core_id, ) values, sticks = accessor.read_scattered(addrs, idx_view.dtype) - per_view_values[iv_idx] = values + if values.size and (values < 0).any(): + raise IndexError( + f"indirect index {int(values.min())} from sub " + f"{sub_i} is negative" + ) + per_sub_values[sub_i] = values if sticks is not None: total_sticks += sticks - return per_view_values, total_sticks + return per_sub_values, total_sticks + + +def _gen_offsets_vso_space_via_broadcast( + iat: "IndirectAccessTile", + idx_values_map: dict, + indirect_subs: list, + dep_vars: set, dep_var_list: list, dep_extents: list, +) -> np.ndarray: + """K→N broadcast: K index values + direct aranges → N flat byte offsets. + + Indirect subs: K-element arrays placed along dep-var axes. + Direct subs: arange placed along that dim's axis. + Numpy broadcasting crosses these 1-D axes into iter_shape (all dims), + weighted by parent strides. + + Output: 1-D int64 array, length N = product of all dim extents. + """ + vss = iat.variables_space_set + tile_ref = iat.parent_ref.to_tile_ref() + parent_strides = np.asarray(tile_ref.strides, dtype=np.int64) + + vss_dim_ranges = [np.arange(int(vss.lo[d]), int(vss.hi[d]), dtype=np.int64) + for d in range(vss.n_dims)] + + # --- Linearize K-dimensional dep-var coordinates into flat 0..K-1 indices --- + if dep_vars: + dep_meshgrid = np.meshgrid( + *[np.arange(e, dtype=np.int64) for e in dep_extents], + indexing='ij', + ) + dep_strides_arr = np.ones(len(dep_var_list), dtype=np.int64) + for i in range(len(dep_var_list) - 2, -1, -1): + dep_strides_arr[i] = dep_strides_arr[i + 1] * dep_extents[i + 1] + dep_flat_idx = sum(g * s for g, s in zip(dep_meshgrid, dep_strides_arr)) + else: + dep_flat_idx = None + + # --- Scatter K index values into n_dims-shaped grids (one per indirect sub) --- + # Each grid has extent only along dep-var axes, size-1 elsewhere (broadcasts) + indirect_coord_grids = {} + for sub_i in range(len(indirect_subs)): + idx_values_arr = idx_values_map[sub_i] + if dep_vars and dep_flat_idx is not None: + broadcast_shape = [1] * vss.n_dims + for d_pos, d in enumerate(dep_var_list): + broadcast_shape[d] = dep_extents[d_pos] + sub_grid = idx_values_arr[dep_flat_idx.ravel()].reshape(broadcast_shape).astype(np.int64) + else: + sub_grid = np.full([1] * vss.n_dims, int(idx_values_arr[0]), dtype=np.int64) + indirect_coord_grids[sub_i] = sub_grid + + # --- Accumulate weighted coordinates: offset += coord_grid * stride --- + # numpy broadcasting expands each 1-D or K-D grid to iter_shape (shape of the VSS iteration space) + iter_shape = tuple(int(vss.hi[d]) - int(vss.lo[d]) for d in range(vss.n_dims)) + offsets = np.zeros(iter_shape, dtype=np.int64) + + sub_idx = 0 + for dim_i, sub_d in enumerate(iat.dim_subscripts): + kind = sub_d["kind"] + s = parent_strides[dim_i] + if kind == "indirect": + offsets = offsets + indirect_coord_grids[sub_idx] * s + sub_idx += 1 + elif kind == "direct": + # Direct dim: reshape 1-D range to broadcast along its axis + var_idx = sub_d["var_index"] + range_direct_dim = vss_dim_ranges[var_idx] + shape_for_broadcast = [1] * vss.n_dims + shape_for_broadcast[var_idx] = len(range_direct_dim) + offsets = offsets + range_direct_dim.reshape(shape_for_broadcast) * s + + return offsets.ravel() + + +def _compute_blocked_indirect_offsets( + context: CoreContext, iat: "IndirectAccessTile", + info: tuple, +) -> Tuple[np.ndarray, int]: + """Compute element-wise linearized offsets via the blocked-indirect broadcast path. + + Reads K index values from HBM (small DMA), then broadcasts them into + N flat offsets via numpy meshgrid — no Python per-point loop. + + Returns (offsets, idx_sticks). + """ + indirect_subs, dep_vars, dep_var_list, dep_extents = info + + # Step 1: prepare dep-var subspace, read K index values from HBM + points = _prepare_dep_var_sub_space(iat, dep_vars, dep_var_list) + idx_values_map, idx_sticks = _runtime_read_and_expand_sub_space( + context, iat, points, indirect_subs, + ) + + # Step 2: broadcast K index values → N flat offsets + offsets = _gen_offsets_vso_space_via_broadcast( + iat, idx_values_map, indirect_subs, + dep_vars, dep_var_list, dep_extents, + ) + return offsets, idx_sticks + + + + +def _enumerate_in_vso_order(iat: "IndirectAccessTile") -> List[Tuple[int, ...]]: + """Enumerate variable-space points in ``variables_space_order``-permuted order. + + Identity (or absent) ``vso`` returns the natural row-major enumeration; + otherwise points are sorted by ``vso.eval(pt)`` per RFC 0682 §473. + + Both :func:`_resolve_idx_reads` and :func:`_build_indirect_coords` route + through this so their pt iteration stays in lockstep — they consume + ``idx_values`` positionally, so any divergence would silently mismatch + indirect dims to coords. Callers are expected to have already rejected + non-permutation ``vso`` upstream; this function trusts the guard. + """ + points = iat.variables_space_set.enumerate(iat.shape) + vso = iat.variables_space_order + if vso is not None and not vso.is_identity(): + points = sorted(points, key=lambda pt: vso.eval(pt)) + return points + + +def _resolve_idx_reads( + context: CoreContext, iat: "IndirectAccessTile", +) -> Tuple[Dict[int, np.ndarray], int]: + """Read every idx-tensor value the IAT enumeration needs (general path). + + Returns ``(per_sub_values, total_sticks)`` keyed by subscription index. + Delegates to :func:`_runtime_read_and_expand_sub_space` with the full VSO-ordered + enumeration. + + Note: :func:`_build_indirect_coords` consumes the dict by + ``index_view_idx`` lookup, which is correct only when each indirect sub + uses a distinct view (the general-path invariant — shared-view IATs + route to the blocked-indirect fast path instead). + """ + points = _enumerate_in_vso_order(iat) + indirect_subs = [s for s in iat.dim_subscripts if s.get("kind") == "indirect"] + return _runtime_read_and_expand_sub_space(context, iat, points, indirect_subs) def _build_indirect_coords( @@ -285,8 +537,9 @@ def _build_indirect_coords( * ``direct`` dims read directly from the variable-space point. * ``direct_expr`` dims evaluate a quasi-affine expression over the point. * ``indirect`` dims consume the next pre-resolved value from - ``idx_values[iv_idx]`` (set up by :func:`_resolve_idx_reads` in the - same pt-major, dim-minor order). + ``idx_values[sub_i]`` (set up by :func:`_resolve_idx_reads` in the + same pt-major, dim-minor order; works because each sub uses a distinct + view on the general path). Raises ``IndexError`` on a negative idx value — NumPy fancy-indexing silently wraps negatives, so we reject them here. The check survives @@ -296,11 +549,12 @@ def _build_indirect_coords( construction stays in lockstep (guard symmetry). """ points = _enumerate_in_vso_order(iat) - idx_iters = {iv_idx: iter(values) for iv_idx, values in idx_values.items()} + idx_iters = {sub_i: iter(values) for sub_i, values in idx_values.items()} coords: List[Tuple[int, ...]] = [] for pt in points: coord: List[int] = [] + indirect_counter = 0 for sub in iat.dim_subscripts: kind = sub["kind"] if kind == "direct": @@ -308,14 +562,14 @@ def _build_indirect_coords( elif kind == "direct_expr": coord.append(eval_subscript_expr(sub["subscript"], pt)) elif kind == "indirect": - iv_idx = sub["index_view_idx"] - raw_idx = int(next(idx_iters[iv_idx])) + raw_idx = int(next(idx_iters[indirect_counter])) if raw_idx < 0: raise IndexError( f"indirect index {raw_idx} from " - f"{iat.index_views[iv_idx]} is negative" + f"{iat.index_views[sub['index_view_idx']]} is negative" ) coord.append(raw_idx) + indirect_counter += 1 else: raise ValueError(f"Unknown indirect subscript kind: {kind}") coords.append(tuple(coord)) @@ -471,6 +725,7 @@ def load( context: CoreContext, tile_ref: TileRef, coords: Optional[List[Tuple[int, ...]]] = None, + offsets: Optional[np.ndarray] = None, result_shape: Optional[Tuple[int, ...]] = None, ) -> Tile: """Load data from HBM or LX into LX and return a Tile. @@ -479,39 +734,24 @@ def load( - HBM source → DMA read from HBM, write into LX scratchpad. - LX source → logical copy within LX (no physical movement). - When *coords* is given (coordinate-set path), gathers only the - elements at those local coordinates and reshapes to *result_shape*. - When *coords* is None, loads the full tile described by tile_ref - (contiguous or strided). + Three dispatch modes (checked in order): + 1. *offsets* — pre-computed flat element offsets (blocked-indirect fast + path). Skips coordinate linearization entirely. + 2. *coords* — gathers elements at those local coordinates. + 3. Neither — loads the full tile (contiguous or strided). A single ``mem.read`` covers the entire element footprint; no per-element dict scans occur. - Example — loading column 2 of a 4×4 f16 matrix (strided, coords=None):: - - # Parent 4×4 allocation at base_ptr=0x1000, values 0..15 - # tile_ref for column 2: base_ptr=0x1004, shape=(4,), strides=[4] - # flat offsets: [0*4, 1*4, 2*4, 3*4] = [0, 4, 8, 12] - # span = 13 (max offset + 1) - # mem.read(0x1004, 13) -> [2,3,4,5,6,7,8,9,10,11,12,13,14] - # gathered = flat[[0,4,8,12]] = [2, 6, 10, 14] ✓ - - Example — upper-triangular load from a 4×4 tile (coords provided):: - - # tile_ref: base_ptr=0x1000, shape=(4,4), strides=[4,1] - # coords = [(0,0),(0,1),...,(3,3)] — 10 upper-tri tuples - # flat offsets = [0*4+0, 0*4+1, ..., 3*4+3] = [0,1,2,3,5,6,7,10,11,15] - # span = 16 - # mem.read(0x1000, 16) -> flat 0..15 - # gathered = flat[[0,1,2,3,5,6,7,10,11,15]] = [0,1,2,3,5,6,7,10,11,15] - Args: context: Core execution context tile_ref: Tile reference (memref) describing source coords: Optional list of local coordinate tuples to gather. Each tuple is 0-based within tile_ref.shape. - result_shape: Output shape when coords is given; defaults to - tile_ref.shape when coords is None. + offsets: Optional pre-computed flat element offsets (int64 ndarray). + Mutually exclusive with coords. + result_shape: Output shape; defaults to tile_ref.shape when + neither coords nor offsets is given. Returns: Tile value (tensor) loaded into LX @@ -519,6 +759,18 @@ def load( mgr = _MemAccessor(context, tile_ref.memref.memory_space, tile_ref.base_ptr, tile_ref.memref.lx_core_id) stick_bytes = mgr.stick_bytes + # Pre-computed offsets path (blocked-indirect fast path). + if offsets is not None: + bpe = _bytes_per_elem(tile_ref.dtype) + unique_sticks = _MemAccessor.count_sticks_array( + tile_ref.memref.memory_space, tile_ref.base_ptr, offsets, bpe, + ) + gathered = mgr.read(len(offsets), tile_ref.dtype, offsets=offsets) + out_shape = result_shape if result_shape is not None else tile_ref.shape + data = gathered.reshape(out_shape) + MemoryOps._write_to_lx(context, data) + return Tile(data, tile_ref.dtype, out_shape, unique_sticks) + # Fast path: contiguous tile, no coord filtering — single dict-key read. if coords is None and MemoryOps._is_contiguous(tile_ref.shape, tile_ref.strides): n = int(np.prod(tile_ref.shape)) @@ -535,15 +787,12 @@ def load( unique_sticks = None return Tile(data, tile_ref.dtype, tile_ref.shape, unique_sticks) - # Strided or coord-set path: linearize coords, single read, numpy fancy-index. + # Strided or coord-set path: linearize coords → sparse read via offsets. offsets, unique_sticks = MemoryOps._flat_memory_offsets( tile_ref.base_ptr, tile_ref.shape, tile_ref.strides, tile_ref.dtype, coords, stick_bytes=stick_bytes ) - span = int(offsets.max()) + 1 if offsets.size else 1 - flat = mgr.read(span, tile_ref.dtype) - - gathered = flat[offsets] + gathered = mgr.read(len(offsets), tile_ref.dtype, offsets=offsets) out_shape = result_shape if result_shape is not None else tile_ref.shape data = gathered.reshape(out_shape) @@ -556,46 +805,47 @@ def store( tile: Tile, tile_ref: TileRef, coords: Optional[List[Tuple[int, ...]]] = None, + offsets: Optional[np.ndarray] = None, ) -> int: """Store tile data to HBM or LX. - HBM target → DMA write from LX to HBM. - LX target → write directly to LX. - When *coords* is given (coordinate-set path), scatters tile elements - to those local coordinates via a read-modify-write on the allocation. - When *coords* is None, stores the full tile (contiguous or strided). - - Source data layout: ``tile.data`` is read in C-order via - ``numpy.ndarray.flatten()``, which always returns a contiguous copy. - Non-contiguous source arrays are handled internally — callers do not - need to pre-``ascontiguousarray`` the tile. When *coords* is supplied, - ``coords[i]`` receives the i-th element of ``tile.data`` in C-order. - - A single ``mem.read`` + ``mem.write`` covers the entire footprint; - no per-element dict scans occur. + Three dispatch modes (checked in order): + 1. *offsets* — pre-computed flat element offsets (blocked-indirect fast + path). Skips coordinate linearization entirely. + 2. *coords* — scatters tile elements to those coordinates via + read-modify-write on the allocation. + 3. Neither — stores the full tile (contiguous or strided). Args: context: Core execution context tile: Tile value (tensor data) to store tile_ref: Tile reference (memref) describing destination coords: Optional list of local coordinate tuples to scatter into. + offsets: Optional pre-computed flat element offsets (int64 ndarray). + Mutually exclusive with coords. Returns: ``unique_sticks`` (int) — the number of distinct 128-byte HBM - sticks the write touches. ``0`` for LX destinations (no stick - concept; LX HBM traffic is zero by definition). The dialect - handler returns this value so :meth:`LatencyTracker._data_size` - charges HBM traffic at stick granularity - (``unique_sticks * STICK_BYTES``) instead of the source tile's - logical ``nbytes``, which would undercount scatter writes. + sticks the write touches. ``0`` for LX destinations. """ mgr = _MemAccessor(context, tile_ref.memref.memory_space, tile_ref.base_ptr, tile_ref.memref.lx_core_id) stick_bytes = mgr.stick_bytes + # Pre-computed offsets path (blocked-indirect fast path). + if offsets is not None: + bpe = _bytes_per_elem(tile_ref.dtype) + unique_sticks = _MemAccessor.count_sticks_array( + tile_ref.memref.memory_space, tile_ref.base_ptr, offsets, bpe, + ) + mgr.write(tile.data.ravel(), offsets=offsets) + return unique_sticks if unique_sticks is not None else 0 + # Fast path: contiguous tile, no coord filtering — single dict-key write. if coords is None and MemoryOps._is_contiguous(tile_ref.shape, tile_ref.strides): - mgr.write(tile.data.ravel()) # write reads it (copies into store) — view is fine + mgr.write(tile.data.ravel()) if not stick_bytes: return 0 n = int(np.prod(tile_ref.shape)) @@ -606,15 +856,12 @@ def store( - tile_ref.base_ptr // stick_bytes ) - # Strided or coord-set path: read-modify-write via scatter offsets. + # Strided or coord-set path: sparse write via offsets. offsets, unique_sticks = MemoryOps._flat_memory_offsets( tile_ref.base_ptr, tile_ref.shape, tile_ref.strides, tile_ref.dtype, coords, stick_bytes=stick_bytes, ) - span = int(offsets.max()) + 1 if offsets.size else 1 - flat = mgr.read(span, tile_ref.dtype) - flat[offsets] = tile.data.ravel() # RHS read-only scatter source — view is fine - mgr.write(flat) + mgr.write(tile.data.ravel(), offsets=offsets) return unique_sticks if unique_sticks is not None else 0 @staticmethod @@ -643,16 +890,23 @@ def indirect_load( f"dimensions; got non-permutation map: {vso.source}" ) - # Resolve every idx-tensor read up front: one accessor per index - # view, one read_scattered call, sticks deduped inside the accessor. - # Both helpers route their pt enumeration through - # _enumerate_in_vso_order, so non-identity vso permutes the - # iteration order consistently across idx reads and coord build - # (RFC 0682 §473). + out_shape = result_shape if result_shape is not None else iat.shape + + # Fast path: blocked-indirect patterns (MoE, paged attention) where the + # index lookup depends on a small subset of iteration variables. + # Bypasses the O(N) Python loops in _resolve_idx_reads / _build_indirect_coords. + block_info = _analyze_blocked_indirect(iat) + if block_info is not None: + offsets, idx_sticks = _compute_blocked_indirect_offsets(context, iat, block_info) + tile_ref = iat.parent_ref.to_tile_ref() + result = MemoryOps.load(context, tile_ref, offsets=offsets, result_shape=out_shape) + result.index_unique_sticks = idx_sticks + return result + + # General path: O(N) Python-loop idx reads + coord build. idx_values, idx_unique_sticks = _resolve_idx_reads(context, iat) coords = _build_indirect_coords(iat, idx_values) - out_shape = result_shape if result_shape is not None else iat.shape result = MemoryOps.load( context, iat.parent_ref.to_tile_ref(), coords=coords, result_shape=out_shape, @@ -660,6 +914,55 @@ def indirect_load( result.index_unique_sticks = idx_unique_sticks return result + @staticmethod + def indirect_store( + context: CoreContext, + tile: Tile, + iat: "IndirectAccessTile", + ) -> int: + """Store data using an indirect access tile (scatter pattern). + + Mirror of :meth:`indirect_load`. Enumerates the variable space, + resolves each coordinate tuple (direct dims use the variable value, + indirect dims look up the index in an index memref), then delegates + to :meth:`store`. + + Returns: + Total ``unique_sticks`` touched on HBM — sum of the parent + tile's destination sticks (from :meth:`store`) and the + idx-side sticks (from :func:`_resolve_idx_reads`). + """ + # MLIR type system should already enforce shape match; raise here so a + # mismatch surfaces clearly instead of as an opaque NumPy shape error. + if tuple(tile.shape) != tuple(iat.shape): + raise ValueError( + f"indirect_store: source tile shape {tuple(tile.shape)} does not " + f"match IAT shape {tuple(iat.shape)}" + ) + + vso = iat.variables_space_order + if vso is not None and not vso.is_permutation(): + raise ValueError( + f"indirect_store: variables_space_order must permute its input " + f"dimensions; got non-permutation map: {vso.source}" + ) + + # Fast path: blocked-indirect patterns. + block_info = _analyze_blocked_indirect(iat) + if block_info is not None: + offsets, idx_sticks = _compute_blocked_indirect_offsets(context, iat, block_info) + tile_ref = iat.parent_ref.to_tile_ref() + data_sticks = MemoryOps.store(context, tile, tile_ref, offsets=offsets) + return data_sticks + idx_sticks + + # General path: O(N) Python-loop idx reads + coord build. + idx_values, idx_unique_sticks = _resolve_idx_reads(context, iat) + coords = _build_indirect_coords(iat, idx_values) + data_sticks = MemoryOps.store( + context, tile, iat.parent_ref.to_tile_ref(), coords=coords, + ) + return data_sticks + idx_unique_sticks + # ------------------------------------------------------------------ # Distributed memory views (RFC 0682 §3.3) # @@ -874,15 +1177,13 @@ def distributed_load( survivor.base_ptr, survivor.shape, survivor.strides, survivor.dtype, local_coords, stick_bytes=mgr.stick_bytes, ) - span = int(offsets.max()) + 1 if offsets.size else 1 - flat = mgr.read(span, survivor.dtype) - # Vectorized scatter: per-dimension index arrays → one fancy-index write. + gathered = mgr.read(len(offsets), survivor.dtype, offsets=offsets) out_idx = tuple( np.fromiter((c[d] for c in access_coords), dtype=np.intp, count=len(access_coords)) for d in range(ndim) ) - out[out_idx] = flat[offsets] + out[out_idx] = gathered if unique_sticks is not None: total_unique_sticks += unique_sticks @@ -946,71 +1247,13 @@ def distributed_store( survivor.base_ptr, survivor.shape, survivor.strides, survivor.dtype, local_coords, stick_bytes=mgr.stick_bytes, ) - span = int(offsets.max()) + 1 if offsets.size else 1 - flat = mgr.read(span, survivor.dtype) - # Vectorized gather/scatter: per-dimension index arrays → one fancy-index read+write. src_idx = tuple( np.fromiter((c[d] for c in access_coords), dtype=np.intp, count=len(access_coords)) for d in range(ndim) ) - flat[offsets] = tile.data[src_idx] - mgr.write(flat) + mgr.write(tile.data[src_idx], offsets=offsets) if unique_sticks is not None: total_unique_sticks += unique_sticks return total_unique_sticks - - @staticmethod - def indirect_store( - context: CoreContext, - tile: Tile, - iat: "IndirectAccessTile", - ) -> int: - """Store data using an indirect access tile (scatter pattern). - - Mirror of :meth:`indirect_load`. Enumerates the variable space, - resolves each coordinate tuple (direct dims use the variable value, - indirect dims look up the index in an index memref), then delegates - to :meth:`store`. - - Coordinate collisions (multiple source elements mapping to the same - destination coordinate) are *implementation-defined*; the current - behavior is last-writer-wins via NumPy fancy-index assignment. - - Returns: - Total ``unique_sticks`` touched on HBM — sum of the parent - tile's destination sticks (from :meth:`store`) and the - idx-side sticks (from :func:`_resolve_idx_reads`). ``0`` when - both the parent and every idx view live in LX (no HBM - traffic). Returned via the dialect handler as the op result - so :meth:`LatencyTracker._data_size` can charge stick-granular - HBM cost — guard symmetry with :meth:`indirect_load`, which - stamps the same totals on the result Tile. - """ - # MLIR type system should already enforce shape match; raise here so a - # mismatch surfaces clearly instead of as an opaque NumPy shape error. - if tuple(tile.shape) != tuple(iat.shape): - raise ValueError( - f"indirect_store: source tile shape {tuple(tile.shape)} does not " - f"match IAT shape {tuple(iat.shape)}" - ) - - vso = iat.variables_space_order - if vso is not None and not vso.is_permutation(): - raise ValueError( - f"indirect_store: variables_space_order must permute its input " - f"dimensions; got non-permutation map: {vso.source}" - ) - - # Resolve idx reads (returns idx_unique_sticks: int, 0 for all-LX - # views) and delegate the data write to MemoryOps.store (returns - # int: HBM stick count, 0 for LX). Both helpers enumerate via - # _enumerate_in_vso_order so non-identity vso permutes the - # iteration order consistently with indirect_load (RFC 0682 §473). - idx_values, idx_unique_sticks = _resolve_idx_reads(context, iat) - coords = _build_indirect_coords(iat, idx_values) - data_sticks = MemoryOps.store( - context, tile, iat.parent_ref.to_tile_ref(), coords=coords, - ) - return data_sticks + idx_unique_sticks diff --git a/tests/test_blocked_indirect_fast_path.py b/tests/test_blocked_indirect_fast_path.py new file mode 100644 index 0000000..c59c0f0 --- /dev/null +++ b/tests/test_blocked_indirect_fast_path.py @@ -0,0 +1,1031 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the blocked-indirect fast path in indirect_load / indirect_store. + +Notation: W[E[e], m, n] means "for each e, look up row E[e] from W across +all m, n." See _analyze_blocked_indirect docstring for full terminology +(index view, dep_var, blocking factor, fast vs general path). + +Covers: + - Classifier gating: accepts qualifying patterns, rejects others + - Load correctness: 1-indirect, 2-indirect, compound idx_exprs, direct_expr + - Store correctness + - Equivalence with general path + - Shared-view patterns (multiple indirect subs on the same index array) + - Edge cases +""" + +import numpy as np +import pytest + +from ktir_cpu.affine import BoxSet +from ktir_cpu.ir_types import MemRef, IndirectAccessTile, Tile +from ktir_cpu.grid import CoreContext +from ktir_cpu.memory import HBMSimulator, LXScratchpad +from ktir_cpu.ops.memory_ops import ( + MemoryOps, + _expr_dependent_vars, + _analyze_blocked_indirect, + _resolve_idx_reads, + _build_indirect_coords, +) +from ktir_cpu.dtypes import bytes_per_elem +from ktir_cpu.parser_ast import parse_affine_map + + +# --------------------------------------------------------------------------- +# Test helpers +# +# _isub(view_idx, *dims) — build an indirect subscript referencing a view +# _dsub(var_idx) — build a direct subscript for a variables-space dim +# _make_iat(...) — assemble a full IAT from parts (parent, shape, subs) +# _alloc_idx(hbm, vals) — write an i32 index array to HBM, return its MemRef +# --------------------------------------------------------------------------- + +_BPE_F16 = bytes_per_elem("f16") +_BPE_I32 = bytes_per_elem("i32") + + +def _make_context(): + hbm = HBMSimulator() + lx = LXScratchpad(size_mb=64) + return CoreContext(core_id=0, grid_pos=(0, 0, 0), lx=lx, hbm=hbm) + + +def _alloc_hbm(hbm, data, dtype): + """Allocate in HBM, write data, return element-index base_ptr.""" + bpe = bytes_per_elem(dtype) + stick = hbm.allocate(data.nbytes) + hbm.write(stick, data) + return (stick * HBMSimulator.STICK_BYTES) // bpe, stick + + +def _alloc_idx(hbm, indices): + """Allocate i32 index array in HBM, return MemRef.""" + data = np.asarray(indices, dtype=np.int32) + ptr, _ = _alloc_hbm(hbm, data, "i32") + return MemRef(base_ptr=ptr, shape=(len(data),), strides=[1], + memory_space="HBM", dtype="i32") + + +def _isub(view_idx, *dims): + """Indirect subscript: read index_views[view_idx] at given iteration dims.""" + return {"kind": "indirect", "index_view_idx": view_idx, + "idx_exprs": [("dim", d) for d in dims]} + + +def _dsub(var_idx): + """Direct subscript: contiguous range over variables-space dim var_idx.""" + return {"kind": "direct", "var_index": var_idx} + + +def _make_iat(parent_ref, shape, dim_subscripts, index_views, vso=None): + """Build an IndirectAccessTile with BoxSet(lo=0, hi=shape).""" + vss = BoxSet(lo=tuple(0 for _ in shape), hi=shape) + return IndirectAccessTile( + parent_ref=parent_ref, shape=shape, + dim_subscripts=dim_subscripts, index_views=index_views, + variables_space_set=vss, variables_space_order=vso, + ) + + +# --------------------------------------------------------------------------- +# _expr_dependent_vars: identifies which iteration dims an index expression +# depends on. This determines K (number of unique lookups) for the fast path. +# --------------------------------------------------------------------------- + +class TestExprDependentVars: + def test_simple_dim(self): + assert _expr_dependent_vars(("dim", 0)) == {0} + assert _expr_dependent_vars(("dim", 2)) == {2} + + def test_const(self): + assert _expr_dependent_vars(("const", 42)) == set() + + def test_ssa(self): + assert _expr_dependent_vars(("ssa", "%grid0")) == set() + + def test_add_two_dims(self): + expr = ("add", ("dim", 0), ("dim", 1)) + assert _expr_dependent_vars(expr) == {0, 1} + + def test_add_dim_const(self): + expr = ("add", ("ssa", "%c0"), ("dim", 0)) + assert _expr_dependent_vars(expr) == {0} + + def test_floordiv(self): + expr = ("floordiv", ("dim", 2), 64) + assert _expr_dependent_vars(expr) == {2} + + def test_mod(self): + expr = ("mod", ("dim", 2), 64) + assert _expr_dependent_vars(expr) == {2} + + def test_mul(self): + expr = ("mul", 4, ("dim", 1)) + assert _expr_dependent_vars(expr) == {1} + + def test_compound_paged_attn(self): + expr1 = ("const", 0) + expr2 = ("add", ("ssa", "%bt_idx"), ("dim", 0)) + assert _expr_dependent_vars(expr1) == set() + assert _expr_dependent_vars(expr2) == {0} + + +# --------------------------------------------------------------------------- +# Classifier gating: _analyze_blocked_indirect decides whether an IAT +# qualifies for the fast path (returns info tuple) or must fall through +# to the general per-point path (returns None). +# --------------------------------------------------------------------------- + +class TestBlockedIndirectGating: + def test_accepted_1_indirect(self): + """X[IDX[e], m, n] — 1 indirect + 2 direct, ratio=8192× → accepted.""" + x_memref = MemRef(base_ptr=0, shape=(128, 64, 128), strides=[8192, 128, 1], + memory_space="HBM", dtype="f16") + idx_memref = MemRef(base_ptr=10000, shape=(8,), strides=[1], + memory_space="HBM", dtype="i32") + iat = _make_iat( + x_memref, (8, 64, 128), + [_isub(0, 0), _dsub(1), _dsub(2)], + [idx_memref], + ) + assert _analyze_blocked_indirect(iat) is not None + + def test_accepted_2_indirect(self): + """W[E[e], H[h], m, n] — 2 indirect + 2 direct → accepted.""" + data_memref = MemRef(base_ptr=0, shape=(8, 4, 16, 32), + strides=[2048, 512, 32, 1], + memory_space="HBM", dtype="f16") + e_memref = MemRef(base_ptr=5000, shape=(3,), strides=[1], + memory_space="HBM", dtype="i32") + h_memref = MemRef(base_ptr=6000, shape=(2,), strides=[1], + memory_space="HBM", dtype="i32") + iat = _make_iat( + data_memref, (3, 2, 16, 32), + [_isub(0, 0), _isub(1, 1), _dsub(2), _dsub(3)], + [e_memref, h_memref], + ) + # unique=3*2=6, total=3*2*16*32=3072, ratio=512× → qualifies + assert _analyze_blocked_indirect(iat) is not None + + def test_rejected_no_direct_dims(self): + """X[IDX1[i], IDX2[j]] — all indirect, no block → rejected.""" + x_memref = MemRef(base_ptr=0, shape=(4, 4), strides=[4, 1], + memory_space="HBM", dtype="f16") + idx1_memref = MemRef(base_ptr=1000, shape=(4, 4), strides=[4, 1], + memory_space="HBM", dtype="i32") + idx2_memref = MemRef(base_ptr=2000, shape=(4, 4), strides=[4, 1], + memory_space="HBM", dtype="i32") + dim_subscripts = [ + {"kind": "indirect", "index_view_idx": 0, "idx_exprs": [("dim", 0), ("dim", 1)]}, + {"kind": "indirect", "index_view_idx": 1, "idx_exprs": [("dim", 0), ("dim", 1)]}, + ] + iat = _make_iat(x_memref, (4, 4), dim_subscripts, [idx1_memref, idx2_memref]) + assert _analyze_blocked_indirect(iat) is None + + @pytest.mark.parametrize("unique_rows,direct_cols,expected", [ + (16, 4, False), # ratio=4×, well below threshold + ( 2, 15, False), # ratio=7.5×, just below threshold (2×16=32 > 30) + ( 1, 16, True), # ratio=16×, exactly at threshold (1×16=16 ≤ 16) + ]) + def test_ratio_threshold(self, unique_rows, direct_cols, expected): + """X[IDX[e], col] — varies the unique:total ratio around the 16× threshold.""" + x_memref = MemRef(base_ptr=0, shape=(unique_rows, direct_cols), + strides=[direct_cols, 1], + memory_space="HBM", dtype="f16") + idx_memref = MemRef(base_ptr=1000, shape=(unique_rows,), strides=[1], + memory_space="HBM", dtype="i32") + iat = _make_iat( + x_memref, (unique_rows, direct_cols), + [_isub(0, 0), _dsub(1)], + [idx_memref], + ) + assert (_analyze_blocked_indirect(iat) is not None) is expected + + +# --------------------------------------------------------------------------- +# Load correctness: each test constructs an IAT representing a real workload +# pattern, runs indirect_load (which takes the fast path), and compares the +# result against a manually-computed expected array. +# --------------------------------------------------------------------------- + +class TestBlockedIndirectLoad: + """Fast-path load produces correct data across all supported patterns. + + Pattern notation (test docstrings): + X[IDX[e], m, n] → 1 indirect dim (e selects rows), 2 direct (m, n) + cache[BT[0,d0], ...] → compound idx_expr (2D view, constant + dim) + W[E[e], H[h], m, n] → 2 independent indirect dims + 2 direct + """ + + def test_moe_1i_2d(self): + """X[IDX[e], M, N] — 8 experts from 128×8×16 weight tensor.""" + ctx = _make_context() + hbm = ctx.hbm + + num_experts, M, N = 128, 8, 16 + x_data = np.random.randn(num_experts * M * N).astype(np.float16) + x_base_ptr, _ = _alloc_hbm(hbm, x_data, "f16") + + selected = np.array([0, 15, 33, 64, 77, 99, 111, 127], dtype=np.int32) + idx_memref = _alloc_idx(hbm, selected) + + x_memref = MemRef(base_ptr=x_base_ptr, shape=(num_experts, M, N), + strides=[M * N, N, 1], memory_space="HBM", dtype="f16") + iat = _make_iat( + x_memref, (8, M, N), + [_isub(0, 0), _dsub(1), _dsub(2)], + [idx_memref], + ) + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + expected = x_data.reshape(num_experts, M, N)[selected, :, :] + np.testing.assert_array_equal(tile.data, expected) + assert tile.index_unique_sticks == 1 # 8 i32 elements = 32 bytes < STICK_BYTES + + def test_paged_attn_compound_idx(self): + """cache[BT[0, d0], d1, d2, d3] — compound idx_exprs, 1i + 3d.""" + ctx = _make_context() + hbm = ctx.hbm + + n_pages, n_heads, block_size, head_dim = 8, 4, 2, 16 + cache_data = np.arange(n_pages * n_heads * block_size * head_dim, dtype=np.float16) + cache_base_ptr, _ = _alloc_hbm(hbm, cache_data, "f16") + + bt_data = np.array([5, 2, 7, 0], dtype=np.int32) + bt_base_ptr, _ = _alloc_hbm(hbm, bt_data, "i32") + + cache_memref = MemRef( + base_ptr=cache_base_ptr, + shape=(n_pages, n_heads, block_size, head_dim), + strides=[n_heads * block_size * head_dim, block_size * head_dim, head_dim, 1], + memory_space="HBM", dtype="f16", + ) + bt_memref = MemRef(base_ptr=bt_base_ptr, shape=(1, 4), strides=[4, 1], + memory_space="HBM", dtype="i32") + + dim_subscripts = [ + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("const", 0), ("dim", 0)]}, + _dsub(1), _dsub(2), _dsub(3), + ] + iat = _make_iat( + cache_memref, (4, n_heads, block_size, head_dim), + dim_subscripts, [bt_memref], + ) + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + cache_arr = cache_data.reshape(n_pages, n_heads, block_size, head_dim) + expected = cache_arr[bt_data, :, :, :] + np.testing.assert_array_equal(tile.data, expected) + + def test_sparse_attn_2i_1d(self): + """cache[page_idx[b], token_idx[t], d] — 2 indirect + 1 direct.""" + ctx = _make_context() + hbm = ctx.hbm + + n_pages, n_tokens, hidden = 8, 6, 32 + n_sel_p, n_sel_t = 4, 3 + + data = np.arange(n_pages * n_tokens * hidden, dtype=np.float16) + base_ptr, _ = _alloc_hbm(hbm, data, "f16") + + page_sel = np.sort(np.random.choice(n_pages, n_sel_p, replace=False)).astype(np.int32) + page_memref = _alloc_idx(hbm, page_sel) + + token_sel = np.sort(np.random.choice(n_tokens, n_sel_t, replace=False)).astype(np.int32) + token_memref = _alloc_idx(hbm, token_sel) + + data_memref = MemRef(base_ptr=base_ptr, shape=(n_pages, n_tokens, hidden), + strides=[n_tokens * hidden, hidden, 1], + memory_space="HBM", dtype="f16") + iat = _make_iat( + data_memref, (n_sel_p, n_sel_t, hidden), + [_isub(0, 0), _isub(1, 1), _dsub(2)], + [page_memref, token_memref], + ) + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + full = data.reshape(n_pages, n_tokens, hidden) + expected = full[np.ix_(page_sel, token_sel, np.arange(hidden))] + np.testing.assert_array_equal(tile.data, expected) + + def test_multi_head_2i_2d(self): + """W[E[e], H[h], m, n] — 2 indirect + 2 direct.""" + ctx = _make_context() + hbm = ctx.hbm + n_exp, n_h, M, N = 8, 4, 16, 32 + n_sel_e, n_sel_h = 3, 2 + + data = np.arange(n_exp * n_h * M * N, dtype=np.float16) + base, _ = _alloc_hbm(hbm, data, "f16") + + e_sel = np.array([1, 3, 7], dtype=np.int32) + e_memref = _alloc_idx(hbm, e_sel) + + h_sel = np.array([0, 3], dtype=np.int32) + h_memref = _alloc_idx(hbm, h_sel) + + data_memref = MemRef(base_ptr=base, shape=(n_exp, n_h, M, N), + strides=[n_h * M * N, M * N, N, 1], + memory_space="HBM", dtype="f16") + iat = _make_iat( + data_memref, (n_sel_e, n_sel_h, M, N), + [_isub(0, 0), _isub(1, 1), _dsub(2), _dsub(3)], + [e_memref, h_memref], + ) + tile = MemoryOps.indirect_load(ctx, iat) + + full = data.reshape(n_exp, n_h, M, N) + expected = full[np.ix_(e_sel, h_sel, np.arange(M), np.arange(N))] + np.testing.assert_array_equal(tile.data, expected) + assert tile.index_unique_sticks == 2 # e_sel: 1 stick; h_sel: 1 stick + + def test_direct_expr(self): + """X[IDX[e], (2*m+1)] — indirect + direct_expr: rejected by classifier, + falls through to general path, result still correct.""" + ctx = _make_context() + hbm = ctx.hbm + + x_data = np.arange(64 * 128, dtype=np.float16) + x_base_ptr, _ = _alloc_hbm(hbm, x_data, "f16") + + idx_data = np.array([3, 7, 50, 63], dtype=np.int32) + idx_memref = _alloc_idx(hbm, idx_data) + + x_memref = MemRef(base_ptr=x_base_ptr, shape=(64, 128), strides=[128, 1], + memory_space="HBM", dtype="f16") + + dim_subscripts = [ + _isub(0, 0), + {"kind": "direct_expr", "subscript": ("add", ("mul", 2, ("dim", 1)), ("const", 1))}, + ] + vss = BoxSet(lo=(0, 0), hi=(4, 60)) + iat = IndirectAccessTile( + parent_ref=x_memref, shape=(4, 60), + dim_subscripts=dim_subscripts, index_views=[idx_memref], + variables_space_set=vss, variables_space_order=None, + ) + # direct_expr causes classifier rejection + assert _analyze_blocked_indirect(iat) is None + tile = MemoryOps.indirect_load(ctx, iat) + + x_arr = x_data.reshape(64, 128) + expected = np.zeros((4, 60), dtype=np.float16) + for e in range(4): + for m in range(60): + expected[e, m] = x_arr[idx_data[e], 2 * m + 1] + np.testing.assert_array_equal(tile.data, expected) + + +# --------------------------------------------------------------------------- +# Store correctness: the fast path also handles indirect_store (scatter). +# Writes should land only at the indirectly-selected positions. +# --------------------------------------------------------------------------- + +class TestBlockedIndirectStore: + def test_scatter_write(self): + """W[E[e], H[h], m, n] = tile — verifies scatter writes back correctly.""" + ctx = _make_context() + hbm = ctx.hbm + n_exp, n_h, M, N = 8, 4, 16, 32 + n_sel_e, n_sel_h = 3, 2 + + data = np.zeros(n_exp * n_h * M * N, dtype=np.float16) + base, stick = _alloc_hbm(hbm, data, "f16") + + e_sel = np.array([1, 3, 7], dtype=np.int32) + e_memref = _alloc_idx(hbm, e_sel) + + h_sel = np.array([0, 3], dtype=np.int32) + h_memref = _alloc_idx(hbm, h_sel) + + data_memref = MemRef(base_ptr=base, shape=(n_exp, n_h, M, N), + strides=[n_h * M * N, M * N, N, 1], + memory_space="HBM", dtype="f16") + iat = _make_iat( + data_memref, (n_sel_e, n_sel_h, M, N), + [_isub(0, 0), _isub(1, 1), _dsub(2), _dsub(3)], + [e_memref, h_memref], + ) + + write_data = np.ones((n_sel_e, n_sel_h, M, N), dtype=np.float16) * 42.0 + write_tile = Tile(write_data, "f16", write_data.shape, 0) + MemoryOps.indirect_store(ctx, write_tile, iat) + + full = hbm.read(stick, n_exp * n_h * M * N, "f16").reshape(n_exp, n_h, M, N) + for ei in e_sel: + for hi in h_sel: + np.testing.assert_array_equal(full[ei, hi], 42.0) + for ei in range(n_exp): + for hi in range(n_h): + if ei not in e_sel or hi not in h_sel: + np.testing.assert_array_equal(full[ei, hi], 0.0) + + +# --------------------------------------------------------------------------- +# Sparse write primitives: unit tests for the underlying HBM/LX write(offsets=) +# ops (these are the building blocks that indirect_store's fast path uses). +# --------------------------------------------------------------------------- + +class TestSparseWrite: + """write(offsets=) writes only the targeted offsets, leaving the rest untouched.""" + + def test_hbm_write_sparse(self): + """Write into a few elements of a larger HBM allocation via offsets.""" + hbm = HBMSimulator() + data = np.zeros(64, dtype=np.float16) + stick = hbm.allocate(data.nbytes) + hbm.write(stick, data) + + offsets = np.array([5, 17, 42, 63], dtype=np.int64) + values = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16) + hbm.write(stick, values, offsets=offsets) + + result = hbm.read(stick, 64, "f16") + for o, v in zip(offsets, values): + assert result[o] == v + untouched = np.delete(np.arange(64), offsets) + np.testing.assert_array_equal(result[untouched], 0.0) + + def test_lx_write_sparse(self): + """Write into LX scratchpad allocation via offsets.""" + lx = LXScratchpad(size_mb=1) + data = np.zeros(32, dtype=np.float16) + lx.write(0, data) + + offsets = np.array([0, 15, 31], dtype=np.int64) + values = np.array([10.0, 20.0, 30.0], dtype=np.float16) + lx.write(0, values, offsets=offsets) + + result = lx.read(0, 32, "f16", offsets=np.arange(32, dtype=np.int64)) + for o, v in zip(offsets, values): + assert result[o] == v + untouched = np.delete(np.arange(32), offsets) + np.testing.assert_array_equal(result[untouched], 0.0) + + def test_write_read_sparse_roundtrip(self): + """write(offsets=) then read(offsets=) at same offsets roundtrips.""" + hbm = HBMSimulator() + data = np.random.randn(128).astype(np.float16) + stick = hbm.allocate(data.nbytes) + hbm.write(stick, data) + + offsets = np.array([10, 50, 100, 127], dtype=np.int64) + new_vals = np.array([99.0, 88.0, 77.0, 66.0], dtype=np.float16) + hbm.write(stick, new_vals, offsets=offsets) + + gathered = hbm.read(stick, len(offsets), "f16", offsets=offsets) + np.testing.assert_array_equal(gathered, new_vals) + + +# --------------------------------------------------------------------------- +# Equivalence: the fast path must produce bit-exact results compared to the +# general per-point path. This is the primary correctness oracle — if the +# fast path ever diverges, this test catches it. +# --------------------------------------------------------------------------- + +class TestBlockedIndirectMatchesGeneral: + def test_fast_equals_general(self): + """Fast path result is bit-exact with general inspector-executor.""" + ctx = _make_context() + hbm = ctx.hbm + + n_pages, n_tokens, hidden = 8, 6, 32 + n_sel_p, n_sel_t = 4, 3 + + data = np.arange(n_pages * n_tokens * hidden, dtype=np.float16) + base_ptr, _ = _alloc_hbm(hbm, data, "f16") + + page_sel = np.sort(np.random.choice(n_pages, n_sel_p, replace=False)).astype(np.int32) + page_memref = _alloc_idx(hbm, page_sel) + + token_sel = np.sort(np.random.choice(n_tokens, n_sel_t, replace=False)).astype(np.int32) + token_memref = _alloc_idx(hbm, token_sel) + + data_memref = MemRef(base_ptr=base_ptr, shape=(n_pages, n_tokens, hidden), + strides=[n_tokens * hidden, hidden, 1], + memory_space="HBM", dtype="f16") + iat = _make_iat( + data_memref, (n_sel_p, n_sel_t, hidden), + [_isub(0, 0), _isub(1, 1), _dsub(2)], + [page_memref, token_memref], + ) + + ctx.lx.memory.clear() + ctx.lx.next_ptr = 0 + fast_tile = MemoryOps.indirect_load(ctx, iat) + + ctx.lx.memory.clear() + ctx.lx.next_ptr = 0 + idx_values, _ = _resolve_idx_reads(ctx, iat) + coords = _build_indirect_coords(iat, idx_values) + general_tile = MemoryOps.load(ctx, iat.parent_ref.to_tile_ref(), + coords=coords, result_shape=iat.shape) + + np.testing.assert_array_equal(fast_tile.data, general_tile.data) + + +# --------------------------------------------------------------------------- +# Non-identity VSO (variables_space_order): when the iteration order differs +# from the natural dimension order (e.g., iterating columns-first), the +# meshgrid broadcast would produce wrong results → classifier rejects. +# The general path handles this correctly via explicit per-point evaluation. +# --------------------------------------------------------------------------- + +class TestBlockedIndirectPermutedVSO: + """Permuted VSO is rejected by _analyze_blocked_indirect, falling + through to the general inspector-executor path. + """ + + def test_permuted_vso_rejected_and_general_path_correct(self): + """W[E[e], n] with vso=(d1,d0): rejected by classifier, general path correct.""" + ctx = _make_context() + hbm = ctx.hbm + + n_exp, N, n_sel_e = 64, 128, 4 + data = np.arange(n_exp * N, dtype=np.float16) + base, _ = _alloc_hbm(hbm, data, "f16") + + e_sel = np.array([5, 17, 42, 63], dtype=np.int32) + idx_memref = _alloc_idx(hbm, e_sel) + + data_memref = MemRef(base_ptr=base, shape=(n_exp, N), strides=[N, 1], + memory_space="HBM", dtype="f16") + vso = parse_affine_map("affine_map<(d0, d1) -> (d1, d0)>") + assert not vso.is_identity() + + vss = BoxSet(lo=(0, 0), hi=(n_sel_e, N)) + iat = IndirectAccessTile( + parent_ref=data_memref, shape=(n_sel_e, N), + dim_subscripts=[_isub(0, 0), _dsub(1)], + index_views=[idx_memref], + variables_space_set=vss, variables_space_order=vso, + ) + + # classifier rejects non-identity VSO + assert _analyze_blocked_indirect(iat) is None + + # indirect_load falls through to general path + ctx.lx.memory.clear() + ctx.lx.next_ptr = 0 + result_tile = MemoryOps.indirect_load(ctx, iat) + + # explicit general path for reference + ctx.lx.memory.clear() + ctx.lx.next_ptr = 0 + idx_values, _ = _resolve_idx_reads(ctx, iat) + coords = _build_indirect_coords(iat, idx_values) + general_tile = MemoryOps.load(ctx, iat.parent_ref.to_tile_ref(), + coords=coords, result_shape=iat.shape) + + np.testing.assert_array_equal(result_tile.data, general_tile.data) + + +# --------------------------------------------------------------------------- +# Edge cases: degenerate inputs that should not crash +# --------------------------------------------------------------------------- + +class TestBlockedIndirectEdgeCases: + def test_empty_iteration_space(self): + """Zero-extent variable space should not crash.""" + ctx = _make_context() + hbm = ctx.hbm + + x_data = np.arange(64, dtype=np.float16) + x_base_ptr, _ = _alloc_hbm(hbm, x_data, "f16") + + idx_data = np.array([], dtype=np.int32) + idx_stick = hbm.allocate(max(idx_data.nbytes, 4)) + idx_base_ptr = (idx_stick * HBMSimulator.STICK_BYTES) // _BPE_I32 + + x_memref = MemRef(base_ptr=x_base_ptr, shape=(64, 4), strides=[4, 1], + memory_space="HBM", dtype="f16") + idx_memref = MemRef(base_ptr=idx_base_ptr, shape=(0,), strides=[1], + memory_space="HBM", dtype="i32") + + iat = _make_iat( + x_memref, (0, 4), + [_isub(0, 0), _dsub(1)], + [idx_memref], + ) + tile = MemoryOps.indirect_load(ctx, iat) + assert tile.data.size == 0 + + +# --------------------------------------------------------------------------- +# index_unique_sticks: the fast path counts how many HBM "sticks" (128-byte +# aligned cache lines) the index reads touch. The latency estimator uses +# this to model index-side memory traffic separately from data traffic. +# --------------------------------------------------------------------------- + +class TestBlockedIndirectIndexUniqueSticks: + """indirect_load populates Tile.index_unique_sticks for the estimator.""" + + def test_multi_stick_index_read(self): + """33 i32 index elements (132 bytes) cross a stick boundary → index_unique_sticks == 2. + + With STICK_BYTES=128 and bpe_i32=4: addresses e*4 for e in 0..32 span + bytes 0..128. Byte 128 lands on the next stick, so the set has 2 entries. + """ + ctx = _make_context() + hbm = ctx.hbm + + # 33 indirect * 32 direct = 1056 total, unique=33, ratio=32× > 16× → qualifies + num_experts, M = 256, 32 + x_data = np.arange(num_experts * M, dtype=np.float16) + x_base_ptr, _ = _alloc_hbm(hbm, x_data, "f16") + + # 33 * 4 = 132 bytes: elements 0-31 in stick N, element 32 in stick N+1 + idx_data = np.arange(33, dtype=np.int32) + idx_memref = _alloc_idx(hbm, idx_data) + + x_memref = MemRef(base_ptr=x_base_ptr, shape=(num_experts, M), + strides=[M, 1], memory_space="HBM", dtype="f16") + iat = _make_iat( + x_memref, (33, M), + [_isub(0, 0), _dsub(1)], + [idx_memref], + ) + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + assert tile.index_unique_sticks == 2 + + +# --------------------------------------------------------------------------- +# OOB handling: when indirect indices point outside the parent allocation, +# read(offsets=) returns zero (safe default) and write(offsets=) silently +# drops the write. This prevents crashes from stale or out-of-range index +# arrays. +# --------------------------------------------------------------------------- + +class TestSparseOOB: + """Verify that read(offsets=) zero-pads OOB and write(offsets=) drops them.""" + + def test_read_oob_returns_zero(self): + """Offsets past allocation end return zero.""" + from ktir_cpu.memory import _read_flat + memory = {0x1000: np.arange(10, dtype=np.float16)} + offsets = np.array([0, 5, 9, 10, 11], dtype=np.int64) + result = _read_flat(memory, 0x1000, len(offsets), np.float16, 2, offsets=offsets) + expected = np.array([0, 5, 9, 0, 0], dtype=np.float16) + np.testing.assert_array_equal(result, expected) + + def test_read_all_inbounds(self): + """All-inbounds path returns correct values (no OOB branch).""" + from ktir_cpu.memory import _read_flat + memory = {0x1000: np.arange(10, dtype=np.float16)} + offsets = np.array([0, 3, 7, 9], dtype=np.int64) + result = _read_flat(memory, 0x1000, len(offsets), np.float16, 2, offsets=offsets) + expected = np.array([0, 3, 7, 9], dtype=np.float16) + np.testing.assert_array_equal(result, expected) + + def test_write_oob_dropped(self): + """OOB offsets are silently dropped; inbounds writes land.""" + from ktir_cpu.memory import _write_flat + memory = {0x1000: np.zeros(10, dtype=np.float16)} + data = np.array([99, 88, 77], dtype=np.float16) + offsets = np.array([0, 10, 5], dtype=np.int64) + _write_flat(memory, 0x1000, data, offsets=offsets) + assert memory[0x1000][0] == 99 + assert memory[0x1000][5] == 77 + assert memory[0x1000][1] == 0 # untouched + + def test_write_all_inbounds(self): + """All-inbounds sparse write lands correctly.""" + from ktir_cpu.memory import _write_flat + memory = {0x1000: np.zeros(10, dtype=np.float16)} + data = np.array([11, 22, 33], dtype=np.float16) + offsets = np.array([1, 4, 8], dtype=np.int64) + _write_flat(memory, 0x1000, data, offsets=offsets) + assert memory[0x1000][1] == 11 + assert memory[0x1000][4] == 22 + assert memory[0x1000][8] == 33 + + +# --------------------------------------------------------------------------- +# Shared view: two or more dim_subscripts entries that reference the same +# index_view_idx — i.e., multiple subscriptions reading from one index array, +# possibly with different idx_exprs. Real examples: +# - Coordinate table: B is (K,3), access B[i,0], B[i,1], B[i,2] for 3D coords +# - Shifted window: B is 1D, access B[e] and B[e+1] for adjacent pairs +# - Diagonal: B[e], B[e] (degenerate — same value used in two parent dims) +# +# These patterns require per-subscription-expression re-keying (not per-view) +# so that each subscript gets its own K-element value array for broadcast. +# --------------------------------------------------------------------------- + +class TestBlockedIndirectSharedView: + """Fast path handles shared views correctly via per-sub re-keying.""" + + def test_shared_view_shifted(self): + """A[B[e], B[e+1], n] — shifted window into 1D index array.""" + ctx = _make_context() + hbm = ctx.hbm + + n_rows, n_cols, N = 16, 16, 32 + data = np.arange(n_rows * n_cols * N, dtype=np.float16) + data_ptr, _ = _alloc_hbm(hbm, data, "f16") + + K = 4 + idx_data = np.array([2, 5, 9, 13, 7], dtype=np.int32) + idx_ptr, _ = _alloc_hbm(hbm, idx_data, "i32") + idx_memref = MemRef(base_ptr=idx_ptr, shape=(5,), strides=[1], + memory_space="HBM", dtype="i32") + + data_memref = MemRef(base_ptr=data_ptr, shape=(n_rows, n_cols, N), + strides=[n_cols * N, N, 1], + memory_space="HBM", dtype="f16") + + # VSS has 2 dims: dim 0 = dep_var e (K), dim 1 = direct n (N) + dim_subscripts = [ + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0)]}, + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("add", ("dim", 0), ("const", 1))]}, + _dsub(1), + ] + iat = _make_iat(data_memref, (K, N), dim_subscripts, [idx_memref]) + + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + arr = data.reshape(n_rows, n_cols, N) + expected = np.zeros((K, N), dtype=np.float16) + for e in range(K): + expected[e, :] = arr[idx_data[e], idx_data[e + 1], :] + np.testing.assert_array_equal(tile.data, expected) + + def test_shared_view_2d_columns(self): + """A[B[i,0], B[i,1], B[i,2], n] — 2D coord table, different columns.""" + ctx = _make_context() + hbm = ctx.hbm + + d0, d1, d2, N = 8, 8, 8, 32 + data = np.arange(d0 * d1 * d2 * N, dtype=np.float16) + data_ptr, _ = _alloc_hbm(hbm, data, "f16") + + K = 3 + coord_table = np.array([[1, 3, 5], [2, 7, 0], [4, 1, 6]], dtype=np.int32) + coord_ptr, _ = _alloc_hbm(hbm, coord_table, "i32") + coord_memref = MemRef(base_ptr=coord_ptr, shape=(K, 3), strides=[3, 1], + memory_space="HBM", dtype="i32") + + data_memref = MemRef(base_ptr=data_ptr, shape=(d0, d1, d2, N), + strides=[d1 * d2 * N, d2 * N, N, 1], + memory_space="HBM", dtype="f16") + + # VSS has 2 dims: dim 0 = dep_var i (K), dim 1 = direct n (N) + dim_subscripts = [ + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0), ("const", 0)]}, + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0), ("const", 1)]}, + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0), ("const", 2)]}, + _dsub(1), + ] + iat = _make_iat(data_memref, (K, N), dim_subscripts, [coord_memref]) + + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + arr = data.reshape(d0, d1, d2, N) + expected = np.zeros((K, N), dtype=np.float16) + for i in range(K): + expected[i, :] = arr[coord_table[i, 0], coord_table[i, 1], coord_table[i, 2], :] + np.testing.assert_array_equal(tile.data, expected) + + def test_shared_view_same_expr(self): + """A[B[e], B[e], m, n] — degenerate same-expr diagonal access.""" + ctx = _make_context() + hbm = ctx.hbm + + n_rows, n_cols, M, N = 8, 8, 4, 16 + data = np.arange(n_rows * n_cols * M * N, dtype=np.float16) + data_ptr, _ = _alloc_hbm(hbm, data, "f16") + + K = 3 + idx_data = np.array([1, 5, 7], dtype=np.int32) + idx_memref = _alloc_idx(hbm, idx_data) + + data_memref = MemRef(base_ptr=data_ptr, shape=(n_rows, n_cols, M, N), + strides=[n_cols * M * N, M * N, N, 1], + memory_space="HBM", dtype="f16") + + dim_subscripts = [ + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0)]}, + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0)]}, + _dsub(1), _dsub(2), + ] + iat = _make_iat(data_memref, (K, M, N), dim_subscripts, [idx_memref]) + + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + arr = data.reshape(n_rows, n_cols, M, N) + expected = np.zeros((K, M, N), dtype=np.float16) + for e in range(K): + expected[e, :, :] = arr[idx_data[e], idx_data[e], :, :] + np.testing.assert_array_equal(tile.data, expected) + + +# --------------------------------------------------------------------------- +# Correctness gaps between fast path and general path. +# +# Three areas where the two paths diverge in capability or behavior: +# 1. Non-zero vss.lo — fast path starts dep-var iteration at lo[d], not 0 +# 2. Negative indices — _read_flat(offsets=) wraps silently; upstream guard prevents +# 3. Shared-view below threshold — general path crashes (StopIteration) +# --------------------------------------------------------------------------- + +class TestBlockedIndirectCorrectnessGaps: + """Tests for edge cases where fast path and general path diverge.""" + + @staticmethod + def _build_1d_iat(hbm, num_experts, N, idx_data, lo=None, hi=None): + """Helper: allocate W[E[e], n] with optional non-zero lo.""" + data = np.arange(num_experts * N, dtype=np.float16) + base, _ = _alloc_hbm(hbm, data, "f16") + idx_memref = _alloc_idx(hbm, idx_data) + data_memref = MemRef(base_ptr=base, shape=(num_experts, N), + strides=[N, 1], memory_space="HBM", dtype="f16") + if lo is None: + return _make_iat(data_memref, (len(idx_data), N), + [_isub(0, 0), _dsub(1)], [idx_memref]), data + K = hi[0] - lo[0] + vss = BoxSet(lo=lo, hi=hi) + iat = IndirectAccessTile( + parent_ref=data_memref, shape=(K, N), + dim_subscripts=[_isub(0, 0), _dsub(1)], + index_views=[idx_memref], + variables_space_set=vss, variables_space_order=None, + ) + return iat, data + + # --- Non-zero vss.lo --- + + def test_nonzero_lo_1d_indirect(self): + """W[E[e], n] with lo=(2,0) reads indices 2..5, not 0..3.""" + ctx = _make_context() + N = 16 + idx_data = np.array([10, 3, 7, 20, 15, 1, 28, 5], dtype=np.int32) + iat, data = self._build_1d_iat(ctx.hbm, 32, N, idx_data, + lo=(2, 0), hi=(6, N)) + + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + arr = data.reshape(32, N) + expected = np.stack([arr[idx_data[e]] for e in range(2, 6)]) + np.testing.assert_array_equal(tile.data, expected) + + def test_nonzero_lo_2d_indirect(self): + """cache[P[p], T[t], h] with lo=(1,2,0) — two indirect dims with non-zero lo.""" + ctx = _make_context() + n_pages, n_tokens, H = 16, 8, 16 # H≥16 passes gate 3 + data = np.arange(n_pages * n_tokens * H, dtype=np.float16) + base, _ = _alloc_hbm(ctx.hbm, data, "f16") + + page_idx = np.array([0, 9, 3, 14, 7, 11], dtype=np.int32) + token_idx = np.array([1, 5, 0, 7, 2, 6, 3, 4], dtype=np.int32) + page_memref = _alloc_idx(ctx.hbm, page_idx) + token_memref = _alloc_idx(ctx.hbm, token_idx) + + data_memref = MemRef(base_ptr=base, shape=(n_pages, n_tokens, H), + strides=[n_tokens * H, H, 1], + memory_space="HBM", dtype="f16") + vss = BoxSet(lo=(1, 2, 0), hi=(4, 6, H)) + iat = IndirectAccessTile( + parent_ref=data_memref, shape=(3, 4, H), + dim_subscripts=[_isub(0, 0), _isub(1, 1), _dsub(2)], + index_views=[page_memref, token_memref], + variables_space_set=vss, variables_space_order=None, + ) + + assert _analyze_blocked_indirect(iat) is not None + tile = MemoryOps.indirect_load(ctx, iat) + + arr = data.reshape(n_pages, n_tokens, H) + expected = np.zeros((3, 4, H), dtype=np.float16) + for pi, p in enumerate(range(1, 4)): + for ti, t in enumerate(range(2, 6)): + expected[pi, ti, :] = arr[page_idx[p], token_idx[t], :] + np.testing.assert_array_equal(tile.data, expected) + + def test_nonzero_lo_store_roundtrip(self): + """Store with non-zero lo, then load — verifies store uses lo correctly.""" + ctx = _make_context() + N = 16 # ≥16 passes gate 3 + idx_data = np.array([10, 3, 7, 20, 15, 1], dtype=np.int32) + iat, _ = self._build_1d_iat(ctx.hbm, 32, N, idx_data, + lo=(2, 0), hi=(5, N)) + + assert _analyze_blocked_indirect(iat) is not None + write_data = np.arange(3 * N, dtype=np.float16).reshape(3, N) + 100 + write_tile = Tile(data=write_data, shape=(3, N), dtype="f16") + ctx.lx.memory.clear() + ctx.lx.next_ptr = 0 + MemoryOps.indirect_store(ctx, write_tile, iat) + + ctx.lx.memory.clear() + ctx.lx.next_ptr = 0 + result = MemoryOps.indirect_load(ctx, iat) + np.testing.assert_array_equal(result.data, write_data) + + # --- Negative index guard --- + + def test_negative_index_in_idx_array_raises(self): + """_runtime_read_and_expand_sub_space raises IndexError on idx < 0.""" + ctx = _make_context() + N = 16 # ≥16 passes gate 3 + idx_data = np.array([3, -1, 7], dtype=np.int32) + iat, _ = self._build_1d_iat(ctx.hbm, 16, N, idx_data) + + assert _analyze_blocked_indirect(iat) is not None + with pytest.raises(IndexError, match="negative"): + MemoryOps.indirect_load(ctx, iat) + + def test_sparse_read_wraps_negative(self): + """_read_flat(offsets=) wraps negative offsets (NumPy behavior) — no guard.""" + from ktir_cpu.memory import _read_flat + memory = {0x1000: np.array([10, 20, 30, 40, 50], dtype=np.float16)} + offsets = np.array([0, -1, 2], dtype=np.int64) + result = _read_flat(memory, 0x1000, len(offsets), np.float16, 2, offsets=offsets) + assert result[0] == 10 + assert result[1] == 50 # wrap-around: flat[-1] = last element + assert result[2] == 30 + + # --- Shared-view below threshold (general path handles it correctly) --- + + def test_shared_view_below_threshold_general_path(self): + """A[B[e], B[e], n]: shared view falls to general path, produces diagonal.""" + ctx = _make_context() + n_rows, n_cols, N = 8, 8, 2 # N=2 → blocking factor < 16 + data = np.arange(n_rows * n_cols * N, dtype=np.float16) + data_ptr, _ = _alloc_hbm(ctx.hbm, data, "f16") + + K = 3 + idx_data = np.array([1, 5, 7], dtype=np.int32) + idx_memref = _alloc_idx(ctx.hbm, idx_data) + data_memref = MemRef(base_ptr=data_ptr, shape=(n_rows, n_cols, N), + strides=[n_cols * N, N, 1], + memory_space="HBM", dtype="f16") + dim_subscripts = [ + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0)]}, + {"kind": "indirect", "index_view_idx": 0, + "idx_exprs": [("dim", 0)]}, + _dsub(1), + ] + iat = _make_iat(data_memref, (K, N), dim_subscripts, [idx_memref]) + + assert _analyze_blocked_indirect(iat) is None + tile = MemoryOps.indirect_load(ctx, iat) + + arr = data.reshape(n_rows, n_cols, N) + expected = np.stack([arr[idx_data[e], idx_data[e]] for e in range(K)]) + np.testing.assert_array_equal(tile.data, expected) + + def test_distinct_view_below_threshold_uses_general_path(self): + """A[B[e], C[e], n]: distinct views work fine on general path.""" + ctx = _make_context() + d0, d1, N = 8, 8, 4 # N=4 → blocking factor < 16 + data = np.arange(d0 * d1 * N, dtype=np.float16) + data_ptr, _ = _alloc_hbm(ctx.hbm, data, "f16") + + idx_b = np.array([1, 5, 7], dtype=np.int32) + idx_c = np.array([3, 0, 6], dtype=np.int32) + b_memref = _alloc_idx(ctx.hbm, idx_b) + c_memref = _alloc_idx(ctx.hbm, idx_c) + data_memref = MemRef(base_ptr=data_ptr, shape=(d0, d1, N), + strides=[d1 * N, N, 1], + memory_space="HBM", dtype="f16") + iat = _make_iat(data_memref, (3, N), + [_isub(0, 0), _isub(1, 0), _dsub(1)], + [b_memref, c_memref]) + + assert _analyze_blocked_indirect(iat) is None + tile = MemoryOps.indirect_load(ctx, iat) + arr = data.reshape(d0, d1, N) + expected = np.stack([arr[idx_b[e], idx_c[e]] for e in range(3)]) + np.testing.assert_array_equal(tile.data, expected)