mlir: add gpu_block and gpu_thread, gpu_lane, gpu_warp primitive

src/xtc/backends/jir/JIRScheduler.py

@@ -361,6 +361,26 @@ def distributed_buffer_at(
         # TODO: not implemented for now
         pass
 
+    @override
+    def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
+    @override
+    def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # 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/MlirNodeScheduler.py

@@ -109,6 +109,26 @@ def fuse_producer_at(
     def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None:
         self._plain_sch.fuse_consumer_at(axis, root)
 
+    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._plain_sch.gpu_thread(axes, root)
+
+    def map_gpu_blocks(self, axes: list[str], root: str = DEFAULT_ROOT):
+        assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu block"
+        assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu block"
+        self._plain_sch.gpu_block(axes, root)
+
+    def map_gpu_lanes(self, axes: list[str], root: str = DEFAULT_ROOT):
+        assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu lane"
+        assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu lane"
+        self._plain_sch.gpu_lane(axes, root)
+
+    def map_gpu_warps(self, axes: list[str], root: str = DEFAULT_ROOT):
+        assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu warp"
+        assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu warp"
+        self._plain_sch.gpu_warp(axes, root)
+
     def get_node_schedule(self) -> MlirNodeSchedule:
         plain_schedule = self._plain_sch.get_plain_schedule()
         return MlirNodeSchedule(**asdict(plain_schedule))

src/xtc/backends/mlir/MlirScheduler.py

@@ -218,6 +218,22 @@ def distributed_buffer_at(
             axis, input_idx, memory_axes, root=root
         )
 
+    @override
+    def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        self._current_scheduler.map_gpu_lanes(axes, root=root)
+
+    @override
+    def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        self._current_scheduler.map_gpu_warps(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_schedule = self._current_scheduler.get_node_schedule()

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

@@ -442,25 +442,48 @@ 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",
             "buffer-results-to-out-params",
+            # GPU to LLVM pipeline
+            "func.func(gpu-eliminate-barriers)",
+            "convert-nvgpu-to-nvvm",
+            "gpu-kernel-outlining",
+            "gpu-launch-sink-index-computations",
+            "convert-vector-to-scf",
+            "convert-vector-to-llvm{vector-contract-lowering=outerproduct}",
+            "nvgpu-optimize-shared-memory",
+            "convert-scf-to-cf",
+            "convert-nvvm-to-llvm",
             "convert-func-to-llvm{use-bare-ptr-memref-call-conv=true}",
-            "gpu-lower-to-nvvm-pipeline{cubin-chip="
+            "expand-strided-metadata",
+            "nvvm-attach-target{chip="
             + sm_arch
-            + " cubin-features=+ptx"
+            + " features=+ptx"
             + "".join(ptx_version.split("."))
-            + " opt-level=3}",
+            + " O=3 fast=true}",
+            "lower-affine",
+            "convert-arith-to-llvm",
+            "convert-index-to-llvm",
+            "canonicalize",
+            "cse",
+            "reconcile-unrealized-casts",
+            "gpu.module(convert-gpu-to-nvvm{use-bare-ptr-memref-call-conv=true})",
+            "gpu.module(canonicalize)",
+            "gpu.module(cse)",
+            "gpu.module(reconcile-unrealized-casts)",
+            "gpu-to-llvm{use-bare-pointers-for-host=true use-bare-pointers-for-kernels=true}",
+            "gpu.module(reconcile-unrealized-casts)",
+            "reconcile-unrealized-casts",
+            "convert-math-to-llvm",
+            "gpu.module(reconcile-unrealized-casts)",
+            "reconcile-unrealized-casts",
+            "gpu-module-to-binary",
+            "canonicalize",
+            "cse",
+            "reconcile-unrealized-casts",
         ]
 
     def run(self, sm_arch: str, ptx_version: str) -> None:

src/xtc/backends/tvm/TVMScheduler.py

@@ -378,6 +378,26 @@ def distributed_buffer_at(
         # 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
+
+    @override
+    def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
+    @override
+    def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
+    @override
+    def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
     def _get_plain_schedule(self) -> PlainNodeSchedule:
         return self._plain_sch.get_plain_schedule()
 

src/xtc/itf/schd/scheduler.py

@@ -319,6 +319,77 @@ def distributed_buffer_at(
         """
         ...
 
+    @abstractmethod
+    def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        """Maps iteration axes to GPU lane dimensions.
+
+        The mapping is positional:
+            - axes[0] -> linear_dim_0
+            - axes[1] -> linear_dim_1
+            - axes[2] -> linear_dim_2
+        At most 3 axes may be provided. If there are less than 3 axes,
+        only the corresponding lane dimensions are assigned.
+        By default lane are mapped on 32 id.
+
+        Args:
+            axes: list of maximum 3 element that are maps to the dimension x, y, z respectively
+            root: the parent split (or the operator's absolute root)
+        """
+        ...
+
+    @abstractmethod
+    def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        """Maps iteration axes to GPU warp dimensions.
+
+        The mapping is positional:
+            - axes[0] -> warp for threadIdx.x
+            - axes[1] -> warp for threadIdx.y
+            - axes[2] -> warp for threadIdx.z
+        At most 3 axes may be provided. If there are less than 3 axes,
+        only the corresponding warp dimensions are assigned.
+        It is preferable to map warp only on the x axis.
+        The thread size of the warp should be a multiple of 32.
+
+        Args:
+            axes: list of maximum 3 element that are maps to the dimension x, y, z respectively
+            root: the parent split (or the operator's absolute root)
+        """
+        ...
+
+    @abstractmethod
+    def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        """Maps iteration axes to GPU thread dimensions.
+
+        The mapping is positional:
+            - axes[0] -> threadIdx.x
+            - axes[1] -> threadIdx.y
+            - axes[2] -> threadIdx.z
+        At most 3 axes may be provided. If there are less than 3 axes,
+        only the corresponding thread dimensions are assigned.
+
+        Args:
+            axes: list of maximum 3 element that are maps to the dimension x, y, z respectively
+            root: the parent split (or the operator's absolute root)
+        """
+        ...
+
+    @abstractmethod
+    def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
+        """
+        Maps iteration axes to GPU block dimensions.
+        The mapping is positional:
+            - axes[0] -> blockIdx.x
+            - axes[1] -> blockIdx.y
+            - axes[2] -> blockIdx.z
+        At most 3 axes may be provided. If there are less than 3 axes,
+        only the corresponding block dimensions are assigned.
+
+        Args:
+            axes: list of maximum 3 element that are maps to the dimension x, y, z respectively
+            root: the parent split (or the operator's absolute root)
+        """
+        ...
+
     @abstractmethod
     def get_loop_nest(self) -> LoopNest:
         """Return a LoopNest representation of the current schedule.

src/xtc/schedules/descript.py

@@ -401,6 +401,18 @@ def _apply_annotations(
         if annotations.fuse_consumer:
             node.fuse_consumer_at.append(loop_name)
 
+        if annotations.gpu_lane is not None:
+            node.gpu_lane[loop_name] = annotations.gpu_lane
+
+        if annotations.gpu_warp is not None:
+            node.gpu_warp[loop_name] = annotations.gpu_warp
+
+        if annotations.gpu_block is not None:
+            node.gpu_block[loop_name] = annotations.gpu_block
+
+        if annotations.gpu_thread is not None:
+            node.gpu_thread[loop_name] = annotations.gpu_thread
+
     def _check_splitting_intervals(
         self,
         item: SplitDecl,
@@ -525,7 +537,6 @@ def _apply_loop_nest(self, loop_nest: LoopNest, scheduler: Scheduler) -> None:
     def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None:
         """Recursively apply a LoopNestNode and its children to the scheduler."""
         root = node.root
-
         for d, s in node.splits.items():
             scheduler.split(d, s, root=root)
 
@@ -548,6 +559,35 @@ def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None:
 
         for axis in node.fuse_consumer_at:
             scheduler.fuse_consumer_at(axis, root=root)
+
+        if node.gpu_lane:
+            sorted_keys = sorted(
+                (k for k, v in node.gpu_lane.items() if v is not None),
+                key=lambda k: node.gpu_lane[k],
+            )
+            scheduler.gpu_lane(sorted_keys, root=root)
+
+        if node.gpu_warp:
+            sorted_keys = sorted(
+                (k for k, v in node.gpu_warp.items() if v is not None),
+                key=lambda k: node.gpu_warp[k],
+            )
+            scheduler.gpu_thread(sorted_keys, root=root)
+
+        if node.gpu_block:
+            sorted_keys = sorted(
+                (k for k, v in node.gpu_block.items() if v is not None),
+                key=lambda k: node.gpu_block[k],
+            )
+            scheduler.gpu_block(sorted_keys, root=root)
+
+        if node.gpu_thread:
+            sorted_keys = sorted(
+                (k for k, v in node.gpu_thread.items() if v is not None),
+                key=lambda k: node.gpu_thread[k],
+            )
+            scheduler.gpu_thread(sorted_keys, root=root)
+
         # Recursively apply children
         for child in node.children:
             self._apply_node(child, scheduler)

src/xtc/schedules/loop_nest.py

@@ -118,6 +118,10 @@ class LoopNestNode(Node["LoopNestNode"]):
     pack_at: dict[str, tuple[int, str | None, bool]] = field(default_factory=dict)
     fuse_producer_at: dict[str, int] = field(default_factory=dict)
     fuse_consumer_at: list[str] = field(default_factory=list)
+    gpu_lane: dict[str, int] = field(default_factory=dict)
+    gpu_warp: dict[str, int] = field(default_factory=dict)
+    gpu_block: dict[str, int] = field(default_factory=dict)
+    gpu_thread: dict[str, int] = field(default_factory=dict)
 
     def pretty_print(self, indent: int = 0) -> str:
         """Return a human-readable representation of the loop nest.
@@ -245,6 +249,14 @@ def _add_annotations(self, line: str, loop_name: str) -> str:
             annotations.append(f"fuse_producer({prod_idx})")
         if loop_name in self.fuse_consumer_at:
             annotations.append("fuse_consumer")
+        if loop_name in self.gpu_lane:
+            annotations.append(f"gpu_lane({self.gpu_lane[loop_name]})")
+        if loop_name in self.gpu_warp:
+            annotations.append(f"gpu_warp({self.gpu_warp[loop_name]})")
+        if loop_name in self.gpu_block:
+            annotations.append(f"gpu_block({self.gpu_block[loop_name]})")
+        if loop_name in self.gpu_thread:
+            annotations.append(f"gpu_thread({self.gpu_thread[loop_name]})")
         if annotations:
             line += "  // " + ", ".join(annotations)
         return line

src/xtc/schedules/parameter_loop_nest.py

@@ -111,6 +111,12 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]):
         fuse_producer_at: Producer fusion configuration per axis. Maps axis
             names to producer indices.
         fuse_consumer_at: List of axes where the output consumer is fused.
+        gpu_block: Maps loops to block id, that loop need to be parallelize
+        gpu_thread: Maps loops to thread id, that loop need to be parallelize
+        gpu_lane: Maps loops to lane id
+        gpu_warp: Maps loops to warp id
+        gpu_block: Maps loops to block id
+        gpu_thread: Maps loops to thread id
     """
 
     root: str
@@ -127,6 +133,10 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]):
     fuse_producer_at: dict[str, int] = field(default_factory=dict)
     fuse_consumer_at: list[str] = field(default_factory=list)
     constraints: list[str] = field(default_factory=list)
+    gpu_lane: dict[str, int] = field(default_factory=dict)
+    gpu_warp: dict[str, int] = field(default_factory=dict)
+    gpu_block: dict[str, int] = field(default_factory=dict)
+    gpu_thread: dict[str, int] = field(default_factory=dict)
 
     def apply_sample(self, sample: dict[str, int]) -> LoopNestNode:
         """
@@ -174,6 +184,10 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode:
             if self.split_origin is not None
             else None
         )
+        gpu_warp = self.gpu_warp
+        gpu_lane = self.gpu_lane
+        gpu_block = self.gpu_block
+        gpu_thread = self.gpu_thread
         return LoopNestNode(
             root=root,
             tiles=tiles,
@@ -188,6 +202,10 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode:
             fuse_consumer_at=fuse_consumer_at,
             children=children,
             split_origin=split_origin,
+            gpu_lane=gpu_lane,
+            gpu_warp=gpu_warp,
+            gpu_block=gpu_block,
+            gpu_thread=gpu_thread,
         )
 
     def pretty_print(self, indent: int = 0) -> str:
@@ -316,6 +334,14 @@ def _add_annotations(self, line: str, loop_name: str) -> str:
             annotations.append(f"fuse_producer({prod_idx})")
         if loop_name in self.fuse_consumer_at:
             annotations.append("fuse_consumer")
+        if loop_name in self.gpu_lane:
+            annotations.append(f"gpu_lane({self.gpu_lane[loop_name]})")
+        if loop_name in self.gpu_warp:
+            annotations.append(f"gpu_warp({self.gpu_warp[loop_name]})")
+        if loop_name in self.gpu_block:
+            annotations.append(f"gpu_block({self.gpu_block[loop_name]})")
+        if loop_name in self.gpu_thread:
+            annotations.append(f"gpu_thread({self.gpu_thread[loop_name]})")
         if annotations:
             line += "  // " + ", ".join(annotations)
         return line