From 1a8b8baf99b2493a16f50d841e8f24b64eebdace Mon Sep 17 00:00:00 2001 From: Karthik Lakshmanan Date: Mon, 31 Aug 2026 00:27:33 +0530 Subject: [PATCH] Added fix for tiling handling. Fixes correctness issue with autotune codegen mapping. --- .../mlir/mlir_codegen_backend.py | 3 ++- PyTorchSimFrontend/mlir/mlir_common.py | 26 ++++++++++++++++--- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 492b7416..ece89a67 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/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 + diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 5cde19eb..f3b6bcd1 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -420,7 +420,9 @@ 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 = ( @@ -428,15 +430,27 @@ def trim_large_tail(self, ranges: list[int]): 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 \ No newline at end of file + loops[0].simd = loops[par_depth - 1].simd