[Frontend] Added bug fix for autotune correctness issues.

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -1025,7 +1025,7 @@ def make_choices(self, nodes, kernel_name):
                     bench_runner = self.run_bench(nodes, kernel_name, src_code)
                     choices.append((bench_runner, src_code, meta_code, self.kernel_group.tile_desc.get_tile_size(), self.kernel_group.tile_desc.vmap.vlane_stride))
                     prevent_infinite_loop += 1
-        self.kernel_group.tile_desc.prev_tail_threshold = prev_tail_threshold
+        self.kernel_group.tile_desc.tail_ratio_threshold = prev_tail_threshold
         return choices
 
     def autotune(self, *args):
@@ -1583,3 +1583,4 @@ def convert_indirect_indexing(self, index :sympy.Expr):
             if mlir_dtype != "index":
                 out = ops.index_cast(out, "index")
         return index + sympy.Symbol(str(out)), compute_dependecy
+

PyTorchSimFrontend/mlir/mlir_common.py

@@ -420,23 +420,37 @@ def trim_large_tail(self, ranges: list[int]):
                 BETA = 0
 
             padding_ratio = TileAdjustMixin.get_padding_ratio(tile_range, dim_range)
-            if padding_ratio < self.tail_ratio_threshold:
+            if padding_ratio == 0:
+                continue
+            if padding_ratio < self.tail_ratio_threshold and not constraint.must_divide_dim:
                 continue
             best_tile = tile_range
             best_cost = (
                 ALPHA * padding_ratio +
                 BETA * (dim_range / tile_range)
             )
 
-            min_tile = 1
+            # Candidates below the axis granularity are not representable
+            min_tile = max(1, constraint.multiple_of)
             for candidate in range(tile_range - 1, min_tile - 1, -1):
-                new_candidate = constraint.adjust(tile_range, candidate, dim_range)
+                try:
+                    new_candidate = constraint.adjust(tile_range, candidate, dim_range)
+                except extension_codecache.TileSizeError:
+                    continue
                 ratio = TileAdjustMixin.get_padding_ratio(new_candidate, dim_range)
                 iter_penalty = (dim_range / new_candidate)
 
                 cost = ALPHA * ratio + BETA * iter_penalty
                 if cost < best_cost:
                     best_tile, best_cost = new_candidate, cost
+
+            if constraint.must_divide_dim and dim_range % best_tile:
+                # No tile both divides the dimension and respects the axis
+                # granularity. Overhang with this tile would corrupt data.
+                best_tile = next(c for c in range(tile_range, 0, -1)
+                                 if dim_range % c == 0)
+                if i == self.vmap.vlane_split_axis and best_tile % self.vmap.vlane_stride:
+                    self.vmap.vlane_stride = 1
             self._tile_size[i] = best_tile
 
     def select_vlane_axis(self):
@@ -533,6 +547,10 @@ def __init__(self, tile_size, vector_lane, vlane_split_axis=None, vlane_stride=N
         self._tile_size = list(tile_size)
         self._tile_stride = None
         self.tile_constraint = [TileConstraint(vlane_stride if idx == vlane_split_axis else 1) for idx, _ in enumerate(tile_size)]
+
+        for constraint in self.tile_constraint:
+            if constraint.multiple_of == 1:
+                constraint.must_divide_dim = True
         self.tile_axis_order = list(range(len(tile_size)))
         self.update_tile_stride()
 
@@ -1160,4 +1178,4 @@ def mark_parallel(self, par_depth):
         loops[0].parallel = par_depth
         for i in range(1, par_depth):
             loops[i].collapsed = True
-        loops[0].simd = loops[par_depth - 1].simd
+        loops[0].simd = loops[par_depth - 1].simd