Pointwise ops on the C++ trace pipeline

.github/workflows/pytorchsim_test.yml

@@ -57,6 +57,24 @@ jobs:
             -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \
             ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/elementwise/test_transcendental.py
 
+  test_pointwise:
+    name: Run test_pointwise.py
+    runs-on: ubuntu-latest
+    steps:
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Run test_pointwise.py
+        run: |
+          echo "Running test_pointwise.py"
+          docker run --rm \
+            -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \
+            ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/elementwise/test_pointwise.py
+
   test_activation:
     name: Run test_activation.py
     runs-on: ubuntu-latest

PyTorchSimFrontend/mlir/mlir_cat_template.py

@@ -209,6 +209,22 @@ def _compute_excluded_dims(self, tile_sizes: list) -> list:
             tile_sizes[idx] = 1
         return excluded
 
+    @staticmethod
+    def _largest_divisor_leq(extent, cap):
+        """Largest divisor of ``extent`` that does not exceed ``cap`` (>= 1).
+
+        The concat-dim copy loop steps by the per-input tile over ``[0, extent)``.
+        If the tile does not divide ``extent`` the final iteration is ragged and,
+        because the DMA carries no remainder mask, over-reads the input and
+        over-writes the output. Snapping the tile down to a divisor keeps every
+        iteration full-width and in-bounds.
+        """
+        cap = min(cap, extent)
+        for d in range(cap, 0, -1):
+            if extent % d == 0:
+                return d
+        return 1
+
     def _calculate_input_tile_sizes(self, kernel, input_sizes, tile_sizes, num_inputs, rank, precision_bytes):
         """Calculate tile sizes along the concat dimension for each input."""
         non_dim_tile_elements = math.prod(tile_sizes) if tile_sizes else 1
