[Fix] WIP

Author: YWHyuk

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -1,5 +1,6 @@
 import contextlib
 import sympy
+import sys
 import re
 import os
 from functools import reduce
@@ -374,9 +375,51 @@ def _convert_sympy_to_mlir_expr(self, expr, sorted_args):
             expr = expr.replace(target_arg, new_arg)
             indices.append(str(new_arg))
 
-        expr_str = str(expr)
-        if "//" in expr_str:
-            expr_str = expr_str.replace("//", " floordiv ")
+        # Convert ModularIndexing and FloorDiv to sympy expressions
+        # ModularIndexing(x, y, z) means (x // y) % z -> Mod(FloorDiv(x, y), z)
+        # FloorDiv(x, y) means x // y -> will be converted to floordiv in string representation
+        # Use preorder_traversal to find all instances
+        replacements = {}
+        for sub in sympy.preorder_traversal(expr):
+            if isinstance(sub, ModularIndexing):
+                # Convert ModularIndexing to Mod(FloorDiv(...), ...)
+                if sub.args[1] != 1:
+                    floor_div = FloorDiv(sub.args[0], sub.args[1])
+                else:
+                    floor_div = sub.args[0]
+                mod_expr = sympy.Mod(floor_div, sub.args[2])
+                replacements[sub] = mod_expr
+            elif isinstance(sub, FloorDiv):
+                # Keep FloorDiv as is, will be handled in custom string conversion
+                # We need to mark it for special handling
+                pass
+
+        # Apply replacements
+        for old_expr, new_expr in replacements.items():
+            expr = expr.subs(old_expr, new_expr)
+
+        # Custom string conversion for MLIR affine expressions
+        def mlir_str(expr):
+            """Convert sympy expression to MLIR affine expression string"""
+            if isinstance(expr, FloorDiv):
+                return f"({mlir_str(expr.args[0])} floordiv {mlir_str(expr.args[1])})"
+            elif isinstance(expr, sympy.Mod):
+                return f"({mlir_str(expr.args[0])} mod {mlir_str(expr.args[1])})"
+            elif isinstance(expr, sympy.Add):
+                terms = [mlir_str(term) for term in expr.args]
+                return " + ".join(terms)
+            elif isinstance(expr, sympy.Mul):
+                factors = [mlir_str(factor) for factor in expr.args]
+                return " * ".join(factors)
+            elif isinstance(expr, sympy.Symbol):
+                return str(expr)
+            elif expr.is_number:
+                return str(expr)
+            else:
+                # Fallback to string representation
+                return str(expr)
+
+        expr_str = mlir_str(expr)
         return expr_str, indices
 
     def parse_indices(self, expr, comments="", indices=None, indirect_dims=[]) -> common.CSEVariable:
@@ -470,6 +513,10 @@ def load(self, name: str, index: sympy.Expr):
         tile_numel_per_lane = local_tile_desc.get_numel_per_lane()
         tile_shape = local_tile_desc.get_mlir_shape(mlir_dtype)
         tile_stride = local_tile_desc.get_tile_stride()
+
+        if len(dram_stride) != len(tile_stride):
+            print("here")
+
         # Compute vector unit size
         vshape = self.kernel_group.tile_desc.get_mlir_vshape(mlir_dtype)
         compute_vec_size = self.kernel_group.tile_desc.get_compute_vec_size()
@@ -1157,7 +1204,8 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe
                 sorted_constraints = sorted(axis_constraints, key=lambda c: int(c.args[1]))
                 for constraint in sorted_constraints[1:]:
                     index = index.replace(constraint.original_expr, 0)
-
+        if self.is_modular_indexing(index):
+            print("here")
         # Calculate dram stride
         dram_stride = [0] * local_tile_desc.get_nr_dim()
         if index.is_Symbol:
@@ -1167,13 +1215,20 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe
             pass
         else:
             dram_dict = defaultdict(list)
+            implicit_dim_divisors = defaultdict(lambda: sys.maxsize)
             # Assume that div will have high priority than mod
             for arg in index.as_ordered_terms():
                 coeff, dim = arg.as_coeff_mul()
                 if len(dim) == 0:
                     continue
                 real_dim = list(dim[0].free_symbols)[0]
-                dram_dict[str(real_dim)].append(coeff)
+                if dim[0].has(ModularIndexing):
+                    if dim[0].args[1] < implicit_dim_divisors[str(real_dim)]:
+                        implicit_dim_divisors[str(real_dim)] = dim[0].args[1]
+                        dram_dict[str(real_dim)] = [coeff]
+                else:
+                    dram_dict[str(real_dim)].append(coeff)
+
             # Add missing dims if not added
             max_dim = len(self.ranges) if not store_reduction else len(self.ranges) - 1
             for i in range(max_dim):

PyTorchSimFrontend/mlir/mlir_common.py

@@ -504,7 +504,7 @@ def __init__(self, tile_size, vector_lane, vlane_split_axis=None, vlane_stride=N
             vlane_stride=vlane_stride
         )
 
-        self.implicit_dim_size = None
+        self.implicit_dim_size = dict()
         self.nr_rdim = 0
         self.offset = sympy.Integer(0) # Dram offset
 
@@ -654,6 +654,10 @@ def reduction(self, dtype, src_dtype, reduction_type, value):
     def indirect_indexing(self, index_var, size, check, wrap_neg):
         raise NotImplementedError()
 
+    def check_bounds(self, index, size, lower, upper):
+        return
+        #raise NotImplementedError()
+
     def codegen_global_init(self):
         raise NotImplementedError()
 
@@ -918,6 +922,10 @@ def indirect_indexing(index_var, size, check=True, wrap_neg=True):
                 # Skip CSE since this doesn't return an expression
                 return self.indirect_indexing(index_var, size, check, wrap_neg)
 
+            @staticmethod
+            def check_bounds(index, size, lower, upper):
+                return self.check_bounds(index, size, lower, upper)
+
             @staticmethod
             def load(name: str, index: sympy.Expr):
                 index = self.rename_indexing(index)
@@ -964,6 +972,10 @@ def store_reduction(name, index, value):
             def reduction(dtype, src_dtype, reduction_type, value):
                 return self.reduction(dtype, src_dtype, reduction_type, value)
 
+            @staticmethod
+            def check_bounds(index, size, lower, upper):
+                return self.check_bounds(index, size, lower, upper)
+
             @staticmethod
             def _index_expr(tile_size, buffer, renamed_expression, index):
                 return self._index_expr(tile_size, buffer, renamed_expression, index)