[Frontend] DMA transfer op + aligned-only decompose pass for >4D access

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]

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

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,
 ]
 

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)

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,"