@@ -218,7 +234,9 @@ def _calculate_input_tile_sizes(self, kernel, input_sizes, tile_sizes, num_input
         input_tile_sizes_dim = []
         for i in range(num_inputs):
             if extra_concat > 0 and non_dim_tile_elements > 0:
-                tile_dim = min(input_sizes[i][self.dim], extra_concat)
+                # Snap to a divisor of the input's concat extent so the copy loop
+                # never emits a ragged (unmasked) tail tile.
+                tile_dim = self._largest_divisor_leq(input_sizes[i][self.dim], extra_concat)
                 extra_concat -= tile_dim
             else:
                 tile_dim = 1

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -848,7 +848,9 @@ def _index_expr(self, tile_desc, renamed_expression, index, base_vector_index):
             if idx == tile_desc.vmap.vlane_split_axis: # Need to add vector lane offset
                 stride_dim = ops.remainder(dim, vlane_stride_vec)
                 outer_dim = ops.remainder(ops.truncdiv(dim, vlane_stride_vec), vlane_outer_vec)
-                dim = ops.add(stride_dim, ops.mul(outer_dim, nr_vector_lane_vec))
+                # Next sublane-row stride is vector_lane*vlane_stride, not vector_lane alone.
+                row_stride_vec = ops.mul(nr_vector_lane_vec, vlane_stride_vec)
+                dim = ops.add(stride_dim, ops.mul(outer_dim, row_stride_vec))
 
                 with self.override_buffer_cse(buffer=self.const_buffer, cse=self.const_cse):
                     vlane_offset = ops.vlane_offset(vlane_vec, vlane_vec, attributes={"vlane_offset": offset}, comment="vlane offset")
@@ -1399,38 +1401,26 @@ def _masked_fill_bits(self, dtype, index):
         return bits & ((1 << (t.element_size() * 8)) - 1)
 
     def _masked_bounds(self, name, index, dram_stride, local_tile_desc, is_load, buffer, local_dims):
-        """Per tile-axis [low, high) clamp for a masked DMA. Per axis the valid GLOBAL range
-        is [glo, ghi) (store: [0, output_extent); padded load: [pad, pad + input_extent)), and
-        _emit_clamp turns it into the tile-local low/high SSA vars. See
-        docs/design/masked-dma-descriptor.md. Returns [(tile_axis, low_var, high_var), ...].
+        """Per tile-axis [low, high) clamp for a masked DMA -- ONLY the trailing tail of a
+        non-dividing loop extent (valid GLOBAL range per axis is [0, loop_extent)). _emit_clamp
+        turns it into the tile-local low/high SSA vars. Returns [(tile_axis, low_var, high_var)].
+
+        A padded load reads out-of-bounds positions, but we do NOT clamp those here: the
+        consumer already yields the correct value at pad positions (a compute-side arith.select
+        for a padded gather, or the pad op's own fill). Reverse-engineering per-dim padding from
+        the single flat index offset is ill-posed -- the offset mixes the tap shift with the
+        padding, and under channels_last the stride-1 (channel) axis absorbs the offset
+        remainder, yielding an impossible clamp (e.g. [4, 8) on a size-2 channel tile) that
+        zeroed the whole load and made channels_last depthwise conv 99.9% wrong. See
+        docs/design/masked-dma-descriptor.md.
         """
         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])
+            glo, ghi = 0, int(self.ranges[k])
             axes.append((d, iv, glo, ghi, int(tile_size[d])))
         return self._emit_clamp(axes, buffer)
 
@@ -1646,6 +1636,12 @@ def convert_indirect_indexing(self, index :sympy.Expr):
                 if arg.is_Mul and arg.args[0].is_number:
                     offset_stride = int(arg.args[0])
                 index = index.replace(arg, 0)
+            # A bare indirect index (e.g. x[idx]: index IS the symbol, so index.args is empty)
+            # is not caught by the loop above. Zero any remaining indirect symbol so the dram
+            # base offset does not reference the loop-local index value (the per-position gather
+            # is carried by the offset spad); otherwise the DMA offset affine.apply uses %sym
+            # before it is defined.
+            index = index.subs({s: 0 for s in index.free_symbols if str(s) in self.indirect_symbols})
             sram_var, _, _, _, tile_shape, _ = self.spad_buffer_dict[first_dim]
             return index, (sram_var, tile_shape, offset_stride)
 

PyTorchSimFrontend/mlir/mlir_common.py

@@ -411,7 +411,8 @@ def init_tile_size(ranges, vlane_stride, vector_lane):
             return [1]
         tile_size = [1] * nr_dim
         if nr_dim == 1:
-            tile_size[0] = 1 if ranges[0] == 1 else 2 * vlane_stride * vector_lane
+            # Cap to extent so the tile never over-reads the DRAM buffer (extent 128 vs tile 512 -> 384 garbage folded into the reduction). No-op when extent >= tile.
+            tile_size[0] = 1 if ranges[0] == 1 else min(2 * vlane_stride * vector_lane, ranges[0])
         elif nr_dim == 2:
             tile_size[-1] = vlane_stride * vector_lane
             tile_size[-2] = 2 * vector_lane

PyTorchSimFrontend/mlir/mlir_lowering.py

@@ -324,6 +324,19 @@ def _mlir_custom_sort_default(
         stable=stable_required,
     )
     sorted_values = mlir_template.generate(template_buffer_node=value).output_node()
+
+    def _unwrap(t):
+        t = t.data if isinstance(t, ir.TensorBox) else t
+        t = t.data if isinstance(t, ir.StorageBox) else t
+        return t
+    # The sort kernel writes `indices` in place (mlir_sort_template make_inplace), but only
+    # `sorted_values` is a scheduler-visible output. Advertise the indices write as a
+    # MutationOutput OWNED by the sort op (see MLIRTemplateBuffer) so the scheduler keeps the
+    # kernel alive when only indices are consumed (argsort: the values output is dead and
+    # would otherwise DCE the whole sort).
+    sort_op = _unwrap(sorted_values)
+    idx_buf = _unwrap(indices)
+    sort_op.add_mutation_output(ir.MutationOutput(idx_buf.get_layout(), idx_buf, sort_op))
     return sorted_values, indices