schedule: Add more restriction to descript gpu

src/xtc/schedules/loop_nest.py

@@ -386,7 +386,7 @@ def check(self):
         self._check_external_consistency()
         self._check_tiling_consistency(info)
         self._check_sizes(info)
-        self._check_gpu_consistency()
+        self._check_gpu_consistency(info)
 
     def _check_use_defined_dims(self, info: LoopInfo):
         for dim in self.abstract_dims:
@@ -507,7 +507,7 @@ def _check_sizes(self, info: LoopInfo):
                             f'`{{"unroll" = {unroll_factor}}}`: unroll factor should be smaller than {loop_size}.'
                         )
 
-    def _check_gpu_consistency(self) -> None:
+    def _check_gpu_consistency(self, info: LoopInfo) -> None:
         for sched in self.nodes:
             gpu_sets = {
                 "gpu_block": set(sched.gpu_block.keys()),
@@ -529,10 +529,28 @@ def _check_gpu_consistency(self) -> None:
             has_thread_or_lane_or_warp = (
                 bool(sched.gpu_thread) or bool(sched.gpu_lane) or bool(sched.gpu_warp)
             )
-            if has_block and not has_thread_or_lane_or_warp:
+            if not has_block and has_thread_or_lane_or_warp:
                 raise ScheduleValidationError(
-                    "gpu_block requires either gpu_thread or gpu_lane or gpu_warp to be specified."
+                    "Need gpu_block to be specified for either gpu_thread or gpu_lane or gpu_warp."
                 )
+            gpu_block_set = gpu_sets["gpu_block"]
+            if not all(prim in info.dims for prim in gpu_block_set):
+                raise ScheduleValidationError(
+                    "Need gpu_block to be an axis and not a tile"
+                )
+
+            for gpu_name in primitive_names[1:]:
+                if not all(
+                    info.tiles_to_axis.get(prim, None) is not None
+                    for prim in gpu_sets[gpu_name]
+                ):
+                    raise ScheduleValidationError(f"{gpu_name} need to be a tile")
+            # We need to check if there is a loop above the generated kernel
+            for prim in gpu_block_set:
+                if sched.interchange.index(prim) >= len(gpu_block_set):
+                    raise ScheduleValidationError(
+                        "gpu_block needs to be in the most outermost loop"
+                    )
 
     @staticmethod
     def _must_be_smaller_routine(

src/xtc/schedules/parameter_loop_nest.py

@@ -473,7 +473,7 @@ def check(self):
         self._check_vectorization_consistency()
         self._check_tiling_consistency(info)
         self._check_sizes(info)
-        self._check_gpu_consistency()
+        self._check_gpu_consistency(info)
 
     def apply_sample(self, sample: dict[str, int]) -> LoopNest:
         """
@@ -588,7 +588,7 @@ def _check_sizes(self, info: ParameterLoopInfo):
                             f'`{{"unroll" = {unroll_factor}}}`: unroll factor should be smaller than {loop_size}.'
                         )
 
-    def _check_gpu_consistency(self) -> None:
+    def _check_gpu_consistency(self, info: ParameterLoopInfo) -> None:
         for sched in self.nodes:
             gpu_sets = {
                 "gpu_block": set(sched.gpu_block.keys()),
@@ -610,10 +610,32 @@ def _check_gpu_consistency(self) -> None:
             has_thread_or_lane_or_warp = (
                 bool(sched.gpu_thread) or bool(sched.gpu_lane) or bool(sched.gpu_warp)
             )
-            if has_block and not has_thread_or_lane_or_warp:
+            if not has_block and has_thread_or_lane_or_warp:
                 raise ScheduleValidationError(
-                    "gpu_block requires either gpu_thread or gpu_lane or gpu_warp to be specified."
+                    "Need gpu_block to be specified for either gpu_thread or gpu_lane or gpu_warp."
                 )
+            if not has_block and has_thread_or_lane_or_warp:
+                raise ScheduleValidationError(
+                    "Need gpu_block to be specified for either gpu_thread or gpu_lane or gpu_warp."
+                )
+            gpu_block_set = gpu_sets["gpu_block"]
+            if not all(prim in info.dims for prim in gpu_block_set):
+                raise ScheduleValidationError(
+                    "Need gpu_block to be an axis and not a tile"
+                )
+
+            for gpu_name in primitive_names[1:]:
+                if not all(
+                    info.tiles_to_axis.get(prim, None) is not None
+                    for prim in gpu_sets[gpu_name]
+                ):
+                    raise ScheduleValidationError(f"{gpu_name} need to be a tile")
+            # We need to check if there is a loop above the generated kernel
+            for prim in gpu_block_set:
+                if sched.interchange.index(prim) >= len(gpu_block_set):
+                    raise ScheduleValidationError(
+                        "gpu_block needs to be in the most outermost loop"
+                    )
 
     @staticmethod
     def _must_be_smaller_routine(

tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py

@@ -72,15 +72,20 @@
 # CHECK-NEXT:      transform.annotate %loops_7 "./k1" : !transform.any_op
 # CHECK-NEXT:      %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 # CHECK-NEXT:      transform.annotate %loops_9 "./i2" : !transform.any_op
-# CHECK-NEXT:      transform.include @_vecto failures(suppress) (%tiled_linalg_op_8) : (!transform.any_op) -> ()
-# CHECK-NEXT:      transform.loop.unroll %loops_9 {factor = 2 : i64} : !transform.any_op
-# CHECK-NEXT:      %2 = transform.get_parent_op %forall_op {isolated_from_above} : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      %2 = transform.get_parent_op %tiled_linalg_op_8 : (!transform.any_op) -> !transform.any_op
 # CHECK-NEXT:      transform.apply_patterns to %2 {
+# CHECK-NEXT:        transform.apply_patterns.xtc.fold_unit_extent_dims_via_slices_for_vectorization
+# CHECK-NEXT:      } : !transform.any_op
+# CHECK-NEXT:      %3 = transform.structured.match interface{LinalgOp} in %2 : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      transform.include @_vecto failures(suppress) (%3) : (!transform.any_op) -> ()
+# CHECK-NEXT:      transform.loop.unroll %loops_9 {factor = 2 : i64} : !transform.any_op
+# CHECK-NEXT:      %4 = transform.get_parent_op %forall_op {isolated_from_above} : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      transform.apply_patterns to %4 {
 # CHECK-NEXT:        transform.apply_patterns.vector.reduction_to_contract
 # CHECK-NEXT:        transform.apply_patterns.vector.transfer_permutation_patterns
 # CHECK-NEXT:      } : !transform.any_op
-# CHECK-NEXT:      %3 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op
-# CHECK-NEXT:      %4 = transform.gpu.map_nested_forall_to_threads %3 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      %5 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      %6 = transform.gpu.map_nested_forall_to_threads %5 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op
 # CHECK-NEXT:      transform.yield 
 # CHECK-NEXT:    }
 # CHECK-NEXT:  }

tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py

@@ -71,14 +71,19 @@
 # CHECK-NEXT:      transform.annotate %forall_op_5 "./j2" : !transform.any_op
 # CHECK-NEXT:      %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_op_4 tile_sizes [0, 0, 16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 # CHECK-NEXT:      transform.annotate %loops_7 "./k" : !transform.any_op
-# CHECK-NEXT:      transform.include @_vecto failures(suppress) (%tiled_linalg_op_6) : (!transform.any_op) -> ()
-# CHECK-NEXT:      %2 = transform.get_parent_op %forall_op {isolated_from_above} : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      %2 = transform.get_parent_op %tiled_linalg_op_6 : (!transform.any_op) -> !transform.any_op
 # CHECK-NEXT:      transform.apply_patterns to %2 {
+# CHECK-NEXT:        transform.apply_patterns.xtc.fold_unit_extent_dims_via_slices_for_vectorization
+# CHECK-NEXT:      } : !transform.any_op
+# CHECK-NEXT:      %3 = transform.structured.match interface{LinalgOp} in %2 : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      transform.include @_vecto failures(suppress) (%3) : (!transform.any_op) -> ()
+# CHECK-NEXT:      %4 = transform.get_parent_op %forall_op {isolated_from_above} : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      transform.apply_patterns to %4 {
 # CHECK-NEXT:        transform.apply_patterns.vector.reduction_to_contract
 # CHECK-NEXT:        transform.apply_patterns.vector.transfer_permutation_patterns
 # CHECK-NEXT:      } : !transform.any_op
-# CHECK-NEXT:      %3 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op
-# CHECK-NEXT:      %4 = transform.gpu.map_nested_forall_to_threads %3 block_dims = [64, 1, 1] : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      %5 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op
+# CHECK-NEXT:      %6 = transform.gpu.map_nested_forall_to_threads %5 block_dims = [64, 1, 1] : (!transform.any_op) -> !transform.any_op
 # CHECK-NEXT:      transform.yield 
 # CHECK-NEXT:    }
 # CHECK-NEXT:  }

tests/filecheck/schedules/test_descript_gpu_validation.py

@@ -0,0 +1,96 @@
+# REQUIRES: mlir-target=nvgpu
+# RUN: not python %s --overlap 2>&1 | filecheck %s --check-prefix=CHECK-OVERLAP
+# RUN: not python %s --no-block 2>&1 | filecheck %s --check-prefix=CHECK-NO-BLOCK
+# RUN: not python %s --block-tile 2>&1 | filecheck %s --check-prefix=CHECK-BLOCK-TILE
+# RUN: not python %s --thread-axis 2>&1 | filecheck %s --check-prefix=CHECK-THREAD-AXIS
+# RUN: not python %s --block-not-outermost 2>&1 | filecheck %s --check-prefix=CHECK-OUTERMOST
+
+import sys
+import xtc.graphs.xtc.op as O
+from xtc.backends.mlir import Backend
+from xtc.schedules.descript import descript_scheduler
+
+I, J, K, dtype = 16, 32, 64, "float32"
+a = O.tensor((I, K), dtype, name="A")
+b = O.tensor((K, J), dtype, name="B")
+
+with O.graph(name="matmul") as gb:
+    O.matmul(a, b, name="C")
+
+graph = gb.graph
+
+
+def make_descript_scheduler(spec):
+    impl = Backend(graph)
+    sch = impl.get_scheduler()
+    descript_scheduler(
+        scheduler=sch,
+        node_name="C",
+        abstract_dims=["I", "J", "K"],
+        spec=spec,
+    )
+    print("ok")
+
+if "--overlap" in sys.argv:
+    # The same loop mapped to two gpu primitives
+    make_descript_scheduler(
+        {
+            "I": {"gpu_block": 0, "gpu_thread": 0},
+            "J": {},
+            "K": {},
+        }
+    )
+
+# CHECK-OVERLAP: Loops I appear in both gpu_block and gpu_thread.
+
+elif "--no-block" in sys.argv:
+    # Threads need at least one block-mapped loop
+    make_descript_scheduler(
+        {
+            "I": {"gpu_thread": 0},
+            "J": {},
+            "K": {},
+        }
+    )
+
+# CHECK-NO-BLOCK: Need gpu_block to be specified for either gpu_thread or gpu_lane or gpu_warp.
+
+elif "--block-tile" in sys.argv:
+    # Blocks need to be mapped to base axes, not tiles
+    make_descript_scheduler(
+        {
+            "I": {},
+            "I#4": {"gpu_block": 0},
+            "J": {},
+            "J#8": {"gpu_thread": 0},
+            "K": {},
+        }
+    )
+
+# CHECK-BLOCK-TILE: Need gpu_block to be an axis and not a tile
+
+elif "--thread-axis" in sys.argv:
+    # Threads need to be mapped to tiles, not base axes
+    make_descript_scheduler(
+        {
+            "I": {"gpu_block": 0},
+            "J": {"gpu_thread": 0},
+            "K": {},
+        }
+    )
+
+# CHECK-THREAD-AXIS: gpu_thread need to be a tile
+
+elif "--block-not-outermost" in sys.argv:
+    # The block-mapped loops need to come first in the loop order
+    make_descript_scheduler(
+        {
+            "K": {},
+            "I": {"gpu_block": 0},
+            "J": {"gpu_block": 1},
+            "I#4": {"gpu_thread": 0},
+            "J#8": {"gpu_thread": 1},
+        }
+    )
+
+# CHECK-OUTERMOST: gpu_block needs to be in the most outermost loop