[Frontend] Represent kernel body as ordered steps; multi-step indirect

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

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):

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

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

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

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

PyTorchSimFrontend/mlir/mlir_template.py

@@ -940,9 +940,43 @@ 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 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 len(_syms) != 1 or _d >= len(_tsize):
+                        continue
+                    _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, _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,
                                           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()
 
@@ -1075,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):

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.