From c83acf8803679ee5acbb67ca53757f4da2cafeb2 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 26 Jun 2026 22:12:22 +0900 Subject: [PATCH 01/13] [Frontend] Carry gather/scatter offset as an explicit transfer descriptor Replace the affine.apply{indirect_access} symbol smuggle with an explicit offset descriptor. convert_indirect_indexing returns the offset spad instead of folding sympy.Symbol(out) into the index; emit_transfer carries it as a togsim.transfer operand; decompose_transfer lifts that operand to a memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start is a registered op and rejects an extra operand, but accepts the attribute); lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4 (drops _find_indirect); build_skeleton adds the offset spad to the gather DMA read_bufs so the offset-build -> gather dependency edge forms in the trace. The index stays clean (base only). Validated on both paths: Spike functional (computed-index gather + scatter allclose, pointwise/reduce regression 0) and the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now includes the offset spad). Indirect addressing (scattered-DMA timing) in the new trace path is a separate gap tracked in issue #284. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- .../mlir/mlir_codegen_backend.py | 37 +++++++++---------- .../mlir/passes/build_skeleton.py | 5 +++ .../mlir/passes/decompose_transfer.py | 7 +++- .../mlir/passes/lower_dma_to_gemmini.py | 35 +++++++----------- 4 files changed, 42 insertions(+), 42 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index a7bc914e..0b8330ce 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -566,7 +566,7 @@ def parse_index_list(self, expr_list:list, offset=sympy.Number(0)) -> common.CSE return index def load(self, name: str, index: sympy.Expr): - index, _ = self.convert_indirect_indexing(index) + index, offset_desc = self.convert_indirect_indexing(index) padding = self.get_padding_type() # Extract dram info @@ -591,7 +591,7 @@ def load(self, name: str, index: sympy.Expr): compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) code = self.emit_transfer("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, int(padding)) + dram_shape, tile_shape, dram_stride, tile_stride, int(padding), offset=offset_desc) self.cse.generate(self.dma_loads, code, assignment = False) # FIXME: assignment = False does not support caching with self.override_buffer_cse(buffer=self.loads): @@ -602,6 +602,7 @@ def load(self, name: str, index: sympy.Expr): def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs): dtype = V.graph.get_dtype(name) mlir_dtype = mlir_common.DTYPE_TO_MLIR[dtype] + offset_desc = None # Handle scatter store if "tmp" in str(index): @@ -613,7 +614,7 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) if mode == "atomic_add": loaded_value = ops.load(name, index) value = ops.add(loaded_value, value) - index, _ = self.convert_indirect_indexing(index) + index, offset_desc = self.convert_indirect_indexing(index) dram_var = self.kernel_group.args.output(name) # Prepare dma instruction @@ -655,7 +656,7 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) # Generate DMA instruction code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, 0) + dram_shape, tile_shape, dram_stride, tile_stride, 0, offset=offset_desc) self.dma_stores.writeline(common.DeferredLine(name, code)) def reduction(self, dtype, src_dtype, reduction_type, value): @@ -1358,7 +1359,7 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, dram_index_var, sram_var, sram_index_var, dram_shape, tile_shape, dram_stride, tile_stride, padding, - subtile_size=None, async_type=None): + subtile_size=None, async_type=None, offset=None): """Emit a generic togsim.transfer op for a DMA whose access exceeds the 4D Gemmini descriptor limit. Carries the full N-D access (dram/tile strides + shapes) plus the SSA operands a memref.dma_start needs @@ -1399,12 +1400,16 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp if subtile_size: av = int(async_type) if async_type is not None else 1 attrs += f', subtile_size = {list(subtile_size)}, async = {av} : i64' - # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride - return ( - f'"togsim.transfer"(%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' - f'%{tag}, %{dma_type}, %{vst}) {{{attrs}}} : ' - f'({dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index) -> ()' - ) + # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride [, offset spad] + operands = (f'%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' + f'%{tag}, %{dma_type}, %{vst}') + optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index' + if offset is not None: # indirect: per-position offset spad (decompose lifts it to a symbol attr) + offset_buf, offset_type = offset + operands += f', %{offset_buf}' + optypes += f', {offset_type}' + attrs += ', indirect = true' + return f'"togsim.transfer"({operands}) {{{attrs}}} : ({optypes}) -> ()' def allocate_sram_buffer(self, dtype, dram_name, tile_desc, raw_index, buffer=None, forced_name=None): c_type = mlir_common.DTYPE_TO_C[dtype] @@ -1547,13 +1552,7 @@ def convert_indirect_indexing(self, index :sympy.Expr): spad_vars[first_dim] = ops.add(spad_vars[first_dim], var) sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[first_dim] - mlir_dtype = vshape.split("x")[1][:-1] with self.override_buffer_cse(buffer=self.dma_loads): ops._store(spad_vars[first_dim], sram_var, sram_index_var, tile_shape) - - mlir_dtype = self.var_info[spad_vars[first_dim]][1] - with self.override_buffer_cse(buffer=self.dma_loads): - out = ops._load(1, mlir_dtype, sram_var, sram_index_var, tile_shape) - if mlir_dtype != "index": - out = ops.index_cast(out, "index") - return index + sympy.Symbol(str(out)), compute_dependecy + # Clean base index + the offset spad as an explicit transfer descriptor + return index, (sram_var, tile_shape) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index 4c3d89cb..4174231d 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -441,6 +441,11 @@ def _emit_one_dma(ctx, op, node, builder, bufs, tags): spad_id = bufs.of(dep._global_of(f["src"] if node.is_write else f["dst"])) read_bufs = [spad_id] if node.is_write else [] write_bufs = [] if node.is_write else [spad_id] + if "indirect_offset" in op.attributes: # gather/scatter reads the offset spad -> dep on its build + from mlir.ir import FlatSymbolRefAttr + off_id = bufs.of(FlatSymbolRefAttr(op.attributes["indirect_offset"]).value) + if off_id not in read_bufs: + read_bufs = read_bufs + [off_id] tag_id = tags.bind(_value_key(f["tag"]), spad_id) _emit_dma(ctx, node, tag_id, dram_index, tag_index, read_bufs, write_bufs) diff --git a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py index 10b2edfb..dfb67bb5 100644 --- a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py +++ b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py @@ -91,7 +91,10 @@ def run(module, vectorlane=128, **_): targets.append(op.operation) for op in targets: - dram, dram_idx, sram, sram_idx, tag, dma_type, vst = op.operands + op_operands = list(op.operands) + dram, dram_idx, sram, sram_idx, tag, dma_type, vst = op_operands[:7] + # indirect: offset spad operand -> lift to a symbol attr (memref.dma_start can't take the operand) + offset_sym = op_operands[7].owner.attributes["name"] if len(op_operands) > 7 else None kind = op.attributes["dma_kind"].value # StringAttr -> "MVIN"/"MVOUT" vlane_axis = IntegerAttr(op.attributes["vlane_split_axis"]).value dram_stride = _int_array(op.attributes["dram_stride"]) @@ -127,6 +130,8 @@ def _emit(sram_mem, sram_indices, dram_idx_val, vsa_val, dr_attr, tl_attr, st_at if st_attr is not None: attrs["subtile_size"] = st_attr attrs["async"] = async_attr + if offset_sym is not None: + attrs["indirect_offset"] = offset_sym Operation.create( "memref.dma_start", results=[], operands=operands, attributes=attrs) diff --git a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py index 998a6db5..51204b50 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py @@ -58,12 +58,18 @@ def run(module, timing=False): memref.dma_wait is erased in both modes (matches C++ DmaWaitOpLowering). """ from mlir.ir import (InsertionPoint, Operation, IntegerType, IndexType, - IntegerAttr, MemRefType) + IntegerAttr, MemRefType, FlatSymbolRefAttr, TypeAttr) from mlir.dialects import llvm, arith, memref i64 = IntegerType.get_signless(64) idx = IndexType.get() + # memref.global symbol -> type, to resolve the indirect_offset spad + sym2type = {} + for g in module.operation.regions[0].blocks[0].operations: + if g.operation.name == "memref.global": + sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + def const_int(val): return IntegerAttr(val.owner.attributes["value"]).value @@ -119,9 +125,8 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): is_mvin = dma_type in (MVIN, MVIN2, MVIN3) elem_bytes = _elem_bytes(src_ty.element_type) - # Indirect (gather): the gather-side indices are src for mvin, dst for mvout. - gather_idx = src_idx if is_mvin else dst_idx - indirect, indirect_memref = _find_indirect(gather_idx) + # Indirect (gather): offset spad referenced by the indirect_offset symbol attr + indirect = "indirect_offset" in op.attributes tile_shape = _subtile(op) if tile_shape is None: @@ -155,9 +160,12 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) if indirect: # CONFIG4: rs1 = indirect index-spad base address, rs2 = (elem_size<<16)|stride(1) + offset_sym = FlatSymbolRefAttr(op.attributes["indirect_offset"]).value + off_ty = sym2type[offset_sym] + indirect_memref = memref.GetGlobalOp(off_ty, offset_sym).result ind_base = memref.ExtractAlignedPointerAsIndexOp(indirect_memref).result ind_addr = arith.IndexCastOp(i64, ind_base).result - ind_esize = _elem_bytes(MemRefType(indirect_memref.type).element_type) + ind_esize = _elem_bytes(off_ty.element_type) asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (1 & 0xFFFF))) asm(dma_type, dram_addr, spad_addr) op.erase() @@ -189,23 +197,6 @@ def _elem_bytes(elem_type): return max(bits, 8) // 8 -def _find_indirect(indices): - """If a gather index is an affine.apply{indirect_access} whose operands include - index_cast(affine.load(%spad)), return (True, %spad memref); else (False, None).""" - for idx in indices: - ap = idx.owner - if getattr(ap, "name", None) != "affine.apply" or "indirect_access" not in ap.attributes: - continue - for operand in ap.operands: - ic = operand.owner - if getattr(ic, "name", None) != "arith.index_cast": - continue - ld = ic.operands[0].owner - if getattr(ld, "name", None) == "affine.load": - return True, ld.operands[0] # affine.load operand 0 == the index spad memref - return False, None - - def lower_text(text): if OP_NAME not in text: return text From b9adc10881233fb174e4fb54fd9ee5b246e7a960 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 26 Jun 2026 23:12:56 +0900 Subject: [PATCH 02/13] [Frontend] Clean up indirect access: symbol-set detection, CONFIG4 stride, compute-loop multi-dim offset Three refinements on top of the explicit offset descriptor: 1. Detect indirect access by an explicit symbol set instead of an "tmp" substring match. indirect_indexing now records str(index_var) in self.indirect_symbols; a _has_indirect(expr) helper tests the index free_symbols against it; the former "tmp"-string sites (store, get_dma_info, convert entry/indirect_dims/stride) use it. Removes the now-dead _find_indirect from lower_dma_to_gemmini. 2. Single indirect dim: pass the raw index spad and let the MVIN apply the gather stride per position (CONFIG4 offset_stride) instead of a vector_load/muli/vector_store round-trip that pre-multiplied the stride. emit_transfer carries offset_stride; decompose copies the attribute; lower programs CONFIG4 with it (was hardcoded to 1). 3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided indices, which a single CONFIG4 channel cannot do, so the sum stays in the kernel -- but build it in the compute loop (chunked by compute_vec_size, not a single tile-wide vector) and store it to a DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix) is not clobbered. push_step separates the offset-build loop from the gather that reads it. Adds multi-dim gather and index-reuse cases to test_indirect_access. Validated: indirect/scatter/embedding + the two new multi-dim cases pass; add/matmul/softmax regression 0. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- .../mlir/mlir_codegen_backend.py | 66 ++++++++++++------- .../mlir/passes/decompose_transfer.py | 1 + .../mlir/passes/lower_dma_to_gemmini.py | 3 +- tests/ops/misc/test_indirect_access.py | 27 ++++++++ 4 files changed, 74 insertions(+), 23 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 0b8330ce..db9875c9 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -384,6 +384,7 @@ def __init__(self, kernel_group, reason=None): self.welford_reduce_out = None self.reduce_iterator = {} self.spad_buffer_dict = dict() + self.indirect_symbols = set() # CSE-var names bound as indirect indices self.base_vector_initialized = False self.loop_size = None @@ -605,7 +606,7 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) offset_desc = None # Handle scatter store - if "tmp" in str(index): + if self._has_indirect(index): # Convert the output buffer type to the inplace buffer arg_name = V.graph.scheduler.mutation_real_name.get(name, name) if arg_name not in self.kernel_group.args.inplace_buffers: @@ -788,8 +789,12 @@ def store_reduction(self, name, index, value): self.reductions_suffix.writeline(common.DeferredLine(name, code)) def indirect_indexing(self, index_var, size, check=True, wrap_neg=True): + self.indirect_symbols.add(str(index_var)) # record the bound indirect symbol return str(index_var) + def _has_indirect(self, expr): + return any(s.name in self.indirect_symbols for s in expr.free_symbols) + def _index_expr(self, tile_desc, renamed_expression, index, base_vector_index): # In case of index expr, dimension size should be divisible by tile size if not self.kernel_group.tile_desc.is_dim_dividable(self.ranges): @@ -1227,7 +1232,7 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe """ # Use loads as default if buffer is None: - buffer = self.applys if "tmp" not in str(index) else self.dma_loads + buffer = self.applys if not self._has_indirect(index) else self.dma_loads # TODO. kg_tile_desc = self.kernel_group.tile_desc @@ -1237,7 +1242,7 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe total_dims = [int(str(i)[5:]) for i in self.itervars] local_tile_desc = mlir_common.MLIRMultiDimTile([1], self.vector_lane) local_dims.sort() # Assume that smaller index is placed in the outer loop - indirect_syms = [s for s in index.free_symbols if "tmp" in s.name] + indirect_syms = [s for s in index.free_symbols if s.name in self.indirect_symbols] index = index.subs({s: 0 for s in indirect_syms}, simultaneous=True) indirect_dims = [f"{i}" for i in indirect_syms] @@ -1405,10 +1410,10 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp f'%{tag}, %{dma_type}, %{vst}') optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index' if offset is not None: # indirect: per-position offset spad (decompose lifts it to a symbol attr) - offset_buf, offset_type = offset + offset_buf, offset_type, offset_stride = offset operands += f', %{offset_buf}' optypes += f', {offset_type}' - attrs += ', indirect = true' + attrs += f', indirect = true, offset_stride = {int(offset_stride)} : i64' return f'"togsim.transfer"({operands}) {{{attrs}}} : ({optypes}) -> ()' def allocate_sram_buffer(self, dtype, dram_name, tile_desc, raw_index, buffer=None, forced_name=None): @@ -1490,7 +1495,7 @@ def get_mask(self): return mask_shape, mask_var def convert_indirect_indexing(self, index :sympy.Expr): - if "tmp" not in str(index): + if not self._has_indirect(index): return index, None # Note: In case of indirect indexing, dimensions should be divisible by tile size @@ -1501,12 +1506,11 @@ def convert_indirect_indexing(self, index :sympy.Expr): raise mlir_common.RecompileSignal(f"Indirect access (tile size {self.kernel_group.tile_desc.get_tile_size()} is not divisible by {self.ranges})") # Process start - indirect_dims = [str(dim) for dim in index.free_symbols if "tmp" in str(dim)] + indirect_dims = [str(dim) for dim in index.free_symbols if str(dim) in self.indirect_symbols] indirect_dims.sort() first_dim = indirect_dims[0] spad_vars = dict() compute_dependecy = any([target_dim not in self.spad_buffer_dict for target_dim in indirect_dims]) - # Store each newly-produced indirect index into spad, in its producing step for target_dim in indirect_dims: if target_dim in self.spad_buffer_dict: @@ -1529,17 +1533,30 @@ def convert_indirect_indexing(self, index :sympy.Expr): if compute_dependecy: self.push_step() - # Build the offset (outer ops) in the current step, reading indices back from spad + # Single indirect dim: the raw index IS the offset; the MVIN applies offset_stride (CONFIG4) + if len(indirect_dims) == 1: + offset_stride = 1 + for arg in list(index.args): + if not self._has_indirect(arg): + continue + if arg.is_Mul and arg.args[0].is_number: + offset_stride = int(arg.args[0]) + index = index.replace(arg, 0) + sram_var, _, _, _, tile_shape, _ = self.spad_buffer_dict[first_dim] + return index, (sram_var, tile_shape, offset_stride) + + # Multi indirect dim: sum the strided indices in the compute loop (chunked by compute_vec_size) + local_tile_desc = self.kernel_group.tile_desc + compute_vec_size = local_tile_desc.get_compute_vec_size() for target_dim in indirect_dims: - sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[target_dim] + sram_var, _, _, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[target_dim] mlir_dtype = vshape.split("x")[1][:-1] - with self.override_buffer_cse(buffer=self.dma_loads): - out = ops._load(tile_numel_per_lane, mlir_dtype, sram_var, sram_index_var, tile_shape) - spad_vars[target_dim] = out - - with self.override_buffer_cse(buffer=self.dma_loads): + compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) + with self.override_buffer_cse(buffer=self.loads): + spad_vars[target_dim] = ops._load(compute_vec_size, mlir_dtype, sram_var, compute_index_var, tile_shape) + with self.override_buffer_cse(buffer=self.compute): for arg in index.args: - if "tmp" not in str(arg): + if not self._has_indirect(arg): continue if arg.is_Mul and arg.args[0].is_number: coeff_dtype = self.var_info[spad_vars[str(arg.args[1])]][1] @@ -1550,9 +1567,14 @@ def convert_indirect_indexing(self, index :sympy.Expr): if dim == first_dim: continue spad_vars[first_dim] = ops.add(spad_vars[first_dim], var) - - sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[first_dim] - with self.override_buffer_cse(buffer=self.dma_loads): - ops._store(spad_vars[first_dim], sram_var, sram_index_var, tile_shape) - # Clean base index + the offset spad as an explicit transfer descriptor - return index, (sram_var, tile_shape) + # Summed offset goes to a DEDICATED spad (not an index buffer) to avoid clobbering a live index + var_info = [v for k, v in self.var_info.items() if str(k) == first_dim][0] + dtype = mlir_common.MLIR_TO_DTYPE[var_info[1]] + off_shape = local_tile_desc.get_mlir_shape(var_info[1]) + off_sram, off_index = self.get_scratchpad_buffer( + dtype, "indirect_offset_" + first_dim, local_tile_desc, "indirect_offset_" + first_dim) + off_compute_index = ",".join(off_index.split(",")[:-1] + [f"%{self.compute_idx}"]) + with self.override_buffer_cse(buffer=self.stores): + ops._store(spad_vars[first_dim], off_sram, off_compute_index, off_shape) + self.push_step() # offset-build compute loop must finish before the gather reads it + return index, (off_sram, off_shape, 1) diff --git a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py index dfb67bb5..0bf04e30 100644 --- a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py +++ b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py @@ -132,6 +132,7 @@ def _emit(sram_mem, sram_indices, dram_idx_val, vsa_val, dr_attr, tl_attr, st_at attrs["async"] = async_attr if offset_sym is not None: attrs["indirect_offset"] = offset_sym + attrs["offset_stride"] = op.attributes["offset_stride"] Operation.create( "memref.dma_start", results=[], operands=operands, attributes=attrs) diff --git a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py index 51204b50..5ca842c1 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py @@ -166,7 +166,8 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): ind_base = memref.ExtractAlignedPointerAsIndexOp(indirect_memref).result ind_addr = arith.IndexCastOp(i64, ind_base).result ind_esize = _elem_bytes(off_ty.element_type) - asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (1 & 0xFFFF))) + off_stride = IntegerAttr(op.attributes["offset_stride"]).value + asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (off_stride & 0xFFFF))) asm(dma_type, dram_addr, spad_addr) op.erase() diff --git a/tests/ops/misc/test_indirect_access.py b/tests/ops/misc/test_indirect_access.py index f64fe50d..ae3a4ed4 100644 --- a/tests/ops/misc/test_indirect_access.py +++ b/tests/ops/misc/test_indirect_access.py @@ -52,6 +52,31 @@ def scatter_only(out, token_indices, weighted_output): res = opt_fn(out, token_indices, weighted_output) test_result("ScatterAdd(index_add_)", res, cpu_out) +def test_multidim_indirect(device, size=(64, 64), n=256): + torch.manual_seed(0) + def gather2d(x, ix, iy): + return x[ix, iy] + 1.0 + x = torch.randn(size, dtype=torch.float32).to(device=device) + ix = torch.randint(0, size[0], [n]).to(device=device) + iy = torch.randint(0, size[1], [n]).to(device=device) + opt_fn = torch.compile(dynamic=False)(gather2d) + res = opt_fn(x, ix, iy) + out = gather2d(x.cpu(), ix.cpu(), iy.cpu()) + test_result("Multi-dim Indirect (x[ix,iy])", res, out) + +def test_multidim_indirect_index_reuse(device, size=(64, 64), n=256): + torch.manual_seed(0) + def gather_reuse(x, ix, iy): + # ix is reused after the gather -> the offset spad must not clobber the index spad + return x[ix, iy] + ix.float() + x = torch.randn(size, dtype=torch.float32).to(device=device) + ix = torch.randint(0, size[0], [n]).to(device=device) + iy = torch.randint(0, size[1], [n]).to(device=device) + opt_fn = torch.compile(dynamic=False)(gather_reuse) + res = opt_fn(x, ix, iy) + out = gather_reuse(x.cpu(), ix.cpu(), iy.cpu()) + test_result("Multi-dim Indirect index reuse (x[ix,iy]+ix)", res, out) + def test_scatter_full(device, size=(128, 128)): def vectoradd(a, idx, b): a[idx, :] = b @@ -71,4 +96,6 @@ def vectoradd(a, idx, b): test_scatter_full(device, size=(2048, 2048)) test_scatter_add(device) test_indirect_vectoradd(device) + test_multidim_indirect(device) + test_multidim_indirect_index_reuse(device) #test_embedding(device, 1024, 2048) \ No newline at end of file From 48194d25b712c290ffbc6449db11302b00c7ee95 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 15:35:41 +0900 Subject: [PATCH 03/13] [Frontend] Drop memref.dma_start: lower togsim.transfer directly to Gemmini memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers directly to Gemmini. - New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather. - togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices. - memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering, erased by the merged lowering). No memref.* DMA ops remain. - decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its DRAM-grow half is what the future masked-DMA clamp replaces). - dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer. Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv (padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0 passes; wrong values, not a crash). Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- PyTorchSimFrontend/extension_codecache.py | 4 +- .../mlir/mlir_codegen_backend.py | 6 +- PyTorchSimFrontend/mlir/passes/__init__.py | 4 +- .../mlir/passes/build_skeleton.py | 60 ++-- .../mlir/passes/dma_fine_grained.py | 96 ++++--- .../mlir/passes/lower_to_llvm.py | 8 +- .../mlir/passes/lower_to_vcix.py | 40 +-- .../mlir/passes/lower_transfer_to_gemmini.py | 257 ++++++++++++++++++ docs/design/transfer-direct-lowering.md | 66 +++++ 9 files changed, 455 insertions(+), 86 deletions(-) create mode 100644 PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py create mode 100644 docs/design/transfer-direct-lowering.md diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 5e6857ee..83503c31 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -51,7 +51,7 @@ def mlir_compile_command(filename, vectorlane_size, vlen=256): return [re.sub(r"[ \n]+", " ", f""" {extension_config.CONFIG_TORCHSIM_LLVM_PATH}/mlir-opt \ - -test-loop-padding \ + -test-loop-padding --allow-unregistered-dialect \ {'--mlir-print-ir-after-all' if extension_config.CONFIG_TORCHSIM_DUMP_MLIR_IR else ''} \ {filename}.mlir -o {filename}_padded.mlir """, @@ -86,7 +86,7 @@ def mlir_gem5_compile_command(filename, sample_filename, vectorlane_size, vlen=2 return [re.sub(r"[ \n]+", " ", f""" {extension_config.CONFIG_TORCHSIM_LLVM_PATH}/mlir-opt \ - -test-loop-padding='timing_mode=1' \ + -test-loop-padding='timing_mode=1' --allow-unregistered-dialect \ {'--mlir-print-ir-after-all' if extension_config.CONFIG_TORCHSIM_DUMP_MLIR_IR else ''} \ {filename}.mlir -o {sample_filename}_padded.mlir """, diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index db9875c9..816f88bc 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -1405,10 +1405,10 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp if subtile_size: av = int(async_type) if async_type is not None else 1 attrs += f', subtile_size = {list(subtile_size)}, async = {av} : i64' - # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride [, offset spad] + # operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vlane_stride [, offset spad] operands = (f'%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' - f'%{tag}, %{dma_type}, %{vst}') - optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index' + f'%{tag}, %{zero_cse}, %{dma_type}, %{vst}') + optypes = f'{dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index, index' if offset is not None: # indirect: per-position offset spad (decompose lifts it to a symbol attr) offset_buf, offset_type, offset_stride = offset operands += f', %{offset_buf}' diff --git a/PyTorchSimFrontend/mlir/passes/__init__.py b/PyTorchSimFrontend/mlir/passes/__init__.py index ab3cdcd3..199f4c7d 100644 --- a/PyTorchSimFrontend/mlir/passes/__init__.py +++ b/PyTorchSimFrontend/mlir/passes/__init__.py @@ -41,9 +41,9 @@ def _ensure_mlir_bindings_on_path(): # Module rewrite passes around the one remaining mlir-opt pass (-test-loop-padding). # Each exposes MARKERS + run(module, **opts); run_module_passes parses once per phase. -# decompose_transfer first: togsim.transfer -> memref.dma_start (downstream expects it). +# togsim.transfer stays through the pipeline (no more memref.dma_start): it lowers +# directly to Gemmini at the end (lower_transfer_to_gemmini); loop-padding runs opaquely. PRE_OPT_PASSES = [ - decompose_transfer, lower_vlane_idx, ] # fine-grained first: splits the matmul DMAs that the vcix lowering then reads. diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index 4174231d..cc0fbd26 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -6,7 +6,7 @@ skeleton with the data computation replaced by calls to the event-based runtime API. This pass performs that reduction at the MLIR level: - * `memref.dma_start` -> `togsim.dma(...) {tag_id, is_async, ...}` carrying the + * `togsim.transfer` -> `togsim.dma(...) {tag_id, is_async, ...}` carrying the runtime tag index operand (`%tag[%idx]`). * `memref.dma_wait` -> `togsim.memory_barrier(tag_idx) {tag_id, write_bufs}`, the explicit async-DMA sync. It pairs with its dma by @@ -47,7 +47,7 @@ ) #: Marker op names for the passes/__init__ fast-path (skip parsing if absent). -MARKERS = ("memref.dma_start", "memref.dma_wait") +MARKERS = ("togsim.transfer", "memref.dma_wait") #: Ops the DCE must never remove (loops, terminators, our API ops). _KEEP = { @@ -75,7 +75,7 @@ def _arg_id_of(base_addr): def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_bufs): - """Insert a `togsim.dma` before the original `memref.dma_start`. + """Insert a `togsim.dma` before the original `togsim.transfer`. `tag_id` is the identity of this DMA's tag memref. An async DMA pairs with its `togsim.memory_barrier` (the original dma_wait) by the RUNTIME tag slot @@ -84,7 +84,7 @@ def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_buf only a runtime key can pair iteration i's dma with iteration i's wait. `dram_index` is the original linear DRAM index Value (the `affine.apply` - result that indexed the tensor in the `memref.dma_start`) -- carried as an + result that indexed the tensor in the `togsim.transfer`) -- carried as an operand so the DCE keeps the address arithmetic live and the C4 lowering can compute the real `base_addr = base[arg_id] + index*elem` (P3, approach A). @@ -425,25 +425,51 @@ def _emit_computes(ctx, builder, bufs): return n +def _transfer_fields(op): + """Decode a `togsim.transfer`'s fixed operands by position. + + Layout (see mlir_codegen_backend.emit_transfer / lower_transfer_to_gemmini): + operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst + [, offset_spad] # 8 or 9 operands + Unlike the old `memref.dma_start`, dram/sram are FIXED (not direction-swapped): + the DRAM side is always operand[0]/[1], the SRAM spad always operand[2], the + runtime tag slot always operand[4] (tag memref) + operand[5] (tag_idx). The + optional indirect-offset spad is operand[8]; its owning `memref.get_global` + carries the offset symbol name in its "name" attribute (matching + lower_transfer_to_gemmini's offset_sym derivation).""" + operands = list(op.operands) + offset = operands[8] if len(operands) > 8 else None + return { + "dram": operands[0], "dram_idx": operands[1], + "sram": operands[2], "sram_idx": operands[3], + "tag": operands[4], "tag_idx": operands[5], + "dma_type": operands[6], "vst": operands[7], + "offset": offset, + } + + def _emit_one_dma(ctx, op, node, builder, bufs, tags): - """Rewrite one memref.dma_start as togsim.dma. A load reads DRAM and writes + """Rewrite one togsim.transfer as togsim.dma. A load reads DRAM and writes its SRAM spad; a store reads the spad and writes DRAM -- which sets the read/write buffer that drives the dependency edge (sec 10). The tag memref is bound to a tag_id (with its loaded buffer) so the paired memory_barrier finds it by the runtime tag slot.""" from . import dep_analysis as dep # lazy: dep_analysis imports build_skeleton - f = builder._dma_start_fields(op) - dram_indices = f["dst_indices"] if node.is_write else f["src_indices"] - dram_index = dram_indices[0] if dram_indices else None - tag_indices = f["tag_indices"] - tag_index = tag_indices[0] if tag_indices else None - # the spad is the SRAM side of the copy: dst for a load, src for a store. - spad_id = bufs.of(dep._global_of(f["src"] if node.is_write else f["dst"])) + f = _transfer_fields(op) + # dram/sram are fixed operands now (not direction-swapped): the DRAM index is + # always operand[1], the SRAM spad always operand[2]. Direction (read/write) + # comes from the node (dma_kind attr / dma_type value). + dram_index = f["dram_idx"] + tag_index = f["tag_idx"] # single runtime tag slot operand (%tag[%idx]) + spad_id = bufs.of(dep._global_of(f["sram"])) read_bufs = [spad_id] if node.is_write else [] write_bufs = [] if node.is_write else [spad_id] - if "indirect_offset" in op.attributes: # gather/scatter reads the offset spad -> dep on its build - from mlir.ir import FlatSymbolRefAttr - off_id = bufs.of(FlatSymbolRefAttr(op.attributes["indirect_offset"]).value) + if f["offset"] is not None: # gather/scatter reads the offset spad -> dep on its build + # the offset symbol name lives on the offset operand's owning get_global + # ("name" attr), the same place lower_transfer_to_gemmini reads it. + off_owner = f["offset"].owner + off_sym = str(off_owner.attributes["name"]).strip('@" ') + off_id = bufs.of(off_sym) if off_id not in read_bufs: read_bufs = read_bufs + [off_id] tag_id = tags.bind(_value_key(f["tag"]), spad_id) @@ -470,7 +496,7 @@ def _emit_one_wait(ctx, op, tags): def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): - """Step 2: rewrite memref.dma_start -> togsim.dma and memref.dma_wait -> + """Step 2: rewrite togsim.transfer -> togsim.dma and memref.dma_wait -> togsim.memory_barrier in program order. An async dma and its barrier are paired by the RUNTIME tag slot (tag_id + tag index), not a compile-time id: one static dma op runs per loop iteration with a different `%tag[%idx]`, so @@ -481,7 +507,7 @@ def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): n_dma = n_wait = 0 for op in list(walk_ops(block)): name = op.operation.name - if name == "memref.dma_start": + if name == "togsim.transfer": node = dma_by_op.get(id(op.operation)) if node is None: continue diff --git a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py index d7571d2b..22eb999e 100644 --- a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py +++ b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py @@ -11,13 +11,16 @@ IRMapping; the MLIR Python bindings expose no IRMapping, so this port builds the fused nest directly and emits each DMA inside it using the fused induction vars (equivalence target: same loop structure / counts, same offset maps, same -dma_start operands+attrs -- validated against mlir-opt -dma-fine-grained and the -end-to-end gemm/conv/model tests, not byte-exact SSA text). +togsim.transfer operands+attrs -- validated against mlir-opt -dma-fine-grained and +the end-to-end gemm/conv/model tests, not byte-exact SSA text). -Operates on the customized memref.dma_start convention (see lower_dma_to_gemmini): -operands = src, *src_idx, dst, *dst_idx, num_elements(dma_type), tag, *tag_idx, -stride(=vlane_split_axis), num_elements_per_stride(=vlane_stride). MVIN dma_type in -{2,1,14}; tile shape = dst shape for MVIN. +Operates on the togsim.transfer convention (see mlir_codegen_backend.emit_transfer +and lower_transfer_to_gemmini): operands = dram, dram_idx, sram, sram_idx, tag, +tag_idx, dma_type, vst(=vlane_stride)[, offset]; attrs = dma_kind, vlane_split_axis +(i64), dram_stride[], tile_stride[], padding, [subtile_size, async]. Direction is +derived from dma_kind / dma_type: MVIN => src=dram, dst=sram; MVOUT => src=sram, +dst=dram. tile shape = the sram memref shape for BOTH directions. MVIN dma_type in +{2,1,14}. Pipeline entry point: run_fine_grained(in_path, out_path, vectorlane). """ @@ -62,25 +65,37 @@ def _is_block_arg(v): class _Dma: - """Positional view of a customized memref.dma_start op.""" + """Positional view of a togsim.transfer op. + + operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst[, offset] + Direction from dma_kind / dma_type: MVIN => src=dram, dst=sram (MVOUT swaps). + self.src_idx is a single-element list holding the base DRAM idx for MVIN (the + SRAM idx for MVOUT); tile_shape is always the sram memref shape. + """ def __init__(self, op): self.op = op operands = list(op.operands) - src_rank = len(ir.MemRefType(operands[0].type).shape) - i = 0 - self.src = operands[i]; i += 1 - self.src_idx = operands[i:i + src_rank]; i += src_rank - self.dst = operands[i]; i += 1 - dst_rank = len(ir.MemRefType(self.dst.type).shape) - self.dst_idx = operands[i:i + dst_rank]; i += dst_rank - self.num_elements = operands[i]; i += 1 - self.tag = operands[i]; i += 1 - tag_rank = len(ir.MemRefType(self.tag.type).shape) - self.tag_idx = operands[i:i + tag_rank]; i += tag_rank - self.stride = operands[i]; i += 1 # = vlane_split_axis - self.num_elements_per_stride = operands[i] # = vlane_stride - self.src_rank, self.dst_rank, self.tag_rank = src_rank, dst_rank, tag_rank + # dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst[, offset] + self.dram = operands[0] + self.dram_idx = operands[1] + self.sram = operands[2] + self.sram_idx = operands[3] + self.tag = operands[4] + self.tag_idx = operands[5] + self.num_elements = operands[6] # = dma_type const operand + self.num_elements_per_stride = operands[7] # = vlane_stride (vst) + + self.sram_rank = len(ir.MemRefType(self.sram.type).shape) + # Direction: MVIN reads dram -> sram; MVOUT writes sram -> dram. + if self.is_mvin: + self.src, self.dst = self.dram, self.sram + self.src_idx = [self.dram_idx] + else: + self.src, self.dst = self.sram, self.dram + self.src_idx = [self.sram_idx] + self.src_rank = len(ir.MemRefType(self.src.type).shape) + self.dst_rank = len(ir.MemRefType(self.dst.type).shape) @property def dma_type(self): @@ -92,21 +107,21 @@ def is_mvin(self): @property def vlane_split_axis(self): - return _const_int(self.stride) + return ir.IntegerAttr(self.op.attributes["vlane_split_axis"]).value @property def vlane_stride(self): return _const_int(self.num_elements_per_stride) & 0x7FFF def tile_shape(self): - mt = ir.MemRefType((self.dst if self.is_mvin else self.src).type) - return list(mt.shape) + return list(ir.MemRefType(self.sram.type).shape) def subtile_size(self): return attr_i64_array(self.op, "subtile_size", default=[]) def sram_stride(self): - return attr_i64_array(self.op, "sram_stride", default=[]) + # togsim.transfer names the spad stride "tile_stride". + return attr_i64_array(self.op, "tile_stride", default=[]) def dram_stride(self): return attr_i64_array(self.op, "dram_stride", default=[]) @@ -200,10 +215,13 @@ def _apply(map_, operands, ip): def _dma_attrs(dma): - """Mirror getDmaAttrs: keep subtile/sram/dram strides, set async + fine_grained.""" + """Build the emitted togsim.transfer's attrs: copy dma_kind, vlane_split_axis, + dram_stride, tile_stride (the spad stride), subtile_size and padding straight + from the source op; set async (BoolAttr) + fine_grained (BoolAttr true).""" attrs = {} op = dma.op - for k in ("subtile_size", "sram_stride", "dram_stride"): + for k in ("dma_kind", "vlane_split_axis", "dram_stride", "tile_stride", + "subtile_size", "padding"): if k in op.attributes: attrs[k] = op.attributes[k] attrs["async"] = ir.BoolAttr.get(dma.is_async()) @@ -212,26 +230,20 @@ def _dma_attrs(dma): def _emit_dma(dma, ivs, vectorlane, ip): - """Emit one fine-grained memref.dma_start at `ip`, indexed by `ivs` (the fused + """Emit one fine-grained togsim.transfer at `ip`, indexed by `ivs` (the fused induction vars for this DMA's dims, in dim order).""" - idx_ty = ir.IndexType.get() - zero = _const_index(0, ip) - dram_off = _apply(_build_dram_map(dma), ivs, ip) - src_idx0 = dma.src_idx[0] - dram_idx = _apply(_sum_map(), [dram_off, src_idx0], ip) + # DRAM base index = the original transfer's dram_idx operand. + dram_idx = _apply(_sum_map(), [dram_off, dma.dram_idx], ip) + # SRAM offset is a SINGLE linear sram_idx operand (row-major stride 1). sram_off = _apply(_build_sram_map(dma, vectorlane), ivs, ip) + # Per-subtile tag index (required for async DMA<->barrier pairing downstream). tag_idx = _apply(_build_tag_map(dma, list(range(len(dma.tile_shape())))), ivs, ip) - # SRAM indices: zeros except the last = sram offset (mirror sramIndices.back()). - sram_indices = [zero] * dma.dst_rank - sram_indices[-1] = sram_off - - operands = [dma.src, dram_idx, dma.dst, *sram_indices, - dma.num_elements, dma.tag, tag_idx, - dma.stride, dma.num_elements_per_stride] - ir.Operation.create("memref.dma_start", results=[], operands=operands, + operands = [dma.dram, dram_idx, dma.sram, sram_off, + dma.tag, tag_idx, dma.num_elements, dma.num_elements_per_stride] + ir.Operation.create("togsim.transfer", results=[], operands=operands, attributes=_dma_attrs(dma), ip=ip) @@ -320,7 +332,7 @@ def _run_func(func, vectorlane): name = op.operation.name if name == "linalg.matmul" and matmul is None: matmul = op - elif name == "memref.dma_start": + elif name == "togsim.transfer": dmas.append(op) if matmul is None: return diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py index ad287499..4629df29 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py @@ -50,15 +50,15 @@ def run_standard_lowering(in_path, out_path=None, timing=False): out_path = in_path from mlir.ir import Context, Module, Location from mlir.passmanager import PassManager - from . import lower_dma_to_gemmini + from . import lower_transfer_to_gemmini ctx = Context() ctx.allow_unregistered_dialects = True with ctx, Location.unknown(): with open(in_path) as f: module = Module.parse(f.read()) - # Imperative Python pass: memref.dma_start/dma_wait -> Gemmini asm (replaces - # the C++ test-memref-to-gemmini), then the registered standard lowering. - lower_dma_to_gemmini.run(module, timing=timing) + # Imperative Python pass: togsim.transfer -> Gemmini asm directly (no + # memref.dma_start intermediate), then the registered standard lowering. + lower_transfer_to_gemmini.run(module, timing=timing) PassManager.parse(STANDARD_PIPELINE, ctx).run(module.operation) with open(out_path, "w") as f: f.write(str(module)) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py index df124d00..d32ada38 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py @@ -255,8 +255,9 @@ def _transfer_write(value, dest, indices): def _dma_wait(tag, idx, num_elements): - from mlir.dialects import memref - memref.DmaWaitOp(tag, [idx], num_elements) + # togsim-level counterpart of the async-DMA barrier (paired with togsim.transfer), + # instead of memref.dma_wait -- no memref.* DMA ops remain in the pipeline. + ir.Operation.create("togsim.wait", results=[], operands=[tag, idx, num_elements]) def _vcix(name, operands, result_tys, attrs): @@ -277,22 +278,29 @@ def _reaches(value, target): class _DmaView: - """Positional view of a customized memref.dma_start (see lower_dma_to_gemmini).""" + """Positional view of a togsim.transfer op (see emit_transfer / + lower_transfer_to_gemmini). + + New operand layout: (dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, + vst [, offset]). Direction (which of dram/sram is source vs dest) comes from + the `dma_kind` attr ("MVIN"|"MVOUT"), not from operand position. To keep the + downstream _dram_is_write(src, dst) / `sram = d.dst` logic working (it + distinguishes MVIN loads from MVOUT stores by memory space), src/dst are + exposed direction-swapped: MVIN -> src=dram, dst=sram; MVOUT -> src=sram, + dst=dram.""" def __init__(self, op): self.op = op operands = list(op.operands) - src_rank = len(ir.MemRefType(operands[0].type).shape) - i = 0 - self.src = operands[i]; i += 1 - i += src_rank - self.dst = operands[i]; i += 1 - dst_rank = len(ir.MemRefType(self.dst.type).shape) - i += dst_rank - i += 1 # num_elements - self.tag = operands[i]; i += 1 - tag_rank = len(ir.MemRefType(self.tag.type).shape) - self.tag_idx = operands[i:i + tag_rank] + dram = operands[0] + sram = operands[2] + self.tag = operands[4] + self.tag_idx = [operands[5]] + kind = ir.StringAttr(op.attributes["dma_kind"]).value + if kind == "MVOUT": + self.src, self.dst = sram, dram + else: # MVIN (load) + self.src, self.dst = dram, sram def subtile_size(self): a = self.op.attributes @@ -359,7 +367,7 @@ def a64(v): return ir.IntegerAttr.get(i64, v) subtileM, subtileN, subtileK = M, N, K a_subk = b_subk = None # Mirror the C++ isAInitialized / isBInitialized flags: an operand is - # "initialized" either by an MVIN dma_start (tag found below) or by a + # "initialized" either by an MVIN togsim.transfer (tag found below) or by a # preceding affine.vector_store into its root memref (the fused case, e.g. # SDPA scores.V where B is the softmax output produced in-place, not DMAed). isAInit = isBInit = False @@ -380,7 +388,7 @@ def _root(v): elif dest == rootB: isBInit = True continue - if o.operation.name != "memref.dma_start": + if o.operation.name != "togsim.transfer": continue d = _DmaView(o.operation) dram, is_write = _dram_is_write(d.src, d.dst) diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py new file mode 100644 index 00000000..037bc152 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -0,0 +1,257 @@ +"""Lower togsim.transfer DIRECTLY to Gemmini RISC-V inline asm (no memref.dma_start). + +Merges decompose_transfer (the <=4D Gemmini-limit handling: drop unit dims / +collapse / >4D affine.for peel with lane-banked SRAM offset) with +lower_dma_to_gemmini (the CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm emission). +togsim.transfer is unregistered so it carries every runtime descriptor as an +operand -- including the future masked-clamp low/high vectors -- which a registered +memref.dma_start cannot. See docs/design/transfer-direct-lowering.md. + +timing=False: emit the gemmini asm. timing=True: erase the transfer (the TOG carries +DMA timing; the cycle binary needs no asm). +""" + +OP_NAME = "togsim.transfer" +WAIT_NAME = "togsim.wait" +MARKERS = (OP_NAME, WAIT_NAME) + +from ._mlir_util import walk_ops +from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS +from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation + +CONFIG, CONFIG2, CONFIG3, CONFIG4 = 0, 4, 5, 6 +MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 +CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} +MAX_TENSOR_DIM = 4 + + +def run(module, timing=False, vectorlane=128, **_): + from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, + IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, + DenseI32ArrayAttr, StridedLayoutAttr, AffineMap, AffineMapAttr, + AffineExpr, BoolAttr, FlatSymbolRefAttr, TypeAttr) + from mlir.dialects import affine, llvm, arith, memref + i64 = IntegerType.get_signless(64) + idx_ty = IndexType.get() + + sym2type = {} + for g in module.operation.regions[0].blocks[0].operations: + if g.operation.name == "memref.global": + sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + + def i64_const(value): + return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result + + def asm(func7, rs1, rs2): + llvm.InlineAsmOp(None, [rs1, rs2], _asm(func7), CONSTRAINTS, + has_side_effects=True, asm_dialect=0) + + def elem_addr_i64(memref_val, indices, mtype, elem_bytes): + base = memref.ExtractAlignedPointerAsIndexOp(memref_val).result + strides = _row_major_strides(list(mtype.shape)) + off = None + for k, ival in enumerate(indices): + if strides[k] == 0: + continue + term = ival + if strides[k] != 1: + term = arith.MulIOp(ival, arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, strides[k])).result).result + off = term if off is None else arith.AddIOp(off, term).result + if off is not None: + byte = arith.MulIOp(off, arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, elem_bytes)).result).result + base = arith.AddIOp(base, byte).result + return arith.IndexCastOp(i64, base).result + + targets, waits = [], [] + for region in module.operation.regions: + for b in region.blocks: + for op in walk_ops(b): + if op.operation.name == OP_NAME: + targets.append(op.operation) + elif op.operation.name == WAIT_NAME: + waits.append(op.operation) + + for op in waits: # togsim.wait: erase in both modes (the barrier is a sync marker) + op.erase() + + for op in targets: + op_operands = list(op.operands) + dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst = op_operands[:8] + offset_sym = (op_operands[8].owner.attributes["name"] if len(op_operands) > 8 else None) + kind = op.attributes["dma_kind"].value + dma_type_val = _const_int(dma_type) # MVIN(2)/MVIN2(1)/MVIN3(14)/MVOUT(3) + is_mvin = dma_type_val in (MVIN, MVIN2, MVIN3) + vlane_axis = IntegerAttr(op.attributes["vlane_split_axis"]).value + dram_stride = _int_array(op.attributes["dram_stride"]) + tile_stride = _int_array(op.attributes["tile_stride"]) + vlane_stride = _const_int(vst, 1) + try: + subtile = _int_array(op.attributes["subtile_size"]) + except KeyError: + subtile = None + + if timing: + op.erase() + continue + + sram_ty = MemRefType(sram.type) + elem, space = sram_ty.element_type, sram_ty.memory_space + elem_bytes = _elem_bytes(sram_ty.element_type) + dram_ty = MemRefType(dram.type) + tile_shape = list(sram_ty.shape) + eff = [i for i, e in enumerate(tile_shape) if e > 1] + indirect = offset_sym is not None + + def _const(v): + return arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, v)).result + + def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, + desc_dram_strides, desc_spad_strides, subtile_shape): + cfg_shape = subtile_shape if subtile_shape is not None else desc_shape + expand = MAX_TENSOR_DIM - len(cfg_shape) + shape4 = [1] * expand + list(cfg_shape) + dram4 = [0] * expand + list(desc_dram_strides) + spad4 = [0] * expand + list(desc_spad_strides) + vsa4 = vsa_val + expand + config_type = CONFIG_TYPE[dma_type_val] + sram_c = MemRefType(sram_mem.type) + dram_addr = elem_addr_i64(dram, [dram_idx_val], dram_ty, elem_bytes) + spad_addr = elem_addr_i64(sram_mem, sram_indices, sram_c, elem_bytes) + cfg_rs1 = i64_const(((shape4[0] & 0xFFFF) << 48) | ((shape4[1] & 0xFFFF) << 32) + | ((shape4[2] & 0xFFFF) << 16) | (shape4[3] & 0xFFFF)) + cfg_rs2 = i64_const((vlane_stride << 32) | ((config_type & 0x3) << 17) + | ((1 if indirect else 0) << 16) + | ((vsa4 & 0x3) << 14) | elem_bytes) + asm(CONFIG, cfg_rs1, cfg_rs2) + asm(CONFIG2, i64_const((dram4[0] << 32) | (dram4[1] & 0xFFFFFFFF)), + i64_const((dram4[2] << 32) | (dram4[3] & 0xFFFFFFFF))) + asm(CONFIG3, i64_const((spad4[0] << 32) | (spad4[1] & 0xFFFFFFFF)), + i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) + if indirect: + sym = FlatSymbolRefAttr(offset_sym).value if not isinstance(offset_sym, str) else offset_sym + off_ty = sym2type[sym] + ind_base = memref.ExtractAlignedPointerAsIndexOp(memref.GetGlobalOp(off_ty, sym).result).result + ind_addr = arith.IndexCastOp(i64, ind_base).result + ind_esize = _elem_bytes(off_ty.element_type) + off_stride = IntegerAttr(op.attributes["offset_stride"]).value + asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (off_stride & 0xFFFF))) + asm(dma_type_val, dram_addr, spad_addr) + + if offset_sym is not None: + offset_sym = FlatSymbolRefAttr(offset_sym).value if not isinstance(offset_sym, str) else offset_sym + + if len(tile_shape) <= 4: + with InsertionPoint(op): + # sram offset is linear (row-major stride 1) -> last index only; others 0. + sidx = [_const(0)] * (len(tile_shape) - 1) + [sram_idx] + _emit_asm(sram, sidx, dram_idx, vlane_axis, + tile_shape, dram_stride, tile_stride, subtile) + op.erase() + continue + + if len(eff) <= 4: + groups, target = _squeeze_reassociation(tile_shape) + reassoc = ArrayAttr.get( + [ArrayAttr.get([IntegerAttr.get(i64, d) for d in g]) for g in groups]) + collapsed_ty = MemRefType.get(target, elem, memory_space=space) + keep = [next((d for d in g if tile_shape[d] > 1), g[-1]) for g in groups] + dr = [dram_stride[i] for i in keep] + tl = [tile_stride[i] for i in keep] + st = [subtile[i] for i in keep] if subtile is not None else None + new_vlane = next(gi for gi, g in enumerate(groups) if vlane_axis in g) + with InsertionPoint(op): + sram_c = Operation.create( + "memref.collapse_shape", results=[collapsed_ty], operands=[sram], + attributes={"reassociation": reassoc}).results[0] + sidx = [_const(0)] * (len(target) - 1) + [sram_idx] + _emit_asm(sram_c, sidx, dram_idx, new_vlane, target, dr, tl, st) + op.erase() + continue + + # >4 effective dims: affine.for peel (mirrors decompose_transfer peel path) + peeled, inner = eff[:-4], eff[-4:] + ndim = len(tile_shape) + inner_shape = [tile_shape[d] for d in inner] + inner_strides = [tile_stride[d] for d in inner] + dr = [dram_stride[d] for d in inner] + tl = [tile_stride[d] for d in inner] + st = [subtile[d] for d in inner] if subtile is not None else None + if vlane_axis in inner: + new_vlane = inner.index(vlane_axis) + elif vlane_axis in peeled: + raise NotImplementedError( + f"vlane split axis {vlane_axis} peeled into the outer loop nest") + else: + new_vlane = 0 + split_extent = tile_shape[vlane_axis] + nr_outerloop = max( + (split_extent + vectorlane * vlane_stride - 1) // (vectorlane * vlane_stride), 1) + new_size = nr_outerloop * vlane_stride + target_stride = tile_stride[vlane_axis] + + def _phys(d): + s = tile_stride[d] + return s // split_extent * new_size if s > target_stride else s + + static_sizes = [1] * ndim + for d in inner: + static_sizes[d] = tile_shape[d] + res_ty = MemRefType.get( + inner_shape, elem, + layout=StridedLayoutAttr.get(0, inner_strides), memory_space=space) + + cur_ip = InsertionPoint(op) + ivs = [] + for d in peeled: + floop = affine.AffineForOp(0, tile_shape[d], 1, ip=cur_ip) + floop.operation.attributes["inner_loop"] = BoolAttr.get(True) + ivs.append(floop.induction_variable) + with InsertionPoint(floop.body): + affine.AffineYieldOp([]) + cur_ip = InsertionPoint.at_block_terminator(floop.body) + + npeel = len(peeled) + with cur_ip: + sub = Operation.create( + "memref.subview", results=[res_ty], operands=[sram], + attributes={"static_offsets": DenseI64ArrayAttr.get([0] * ndim), + "static_sizes": DenseI64ArrayAttr.get(static_sizes), + "static_strides": DenseI64ArrayAttr.get([1] * ndim), + "operandSegmentSizes": DenseI32ArrayAttr.get([1, 0, 0, 0])} + ).results[0] + sram_expr = AffineExpr.get_dim(0) * _phys(peeled[0]) + for k in range(1, npeel): + sram_expr = sram_expr + AffineExpr.get_dim(k) * _phys(peeled[k]) + sram_off_val = Operation.create( + "affine.apply", results=[idx_ty], operands=list(ivs), + attributes={"map": AffineMapAttr.get(AffineMap.get(npeel, 0, [sram_expr]))} + ).results[0] + dram_expr = AffineExpr.get_dim(0) + for k in range(npeel): + dram_expr = dram_expr + AffineExpr.get_dim(k + 1) * dram_stride[peeled[k]] + dram_idx_val = Operation.create( + "affine.apply", results=[idx_ty], operands=[dram_idx, *ivs], + attributes={"map": AffineMapAttr.get(AffineMap.get(npeel + 1, 0, [dram_expr]))} + ).results[0] + zero = _const(0) + _emit_asm(sub, [zero, zero, zero, sram_off_val], dram_idx_val, new_vlane, + inner_shape, dr, tl, st) + op.erase() + + +def lower_text(text, timing=False): + if OP_NAME not in text: + return text + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(text) + run(m, timing=timing) + return str(m) + + +if __name__ == "__main__": + import sys + out = lower_text(open(sys.argv[1]).read()) + (open(sys.argv[2], "w").write(out) if len(sys.argv) > 2 else sys.stdout.write(out)) diff --git a/docs/design/transfer-direct-lowering.md b/docs/design/transfer-direct-lowering.md new file mode 100644 index 00000000..7e693d91 --- /dev/null +++ b/docs/design/transfer-direct-lowering.md @@ -0,0 +1,66 @@ +# Direct togsim.transfer -> Gemmini lowering (drop memref.dma_start) + +## Why +`togsim.transfer` currently lowers `decompose_transfer` -> `memref.dma_start` -> +... -> `lower_dma_to_gemmini` -> gemmini asm. `memref.dma_start` is a REGISTERED +MLIR op with a fixed operand list, so it cannot carry extra RUNTIME operands. The +indirect offset worked around this with a static symbol attribute; the masked-DMA +`low`/`high` clamp bounds are RUNTIME index values (dynamic shapes) and cannot be +attributes -- so they cannot ride on `memref.dma_start` at all. + +Fix: keep `togsim.transfer` (unregistered -> arbitrary operands) alive through the +pipeline and lower it DIRECTLY to gemmini, carrying every runtime descriptor +(dram_idx, vlane, offset, low/high) as operands. This also lets `test-loop-padding` +retire later (the clamp replaces it), collapsing the mlir-opt step. + +## Target transfer op (superset, all optional beyond the base) +``` +togsim.transfer(dram, dram_idx, sram, sram_idx, tag, dma_type, vst, + [offset_spad], # indirect (gather/scatter) + [low_vec, high_vec]) # masked clamp: per-dim valid [low, high) + attrs: dma_kind, dram_stride[], tile_stride[], padding, [subtile_size, async] +``` +- no `low`/`high` -> full transfer `[0, memref_dim)` (current behaviour). +- no `offset` -> base address (current behaviour). +- `low`/`high` are runtime index operands (vectors length ndim), NOT attributes. + +## Semantics of the clamp (step 2; recorded here for the target) +dest tile stays FULL (banking unchanged -> no vlane shift); only the SOURCE sub-box +`[low, high)` per dim is read; dest positions outside it get `fill` (= ops.masked +`other`). This is "source subview + full dest + fill", not a dest subview. + +## memref.dma_start consumers to port (the scope -- all-or-nothing) +1. `lower_dma_to_gemmini` (dma_start -> asm): consume togsim.transfer; ABSORB + decompose_transfer's 4D handling (<=4D direct / collapse unit dims / >4D + affine.for peel + subview + lane-banked phys offset). +2. `dma_fine_grained` (matmul subtile split, marker=subtile_size): split + togsim.transfer instead of dma_start (same loop structure/offsets). +3. `lower_to_vcix._DmaView` (positional read for compute coord): read + togsim.transfer operand layout. +4. `test-loop-padding` (C++ mlir-opt): verify it tolerates an unregistered + togsim.transfer in the module (allow-unregistered) and pads loops as before. + +## Staging (step 1 = behaviour-preserving; regression 0; low/high NOT added yet) +- 1i. Draft merged `lower_transfer_to_gemmini` (transfer -> asm + 4D handling) in a + NEW file. A/B validate its gemmini asm against the current + decompose->...->lower chain on sample post-vcix MLIR (gemm, conv, pad, + pointwise, gather). NOT wired into the live pipeline yet. +- 1ii. Port `dma_fine_grained` + `_DmaView` to the togsim.transfer operand layout. +- 1iii. Remove `decompose_transfer` from PRE_OPT; rewire so togsim.transfer survives + to the merged lowering. Confirm test-loop-padding still runs (unregistered). +- 1iv. Full regression (the CI allowlist locally): add/matmul/conv/attention/pad/ + gather/models. Gate = regression 0. + +## Step 2 (feature, on top of step 1) +- Add `low`/`high` operands to `emit_transfer` (from load index-affine intersect + operand shape, D2) + `fill`; merged lowering emits the clamp into the gemmini + CONFIG/MVIN; Spike MVIN gates per-position `low[d] <= idx[d] < high[d]` else fill. + Drop the ops.masked `where` for pure-load bodies; keep it for compute-mixed. +- Gate: constant_pad correct with divisibility ON (regression 0), then P5 removes + divisibility -> mobilenet wrapper3 holes/segfault fixed. + +## Validation harness +A/B equivalence: for each stage, run the OLD chain and the NEW pass on the same +post-vcix MLIR dump and diff the emitted gemmini asm (structurally: same +CONFIG/CONFIG2/3/[4]/MVIN sequence, same packed rs1/rs2, same loop nests). Then the +end-to-end Spike allclose + timing tests. From cef1c5b447681df0ef96971947befdf4749a9da5 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 15:32:22 +0900 Subject: [PATCH 04/13] [Frontend] Port the trace path (build_tog) to togsim.transfer/togsim.wait Finishes the memref.dma_start removal for the trace/timing path. build_tog's TogBuilder now decodes togsim.transfer (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) and togsim.wait (tag,tag_idx,num_elements) instead of memref.dma_start/dma_wait: op dispatch, _dma_start_fields (positional decode + direction from dma_kind / dma_type), the tag-user base-address scan, and the wait operand reads. TOGDMANode.tile_stride keeps reading the dram_stride attr (unchanged semantics). build_skeleton's wait marker/dispatch move to togsim.wait so barriers still become togsim.memory_barrier. dep_analysis needs no change (the decoder keeps the same field keys). Trace/timing (TOGSim, timing_mode=1) validated: add 5/5, matmul 11/11 (async DMA + tag_idx barrier pairing). No memref.* DMA ops remain in either path. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- .../mlir/passes/build_skeleton.py | 4 +- PyTorchSimFrontend/mlir/passes/build_tog.py | 93 ++++++++++++------- 2 files changed, 61 insertions(+), 36 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index cc0fbd26..2edbfec3 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -47,7 +47,7 @@ ) #: Marker op names for the passes/__init__ fast-path (skip parsing if absent). -MARKERS = ("togsim.transfer", "memref.dma_wait") +MARKERS = ("togsim.transfer", "togsim.wait") #: Ops the DCE must never remove (loops, terminators, our API ops). _KEEP = { @@ -514,7 +514,7 @@ def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): _emit_one_dma(ctx, op, node, builder, bufs, tags) originals.append(op) n_dma += 1 - elif name == "memref.dma_wait": + elif name == "togsim.wait": if _emit_one_wait(ctx, op, tags): n_wait += 1 originals.append(op) diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index ae515010..5a40feec 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -608,11 +608,11 @@ def bool_true(k): self.print_operation(inner, loop_node) return - if name == "memref.dma_start": + if name == "togsim.transfer": self._handle_dma_start(op, node) return - if name == "memref.dma_wait": + if name == "togsim.wait": self._handle_dma_wait(op, node) return @@ -716,36 +716,60 @@ def _handle_compute(self, op, node): return self._append_vector_compute(node, op) - # ---- dma_start ---- - def _dma_start_fields(self, op): - """Decode memref.dma_start operands by memref ranks. + # ---- transfer (formerly memref.dma_start) ---- + @staticmethod + def _transfer_is_load(op, dma_type_operand): + """True if a `togsim.transfer` is a load (DRAM -> SRAM), False for a store. - Layout: src[srcIdx], dst[dstIdx], numElements, tag[tagIdx], stride, - numElementsPerStride. - """ + Prefer the `dma_kind` attr ("MVIN"/"MVIN2"/"MVIN3" load, "MVOUT" store); + fall back to the `dma_type` operand constant (MVIN=2/MVIN2=1/MVIN3=14 are + loads, MVOUT=3 is a store).""" + oper = op.operation + if "dma_kind" in oper.attributes: + try: + kind = ir.StringAttr(oper.attributes["dma_kind"]).value + except Exception: + kind = str(oper.attributes["dma_kind"]).strip('"') + if kind: + return kind.upper().startswith("MVIN") + c = _const_index_value(dma_type_operand) + if c is not None: + return c in (1, 2, 14) + return True + + def _dma_start_fields(self, op): + """Decode a `togsim.transfer` into the legacy src/dst view. + + togsim.transfer operand layout (mirrors build_skeleton._transfer_fields / + lower_transfer_to_gemmini): + dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst[, offset] + The DRAM side is always operand[0]/[1], the SRAM spad operand[2]/[3], the + runtime tag slot operand[4] (tag memref) + operand[5] (tag_idx). The + optional indirect-offset spad is operand[8]. + + Direction (from dma_kind / dma_type) decides the src/dst mapping so the + rest of build_tog keeps the old memref.dma_start convention: for a load + the DRAM side is the SOURCE and the SRAM spad the DEST; a store reverses + it. src/dst therefore still carry the right memory spaces (DRAM=0, + SRAM=1) that `_handle_dma_start` recovers is_write from.""" operands = list(op.operands) - i = 0 - src = operands[i] - src_type = src.type - src_rank = len(ir.MemRefType(src_type).shape) - i += 1 - src_indices = operands[i:i + src_rank] - i += src_rank - dst = operands[i] - dst_type = dst.type - dst_rank = len(ir.MemRefType(dst_type).shape) - i += 1 - dst_indices = operands[i:i + dst_rank] - i += dst_rank - i += 1 # numElements - tag = operands[i] - tag_rank = len(ir.MemRefType(tag.type).shape) - i += 1 - tag_indices = operands[i:i + tag_rank] + dram, dram_idx = operands[0], operands[1] + sram, sram_idx = operands[2], operands[3] + tag, tag_idx = operands[4], operands[5] + dma_type = operands[6] + offset = operands[8] if len(operands) > 8 else None + + if self._transfer_is_load(op, dma_type): # DRAM -> SRAM + src, src_idx = dram, dram_idx + dst, dst_idx = sram, sram_idx + else: # SRAM -> DRAM + src, src_idx = sram, sram_idx + dst, dst_idx = dram, dram_idx return { - "src": src, "src_type": src_type, "src_indices": src_indices, - "dst": dst, "dst_type": dst_type, "dst_indices": dst_indices, - "tag": tag, "tag_indices": tag_indices, + "src": src, "src_type": src.type, "src_indices": [src_idx], + "dst": dst, "dst_type": dst.type, "dst_indices": [dst_idx], + "tag": tag, "tag_indices": [tag_idx], + "dma_type": dma_type, "offset": offset, } def _handle_dma_start(self, op, node): @@ -776,7 +800,7 @@ def _handle_dma_start(self, op, node): tile_size = [int(x) for x in tile_shape] tile_stride = [] - ds = _int_array_attr(oper, "dram_stride") + ds = _int_array_attr(oper, "dram_stride") # TOGDMANode.tile_stride keeps the DRAM stride (as before) if ds: tile_stride = list(ds) @@ -857,10 +881,11 @@ def _handle_dma_start(self, op, node): # ---- dma_wait ---- def _handle_dma_wait(self, op, node): oper = op.operation + # togsim.wait operands: (tag[0], tag_idx[1], num_elements[2]). tag_idx is + # now a single operand (memref.dma_wait had a variadic [tag_idx]). operands = list(oper.operands) tag = operands[0] - tag_rank = len(ir.MemRefType(tag.type).shape) - tag_indices = operands[1:1 + tag_rank] + tag_indices = operands[1:2] tag_index_list = [] tag_stride_list = [] @@ -880,11 +905,11 @@ def _handle_dma_wait(self, op, node): tag_stride_list = _collect_coefficients(amap.results[0]) tag_divider_list = _collect_dividers(amap.results[0]) - # base address: scan users of tag memref for a dma_start. + # base address: scan users of tag memref for a togsim.transfer. address = "arg" for use in tag.uses: user = use.owner - if user.name == "memref.dma_start": + if user.name == "togsim.transfer": f = self._dma_start_fields(user) dst_space = _memref_space(f["dst_type"]) src_space = _memref_space(f["src_type"]) From 3a91b023e18d2d2a7059f26a937463dbb29d31e6 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 15:43:28 +0900 Subject: [PATCH 05/13] [Frontend] Detect the indirect offset operand by the `indirect` attr, not operand count len(operands) > 8 is ambiguous once the masked-DMA clamp adds more optional operands (low/high after offset). Gate on the `indirect` attr instead, so offset vs mask operands are disambiguated by explicit flags. Behavior-preserving: gather/add pass on both the functional (Spike) and trace (TOGSim) paths. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- PyTorchSimFrontend/mlir/passes/build_skeleton.py | 2 +- PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index 2edbfec3..c024a366 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -438,7 +438,7 @@ def _transfer_fields(op): carries the offset symbol name in its "name" attribute (matching lower_transfer_to_gemmini's offset_sym derivation).""" operands = list(op.operands) - offset = operands[8] if len(operands) > 8 else None + offset = operands[8] if "indirect" in op.attributes else None return { "dram": operands[0], "dram_idx": operands[1], "sram": operands[2], "sram_idx": operands[3], diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py index 037bc152..6bc9f6fc 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -77,7 +77,7 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): for op in targets: op_operands = list(op.operands) dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vst = op_operands[:8] - offset_sym = (op_operands[8].owner.attributes["name"] if len(op_operands) > 8 else None) + offset_sym = (op_operands[8].owner.attributes["name"] if "indirect" in op.attributes else None) kind = op.attributes["dma_kind"].value dma_type_val = _const_int(dma_type) # MVIN(2)/MVIN2(1)/MVIN3(14)/MVOUT(3) is_mvin = dma_type_val in (MVIN, MVIN2, MVIN3) From b77c93d15cbe59eafdc994a11150ea51f1c45cb1 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 15:57:45 +0900 Subject: [PATCH 06/13] [Frontend] Fix stale build_skeleton docstrings (memref.dma_wait -> togsim.wait) The trace producer now reads togsim.wait (not memref.dma_wait); update the pass docstring and comments to match. No code change. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- PyTorchSimFrontend/mlir/passes/build_skeleton.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index c024a366..5a8c35b7 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -8,7 +8,7 @@ * `togsim.transfer` -> `togsim.dma(...) {tag_id, is_async, ...}` carrying the runtime tag index operand (`%tag[%idx]`). - * `memref.dma_wait` -> `togsim.memory_barrier(tag_idx) {tag_id, write_bufs}`, + * `togsim.wait` -> `togsim.memory_barrier(tag_idx) {tag_id, write_bufs}`, the explicit async-DMA sync. It pairs with its dma by the RUNTIME tag slot (tag_id + the tag index), not a compile-time id: one static dma op runs once per loop @@ -78,7 +78,7 @@ def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_buf """Insert a `togsim.dma` before the original `togsim.transfer`. `tag_id` is the identity of this DMA's tag memref. An async DMA pairs with - its `togsim.memory_barrier` (the original dma_wait) by the RUNTIME tag slot + its `togsim.memory_barrier` (the original togsim.wait) by the RUNTIME tag slot -- (tag_id, tag_index) -- not a compile-time identifier: one static dma op runs once per loop iteration, each with a different runtime `%tag[%idx]` slot, so only a runtime key can pair iteration i's dma with iteration i's wait. @@ -119,7 +119,7 @@ def _emit_dma(ctx, dma_node, tag_id, dram_index, tag_index, read_bufs, write_buf def _emit_memory_bar(ctx, anchor_op, tag_id, tag_index, write_bufs): """Insert a `togsim.memory_barrier` before `anchor_op` -- the explicit - async-DMA sync that was the original `memref.dma_wait`. It pairs with its + async-DMA sync that was the original `togsim.wait`. It pairs with its async `togsim.dma` by the RUNTIME tag slot (tag_id + tag_index), and carries the SRAM buffer that dma loaded so consumers gate on data-arrival, not on the async dma's issue-complete.""" @@ -146,7 +146,7 @@ def _flatten_add(expr): def _neg_coeff_dim(summand): """If `summand` is `dim * c` with a negative constant `c`, return that dim's position; else None. lower_to_vcix tags each accumulation (reduction) loop var - with coefficient -1 in the dma_wait tag index -- a SENTINEL marking the + with coefficient -1 in the togsim.wait tag index -- a SENTINEL marking the reduction axis, not an arithmetic offset (legacy TileGraphParser skips stride -1 for the same reason).""" if not ir.AffineMulExpr.isinstance(summand): @@ -389,7 +389,7 @@ def of(self, name): class _TagIds: """Identity of a DMA's tag memref -> stable small int, plus the SRAM buffer that tag's async DMA loads. An async dma and its memory_barrier (the original - dma_wait) share a tag memref; this assigns it a tag_id (so the runtime can + togsim.wait) share a tag memref; this assigns it a tag_id (so the runtime can pair them by the runtime tag slot) and remembers the loaded buffer so the barrier can release it to consumers. Pairing is by tag, never a static id.""" @@ -477,7 +477,7 @@ def _emit_one_dma(ctx, op, node, builder, bufs, tags): def _emit_one_wait(ctx, op, tags): - """Rewrite one memref.dma_wait as togsim.memory_barrier -- the explicit + """Rewrite one togsim.wait as togsim.memory_barrier -- the explicit async-DMA sync already in the IR. Paired with its dma by the tag memref (tag_id) and the runtime tag index; carries the buffer the dma loaded. Returns True iff emitted (a wait whose tag no dma used is dropped).""" @@ -496,7 +496,7 @@ def _emit_one_wait(ctx, op, tags): def _emit_dmas_and_waits(ctx, block, builder, dma_by_op, bufs): - """Step 2: rewrite togsim.transfer -> togsim.dma and memref.dma_wait -> + """Step 2: rewrite togsim.transfer -> togsim.dma and togsim.wait -> togsim.memory_barrier in program order. An async dma and its barrier are paired by the RUNTIME tag slot (tag_id + tag index), not a compile-time id: one static dma op runs per loop iteration with a different `%tag[%idx]`, so From 0581f71d78058207d0012e9c04882987c127e6ea Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 17:18:58 +0900 Subject: [PATCH 07/13] [Frontend] Lower togsim.transfer to a DMA descriptor + CONFIG_DESC (no CONFIG/2/3/4) Instead of packing the transfer params into four CONFIG asm instructions, build a 144-byte DMA descriptor as an 18xi64 memref.global (.data, non-constant so runtime fields can be written) and emit a single CONFIG_DESC that hands its address to the MVIN/MVOUT (which read the struct, matching the Spike-side torchsim_config_desc). This is the vehicle for the masked-DMA low/high clamp: dim_low/dim_high/fill are descriptor fields (currently defaulted, so behavior is unchanged), not new special registers. - _pack_desc encodes the fields into the struct's i64 slots (dim_size/low/high/mm_stride/ spad_stride/element_size/vlane/config_type/flags/indirect stride+esize). - descriptor globals are deduped by content; the indirect offset spad address is a runtime value, stored into slot 15 before the transfer (flags.bit0 = indirect). Requires the Spike torchsim_config_desc instruction (riscv-isa-sim). Validated functional (Spike) + trace (TOGSim): add, matmul (subtile+async), softmax, gather/scatter (indirect runtime store), constant_pad all pass. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- .../mlir/passes/lower_transfer_to_gemmini.py | 80 +++++++++++++++---- 1 file changed, 65 insertions(+), 15 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py index 6bc9f6fc..984a39c5 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -19,8 +19,9 @@ from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation -CONFIG, CONFIG2, CONFIG3, CONFIG4 = 0, 4, 5, 6 +CONFIG, CONFIG2, CONFIG3, CONFIG4, CONFIG_DESC = 0, 4, 5, 6, 7 MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 +DESC_SLOTS = 18 # 144-byte DMA descriptor as 18 x i64 (see project-dma-descriptor) CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} MAX_TENSOR_DIM = 4 @@ -29,16 +30,70 @@ def run(module, timing=False, vectorlane=128, **_): from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, DenseI32ArrayAttr, StridedLayoutAttr, AffineMap, AffineMapAttr, - AffineExpr, BoolAttr, FlatSymbolRefAttr, TypeAttr) + AffineExpr, BoolAttr, FlatSymbolRefAttr, TypeAttr, + DenseElementsAttr, RankedTensorType, StringAttr, UnitAttr) from mlir.dialects import affine, llvm, arith, memref i64 = IntegerType.get_signless(64) idx_ty = IndexType.get() + module_top = module.operation.regions[0].blocks[0] sym2type = {} - for g in module.operation.regions[0].blocks[0].operations: + for g in module_top.operations: if g.operation.name == "memref.global": sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + desc_ty = MemRefType.get([DESC_SLOTS], i64) + desc_globals = {} # slots tuple -> global sym name (dedup) + desc_counter = [0] + + def _pack_desc(shape4, dram4, spad4, elem_bytes, vlstride, vsa4, cfg_type, indirect, + ind_stride, ind_esize): + # 18 i64 slots matching the C dma_descriptor byte layout (little-endian). + lo = [0, 0, 0, 0] + hi = list(shape4) + def s2(a, b): return (a & 0xFFFFFFFF) | ((b & 0xFFFFFFFF) << 32) + m = 0xFFFFFFFFFFFFFFFF + slots = [ + s2(shape4[0], shape4[1]), s2(shape4[2], shape4[3]), + s2(lo[0], lo[1]), s2(lo[2], lo[3]), + s2(hi[0], hi[1]), s2(hi[2], hi[3]), + dram4[0] & m, dram4[1] & m, dram4[2] & m, dram4[3] & m, + spad4[0] & m, spad4[1] & m, spad4[2] & m, spad4[3] & m, + (elem_bytes & 0xFFFF) | ((vlstride & 0xFFFF) << 16) | ((vsa4 & 0xFF) << 32) + | ((cfg_type & 0xFF) << 40) | (((1 if indirect else 0) & 0xFFFF) << 48), + 0, # +120 indirect_addr (runtime) + (ind_stride & 0xFFFF) | ((ind_esize & 0xFFFF) << 16), # +128 + 0, # +136 fill (runtime/step3) + ] + return slots + + def _desc_global(slots): + key = tuple(slots) + if key not in desc_globals: + name = f"dma_desc_{desc_counter[0]}" + desc_counter[0] += 1 + import numpy as np + init = DenseElementsAttr.get( + np.array([_i64_signed(v) for v in slots], dtype=np.int64), + type=RankedTensorType.get([DESC_SLOTS], i64)) + with InsertionPoint.at_block_begin(module_top): + Operation.create("memref.global", results=[], operands=[], attributes={ + "sym_name": StringAttr.get(name), + "sym_visibility": StringAttr.get("private"), + "type": TypeAttr.get(desc_ty), + "initial_value": init}) + desc_globals[key] = name + return desc_globals[key] + + def desc_ptr_and_store_indirect(slots, ind_addr_i64): + # get_global -> byte pointer; for indirect, store the runtime addr into slot 15. + name = _desc_global(slots) + g = memref.GetGlobalOp(desc_ty, name).result + if ind_addr_i64 is not None: + memref.StoreOp(ind_addr_i64, g, [arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, 15)).result]) + base = memref.ExtractAlignedPointerAsIndexOp(g).result + return arith.IndexCastOp(i64, base).result + def i64_const(value): return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result @@ -117,24 +172,19 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, sram_c = MemRefType(sram_mem.type) dram_addr = elem_addr_i64(dram, [dram_idx_val], dram_ty, elem_bytes) spad_addr = elem_addr_i64(sram_mem, sram_indices, sram_c, elem_bytes) - cfg_rs1 = i64_const(((shape4[0] & 0xFFFF) << 48) | ((shape4[1] & 0xFFFF) << 32) - | ((shape4[2] & 0xFFFF) << 16) | (shape4[3] & 0xFFFF)) - cfg_rs2 = i64_const((vlane_stride << 32) | ((config_type & 0x3) << 17) - | ((1 if indirect else 0) << 16) - | ((vsa4 & 0x3) << 14) | elem_bytes) - asm(CONFIG, cfg_rs1, cfg_rs2) - asm(CONFIG2, i64_const((dram4[0] << 32) | (dram4[1] & 0xFFFFFFFF)), - i64_const((dram4[2] << 32) | (dram4[3] & 0xFFFFFFFF))) - asm(CONFIG3, i64_const((spad4[0] << 32) | (spad4[1] & 0xFFFFFFFF)), - i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) + ind_addr = None + ind_stride = ind_esize = 0 if indirect: sym = FlatSymbolRefAttr(offset_sym).value if not isinstance(offset_sym, str) else offset_sym off_ty = sym2type[sym] ind_base = memref.ExtractAlignedPointerAsIndexOp(memref.GetGlobalOp(off_ty, sym).result).result ind_addr = arith.IndexCastOp(i64, ind_base).result ind_esize = _elem_bytes(off_ty.element_type) - off_stride = IntegerAttr(op.attributes["offset_stride"]).value - asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (off_stride & 0xFFFF))) + ind_stride = IntegerAttr(op.attributes["offset_stride"]).value + slots = _pack_desc(shape4, dram4, spad4, elem_bytes, vlane_stride, vsa4, + config_type, indirect, ind_stride, ind_esize) + desc_ptr = desc_ptr_and_store_indirect(slots, ind_addr) + asm(CONFIG_DESC, desc_ptr, i64_const(0)) asm(dma_type_val, dram_addr, spad_addr) if offset_sym is not None: From 80db6b24505cc83745d5ecdbd99fb4f2fdea1bdb Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 6 Jul 2026 17:41:21 +0900 Subject: [PATCH 08/13] [Build] Bump spike to v1.0.4 (DMA descriptor / CONFIG_DESC) The togsim.transfer descriptor lowering requires the spike torchsim_config_desc instruction (riscv-isa-sim PR #5, tag v1.0.4). Track the new tag so build-from-source picks up the descriptor-capable spike. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- thirdparty/github-releases.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/github-releases.json b/thirdparty/github-releases.json index 390b68fb..792aa23c 100644 --- a/thirdparty/github-releases.json +++ b/thirdparty/github-releases.json @@ -13,7 +13,7 @@ }, "spike": { "repository": "PSAL-POSTECH/riscv-isa-sim", - "release_tag": "v1.0.3", + "release_tag": "v1.0.4", "asset_name": "spike-release.tar.gz" } } From fb6bed47411ce3d80e5218d92d5efaa19af76663 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 11:34:03 +0900 Subject: [PATCH 09/13] [Build] Bump spike to v1.0.5 (masked-DMA clamp + MVOUT accumulate) The masked-DMA clamp needs spike to read the descriptor dim_low/dim_high + desc_fill, and index_add needs the MVOUT accumulate flag. Pin v1.0.5. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- thirdparty/github-releases.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/github-releases.json b/thirdparty/github-releases.json index 792aa23c..2618c236 100644 --- a/thirdparty/github-releases.json +++ b/thirdparty/github-releases.json @@ -13,7 +13,7 @@ }, "spike": { "repository": "PSAL-POSTECH/riscv-isa-sim", - "release_tag": "v1.0.4", + "release_tag": "v1.0.5", "asset_name": "spike-release.tar.gz" } } From 906884a17fc109537c240362939ee1e25d0d9dd8 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 11:34:03 +0900 Subject: [PATCH 10/13] [Frontend] Masked-DMA dynamic clamp: drop tile divisibility, fill non-dividing tails Drop the index_expr / indirect tile-divisibility recompile so pad_vlane_tile's lane-aligned tiles are kept instead of shrunk to a divisor of the dim. A non-dividing tile's per-tile remainder is handled by a masked-DMA clamp: - _emit_clamp emits per tile axis a dynamic [low, high) as affine.min/max of that axis' loop iv (the tile base), so the last partial tile and the pad borders fall out per iteration; low = max(0, glo - base), high = min(tile, ghi - base). The pointwise path (_masked_bounds) derives [glo, ghi) per dim (MVOUT [0, output); MVIN of a padded input [pad, pad+input_extent)); the MLIR templates (def_dma_op, e.g. gemm) derive it from the DRAM extent so a matmul tile padding past M/N/K is filled with 0 instead of MAC-ing garbage. - emit_transfer carries the (low, high) index operands + masked_axes; lower_transfer_to_gemmini writes the runtime low/high into the descriptor dim_low/dim_high through an i32-view (byte layout unchanged) + masked flag + desc_fill. One desc_ptr helper builds every descriptor: a fully static one is a deduped i64 global, while any runtime-written field (masked dim_low/dim_high and/or the indirect address, stored as two i32 words) uses a unique i32-view global, so a single descriptor carries both a masked clamp and a gather/scatter, and indirect DMAs drop their divisibility recompile too. - The box-excluded fill is the consuming reduction identity (reduction_init: sum 0, prod 1, max -inf, min +inf; 0 otherwise -- matmul MAC and stores use 0). Log-sum- exp exception (softmax / log_softmax): a sum whose node came from exp reduces exp(input), so the PRIMARY input (indexed by a reduction itervar) is filled with -inf; broadcast operands (the subtracted max) keep their finite identity, else input - (-inf) = +inf -> exp = inf -> NaN. - index_add: instead of the compute gather-add-overwrite (which loses duplicate indices landing in the same non-dividing tile), the MVOUT accumulates out[idx] += val via an accumulate flag; the DMA processes positions sequentially so duplicates accumulate correctly regardless of tile size. acc_float selects float vs integer add (element_size alone is ambiguous). With no caller left, is_dim_dividable / adjust_tile_to_divisible and the TileConstraint.must_divide_dim machinery (trim_large_tail bias, adjust loop) are removed. Fixes elementwise, constant_pad (mobilenet ragged vlane holes), reductions (sum/max/min/prod/softmax/log_softmax/var/std/layernorm/mean), gather/scatter (index_add with duplicates) and matmul/addmm on non-dividing shapes. Requires spike with the descriptor dim_low/dim_high clamp, desc_fill and the MVOUT accumulate flag. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- .../mlir/mlir_codegen_backend.py | 149 +++++++++++++++--- PyTorchSimFrontend/mlir/mlir_common.py | 52 +----- PyTorchSimFrontend/mlir/mlir_template.py | 21 ++- .../mlir/passes/lower_transfer_to_gemmini.py | 104 +++++++++--- 4 files changed, 235 insertions(+), 91 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 816f88bc..e0edbe2e 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -591,8 +591,10 @@ def load(self, name: str, index: sympy.Expr): sram_var, sram_index_var = self.get_scratchpad_buffer(dtype, name, local_tile_desc, index) compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) + masked_bounds = self._masked_bounds(name, index, dram_stride, local_tile_desc, is_load=True, buffer=self.dma_loads) code = self.emit_transfer("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, int(padding), offset=offset_desc) + dram_shape, tile_shape, dram_stride, tile_stride, int(padding), offset=offset_desc, + masked_bounds=masked_bounds, masked_fill=self._masked_fill_bits(dtype, index)) self.cse.generate(self.dma_loads, code, assignment = False) # FIXME: assignment = False does not support caching with self.override_buffer_cse(buffer=self.loads): @@ -606,15 +608,18 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) offset_desc = None # Handle scatter store + accumulate = False if self._has_indirect(index): # Convert the output buffer type to the inplace buffer arg_name = V.graph.scheduler.mutation_real_name.get(name, name) if arg_name not in self.kernel_group.args.inplace_buffers: self.kernel_group.args.make_inplace(arg_name, arg_name) - if mode == "atomic_add": - loaded_value = ops.load(name, index) - value = ops.add(loaded_value, value) + # index_add: let the MVOUT do out[idx] += val. The DMA processes positions + # sequentially, so duplicate indices accumulate correctly regardless of tile + # size -- unlike the compute gather-add-overwrite, which loses duplicates that + # land in the same (non-dividing) tile. + accumulate = (mode == "atomic_add") index, offset_desc = self.convert_indirect_indexing(index) dram_var = self.kernel_group.args.output(name) @@ -656,8 +661,10 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) sram_index_var = self.spad_buffer_dict[str(value)][3] # Generate DMA instruction + masked_bounds = self._masked_bounds(name, index, dram_stride, local_tile_desc, is_load=False, buffer=self.dma_stores) code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, 0, offset=offset_desc) + dram_shape, tile_shape, dram_stride, tile_stride, 0, offset=offset_desc, masked_bounds=masked_bounds, + accumulate=accumulate, acc_float=dtype.is_floating_point) self.dma_stores.writeline(common.DeferredLine(name, code)) def reduction(self, dtype, src_dtype, reduction_type, value): @@ -796,14 +803,6 @@ def _has_indirect(self, expr): return any(s.name in self.indirect_symbols for s in expr.free_symbols) def _index_expr(self, tile_desc, renamed_expression, index, base_vector_index): - # In case of index expr, dimension size should be divisible by tile size - if not self.kernel_group.tile_desc.is_dim_dividable(self.ranges): - new_tile_size = self.kernel_group.tile_desc.adjust_tile_to_divisible(self.ranges) - prior_tile_size, prior_ranges = self.kernel_group.tile_desc.get_tile_size(), self.ranges - self.kernel_group.tile_desc.set_tile_size(new_tile_size) - self.reset("recompile") - raise mlir_common.RecompileSignal(f"Index access (tile size {prior_tile_size} is not divisible by {prior_ranges})") - tile_size_per_lane = tile_desc.get_tile_size_per_lane() compute_vec_size = tile_desc.get_compute_vec_size() strides = tile_desc.get_tile_stride_per_lane() @@ -1359,12 +1358,111 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe if len(self.itervars) == 1 and self.reduction_depth == 0: # In case of reduction loop only case, we will add dummy loop so shift it once dram_stride = [0] + dram_stride[:-1] + + # Stash the tile-axis -> loop-dim map so load()/store() can build the masked-DMA + # per-dim [low, high) clamp (needs the loop iv of each tile axis). See _masked_bounds. + self._last_local_dims = local_dims return local_tile_desc, index_var, dram_stride + _FILL_BITVIEW = {torch.float32: torch.int32, torch.float16: torch.int16, + torch.bfloat16: torch.int16, torch.float64: torch.int64} + + def _masked_fill_bits(self, dtype, index): + """Raw bits for the masked-DMA tail fill = the consuming reduction's identity + (reduction_init: sum->0, prod->1, max->-inf, min->+inf), in the LOAD dtype; 0 for + a non-reduction load. So a non-dividing reduction's tail is the reduction identity, + not garbage/0. See Reduction.default_accumulator (PyTorch's canonical table). + + Log-sum-exp exception (softmax): a sum whose node came from exp reduces exp(input), + so the PRIMARY input's tail must be -inf (exp(-inf)=0), not the sum identity 0. + Only the primary reduced tensor (indexed by a reduction itervar) is overridden; + broadcast operands (e.g. the subtracted max) keep their finite identity, else + input - (-inf) = +inf -> exp = inf -> NaN.""" + node = getattr(self, "current_node", None) + if node is None or getattr(node, "node", None) is None or not node.node.get_reduction_type(): + return 0 + rtype = node.node.get_reduction_type() + is_primary = bool(set(self.itervars[self.reduction_depth:]) & index.free_symbols) + if rtype == "sum" and is_primary and any("exp" in o.name for o in node.node.origins): + init = "-inf" + else: + init = reduction_init(rtype, dtype) + val = {"-inf": float("-inf"), "inf": float("inf")}.get(init, init) + if isinstance(val, str): # e.g. welford_reduce -> "0.0" + val = float(val) + t = torch.tensor(val, dtype=dtype) + view = self._FILL_BITVIEW.get(dtype) + bits = int(t.view(view).item()) if view is not None else int(t.item()) + return bits & ((1 << (t.element_size() * 8)) - 1) + + def _masked_bounds(self, name, index, dram_stride, local_tile_desc, is_load, buffer): + """Per tile-axis dynamic [low, high) clamp for a masked DMA. + + For tile axis d (loop dim k, loop iv %index_k = the tile base), the valid GLOBAL + range is [glo, ghi): MVOUT writes [0, output_extent); MVIN reads the padded input, + valid where the input coord is in-bounds -> [pad_d, pad_d + input_extent_d). The + tile-LOCAL clamp is then low = max(0, glo - base), high = min(tile, ghi - base), + emitted as affine.max/affine.min of the loop iv so the last (partial) tile and the + pad borders fall out naturally. Returns [(tile_axis, low_var, high_var), ...] for + the non-trivial axes; low/high are index SSA vars. + """ + local_dims = self._last_local_dims + tile_size = local_tile_desc.get_tile_size() + const = int(index.as_coeff_Add()[0]) if not index.is_number else int(index) + # index const = -sum(pad_d * dram_stride[d]); recover per-dim pad greedily (desc stride). + pad = [0] * len(tile_size) + if is_load and const < 0: + rem = -const + for d in sorted(range(len(dram_stride)), key=lambda x: -dram_stride[x]): + s = dram_stride[d] + if s > 0 and d < len(pad): + pad[d] = rem // s + rem -= pad[d] * s + in_shape, in_stride = self.buffer_types[name][2], self.buffer_types[name][3] + axes = [] + for d, k in enumerate(local_dims): + if d >= len(tile_size) or k >= len(self.ranges): + continue + iv = str(self.itervars[k]) + # A padded load (const < 0, i.e. index shifted by -pad) reads a smaller input: + # clamp per-dim to [pad_d, pad_d + input_extent_d). Every other DMA (stores and + # non-padded/contiguous loads, including collapsed tensors) is valid over the + # whole loop extent -> only the trailing tail needs clamping. + if is_load and const < 0: + ext = next((int(in_shape[j]) for j, st in enumerate(in_stride) + if int(st) == int(dram_stride[d])), None) + glo, ghi = (pad[d], pad[d] + ext) if ext is not None else (0, int(self.ranges[k])) + else: + glo, ghi = 0, int(self.ranges[k]) + axes.append((d, iv, glo, ghi, int(tile_size[d]))) + return self._emit_clamp(axes, buffer) + + def _emit_clamp(self, axes, buffer): + """Emit the per-axis dynamic clamp. axes: [(tile_axis, base_iv, glo, ghi, tile), ...] + -- low = max(0, glo - base), high = min(tile, ghi - base) as affine.max/affine.min of + the loop iv so the last partial tile and the pad borders fall out per iteration. + Returns [(tile_axis, low_var, high_var), ...] for the non-trivial axes only.""" + result = [] + for d, iv, glo, ghi, tile in axes: + if glo == 0 and ghi % tile == 0: # every tile fully valid -> no clamp + continue + high_var = self.apply_cse.generate( + buffer, f"affine.min affine_map<(d0) -> ({tile}, {ghi} - d0)>(%{iv})") + self.register_var_info(high_var, [1, "index"]) + if glo > 0: + low_var = self.apply_cse.generate( + buffer, f"affine.max affine_map<(d0) -> (0, {glo} - d0)>(%{iv})") + self.register_var_info(low_var, [1, "index"]) + else: + low_var = self.get_const_cse(0) + result.append((d, low_var, high_var)) + return result + def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, dram_index_var, sram_var, sram_index_var, dram_shape, tile_shape, dram_stride, tile_stride, padding, - subtile_size=None, async_type=None, offset=None): + subtile_size=None, async_type=None, offset=None, masked_bounds=None, masked_fill=0, + accumulate=False, acc_float=False): """Emit a generic togsim.transfer op for a DMA whose access exceeds the 4D Gemmini descriptor limit. Carries the full N-D access (dram/tile strides + shapes) plus the SSA operands a memref.dma_start needs @@ -1405,6 +1503,10 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp if subtile_size: av = int(async_type) if async_type is not None else 1 attrs += f', subtile_size = {list(subtile_size)}, async = {av} : i64' + if accumulate: # index_add: MVOUT does out[idx] += val (float or integer add) + attrs += f', accumulate = true' + if acc_float: + attrs += f', acc_float = true' # operands: dram, dram_idx, sram, sram_idx, tag, tag_idx, dma_type, vlane_stride [, offset spad] operands = (f'%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' f'%{tag}, %{zero_cse}, %{dma_type}, %{vst}') @@ -1414,6 +1516,18 @@ def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtyp operands += f', %{offset_buf}' optypes += f', {offset_type}' attrs += f', indirect = true, offset_stride = {int(offset_stride)} : i64' + # masked-DMA dynamic clamp: append (low, high) index operands per clamped tile axis; + # masked_axes names the tile axis of each pair so the lowering writes the runtime + # values into the descriptor's dim_low/dim_high before the DMA. See _masked_bounds. + if masked_bounds: + axes = [d for d, _lo, _hi in masked_bounds] + for _d, lo, hi in masked_bounds: + operands += f', %{lo}, %{hi}' + optypes += ', index, index' + attrs += f', masked_axes = {axes}' + # box-excluded positions are filled with the consuming reduction's identity + # (0/1/-inf/+inf, per dtype); 0 for non-reduction loads. See _masked_fill_bits. + attrs += f', masked_fill = {int(masked_fill)} : i64' return f'"togsim.transfer"({operands}) {{{attrs}}} : ({optypes}) -> ()' def allocate_sram_buffer(self, dtype, dram_name, tile_desc, raw_index, buffer=None, forced_name=None): @@ -1498,13 +1612,6 @@ def convert_indirect_indexing(self, index :sympy.Expr): if not self._has_indirect(index): return index, None - # Note: In case of indirect indexing, dimensions should be divisible by tile size - if not self.kernel_group.tile_desc.is_dim_dividable(self.ranges): - new_tile_size = self.kernel_group.tile_desc.adjust_tile_to_divisible(self.ranges) - self.kernel_group.tile_desc.set_tile_size(new_tile_size) - self.reset("recompile") - raise mlir_common.RecompileSignal(f"Indirect access (tile size {self.kernel_group.tile_desc.get_tile_size()} is not divisible by {self.ranges})") - # Process start indirect_dims = [str(dim) for dim in index.free_symbols if str(dim) in self.indirect_symbols] indirect_dims.sort() diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index a70d1c7d..bfd78fd9 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -50,7 +50,7 @@ torch.int8: "i8", torch.uint8: "i8", torch.bool: "i8", - torch.bfloat16: "bf16", + torch.bfloat16: "bf16", # bf16 unsupported (Spike/codegen have no bf16); left as-is } MLIR_TO_DTYPE = { @@ -324,41 +324,6 @@ def apply_divisor(self, axis: int, divisor: int, mode: str = "split"): raise ValueError(f"Unknown mode: {mode}. Supported: 'pad', 'split'.") - def is_dim_dividable(self, dim_sizes: list[int]) -> bool: - if len(dim_sizes) != len(self._tile_size): - raise ValueError("dim_sizes must match the tile size dimensions") - - dim_sizes_cpy = list(dim_sizes) - axis, stride = self.vmap.vlane_split_axis, self.vmap.vlane_stride - remain = dim_sizes_cpy[axis] % stride - if remain: - dim_sizes_cpy[axis] += stride - remain - - return all(d % t == 0 for d, t in zip(dim_sizes_cpy, self._tile_size)) - - def adjust_tile_to_divisible(self, dim_sizes: list[int]) -> list[int]: - """Adjust current tile to be divisible by given dimensions.""" - if len(dim_sizes) != len(self._tile_size): - raise ValueError("dim_sizes must match the tile size dimensions") - - def _adjust_one(dim_size, tile_size): - for candidate in range(tile_size, 0, -1): - if dim_size % candidate == 0: - return candidate - return 1 - - candidate_tile_size = [_adjust_one(d, t) for d, t in zip(dim_sizes, self._tile_size)] - for i in range(len(candidate_tile_size)): - self.tile_constraint[i].must_divide_dim = True - - axis, stride = self.vmap.vlane_split_axis, self.vmap.vlane_stride - remain = candidate_tile_size[axis] % stride - - if remain: - # #201: relax vlane_stride constraints - self.vmap.vlane_stride = 1 - return candidate_tile_size - def scale_tile_dim(self, axis, dim_sz, scale_factor=2): axis_constrinat = self.tile_constraint[axis] current_sz = self._tile_size[axis] @@ -395,8 +360,6 @@ def trim_large_tail(self, ranges: list[int]): constraint = self.tile_constraint[i] if constraint.fixed: continue - elif constraint.must_divide_dim: - BETA = 0 padding_ratio = TileAdjustMixin.get_padding_ratio(tile_range, dim_range) if padding_ratio < self.tail_ratio_threshold: @@ -469,23 +432,14 @@ def get_padding_ratio(tile_range: int, dim_range: int) -> float: @dataclass class TileConstraint: multiple_of: int = 1 - must_divide_dim: bool = False fixed: bool = False def adjust(self, old: int, new: int, dim: int) -> int: if self.fixed: return old # Fixed tile size - tail = new % self.multiple_of - new -= tail - if not self.must_divide_dim: - return max(new, self.multiple_of) - - while new > 0: - if dim % new == 0: - return new - new -= self.multiple_of - raise extension_codecache.TileSizeError("Cannot find suitable tile size under the given constraints.") + new -= new % self.multiple_of + return max(new, self.multiple_of) class MLIRMultiDimTile(TileAdjustMixin): def __init__(self, tile_size, vector_lane, vlane_split_axis=None, vlane_stride=None, forced_vec_size=None): diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index cd3ca018..40dd10bf 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -940,9 +940,28 @@ def generate_dma_code(): zero_cse = self.get_const_cse(0, "index") sram_index_var = ", ".join([f"%{str(zero_cse)}"]*tile_desc.get_nr_dim()) + # masked-DMA clamp: per tile dim, clamp to the real DRAM extent so a tile + # that pads past M/N/K (matmul) or a dim is filled with 0 instead of MAC-ing + # garbage. index_list[d] = iv * stride -> base iv; the extent is the DRAM dim + # matched by that tile dim's stride (tile order != layout order for conv). + # _emit_clamp skips dividing dims (no-op). + _dram_size, _dram_str = node_layout.size, node_layout.stride + _tsize = tile_desc.get_tile_size() + _axes = [] + for _d, _idx in enumerate(index_list): + _syms = list(_idx.free_symbols) + if not _syms or _d >= len(_tsize): + continue + _ext = next((int(_dram_size[j]) for j, st in enumerate(_dram_str) + if _dram_size[j].is_number and int(st) == int(_dram_stride[_d])), None) + if _ext is not None: + _axes.append((_d, str(_syms[0]), 0, _ext, int(_tsize[_d]))) + masked_bounds = self._emit_clamp(_axes, local_code) + code = self.emit_transfer(dma_type, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, dram_shape, tile_shape, _dram_stride, sram_strides, int(padding), - subtile_size=subtile_size if subtile_size else None, async_type=async_type) + subtile_size=subtile_size if subtile_size else None, async_type=async_type, + masked_bounds=masked_bounds, masked_fill=0) local_code.writeline(code) return textwrap.indent(local_code.getvalue(), " "*indent_size).strip() diff --git a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py index 984a39c5..316454f5 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py +++ b/PyTorchSimFrontend/mlir/passes/lower_transfer_to_gemmini.py @@ -2,7 +2,7 @@ Merges decompose_transfer (the <=4D Gemmini-limit handling: drop unit dims / collapse / >4D affine.for peel with lane-banked SRAM offset) with -lower_dma_to_gemmini (the CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm emission). +lower_dma_to_gemmini (the CONFIG_DESC/MVIN|MVOUT asm emission). togsim.transfer is unregistered so it carries every runtime descriptor as an operand -- including the future masked-clamp low/high vectors -- which a registered memref.dma_start cannot. See docs/design/transfer-direct-lowering.md. @@ -19,7 +19,7 @@ from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation -CONFIG, CONFIG2, CONFIG3, CONFIG4, CONFIG_DESC = 0, 4, 5, 6, 7 +CONFIG_DESC = 7 # func7 for the TMA-style descriptor pointer (replaces CONFIG/2/3/4) MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 DESC_SLOTS = 18 # 144-byte DMA descriptor as 18 x i64 (see project-dma-descriptor) CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} @@ -42,15 +42,23 @@ def run(module, timing=False, vectorlane=128, **_): if g.operation.name == "memref.global": sym2type[g.attributes["sym_name"].value] = MemRefType(TypeAttr(g.attributes["type"]).value) + i32 = IntegerType.get_signless(32) desc_ty = MemRefType.get([DESC_SLOTS], i64) + desc_i32_ty = MemRefType.get([DESC_SLOTS * 2], i32) # same bytes, i32-addressable (masked) desc_globals = {} # slots tuple -> global sym name (dedup) desc_counter = [0] + def _i32_signed(v): + v &= 0xFFFFFFFF + return v - (1 << 32) if v >= (1 << 31) else v + def _pack_desc(shape4, dram4, spad4, elem_bytes, vlstride, vsa4, cfg_type, indirect, - ind_stride, ind_esize): + ind_stride, ind_esize, hi=None, accumulate=False, acc_float=False): # 18 i64 slots matching the C dma_descriptor byte layout (little-endian). lo = [0, 0, 0, 0] - hi = list(shape4) + if hi is None: + hi = list(shape4) + masked = any(h < s for h, s in zip(hi, shape4)) # dim_high clamps below extent def s2(a, b): return (a & 0xFFFFFFFF) | ((b & 0xFFFFFFFF) << 32) m = 0xFFFFFFFFFFFFFFFF slots = [ @@ -60,7 +68,9 @@ def s2(a, b): return (a & 0xFFFFFFFF) | ((b & 0xFFFFFFFF) << 32) dram4[0] & m, dram4[1] & m, dram4[2] & m, dram4[3] & m, spad4[0] & m, spad4[1] & m, spad4[2] & m, spad4[3] & m, (elem_bytes & 0xFFFF) | ((vlstride & 0xFFFF) << 16) | ((vsa4 & 0xFF) << 32) - | ((cfg_type & 0xFF) << 40) | (((1 if indirect else 0) & 0xFFFF) << 48), + | ((cfg_type & 0xFF) << 40) + | (((1 if indirect else 0) | (2 if masked else 0) + | (4 if accumulate else 0) | (8 if acc_float else 0)) << 48), 0, # +120 indirect_addr (runtime) (ind_stride & 0xFFFF) | ((ind_esize & 0xFFFF) << 16), # +128 0, # +136 fill (runtime/step3) @@ -85,14 +95,46 @@ def _desc_global(slots): desc_globals[key] = name return desc_globals[key] - def desc_ptr_and_store_indirect(slots, ind_addr_i64): - # get_global -> byte pointer; for indirect, store the runtime addr into slot 15. - name = _desc_global(slots) - g = memref.GetGlobalOp(desc_ty, name).result - if ind_addr_i64 is not None: - memref.StoreOp(ind_addr_i64, g, [arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, 15)).result]) - base = memref.ExtractAlignedPointerAsIndexOp(g).result - return arith.IndexCastOp(i64, base).result + def desc_ptr(slots, masked_pairs=(), fill=0, ind_addr_i64=None): + """Byte pointer to the DMA descriptor global. A fully static descriptor (no + masked clamp, no indirect address) is a deduped i64 global. Anything with a + runtime-written field -- per-dim dim_low/dim_high (masked) and/or the indirect + address -- gets a UNIQUE i32-view global (same byte layout, i32-addressable) so + those int32/i64 fields can be stored individually.""" + import numpy as np + if not masked_pairs and ind_addr_i64 is None: + g = memref.GetGlobalOp(desc_ty, _desc_global(slots)).result + return arith.IndexCastOp(i64, memref.ExtractAlignedPointerAsIndexOp(g).result).result + slots = list(slots) + if masked_pairs: + slots[14] |= (2 << 48) # masked flag (+118 bit1) + slots[17] = fill & 0xFFFFFFFFFFFFFFFF # +136 fill: box-excluded positions + words = [] + for v in slots: + v &= 0xFFFFFFFFFFFFFFFF + words.append(_i32_signed(v & 0xFFFFFFFF)) + words.append(_i32_signed((v >> 32) & 0xFFFFFFFF)) + name = f"dma_desc_{desc_counter[0]}" + desc_counter[0] += 1 + init = DenseElementsAttr.get(np.array(words, dtype=np.int32), + type=RankedTensorType.get([DESC_SLOTS * 2], i32)) + with InsertionPoint.at_block_begin(module_top): + Operation.create("memref.global", results=[], operands=[], attributes={ + "sym_name": StringAttr.get(name), + "sym_visibility": StringAttr.get("private"), + "type": TypeAttr.get(desc_i32_ty), + "initial_value": init}) + g = memref.GetGlobalOp(desc_i32_ty, name).result + def store_word(v32, i32_idx): + memref.StoreOp(v32, g, [arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, i32_idx)).result]) + for axis4, low_val, high_val in masked_pairs: # dim_low i32 idx 4 (+16B), dim_high 8 (+32B) + store_word(arith.IndexCastOp(i32, low_val).result, 4 + axis4) + store_word(arith.IndexCastOp(i32, high_val).result, 8 + axis4) + if ind_addr_i64 is not None: # indirect_addr (i64) at +120 -> words 30/31 + store_word(arith.TruncIOp(i32, ind_addr_i64).result, 30) + shamt = arith.ConstantOp(i64, IntegerAttr.get(i64, 32)).result + store_word(arith.TruncIOp(i32, arith.ShRUIOp(ind_addr_i64, shamt).result).result, 31) + return arith.IndexCastOp(i64, memref.ExtractAlignedPointerAsIndexOp(g).result).result def i64_const(value): return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result @@ -144,6 +186,20 @@ def elem_addr_i64(memref_val, indices, mtype, elem_bytes): subtile = _int_array(op.attributes["subtile_size"]) except KeyError: subtile = None + # masked-DMA dynamic clamp: masked_axes lists the clamped tile axes; the trailing + # (low, high) index operands (2 per axis) are runtime-stored into dim_low/dim_high. + if "masked_axes" in op.attributes: + masked_axes = _int_array(op.attributes["masked_axes"]) + n_base = 9 if "indirect" in op.attributes else 8 + mvals = op_operands[n_base:] + masked_pairs = [(masked_axes[i], mvals[2 * i], mvals[2 * i + 1]) + for i in range(len(masked_axes))] + masked_fill = IntegerAttr(op.attributes["masked_fill"]).value + else: + masked_pairs = [] + masked_fill = 0 + accumulate = "accumulate" in op.attributes # index_add: MVOUT out[idx] += val + acc_float = "acc_float" in op.attributes if timing: op.erase() @@ -161,7 +217,7 @@ def _const(v): return arith.ConstantOp(idx_ty, IntegerAttr.get(idx_ty, v)).result def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, - desc_dram_strides, desc_spad_strides, subtile_shape): + desc_dram_strides, desc_spad_strides, subtile_shape, desc_masked_pairs=None): cfg_shape = subtile_shape if subtile_shape is not None else desc_shape expand = MAX_TENSOR_DIM - len(cfg_shape) shape4 = [1] * expand + list(cfg_shape) @@ -182,9 +238,11 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, ind_esize = _elem_bytes(off_ty.element_type) ind_stride = IntegerAttr(op.attributes["offset_stride"]).value slots = _pack_desc(shape4, dram4, spad4, elem_bytes, vlane_stride, vsa4, - config_type, indirect, ind_stride, ind_esize) - desc_ptr = desc_ptr_and_store_indirect(slots, ind_addr) - asm(CONFIG_DESC, desc_ptr, i64_const(0)) + config_type, indirect, ind_stride, ind_esize, + accumulate=accumulate, acc_float=acc_float) + pairs4 = [(expand + d, lo, hi) for d, lo, hi in (desc_masked_pairs or ())] # 4D-expand axis + desc_ptr_val = desc_ptr(slots, pairs4, masked_fill, ind_addr) + asm(CONFIG_DESC, desc_ptr_val, i64_const(0)) asm(dma_type_val, dram_addr, spad_addr) if offset_sym is not None: @@ -195,7 +253,7 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, # sram offset is linear (row-major stride 1) -> last index only; others 0. sidx = [_const(0)] * (len(tile_shape) - 1) + [sram_idx] _emit_asm(sram, sidx, dram_idx, vlane_axis, - tile_shape, dram_stride, tile_stride, subtile) + tile_shape, dram_stride, tile_stride, subtile, masked_pairs) op.erase() continue @@ -208,13 +266,16 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, dr = [dram_stride[i] for i in keep] tl = [tile_stride[i] for i in keep] st = [subtile[i] for i in keep] if subtile is not None else None + # remap each masked tile axis to its collapsed group index. + mp = [(next(gi for gi, g in enumerate(groups) if d in g), lo, hi) + for d, lo, hi in masked_pairs] new_vlane = next(gi for gi, g in enumerate(groups) if vlane_axis in g) with InsertionPoint(op): sram_c = Operation.create( "memref.collapse_shape", results=[collapsed_ty], operands=[sram], attributes={"reassociation": reassoc}).results[0] sidx = [_const(0)] * (len(target) - 1) + [sram_idx] - _emit_asm(sram_c, sidx, dram_idx, new_vlane, target, dr, tl, st) + _emit_asm(sram_c, sidx, dram_idx, new_vlane, target, dr, tl, st, mp) op.erase() continue @@ -226,6 +287,9 @@ def _emit_asm(sram_mem, sram_indices, dram_idx_val, vsa_val, desc_shape, dr = [dram_stride[d] for d in inner] tl = [tile_stride[d] for d in inner] st = [subtile[d] for d in inner] if subtile is not None else None + if any(d in peeled for d, _lo, _hi in masked_pairs): + raise NotImplementedError("masked-DMA clamp on a peeled (outer-loop) axis") + mp = [(inner.index(d), lo, hi) for d, lo, hi in masked_pairs if d in inner] if vlane_axis in inner: new_vlane = inner.index(vlane_axis) elif vlane_axis in peeled: @@ -285,7 +349,7 @@ def _phys(d): ).results[0] zero = _const(0) _emit_asm(sub, [zero, zero, zero, sram_off_val], dram_idx_val, new_vlane, - inner_shape, dr, tl, st) + inner_shape, dr, tl, st, mp) op.erase() From b9bafbd5526b6744f39f3ffa73edc1964992b7ab Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 13:16:25 +0900 Subject: [PATCH 11/13] [Frontend] Extend the masked-DMA clamp to conv (explicit extents + fine-grained) Conv on a non-dividing reduction (K = in_ch * kh * kw, e.g. 3*7*7=147 > systolic size) MAC-ed garbage from the tail of the K tile. Two fixes: - def_dma_op now takes the per-dim clamp extent from the kernel's loop_extents {loop_iv_name: loop_bound}, which each conv template declares once (the template's ranges/itervars equivalent). The stride-match extent inference is unreliable for conv because the gemm reads a REPACKED weight (O_C contiguous, stride 1) whose strides do not match the declared node_layout, so the inferred extent is wrong (e.g. O_C=64 read as 8); gemm (no loop_extents) keeps the inference. Only a dim whose index is a single iv is clamped -- an im2col spatial axis (o_h + k_h) has no single base and is left unclamped, which is correct because that input is pre-padded. All four conv templates (base, mt, sb, sbs) declare loop_extents; note the mt template fuses channel and kernel-width so its tile_k loops 0..I_C*K_W. - dma_fine_grained now PROPAGATES the masked (and indirect / accumulate) attrs + the low/high operands when it splits a tile into subtiles, and REMAPS each clamp per subtile: a subtile at offset iv*sub sees high = min(sub, high - iv*sub), low = max(0, low - iv*sub) (out-of-window subtiles get a negative high / low > sub and Spike skips them). Previously it dropped the clamp, so the conv clamp never reached Spike; gemm was unaffected because its DMAs do not go through this pass. Fixes conv2d on non-dividing in_channels / kernels (test_conv2d 13/13), including fused pointwise epilogues (conv + relu/scale/bias/sigmoid). No regression: matmul/addmm, bmm, elementwise, reduce, softmax, index_add all pass. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- .../mlir/mlir_conv_mt_template.py | 4 +++ .../mlir/mlir_conv_sb_template.py | 3 ++ .../mlir/mlir_conv_sbs_template.py | 3 ++ PyTorchSimFrontend/mlir/mlir_conv_template.py | 3 ++ PyTorchSimFrontend/mlir/mlir_template.py | 29 ++++++++++++----- .../mlir/passes/dma_fine_grained.py | 31 ++++++++++++++++++- 6 files changed, 65 insertions(+), 8 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py index 8b8288a8..7964da0f 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py @@ -143,6 +143,10 @@ def render(self, TOG_latency = O_W if TILE_M > O_W else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + # This template fuses channel and kernel-width into tile_k (loops 0..I_C*K_W). + kernel.loop_extents = {"tile_m": BATCH, "tile_n": O_C, "o_h": O_H, "o_w": O_W, + "k_h": K_H, "tile_k": I_C * K_W} # Prepare tile descriptors vlane_stride = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py index 92efff66..dfca23ec 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py @@ -144,6 +144,9 @@ def render(self, TOG_latency = O_W if TILE_M > O_W else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + kernel.loop_extents = {"tile_n": O_C, "o_h": O_H, "tile_m": O_W, + "k_h": K_H, "k_w": K_W, "tile_k": I_C} # Prepare tile descriptors vlane_stride = 1 vlane_split_axis = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py index dfd418d9..f1a42964 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py @@ -144,6 +144,9 @@ def render(self, TOG_latency = O_W if TILE_M > O_W else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + kernel.loop_extents = {"tile_n": O_C, "o_h": O_H, "tile_m": O_W, + "k_h": K_H, "k_w": K_W, "tile_k": I_C} # Prepare tile descriptors vlane_stride = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_conv_template.py b/PyTorchSimFrontend/mlir/mlir_conv_template.py index 178ba7c6..8bb64d48 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_template.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_template.py @@ -147,6 +147,9 @@ def render(self, TOG_latency = BATCH if TILE_M > BATCH else TILE_M TOG_latency = 8 if TOG_latency < 8 else TOG_latency kernel.loop_size = [TOG_latency, TILE_N, TILE_K] + # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op). + kernel.loop_extents = {"tile_m": BATCH, "tile_n": O_C, "o_h": O_H, "o_w": O_W, + "k_h": K_H, "k_w": K_W, "tile_k": I_C} # Prepare tile descriptors vlane_stride = 1 diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 40dd10bf..488de022 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -942,20 +942,35 @@ def generate_dma_code(): # masked-DMA clamp: per tile dim, clamp to the real DRAM extent so a tile # that pads past M/N/K (matmul) or a dim is filled with 0 instead of MAC-ing - # garbage. index_list[d] = iv * stride -> base iv; the extent is the DRAM dim - # matched by that tile dim's stride (tile order != layout order for conv). - # _emit_clamp skips dividing dims (no-op). + # garbage. index_list[d] = iv * stride -> base iv; the extent comes from the + # kernel's loop_extents {iv_name: loop_bound} (the template's ranges/itervars + # equivalent -- needed for conv, whose repacked-weight strides do not match + # node_layout), else it is matched by that tile dim's stride (gemm). + # _emit_clamp skips dividing dims (no-op). Only a dim whose index is a single + # iv is clampable: the clamp base is that iv, and high/low are affine in it. + # A dim indexed by a SUM of ivs (e.g. an im2col spatial axis o_h + k_h) has + # no single base, so it is left unclamped -- correct here because such inputs + # are pre-padded (materialized), not masked. + # FIXME: loop_extents is declared by hand in each conv template (verbose, + # easy to drift from the affine.for bounds). Fold it into the template + # overhaul -- e.g. record each loop's (iv -> bound) as the affine.for is + # emitted, so def_dma_op derives extents the way pointwise uses ranges/itervars. + _loop_ext = getattr(self, "loop_extents", None) _dram_size, _dram_str = node_layout.size, node_layout.stride _tsize = tile_desc.get_tile_size() _axes = [] for _d, _idx in enumerate(index_list): _syms = list(_idx.free_symbols) - if not _syms or _d >= len(_tsize): + if len(_syms) != 1 or _d >= len(_tsize): continue - _ext = next((int(_dram_size[j]) for j, st in enumerate(_dram_str) - if _dram_size[j].is_number and int(st) == int(_dram_stride[_d])), None) + _iv = str(_syms[0]) + if _loop_ext: + _ext = _loop_ext.get(_iv) # explicit: clamp only known ivs + else: + _ext = next((int(_dram_size[j]) for j, st in enumerate(_dram_str) + if _dram_size[j].is_number and int(st) == int(_dram_stride[_d])), None) if _ext is not None: - _axes.append((_d, str(_syms[0]), 0, _ext, int(_tsize[_d]))) + _axes.append((_d, _iv, 0, int(_ext), int(_tsize[_d]))) masked_bounds = self._emit_clamp(_axes, local_code) code = self.emit_transfer(dma_type, vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, diff --git a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py index 22eb999e..b9b482d6 100644 --- a/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py +++ b/PyTorchSimFrontend/mlir/passes/dma_fine_grained.py @@ -85,6 +85,14 @@ def __init__(self, op): self.tag_idx = operands[5] self.num_elements = operands[6] # = dma_type const operand self.num_elements_per_stride = operands[7] # = vlane_stride (vst) + # trailing operands: [offset (indirect)] then (low, high) per masked axis. + _extra = operands[8:] + self.offset = None + if "indirect" in op.attributes: + self.offset = _extra[0] + _extra = _extra[1:] + self.masked_axes = attr_i64_array(op, "masked_axes", default=[]) + self.masked_ops = _extra # low0, high0, low1, high1, ... self.sram_rank = len(ir.MemRefType(self.sram.type).shape) # Direction: MVIN reads dram -> sram; MVOUT writes sram -> dram. @@ -221,7 +229,8 @@ def _dma_attrs(dma): attrs = {} op = dma.op for k in ("dma_kind", "vlane_split_axis", "dram_stride", "tile_stride", - "subtile_size", "padding"): + "subtile_size", "padding", "masked_axes", "masked_fill", "indirect", + "offset_stride", "accumulate", "acc_float"): if k in op.attributes: attrs[k] = op.attributes[k] attrs["async"] = ir.BoolAttr.get(dma.is_async()) @@ -229,6 +238,18 @@ def _dma_attrs(dma): return attrs +def _remap_bound(bound, iv, sub, is_high, ip): + """The masked low/high are full-tile-local; shift to this subtile: subtile position + p maps to full-tile p + iv*sub, so subtile-local high = min(sub, high - iv*sub) and + low = max(0, low - iv*sub). Out-of-window (neg high / low>sub) -> Spike skips all.""" + from mlir.dialects import affine + d0, d1 = ir.AffineDimExpr.get(0), ir.AffineDimExpr.get(1) + edge = ir.AffineConstantExpr.get(sub if is_high else 0) + m = ir.AffineMap.get(2, 0, [edge, d0 - d1 * sub]) + op = affine.AffineMinOp if is_high else affine.AffineMaxOp + return op(m, [bound, iv], ip=ip).result + + def _emit_dma(dma, ivs, vectorlane, ip): """Emit one fine-grained togsim.transfer at `ip`, indexed by `ivs` (the fused induction vars for this DMA's dims, in dim order).""" @@ -243,6 +264,14 @@ def _emit_dma(dma, ivs, vectorlane, ip): operands = [dma.dram, dram_idx, dma.sram, sram_off, dma.tag, tag_idx, dma.num_elements, dma.num_elements_per_stride] + if dma.offset is not None: + operands.append(dma.offset) + # masked low/high are full-tile-local -> remap each per THIS subtile's offset. + sub = dma.subtile_size() + for i, axis in enumerate(dma.masked_axes): + low, high = dma.masked_ops[2 * i], dma.masked_ops[2 * i + 1] + operands.append(_remap_bound(low, ivs[axis], sub[axis], False, ip)) + operands.append(_remap_bound(high, ivs[axis], sub[axis], True, ip)) ir.Operation.create("togsim.transfer", results=[], operands=operands, attributes=_dma_attrs(dma), ip=ip) From 6b2b492f5579dc7af5247720eade9d3b4af7f0c7 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 14:21:28 +0900 Subject: [PATCH 12/13] [Frontend] Clamp the fused-epilogue output store on non-dividing shapes A fused template epilogue (matmul/conv + relu/add/gelu/...) stores its result via store_epilogue, which -- unlike the plain store() -- emitted no masked-DMA clamp. On a non-dividing output the store then wrote the systolic pad region past the real extent, corrupting the result (fused matmul+relu at 10x10x10 was off by ~7). store_epilogue now clamps each output tile dim to [0, extent): the extent of every loop iv is ranges[k], and the tile dims are indexed by dim_aliasing (tile-dim order), so a dim whose iv extent does not divide the tile is clamped. Keyed by iv name, it is naming-agnostic -- it covers both the index-N matmul epilogue and the c0/tile_n/... conv epilogue -- and _emit_clamp skips dividing dims, so a dividing output is a no-op. Fixes fused matmul/conv epilogues on non-dividing shapes (matmul+relu/add/gelu, conv+relu/scale/bias/sigmoid, non-dividing K and non-dividing output). No regression: test_matmul, matmul_activation, addmm_residual, matmul_scalar all pass. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- PyTorchSimFrontend/mlir/mlir_template.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 488de022..462dc93e 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -1109,9 +1109,26 @@ def store_epilogue(self, name: str, index: sympy.Expr, value, *args, **kwargs): with self.override_buffer_cse(buffer=self.stores): ops._store(value, sram_var, compute_index_var, tile_shape, buffer_name=buffer_name) - # Generate DMA instruction + # Generate DMA instruction. Clamp the output tail, else a non-dividing epilogue + # store writes the systolic pad region past the real output extent (fused matmul + + # relu/add on a non-dividing shape). The extent of each loop iv is ranges[k]; the + # tile dims are indexed by dim_aliasing (tile-dim order), so a dim whose iv extent + # does not divide the tile is clamped to [0, extent). Naming-agnostic (works for the + # index-N matmul epilogue and the c0/tile_n/... conv epilogue alike); _emit_clamp + # skips dividing dims, so a dividing output is a no-op. + iv_extent = {} + for _v, _r in zip(self.itervars, self.ranges): + try: + iv_extent[str(_v)] = int(_r) + except (TypeError, ValueError): + pass + _tile = self.kernel_group.tile_desc.get_tile_size() + _axes = [(d, iv, 0, iv_extent[iv], int(_tile[d])) + for d, iv in enumerate(self.dim_aliasing.values()) + if d < len(_tile) and iv in iv_extent] + masked_bounds = self._emit_clamp(_axes, self.dma_stores) code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, dram_stride, tile_stride, 0) + dram_shape, tile_shape, dram_stride, tile_stride, 0, masked_bounds=masked_bounds) self.dma_stores.writeline(DeferredLine(name, code)) def reduction_epilogue(self, dtype, src_dtype, reduction_type, value): From 5c90140856748cc9b776ec8ddbf14f8395c26359 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 7 Jul 2026 11:59:43 +0900 Subject: [PATCH 13/13] [Test] Golden unit test for the lower_transfer_to_gemmini pass Hand-write a masked MVOUT togsim.transfer, run lower_transfer_to_gemmini in-process (no torch, no spike; skipped without the MLIR bindings), and compare to an embedded golden -- the exact lowered IR (MLIR SSA numbering is deterministic for a fixed input): the i32-view descriptor global (dim_size / config_type / masked flag packed), the runtime dim_low/dim_high stores at i32 idx 6/10 for masked_axes=[2], and the CONFIG_DESC + MVOUT asm. A second case adds subtile_size=[1,1,4,8] and shows the descriptor dim_size become the subtile (config) shape. Timing mode erases the transfer. Under tests/lowering/ so further per-pass goldens have a home. Regenerate the golden and review the diff on an intended lowering change. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN --- tests/lowering/__init__.py | 0 .../test_lower_transfer_to_gemmini.py | 141 ++++++++++++++++++ 2 files changed, 141 insertions(+) create mode 100644 tests/lowering/__init__.py create mode 100644 tests/lowering/test_lower_transfer_to_gemmini.py diff --git a/tests/lowering/__init__.py b/tests/lowering/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/lowering/test_lower_transfer_to_gemmini.py b/tests/lowering/test_lower_transfer_to_gemmini.py new file mode 100644 index 00000000..6b7bf1d5 --- /dev/null +++ b/tests/lowering/test_lower_transfer_to_gemmini.py @@ -0,0 +1,141 @@ +"""Unit tests for the lower_transfer_to_gemmini pass (no torch, no Spike). + +Hand-write a `togsim.transfer`, run the pass in-process, and assert the lowered IR +matches an embedded golden -- the exact expected output (MLIR SSA numbering is +deterministic for a fixed input). If a lowering change is intended, regenerate the +golden and review the diff. Skipped without the MLIR bindings. +""" +import importlib.util +import textwrap + +import pytest + +_MLIR = importlib.util.find_spec("mlir") is not None +pytestmark = pytest.mark.skipif(not _MLIR, reason="MLIR Python bindings not installed") + + +def _lower_transfer(ir_text, timing=False): + from mlir.ir import Context, Module, Location + from PyTorchSimFrontend.mlir.passes import lower_transfer_to_gemmini as L + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(ir_text) + L.run(m, timing=timing) + return str(m).strip() + + +# INPUT: one masked MVOUT (store an SRAM tile back to DRAM) of a [1,2,8,8] tile. +# The togsim.transfer operands are positional (see emit_transfer): +# %arg1 DRAM base buffer %c0 dram_idx (element offset into it) +# %s SRAM tile (the [1,2,8,8] spad global) %c0 sram_idx +# %alloc async tag buffer %c0 tag_idx +# %c3 dma_type (3 = MVOUT) %c1 vlane_stride +# %c0, %c7 the masked clamp (low, high) -- ONE (low, high) pair per masked axis. +# Attrs: dram_stride/tile_stride = row-major [1,2,8,8] strides; vlane_split_axis=3; +# masked_axes=[2] clamps tile axis 2 to [low, high) = [0, 7); masked_fill=0. So on +# store, axis-2 positions >= 7 (the ragged tail) are skipped instead of written. +# masked_axes is a SPARSE overlay, NOT one pair per rank: the descriptor always carries +# all 4 axes' dim_low/dim_high, defaulting to [0, dim_size) (= no clamp). Only axes that +# are actually ragged are listed (the frontend skips axes that divide the tile evenly), +# so a single-axis clamp on a 4D tile is the normal case, not a rank mismatch. +_MASKED_MVOUT = """ +module { + memref.global @buf1_spad : memref<1x2x8x8xf32, 1> + func.func @k(%arg1: memref<2048xf32>) { + %c0 = arith.constant 0 : index + %c7 = arith.constant 7 : index + %alloc = memref.alloc() : memref<1xi32> + %s = memref.get_global @buf1_spad : memref<1x2x8x8xf32, 1> + %c3 = arith.constant 3 : index + %c1 = arith.constant 1 : index + "togsim.transfer"(%arg1, %c0, %s, %c0, %alloc, %c0, %c3, %c1, %c0, %c7) {dma_kind="MVOUT", dram_stride=[128,64,8,1], tile_stride=[128,64,8,1], vlane_split_axis=3:i64, masked_axes=[2], masked_fill=0:i64} : (memref<2048xf32>, index, memref<1x2x8x8xf32,1>, index, memref<1xi32>, index, index, index, index, index) -> () + return + } +}""" + +# OUTPUT (golden): the transfer lowers to a DMA descriptor + two RISC-V custom insns. +# @dma_desc_0 is the descriptor as 36 i32 (= 18 packed i64 slots); the dense init reads: +# words[0:4] = [1,2,8,8] dim_size (the tile / config shape) +# words[4:8] = [0,0,0,0] dim_low (default; runtime-overwritten below) +# words[8:12] = [1,2,8,8] dim_high (defaults to dim_size; runtime-overwritten) +# words[12:20] = dram_stride 128,64,8,1 words[20:28] = tile_stride 128,64,8,1 (i64 pairs) +# words[28:30] = 65540,131843 = the packed flags slot: elem_bytes 4 | vlane_stride 1 | +# vlane_split_axis 3 | config_type 3 (MVOUT) | flags 2 (masked) +# words[30:36] = 0 (masked_fill / indirect, unused here) +# Then the masked clamp is written at RUN TIME into the descriptor: low=%c0 -> i32 idx 6 +# (dim_low base 4 + axis 2), high=%c7 -> i32 idx 10 (dim_high base 8 + axis 2). The rest +# is address arithmetic building the DRAM byte addr (%3) and SRAM byte addr (%12), then two +# `.insn r CUSTOM_1` ops: func7=7 = CONFIG_DESC (hands the descriptor pointer to the DMA +# engine), func7=3 = MVOUT (hands it the DRAM + SRAM addresses -> issues the store). +_MASKED_MVOUT_LOWERED = textwrap.dedent("""\ + module { + memref.global "private" @dma_desc_0 : memref<36xi32> = dense<[1, 2, 8, 8, 0, 0, 0, 0, 1, 2, 8, 8, 128, 0, 64, 0, 8, 0, 1, 0, 128, 0, 64, 0, 8, 0, 1, 0, 65540, 131843, 0, 0, 0, 0, 0, 0]> + memref.global @buf1_spad : memref<1x2x8x8xf32, 1> + func.func @k(%arg0: memref<2048xf32>) { + %c0 = arith.constant 0 : index + %c7 = arith.constant 7 : index + %alloc = memref.alloc() : memref<1xi32> + %0 = memref.get_global @buf1_spad : memref<1x2x8x8xf32, 1> + %c3 = arith.constant 3 : index + %c1 = arith.constant 1 : index + %c0_0 = arith.constant 0 : index + %intptr = memref.extract_aligned_pointer_as_index %arg0 : memref<2048xf32> -> index + %c4 = arith.constant 4 : index + %1 = arith.muli %c0, %c4 : index + %2 = arith.addi %intptr, %1 : index + %3 = arith.index_cast %2 : index to i64 + %intptr_1 = memref.extract_aligned_pointer_as_index %0 : memref<1x2x8x8xf32, 1> -> index + %c128 = arith.constant 128 : index + %4 = arith.muli %c0_0, %c128 : index + %c64 = arith.constant 64 : index + %5 = arith.muli %c0_0, %c64 : index + %6 = arith.addi %4, %5 : index + %c8 = arith.constant 8 : index + %7 = arith.muli %c0_0, %c8 : index + %8 = arith.addi %6, %7 : index + %9 = arith.addi %8, %c0 : index + %c4_2 = arith.constant 4 : index + %10 = arith.muli %9, %c4_2 : index + %11 = arith.addi %intptr_1, %10 : index + %12 = arith.index_cast %11 : index to i64 + %13 = memref.get_global @dma_desc_0 : memref<36xi32> + %14 = arith.index_cast %c0 : index to i32 + %c6 = arith.constant 6 : index + memref.store %14, %13[%c6] : memref<36xi32> + %15 = arith.index_cast %c7 : index to i32 + %c10 = arith.constant 10 : index + memref.store %15, %13[%c10] : memref<36xi32> + %intptr_3 = memref.extract_aligned_pointer_as_index %13 : memref<36xi32> -> index + %16 = arith.index_cast %intptr_3 : index to i64 + %c0_i64 = arith.constant 0 : i64 + llvm.inline_asm has_side_effects asm_dialect = att ".insn r CUSTOM_1, 0x3, 7, x0, $0, $1", "r,r,~{dirflag},~{fpsr},~{flags}" %16, %c0_i64 : (i64, i64) -> () + llvm.inline_asm has_side_effects asm_dialect = att ".insn r CUSTOM_1, 0x3, 3, x0, $0, $1", "r,r,~{dirflag},~{fpsr},~{flags}" %3, %12 : (i64, i64) -> () + return + } + }""") + + +def test_masked_transfer_lowering_golden(): + assert _lower_transfer(_MASKED_MVOUT) == _MASKED_MVOUT_LOWERED + + +# Same transfer with subtile_size=[1,1,4,8] (the H axis is subtiled 8 -> 4). The +# descriptor's dim_size is the SUBTILE (config) shape, not the full tile; the masked +# axis + clamp stores are unchanged (a 4D subtile -> expand 0 -> same idx 6/10). +_MASKED_MVOUT_SUBTILE = _MASKED_MVOUT.replace( + 'vlane_split_axis=3:i64,', 'vlane_split_axis=3:i64, subtile_size=[1,1,4,8],') + +_MASKED_MVOUT_SUBTILE_LOWERED = _MASKED_MVOUT_LOWERED.replace( + "dense<[1, 2, 8, 8, 0, 0, 0, 0, 1, 2, 8, 8,", + "dense<[1, 1, 4, 8, 0, 0, 0, 0, 1, 1, 4, 8,") + + +def test_masked_transfer_subtile_lowering_golden(): + assert _lower_transfer(_MASKED_MVOUT_SUBTILE) == _MASKED_MVOUT_SUBTILE_LOWERED + + +def test_timing_mode_erases_transfer(): + # In timing mode the TOG carries DMA timing -> the transfer + descriptor are erased. + out = _lower_transfer(_MASKED_MVOUT, timing=True) + assert "togsim.transfer" not in out and "dma_desc" not in out