[Fix] Fusion axis mechanism change

Author: YWHyuk

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -313,7 +313,9 @@ def __init__(self, kernel_group, reason=None):
         self.base_vector_initialized = False
 
     def reset(self, reason):
+        save = self.exit_stack, self._nested_context_depth
         self.__init__(self.kernel_group, reason=reason)
+        self.exit_stack, self._nested_context_depth = save
 
     # padding type 0: zero-padding 1: negative-padding(-inf) ...
     def get_padding_type(self):
@@ -395,17 +397,11 @@ def parse_indices(self, expr, comments="", indices=None, indirect_dims=[]) -> co
 
         # Convert sympy expression to affine map expression
         expr_str, indices = self._convert_sympy_to_mlir_expr(expr, sorted_args)
-
-        # Extract index var
-        if len(indirect_dims):
-            comments = "{indirect_access} " + comments # Add indirect access attribute
         indirect_args = [f"%{i}" for i in indirect_dims]
         # Create affine.apply operation
         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"):
-            indices = [self.dim_aliasing.get(index, index) for index in indices]
         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_ops.py

@@ -1186,7 +1186,7 @@ def affine_apply(map_var, indices, indirect_dims=None, comment=None, *args, **kw
         # Add indirect dimensions if provided
         if indirect_dims:
             indirect_str = ", ".join(indirect_dims)
-            op_str += f"[{indirect_str}]"
+            op_str += f"[{indirect_str}] {{indirect_access}}"
         if comment:
             op_str += f" // {comment}"
         return op_str, [1, "index"]

PyTorchSimFrontend/mlir/mlir_scheduling.py

@@ -178,7 +178,7 @@ def can_fuse_horizontal(self, node1, node2):
                 return False
 
             size_match = node1.get_nodes()[0].node.get_numel() == reduce(operator.mul, node2.get_nodes()[0].node.get_size(), 1) * reduce(operator.mul, node2.get_nodes()[0].node.get_reduction_size(), 1)
-            target_symbol = symbols("r0")
+            target_symbol = symbols("r0_0")
             try:
                 stride = [i.strip()[:-1].split(",")[-1].strip() for i in str(node2.get_nodes()[0].node).split("\n") if "r0" in i][1]
                 stride = int(sympify(stride).coeff(target_symbol))

PyTorchSimFrontend/mlir/mlir_template.py

@@ -429,7 +429,7 @@ def codegen_template_code(self, render, template_node, prologue_nodes, epilogue_
                     ).group
                     prologue_tile_desc = kernel.set_tile_size(kernel.prologue_info, prologue=True)
                     kernel.kernel_group.set_tile_info(prologue_tile_desc)
-                    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 prologue_nodes:
                         # Reuse created spad
                         read_list = sorted([i.name for i in node.read_writes.reads])
@@ -469,10 +469,11 @@ 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))
 
+        with self as kernel:
             src_code = (
                 partial_code
                 if isinstance(partial_code, str)
@@ -855,7 +856,7 @@ def load_epilogue(self, name: str, index: sympy.Expr):
         # Want to use tile_desc from epilogue_info
         with self.override_buffer_cse(buffer=self.applys, cse=self.apply_cse):
             index_var = self.parse_indices(index)
-        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.keys()]
+        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.values()]
         vlane_split_axis = self.kernel_group.tile_desc.vmap.vlane_split_axis
         vlane_stride = self.kernel_group.tile_desc.vmap.vlane_stride
         tile_shape = self.kernel_group.tile_desc.get_mlir_shape(mlir_dtype)
@@ -892,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):
@@ -908,7 +908,7 @@ def store_epilogue(self, name: str, index: sympy.Expr, value, *args, **kwargs):
 
         with self.override_buffer_cse(buffer=self.applys, cse=self.apply_cse):
             index_var = self.parse_indices(index)
-        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.keys()]
+        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.values()]
         vlane_split_axis = self.kernel_group.tile_desc.vmap.vlane_split_axis
         vlane_stride = self.kernel_group.tile_desc.vmap.vlane_stride
         tile_shape = self.kernel_group.tile_desc.get_mlir_shape(mlir_dtype)
@@ -1012,7 +1012,7 @@ def store_reduction_epilogue(self, name, index, value):
 
         with self.override_buffer_cse(buffer=self.reductions_suffix, cse=self.apply_cse):
             index_var = self.parse_indices(index, comments="// Store reduction")
-        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.keys()][:-1] # Assume that there is only one reduction axis
+        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.values()][:-1] # Assume that there is only one reduction axis
         vlane_split_axis = self.kernel_group.tile_desc.vmap.vlane_split_axis
         vlane_stride = self.kernel_group.tile_desc.vmap.vlane_stride
 
@@ -1123,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})"