[Frontend] Introduce recompile signal + Support floordiv pattern

Author: YWHyuk

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -15,7 +15,7 @@
     is_welford_reduction,
     sympy_product
 )
-from torch.utils._sympy.functions import ModularIndexing
+from torch.utils._sympy.functions import ModularIndexing, FloorDiv
 import PyTorchSimFrontend.extension_codecache as extension_codecache
 
 from PyTorchSimFrontend import extension_config
@@ -260,10 +260,10 @@ def binary_elementwise_common(operand1, operand2, var_info):
                 operand2 = ops.to_dtype(operand2, op_type1[1], var_info)
                 op_type2 = var_info[operand2]
             elif op_type1[1][0] == op_type2[1][0]:
-                if int(op_type1[1][1:]) > int(op_type2[1][1:]):
+                if mlir_common.MLIR_TO_BIT[op_type1[1]] > mlir_common.MLIR_TO_BIT[op_type2[1]]:
                    operand2 = ops.ext(operand2, op_type1[1])
                    op_type2 = var_info[operand2]
-                elif int(op_type1[1][1:]) < int(op_type2[1][1:]):
+                elif mlir_common.MLIR_TO_BIT[op_type1[1]] < mlir_common.MLIR_TO_BIT[op_type2[1]]:
                    operand1 = ops.ext(operand1, op_type2[1])
                    op_type1 = var_info[operand1]
             else:
@@ -348,17 +348,21 @@ def maximum(operand1, operand2, *args, var_info=None, **kwargs):
     @staticmethod
     def to_dtype(operand, dst_mlir_dtype, *args, var_info=None, **kwargs):
         src_mlir_dtype = var_info[operand][1]
+        if src_mlir_dtype == "index":
+            operand = ops.index_cast(operand, "i64", var_info=var_info)
+            src_mlir_dtype = var_info[operand][1]
+
         tile_size = var_info[operand][0]
         if isinstance(dst_mlir_dtype, torch.dtype):
             dst_mlir_dtype = mlir_common.DTYPE_TO_MLIR[dst_mlir_dtype]
-        dst_bits = int(dst_mlir_dtype[1:])
-        src_bits = int(src_mlir_dtype[1:])
+        dst_bits = mlir_common.MLIR_TO_BIT[dst_mlir_dtype]
+        src_bits = mlir_common.MLIR_TO_BIT[src_mlir_dtype]
         shape = f"vector<{tile_size}x{dst_mlir_dtype}>" if tile_size > 1 else dst_mlir_dtype
         src_shape = f"vector<{tile_size}x{src_mlir_dtype}>" if tile_size > 1 else src_mlir_dtype
         if dst_mlir_dtype[0] == "i" and src_mlir_dtype[0] == "f":
-            return f"arith.fptoui%{operand} : {src_shape} to {shape}", [tile_size, dst_mlir_dtype]
+            return f"arith.fptoui %{operand} : {src_shape} to {shape}", [tile_size, dst_mlir_dtype]
         if dst_mlir_dtype[0] == "f" and src_mlir_dtype[0] == "i":
-            return f"arith.uitofp%{operand} : {src_shape} to {shape}", [tile_size, dst_mlir_dtype]
+            return f"arith.uitofp %{operand} : {src_shape} to {shape}", [tile_size, dst_mlir_dtype]
         if dst_mlir_dtype[0] == "i":
             if dst_bits > src_bits:
                 return f"arith.extui %{operand} : {src_shape} to {shape}", [tile_size, dst_mlir_dtype]
@@ -955,6 +959,8 @@ def parse_indices(self, expr, buffer=None, comments="", indirect_dims=[]) -> com
         # Extract index var
         indirect_args = [f"%{i}" for i in indirect_dims]
         expr_str = str(expr)
+        if "//" in expr_str:
+            expr_str = expr_str.replace("//", " floordiv ")
         args = ", ".join(map(str, indices))
         map_var = self.map_cse.generate(self.global_vars, f"affine_map<({args})[{','.join(indirect_dims)}] -> ({expr_str})>")
         args = ", ".join([f"%{i}" for i in indices])
@@ -1063,6 +1069,9 @@ def store(self, name: str, index: sympy.Expr, value, *args, **kwargs):
         vshape = self.kernel_group.tile_desc.get_mlir_vshape(mlir_dtype)
         compute_vec_size = self.kernel_group.tile_desc.get_compute_vec_size()
         require_store = True
+        if compute_vec_size < self.var_info[value][0]:
+            value = self.cse.generate(self.stores, f"vector.extract_strided_slice  %{value} {{offsets = [0], sizes = [{compute_vec_size}], strides = [1]}}: vector<{self.var_info[value][0]}x{self.var_info[value][1]}> to {vshape}")
+            self.register_var_info(value, [compute_vec_size, mlir_dtype])
 
         if str(value) in self.spad_buffer_dict:
             # Todo. If tile_size is not same (i.e., view operation), we can't apply peephole optimization easily
@@ -1680,6 +1689,40 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe
             sorted_keys = sorted(dram_dict.keys())
             dram_stride = sum((dram_dict[key] for key in sorted_keys), [])
 
