[Frontend] Lower togsim.transfer directly to Gemmini (drop memref.dma_start)

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_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,15 +1503,31 @@ 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]
+        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}, %{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}'
             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()

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