diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index b163ad1a..19ae3af5 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -322,6 +322,10 @@ def __init__(self, kernel_group, reason=None): self.spad_buffer_dict = dict() self.base_vector_initialized = False self.loop_size = None + # Set by get_dma_info when a DMA access cannot fit one <=4D Gemmini + # descriptor; load()/store() then emit a togsim.transfer for the + # decompose pass to peel into a loop of <=4D dma_start. + self._dma_needs_transfer = False def reset(self, reason): save = self.exit_stack, self._nested_context_depth @@ -537,9 +541,14 @@ def load(self, name: str, index: sympy.Expr): compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) # MVIN Encoding - attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, int(padding)) - code = self.get_dma_code("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, attribute) + if self._dma_needs_transfer: + self._dma_needs_transfer = False + 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)) + else: + attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, int(padding)) + code = self.get_dma_code("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, + dram_shape, tile_shape, attribute) self.cse.generate(dma_buffer, code, assignment = False) # FIXME: assignment = False does not support caching if not comptute_depedency: @@ -608,9 +617,14 @@ 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 - attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, 0) - code = self.get_dma_code("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, attribute) + if self._dma_needs_transfer: + self._dma_needs_transfer = False + 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) + else: + attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, 0) + code = self.get_dma_code("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, + dram_shape, tile_shape, attribute) self.dma_stores.writeline(common.DeferredLine(name, code)) def reduction(self, dtype, src_dtype, reduction_type, value): @@ -1243,7 +1257,13 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride else: - raise NotImplementedError("Currently not implemented... ;)") + # >4D access: one Gemmini DMA descriptor (<=4D) cannot represent this. + # Build the full N-D tile and flag it for togsim.transfer; the decompose + # pass peels the excess dims into a loop of <=4D memref.dma_start. + local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) + local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis + local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride + self._dma_needs_transfer = True if len(implicit_local_dims)!=0 and len(local_dims) != len(implicit_local_dims) and self.is_modular_indexing(index): for axis_constraints in self.kernel_group.tile_desc.implicit_dim_size.values(): @@ -1426,6 +1446,53 @@ def get_dma_code(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype return f"memref.dma_start {src_operand}, {dst_operand}, %{dma_type}, {tag_var}, {dma_attribute} : {src_shape}, {dst_shape}, {tag_shape} {attribute}" + 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): + """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 + (dma_type / vlane_split_axis / vlane_stride), so the decompose pass + (passes/decompose_transfer.py) is purely mechanical: it peels the excess + dims into a loop of <=4D memref.dma_start, reusing these operands. + + The operand prep mirrors get_dma_code (dma_type enum via the read/write + cache+counter, vlane consts via CSE) so the transfer is self-contained; + togsim is an unregistered dialect -> generic form. + """ + dma_key = (vlane_split_axis, vlane_stride, mlir_dtype) + if dma_type_name == "MVIN" and dma_key in self.dma_read_cache: + dma_type, vsa, vst = self.dma_read_cache[dma_key] + elif dma_type_name == "MVOUT" and dma_key in self.dma_write_cache: + dma_type, vsa, vst = self.dma_write_cache[dma_key] + else: + vsa = self.get_const_cse(vlane_split_axis) + vst = self.get_const_cse(vlane_stride) + if dma_type_name == "MVIN": + dma_type = self.get_const_cse(DMA_TYPE[f"{dma_type_name}{self.dma_read_counter}"]) + self.dma_read_counter += 1 + self.dma_read_cache[dma_key] = [dma_type, vsa, vst] + else: + dma_type = self.get_const_cse(DMA_TYPE[f"{dma_type_name}{self.dma_write_counter}"]) + self.dma_write_cache[dma_key] = [dma_type, vsa, vst] + tag = self.get_tag_cse() + zero_cse = self.get_const_cse(0) + # vlane_split_axis is carried as a VALUE attr (not an SSA operand) because the + # decompose pass must remap it: collapsing unit tile dims renumbers the axes, + # so the descriptor's vlane axis index changes and the pass rebuilds the const. + attrs = ( + f'dma_kind = "{dma_type_name}", ' + f'vlane_split_axis = {int(vlane_split_axis)} : i64, ' + f'dram_stride = {dram_stride}, tile_stride = {tile_stride}, ' + f'padding = {int(padding)} : 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) -> ()' + ) + 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] mlir_dtype = mlir_common.DTYPE_TO_MLIR[dtype] diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 734ca967..f73d818e 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -472,29 +472,25 @@ def apply_constraints(self, constraints, ranges): @staticmethod def init_tile_size(ranges, vlane_stride, vector_lane): + # Logical tile init for ANY rank. Only the innermost dims carry the + # vectorized tile; all further-outer dims stay 1. The physical Gemmini DMA + # descriptor is <=4D -- a higher-rank logical tile is mapped onto <=4D + # descriptors by togsim.transfer + the decompose pass (logical/physical + # tile split), so no rank cap here. nr_dim = len(ranges) + if nr_dim == 0: # scalar + return [1] tile_size = [1] * nr_dim - if len(tile_size) == 2: + if nr_dim == 1: + tile_size[0] = 1 if ranges[0] == 1 else 2 * vlane_stride * vector_lane + elif nr_dim == 2: tile_size[-1] = vlane_stride * vector_lane tile_size[-2] = 2 * vector_lane - elif len(tile_size) == 0: # Scalar - tile_size = [1] - ranges = [1] - elif len(tile_size) == 1 and ranges[0]==1: - tile_size[0] = 1 - elif len(tile_size) == 1: - tile_size[0] = 2 * vlane_stride * vector_lane - elif len(tile_size) == 3: + else: # 3D and up (general) tile_size[-1] = vector_lane tile_size[-2] = 4 * vector_lane tile_size[-3] = 2 - elif len(tile_size) == 4: - tile_size[-1] = vector_lane - tile_size[-2] = 4 * vector_lane - tile_size[-3] = 2 - tile_size[-4] = 1 - else: - raise NotImplementedError("dummy tile size fail!") + # tile_size[:-3] stay 1 (subsumes the old 4D [-4]=1 and any higher rank) return tile_size @staticmethod diff --git a/PyTorchSimFrontend/mlir/passes/__init__.py b/PyTorchSimFrontend/mlir/passes/__init__.py index 1ab47ee8..be5533ac 100644 --- a/PyTorchSimFrontend/mlir/passes/__init__.py +++ b/PyTorchSimFrontend/mlir/passes/__init__.py @@ -10,10 +10,14 @@ run(module) (mutates the Module in place), and append it to PASSES below. """ from . import lower_vlane_idx +from . import decompose_transfer from .lower_to_llvm import run_standard_lowering # noqa: F401 (re-exported) # Ordered passes applied to each kernel .mlir before mlir-opt. +# decompose_transfer first: it lowers togsim.transfer -> memref.dma_start, which +# downstream passes (and the gemmini lowering) expect. PASSES = [ + decompose_transfer, lower_vlane_idx, ] diff --git a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py new file mode 100644 index 00000000..f606b95e --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py @@ -0,0 +1,200 @@ +"""Python out-of-line MLIR pass: decompose togsim.transfer -> <=4D memref.dma_start. + +A togsim.transfer carries a per-axis affine DMA whose descriptor rank may exceed +the 4D Gemmini limit. This pass is a **pure mechanical rank peel** of that +already-affine access (see docs/dma-transfer-lowering.md, "aligned-only peel"): + + - drop unit (extent-1) tile dims: they contribute no descriptor axis; + - if the remaining (effective) rank <= 4 -> emit one customized + memref.dma_start, reusing the transfer's operands (fast path); + - if effective rank > 4 -> peel the outer dims into a loop, adjusting the + base index by stride*iv per iteration, inner descriptor <=4D. + +It does NO floor/mod linearization (aligned split happens upstream at the +scheduling layer) and NO relayout (misaligned access is copy-inserted at the +graph level). A transfer whose access is not per-axis affine is a contract +violation -- but by construction codegen only emits affine transfers. + +togsim.transfer operands (see emit_transfer): + (dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_split_axis, vlane_stride) +attrs: dma_kind ("MVIN"/"MVOUT"), dram_stride[], tile_stride[], padding. + +memref.dma_start (customized) operands: + src[idx], dst[idx], dma_type, tag[idx], vlane_split_axis, vlane_stride + : src_memref, dst_memref, memref<1xi32> {dram_stride, sram_stride, padding} + +Pass interface (passes/__init__.py): MARKERS + run(module). +""" + +OP_NAME = "togsim.transfer" +MARKERS = (OP_NAME,) + + +def _iter_ops(block): + for op in list(block.operations): + yield op + for region in op.operation.regions: + for b in region.blocks: + yield from _iter_ops(b) + + +def _int_array(attr): + from mlir.ir import ArrayAttr, IntegerAttr + return [IntegerAttr(a).value for a in ArrayAttr(attr)] + + +def _squeeze_reassociation(shape): + """Group source dims so each group's product is one effective (non-unit) dim; + unit dims attach to a neighbor. Returns (groups, target_shape).""" + groups, cur = [], [] + for i, e in enumerate(shape): + cur.append(i) + if e > 1: + groups.append(cur) + cur = [] + if cur: # trailing unit dims + if groups: + groups[-1] += cur + else: + groups.append(cur) # all-ones -> single dim of size 1 + import math + target = [math.prod(shape[d] for d in g) for g in groups] + return groups, target + + +def run(module): + """Lower every togsim.transfer in `module`, in place. Context must be active.""" + import itertools + from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, + IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, + StridedLayoutAttr) + i64 = IntegerType.get_signless(64) + idx_ty = IndexType.get() + + targets = [] + for region in module.operation.regions: + for b in region.blocks: + for op in _iter_ops(b): + if op.operation.name == OP_NAME: + targets.append(op.operation) + + for op in targets: + dram, dram_idx, sram, sram_idx, tag, dma_type, vst = op.operands + 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"]) + tile_stride = _int_array(op.attributes["tile_stride"]) + padding = op.attributes["padding"] + + sram_ty = MemRefType(sram.type) + elem, space = sram_ty.element_type, sram_ty.memory_space + tile_shape = list(sram_ty.shape) + # effective (non-unit) dims carry the descriptor; unit dims drop out. + eff = [i for i, e in enumerate(tile_shape) if e > 1] + + def _const(v): + return Operation.create( + "arith.constant", results=[idx_ty], + attributes={"value": IntegerAttr.get(idx_ty, v)}).results[0] + + def _emit(sram_mem, sram_indices, dram_idx_val, vsa_val, dr_attr, tl_attr): + vsa = _const(vsa_val) + if kind == "MVIN": + operands = [dram, dram_idx_val, sram_mem, *sram_indices, + dma_type, tag, sram_idx, vsa, vst] + else: + operands = [sram_mem, *sram_indices, dram, dram_idx_val, + dma_type, tag, sram_idx, vsa, vst] + Operation.create( + "memref.dma_start", results=[], operands=operands, + attributes={"dram_stride": dr_attr, "sram_stride": tl_attr, + "padding": padding}) + + if len(eff) <= 4: + # Fast path: drop unit dims so the descriptor reaches <=4D. The customized + # dma_start convention requires SRAM rank == #indices == len(sram_stride), + # so collapse the unit tile dims away. DRAM stays flat rank-1 (its N-D + # structure is in dram_stride). + 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 = [g[-1] for g in groups] # the non-unit dim in each group + dr_attr = ArrayAttr.get([IntegerAttr.get(i64, dram_stride[i]) for i in keep]) + tl_attr = ArrayAttr.get([IntegerAttr.get(i64, tile_stride[i]) for i in keep]) + # Remap vlane axis to the collapsed-dim index (the group containing it). + 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] + _emit(sram_c, [sram_idx] * len(target), dram_idx, new_vlane, + dr_attr, tl_attr) + op.erase() + continue + + # Peel path: >4 effective dims. Keep the inner 4 as the <=4D descriptor and + # peel the outer (len-4) effective dims into a fully-unrolled set of slices + # (one descriptor per outer index combo; base advances by stride*idx). The + # SRAM slice is a rank-reduced memref.subview at the slice offset; DRAM base + # is dram_idx + constant. Unrolling (vs scf.for) keeps the slice offsets + # static so no per-iteration index arithmetic on the SRAM side is needed. + # + # NOTE: currently unreachable -- init_tile_size caps non-unit tile dims at 3, + # so eff <= 3 in practice. Implemented for completeness / future tilings and + # validated only in isolation (passes/decompose_transfer.py CLI / lower_text). + 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_attr = ArrayAttr.get([IntegerAttr.get(i64, dram_stride[d]) for d in inner]) + tl_attr = ArrayAttr.get([IntegerAttr.get(i64, tile_stride[d]) for d in inner]) + # the vlane axis must survive into the inner descriptor (it is the lane dim). + new_vlane = inner.index(vlane_axis) if vlane_axis in inner else 0 + for combo in itertools.product(*[range(tile_shape[d]) for d in peeled]): + static_offsets = [0] * ndim + static_sizes = [1] * ndim + for k, d in enumerate(peeled): + static_offsets[d] = combo[k] + for d in inner: + static_sizes[d] = tile_shape[d] + sram_off = sum(combo[k] * tile_stride[peeled[k]] for k in range(len(peeled))) + dram_off = sum(combo[k] * dram_stride[peeled[k]] for k in range(len(peeled))) + res_ty = MemRefType.get( + inner_shape, elem, + layout=StridedLayoutAttr.get(sram_off, inner_strides), memory_space=space) + with InsertionPoint(op): + sub = Operation.create( + "memref.subview", results=[res_ty], operands=[sram], + attributes={"static_offsets": DenseI64ArrayAttr.get(static_offsets), + "static_sizes": DenseI64ArrayAttr.get(static_sizes), + "static_strides": DenseI64ArrayAttr.get([1] * ndim), + "operandSegmentSizes": DenseI64ArrayAttr.get([1, 0, 0, 0])} + ).results[0] + dram_idx_val = dram_idx if dram_off == 0 else Operation.create( + "arith.addi", results=[idx_ty], + operands=[dram_idx, _const(dram_off)]).results[0] + _emit(sub, [sram_idx] * 4, dram_idx_val, new_vlane, dr_attr, tl_attr) + op.erase() + + +def lower_text(text: str) -> str: + """Parse `text`, run this pass, return the printed module. CLI/testing helper.""" + 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) + return str(m) + + +if __name__ == "__main__": + import sys + out = lower_text(open(sys.argv[1]).read()) + if len(sys.argv) > 2: + open(sys.argv[2], "w").write(out) + else: + sys.stdout.write(out) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py index 19644b28..379fbac9 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py @@ -19,6 +19,7 @@ "convert-linalg-to-loops," "convert-vector-to-scf{full-unroll=true}," "lower-affine," + "expand-strided-metadata," # decompose memref.collapse_shape/subview before LLVM "finalize-memref-to-llvm," "func.func(lower-vector-multi-reduction)," "convert-vector-to-llvm," diff --git a/docs/dma-transfer-lowering.md b/docs/dma-transfer-lowering.md index c046c46a..c4d6f29c 100644 --- a/docs/dma-transfer-lowering.md +++ b/docs/dma-transfer-lowering.md @@ -133,37 +133,55 @@ Design choices: lowers to subview+scf, not our descriptors -- so a custom op modeled on linalg's design, reusing AffineMap utilities. -### Decomposition pass (contract) +### Decomposition pass (contract): aligned-only mechanical peel + +> **Scope decision (narrowed).** This pass is a **pure mechanical rank peel** of an +> already-affine access. It does **not** linearize floor/mod and does **not** do +> relayout. Those two responsibilities moved upstream (see "Division of labor" +> below): aligned floor/mod is removed by **axis splitting at the Inductor +> scheduling layer** (`axis-split-scheduling.md`), and misaligned access is +> resolved by **graph-level copy insertion**. So every `togsim.transfer` that +> reaches this pass is guaranteed per-axis affine; the only thing left is that its +> rank may exceed the 4D Gemmini descriptor. The DMA descriptor is an **affine map of rank <= 4 with integer strides** -(`base + sum_i stride_i * idx_i`). Decide by **rank after linearization**, NOT by -the presence of floordiv/mod: - -1. **Linearize** `src_map`: rewrite each `floordiv c` / `mod c` on an iteration dim - into a split pair (`idx = outer*c + inner`), which is purely linear in the new - dims. (This is exactly what `apply_divisor("split")` already does.) Let `D` be - the resulting affine rank. -2. **`D <= 4`** -> emit **one** customized `memref.dma_start`; the split dims become - the descriptor's <=4D shape/strides. Identical to today's output (fast path). - floordiv/mod that still fits in <=4D after splitting stays here -- it is *not* a - peel trigger. -3. **`D > 4`** (not expressible as a single linear combination) -> express it as a - **combination of linear combinations**: peel `D - 4` dims into an outer - `affine.for`; each iteration computes a base with `affine.apply` (the peeled - dims' linear, incl. split-derived, contribution) and issues the inner <=4D - affine descriptor. SRAM offsets are computed symmetrically in the same loop. -4. If the estimated descriptor count is pathological -> fall back to **relayout**. - -Genuinely non-affine access (data-dependent / indirect / gather -- an index that -comes from a loaded value and cannot be linearized by splitting) is **out of scope** -for this pass; it stays on the indirect-indexing path (or a relayout). - -The decision point maps onto existing code: codegen already splits floordiv/mod via -`apply_divisor` and raises `NotImplementedError` at >4D (`get_dma_info`). That exact -site becomes "emit `togsim.transfer`" instead of dying, and the recompile/tile --forcing dance is unnecessary because the outer peel loop's `ceil` bound absorbs +(`base + sum_i stride_i * idx_i`). The pass sees affine input (rank `D`) and: + +1. **`D <= 4`** -> emit **one** customized `memref.dma_start`; the dims become the + descriptor's <=4D shape/strides. Identical to today's output (fast path). +2. **`D > 4`** -> peel `D - 4` dims into an outer `affine.for`; each iteration + computes a base with `affine.apply` (the peeled dims' linear contribution) and + issues the inner <=4D affine descriptor. SRAM offsets are computed symmetrically + in the same loop. + +That is the whole pass. There is **no linearization step** (upstream guarantees +affine) and **no relayout fallback** (upstream graph copy handles misalignment). + +**Fail loud, not silent.** If the pass encounters floor/mod that does not reduce to +per-axis affine (misaligned), or a genuinely non-affine / indirect / gather index, +that is a **contract violation** -- upstream did not normalize it. The pass +**asserts/errors** rather than silently inserting a relayout. A silent in-pass copy +would be a hidden performance cliff and would duplicate, at the wrong layer, a +global layout decision only the graph can make correctly. + +The decision point maps onto existing code: `get_dma_info` already raises at >4D. +That exact site becomes "emit `togsim.transfer`" (done, Phase 1), and this pass +consumes it. The recompile/tile-forcing dance is unnecessary because (a) aligned +floor/mod is gone before codegen and (b) the outer peel loop's `ceil` bound absorbs non-divisible remainders. +### Division of labor (the affine-only contract) + +| floor/mod source | handled by | cost | layer | +|---|---|---|---| +| aligned (single axis, divisor \| extent; group norm, broadcast) | axis split | free | Inductor scheduling | +| misaligned (uneven cat, non-factor reshape, multi-axis arg) | copy insertion | copy | FX graph | +| affine but rank > 4 (e.g. 5D permute) | mechanical peel | free | **this pass** | +| data-dependent / indirect / gather | indirect-indexing path | -- | out of scope | + +Only the third row is this pass. The first two produce the affine-only invariant +this pass relies on. + ### Relationship to memref-to-gemmini (ISA lowering) -- keep separate `memref.dma_start` is the boundary, not the endpoint. The layering is: @@ -198,7 +216,10 @@ Ramulator). Rules: peeled extents). 2. Keep the inner descriptor **as large and contiguous as possible** (maximize bytes per descriptor). -3. If even the best peel is pathological, fall back to **relayout**. + +(A pathological peel is not this pass's problem to fix: it means the operand's +layout is bad, which is a graph-level layout/copy decision, not an in-pass +relayout.) ### Placement: hybrid (least burden) @@ -207,7 +228,7 @@ fast); keep the C++ pass purely mechanical. | Step | Where | |---|---| -| peel-plan decision (which dims, count estimate, peel vs relayout) | Python | +| peel-plan decision (which dims to peel, count estimate) | Python | | encode plan as op attributes | Python -> MLIR | | emit `scf.for { customized dma_start }` per the plan | C++ pass | @@ -240,7 +261,8 @@ The cost model can migrate into C++ later if desired. that is actually broken (DMA decomposition) into a lowering pass, without the full linalg rewrite. - **Cost-aware, so modeled performance is protected.** Peel small/outer, keep inner - contiguous, relayout for pathological cases. + contiguous. Pathological layouts are fixed upstream (graph copy), not by an + in-pass relayout. ## Migration strategy @@ -250,12 +272,13 @@ The cost model can migrate into C++ later if desired. maps and vlane attributes it already computes. 3. Implement `decompose-transfer` with the fast path first (<=4D affine -> one `dma_start`), proving **bit-identical output** to today on a smoke test. -4. Add the peel path for floor/mod / >4D; validate end-to-end through all three +4. Add the **affine** peel path for >4D; validate end-to-end through all three simulators (the loop-of-descriptors must satisfy the TOG / Spike / gem5 - contract). -5. Add the relayout fallback gated by the cost estimate. -6. Remove the `get_dma_info` recompile branches once the pass covers their cases; - use the failure ledger + assert-only `TestLoopPadding` to confirm nothing + contract). Make the pass **assert** on any non-affine residue (contract guard). +5. Land the upstream producers of the affine-only invariant: aligned axis split at + scheduling (`axis-split-scheduling.md`) and misaligned graph copy insertion. +6. Remove the `get_dma_info` recompile branches once the pass + upstream cover their + cases; use the failure ledger + assert-only `TestLoopPadding` to confirm nothing regresses before deleting. ## Relationship to Plan A and Plan B @@ -275,12 +298,169 @@ The cost model can migrate into C++ later if desired. Python, pass is mechanical). - **TOG / Spike / gem5 contract on a loop of descriptors.** If TOG generation assumes "one DMA = one node," the loop form needs handling. Validate at step 4. -- **Cost model accuracy** for peel-vs-relayout; start with a simple - descriptor-count threshold and refine against measured cycles. -- **Dynamic shapes**: `iter_bounds` as SSA operands and symbolic-divisor floor/mod - (semi-affine) must be handled by the pass. -- **Relayout fallback** needs a scratch buffer and a copy kernel; account for its - memory and cycle cost in the decision. +- **Cost model accuracy** for the peel plan; start with a simple descriptor-count + threshold and refine against measured cycles. +- **Dynamic shapes**: `iter_bounds` as SSA operands; affine peel must handle + symbolic outer extents. (Symbolic-divisor floor/mod normalization is an upstream + concern, not this pass.) +- **Upstream completeness.** The pass's fail-loud contract is only safe if the + upstream producers (axis split + graph copy) actually normalize every misaligned + case. Until they do, the assert may fire on real models -- track which ops trip it + as the work-list for the upstream passes. - **Async / tag management across the peel loop**: double-buffering / compute overlap must survive decomposition (e.g. keep the inner large DMA async, sequence the outer peel). + +## Appendix: alignment theory (when floor/mod is statically decomposable) + +This section records the math that decides, for a given DMA access, whether the +non-affine `FloorDiv`/`ModularIndexing` terms can be peeled into a *static* loop +of affine descriptors (free) or require a data movement (relayout / copy). + +### Setup + +A Gemmini-style descriptor addresses an element as + + addr(idx) = base + Σ_k stride_k · idx_k (integer strides, rank <= 4) + +i.e. each loop index `idx_k` contributes a **constant** stride. A DMA is +statically decomposable iff every index term it reads has constant stride over +the rectangular tile domain. Inductor index expressions, after fusion/view, carry +`FloorDiv(x, y)` and `ModularIndexing(x, y, z)` of the *flattened* loop variable +`x`. The question is when those reduce to constant-stride axes. + +### Mixed-radix decomposition + +Write the flattened index `x` (extent `E`) in mixed radix. For a `ModularIndexing` +with inner period `y` and modulus `z`, decompose uniquely as + + x = o·(y·z) + m·y + r, with 0 <= r < y, 0 <= m < z, o >= 0 + +Then `FloorDiv(x, y) = o·z + m`, and `ModularIndexing(x, y, z) = m`. Each of +`o, m, r` is a separate **implicit axis** with a constant per-axis stride — +*provided the axis boundaries do not move across the tile*. That holds iff the +period divides the extent it partitions: + +- `ModularIndexing(x, y, z)` is a valid rectangular axis **iff y·z | E**. +- `FloorDiv(x, y)` is a valid rectangular axis **iff y | E**. + +**Aligned** = the divisor (and modular period `y·z`) divides the extent, so the +wrap point lands on a fixed axis boundary -> constant stride -> peelable for free. +**Misaligned** = the wrap point falls at a loop-value-dependent position inside the +descriptor (e.g. uneven `cat`, ragged split) -> the stride is not constant -> +**not** statically decomposable; only a relayout (physical copy) fixes it. + +### One loop axis -> several implicit axes (complex fusion) + +When fusion merges many dims into one flattened loop variable, a *single* loop +axis can expand into **several** implicit axes through nested floor/mod, e.g. + + x in [0, D0·D1·D2): + a = FloorDiv(x, D1·D2) # outer + b = ModularIndexing(x, D2, D1) # middle + c = ModularIndexing(x, 1, D2) # inner + +That is three implicit descriptor axes coming from one loop axis. This is the +general case the un-flatten must handle: it is **not** limited to splitting one +axis into two. Key consequences: + +1. **The loop's own factorization is always aligned.** When the implicit axes + come from re-reading the loop's *own* contiguous factorization (the common + fusion case -- Inductor flattens contiguous dims then a consumer reads them + back via floor/mod), every period divides by construction (`D1·D2 | D0·D1·D2`, + etc.). So these un-flatten splits are **free** -- they just add descriptor + axes, never a copy. +2. **Rank blows past 4 fast.** k implicit axes per loop axis, across multiple + operands, means the descriptor rank exceeds the 4D Gemmini limit very quickly. + This is exactly why `togsim.transfer` + the peel pass matters *more* under + complex fusion, independent of any misalignment. The >4D branch in + `get_dma_info` already routes these to `togsim.transfer`. +3. **Misalignment is still only from non-factor views.** An implicit axis is + misaligned only when its period does not divide the extent -- i.e. the view + does not factor along the loop's factorization (uneven `cat`, ragged split, + group sizes that don't divide the channel count). Those, and only those, need + relayout. + +### Case-handling summary + +| Source of floor/mod | Aligned? | Handling | Cost | +|------------------------------------------------|----------|-----------------------------------|------| +| Broadcast / dim-merge (`[N,1]->[N,M]`, `i//M`) | always | un-merge (split loop axis back) | free | +| Reshape along the loop's own factorization | yes (`y·z\|E`) | un-flatten split, then peel for rank | free | +| >4D logical tile from complex fusion | yes | `togsim.transfer` -> peel into <=4D loop | free (extra DMA nodes) | +| Uneven `cat`, ragged split, non-dividing group | no | graph copy insertion (relayout, upstream) | copy = TPU `concatenate` | + +The TPU/XLA model is the reference: express only aligned views as +descriptor/bitcast (free reshape); never put a misaligned access in the +descriptor -- insert a copy (relayout) instead. Plan A (graph-level +force-contiguous / pad-to-granule, like XLA copy-insertion) is the upstream lever +that *reduces how often* the misaligned branch fires, keeping codegen affine-only. + +## Implementation status (Phase 1: codegen emission) + +Landed on branch `dma-transfer/codegen` (worktree), emission only -- the +decompose pass is deferred until explicitly signalled. A >4D access now emits a +`togsim.transfer` instead of hard-failing; without the pass it does not yet run +end-to-end (expected). + +- **`mlir_common.py` `init_tile_size`** generalized to any rank. Logical tile is + separated from the physical (<=4D) descriptor: only the innermost dims carry the + vectorized tile, all further-outer dims stay 1, and there is no rank cap. The + `nr_dim >= 3` formula reproduces the old 3D/4D values exactly (the old `[-4]=1` + is subsumed by "outer dims stay 1"); scalar/1D/2D keep their special cases. This + removes the old `raise NotImplementedError("dummy tile size fail!")` that + conflated logical and physical tile rank. +- **`mlir_codegen_backend.py`**: + - `__init__` adds `self._dma_needs_transfer = False`. + - `get_dma_info` >4D `else` branch (was + `raise NotImplementedError("Currently not implemented... ;)")`) now builds the + full N-D tile (`set_tile_size`, vlane split/stride) and sets + `self._dma_needs_transfer = True`. + - `emit_transfer(...)` emits the generic-form `"togsim.transfer"(...)` op + carrying `dma_kind`, `vlane_split_axis`, `vlane_stride`, `dram_stride`, + `tile_stride`, `padding`, with operands `(dram, dram_idx, sram, 0, tag)`. + `togsim` is an unregistered dialect, hence generic form. + - `load()` (MVIN) and `store()` (MVOUT) check the flag: if set, reset it and + call `emit_transfer`; otherwise the existing `get_dma_code` path is unchanged. + So aligned <=4D DMAs are **bit-identical** to before; only >4D accesses change. + +Validated: the 5D permute smoke test (`x.permute(4,3,2,1,0).contiguous() + 1.0`) +now emits MVIN/MVOUT `togsim.transfer` with 5D `dram_stride [1,6,30,120,360]` and a +`memref<1x1x2x4x2xf32,1>` tile, instead of crashing in `init_tile_size` or the +`get_dma_info` >4D branch. + +### Phase 2: aligned-only peel pass (landed: unit-collapse path) + +`passes/decompose_transfer.py` (registered in `passes/__init__.py`, runs before +`lower_vlane_idx`) lowers each `togsim.transfer` to a customized `memref.dma_start`: + +- **Unit-dim collapse (done, validated).** Drop extent-1 tile dims so the + descriptor reaches <=4D. The SRAM (spad) memref is collapsed to the effective + rank via `memref.collapse_shape` (the customized `dma_start` convention requires + SRAM rank == #indices == len(sram_stride)); DRAM stays flat rank-1 with its N-D + structure in `dram_stride`. The `vlane_split_axis` is **remapped** from the + original tile-dim index to the collapsed-dim index and rematerialized as a const + (carried as a value attr precisely so the pass can remap it). +- Supporting changes: `emit_transfer` now carries the SSA operands a `dma_start` + needs (`dma_type`, `vlane_stride`) + the `vlane_split_axis` value attr, so the + pass is mechanical. `lower_to_llvm.py` gains `expand-strided-metadata` to lower + `collapse_shape`. + +Validated end-to-end (Gem5 + Spike + TOGSim, `allclose=True`) on the 5D permute +`x.permute(4,3,2,1,0).contiguous() + 1.0`; no regression on 2D/3D/elementwise. + +- **Genuine >4 effective rank (done, isolation-validated).** When >4 *non-unit* + dims survive, the pass keeps the inner 4 as the <=4D descriptor and peels the + outer dims by **full unrolling**: one descriptor per outer-index combo, the SRAM + slice a rank-reduced `memref.subview` at the static slice offset, the DRAM base + `dram_idx + constant`. Unrolling (vs `scf.for`) keeps slice offsets static, so no + per-iteration SRAM index arithmetic is needed. **Currently unreachable**: + `init_tile_size` caps non-unit tile dims at 3 (effective rank <= 3 in practice), + so this path is exercised only in isolation (`lower_text` / the module CLI), not + through the full pipeline. Implemented for completeness and future tilings. + +The input stays per-axis affine by upstream guarantee, so both paths are pure +mechanical peeling. A non-affine residue is a contract violation (aligned floor/mod +removal lives in `axis-split-scheduling.md`, misaligned relayout in graph copy +insertion -- see "Division of labor"); a genuinely non-affine / indirect index +would surface as a build failure here rather than being silently relaid out.