+        # Support floordiv pattern
+        # FIXME. How to integrate implicit dims and floordiv?
+        # This was introduced to support GroupNorm
+        if index.has(FloorDiv) and not index.has(ModularIndexing):
+            dim_divisor = [1] * len(local_dims)
+            for sub in sympy.preorder_traversal(index):
+                if isinstance(sub, FloorDiv):
+                    if not str(sub.args[0]).startswith("index"):
+                        continue
+                    dim_idx = int((str(sub.args[0])[5:]))
+                    if int(self.kernel_group.tile_desc.get_tile_size()[dim_idx] % sub.args[1]) != 0:
+                        # In this case, need to recompile
+                        original_size = self.kernel_group.tile_desc.get_tile_size()[dim_idx]
+                        divisor = sub.args[1]
+                        new_size = ((original_size + divisor - 1) // divisor) * divisor
+                        new_tile_sizes = list(self.kernel_group.tile_desc.get_tile_size())
+                        new_tile_sizes[dim_idx] = new_size
+                        self.kernel_group.tile_desc.set_tile_size(new_tile_sizes)
+
+                        # Send recompile signal
+                        self.reset("recompile")
+                        raise mlir_common.RecompileSignal(f"Tile size {self.kernel_group.tile_desc.get_tile_size()[dim_idx]} is not divisible by {sub.args[1]}")
+                    dim_divisor[dim_idx] = sub.args[1]
+
+            # Update dram_stride, just insert 0 next to target dim
+            offset = 0
+            for dim_idx, divisor in enumerate(dim_divisor):
+                if divisor == 1:
+                    continue
+                dram_stride.insert(dim_idx+offset+1, 0)
+                local_tile_desc.apply_divisor(dim_idx+offset, divisor, "pad")
+                local_tile_desc.apply_divisor(dim_idx+offset, divisor, "split")
+                offset = offset+1
+
         # FIXME. It will be nice to modify node instead of this exception handling...
         if len(self.itervars) == 1 and self.reduction_depth == 0:
             # In case of reduction loop only case, we will add dummy loop so shift it once
@@ -1810,7 +1853,11 @@ def convert_indirect_indexing(self, index :sympy.Expr):
 
         # Load indirect operands
         for target_dim in indirect_dims:
-            sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[target_dim]
+            if target_dim in self.spad_buffer_dict:
+                sram_var, _, tile_numel_per_lane, sram_index_var, tile_shape, vshape = self.spad_buffer_dict[target_dim]
+            else:
+                raise NotImplementedError("TODO.")
+
             mlir_dtype = vshape.split("x")[1][:-1]
             vshape = f"vector<{tile_numel_per_lane}x{mlir_dtype}>" # FIXME. Maybe require fine grain compute...
             if tile_numel_per_lane > 1:

PyTorchSimFrontend/mlir/mlir_common.py

@@ -61,6 +61,19 @@
     torch.bfloat16: "bfloat16",
 }
 
