mlir: add gpu_block and gpu_thread primitive

src/xtc/backends/jir/JIRScheduler.py

@@ -346,6 +346,16 @@ def distributed_buffer_at(
         # TODO: not implemented for now
         pass
 
+    @override
+    def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
+    @override
+    def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
     def get_schedule_str(self) -> str:
         return str(JIRSchedule(scheduler=self))
 

src/xtc/backends/mlir/MlirCompilerPasses.py

@@ -23,13 +23,19 @@
     MatchInterfaceEnum,
     FuseIntoContainingOp,
 )
+from mlir.dialects.transform.gpu import (
+    MapForallToBlocks,
+    MapNestedForallToThreads,
+)
 from mlir.dialects.transform.loop import loop_unroll
 from mlir.dialects.transform import SplitHandleOp
 from mlir.ir import (
     Location,
     InsertionPoint,
     UnitAttr,
     OpResult,
+    Attribute,
+    ArrayAttr,
 )
 from mlir.passmanager import PassManager
 from mlir.ir import Module
@@ -52,6 +58,7 @@
 _VECTO_SEQ_NAME = "_vecto"
 _SUPER_VECTORIZE_SEQ_NAME = "_super_vectorize"
 _POST_BUFFERIZE_SEQ_NAME = "_post_bufferize"
+_GPU_DIM = ["x", "y", "z"]
 
 
 @dataclass
@@ -130,6 +137,7 @@ def __init__(
         self._super_vectorize_sequence: NamedSequenceOp | None = None
         self._post_bufferize_sequence: NamedSequenceOp | None = None
         self._named_sequence: NamedSequenceOp | None = None
+        self._gpu_block_order: ArrayAttr | None = None
         self._nodes_schedules = (
             self._mlir_schedule.schedule_impl if self._mlir_schedule is not None else []
         )
@@ -234,6 +242,13 @@ def _generate_scheduling(self) -> OpResult:
                 )
                 if schedule.vectorization or self._always_vectorize:
                     self._post_vectorize(scheduling_state, schedule)
+
+                # GPU mapping
+                if schedule.gpu_blocks:
+                    self._gpu_mapping(
+                        schedule,
+                        scheduling_state,
+                    )
                 handle = scheduling_state.handle
 
         assert handle, "At least 1 operation should have been processed"
