[Fix] indirect access attribute

Author: YWHyuk

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -395,10 +395,6 @@ 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):
@@ -407,8 +403,8 @@ def parse_indices(self, expr, comments="", indices=None, indirect_dims=[]) -> co
         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.get(idx, idx) for idx in indices]
-            sorted_pairs = sorted(indices_as_keys, key=lambda x: str(x))
+            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_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_template.py

@@ -856,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)
@@ -909,7 +909,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)
@@ -1013,7 +1013,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