+MLIR_TO_BIT = {
+    "i1": 1,
+    "i8": 8,
+    "i16": 16,
+    "i32": 32,
+    "i64": 64,
+    "f16": 16,
+    "f32": 32,
+    "f64": 64,
+    "bf16": 16,
+    "index": 64
+}
+
 DTYPE_LOWP_FP = [
     torch.bfloat16,
     torch.float16,
@@ -105,6 +118,14 @@ def ctx():
 
         return ctx()
 
+class RecompileSignal(BaseException):
+    """
+    Exception raised when a recompilation of a kernel or code block is required.
+    """
+    def __init__(self, message="Recompilation requested."):
+        self.message = message
+        super().__init__(self.message)
+
 class MLIRKernelArgs(common.KernelArgs):
     MLIR_ARGS_IN = 0x01
     MLIR_ARGS_OUT = 0x02
@@ -310,7 +331,7 @@ def get_compute_vec_size(self):
         if self.vec_size is not None:
             return self.vec_size
         if self.nr_rdim:
-            assert self.nr_rdim==1
+            assert self.nr_rdim!=0
             val = self.get_numel_per_lane() // self._tile_size[-1]
             if self.get_numel_per_lane() >= val * 8:
                 return val*8
@@ -331,6 +352,44 @@ def get_compute_vec_size(self):
     def div_round_up(size, round_val):
         return (size + round_val - 1) // round_val
 
+    def apply_divisor(self, axis: int, divisor: int, mode: str = "split"):
+        # Apply divisor to tile size at given axis.
+        # This method based on axis order.
+        old_size = self._tile_size[axis]
+        if divisor == 1:
+            return
+        padded = self.div_round_up(old_size, divisor) * divisor
+        outer  = self.div_round_up(old_size, divisor)
+        inner  = divisor
+        if mode == "pad":
+            self._tile_size[axis] = padded
+            self.update_tile_stride()
+            return
+        elif mode == "split":
+            new_sizes = list(self._tile_size)
+            new_sizes[axis] = outer
+            new_sizes.insert(axis + 1, inner)
+            self._tile_size = new_sizes
+
+            # Update tile_axis_order
+            old_order_val = self.tile_axis_order[axis]
+            new_order = list(self.tile_axis_order)
+            new_order.insert(axis + 1, old_order_val + 0.1)
+            sorted_pairs = sorted(
+                zip(range(len(new_order)), new_order),
+                key=lambda x: x[1]
+            )
+            self.tile_axis_order = [idx for idx, _ in sorted_pairs]
+            self.update_tile_stride()
+
+            if self.vlane_split_axis == axis:
+                self.vlane_split_axis = axis
+            elif self.vlane_split_axis > axis:
+                self.vlane_split_axis += 1
+            return
+        else:
+            raise ValueError(f"Unknown mode: {mode}. Supported modes are 'pad' and 'split'.")
+
 class MLIRWrapperKenrelGroup(cpp.KernelGroup):
     def __init__(self):
         super().__init__()
@@ -538,6 +597,8 @@ def dummy_tile_size():
                 dim = int(self.recodegen.split("_")[-1])
                 tile_size = self.kernel_group.tile_desc.get_tile_size() # TODO:
                 tile_size[dim] = tile_size[dim] * 2
+            elif self.recodegen == "recompile":
+                return self.kernel_group.tile_desc
             else:
                 raise NotImplementedError(f"Unknown recodegen reason: {self.recodegen}")
 
@@ -608,26 +669,36 @@ def dummy_tile_size():
         return tile_desc
 
     def codegen_nodes(self, nodes, kernel_name):
-        _, (group, reduction_group) = max(
-            nodes, key=lambda x: int(x.is_reduction())
-        ).group
-
-        # Set node range info
-        vars, reduction_vars = self.set_ranges(group, reduction_group)
-        tile_desc = self.compute_tile_size(nodes, vars, reduction_vars)
-        self.compute_body_loop.size = tile_desc.get_numel_per_lane()
-        self.compute_body_loop.step = tile_desc.get_compute_vec_size()
-        self.kernel_group.set_tile_info(tile_desc)
-
-        _, _, _, self.buffer_types = self.kernel_group.args.mlir_argdefs()
-        with self as kernel:
-            for node in nodes:
-                node.run(vars, reduction_vars)
-        V.graph.removed_buffers |= self.removed_buffers
-        # V.graph.inplaced_to_remove |= self.inplaced_to_remove
-        src_code = self.codegen_kernel(kernel_name=kernel_name)
-        self.meta_kernel()
-        return src_code
+        recompile_try = 0
+        max_retry_compile = 5
+        while True:
+            _, (group, reduction_group) = max(
+                nodes, key=lambda x: int(x.is_reduction())
+            ).group
+
+            # Set node range info
+            vars, reduction_vars = self.set_ranges(group, reduction_group)
+            tile_desc = self.compute_tile_size(nodes, vars, reduction_vars)
+            self.compute_body_loop.size = tile_desc.get_numel_per_lane()
+            self.compute_body_loop.step = tile_desc.get_compute_vec_size()
+            self.kernel_group.set_tile_info(tile_desc)
+            try:
+                _, _, _, self.buffer_types = self.kernel_group.args.mlir_argdefs()
+                with self as kernel:
+                    for node in nodes:
+                        node.run(vars, reduction_vars)
+            except RecompileSignal as e:
+                recompile_try += 1
+                if recompile_try > max_retry_compile:
+                    raise RuntimeError("Failed to compile kernel after multiple attempts.")
+                # Retry compile nodes
+                #print(f"Try recompile({recompile_try}/{max_retry_compile}). Reason: {e}")
+                continue
+            V.graph.removed_buffers |= self.removed_buffers
+            # V.graph.inplaced_to_remove |= self.inplaced_to_remove
+            src_code = self.codegen_kernel(kernel_name=kernel_name)
+            self.meta_kernel()
+            return src_code
 
     def run_bench(self, nodes, kernel_name, src_code):
         _, _, arg_attributes, _ = self.kernel_group.args.mlir_argdefs()

PyTorchSimFrontend/mlir/mlir_scheduling.py

@@ -46,10 +46,15 @@ def can_fuse_with_exceptions(self, node1: BaseSchedulerNode, node2: BaseSchedule
             if (isinstance(base_template_node1[0].node.template, MLIRGemmTemplate) or isinstance(base_template_node1[0].node.template, MLIRBMMTemplate)) and node2.is_reduction():
                 # For matmul/bmm+reduction case
                 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)
-                stride = [i.strip()[:-1].split(",")[-1].strip() for i in str(node2.get_nodes()[0].node).split("\n") if "r0" in i][1]
                 target_symbol = symbols("r0")
+                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))
+                except sympy.core.SympifyError:
+                    return False
+
                 # We can't fuse dim=-1
-                layout_possible = int(sympify(stride).coeff(target_symbol)) != 1
+                layout_possible = stride != 1
                 # Directed linked?
                 dependency_check = node2.get_nodes()[0] in [node.node for node in base_template_node1[0].users]# and len(node2.read_writes.reads)==1
                 dependency_size = all([i.get_numel() == node1.get_nodes()[0].node.get_numel() for i in node2.read_writes.reads])