From eb5b583bca0cb38ec111dff9fdd290d5e323ac45 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Tue, 6 Oct 2026 15:24:26 +0200 Subject: [PATCH] schedule: Add more restriction to descript gpu - Add some gpu validation test - Regen some test --- src/xtc/schedules/loop_nest.py | 26 ++++- src/xtc/schedules/parameter_loop_nest.py | 30 +++++- ...st_matmul_mlir_offload_tensor_vectorise.py | 15 ++- ...matmul_mlir_offload_warp_lane_vectorise.py | 13 ++- .../schedules/test_descript_gpu_validation.py | 96 +++++++++++++++++++ 5 files changed, 163 insertions(+), 17 deletions(-) create mode 100644 tests/filecheck/schedules/test_descript_gpu_validation.py diff --git a/src/xtc/schedules/loop_nest.py b/src/xtc/schedules/loop_nest.py index b13d448c..95a6d0df 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/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( diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index 85a65799..d554741c 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/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( diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py index 9ce23e87..f4320450 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py +++ b/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: } diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py index 41d00861..f9202075 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py +++ b/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: } diff --git a/tests/filecheck/schedules/test_descript_gpu_validation.py b/tests/filecheck/schedules/test_descript_gpu_validation.py new file mode 100644 index 00000000..3aeea2dd --- /dev/null +++ b/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