@@ -315,7 +330,8 @@ def _generate_node_scheduling(
         permutation = schedule.permutation[root]
         if not permutation:
             return sched_state
-
+        gpu_material = True
+        gpu_mat_thread = True
         # Materialize the loops
         for loop_name in permutation:
             # Manage the splits
@@ -351,12 +367,67 @@ def _generate_node_scheduling(
                 self._vectorize(sched_state)
                 break
             elif loop_name in tiles_sizes_by_loops:
-                self._strip_mine(
-                    loop_name=loop_name,
-                    tiling_vector=tiles_sizes_by_loops[loop_name],
-                    schedule=schedule,
-                    sched_state=sched_state,
-                )
+                if loop_name in schedule.gpu_blocks:
+                    tile_vect = [
+                        sum(values)
+                        for values in zip(
+                            *[
+                                tiles_sizes_by_loops[loop]
+                                for loop in schedule.gpu_blocks
+                            ]
+                        )
+                    ]
+                    tile_vect = tile_vect + [0] * (3 - len(tile_vect))
+                    # TODO: Do not work with splitting
+                    position_index = [
+                        permutation.index(loop) for loop in schedule.gpu_blocks
+                    ]
+                    mapping_order = sorted(
+                        range(len(position_index)), key=lambda i: position_index[i]
+                    )
+                    if gpu_material:
+                        self._strip_mine(
+                            loop_name=loop_name,
+                            tiling_vector=tile_vect,
+                            mapping_order=mapping_order,
+                            schedule=schedule,
+                            sched_state=sched_state,
+                        )
+                        gpu_material = False
+                elif loop_name in schedule.gpu_threads:
+                    tile_vect = [
+                        sum(values)
+                        for values in zip(
+                            *[
+                                tiles_sizes_by_loops[loop]
+                                for loop in schedule.gpu_threads
+                            ]
+                        )
+                    ]
+                    tile_vect = tile_vect + [0] * (3 - len(tile_vect))
+                    position_index = [
+                        permutation.index(loop) for loop in schedule.gpu_threads
+                    ]
+                    mapping_order = sorted(
+                        range(len(position_index)), key=lambda i: position_index[i]
+                    )
+                    if gpu_mat_thread:
+                        self._strip_mine(
+                            loop_name=loop_name,
+                            tiling_vector=tile_vect,
+                            mapping_order=mapping_order,
+                            schedule=schedule,
+                            sched_state=sched_state,
+                        )
+                        gpu_mat_thread = False
+                else:
+                    self._strip_mine(
+                        loop_name=loop_name,
+                        tiling_vector=tiles_sizes_by_loops[loop_name],
+                        mapping_order=[],
+                        schedule=schedule,
+                        sched_state=sched_state,
+                    )
                 if loop_name in schedule.distribution:
                     self._distribute_loop(loop_name, schedule, sched_state)
             # Fuse the producers
@@ -494,20 +565,34 @@ def _strip_mine(
         self,
         loop_name: str,
         tiling_vector: list[int],
+        mapping_order: list[int],
         schedule: MlirNodeSchedule,
         sched_state: SchedulingState,
     ) -> OpResult:
-        if loop_name in schedule.parallelization:
-            tiling_command = TileUsingForallOp(
-                sched_state.handle, tile_sizes=tiling_vector
+        attr_array = {}
+        attr_array["tile_sizes"] = tiling_vector
+        if loop_name in schedule.gpu_blocks:
+            attr_array["mapping"] = ArrayAttr.get(
+                [self._get_block_id(index) for index in mapping_order]
+            )
+            self._gpu_block_order = attr_array["mapping"]
+            tiling_command = TileUsingForallOp(sched_state.handle, **attr_array)
+        elif loop_name in schedule.gpu_threads:
+            attr_array["mapping"] = ArrayAttr.get(
+                [self._get_thread_id(index) for index in mapping_order]
             )
+            tiling_command = TileUsingForallOp(sched_state.handle, **attr_array)
+        elif loop_name in schedule.parallelization:
+            tiling_command = TileUsingForallOp(sched_state.handle, **attr_array)
         else:
             tiling_command = TileUsingForOp(sched_state.handle, sizes=tiling_vector)
         # Extract the results
         sched_state.handle = tiling_command.results[0]
         assert len(tiling_command.results) == 2
         new_loop = tiling_command.results[-1]
         sched_state.all_loops[loop_name] = new_loop
+        if loop_name in schedule.gpu_blocks:
+            loop_name = schedule.gpu_blocks[0]
         # Annotate the resulting loop if successfully generated
         transform.AnnotateOp(new_loop, loop_name)
 
@@ -554,11 +639,12 @@ def _post_vectorize(self, sched_state: SchedulingState, schedule: MlirNodeSchedu
             vector.ApplyTransferPermutationPatternsOp()
 
         # the remaining patterns must be applied post-bufferization to work properly
-        if not self._post_bufferize_sequence:
+        if not self._post_bufferize_sequence and not schedule.gpu_blocks:
             with InsertionPoint(transform.ApplyPatternsOp(parent_op).patterns):
                 vector.ApplyLowerOuterProductPatternsOp()
                 vector.ApplyLowerContractionPatternsOp()
-        else:
+        # Do not lower vector contract as it can be useful for gpu optimisation
+        elif self._post_bufferize_sequence and not schedule.gpu_blocks:
             func_name = self._mlir_program.mlir_module.body.operations[0].attributes[
                 "sym_name"
             ]
@@ -676,6 +762,87 @@ def _collect_fused_producers(self, unscheduled_handles: set[str | None]):
 
         return fused_producers
 
+    def _get_thread_id(self, index: int) -> Attribute:
+        ctx = self._mlir_program.mlir_context
+        return Attribute.parse(f"#gpu.thread<{_GPU_DIM[index]}>", context=ctx)
+
+    def _get_block_id(self, index: int) -> Attribute:
+        ctx = self._mlir_program.mlir_context
+        return Attribute.parse(f"#gpu.block<{_GPU_DIM[index]}>", context=ctx)
+
+    def _gpu_mapping(
+        self,
+        schedule: MlirNodeSchedule,
+        sched_state: SchedulingState,
+    ):
+        tiles_sizes_by_loops = self._generate_tiling_insns(schedule)
+        if schedule.gpu_blocks and not self._using_tensors:
+            new_loop = next(
+                (
+                    sched_state.all_loops[loop_name]
+                    for loop_name in schedule.gpu_blocks
+                    if loop_name in sched_state.all_loops
+                ),
+                None,
+            )
+            # Since we know there only 1 non zero number
+            # TODO Find a way to put thread number instead of putting tile size
+            new_loop = MapForallToBlocks(
+                new_loop,
+                generate_gpu_launch=True,
+            ).result
+            if schedule.gpu_threads:
+                block_dims = [
+                    max(tiles_sizes_by_loops[loop_name_block])
+                    // max(tiles_sizes_by_loops[loop_name])
+                    for loop_name, loop_name_block in zip(
+                        schedule.gpu_threads, schedule.gpu_blocks
+                    )
+                ]
+                block_dims = block_dims + [1] * (3 - len(block_dims))
+                MapNestedForallToThreads(
+                    new_loop,
+                    block_dims=block_dims,
+                )
+        elif (
+            schedule.gpu_blocks
+            and self._using_tensors
+            and self._post_bufferize_sequence
+            and self._gpu_block_order is not None
+        ):
+            with (
+                InsertionPoint.at_block_begin(self._post_bufferize_sequence.body),
+                self._mlir_program.mlir_context,
+                self._loc,
+            ):
+                gpu_block_handle = structured_match(
+                    results_=transform.AnyOpType.get(),
+                    target=self._post_bufferize_sequence.bodyTarget,
+                    op_attrs={
+                        schedule.gpu_blocks[0]: UnitAttr.get(),
+                        "mapping": self._gpu_block_order,
+                    },
+                )
+                # Since we know there only 1 non zero number
+                # TODO Find a way to put thread number instead of putting tile size
+                new_loop = MapForallToBlocks(
+                    gpu_block_handle,
+                    generate_gpu_launch=True,
+                ).result
+                if schedule.gpu_threads:
+                    block_dims = [
+                        max(tiles_sizes_by_loops[loop_name_block])
+                        // max(tiles_sizes_by_loops[loop_name])
+                        for loop_name, loop_name_block in zip(
+                            schedule.gpu_threads, schedule.gpu_blocks
+                        )
+                    ]
+                    block_dims = block_dims + [1] * (3 - len(block_dims))
+                    MapNestedForallToThreads(
+                        new_loop,
+                        block_dims=block_dims,
+                    )
+
 
 def find_producer_handles(module: Module, root_handle: str) -> list[str | None]:
     # returns the handles for each operand of the operation specified by root_handle

src/xtc/backends/mlir/MlirNodeScheduler.py

@@ -25,6 +25,8 @@ class MlirNodeSchedule:
     permutation: dict[str, list[str]]
     vectorization: list[str]
     parallelization: list[str]
+    gpu_threads: list[str]
+    gpu_blocks: list[str]
     unrolling: dict[str, int]
     packed_buffers: dict[str, list[int]]
     memory_mesh: dict[str, int]
@@ -86,6 +88,8 @@ def __init__(
         self.permutation: dict[str, list[str]] = {}
         self.vectorization: list[str] = []
         self.parallelization: list[str] = []
+        self.gpu_threads: list[str] = []
+        self.gpu_blocks: list[str] = []
         self.unrolling: dict[str, int] = {}
         self.packed_buffers: dict[str, list[int]] = {}
         self.memory_mesh: dict[str, int] = {}
@@ -113,6 +117,8 @@ def mlir_node_schedule(self) -> MlirNodeSchedule:
             permutation=self.permutation,
             vectorization=self.vectorization,
             parallelization=self.parallelization,
+            gpu_threads=self.gpu_threads,
+            gpu_blocks=self.gpu_blocks,
             unrolling=self.unrolling,
             memory_mesh=self.memory_mesh,
             packed_buffers=self.packed_buffers,
@@ -230,3 +236,13 @@ def fuse_producer_at(
         self, axis: str, input_idx: int, root: str = DEFAULT_ROOT
     ) -> None:
         self.fused.append((make_loop_name(root, axis), input_idx))
+
+    def map_gpu_threads(self, axes: list[str], root: str = DEFAULT_ROOT):
+        assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu thread"
+        assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu thread"
+        self.gpu_threads = [make_loop_name(root, axis) for axis in axes]
+
+    def map_gpu_blocks(self, axes: list[str], root: str = DEFAULT_ROOT):
+        assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu thread"
+        assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu block"
+        self.gpu_blocks = [make_loop_name(root, axis) for axis in axes]

src/xtc/backends/mlir/MlirScheduler.py

@@ -204,6 +204,14 @@ def distributed_buffer_at(
             axis, input_idx, memory_axes, root=root
         )
 
+    @override
+    def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        self._current_scheduler.map_gpu_threads(axes, root=root)
+
+    @override
+    def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        self._current_scheduler.map_gpu_blocks(axes, root=root)
+
     @override
     def get_loop_nest(self) -> LoopNest:
         node_sched = self._current_scheduler

src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py

@@ -442,18 +442,11 @@ def _lowering_pipeline(self, sm_arch: str, ptx_version: str) -> list[str]:
             "cse",
             "sccp",
             # From complex control to the soup of basic blocks
-            "expand-strided-metadata",
-            "scf-forall-to-parallel",
-            "canonicalize",
-            "cse",
-            "sccp",
-            "func.func(gpu-map-parallel-loops)",
-            "convert-parallel-loops-to-gpu",
             "convert-linalg-to-loops",
             "canonicalize",
             "cse",
             "sccp",
-            "convert-vector-to-llvm",
+            "convert-vector-to-llvm{vector-contract-lowering=outerproduct}",
             "buffer-results-to-out-params",
             "convert-func-to-llvm{use-bare-ptr-memref-call-conv=true}",
             "gpu-lower-to-nvvm-pipeline{cubin-chip="

src/xtc/backends/tvm/TVMScheduler.py

@@ -495,6 +495,16 @@ def distributed_buffer_at(
         # TODO: not implemented for now
         pass
 
+    @override
+    def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
+    @override
+    def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
     def _get_plain_schedule(self) -> TVMPlainSchedule:
         return TVMPlainSchedule(
             dims=deepcopy(self.dims),