[Fix] Fusion axis mechanism change

Author: YWHyuk

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -400,12 +400,6 @@ def parse_indices(self, expr, comments="", indices=None, indirect_dims=[]) -> co
         with self.override_buffer_cse(buffer=self.global_vars, cse=self.map_cse):
             map_var = ops.affine_map(indices, expr_str, symbol_names=indirect_dims)
 
-        if hasattr(self, "dim_aliasing"):
-            # Create reverse mapping: value -> key
-            reverse_mapping = {v: k for k, v in self.dim_aliasing.items()}
-            indices_as_keys = [reverse_mapping[idx] for idx in indices]
-            sorted_pairs = sorted(indices_as_keys, key=lambda x: list(self.dim_aliasing.keys()).index(x))
-            indices = [self.dim_aliasing[idx] for idx in sorted_pairs]
         index = ops.affine_apply(map_var, indices, indirect_dims=indirect_args, comment=comments)
         return index
 

PyTorchSimFrontend/mlir/mlir_common.py

@@ -614,7 +614,7 @@ def __init__(self, kernel_group, reason=None):
         self.target_cse_override = contextvars.ContextVar(f"Handler_cse_override_{instance_id}", default=self.cse)
         self._nested_context_depth = 0
 
-    def set_ranges(self, lengths, reduction_lengths):
+    def set_ranges(self, lengths, reduction_lengths, index_names=None):
         if self.call_ranges:
             assert self.call_ranges == tuple(lengths) + tuple(
                 reduction_lengths
@@ -623,7 +623,12 @@ def set_ranges(self, lengths, reduction_lengths):
         else:
             self.call_ranges = tuple(lengths) + tuple(reduction_lengths)
             self.ranges = [self.rename_indexing(x) for x in self.call_ranges]
-            self.itervars = [sympy.Symbol(f"index{n}") for n in range(len(self.ranges))]
+            if index_names is None:
+                self.itervars = [sympy.Symbol(f"index{n}") for n in range(len(self.ranges))]
+            else:
+                assert len(index_names) == len(self.ranges), f"Index names length mismatch: {len(index_names)} != {len(self.ranges)}"
+                self.itervars = [sympy.Symbol(str(n)) for n in index_names]
+
             self.itervar_cses = {str(index) : self.register_var_cse(str(index), 1, "index") for index in self.itervars}
             self.reduction_depth = len(lengths)
         return (
@@ -867,18 +872,22 @@ def rename_indexing(self, index) -> sympy.Expr:
     def override_buffer_cse(self, *, buffer=None, cse=None):
         buffer_override = self.target_buffer_override
         cse_override = self.target_cse_override
-        target_buffer = target_cse = None
+        buffer_token = cse_token = None
         try:
+            # Store tokens for proper restoration in nested contexts
+            # contextvars.set() returns the previous value (token) which can be used for reset()
             if buffer is not None:
-                target_buffer = buffer_override.set(buffer)
+                buffer_token = buffer_override.set(buffer)
             if cse is not None:
-                target_cse = cse_override.set(cse)
+                cse_token = cse_override.set(cse)
             yield self
         finally:
-            if target_cse is not None:
-                cse_override.reset(target_cse)
-            if target_buffer is not None:
-                buffer_override.reset(target_buffer)
+            # Restore using tokens - contextvars automatically handles nested contexts
+            # Each level restores to its own previous value
+            if cse_token is not None:
+                cse_override.reset(cse_token)
+            if buffer_token is not None:
+                buffer_override.reset(buffer_token)
 
     def __enter__(self):
         class CSEProxy:

PyTorchSimFrontend/mlir/mlir_template.py

@@ -469,7 +469,7 @@ def codegen_template_code(self, render, template_node, prologue_nodes, epilogue_
                     _, (group, reduction_group) = max(
                         epilogue_nodes, key=lambda x: int(x.is_reduction())
                     ).group
-                    vars, reduction_vars = kernel.set_ranges(group, reduction_group)
+                    vars, reduction_vars = kernel.set_ranges(group, reduction_group, list(self.dim_aliasing.values()))
                     for node in epilogue_nodes:
                         node.codegen((vars, reduction_vars))
 
@@ -893,7 +893,6 @@ def load_epilogue(self, name: str, index: sympy.Expr):
                     map_var = ops.affine_map(["d0", "d1"], f"d0 + d1*{(self.r_tile_size)}")
                 with self.override_buffer_cse(buffer=self.loads):
                     offset = ops.affine_apply(map_var, [self.compute_idx, self.reduction_loop_idx])
-                #offset = self.cse.generate(self.loads, f"affine.apply affine_map<(d0, d1) -> (d0 + d1*{(self.r_tile_size)})>(%{self.compute_idx}, %{self.reduction_loop_idx})")
                 compute_index_var = ",".join([f"%{zero_var}"] * (self.kernel_group.tile_desc.get_nr_dim()-1) + [f"%{offset}"])
 
             with self.override_buffer_cse(buffer=self.loads):
@@ -1124,22 +1123,6 @@ def set_tile_size(self, template_fusion_info, prologue=False):
                 self.compute_body_loop.step = tile_desc.get_compute_vec_size()
         return tile_desc
 
-    def rename_indexing(self, index) -> sympy.Expr:
-        # First step: replace dim_name with tmp_+dim_aliased_name to avoid circular dependencies
-        # (e.g., {"index0":"index1", "index1":"index0"})
-        tmp_subs = {
-            sympy.Symbol(dim_name): sympy.Symbol("tmp_"+dim_aliased_name)
-            for dim_name, dim_aliased_name in self.dim_aliasing.items()
-        }
-        index = index.subs(tmp_subs)
-        # Second step: replace tmp_+dim_aliased_name with dim_aliased_name
-        final_subs = {
-            sympy.Symbol("tmp_"+dim_aliased_name): sympy.Symbol(dim_aliased_name)
-            for dim_aliased_name in self.dim_aliasing.values()
-        }
-        index = index.subs(final_subs)
-        return index
-
 class MLIRTemplateCaller(CUDATemplateCaller):
     def __str__(self):
         return f"MLIRTemplateCaller(source_file={self.bmreq.source_file})"