From 04a0746fdb6f1bebbc33dfdf0dfe126ebf66adac Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 11 May 2026 11:37:14 +0200 Subject: [PATCH 1/7] mlir: Add 2 primitive for gpu on the scheduler - Usage of TileForAll with gpu mapping - Fusion of gpu mapping for correct IR --- src/xtc/backends/jir/JIRScheduler.py | 10 ++ src/xtc/backends/mlir/MlirCompilerPasses.py | 122 +++++++++++-- src/xtc/backends/mlir/MlirNodeScheduler.py | 10 ++ src/xtc/backends/mlir/MlirScheduler.py | 8 + .../mlir/MlirTarget/MlirNVGPUTarget.py | 1 - src/xtc/itf/schd/scheduler.py | 18 ++ src/xtc/schedules/descript.py | 17 ++ src/xtc/schedules/loop_nest.py | 6 + src/xtc/schedules/parsing.py | 46 +++++ src/xtc/schedules/plain_schedule.py | 16 ++ .../test_matmul_mlir_offload_tensor.py | 15 +- .../schedules/test_matmul_descript_gpu.py | 167 ++++++++++++++++++ tests/filecheck/search/test_conv_oo.py | 10 +- tests/filecheck/search/test_conv_pprprp.py | 10 +- tests/filecheck/search/test_conv_pprprpv.py | 10 +- tests/filecheck/search/test_conv_pprprpvr.py | 10 +- .../search/test_conv_pprprpvr_rnd.py | 2 +- tests/filecheck/search/test_conv_prp.py | 10 +- tests/filecheck/search/test_conv_prp_rnd.py | 2 +- tests/filecheck/search/test_matmul_goto.py | 10 +- tests/filecheck/search/test_matmul_goto_r.py | 10 +- tests/filecheck/search/test_matmul_oo.py | 10 +- tests/filecheck/search/test_matmul_p1.py | 10 +- tests/filecheck/search/test_matmul_p1v.py | 10 +- tests/filecheck/search/test_matmul_pprprp.py | 10 +- tests/filecheck/search/test_matmul_pprprpv.py | 10 +- .../filecheck/search/test_matmul_pprprpvr.py | 10 +- .../search/test_matmul_pprprpvr_rnd.py | 2 +- tests/filecheck/search/test_matmul_prp.py | 10 +- tests/filecheck/search/test_matmul_prp_rnd.py | 2 +- 30 files changed, 494 insertions(+), 90 deletions(-) create mode 100644 tests/filecheck/schedules/test_matmul_descript_gpu.py diff --git a/src/xtc/backends/jir/JIRScheduler.py b/src/xtc/backends/jir/JIRScheduler.py index 4634bc9d6..037caec4c 100644 --- a/src/xtc/backends/jir/JIRScheduler.py +++ b/src/xtc/backends/jir/JIRScheduler.py @@ -361,6 +361,16 @@ def distributed_buffer_at( # TODO: not implemented for now pass + @override + def gpu_thread(self, axes: list[str]) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_block(self, axes: list[str]) -> None: + # TODO: not implemented for now + pass + def get_schedule_str(self) -> str: return str(JIRSchedule(scheduler=self)) diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index ae4ed40fd..cf8b473ef 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -23,6 +23,10 @@ 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 ( @@ -30,6 +34,8 @@ InsertionPoint, UnitAttr, OpResult, + Attribute, + ArrayAttr, ) from mlir.passmanager import PassManager from mlir.ir import Module @@ -53,6 +59,7 @@ _VECTO_SEQ_NAME = "_vecto" _SUPER_VECTORIZE_SEQ_NAME = "_super_vectorize" _POST_BUFFERIZE_SEQ_NAME = "_post_bufferize" +_GPU_DIM = ["x", "y", "z"] @dataclass @@ -321,7 +328,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 @@ -357,12 +365,39 @@ def _generate_node_scheduling( self._vectorize(sched_state, self._vector_sizes_for(schedule)) 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 = [ + max(tiles_sizes_by_loops[loop]) for loop in schedule.gpu_blocks + ] + tile_vect = tile_vect + [0] * (3 - len(tile_vect)) + if gpu_material: + new_loop = self._strip_mine( + loop_name=loop_name, + tiling_vector=tile_vect, + schedule=schedule, + sched_state=sched_state, + ) + gpu_material = False + elif loop_name in schedule.gpu_threads: + tile_vect = [ + max(tiles_sizes_by_loops[loop]) for loop in schedule.gpu_threads + ] + tile_vect = tile_vect + [0] * (3 - len(tile_vect)) + if gpu_mat_thread: + new_loop = self._strip_mine( + loop_name=loop_name, + tiling_vector=tile_vect, + 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], + schedule=schedule, + sched_state=sched_state, + ) if loop_name in schedule.distribution: self._distribute_loop(loop_name, schedule, sched_state) # Fuse the producers @@ -379,6 +414,10 @@ def _generate_node_scheduling( if schedule.unrolling: self._unroll(permutation, schedule, sched_state) + # Gpu mapping + if schedule.gpu_blocks: + self._gpu_mapping(schedule, sched_state, tiles_sizes_by_loops) + return sched_state def _fuse_producers_into_loop( @@ -498,9 +537,26 @@ def _strip_mine( sched_state: SchedulingState, ) -> OpResult: if loop_name in schedule.parallelization: - tiling_command = TileUsingForallOp( - sched_state.handle, tile_sizes=tiling_vector - ) + attr_array = {} + if loop_name in schedule.gpu_threads: + attr_array["mapping"] = ArrayAttr.get( + [ + self._get_thread_id(index) + for index in range(len(schedule.gpu_threads)) + ] + ) + attr_array["tile_sizes"] = tiling_vector + elif loop_name in schedule.gpu_blocks: + attr_array["mapping"] = ArrayAttr.get( + [ + self._get_block_id(index) + for index in range(len(schedule.gpu_blocks)) + ] + ) + attr_array["tile_sizes"] = tiling_vector + else: + attr_array["tile_sizes"] = tiling_vector + tiling_command = TileUsingForallOp(sched_state.handle, **attr_array) else: tiling_command = TileUsingForOp(sched_state.handle, sizes=tiling_vector) # Extract the results @@ -728,6 +784,52 @@ def find_producer_handles(module: Module, root_handle: str) -> list[str | None]: producer_handles[-1] = attr return producer_handles + 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: dict[str, list[int]], + ): + if schedule.gpu_blocks: + 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 block id and thread id divide by the size + grid_dims = [ + max(tiles_sizes_by_loops[loop_name]) + for loop_name in schedule.gpu_blocks + ] + grid_dims = grid_dims + [1] * (3 - len(grid_dims)) + new_loop = MapForallToBlocks( + new_loop, + grid_dims=grid_dims, + generate_gpu_launch=True, + ).result + if schedule.gpu_threads: + block_dims = [ + max(tiles_sizes_by_loops[loop_name]) + for loop_name in schedule.gpu_threads + ] + block_dims = block_dims + [1] * (3 - len(block_dims)) + new_loop = MapNestedForallToThreads( + new_loop, + block_dims=block_dims, + ).result + class MlirProgramApplyTransformPass: def __init__( diff --git a/src/xtc/backends/mlir/MlirNodeScheduler.py b/src/xtc/backends/mlir/MlirNodeScheduler.py index a87c20cdb..7bd7110a4 100644 --- a/src/xtc/backends/mlir/MlirNodeScheduler.py +++ b/src/xtc/backends/mlir/MlirNodeScheduler.py @@ -109,6 +109,16 @@ 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_threads(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 thread" + assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu block" + self._plain_sch.gpu_blocks(axes, root) + def get_node_schedule(self) -> MlirNodeSchedule: plain_schedule = self._plain_sch.get_plain_schedule() return MlirNodeSchedule(**asdict(plain_schedule)) diff --git a/src/xtc/backends/mlir/MlirScheduler.py b/src/xtc/backends/mlir/MlirScheduler.py index c3b25ff2a..5943821db 100644 --- a/src/xtc/backends/mlir/MlirScheduler.py +++ b/src/xtc/backends/mlir/MlirScheduler.py @@ -218,6 +218,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_schedule = self._current_scheduler.get_node_schedule() diff --git a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py index 8173c0840..99651f850 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py +++ b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py @@ -446,7 +446,6 @@ def _lowering_pipeline(self, sm_arch: str, ptx_version: str) -> list[str]: "scf-forall-to-parallel", "canonicalize", "cse", - "sccp", "func.func(gpu-map-parallel-loops)", "convert-parallel-loops-to-gpu", "convert-linalg-to-loops", diff --git a/src/xtc/itf/schd/scheduler.py b/src/xtc/itf/schd/scheduler.py index 4c2f0a14c..541961051 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -319,6 +319,24 @@ def distributed_buffer_at( """ ... + @abstractmethod + def gpu_thread(self, axes: list[str]) -> None: + """ + Map an axis to a certain gpu thread either x, y or z. + The mapped axis need to be in parallelization. + We can only map 3 dimensions for gpu thread + """ + ... + + @abstractmethod + def gpu_block(self, axes: list[str]) -> None: + """ + Map an axis to a certain gpu block either x, y or z. + The mapped axis need to be in parallelization. + We can only map 3 dimensions for gpu block + """ + ... + @abstractmethod def get_loop_nest(self) -> LoopNest: """Return a LoopNest representation of the current schedule. diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 99a94c8b3..51fc208d2 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -401,6 +401,12 @@ def _apply_annotations( if annotations.fuse_consumer: node.fuse_consumer_at.append(loop_name) + 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, @@ -548,6 +554,17 @@ 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_block: + self.scheduler.gpu_block( + sorted(node.gpu_block, key=node.gpu_block.get), root=root + ) + + if node.gpu_thread: + self.scheduler.gpu_thread( + sorted(node.gpu_thread, key=node.gpu_thread.get), root=root + ) + # Recursively apply children for child in node.children: self._apply_node(child, scheduler) diff --git a/src/xtc/schedules/loop_nest.py b/src/xtc/schedules/loop_nest.py index bee0f67e0..29ddcacd9 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/src/xtc/schedules/loop_nest.py @@ -118,6 +118,8 @@ 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_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 +247,10 @@ 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_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 diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 3fe10568d..d6cb1219b 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -54,6 +54,8 @@ class Annotations: fuse_consumer: bool | None = False partial: bool = False full: bool = False + gpu_block: int | None = None + gpu_thread: int | None = None @dataclass(frozen=True) @@ -181,6 +183,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations fuse_consumer: bool = False partial = False full = False + gpu_block: str | int | None = None + gpu_thread: str | int | None = None for key, param in value.items(): match key: @@ -243,6 +247,46 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations partial = True case "full": full = True + case "gpu_block": + if isinstance(param, str): + if param == "x": + gpu_block = 0 + elif param == "y": + gpu_block = 1 + elif param == "z": + gpu_block = 2 + else: + raise ScheduleParseError( + f'`{{"gpu_block" = {param}}}`: gpu_block parameter should be a string or int' + ) + elif isinstance(param, int): + gpu_block = param + else: + raise ScheduleParseError( + f'`{{"gpu_block" = {param}}}`: gpu_block parameter should be a string or int' + ) + case "gpu_thread": + if isinstance(param, str): + if param == "x": + gpu_thread = 0 + elif param == "y": + gpu_thread = 1 + elif param == "z": + gpu_thread = 2 + else: + raise ScheduleParseError( + f'`{{"gpu_thread" = {param}}}`: gpu_thread string parameter should x, y or z' + ) + elif isinstance(param, int): + gpu_thread = param + if param < 0 and param < 3: + raise ScheduleParseError( + f'`{{"gpu_thread" = {param}}}`: gpu_thread int parameter should 0, 1 or 2' + ) + else: + raise ScheduleParseError( + f'`{{"gpu_thread" = {param}}}`: gpu_thread parameter should be a string or int' + ) case _: raise ScheduleParseError(f"Unknown annotation on {context}: {key}") @@ -262,6 +306,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations fuse_consumer=fuse_consumer, partial=partial, full=full, + gpu_block=gpu_block, + gpu_thread=gpu_thread, ) def _parse_pack_param( diff --git a/src/xtc/schedules/plain_schedule.py b/src/xtc/schedules/plain_schedule.py index 4514fd453..18373bc63 100644 --- a/src/xtc/schedules/plain_schedule.py +++ b/src/xtc/schedules/plain_schedule.py @@ -33,6 +33,8 @@ class PlainNodeSchedule: distributed_buffers: dict[str, dict] fused: list[tuple[str, int]] fused_consumers: list[str] + gpu_blocks: list[str] + gpu_threads: list[str] # Optional caller-provided vector sizes, keyed by vectorized axis name. # When an axis has a size, its dimension is vectorized with masking for # non-divisible extents; axes absent from this mapping are vectorized to @@ -116,6 +118,8 @@ def __init__( self.distributed_buffers: dict[str, dict] = {} self.fused: list[tuple[str, int]] = [] self.fused_consumers: list[str] = [] + self.gpu_blocks: list[str] = [] + self.gpu_threads: list[str] = [] def get_plain_schedule(self) -> PlainNodeSchedule: return PlainNodeSchedule( @@ -137,6 +141,8 @@ def get_plain_schedule(self) -> PlainNodeSchedule: distributed_buffers=deepcopy(self.distributed_buffers), fused=deepcopy(self.fused), fused_consumers=deepcopy(self.fused_consumers), + gpu_blocks=deepcopy(self.gpu_blocks), + gpu_threads=deepcopy(self.gpu_threads), vectorization_sizes=deepcopy(self.vectorization_sizes), ) @@ -274,3 +280,13 @@ def fuse_producer_at( def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None: fuse_axis = make_loop_name(root, axis) self.fused_consumers.append(fuse_axis) + + def gpu_thread(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 gpu_block(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] diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py index 7c1f84923..9cbaeb9ac 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py @@ -9,7 +9,7 @@ # Create device gpu = GPUDevice() -I, J, K, dtype = 4, 32, 512, "float32" +I, J, K, dtype = 1024, 1024, 512, "float32" a = O.tensor((I, K), dtype, name="A") # A lives on the host b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator @@ -22,10 +22,14 @@ impl = Backend(graph) sch = impl.get_scheduler() -sch.tile("i", {"i1": 2}) -sch.tile("j", {"j1": 16}) -sch.unroll({"i1": 2}) -sch.parallelize(["i"]) +sch.tile("i", {"i1": 128, "i2": 32}) +sch.tile("j", {"j1": 128, "j2": 32}) +sch.tile("k", {"k1": 64}) +# sch.unroll({"i1": 2}) +sch.parallelize(["i", "j","i1", "j1"]) +sch.gpu_block(["i", "j"]) +sch.gpu_thread(["i1", "j1"]) +sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) sched = sch.schedule() comp = impl.get_compiler( @@ -34,6 +38,7 @@ dump_file="gpu_matmul_mlir_offload_tensor", print_source_ir=True, print_transformed_ir=True, + print_lowered_ir=True, ) module = comp.compile(sched) executor = module.get_executor(validate=True) diff --git a/tests/filecheck/schedules/test_matmul_descript_gpu.py b/tests/filecheck/schedules/test_matmul_descript_gpu.py new file mode 100644 index 000000000..7a1080d50 --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_gpu.py @@ -0,0 +1,167 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend +from xtc.schedules.descript import descript_scheduler + +from xtc.runtimes.accelerator.gpu import GPUDevice + +gpu = GPUDevice() +I, J, K, dtype = 1024, 1024, 512, "float32" +a = O.tensor((I, K), dtype, name="A") +b = O.tensor((K, J), dtype, name="B", device=gpu) + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +descript_scheduler( + scheduler = sch, + node_name = "C", + abstract_dims = ["I","J","K"], + spec = { + "I": {"parallelize": True, "gpu_block": 0}, + "J": {"parallelize": True, "gpu_block": 1}, + "K": {}, + "I#128": {"parallelize": True, "gpu_thread": 0}, + "J#128": {"parallelize": True, "gpu_thread": 1}, + "I#32": {}, + "J#32": {}, + + } +) + +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="matmul_descript_mlir_gpu", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<4x32xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<4x512xf32>, memref<512x32xf32>) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %1 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "C/K" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [2, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "C/I" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 16, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "C/J" : !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 "C/I0" : !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 %loops_3 {isolated_from_above} : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %2 { +# 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: transform.apply_patterns to %2 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<4x32xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<1x16xf32> +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %cst_0 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c4 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 32] [1, 1] : memref<4x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c1 { +# CHECK-NEXT: %subview_1 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst_0 : f32) outs(%subview_1 : memref<1x1xf32, strided<[32, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg0[0, %arg3] [4, 1] [1, 1] : memref<4x512xf32> to memref<4x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_1 = memref.subview %arg1[%arg3, 0] [1, 32] [1, 1] : memref<512x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_2 = memref.subview %arg2[0, 0] [4, 32] [1, 1] : memref<4x32xf32> to memref<4x32xf32, strided<[32, 1]>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c4 step %c2 { +# CHECK-NEXT: %subview_3 = memref.subview %subview[%arg4, 0] [2, 1] [1, 1] : memref<4x1xf32, strided<[512, 1], offset: ?>> to memref<2x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_4 = memref.subview %subview_2[%arg4, 0] [2, 32] [1, 1] : memref<4x32xf32, strided<[32, 1]>> to memref<2x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg5 = %c0 to %c32 step %c16 { +# CHECK-NEXT: %subview_5 = memref.subview %subview_1[0, %arg5] [1, 16] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_6 = memref.subview %subview_4[0, %arg5] [2, 16] [1, 1] : memref<2x32xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %subview_3[%c0, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %subview_6[%c0, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %1 = vector.transfer_read %subview_7[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %2 = vector.transfer_read %subview_5[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> +# CHECK-NEXT: %3 = vector.transfer_read %subview_8[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> +# CHECK-NEXT: %4 = vector.extract %2[0] : vector<16xf32> from vector<1x16xf32> +# CHECK-NEXT: %5 = vector.extract %1[0, 0] : f32 from vector<1x1xf32> +# CHECK-NEXT: %6 = vector.broadcast %5 : f32 to vector<16xf32> +# CHECK-NEXT: %7 = vector.extract %3[0] : vector<16xf32> from vector<1x16xf32> +# CHECK-NEXT: %8 = vector.fma %6, %4, %7 : vector<16xf32> +# CHECK-NEXT: %9 = vector.insert %8, %cst [0] : vector<16xf32> into vector<1x16xf32> +# CHECK-NEXT: vector.transfer_write %9, %subview_8[%c0, %c0] {in_bounds = [true, true]} : vector<1x16xf32>, memref<1x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_3[%c1, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_6[%c1, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %10 = vector.transfer_read %subview_9[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %11 = vector.transfer_read %subview_5[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> +# CHECK-NEXT: %12 = vector.transfer_read %subview_10[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> +# CHECK-NEXT: %13 = vector.extract %11[0] : vector<16xf32> from vector<1x16xf32> +# CHECK-NEXT: %14 = vector.extract %10[0, 0] : f32 from vector<1x1xf32> +# CHECK-NEXT: %15 = vector.broadcast %14 : f32 to vector<16xf32> +# CHECK-NEXT: %16 = vector.extract %12[0] : vector<16xf32> from vector<1x16xf32> +# CHECK-NEXT: %17 = vector.fma %15, %13, %16 : vector<16xf32> +# CHECK-NEXT: %18 = vector.insert %17, %cst [0] : vector<16xf32> into vector<1x16xf32> +# CHECK-NEXT: vector.transfer_write %18, %subview_10[%c0, %c0] {in_bounds = [true, true]} : vector<1x16xf32>, memref<1x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: } {"C/J"} +# CHECK-NEXT: } {"C/I"} +# CHECK-NEXT: } {"C/K"} +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 4x512xfloat32 +# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/search/test_conv_oo.py b/tests/filecheck/search/test_conv_oo.py index 4798c0913..e842e3386 100644 --- a/tests/filecheck/search/test_conv_oo.py +++ b/tests/filecheck/search/test_conv_oo.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 2, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 2, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 3, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 3, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 7, 1] @@ -99,4 +99,4 @@ # CHECK-NEXT: sample 76: [2, 2, 2, 8, 1, 1, 1] # CHECK-NEXT: sample 77: [2, 2, 2, 16, 1, 1, 1] # CHECK-NEXT: stats {'filtered': 78, 'all': 384} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './c1': 1, './s1': 1, './r1': 1, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './c1': 1, './s1': 1, './r1': 1, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprp.py b/tests/filecheck/search/test_conv_pprprp.py index a56b9fb53..aa0763a39 100644 --- a/tests/filecheck/search/test_conv_pprprp.py +++ b/tests/filecheck/search/test_conv_pprprp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 1] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 16, 1, 1, 3] # CHECK-NEXT: stats {'filtered': 100, 'all': 202} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpv.py b/tests/filecheck/search/test_conv_pprprpv.py index 01f3322d9..afecd1a9b 100644 --- a/tests/filecheck/search/test_conv_pprprpv.py +++ b/tests/filecheck/search/test_conv_pprprpv.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 7, 1] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 16, 7, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 32, 1, 1, 1] # CHECK-NEXT: stats {'filtered_vec': 100, 'filtered': 1520, 'all': 4521} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpvr.py b/tests/filecheck/search/test_conv_pprprpvr.py index 337e138aa..a9f78119c 100644 --- a/tests/filecheck/search/test_conv_pprprpvr.py +++ b/tests/filecheck/search/test_conv_pprprpvr.py @@ -20,13 +20,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 7, 1] @@ -128,4 +128,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 32, 1, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 32, 1, 1, 3] # CHECK-NEXT: stats {'filtered_l2': 100, 'filtered_l1': 102, 'filtered_reg': 132, 'filtered_vec': 134, 'filtered': 1918, 'all': 3178} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 2}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 2, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 2}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 2, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpvr_rnd.py b/tests/filecheck/search/test_conv_pprprpvr_rnd.py index 2ce179546..e6f11e898 100644 --- a/tests/filecheck/search/test_conv_pprprpvr_rnd.py +++ b/tests/filecheck/search/test_conv_pprprpvr_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 2, 16, 7, 1, 1] # CHECK-NEXT: sample 19: [1, 2, 1, 1, 2, 1, 2, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: stats {'filtered_l2': 5, 'filtered_l1': 5, 'filtered_reg': 6, 'filtered_vec': 6, 'filtered': 100} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2, './b2': 2, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2, './b2': 2, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_prp.py b/tests/filecheck/search/test_conv_prp.py index afc1a4511..fd1577917 100644 --- a/tests/filecheck/search/test_conv_prp.py +++ b/tests/filecheck/search/test_conv_prp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 2, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 2, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 4] @@ -68,4 +68,4 @@ # CHECK-NEXT: sample 45: [2, 2, 2, 8] # CHECK-NEXT: sample 46: [2, 2, 2, 16] # CHECK-NEXT: stats {'filtered': 47, 'all': 48} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_prp_rnd.py b/tests/filecheck/search/test_conv_prp_rnd.py index 622601c0b..e86b96833 100644 --- a/tests/filecheck/search/test_conv_prp_rnd.py +++ b/tests/filecheck/search/test_conv_prp_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [2, 2, 1, 8] # CHECK-NEXT: sample 19: [2, 2, 1, 4] # CHECK-NEXT: stats {'filtered': 20} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 1}, './f': {'./f1': 4}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=['./b', './h', './w', './f'], unrolling={'./f1': 4, './w1': 1, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 1}, './f': {'./f1': 4}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=['./b', './h', './w', './f'], unrolling={'./f1': 4, './w1': 1, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_goto.py b/tests/filecheck/search/test_matmul_goto.py index 615edb798..ccf37bac2 100644 --- a/tests/filecheck/search/test_matmul_goto.py +++ b/tests/filecheck/search/test_matmul_goto.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 0, 0] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 0, 1] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 0] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 2, 2, 3, 1, 0] # CHECK-NEXT: sample 99: [1, 1, 1, 2, 2, 3, 1, 1] # CHECK-NEXT: stats {'filtered': 100, 'all': 108} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 2, './j2': 2}, './k': {'./k1': 2}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 2, './j2': 2}, './k': {'./k1': 2}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_goto_r.py b/tests/filecheck/search/test_matmul_goto_r.py index b250298df..e76e2cdf7 100644 --- a/tests/filecheck/search/test_matmul_goto_r.py +++ b/tests/filecheck/search/test_matmul_goto_r.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [3, 1, 2, 1, 1, 1, 0, 0] # CHECK-NEXT: sample 1: [3, 1, 2, 1, 1, 1, 0, 1] # CHECK-NEXT: sample 2: [3, 1, 2, 1, 1, 1, 1, 0] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [3, 1, 2, 2, 6, 2, 1, 0] # CHECK-NEXT: sample 99: [3, 1, 2, 2, 6, 2, 1, 1] # CHECK-NEXT: stats {'filtered': 3256, 'all': 6620} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1}, './j': {'./j1': 4, './j2': 2}, './k': {'./k1': 6}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1}, './j': {'./j1': 4, './j2': 2}, './k': {'./k1': 6}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_oo.py b/tests/filecheck/search/test_matmul_oo.py index 6fb3d9a44..0a9dfa591 100644 --- a/tests/filecheck/search/test_matmul_oo.py +++ b/tests/filecheck/search/test_matmul_oo.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [3, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 12, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 12, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 3] @@ -66,4 +66,4 @@ # CHECK-NEXT: sample 43: [7, 8, 1] # CHECK-NEXT: sample 44: [7, 16, 1] # CHECK-NEXT: stats {'filtered': 45, 'all': 144} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_p1.py b/tests/filecheck/search/test_matmul_p1.py index ffa955612..43773867b 100644 --- a/tests/filecheck/search/test_matmul_p1.py +++ b/tests/filecheck/search/test_matmul_p1.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 0] # CHECK-NEXT: sample 1: [1, 1, 1, 1] # CHECK-NEXT: sample 2: [1, 1, 1, 2] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 32, 1, 1] # CHECK-NEXT: sample 99: [1, 32, 1, 4] # CHECK-NEXT: stats {'filtered': 100, 'all': 185} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 32}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 32, './i1': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 32}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 32, './i1': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_p1v.py b/tests/filecheck/search/test_matmul_p1v.py index 92b21f24a..84a94b71d 100644 --- a/tests/filecheck/search/test_matmul_p1v.py +++ b/tests/filecheck/search/test_matmul_p1v.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 16, 1, 1] # CHECK-NEXT: sample 1: [1, 16, 1, 4] # CHECK-NEXT: sample 2: [1, 16, 2, 1] @@ -47,4 +47,4 @@ # CHECK-NEXT: sample 24: [7, 16, 1, 1] # CHECK-NEXT: sample 25: [7, 16, 1, 4] # CHECK-NEXT: stats {'filtered': 154, 'all': 864} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprp.py b/tests/filecheck/search/test_matmul_pprprp.py index 5a31e9712..e79fd3449 100644 --- a/tests/filecheck/search/test_matmul_pprprp.py +++ b/tests/filecheck/search/test_matmul_pprprp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 3] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 16, 2, 6] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 32, 1, 1] # CHECK-NEXT: stats {'filtered': 100, 'all': 121} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpv.py b/tests/filecheck/search/test_matmul_pprprpv.py index 2eb1bf600..30becfda8 100644 --- a/tests/filecheck/search/test_matmul_pprprpv.py +++ b/tests/filecheck/search/test_matmul_pprprpv.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 16, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 16, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 16, 3] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [3, 1, 1, 1, 1, 16, 3] # CHECK-NEXT: sample 99: [3, 1, 1, 1, 1, 16, 4] # CHECK-NEXT: stats {'filtered_vec': 100, 'filtered': 1472, 'all': 3052} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 4}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 4}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 4}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 4}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpvr.py b/tests/filecheck/search/test_matmul_pprprpvr.py index fd0c84aaa..9e11e191d 100644 --- a/tests/filecheck/search/test_matmul_pprprpvr.py +++ b/tests/filecheck/search/test_matmul_pprprpvr.py @@ -20,13 +20,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 16, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 16, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 16, 3] @@ -128,4 +128,4 @@ # CHECK-NEXT: sample 98: [1, 21, 1, 1, 1, 32, 3] # CHECK-NEXT: sample 99: [1, 21, 1, 1, 2, 16, 1] # CHECK-NEXT: stats {'filtered_l2': 100, 'filtered_l1': 105, 'filtered_reg': 115, 'filtered_vec': 154, 'filtered': 2126, 'all': 2749} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 21, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 21, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpvr_rnd.py b/tests/filecheck/search/test_matmul_pprprpvr_rnd.py index 6fa9effc2..5780a7eb6 100644 --- a/tests/filecheck/search/test_matmul_pprprpvr_rnd.py +++ b/tests/filecheck/search/test_matmul_pprprpvr_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [1, 1, 1, 1, 2, 16, 3] # CHECK-NEXT: sample 19: [7, 1, 3, 1, 1, 16, 2] # CHECK-NEXT: stats {'filtered_l2': 2, 'filtered_l1': 2, 'filtered_reg': 3, 'filtered_vec': 3, 'filtered': 70} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 2}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 2}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_prp.py b/tests/filecheck/search/test_matmul_prp.py index 757860aaa..71837f2c8 100644 --- a/tests/filecheck/search/test_matmul_prp.py +++ b/tests/filecheck/search/test_matmul_prp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [3, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1] # CHECK-NEXT: sample 1: [1, 2] # CHECK-NEXT: sample 2: [1, 4] @@ -38,4 +38,4 @@ # CHECK-NEXT: sample 15: [7, 8] # CHECK-NEXT: sample 16: [7, 16] # CHECK-NEXT: stats {'filtered': 17, 'all': 24} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_prp_rnd.py b/tests/filecheck/search/test_matmul_prp_rnd.py index 47a46ed14..aafbc2661 100644 --- a/tests/filecheck/search/test_matmul_prp_rnd.py +++ b/tests/filecheck/search/test_matmul_prp_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [3, 2] # CHECK-NEXT: sample 19: [21, 8] # CHECK-NEXT: stats {'filtered': 19} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21}, './j': {'./j1': 8}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=['./i', './j'], unrolling={'./j1': 8, './i1': 21}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21}, './j': {'./j1': 8}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=['./i', './j'], unrolling={'./j1': 8, './i1': 21}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] From 3de501c5e4731b2daaf9e557b798252bfe629259 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 18 May 2026 12:31:48 +0200 Subject: [PATCH 2/7] Add mapping order + use can use forall for any size of problem now --- src/xtc/backends/mlir/MlirCompilerPasses.py | 24 +++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index cf8b473ef..509ca918b 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -367,26 +367,40 @@ def _generate_node_scheduling( elif loop_name in tiles_sizes_by_loops: if loop_name in schedule.gpu_blocks: tile_vect = [ - max(tiles_sizes_by_loops[loop]) for loop in schedule.gpu_blocks + 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)) if gpu_material: - new_loop = self._strip_mine( + self._strip_mine( loop_name=loop_name, tiling_vector=tile_vect, + mapping_order=[], schedule=schedule, sched_state=sched_state, ) gpu_material = False elif loop_name in schedule.gpu_threads: tile_vect = [ - max(tiles_sizes_by_loops[loop]) for loop in schedule.gpu_threads + 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)) if gpu_mat_thread: - new_loop = self._strip_mine( + self._strip_mine( loop_name=loop_name, tiling_vector=tile_vect, + mapping_order=[], schedule=schedule, sched_state=sched_state, ) @@ -395,6 +409,7 @@ def _generate_node_scheduling( self._strip_mine( loop_name=loop_name, tiling_vector=tiles_sizes_by_loops[loop_name], + mapping_order=[], schedule=schedule, sched_state=sched_state, ) @@ -533,6 +548,7 @@ def _strip_mine( self, loop_name: str, tiling_vector: list[int], + mapping_order: list[int], schedule: MlirNodeSchedule, sched_state: SchedulingState, ) -> OpResult: From 2bc574b73edcf4d0b02acc91579ae960dcae800d Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 18 May 2026 14:38:04 +0200 Subject: [PATCH 3/7] Add mapping order for the gpu thread and block And fix some primitive and type problem --- src/xtc/backends/jir/JIRScheduler.py | 4 +- src/xtc/backends/mlir/MlirCompilerPasses.py | 113 ++++++----- src/xtc/itf/schd/scheduler.py | 4 +- src/xtc/schedules/descript.py | 12 +- src/xtc/schedules/parameter_loop_nest.py | 5 + src/xtc/schedules/parsing.py | 3 +- .../test_matmul_mlir_offload_tensor.py | 191 ++++++++++-------- tests/filecheck/search/test_matmul_prp.py | 37 +--- 8 files changed, 198 insertions(+), 171 deletions(-) diff --git a/src/xtc/backends/jir/JIRScheduler.py b/src/xtc/backends/jir/JIRScheduler.py index 037caec4c..5627ad0ed 100644 --- a/src/xtc/backends/jir/JIRScheduler.py +++ b/src/xtc/backends/jir/JIRScheduler.py @@ -362,12 +362,12 @@ def distributed_buffer_at( pass @override - def gpu_thread(self, axes: list[str]) -> None: + 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]) -> None: + def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: # TODO: not implemented for now pass diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index 509ca918b..6df98556a 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -247,6 +247,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" @@ -376,11 +383,18 @@ def _generate_node_scheduling( ) ] 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=mapping_order, schedule=schedule, sched_state=sched_state, ) @@ -396,11 +410,17 @@ def _generate_node_scheduling( ) ] 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=mapping_order, schedule=schedule, sched_state=sched_state, ) @@ -429,10 +449,6 @@ def _generate_node_scheduling( if schedule.unrolling: self._unroll(permutation, schedule, sched_state) - # Gpu mapping - if schedule.gpu_blocks: - self._gpu_mapping(schedule, sched_state, tiles_sizes_by_loops) - return sched_state def _fuse_producers_into_loop( @@ -556,22 +572,13 @@ def _strip_mine( attr_array = {} if loop_name in schedule.gpu_threads: attr_array["mapping"] = ArrayAttr.get( - [ - self._get_thread_id(index) - for index in range(len(schedule.gpu_threads)) - ] + [self._get_thread_id(index) for index in mapping_order] ) - attr_array["tile_sizes"] = tiling_vector elif loop_name in schedule.gpu_blocks: attr_array["mapping"] = ArrayAttr.get( - [ - self._get_block_id(index) - for index in range(len(schedule.gpu_blocks)) - ] + [self._get_block_id(index) for index in mapping_order] ) - attr_array["tile_sizes"] = tiling_vector - else: - attr_array["tile_sizes"] = tiling_vector + attr_array["tile_sizes"] = tiling_vector tiling_command = TileUsingForallOp(sched_state.handle, **attr_array) else: tiling_command = TileUsingForOp(sched_state.handle, sizes=tiling_vector) @@ -776,30 +783,6 @@ def _collect_fused_producers(self, unscheduled_handles: set[str | None]): return fused_producers - -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 - producer_handles: list[str | None] = [] - root_op = None - for func_op in module.body.operations: - for op in func_op.regions[0].blocks[0].operations: - if root_handle in op.attributes: - root_op = op - break - if root_op: - break - - if not root_op: - return producer_handles - for operand in root_op.operands: - producer_op = operand.owner - producer_handles.append(None) - if producer_op and hasattr(producer_op, "attributes"): - for attr in producer_op.attributes: - if attr.startswith("__xtc_id_"): - producer_handles[-1] = attr - return producer_handles - 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) @@ -812,8 +795,8 @@ def _gpu_mapping( self, schedule: MlirNodeSchedule, sched_state: SchedulingState, - tiles_sizes_by_loops: dict[str, list[int]], ): + tiles_sizes_by_loops = self._generate_tiling_insns(schedule) if schedule.gpu_blocks: new_loop = next( ( @@ -824,27 +807,49 @@ def _gpu_mapping( None, ) # Since we know there only 1 non zero number - # TODO Find a way to put block id and thread id divide by the size - grid_dims = [ - max(tiles_sizes_by_loops[loop_name]) - for loop_name in schedule.gpu_blocks - ] - grid_dims = grid_dims + [1] * (3 - len(grid_dims)) + # TODO Find a way to put thread number instead of putting tile size new_loop = MapForallToBlocks( new_loop, - grid_dims=grid_dims, generate_gpu_launch=True, ).result if schedule.gpu_threads: block_dims = [ - max(tiles_sizes_by_loops[loop_name]) - for loop_name in schedule.gpu_threads + 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)) - new_loop = MapNestedForallToThreads( + MapNestedForallToThreads( new_loop, block_dims=block_dims, - ).result + ) + + +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 + producer_handles: list[str | None] = [] + root_op = None + for func_op in module.body.operations: + for op in func_op.regions[0].blocks[0].operations: + if root_handle in op.attributes: + root_op = op + break + if root_op: + break + + if not root_op: + return producer_handles + for operand in root_op.operands: + producer_op = operand.owner + producer_handles.append(None) + if producer_op and hasattr(producer_op, "attributes"): + for attr in producer_op.attributes: + if attr.startswith("__xtc_id_"): + producer_handles[-1] = attr + return producer_handles + class MlirProgramApplyTransformPass: diff --git a/src/xtc/itf/schd/scheduler.py b/src/xtc/itf/schd/scheduler.py index 541961051..0fa697f95 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -320,7 +320,7 @@ def distributed_buffer_at( ... @abstractmethod - def gpu_thread(self, axes: list[str]) -> None: + def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: """ Map an axis to a certain gpu thread either x, y or z. The mapped axis need to be in parallelization. @@ -329,7 +329,7 @@ def gpu_thread(self, axes: list[str]) -> None: ... @abstractmethod - def gpu_block(self, axes: list[str]) -> None: + def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: """ Map an axis to a certain gpu block either x, y or z. The mapped axis need to be in parallelization. diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 51fc208d2..8ac2f799c 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -556,14 +556,18 @@ def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None: scheduler.fuse_consumer_at(axis, root=root) if node.gpu_block: - self.scheduler.gpu_block( - sorted(node.gpu_block, key=node.gpu_block.get), root=root + 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: - self.scheduler.gpu_thread( - sorted(node.gpu_thread, key=node.gpu_thread.get), root=root + sorted_keys = sorted( + (k for k, v in node.gpu_thread.items() if v is not None), + key=lambda k: node.gpu_block[k], ) + scheduler.gpu_thread(sorted_keys, root=root) # Recursively apply children for child in node.children: diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index 822073fa9..b4e2f9e91 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -111,6 +111,8 @@ 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 """ root: str @@ -127,6 +129,9 @@ 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) + # TODO: make gpu_block and gpu_thread work for parameter loop nest + 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: """ diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index d6cb1219b..4e2bac4b9 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -292,7 +292,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations if partial and full: raise ScheduleParseError(f"{context} has both annotations full and partial") - + assert isinstance(gpu_block, int) or gpu_block is None + assert isinstance(gpu_thread, int) or gpu_thread is None return Annotations( unroll_factor=unroll_factor, unroll_specified=unroll_specified, diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py index 9cbaeb9ac..2a02e1317 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py @@ -25,7 +25,7 @@ sch.tile("i", {"i1": 128, "i2": 32}) sch.tile("j", {"j1": 128, "j2": 32}) sch.tile("k", {"k1": 64}) -# sch.unroll({"i1": 2}) +sch.unroll({"i2": 2}) sch.parallelize(["i", "j","i1", "j1"]) sch.gpu_block(["i", "j"]) sch.gpu_thread(["i1", "j1"]) @@ -38,7 +38,6 @@ dump_file="gpu_matmul_mlir_offload_tensor", print_source_ir=True, print_transformed_ir=True, - print_lowered_ir=True, ) module = comp.compile(sched) executor = module.get_executor(validate=True) @@ -46,10 +45,10 @@ print(f"CODE: {res}") # CHECK: // -----// IR Dump Before transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias, memref.on_device}, %arg2: memref<4x32xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<4x32xf32>) -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<4x512xf32>, memref<512x32xf32>) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { @@ -63,92 +62,122 @@ # CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op # CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op -# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [2, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 16, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_3 "./j" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [32, 32, 0](mapping = [#gpu.thread, #gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_3 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_7 "./i1" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_9 "./j1" : !transform.any_op -# CHECK-NEXT: transform.loop.unroll %loops_7 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# 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: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./j2" : !transform.any_op +# CHECK-NEXT: transform.loop.unroll %loops_9 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op # CHECK-NEXT: transform.yield # CHECK-NEXT: } # CHECK-NEXT: } # CHECK-NEXT: # CHECK-NEXT: // -----// IR Dump After transform //----- // -# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 2)> +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias, memref.on_device}, %arg2: memref<4x32xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index -# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index # CHECK-NEXT: %c1 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg3 = %c0 to %c4 step %c1 { -# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 32] [1, 1] : memref<4x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %c0_0 = arith.constant 0 : index -# CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %c1_1 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg4 = %c0_0 to %c32 step %c1_1 { -# CHECK-NEXT: %subview_2 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_2 : memref<1x1xf32, strided<[32, 1], offset: ?>>) +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c1024_6 = arith.constant 1024 : index +# CHECK-NEXT: %c1_7 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_5 to %c1024_6 step %c1_7 { +# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_8 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) # CHECK-NEXT: } {"./j"} # CHECK-NEXT: } {"./i"} -# CHECK-NEXT: scf.forall (%arg3) in (2) { -# CHECK-NEXT: %0 = affine.apply #map(%arg3) -# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [2, 512] [1, 1] : memref<4x512xf32> to memref<2x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_0 = memref.subview %arg1[0, 0] [512, 32] [1, 1] : memref<512x32xf32> to memref<512x32xf32, strided<[32, 1]>> -# CHECK-NEXT: %subview_1 = memref.subview %arg2[%0, 0] [2, 32] [1, 1] : memref<4x32xf32> to memref<2x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %c0_2 = arith.constant 0 : index -# CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %c16 = arith.constant 16 : index -# CHECK-NEXT: scf.for %arg4 = %c0_2 to %c32 step %c16 { -# CHECK-NEXT: %subview_3 = memref.subview %subview[0, 0] [2, 512] [1, 1] : memref<2x512xf32, strided<[512, 1], offset: ?>> to memref<2x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_4 = memref.subview %subview_0[0, %arg4] [512, 16] [1, 1] : memref<512x32xf32, strided<[32, 1]>> to memref<512x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_5 = memref.subview %subview_1[0, %arg4] [2, 16] [1, 1] : memref<2x32xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %c0_6 = arith.constant 0 : index -# CHECK-NEXT: %c512 = arith.constant 512 : index -# CHECK-NEXT: %c1_7 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg5 = %c0_6 to %c512 step %c1_7 { -# CHECK-NEXT: %subview_8 = memref.subview %subview_3[0, %arg5] [2, 1] [1, 1] : memref<2x512xf32, strided<[512, 1], offset: ?>> to memref<2x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_9 = memref.subview %subview_4[%arg5, 0] [1, 16] [1, 1] : memref<512x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_10 = memref.subview %subview_5[0, 0] [2, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %c0_11 = arith.constant 0 : index +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c4_1 = arith.constant 4 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c8_3 = arith.constant 8 : index +# CHECK-NEXT: %c1_4 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c8_3, %arg11 = %c1_4) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %2 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_y) +# CHECK-NEXT: %subview_9 = memref.subview %subview[%2, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_7[0, %3] [512, 32] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<512x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[%2, %3] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_12 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: scf.for %arg15 = %c0_12 to %c512 step %c64 { +# CHECK-NEXT: %subview_13 = memref.subview %subview_9[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_10[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[1024, 1], offset: ?>> to memref<64x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_11[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_16 = arith.constant 0 : index +# CHECK-NEXT: %c64_17 = arith.constant 64 : index +# CHECK-NEXT: %c1_18 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg16 = %c0_16 to %c64_17 step %c1_18 { +# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_21 = memref.subview %subview_15[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_22 = arith.constant 0 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c1_23 = arith.constant 1 : index # CHECK-NEXT: %c2 = arith.constant 2 : index -# CHECK-NEXT: %c1_12 = arith.constant 1 : index -# CHECK-NEXT: %c2_13 = arith.constant 2 : index -# CHECK-NEXT: %subview_14 = memref.subview %subview_8[%c0_11, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_15 = memref.subview %subview_9[0, 0] [1, 16] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_16 = memref.subview %subview_10[%c0_11, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %c0_17 = arith.constant 0 : index -# CHECK-NEXT: %c16_18 = arith.constant 16 : index -# CHECK-NEXT: %c1_19 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg6 = %c0_17 to %c16_18 step %c1_19 { -# CHECK-NEXT: %subview_27 = memref.subview %subview_14[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_28 = memref.subview %subview_15[0, %arg6] [1, 1] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_29 = memref.subview %subview_16[0, %arg6] [1, 1] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_27, %subview_28 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[32, 1], offset: ?>>) outs(%subview_29 : memref<1x1xf32, strided<[32, 1], offset: ?>>) -# CHECK-NEXT: } {"./j1"} -# CHECK-NEXT: %c1_20 = arith.constant 1 : index -# CHECK-NEXT: %1 = arith.muli %c1_12, %c1_20 : index -# CHECK-NEXT: %2 = arith.addi %c0_11, %1 : index -# CHECK-NEXT: %subview_21 = memref.subview %subview_8[%2, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_22 = memref.subview %subview_9[0, 0] [1, 16] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_23 = memref.subview %subview_10[%2, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %c0_24 = arith.constant 0 : index -# CHECK-NEXT: %c16_25 = arith.constant 16 : index -# CHECK-NEXT: %c1_26 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg6 = %c0_24 to %c16_25 step %c1_26 { -# CHECK-NEXT: %subview_27 = memref.subview %subview_21[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_28 = memref.subview %subview_22[0, %arg6] [1, 1] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_29 = memref.subview %subview_23[0, %arg6] [1, 1] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_27, %subview_28 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[32, 1], offset: ?>>) outs(%subview_29 : memref<1x1xf32, strided<[32, 1], offset: ?>>) -# CHECK-NEXT: } {"./j1"} -# CHECK-NEXT: } {"./k"} -# CHECK-NEXT: } {"./j"} -# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: scf.for %arg17 = %c0_22 to %c32 step %c2 { +# CHECK-NEXT: %subview_24 = memref.subview %subview_19[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_27 = arith.constant 0 : index +# CHECK-NEXT: %c32_28 = arith.constant 32 : index +# CHECK-NEXT: %c1_29 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_27 to %c32_28 step %c1_29 { +# CHECK-NEXT: %subview_37 = memref.subview %subview_24[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_38 = memref.subview %subview_25[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_39 = memref.subview %subview_26[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %c1_30 = arith.constant 1 : index +# CHECK-NEXT: %4 = arith.muli %c1_23, %c1_30 : index +# CHECK-NEXT: %5 = arith.addi %arg17, %4 : index +# CHECK-NEXT: %subview_31 = memref.subview %subview_19[%5, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_33 = memref.subview %subview_21[%5, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_34 = arith.constant 0 : index +# CHECK-NEXT: %c32_35 = arith.constant 32 : index +# CHECK-NEXT: %c1_36 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_34 to %c32_35 step %c1_36 { +# CHECK-NEXT: %subview_37 = memref.subview %subview_31[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_38 = memref.subview %subview_32[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_39 = memref.subview %subview_33[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: } @@ -156,11 +185,11 @@ # CHECK-NEXT: graph: # CHECK-NEXT: name: matmul # CHECK-NEXT: inputs: -# CHECK-NEXT: - %0 : 4x512xfloat32 -# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: - %0 : 1024x512xfloat32 +# CHECK-NEXT: - %1 : 512x1024xfloat32 # CHECK-NEXT: outputs: -# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: - %2 : 1024x1024xfloat32 # CHECK-NEXT: nodes: -# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] # CHECK-NEXT: # CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/search/test_matmul_prp.py b/tests/filecheck/search/test_matmul_prp.py index 71837f2c8..199db2797 100644 --- a/tests/filecheck/search/test_matmul_prp.py +++ b/tests/filecheck/search/test_matmul_prp.py @@ -12,30 +12,13 @@ utils.print_all_opt_schedules(backend, strategy) utils.print_exhaustive_samples(backend, strategy, 100) -# CHECK: schedule O0: [1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] -# CHECK-NEXT: schedule O1: [1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] -# CHECK-NEXT: schedule O2: [1, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] -# CHECK-NEXT: schedule O3: [3, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] -# CHECK-NEXT: sample 0: [1, 1] -# CHECK-NEXT: sample 1: [1, 2] -# CHECK-NEXT: sample 2: [1, 4] -# CHECK-NEXT: sample 3: [1, 8] -# CHECK-NEXT: sample 4: [1, 16] -# CHECK-NEXT: sample 5: [1, 32] -# CHECK-NEXT: sample 6: [3, 1] -# CHECK-NEXT: sample 7: [3, 2] -# CHECK-NEXT: sample 8: [3, 4] -# CHECK-NEXT: sample 9: [3, 8] -# CHECK-NEXT: sample 10: [3, 16] -# CHECK-NEXT: sample 11: [3, 32] -# CHECK-NEXT: sample 12: [7, 1] -# CHECK-NEXT: sample 13: [7, 2] -# CHECK-NEXT: sample 14: [7, 4] -# CHECK-NEXT: sample 15: [7, 8] -# CHECK-NEXT: sample 16: [7, 16] -# CHECK-NEXT: stats {'filtered': 17, 'all': 24} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +<<<<<<< HEAD +======= +>>>>>>> 7fae6dc (Add mapping order for the gpu thread and block) +<<<<<<< HEAD +======= +>>>>>>> 7fae6dc (Add mapping order for the gpu thread and block) +# CHECK: File "/home/ruicesista/Documents/xtc/xtc/tests/filecheck/search/test_matmul_prp.py", line 17 +# CHECK-NEXT: >>>>>>> 7fae6dc (Add mapping order for the gpu thread and block) +# CHECK-NEXT: ^ +# CHECK-NEXT: SyntaxError: invalid decimal literal From e679478f112ba782d006e4aaa3f9ded5a48c438f Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Fri, 19 Jun 2026 16:51:43 +0200 Subject: [PATCH 4/7] nvgpu: gpu work now with tensor now --- src/xtc/backends/mlir/MlirCompilerPasses.py | 44 +++- .../test_matmul_mlir_with_tensor.py | 196 ++++++++++++++++++ 2 files changed, 239 insertions(+), 1 deletion(-) create mode 100644 tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index 6df98556a..829aff96f 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -142,6 +142,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 [] ) @@ -578,6 +579,7 @@ def _strip_mine( attr_array["mapping"] = ArrayAttr.get( [self._get_block_id(index) for index in mapping_order] ) + self._gpu_block_order = attr_array["mapping"] attr_array["tile_sizes"] = tiling_vector tiling_command = TileUsingForallOp(sched_state.handle, **attr_array) else: @@ -587,6 +589,8 @@ def _strip_mine( assert len(tiling_command.results) == 2 new_loop = tiling_command.results[-1] sched_state.all_loops[loop_name] = new_loop + if schedule.gpu_blocks: + loop_name = schedule.gpu_blocks[0] # Annotate the resulting loop if successfully generated transform.AnnotateOp(new_loop, loop_name) @@ -797,7 +801,7 @@ def _gpu_mapping( sched_state: SchedulingState, ): tiles_sizes_by_loops = self._generate_tiling_insns(schedule) - if schedule.gpu_blocks: + if schedule.gpu_blocks and not self._using_tensors: new_loop = next( ( sched_state.all_loops[loop_name] @@ -825,6 +829,44 @@ def _gpu_mapping( 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]: diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py new file mode 100644 index 000000000..0858d6577 --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py @@ -0,0 +1,196 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir.MlirGraphBackend import MlirGraphBackend as Backend + +from xtc.runtimes.accelerator.gpu import GPUDevice + +# Create device +gpu = GPUDevice() + +I, J, K, dtype = 1024, 1024, 512, "float32" +a = O.tensor((I, K), dtype, name="A") # A lives on the host +b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) # C must live on the accelerator + +graph = gb.graph +print(graph) + +impl = Backend(graph, use_tensor_dialect=True) + +sch = impl.get_scheduler() +sch.tile("i", {"i1": 128, "i2": 32}) +sch.tile("j", {"j1": 128, "j2": 32}) +sch.tile("k", {"k1": 64}) +# sch.unroll({"i2": 2}) +sch.parallelize(["i", "j","i1", "j1"]) +sch.gpu_block(["i", "j"]) +sch.gpu_thread(["i1", "j1"]) +sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="gpu_matmul_mlir_offload_tensor", + print_source_ir=True, + print_transformed_ir=True, + print_bufferization_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op +# CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [32, 32, 0](mapping = [#gpu.thread, #gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_3 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# 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: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./j2" : !transform.any_op +# CHECK-NEXT: transform.loop.unroll %loops_9 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c1024_6 = arith.constant 1024 : index +# CHECK-NEXT: %c1_7 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_5 to %c1024_6 step %c1_7 { +# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_8 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c4_1 = arith.constant 4 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c8_3 = arith.constant 8 : index +# CHECK-NEXT: %c1_4 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c8_3, %arg11 = %c1_4) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %2 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_y) +# CHECK-NEXT: %subview_9 = memref.subview %subview[%2, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_7[0, %3] [512, 32] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<512x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[%2, %3] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_12 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: scf.for %arg15 = %c0_12 to %c512 step %c64 { +# CHECK-NEXT: %subview_13 = memref.subview %subview_9[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_10[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[1024, 1], offset: ?>> to memref<64x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_11[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_16 = arith.constant 0 : index +# CHECK-NEXT: %c64_17 = arith.constant 64 : index +# CHECK-NEXT: %c1_18 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg16 = %c0_16 to %c64_17 step %c1_18 { +# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_21 = memref.subview %subview_15[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_22 = arith.constant 0 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c1_23 = arith.constant 1 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index +# CHECK-NEXT: scf.for %arg17 = %c0_22 to %c32 step %c2 { +# CHECK-NEXT: %subview_24 = memref.subview %subview_19[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_27 = arith.constant 0 : index +# CHECK-NEXT: %c32_28 = arith.constant 32 : index +# CHECK-NEXT: %c1_29 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_27 to %c32_28 step %c1_29 { +# CHECK-NEXT: %subview_37 = memref.subview %subview_24[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_38 = memref.subview %subview_25[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_39 = memref.subview %subview_26[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %c1_30 = arith.constant 1 : index +# CHECK-NEXT: %4 = arith.muli %c1_23, %c1_30 : index +# CHECK-NEXT: %5 = arith.addi %arg17, %4 : index +# CHECK-NEXT: %subview_31 = memref.subview %subview_19[%5, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_33 = memref.subview %subview_21[%5, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_34 = arith.constant 0 : index +# CHECK-NEXT: %c32_35 = arith.constant 32 : index +# CHECK-NEXT: %c1_36 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_34 to %c32_35 step %c1_36 { +# CHECK-NEXT: %subview_37 = memref.subview %subview_31[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_38 = memref.subview %subview_32[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_39 = memref.subview %subview_33[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1024x512xfloat32 +# CHECK-NEXT: - %1 : 512x1024xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 1024x1024xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 From 348784f7d17bc8c8f7e40649227d07a778a5755e Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 22 Jun 2026 13:55:53 +0200 Subject: [PATCH 5/7] Remove mandatory parallel for gpu_block and gpu_thread primitive + modify some test Fix annotation the same with gpu_block --- src/xtc/backends/mlir/MlirCompilerPasses.py | 33 +-- src/xtc/backends/mlir/MlirNodeScheduler.py | 4 +- .../mlir/MlirTarget/MlirNVGPUTarget.py | 8 +- src/xtc/itf/schd/scheduler.py | 30 ++- .../test_matmul_mlir_1D_offload_tensor.py | 212 ++++++++++++++++ .../test_matmul_mlir_offload_tensor.py | 139 ++++++----- ...st_matmul_mlir_offload_tensor_vectorise.py | 173 +++++++------ .../test_matmul_mlir_with_tensor.py | 233 ++++++++++++------ .../test_matmul_pmu_counters_gpu.py | 8 +- .../schedules/test_matmul_descript_gpu.py | 172 +++++++------ 10 files changed, 681 insertions(+), 331 deletions(-) create mode 100644 tests/filecheck/backends/target_gpu/test_matmul_mlir_1D_offload_tensor.py diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index 829aff96f..855201d4f 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -569,18 +569,20 @@ def _strip_mine( schedule: MlirNodeSchedule, sched_state: SchedulingState, ) -> OpResult: - if loop_name in schedule.parallelization: - attr_array = {} - if loop_name in schedule.gpu_threads: - attr_array["mapping"] = ArrayAttr.get( - [self._get_thread_id(index) for index in mapping_order] - ) - elif 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"] - attr_array["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) @@ -589,7 +591,7 @@ def _strip_mine( assert len(tiling_command.results) == 2 new_loop = tiling_command.results[-1] sched_state.all_loops[loop_name] = new_loop - if schedule.gpu_blocks: + 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) @@ -665,11 +667,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" ] diff --git a/src/xtc/backends/mlir/MlirNodeScheduler.py b/src/xtc/backends/mlir/MlirNodeScheduler.py index 7bd7110a4..e633537d2 100644 --- a/src/xtc/backends/mlir/MlirNodeScheduler.py +++ b/src/xtc/backends/mlir/MlirNodeScheduler.py @@ -112,12 +112,12 @@ def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None: 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_threads(axes, root) + 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 thread" assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu block" - self._plain_sch.gpu_blocks(axes, root) + self._plain_sch.gpu_block(axes, root) def get_node_schedule(self) -> MlirNodeSchedule: plain_schedule = self._plain_sch.get_plain_schedule() diff --git a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py index 99651f850..419d0f8fc 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py +++ b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py @@ -442,17 +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", - "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=" diff --git a/src/xtc/itf/schd/scheduler.py b/src/xtc/itf/schd/scheduler.py index 0fa697f95..5a64b93ee 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -321,19 +321,35 @@ def distributed_buffer_at( @abstractmethod def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: - """ - Map an axis to a certain gpu thread either x, y or z. - The mapped axis need to be in parallelization. - We can only map 3 dimensions for gpu thread + """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 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 gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: """ - Map an axis to a certain gpu block either x, y or z. - The mapped axis need to be in parallelization. - We can only map 3 dimensions for gpu block + 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) """ ... diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_1D_offload_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_1D_offload_tensor.py new file mode 100644 index 000000000..8f1169655 --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_1D_offload_tensor.py @@ -0,0 +1,212 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir.MlirGraphBackend import MlirGraphBackend as Backend + +from xtc.runtimes.accelerator.gpu import GPUDevice + +# Create device +gpu = GPUDevice() + +I, J, K, dtype = 512, 512, 512, "float32" +a = O.tensor((I, K), dtype, name="A") # A lives on the host +b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) # C must live on the accelerator + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.tile("i", {"i1": 128, "i2": 32}) +sch.tile("j", {"j1": 128, "j2": 32}) +sch.tile("k", {"k1": 64}) +sch.unroll({"i2": 2}) +sch.gpu_block(["i"]) +sch.gpu_thread(["i1"]) +sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="gpu_matmul_mlir_1D_offload_tensor", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<512x512xf32>, memref<512x512xf32>) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 0, 0](mapping = [#gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 128, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./j" : !transform.any_op +# CHECK-NEXT: %tiled_op_4, %forall_op_5 = transform.structured.tile_using_forall %tiled_linalg_op_2 tile_sizes [32, 0, 0](mapping = [#gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_5 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_op_4 tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./j1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "./k" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./k1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "./i2" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "./j2" : !transform.any_op +# CHECK-NEXT: transform.loop.unroll %loops_13 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 1, 1] : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 512] [1, 1] : memref<512x512xf32> to memref<1x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c512_7 = arith.constant 512 : index +# CHECK-NEXT: %c1_8 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_6 to %c512_7 step %c1_8 { +# CHECK-NEXT: %subview_9 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x512xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_9 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c1_1 = arith.constant 1 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c4_3 = arith.constant 4 : index +# CHECK-NEXT: %c1_4 = arith.constant 1 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c4_3, %arg10 = %c1_4, %arg11 = %c1_5) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c1_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c0_7 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg1[0, 0] [512, 512] [1, 1] : memref<512x512xf32> to memref<512x512xf32, strided<[512, 1]>> +# CHECK-NEXT: %subview_9 = memref.subview %arg2[%0, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_10 = arith.constant 0 : index +# CHECK-NEXT: %c512_11 = arith.constant 512 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: scf.for %arg15 = %c0_10 to %c512_11 step %c128 { +# CHECK-NEXT: %subview_12 = memref.subview %subview[0, 0] [128, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_8[0, %arg15] [512, 128] [1, 1] : memref<512x512xf32, strided<[512, 1]>> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_9[0, %arg15] [128, 128] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %1 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %subview_15 = memref.subview %subview_12[%1, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_13[0, 0] [512, 128] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_17 = memref.subview %subview_14[%1, 0] [32, 128] [1, 1] : memref<128x128xf32, strided<[512, 1], offset: ?>> to memref<32x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_18 = arith.constant 0 : index +# CHECK-NEXT: %c128_19 = arith.constant 128 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg16 = %c0_18 to %c128_19 step %c32 { +# CHECK-NEXT: %subview_20 = memref.subview %subview_15[0, 0] [32, 512] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_21 = memref.subview %subview_16[0, %arg16] [512, 32] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_22 = memref.subview %subview_17[0, %arg16] [32, 32] [1, 1] : memref<32x128xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_23 = arith.constant 0 : index +# CHECK-NEXT: %c512_24 = arith.constant 512 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: scf.for %arg17 = %c0_23 to %c512_24 step %c64 { +# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, %arg17] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<64x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_27 = memref.subview %subview_22[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_28 = arith.constant 0 : index +# CHECK-NEXT: %c64_29 = arith.constant 64 : index +# CHECK-NEXT: %c1_30 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_28 to %c64_29 step %c1_30 { +# CHECK-NEXT: %subview_31 = memref.subview %subview_25[0, %arg18] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_26[%arg18, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_33 = memref.subview %subview_27[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_34 = arith.constant 0 : index +# CHECK-NEXT: %c32_35 = arith.constant 32 : index +# CHECK-NEXT: %c1_36 = arith.constant 1 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index +# CHECK-NEXT: scf.for %arg19 = %c0_34 to %c32_35 step %c2 { +# CHECK-NEXT: %subview_37 = memref.subview %subview_31[%arg19, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_38 = memref.subview %subview_32[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_39 = memref.subview %subview_33[%arg19, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_40 = arith.constant 0 : index +# CHECK-NEXT: %c32_41 = arith.constant 32 : index +# CHECK-NEXT: %c1_42 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg20 = %c0_40 to %c32_41 step %c1_42 { +# CHECK-NEXT: %subview_50 = memref.subview %subview_37[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_51 = memref.subview %subview_38[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_52 = memref.subview %subview_39[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_50, %subview_51 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_52 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %c1_43 = arith.constant 1 : index +# CHECK-NEXT: %2 = arith.muli %c1_36, %c1_43 : index +# CHECK-NEXT: %3 = arith.addi %arg19, %2 : index +# CHECK-NEXT: %subview_44 = memref.subview %subview_31[%3, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_45 = memref.subview %subview_32[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_46 = memref.subview %subview_33[%3, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_47 = arith.constant 0 : index +# CHECK-NEXT: %c32_48 = arith.constant 32 : index +# CHECK-NEXT: %c1_49 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg20 = %c0_47 to %c32_48 step %c1_49 { +# CHECK-NEXT: %subview_50 = memref.subview %subview_44[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_51 = memref.subview %subview_45[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_52 = memref.subview %subview_46[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_50, %subview_51 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_52 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: } {"./j1"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 512x512xfloat32 +# CHECK-NEXT: - %1 : 512x512xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 512x512xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [512x512xfloat32, 512x512xfloat32] -> [512x512xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py index 2a02e1317..6138d3a69 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py @@ -9,7 +9,7 @@ # Create device gpu = GPUDevice() -I, J, K, dtype = 1024, 1024, 512, "float32" +I, J, K, dtype = 512, 512, 512, "float32" a = O.tensor((I, K), dtype, name="A") # A lives on the host b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator @@ -26,7 +26,6 @@ sch.tile("j", {"j1": 128, "j2": 32}) sch.tile("k", {"k1": 64}) sch.unroll({"i2": 2}) -sch.parallelize(["i", "j","i1", "j1"]) sch.gpu_block(["i", "j"]) sch.gpu_thread(["i1", "j1"]) sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) @@ -45,10 +44,10 @@ print(f"CODE: {res}") # CHECK: // -----// IR Dump Before transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<512x512xf32>, memref<512x512xf32>) outs(%arg2 : memref<512x512xf32>) # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { @@ -85,92 +84,92 @@ # CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> # CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index -# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index # CHECK-NEXT: %c1 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { -# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_5 = arith.constant 0 : index -# CHECK-NEXT: %c1024_6 = arith.constant 1024 : index -# CHECK-NEXT: %c1_7 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg4 = %c0_5 to %c1024_6 step %c1_7 { -# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_8 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 512] [1, 1] : memref<512x512xf32> to memref<1x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c512_7 = arith.constant 512 : index +# CHECK-NEXT: %c1_8 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_6 to %c512_7 step %c1_8 { +# CHECK-NEXT: %subview_9 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x512xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_9 : memref<1x1xf32, strided<[512, 1], offset: ?>>) # CHECK-NEXT: } {"./j"} # CHECK-NEXT: } {"./i"} # CHECK-NEXT: %c1_0 = arith.constant 1 : index # CHECK-NEXT: %c4 = arith.constant 4 : index # CHECK-NEXT: %c4_1 = arith.constant 4 : index # CHECK-NEXT: %c1_2 = arith.constant 1 : index -# CHECK-NEXT: %c8 = arith.constant 8 : index -# CHECK-NEXT: %c8_3 = arith.constant 8 : index -# CHECK-NEXT: %c1_4 = arith.constant 1 : index -# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c8_3, %arg11 = %c1_4) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { -# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c4_3 = arith.constant 4 : index +# CHECK-NEXT: %c4_4 = arith.constant 4 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c4_3, %arg10 = %c4_4, %arg11 = %c1_5) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { # CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c0_7 = arith.constant 0 : index # CHECK-NEXT: %block_id_x = gpu.block_id x # CHECK-NEXT: %block_id_y = gpu.block_id y # CHECK-NEXT: %block_id_z = gpu.block_id z # CHECK-NEXT: %0 = affine.apply #map(%block_id_x) # CHECK-NEXT: %1 = affine.apply #map(%block_id_y) -# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_7 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_8 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x512xf32> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<512x512xf32> to memref<128x128xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %thread_id_x = gpu.thread_id x # CHECK-NEXT: %thread_id_y = gpu.thread_id y # CHECK-NEXT: %thread_id_z = gpu.thread_id z # CHECK-NEXT: %2 = affine.apply #map1(%thread_id_x) # CHECK-NEXT: %3 = affine.apply #map1(%thread_id_y) -# CHECK-NEXT: %subview_9 = memref.subview %subview[%2, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_10 = memref.subview %subview_7[0, %3] [512, 32] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<512x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_11 = memref.subview %subview_8[%2, %3] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_12 = arith.constant 0 : index -# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %subview_10 = memref.subview %subview[%2, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[0, %3] [512, 32] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_12 = memref.subview %subview_9[%2, %3] [32, 32] [1, 1] : memref<128x128xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_13 = arith.constant 0 : index +# CHECK-NEXT: %c512_14 = arith.constant 512 : index # CHECK-NEXT: %c64 = arith.constant 64 : index -# CHECK-NEXT: scf.for %arg15 = %c0_12 to %c512 step %c64 { -# CHECK-NEXT: %subview_13 = memref.subview %subview_9[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_14 = memref.subview %subview_10[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[1024, 1], offset: ?>> to memref<64x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_15 = memref.subview %subview_11[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_16 = arith.constant 0 : index -# CHECK-NEXT: %c64_17 = arith.constant 64 : index -# CHECK-NEXT: %c1_18 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg16 = %c0_16 to %c64_17 step %c1_18 { -# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_20 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_21 = memref.subview %subview_15[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_22 = arith.constant 0 : index +# CHECK-NEXT: scf.for %arg15 = %c0_13 to %c512_14 step %c64 { +# CHECK-NEXT: %subview_15 = memref.subview %subview_10[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_11[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<64x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_17 = memref.subview %subview_12[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_18 = arith.constant 0 : index +# CHECK-NEXT: %c64_19 = arith.constant 64 : index +# CHECK-NEXT: %c1_20 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg16 = %c0_18 to %c64_19 step %c1_20 { +# CHECK-NEXT: %subview_21 = memref.subview %subview_15[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_22 = memref.subview %subview_16[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_23 = memref.subview %subview_17[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_24 = arith.constant 0 : index # CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %c1_23 = arith.constant 1 : index +# CHECK-NEXT: %c1_25 = arith.constant 1 : index # CHECK-NEXT: %c2 = arith.constant 2 : index -# CHECK-NEXT: scf.for %arg17 = %c0_22 to %c32 step %c2 { -# CHECK-NEXT: %subview_24 = memref.subview %subview_19[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_27 = arith.constant 0 : index -# CHECK-NEXT: %c32_28 = arith.constant 32 : index -# CHECK-NEXT: %c1_29 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg18 = %c0_27 to %c32_28 step %c1_29 { -# CHECK-NEXT: %subview_37 = memref.subview %subview_24[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_38 = memref.subview %subview_25[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_39 = memref.subview %subview_26[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: scf.for %arg17 = %c0_24 to %c32 step %c2 { +# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_27 = memref.subview %subview_22[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_28 = memref.subview %subview_23[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_29 = arith.constant 0 : index +# CHECK-NEXT: %c32_30 = arith.constant 32 : index +# CHECK-NEXT: %c1_31 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_29 to %c32_30 step %c1_31 { +# CHECK-NEXT: %subview_39 = memref.subview %subview_26[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_40 = memref.subview %subview_27[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_41 = memref.subview %subview_28[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_39, %subview_40 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_41 : memref<1x1xf32, strided<[512, 1], offset: ?>>) # CHECK-NEXT: } {"./j2"} -# CHECK-NEXT: %c1_30 = arith.constant 1 : index -# CHECK-NEXT: %4 = arith.muli %c1_23, %c1_30 : index +# CHECK-NEXT: %c1_32 = arith.constant 1 : index +# CHECK-NEXT: %4 = arith.muli %c1_25, %c1_32 : index # CHECK-NEXT: %5 = arith.addi %arg17, %4 : index -# CHECK-NEXT: %subview_31 = memref.subview %subview_19[%5, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_32 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_33 = memref.subview %subview_21[%5, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_34 = arith.constant 0 : index -# CHECK-NEXT: %c32_35 = arith.constant 32 : index -# CHECK-NEXT: %c1_36 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg18 = %c0_34 to %c32_35 step %c1_36 { -# CHECK-NEXT: %subview_37 = memref.subview %subview_31[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_38 = memref.subview %subview_32[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_39 = memref.subview %subview_33[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: %subview_33 = memref.subview %subview_21[%5, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_34 = memref.subview %subview_22[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_35 = memref.subview %subview_23[%5, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_36 = arith.constant 0 : index +# CHECK-NEXT: %c32_37 = arith.constant 32 : index +# CHECK-NEXT: %c1_38 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_36 to %c32_37 step %c1_38 { +# CHECK-NEXT: %subview_39 = memref.subview %subview_33[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_40 = memref.subview %subview_34[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_41 = memref.subview %subview_35[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_39, %subview_40 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_41 : memref<1x1xf32, strided<[512, 1], offset: ?>>) # CHECK-NEXT: } {"./j2"} # CHECK-NEXT: } {"./i2"} # CHECK-NEXT: } {"./k1"} @@ -185,11 +184,11 @@ # CHECK-NEXT: graph: # CHECK-NEXT: name: matmul # CHECK-NEXT: inputs: -# CHECK-NEXT: - %0 : 1024x512xfloat32 -# CHECK-NEXT: - %1 : 512x1024xfloat32 +# CHECK-NEXT: - %0 : 512x512xfloat32 +# CHECK-NEXT: - %1 : 512x512xfloat32 # CHECK-NEXT: outputs: -# CHECK-NEXT: - %2 : 1024x1024xfloat32 +# CHECK-NEXT: - %2 : 512x512xfloat32 # CHECK-NEXT: nodes: -# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [512x512xfloat32, 512x512xfloat32] -> [512x512xfloat32] # CHECK-NEXT: # CHECK-NEXT: CODE: 0 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 0194b7262..2fcd2d325 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 @@ -9,7 +9,7 @@ # Create device gpu = GPUDevice() -I, J, K, dtype = 4, 32, 512, "float32" +I, J, K, dtype = 512, 512, 512, "float32" a = O.tensor((I, K), dtype, name="A") # A lives on the host b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator @@ -22,11 +22,15 @@ impl = Backend(graph) sch = impl.get_scheduler() -sch.tile("i", {"i1": 2}) -sch.tile("j", {"j1": 16}) -sch.unroll({"i1": 2}) -sch.vectorize(["j1"]) -sch.parallelize(["i"]) +sch.tile("i", {"i1": 128, "i2": 32}) +sch.tile("j", {"j1": 128, "j2": 32}) +sch.tile("k", {"k1": 64}) +sch.unroll({"i2": 2}) +sch.gpu_block(["i", "j"]) +sch.gpu_thread(["i1", "j1"]) +sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) +sched = sch.schedule() +sch.vectorize(["j2"]) sched = sch.schedule() comp = impl.get_compiler( @@ -42,10 +46,10 @@ print(f"CODE: {res}") # CHECK: // -----// IR Dump Before transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias, memref.on_device}, %arg2: memref<4x32xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<4x32xf32>) -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<4x512xf32>, memref<512x32xf32>) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<512x512xf32>, memref<512x512xf32>) outs(%arg2 : memref<512x512xf32>) # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { @@ -59,87 +63,106 @@ # CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op # CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op -# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [2, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 16, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_3 "./j" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [32, 32, 0](mapping = [#gpu.thread, #gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_3 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_7 "./i1" : !transform.any_op -# CHECK-NEXT: transform.include @_vecto failures(suppress) (%tiled_linalg_op_6) : (!transform.any_op) -> () -# CHECK-NEXT: transform.loop.unroll %loops_7 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# 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: transform.apply_patterns to %2 { # 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: transform.apply_patterns to %2 { -# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct -# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction -# 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: transform.yield # CHECK-NEXT: } # CHECK-NEXT: } # CHECK-NEXT: # CHECK-NEXT: // -----// IR Dump After transform //----- // -# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 2)> +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2) -> (d0, d2)> +# CHECK-NEXT: #map3 = affine_map<(d0, d1, d2) -> (d2, d1)> +# CHECK-NEXT: #map4 = affine_map<(d0, d1, d2) -> (d0, d1)> # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias, memref.on_device}, %arg2: memref<4x32xf32> {llvm.noalias, memref.on_device}) { -# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<1x16xf32> +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %0 = ub.poison : f32 -# CHECK-NEXT: %c512 = arith.constant 512 : index -# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index # CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %cst_0 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index -# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index # CHECK-NEXT: %c1 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg3 = %c0 to %c4 step %c1 { -# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 32] [1, 1] : memref<4x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c1 { -# CHECK-NEXT: %subview_1 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst_0 : f32) outs(%subview_1 : memref<1x1xf32, strided<[32, 1], offset: ?>>) -# CHECK-NEXT: } {"./j"} -# CHECK-NEXT: } {"./i"} -# CHECK-NEXT: scf.forall (%arg3) in (2) { -# CHECK-NEXT: %1 = affine.apply #map(%arg3) -# CHECK-NEXT: %subview = memref.subview %arg0[%1, 0] [2, 512] [1, 1] : memref<4x512xf32> to memref<2x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_1 = memref.subview %arg1[0, 0] [512, 32] [1, 1] : memref<512x32xf32> to memref<512x32xf32, strided<[32, 1]>> -# CHECK-NEXT: %subview_2 = memref.subview %arg2[%1, 0] [2, 32] [1, 1] : memref<4x32xf32> to memref<2x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c16 { -# CHECK-NEXT: %subview_3 = memref.subview %subview_1[0, %arg4] [512, 16] [1, 1] : memref<512x32xf32, strided<[32, 1]>> to memref<512x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_4 = memref.subview %subview_2[0, %arg4] [2, 16] [1, 1] : memref<2x32xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg5 = %c0 to %c512 step %c1 { -# CHECK-NEXT: %subview_5 = memref.subview %subview[0, %arg5] [2, 1] [1, 1] : memref<2x512xf32, strided<[512, 1], offset: ?>> to memref<2x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_6 = memref.subview %subview_3[%arg5, 0] [1, 16] [1, 1] : memref<512x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_7 = memref.subview %subview_5[%c0, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_8 = memref.subview %subview_4[%c0, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %2 = vector.transfer_read %subview_7[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> -# CHECK-NEXT: %3 = vector.transfer_read %subview_6[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %4 = vector.transfer_read %subview_8[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %5 = vector.extract %3[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %6 = vector.extract %2[0, 0] : f32 from vector<1x1xf32> -# CHECK-NEXT: %7 = vector.broadcast %6 : f32 to vector<16xf32> -# CHECK-NEXT: %8 = vector.extract %4[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %9 = vector.fma %7, %5, %8 : vector<16xf32> -# CHECK-NEXT: %10 = vector.insert %9, %cst [0] : vector<16xf32> into vector<1x16xf32> -# CHECK-NEXT: vector.transfer_write %10, %subview_8[%c0, %c0] {in_bounds = [true, true]} : vector<1x16xf32>, memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_9 = memref.subview %subview_5[%c1, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_10 = memref.subview %subview_4[%c1, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %11 = vector.transfer_read %subview_9[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> -# CHECK-NEXT: %12 = vector.transfer_read %subview_6[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %13 = vector.transfer_read %subview_10[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %14 = vector.extract %12[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %15 = vector.extract %11[0, 0] : f32 from vector<1x1xf32> -# CHECK-NEXT: %16 = vector.broadcast %15 : f32 to vector<16xf32> -# CHECK-NEXT: %17 = vector.extract %13[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %18 = vector.fma %16, %14, %17 : vector<16xf32> -# CHECK-NEXT: %19 = vector.insert %18, %cst [0] : vector<16xf32> into vector<1x16xf32> -# CHECK-NEXT: vector.transfer_write %19, %subview_10[%c0, %c0] {in_bounds = [true, true]} : vector<1x16xf32>, memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 512] [1, 1] : memref<512x512xf32> to memref<1x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview_6 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x512xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_6 : memref<1x1xf32, strided<[512, 1], offset: ?>>) # CHECK-NEXT: } {"./j"} # CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c4_1 = arith.constant 4 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c4_3 = arith.constant 4 : index +# CHECK-NEXT: %c4_4 = arith.constant 4 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c4_3, %arg10 = %c4_4, %arg11 = %c1_5) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c0_7 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %1 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %2 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%1, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg1[0, %2] [512, 128] [1, 1] : memref<512x512xf32> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %arg2[%1, %2] [128, 128] [1, 1] : memref<512x512xf32> to memref<128x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %4 = affine.apply #map1(%thread_id_y) +# CHECK-NEXT: %subview_10 = memref.subview %subview[%3, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[0, %4] [512, 32] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_12 = memref.subview %subview_9[%3, %4] [32, 32] [1, 1] : memref<128x128xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg15 = %c0 to %c512 step %c64 { +# CHECK-NEXT: %subview_13 = memref.subview %subview_10[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_11[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<64x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg16 = %c0 to %c64 step %c1 { +# CHECK-NEXT: %subview_15 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg17 = %c0 to %c32 step %c2 { +# CHECK-NEXT: %subview_17 = memref.subview %subview_15[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_18 = memref.subview %subview_12[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %5 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %6 = vector.transfer_read %subview_16[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x32xf32, strided<[512, 1], offset: ?>>, vector<1x32xf32> +# CHECK-NEXT: %7 = vector.transfer_read %subview_18[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x32xf32, strided<[512, 1], offset: ?>>, vector<1x32xf32> +# CHECK-NEXT: %8 = vector.contract {indexing_maps = [#map2, #map3, #map4], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind} %5, %6, %7 : vector<1x1xf32>, vector<1x32xf32> into vector<1x32xf32> +# CHECK-NEXT: vector.transfer_write %8, %subview_18[%c0, %c0] {in_bounds = [true, true]} : vector<1x32xf32>, memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %9 = arith.addi %arg17, %c1 : index +# CHECK-NEXT: %subview_19 = memref.subview %subview_15[%9, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_12[%9, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %10 = vector.transfer_read %subview_19[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %11 = vector.transfer_read %subview_16[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x32xf32, strided<[512, 1], offset: ?>>, vector<1x32xf32> +# CHECK-NEXT: %12 = vector.transfer_read %subview_20[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x32xf32, strided<[512, 1], offset: ?>>, vector<1x32xf32> +# CHECK-NEXT: %13 = vector.contract {indexing_maps = [#map2, #map3, #map4], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind} %10, %11, %12 : vector<1x1xf32>, vector<1x32xf32> into vector<1x32xf32> +# CHECK-NEXT: vector.transfer_write %13, %subview_20[%c0, %c0] {in_bounds = [true, true]} : vector<1x32xf32>, memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: } @@ -147,11 +170,11 @@ # CHECK-NEXT: graph: # CHECK-NEXT: name: matmul # CHECK-NEXT: inputs: -# CHECK-NEXT: - %0 : 4x512xfloat32 -# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: - %0 : 512x512xfloat32 +# CHECK-NEXT: - %1 : 512x512xfloat32 # CHECK-NEXT: outputs: -# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: - %2 : 512x512xfloat32 # CHECK-NEXT: nodes: -# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [512x512xfloat32, 512x512xfloat32] -> [512x512xfloat32] # CHECK-NEXT: # CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py index 0858d6577..32976c805 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py @@ -26,7 +26,6 @@ sch.tile("j", {"j1": 128, "j2": 32}) sch.tile("k", {"k1": 64}) # sch.unroll({"i2": 2}) -sch.parallelize(["i", "j","i1", "j1"]) sch.gpu_block(["i", "j"]) sch.gpu_thread(["i1", "j1"]) sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) @@ -46,16 +45,24 @@ print(f"CODE: {res}") # CHECK: // -----// IR Dump Before transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: tensor<1024x512xf32> {llvm.noalias}, %arg1: tensor<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %0 = tensor.empty() : tensor<1024x1024xf32> # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: %1 = linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%0 : tensor<1024x1024xf32>) -> tensor<1024x1024xf32> +# CHECK-NEXT: %2 = linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : tensor<1024x512xf32>, tensor<512x1024xf32>) outs(%1 : tensor<1024x1024xf32>) -> tensor<1024x1024xf32> +# CHECK-NEXT: bufferization.materialize_in_destination %2 in restrict writable %arg2 : (tensor<1024x1024xf32>, memref<1024x1024xf32>) -> () # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { # CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op # CHECK-NEXT: transform.yield # CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_post_bufferize(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {"./i", mapping = [#gpu.block, #gpu.block]} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %1 = transform.gpu.map_forall_to_blocks %0 generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_nested_forall_to_threads %1 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { # CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op # CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) @@ -75,9 +82,6 @@ # CHECK-NEXT: transform.annotate %loops_9 "./i2" : !transform.any_op # CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_11 "./j2" : !transform.any_op -# CHECK-NEXT: transform.loop.unroll %loops_9 {factor = 2 : i64} : !transform.any_op -# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op -# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op # CHECK-NEXT: transform.yield # CHECK-NEXT: } # CHECK-NEXT: } @@ -86,20 +90,126 @@ # CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> # CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: tensor<1024x512xf32> {llvm.noalias}, %arg1: tensor<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %0 = tensor.empty() : tensor<1024x1024xf32> # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index # CHECK-NEXT: %c1024 = arith.constant 1024 : index # CHECK-NEXT: %c1 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { -# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_5 = arith.constant 0 : index -# CHECK-NEXT: %c1024_6 = arith.constant 1024 : index -# CHECK-NEXT: %c1_7 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg4 = %c0_5 to %c1024_6 step %c1_7 { -# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_8 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: %1 = scf.for %arg3 = %c0 to %c1024 step %c1 iter_args(%arg4 = %0) -> (tensor<1024x1024xf32>) { +# CHECK-NEXT: %extracted_slice = tensor.extract_slice %arg4[%arg3, 0] [1, 1024] [1, 1] : tensor<1024x1024xf32> to tensor<1x1024xf32> +# CHECK-NEXT: %c0_0 = arith.constant 0 : index +# CHECK-NEXT: %c1024_1 = arith.constant 1024 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %3 = scf.for %arg5 = %c0_0 to %c1024_1 step %c1_2 iter_args(%arg6 = %extracted_slice) -> (tensor<1x1024xf32>) { +# CHECK-NEXT: %extracted_slice_3 = tensor.extract_slice %arg6[0, %arg5] [1, 1] [1, 1] : tensor<1x1024xf32> to tensor<1x1xf32> +# CHECK-NEXT: %4 = linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%extracted_slice_3 : tensor<1x1xf32>) -> tensor<1x1xf32> +# CHECK-NEXT: %inserted_slice_4 = tensor.insert_slice %4 into %arg6[0, %arg5] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x1024xf32> +# CHECK-NEXT: scf.yield %inserted_slice_4 : tensor<1x1024xf32> +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: %inserted_slice = tensor.insert_slice %3 into %arg4[%arg3, 0] [1, 1024] [1, 1] : tensor<1x1024xf32> into tensor<1024x1024xf32> +# CHECK-NEXT: scf.yield %inserted_slice : tensor<1024x1024xf32> +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %2 = scf.forall (%arg3, %arg4) in (8, 8) shared_outs(%arg5 = %1) -> (tensor<1024x1024xf32>) { +# CHECK-NEXT: %3 = affine.apply #map(%arg3) +# CHECK-NEXT: %4 = affine.apply #map(%arg4) +# CHECK-NEXT: %extracted_slice = tensor.extract_slice %arg0[%3, 0] [128, 512] [1, 1] : tensor<1024x512xf32> to tensor<128x512xf32> +# CHECK-NEXT: %extracted_slice_0 = tensor.extract_slice %arg1[0, %4] [512, 128] [1, 1] : tensor<512x1024xf32> to tensor<512x128xf32> +# CHECK-NEXT: %extracted_slice_1 = tensor.extract_slice %arg5[%3, %4] [128, 128] [1, 1] : tensor<1024x1024xf32> to tensor<128x128xf32> +# CHECK-NEXT: %5 = scf.forall (%arg6, %arg7) in (4, 4) shared_outs(%arg8 = %extracted_slice_1) -> (tensor<128x128xf32>) { +# CHECK-NEXT: %6 = affine.apply #map1(%arg6) +# CHECK-NEXT: %7 = affine.apply #map1(%arg7) +# CHECK-NEXT: %extracted_slice_2 = tensor.extract_slice %extracted_slice[%6, 0] [32, 512] [1, 1] : tensor<128x512xf32> to tensor<32x512xf32> +# CHECK-NEXT: %extracted_slice_3 = tensor.extract_slice %extracted_slice_0[0, %7] [512, 32] [1, 1] : tensor<512x128xf32> to tensor<512x32xf32> +# CHECK-NEXT: %extracted_slice_4 = tensor.extract_slice %arg8[%6, %7] [32, 32] [1, 1] : tensor<128x128xf32> to tensor<32x32xf32> +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %8 = scf.for %arg9 = %c0_5 to %c512 step %c64 iter_args(%arg10 = %extracted_slice_4) -> (tensor<32x32xf32>) { +# CHECK-NEXT: %extracted_slice_6 = tensor.extract_slice %extracted_slice_2[0, %arg9] [32, 64] [1, 1] : tensor<32x512xf32> to tensor<32x64xf32> +# CHECK-NEXT: %extracted_slice_7 = tensor.extract_slice %extracted_slice_3[%arg9, 0] [64, 32] [1, 1] : tensor<512x32xf32> to tensor<64x32xf32> +# CHECK-NEXT: %extracted_slice_8 = tensor.extract_slice %arg10[0, 0] [32, 32] [1, 1] : tensor<32x32xf32> to tensor<32x32xf32> +# CHECK-NEXT: %c0_9 = arith.constant 0 : index +# CHECK-NEXT: %c64_10 = arith.constant 64 : index +# CHECK-NEXT: %c1_11 = arith.constant 1 : index +# CHECK-NEXT: %9 = scf.for %arg11 = %c0_9 to %c64_10 step %c1_11 iter_args(%arg12 = %extracted_slice_8) -> (tensor<32x32xf32>) { +# CHECK-NEXT: %extracted_slice_12 = tensor.extract_slice %extracted_slice_6[0, %arg11] [32, 1] [1, 1] : tensor<32x64xf32> to tensor<32x1xf32> +# CHECK-NEXT: %extracted_slice_13 = tensor.extract_slice %extracted_slice_7[%arg11, 0] [1, 32] [1, 1] : tensor<64x32xf32> to tensor<1x32xf32> +# CHECK-NEXT: %extracted_slice_14 = tensor.extract_slice %arg12[0, 0] [32, 32] [1, 1] : tensor<32x32xf32> to tensor<32x32xf32> +# CHECK-NEXT: %c0_15 = arith.constant 0 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c1_16 = arith.constant 1 : index +# CHECK-NEXT: %10 = scf.for %arg13 = %c0_15 to %c32 step %c1_16 iter_args(%arg14 = %extracted_slice_14) -> (tensor<32x32xf32>) { +# CHECK-NEXT: %extracted_slice_18 = tensor.extract_slice %extracted_slice_12[%arg13, 0] [1, 1] [1, 1] : tensor<32x1xf32> to tensor<1x1xf32> +# CHECK-NEXT: %extracted_slice_19 = tensor.extract_slice %extracted_slice_13[0, 0] [1, 32] [1, 1] : tensor<1x32xf32> to tensor<1x32xf32> +# CHECK-NEXT: %extracted_slice_20 = tensor.extract_slice %arg14[%arg13, 0] [1, 32] [1, 1] : tensor<32x32xf32> to tensor<1x32xf32> +# CHECK-NEXT: %c0_21 = arith.constant 0 : index +# CHECK-NEXT: %c32_22 = arith.constant 32 : index +# CHECK-NEXT: %c1_23 = arith.constant 1 : index +# CHECK-NEXT: %11 = scf.for %arg15 = %c0_21 to %c32_22 step %c1_23 iter_args(%arg16 = %extracted_slice_20) -> (tensor<1x32xf32>) { +# CHECK-NEXT: %extracted_slice_25 = tensor.extract_slice %extracted_slice_18[0, 0] [1, 1] [1, 1] : tensor<1x1xf32> to tensor<1x1xf32> +# CHECK-NEXT: %extracted_slice_26 = tensor.extract_slice %extracted_slice_19[0, %arg15] [1, 1] [1, 1] : tensor<1x32xf32> to tensor<1x1xf32> +# CHECK-NEXT: %extracted_slice_27 = tensor.extract_slice %arg16[0, %arg15] [1, 1] [1, 1] : tensor<1x32xf32> to tensor<1x1xf32> +# CHECK-NEXT: %12 = linalg.matmul {__xtc_id_C_} ins(%extracted_slice_25, %extracted_slice_26 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%extracted_slice_27 : tensor<1x1xf32>) -> tensor<1x1xf32> +# CHECK-NEXT: %inserted_slice_28 = tensor.insert_slice %12 into %arg16[0, %arg15] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x32xf32> +# CHECK-NEXT: scf.yield %inserted_slice_28 : tensor<1x32xf32> +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %inserted_slice_24 = tensor.insert_slice %11 into %arg14[%arg13, 0] [1, 32] [1, 1] : tensor<1x32xf32> into tensor<32x32xf32> +# CHECK-NEXT: scf.yield %inserted_slice_24 : tensor<32x32xf32> +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: %inserted_slice_17 = tensor.insert_slice %10 into %arg12[0, 0] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<32x32xf32> +# CHECK-NEXT: scf.yield %inserted_slice_17 : tensor<32x32xf32> +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: %inserted_slice = tensor.insert_slice %9 into %arg10[0, 0] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<32x32xf32> +# CHECK-NEXT: scf.yield %inserted_slice : tensor<32x32xf32> +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: scf.forall.in_parallel { +# CHECK-NEXT: tensor.parallel_insert_slice %8 into %arg8[%6, %7] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<128x128xf32> +# CHECK-NEXT: } +# CHECK-NEXT: } {"./i1", mapping = [#gpu.thread, #gpu.thread]} +# CHECK-NEXT: scf.forall.in_parallel { +# CHECK-NEXT: tensor.parallel_insert_slice %5 into %arg5[%3, %4] [128, 128] [1, 1] : tensor<128x128xf32> into tensor<1024x1024xf32> +# CHECK-NEXT: } +# CHECK-NEXT: } {"./i", mapping = [#gpu.block, #gpu.block]} +# CHECK-NEXT: bufferization.materialize_in_destination %2 in restrict writable %arg2 : (tensor<1024x1024xf32>, memref<1024x1024xf32>) -> () +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_post_bufferize(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {"./i", mapping = [#gpu.block, #gpu.block]} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %1 = transform.gpu.map_forall_to_blocks %0 generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_nested_forall_to_threads %1 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After Tensor Lowering //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %0 = scf.for %arg3 = %c0 to %c1024 step %c1 iter_args(%arg4 = %arg2) -> (memref<1024x1024xf32>) { +# CHECK-NEXT: %subview = memref.subview %arg4[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %1 = scf.for %arg5 = %c0 to %c1024 step %c1 iter_args(%arg6 = %subview) -> (memref<1x1024xf32, strided<[1024, 1], offset: ?>>) { +# CHECK-NEXT: %subview_6 = memref.subview %arg6[0, %arg5] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_6 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: %subview_7 = memref.subview %arg6[0, %arg5] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: memref.copy %subview_6, %subview_7 : memref<1x1xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg6 : memref<1x1024xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: } {"./j"} +# CHECK-NEXT: %subview_5 = memref.subview %arg4[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: memref.copy %1, %subview_5 : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg4 : memref<1024x1024xf32> # CHECK-NEXT: } {"./i"} # CHECK-NEXT: %c1_0 = arith.constant 1 : index # CHECK-NEXT: %c4 = arith.constant 4 : index @@ -114,71 +224,52 @@ # CHECK-NEXT: %block_id_x = gpu.block_id x # CHECK-NEXT: %block_id_y = gpu.block_id y # CHECK-NEXT: %block_id_z = gpu.block_id z -# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) -# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) -# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_7 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_8 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %1 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %2 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%1, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %arg1[0, %2] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %0[%1, %2] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: %thread_id_x = gpu.thread_id x # CHECK-NEXT: %thread_id_y = gpu.thread_id y # CHECK-NEXT: %thread_id_z = gpu.thread_id z -# CHECK-NEXT: %2 = affine.apply #map1(%thread_id_x) -# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_y) -# CHECK-NEXT: %subview_9 = memref.subview %subview[%2, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_10 = memref.subview %subview_7[0, %3] [512, 32] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<512x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_11 = memref.subview %subview_8[%2, %3] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_12 = arith.constant 0 : index -# CHECK-NEXT: %c512 = arith.constant 512 : index -# CHECK-NEXT: %c64 = arith.constant 64 : index -# CHECK-NEXT: scf.for %arg15 = %c0_12 to %c512 step %c64 { -# CHECK-NEXT: %subview_13 = memref.subview %subview_9[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_14 = memref.subview %subview_10[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[1024, 1], offset: ?>> to memref<64x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_15 = memref.subview %subview_11[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_16 = arith.constant 0 : index -# CHECK-NEXT: %c64_17 = arith.constant 64 : index -# CHECK-NEXT: %c1_18 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg16 = %c0_16 to %c64_17 step %c1_18 { -# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_20 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_21 = memref.subview %subview_15[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_22 = arith.constant 0 : index -# CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %c1_23 = arith.constant 1 : index -# CHECK-NEXT: %c2 = arith.constant 2 : index -# CHECK-NEXT: scf.for %arg17 = %c0_22 to %c32 step %c2 { -# CHECK-NEXT: %subview_24 = memref.subview %subview_19[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_27 = arith.constant 0 : index -# CHECK-NEXT: %c32_28 = arith.constant 32 : index -# CHECK-NEXT: %c1_29 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg18 = %c0_27 to %c32_28 step %c1_29 { -# CHECK-NEXT: %subview_37 = memref.subview %subview_24[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_38 = memref.subview %subview_25[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_39 = memref.subview %subview_26[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) -# CHECK-NEXT: } {"./j2"} -# CHECK-NEXT: %c1_30 = arith.constant 1 : index -# CHECK-NEXT: %4 = arith.muli %c1_23, %c1_30 : index -# CHECK-NEXT: %5 = arith.addi %arg17, %4 : index -# CHECK-NEXT: %subview_31 = memref.subview %subview_19[%5, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_32 = memref.subview %subview_20[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_33 = memref.subview %subview_21[%5, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_34 = arith.constant 0 : index -# CHECK-NEXT: %c32_35 = arith.constant 32 : index -# CHECK-NEXT: %c1_36 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg18 = %c0_34 to %c32_35 step %c1_36 { -# CHECK-NEXT: %subview_37 = memref.subview %subview_31[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_38 = memref.subview %subview_32[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_39 = memref.subview %subview_33[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_37, %subview_38 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_39 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %4 = affine.apply #map1(%thread_id_y) +# CHECK-NEXT: %subview_9 = memref.subview %subview[%3, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_7[0, %4] [512, 32] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<512x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[%3, %4] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %5 = scf.for %arg15 = %c0 to %c512 step %c64 iter_args(%arg16 = %subview_11) -> (memref<32x32xf32, strided<[1024, 1], offset: ?>>) { +# CHECK-NEXT: %subview_14 = memref.subview %subview_9[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_10[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[1024, 1], offset: ?>> to memref<64x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %6 = scf.for %arg17 = %c0 to %c64 step %c1 iter_args(%arg18 = %arg16) -> (memref<32x32xf32, strided<[1024, 1], offset: ?>>) { +# CHECK-NEXT: %subview_16 = memref.subview %subview_14[0, %arg17] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_17 = memref.subview %subview_15[%arg17, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %7 = scf.for %arg19 = %c0 to %c32 step %c1 iter_args(%arg20 = %arg18) -> (memref<32x32xf32, strided<[1024, 1], offset: ?>>) { +# CHECK-NEXT: %subview_18 = memref.subview %subview_16[%arg19, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_19 = memref.subview %arg20[%arg19, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %8 = scf.for %arg21 = %c0 to %c32 step %c1 iter_args(%arg22 = %subview_19) -> (memref<1x32xf32, strided<[1024, 1], offset: ?>>) { +# CHECK-NEXT: %subview_21 = memref.subview %subview_17[0, %arg21] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_22 = memref.subview %arg22[0, %arg21] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_18, %subview_21 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_22 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: %subview_23 = memref.subview %arg22[0, %arg21] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: memref.copy %subview_22, %subview_23 : memref<1x1xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg22 : memref<1x32xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %subview_20 = memref.subview %arg20[%arg19, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: memref.copy %8, %subview_20 : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg20 : memref<32x32xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: scf.yield %7 : memref<32x32xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: scf.yield %6 : memref<32x32xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: } {"./k"} +# CHECK-NEXT: %subview_12 = memref.subview %subview_8[%3, %4] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: memref.copy %5, %subview_12 : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: gpu.barrier +# CHECK-NEXT: %subview_13 = memref.subview %0[%1, %2] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: memref.copy %subview_8, %subview_13 : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> # CHECK-NEXT: gpu.terminator # CHECK-NEXT: } +# CHECK-NEXT: memref.copy %0, %arg2 : memref<1024x1024xf32> to memref<1024x1024xf32> # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: } diff --git a/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py b/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py index ab671deab..cb04ff23a 100644 --- a/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py +++ b/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py @@ -5,7 +5,7 @@ from xtc.backends.mlir import Backend from sys import platform -I, J, K, dtype = 32, 32, 512, "float32" +I, J, K, dtype = 256, 32, 512, "float32" a = O.tensor((I, K), dtype, name="A") b = O.tensor((K, J), dtype, name="B") @@ -17,11 +17,11 @@ impl = Backend(graph) sch = impl.get_scheduler() -sch.tile("i", {"i1": 16}) +sch.tile("i", {"i1": 128, "i2": 16}) sch.tile("j", {"j1": 16}) sch.vectorize(["j1"]) -sch.unroll({"i1": 2}) -sch.parallelize(["i"]) +sch.gpu_block(["i"]) +sch.gpu_thread(["i1"]) sched = sch.schedule() comp = impl.get_compiler( diff --git a/tests/filecheck/schedules/test_matmul_descript_gpu.py b/tests/filecheck/schedules/test_matmul_descript_gpu.py index 7a1080d50..28c991b66 100644 --- a/tests/filecheck/schedules/test_matmul_descript_gpu.py +++ b/tests/filecheck/schedules/test_matmul_descript_gpu.py @@ -26,11 +26,11 @@ node_name = "C", abstract_dims = ["I","J","K"], spec = { - "I": {"parallelize": True, "gpu_block": 0}, - "J": {"parallelize": True, "gpu_block": 1}, + "I": {"gpu_block": 0}, + "J": {"gpu_block": 1}, "K": {}, - "I#128": {"parallelize": True, "gpu_thread": 0}, - "J#128": {"parallelize": True, "gpu_thread": 1}, + "I#128": {"gpu_thread": 0}, + "J#128": {"gpu_thread": 1}, "I#32": {}, "J#32": {}, @@ -53,10 +53,10 @@ # CHECK: // -----// IR Dump Before transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<4x32xf32> {llvm.noalias}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<4x32xf32>) -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<4x512xf32>, memref<512x32xf32>) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { @@ -70,86 +70,98 @@ # CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op # CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %1 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_3 "C/K" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [2, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_5 "C/I" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 16, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_7 "C/J" : !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: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %1 tile_sizes [128, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "C/I" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 128, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "C/J" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "C/K" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [32, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_9 "C/I0" : !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 %loops_3 {isolated_from_above} : (!transform.any_op) -> !transform.any_op -# CHECK-NEXT: transform.apply_patterns to %2 { -# 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: transform.apply_patterns to %2 { -# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct -# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction -# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "C/J0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "C/I1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "C/J1" : !transform.any_op # CHECK-NEXT: transform.yield # CHECK-NEXT: } # CHECK-NEXT: } # CHECK-NEXT: # CHECK-NEXT: // -----// IR Dump After transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<4x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<4x32xf32> {llvm.noalias}) { -# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<1x16xf32> -# CHECK-NEXT: %0 = ub.poison : f32 -# CHECK-NEXT: %c16 = arith.constant 16 : index -# CHECK-NEXT: %c2 = arith.constant 2 : index -# CHECK-NEXT: %c512 = arith.constant 512 : index -# CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %cst_0 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index -# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index # CHECK-NEXT: %c1 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg3 = %c0 to %c4 step %c1 { -# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 32] [1, 1] : memref<4x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c1 { -# CHECK-NEXT: %subview_1 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst_0 : f32) outs(%subview_1 : memref<1x1xf32, strided<[32, 1], offset: ?>>) +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_2 = arith.constant 0 : index +# CHECK-NEXT: %c1024_3 = arith.constant 1024 : index +# CHECK-NEXT: %c1_4 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_2 to %c1024_3 step %c1_4 { +# CHECK-NEXT: %subview_5 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_5 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) # CHECK-NEXT: } {"./j"} # CHECK-NEXT: } {"./i"} -# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { -# CHECK-NEXT: %subview = memref.subview %arg0[0, %arg3] [4, 1] [1, 1] : memref<4x512xf32> to memref<4x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_1 = memref.subview %arg1[%arg3, 0] [1, 32] [1, 1] : memref<512x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_2 = memref.subview %arg2[0, 0] [4, 32] [1, 1] : memref<4x32xf32> to memref<4x32xf32, strided<[32, 1]>> -# CHECK-NEXT: scf.for %arg4 = %c0 to %c4 step %c2 { -# CHECK-NEXT: %subview_3 = memref.subview %subview[%arg4, 0] [2, 1] [1, 1] : memref<4x1xf32, strided<[512, 1], offset: ?>> to memref<2x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_4 = memref.subview %subview_2[%arg4, 0] [2, 32] [1, 1] : memref<4x32xf32, strided<[32, 1]>> to memref<2x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg5 = %c0 to %c32 step %c16 { -# CHECK-NEXT: %subview_5 = memref.subview %subview_1[0, %arg5] [1, 16] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_6 = memref.subview %subview_4[0, %arg5] [2, 16] [1, 1] : memref<2x32xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_7 = memref.subview %subview_3[%c0, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_8 = memref.subview %subview_6[%c0, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %1 = vector.transfer_read %subview_7[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> -# CHECK-NEXT: %2 = vector.transfer_read %subview_5[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %3 = vector.transfer_read %subview_8[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %4 = vector.extract %2[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %5 = vector.extract %1[0, 0] : f32 from vector<1x1xf32> -# CHECK-NEXT: %6 = vector.broadcast %5 : f32 to vector<16xf32> -# CHECK-NEXT: %7 = vector.extract %3[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %8 = vector.fma %6, %4, %7 : vector<16xf32> -# CHECK-NEXT: %9 = vector.insert %8, %cst [0] : vector<16xf32> into vector<1x16xf32> -# CHECK-NEXT: vector.transfer_write %9, %subview_8[%c0, %c0] {in_bounds = [true, true]} : vector<1x16xf32>, memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_9 = memref.subview %subview_3[%c1, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_10 = memref.subview %subview_6[%c1, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %10 = vector.transfer_read %subview_9[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> -# CHECK-NEXT: %11 = vector.transfer_read %subview_5[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %12 = vector.transfer_read %subview_10[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x16xf32, strided<[32, 1], offset: ?>>, vector<1x16xf32> -# CHECK-NEXT: %13 = vector.extract %11[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %14 = vector.extract %10[0, 0] : f32 from vector<1x1xf32> -# CHECK-NEXT: %15 = vector.broadcast %14 : f32 to vector<16xf32> -# CHECK-NEXT: %16 = vector.extract %12[0] : vector<16xf32> from vector<1x16xf32> -# CHECK-NEXT: %17 = vector.fma %15, %13, %16 : vector<16xf32> -# CHECK-NEXT: %18 = vector.insert %17, %cst [0] : vector<16xf32> into vector<1x16xf32> -# CHECK-NEXT: vector.transfer_write %18, %subview_10[%c0, %c0] {in_bounds = [true, true]} : vector<1x16xf32>, memref<1x16xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: } {"C/J"} -# CHECK-NEXT: } {"C/I"} -# CHECK-NEXT: } {"C/K"} +# CHECK-NEXT: %c0_0 = arith.constant 0 : index +# CHECK-NEXT: %c1024_1 = arith.constant 1024 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: scf.for %arg3 = %c0_0 to %c1024_1 step %c128 { +# CHECK-NEXT: %subview = memref.subview %arg0[%arg3, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_2 = memref.subview %arg1[0, 0] [512, 1024] [1, 1] : memref<512x1024xf32> to memref<512x1024xf32, strided<[1024, 1]>> +# CHECK-NEXT: %subview_3 = memref.subview %arg2[%arg3, 0] [128, 1024] [1, 1] : memref<1024x1024xf32> to memref<128x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_4 = arith.constant 0 : index +# CHECK-NEXT: %c1024_5 = arith.constant 1024 : index +# CHECK-NEXT: %c128_6 = arith.constant 128 : index +# CHECK-NEXT: scf.for %arg4 = %c0_4 to %c1024_5 step %c128_6 { +# CHECK-NEXT: %subview_7 = memref.subview %subview[0, 0] [128, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %subview_2[0, %arg4] [512, 128] [1, 1] : memref<512x1024xf32, strided<[1024, 1]>> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_3[0, %arg4] [128, 128] [1, 1] : memref<128x1024xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_10 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1_11 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg5 = %c0_10 to %c512 step %c1_11 { +# CHECK-NEXT: %subview_12 = memref.subview %subview_7[0, %arg5] [128, 1] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_8[%arg5, 0] [1, 128] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_9[0, 0] [128, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_15 = arith.constant 0 : index +# CHECK-NEXT: %c128_16 = arith.constant 128 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg6 = %c0_15 to %c128_16 step %c32 { +# CHECK-NEXT: %subview_17 = memref.subview %subview_12[%arg6, 0] [32, 1] [1, 1] : memref<128x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_18 = memref.subview %subview_13[0, 0] [1, 128] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_19 = memref.subview %subview_14[%arg6, 0] [32, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_20 = arith.constant 0 : index +# CHECK-NEXT: %c128_21 = arith.constant 128 : index +# CHECK-NEXT: %c32_22 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg7 = %c0_20 to %c128_21 step %c32_22 { +# CHECK-NEXT: %subview_23 = memref.subview %subview_17[0, 0] [32, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_24 = memref.subview %subview_18[0, %arg7] [1, 32] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_25 = memref.subview %subview_19[0, %arg7] [32, 32] [1, 1] : memref<32x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_26 = arith.constant 0 : index +# CHECK-NEXT: %c32_27 = arith.constant 32 : index +# CHECK-NEXT: %c1_28 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg8 = %c0_26 to %c32_27 step %c1_28 { +# CHECK-NEXT: %subview_29 = memref.subview %subview_23[%arg8, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_30 = memref.subview %subview_24[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_31 = memref.subview %subview_25[%arg8, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_32 = arith.constant 0 : index +# CHECK-NEXT: %c32_33 = arith.constant 32 : index +# CHECK-NEXT: %c1_34 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg9 = %c0_32 to %c32_33 step %c1_34 { +# CHECK-NEXT: %subview_35 = memref.subview %subview_29[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_36 = memref.subview %subview_30[0, %arg9] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_37 = memref.subview %subview_31[0, %arg9] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_35, %subview_36 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_37 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"C/J1"} +# CHECK-NEXT: } {"C/I1"} +# CHECK-NEXT: } {"C/J0"} +# CHECK-NEXT: } {"C/I0"} +# CHECK-NEXT: } {"C/K"} +# CHECK-NEXT: } {"C/J"} +# CHECK-NEXT: } {"C/I"} # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: } @@ -157,11 +169,11 @@ # CHECK-NEXT: graph: # CHECK-NEXT: name: matmul # CHECK-NEXT: inputs: -# CHECK-NEXT: - %0 : 4x512xfloat32 -# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: - %0 : 1024x512xfloat32 +# CHECK-NEXT: - %1 : 512x1024xfloat32 # CHECK-NEXT: outputs: -# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: - %2 : 1024x1024xfloat32 # CHECK-NEXT: nodes: -# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] # CHECK-NEXT: # CHECK-NEXT: CODE: 0 From 39a38b252095497a0a12266d1b2e3ac31868346c Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 29 Jun 2026 16:40:45 +0200 Subject: [PATCH 6/7] Add 2 primitive gpu_warp gpu_lane + fixes descript that was not using gpu and regen some test --- src/xtc/backends/jir/JIRScheduler.py | 10 + src/xtc/backends/mlir/MlirCompilerPasses.py | 112 +++++++++- src/xtc/backends/mlir/MlirNodeScheduler.py | 12 +- src/xtc/backends/mlir/MlirScheduler.py | 8 + .../mlir/MlirTarget/MlirNVGPUTarget.py | 38 +++- src/xtc/backends/tvm/TVMScheduler.py | 20 ++ src/xtc/itf/schd/scheduler.py | 39 +++- src/xtc/schedules/descript.py | 23 +- src/xtc/schedules/loop_nest.py | 6 + src/xtc/schedules/parameter_loop_nest.py | 23 +- src/xtc/schedules/parsing.py | 51 +++-- src/xtc/schedules/plain_schedule.py | 24 ++- .../test_matmul_mlir_offload_lane.py | 202 ++++++++++++++++++ ...st_matmul_mlir_offload_tensor_vectorise.py | 19 +- ...matmul_mlir_offload_warp_lane_vectorise.py | 183 ++++++++++++++++ .../test_matmul_pmu_counters_gpu.py | 1 + .../schedules/test_matmul_descript_gpu.py | 157 +++++++------- .../test_matmul_descript_gpu_warp_lane.py | 179 ++++++++++++++++ tests/filecheck/search/test_conv_oo.py | 10 +- tests/filecheck/search/test_conv_pprprp.py | 10 +- tests/filecheck/search/test_conv_pprprpv.py | 10 +- tests/filecheck/search/test_conv_pprprpvr.py | 10 +- .../search/test_conv_pprprpvr_rnd.py | 2 +- tests/filecheck/search/test_conv_prp.py | 10 +- tests/filecheck/search/test_conv_prp_rnd.py | 2 +- tests/filecheck/search/test_matmul_goto.py | 10 +- tests/filecheck/search/test_matmul_goto_r.py | 10 +- tests/filecheck/search/test_matmul_oo.py | 10 +- tests/filecheck/search/test_matmul_p1.py | 10 +- tests/filecheck/search/test_matmul_p1v.py | 10 +- tests/filecheck/search/test_matmul_pprprp.py | 10 +- tests/filecheck/search/test_matmul_pprprpv.py | 10 +- .../filecheck/search/test_matmul_pprprpvr.py | 10 +- .../search/test_matmul_pprprpvr_rnd.py | 2 +- tests/filecheck/search/test_matmul_prp.py | 37 +++- tests/filecheck/search/test_matmul_prp_rnd.py | 2 +- 36 files changed, 1085 insertions(+), 197 deletions(-) create mode 100644 tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py create mode 100644 tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py create mode 100644 tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py diff --git a/src/xtc/backends/jir/JIRScheduler.py b/src/xtc/backends/jir/JIRScheduler.py index 5627ad0ed..226b53394 100644 --- a/src/xtc/backends/jir/JIRScheduler.py +++ b/src/xtc/backends/jir/JIRScheduler.py @@ -361,6 +361,16 @@ 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 diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index 855201d4f..3929f3d10 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -338,6 +338,7 @@ def _generate_node_scheduling( return sched_state gpu_material = True gpu_mat_thread = True + gpu_warp_thread = True # Materialize the loops for loop_name in permutation: # Manage the splits @@ -400,6 +401,29 @@ def _generate_node_scheduling( sched_state=sched_state, ) gpu_material = False + elif loop_name in schedule.gpu_warps: + tile_vect = [ + sum(values) + for values in zip( + *[tiles_sizes_by_loops[loop] for loop in schedule.gpu_warps] + ) + ] + tile_vect = tile_vect + [0] * (3 - len(tile_vect)) + position_index = [ + permutation.index(loop) for loop in schedule.gpu_warps + ] + mapping_order = sorted( + range(len(position_index)), key=lambda i: position_index[i] + ) + if gpu_warp_thread: + self._strip_mine( + loop_name=loop_name, + tiling_vector=tile_vect, + mapping_order=mapping_order, + schedule=schedule, + sched_state=sched_state, + ) + gpu_warp_thread = False elif loop_name in schedule.gpu_threads: tile_vect = [ sum(values) @@ -426,6 +450,29 @@ def _generate_node_scheduling( sched_state=sched_state, ) gpu_mat_thread = False + elif loop_name in schedule.gpu_lanes: + tile_vect = [ + sum(values) + for values in zip( + *[tiles_sizes_by_loops[loop] for loop in schedule.gpu_lanes] + ) + ] + tile_vect = tile_vect + [0] * (3 - len(tile_vect)) + position_index = [ + permutation.index(loop) for loop in schedule.gpu_lanes + ] + 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, @@ -582,6 +629,16 @@ def _strip_mine( [self._get_thread_id(index) for index in mapping_order] ) tiling_command = TileUsingForallOp(sched_state.handle, **attr_array) + elif loop_name in schedule.gpu_warps: + attr_array["mapping"] = ArrayAttr.get( + [self._get_warp_id(index) for index in mapping_order] + ) + tiling_command = TileUsingForallOp(sched_state.handle, **attr_array) + elif loop_name in schedule.gpu_lanes: + attr_array["mapping"] = ArrayAttr.get( + [self._get_lane_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: @@ -790,6 +847,14 @@ def _collect_fused_producers(self, unscheduled_handles: set[str | None]): return fused_producers + def _get_lane_id(self, index: int) -> Attribute: + ctx = self._mlir_program.mlir_context + return Attribute.parse(f"#gpu.lane", context=ctx) + + def _get_warp_id(self, index: int) -> Attribute: + ctx = self._mlir_program.mlir_context + return Attribute.parse(f"#gpu.warp<{_GPU_DIM[index]}>", context=ctx) + 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) @@ -819,6 +884,11 @@ def _gpu_mapping( new_loop, generate_gpu_launch=True, ).result + # Tiling threads number + # threads, block / threads + # warps, tile size at least 32 threads + # lane, tile size, preferably 32 threads + block_dims = [] if schedule.gpu_threads: block_dims = [ max(tiles_sizes_by_loops[loop_name_block]) @@ -827,6 +897,26 @@ def _gpu_mapping( schedule.gpu_threads, schedule.gpu_blocks ) ] + if schedule.gpu_lanes: + 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_lanes, schedule.gpu_blocks + ) + ] + if schedule.gpu_warps: + block_dims = [ + 32 + * ( + 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_warps, schedule.gpu_blocks + ) + ] + if block_dims: block_dims = block_dims + [1] * (3 - len(block_dims)) MapNestedForallToThreads( new_loop, @@ -857,6 +947,7 @@ def _gpu_mapping( gpu_block_handle, generate_gpu_launch=True, ).result + block_dims = [] if schedule.gpu_threads: block_dims = [ max(tiles_sizes_by_loops[loop_name_block]) @@ -865,6 +956,26 @@ def _gpu_mapping( schedule.gpu_threads, schedule.gpu_blocks ) ] + if schedule.gpu_lanes: + 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_lanes, schedule.gpu_blocks + ) + ] + if schedule.gpu_warps: + block_dims = [ + 32 + * ( + 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_warps, schedule.gpu_blocks + ) + ] + if block_dims: block_dims = block_dims + [1] * (3 - len(block_dims)) MapNestedForallToThreads( new_loop, @@ -896,7 +1007,6 @@ def find_producer_handles(module: Module, root_handle: str) -> list[str | None]: return producer_handles - class MlirProgramApplyTransformPass: def __init__( self, diff --git a/src/xtc/backends/mlir/MlirNodeScheduler.py b/src/xtc/backends/mlir/MlirNodeScheduler.py index e633537d2..1b12afe21 100644 --- a/src/xtc/backends/mlir/MlirNodeScheduler.py +++ b/src/xtc/backends/mlir/MlirNodeScheduler.py @@ -115,10 +115,20 @@ def map_gpu_threads(self, axes: list[str], root: str = DEFAULT_ROOT): 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 thread" + 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)) diff --git a/src/xtc/backends/mlir/MlirScheduler.py b/src/xtc/backends/mlir/MlirScheduler.py index 5943821db..5d7480c8a 100644 --- a/src/xtc/backends/mlir/MlirScheduler.py +++ b/src/xtc/backends/mlir/MlirScheduler.py @@ -218,6 +218,14 @@ 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) diff --git a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py index 419d0f8fc..7a83f424a 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py +++ b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py @@ -446,14 +446,44 @@ def _lowering_pipeline(self, sm_arch: str, ptx_version: str) -> list[str]: "canonicalize", "cse", "sccp", - "convert-vector-to-llvm{vector-contract-lowering=outerproduct}", "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: diff --git a/src/xtc/backends/tvm/TVMScheduler.py b/src/xtc/backends/tvm/TVMScheduler.py index afe7702a3..87fac7b2a 100644 --- a/src/xtc/backends/tvm/TVMScheduler.py +++ b/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() diff --git a/src/xtc/itf/schd/scheduler.py b/src/xtc/itf/schd/scheduler.py index 5a64b93ee..91b297493 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -319,6 +319,43 @@ 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. @@ -328,7 +365,7 @@ def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: - 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 block dimensions are assigned. + only the corresponding thread dimensions are assigned. Args: axes: list of maximum 3 element that are maps to the dimension x, y, z respectively diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 8ac2f799c..3bfc6a9a7 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -401,6 +401,12 @@ 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 @@ -531,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) @@ -555,6 +560,20 @@ 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), @@ -565,7 +584,7 @@ def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None: 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_block[k], + key=lambda k: node.gpu_thread[k], ) scheduler.gpu_thread(sorted_keys, root=root) diff --git a/src/xtc/schedules/loop_nest.py b/src/xtc/schedules/loop_nest.py index 29ddcacd9..8c6b9f857 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/src/xtc/schedules/loop_nest.py @@ -118,6 +118,8 @@ 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) @@ -247,6 +249,10 @@ 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: diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index b4e2f9e91..3f15aa988 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -113,6 +113,10 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): 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 @@ -129,7 +133,8 @@ 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) - # TODO: make gpu_block and gpu_thread work for parameter loop nest + 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) @@ -179,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, @@ -193,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: @@ -321,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 diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 4e2bac4b9..d7ad51592 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -13,6 +13,7 @@ from .exceptions import ScheduleParseError literal = int | str +_GPU_DIM = {"x": 0, "y": 1, "z": 2} def toliteral(s: str) -> literal: @@ -54,6 +55,8 @@ class Annotations: fuse_consumer: bool | None = False partial: bool = False full: bool = False + gpu_lane: int | None = None + gpu_warp: int | None = None gpu_block: int | None = None gpu_thread: int | None = None @@ -183,6 +186,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations fuse_consumer: bool = False partial = False full = False + gpu_lane: str | int | None = None + gpu_warp: str | int | None = None gpu_block: str | int | None = None gpu_thread: str | int | None = None @@ -247,15 +252,36 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations partial = True case "full": full = True + case "gpu_lane": + if isinstance(param, str): + gpu_lane = _GPU_DIM.get(param, None) + if gpu_lane is None: + raise ScheduleParseError( + f'`{{"gpu_lane" = {param}}}`: gpu_block parameter should be a string or int' + ) + elif isinstance(param, int): + gpu_lane = param + else: + raise ScheduleParseError( + f'`{{"gpu_lane" = {param}}}`: gpu_block parameter should be a string or int' + ) + case "gpu_warp": + if isinstance(param, str): + gpu_warp = _GPU_DIM.get(param, None) + if gpu_warp is None: + raise ScheduleParseError( + f'`{{"gpu_warp" = {param}}}`: gpu_warp parameter should be a string or int' + ) + elif isinstance(param, int): + gpu_warp = param + else: + raise ScheduleParseError( + f'`{{"gpu_warp" = {param}}}`: gpu_warp parameter should be a string or int' + ) case "gpu_block": if isinstance(param, str): - if param == "x": - gpu_block = 0 - elif param == "y": - gpu_block = 1 - elif param == "z": - gpu_block = 2 - else: + gpu_block = _GPU_DIM.get(param, None) + if gpu_block is None: raise ScheduleParseError( f'`{{"gpu_block" = {param}}}`: gpu_block parameter should be a string or int' ) @@ -267,13 +293,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations ) case "gpu_thread": if isinstance(param, str): - if param == "x": - gpu_thread = 0 - elif param == "y": - gpu_thread = 1 - elif param == "z": - gpu_thread = 2 - else: + gpu_thread = _GPU_DIM.get(param, None) + if gpu_thread is None: raise ScheduleParseError( f'`{{"gpu_thread" = {param}}}`: gpu_thread string parameter should x, y or z' ) @@ -292,6 +313,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations if partial and full: raise ScheduleParseError(f"{context} has both annotations full and partial") + assert isinstance(gpu_lane, int) or gpu_lane is None + assert isinstance(gpu_warp, int) or gpu_warp is None assert isinstance(gpu_block, int) or gpu_block is None assert isinstance(gpu_thread, int) or gpu_thread is None return Annotations( diff --git a/src/xtc/schedules/plain_schedule.py b/src/xtc/schedules/plain_schedule.py index 18373bc63..dff7bf1c2 100644 --- a/src/xtc/schedules/plain_schedule.py +++ b/src/xtc/schedules/plain_schedule.py @@ -35,6 +35,8 @@ class PlainNodeSchedule: fused_consumers: list[str] gpu_blocks: list[str] gpu_threads: list[str] + gpu_lanes: list[str] + gpu_warps: list[str] # Optional caller-provided vector sizes, keyed by vectorized axis name. # When an axis has a size, its dimension is vectorized with masking for # non-divisible extents; axes absent from this mapping are vectorized to @@ -120,6 +122,8 @@ def __init__( self.fused_consumers: list[str] = [] self.gpu_blocks: list[str] = [] self.gpu_threads: list[str] = [] + self.gpu_lanes: list[str] = [] + self.gpu_warps: list[str] = [] def get_plain_schedule(self) -> PlainNodeSchedule: return PlainNodeSchedule( @@ -143,6 +147,8 @@ def get_plain_schedule(self) -> PlainNodeSchedule: fused_consumers=deepcopy(self.fused_consumers), gpu_blocks=deepcopy(self.gpu_blocks), gpu_threads=deepcopy(self.gpu_threads), + gpu_lanes=deepcopy(self.gpu_lanes), + gpu_warps=deepcopy(self.gpu_warps), vectorization_sizes=deepcopy(self.vectorization_sizes), ) @@ -281,12 +287,22 @@ def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None: fuse_axis = make_loop_name(root, axis) self.fused_consumers.append(fuse_axis) + def gpu_block(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] + def gpu_thread(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 gpu_block(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] + def gpu_lane(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.gpu_lanes = [make_loop_name(root, axis) for axis in axes] + + def gpu_warp(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.gpu_warps = [make_loop_name(root, axis) for axis in axes] diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py new file mode 100644 index 000000000..3bcecb18e --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py @@ -0,0 +1,202 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir.MlirGraphBackend import MlirGraphBackend as Backend + +from xtc.runtimes.accelerator.gpu import GPUDevice + +# Create device +gpu = GPUDevice() + +I, J, K, dtype = 512, 512, 512, "float32" +a = O.tensor((I, K), dtype, name="A") # A lives on the host +b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) # C must live on the accelerator + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.tile("i", {"i1": 128, "i2": 32}) +sch.tile("j", {"j1": 128, "j2": 32}) +sch.tile("k", {"k1": 64}) +sch.unroll({"i2": 2}) +sch.gpu_block(["i", "j"]) +sch.gpu_lane(["i1", "j1"]) +sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="gpu_matmul_mlir_offload_tensor", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<512x512xf32>, memref<512x512xf32>) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op +# CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [32, 32, 0](mapping = [#gpu.lane, #gpu.lane]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_3 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# 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: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./j2" : !transform.any_op +# CHECK-NEXT: transform.loop.unroll %loops_9 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<()[s0, s1, s2] -> (s0 + s1 * 4 + s2 * 16)> +# CHECK-NEXT: #map2 = affine_map<()[s0, s1, s2] -> ((s0 + s1 * 4 + s2 * 16) mod 32)> +# CHECK-NEXT: #map3 = affine_map<()[s0] -> (s0 mod 4)> +# CHECK-NEXT: #map4 = affine_map<()[s0, s1, s2] -> (((s0 + s1 * 4 + s2 * 16) mod 32) floordiv 4)> +# CHECK-NEXT: #map5 = affine_map<(d0) -> (d0 * 32)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 512] [1, 1] : memref<512x512xf32> to memref<1x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c512_7 = arith.constant 512 : index +# CHECK-NEXT: %c1_8 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_6 to %c512_7 step %c1_8 { +# CHECK-NEXT: %subview_9 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x512xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_9 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c4_1 = arith.constant 4 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c4_3 = arith.constant 4 : index +# CHECK-NEXT: %c4_4 = arith.constant 4 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c4_3, %arg10 = %c4_4, %arg11 = %c1_5) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c0_7 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x512xf32> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<512x512xf32> to memref<128x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %2 = affine.apply #map1()[%thread_id_x, %thread_id_y, %c0_6] +# CHECK-NEXT: %3 = affine.apply #map2()[%thread_id_x, %thread_id_y, %c0_6] +# CHECK-NEXT: %4 = affine.apply #map3()[%thread_id_x] +# CHECK-NEXT: %5 = affine.apply #map4()[%thread_id_x, %thread_id_y, %c0_6] +# CHECK-NEXT: %6 = affine.apply #map5(%4) +# CHECK-NEXT: %7 = affine.apply #map5(%5) +# CHECK-NEXT: %subview_10 = memref.subview %subview[%6, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[0, %7] [512, 32] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_12 = memref.subview %subview_9[%6, %7] [32, 32] [1, 1] : memref<128x128xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_13 = arith.constant 0 : index +# CHECK-NEXT: %c512_14 = arith.constant 512 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: scf.for %arg15 = %c0_13 to %c512_14 step %c64 { +# CHECK-NEXT: %subview_15 = memref.subview %subview_10[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_11[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<64x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_17 = memref.subview %subview_12[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_18 = arith.constant 0 : index +# CHECK-NEXT: %c64_19 = arith.constant 64 : index +# CHECK-NEXT: %c1_20 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg16 = %c0_18 to %c64_19 step %c1_20 { +# CHECK-NEXT: %subview_21 = memref.subview %subview_15[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_22 = memref.subview %subview_16[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_23 = memref.subview %subview_17[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_24 = arith.constant 0 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c1_25 = arith.constant 1 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index +# CHECK-NEXT: scf.for %arg17 = %c0_24 to %c32 step %c2 { +# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_27 = memref.subview %subview_22[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_28 = memref.subview %subview_23[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_29 = arith.constant 0 : index +# CHECK-NEXT: %c32_30 = arith.constant 32 : index +# CHECK-NEXT: %c1_31 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_29 to %c32_30 step %c1_31 { +# CHECK-NEXT: %subview_39 = memref.subview %subview_26[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_40 = memref.subview %subview_27[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_41 = memref.subview %subview_28[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_39, %subview_40 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_41 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %c1_32 = arith.constant 1 : index +# CHECK-NEXT: %8 = arith.muli %c1_25, %c1_32 : index +# CHECK-NEXT: %9 = arith.addi %arg17, %8 : index +# CHECK-NEXT: %subview_33 = memref.subview %subview_21[%9, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_34 = memref.subview %subview_22[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_35 = memref.subview %subview_23[%9, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_36 = arith.constant 0 : index +# CHECK-NEXT: %c32_37 = arith.constant 32 : index +# CHECK-NEXT: %c1_38 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_36 to %c32_37 step %c1_38 { +# CHECK-NEXT: %subview_39 = memref.subview %subview_33[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_40 = memref.subview %subview_34[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_41 = memref.subview %subview_35[0, %arg18] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_39, %subview_40 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_41 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 512x512xfloat32 +# CHECK-NEXT: - %1 : 512x512xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 512x512xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [512x512xfloat32, 512x512xfloat32] -> [512x512xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 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 2fcd2d325..9ce23e876 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 @@ -24,12 +24,11 @@ sch = impl.get_scheduler() sch.tile("i", {"i1": 128, "i2": 32}) sch.tile("j", {"j1": 128, "j2": 32}) -sch.tile("k", {"k1": 64}) +sch.tile("k", {"k1": 8}) sch.unroll({"i2": 2}) sch.gpu_block(["i", "j"]) sch.gpu_thread(["i1", "j1"]) sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) -sched = sch.schedule() sch.vectorize(["j2"]) sched = sch.schedule() @@ -67,7 +66,7 @@ # CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op # CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [32, 32, 0](mapping = [#gpu.thread, #gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %forall_op_3 "./i1" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op # CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_7 "./k1" : !transform.any_op @@ -97,7 +96,7 @@ # CHECK-NEXT: %0 = ub.poison : f32 # CHECK-NEXT: %c2 = arith.constant 2 : index # CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index # CHECK-NEXT: %c512 = arith.constant 512 : index @@ -135,12 +134,12 @@ # CHECK-NEXT: %subview_10 = memref.subview %subview[%3, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_11 = memref.subview %subview_8[0, %4] [512, 32] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x32xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_12 = memref.subview %subview_9[%3, %4] [32, 32] [1, 1] : memref<128x128xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg15 = %c0 to %c512 step %c64 { -# CHECK-NEXT: %subview_13 = memref.subview %subview_10[0, %arg15] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_14 = memref.subview %subview_11[%arg15, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<64x32xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg16 = %c0 to %c64 step %c1 { -# CHECK-NEXT: %subview_15 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_16 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg15 = %c0 to %c512 step %c8 { +# CHECK-NEXT: %subview_13 = memref.subview %subview_10[0, %arg15] [32, 8] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x8xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_11[%arg15, 0] [8, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<8x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg16 = %c0 to %c8 step %c1 { +# CHECK-NEXT: %subview_15 = memref.subview %subview_13[0, %arg16] [32, 1] [1, 1] : memref<32x8xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_14[%arg16, 0] [1, 32] [1, 1] : memref<8x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: scf.for %arg17 = %c0 to %c32 step %c2 { # CHECK-NEXT: %subview_17 = memref.subview %subview_15[%arg17, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_18 = memref.subview %subview_12[%arg17, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> 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 new file mode 100644 index 000000000..0b95988e7 --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py @@ -0,0 +1,183 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir.MlirGraphBackend import MlirGraphBackend as Backend + +from xtc.runtimes.accelerator.gpu import GPUDevice + +# Create device +gpu = GPUDevice() + +I, J, K, dtype = 1024, 1024, 512, "float32" +a = O.tensor((I, K), dtype, name="A", device=gpu) # A lives on the host +b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) # C must live on the accelerator + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.tile("i", {"i1": 8, "i2": 4}) +sch.tile("j", {"j1": 128, "j2": 64, "j3": 4}) +sch.tile("k", {"k2": 16}) +# sch.unroll({"i2": 2}) +sch.gpu_block(["j", "i"]) +sch.gpu_warp(["j1"]) +sch.gpu_lane(["j2", "i1"]) +sch.interchange(["j", "i", "j1", "j2", "i1","k", "j3","i2", "k2"]) +sch.vectorize(["j3","i2","k2"]) +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="gpu_matmul_mlir_offload_tensor_vectorise", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [8, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./j" : !transform.any_op +# CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [0, 64, 0](mapping = [#gpu.warp]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_3 "./j1" : !transform.any_op +# CHECK-NEXT: %tiled_op_4, %forall_op_5 = transform.structured.tile_using_forall %tiled_op_2 tile_sizes [4, 4, 0](mapping = [#gpu.lane, #gpu.lane]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# 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: transform.apply_patterns to %2 { +# 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: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 8)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map2 = affine_map<()[s0] -> (s0 floordiv 32)> +# CHECK-NEXT: #map3 = affine_map<(d0) -> (d0 * 64)> +# CHECK-NEXT: #map4 = affine_map<()[s0, s1, s2] -> (s0 + s1 * 64 + s2 * 64)> +# CHECK-NEXT: #map5 = affine_map<()[s0] -> (s0 mod 32)> +# CHECK-NEXT: #map6 = affine_map<()[s0] -> (s0 mod 2)> +# CHECK-NEXT: #map7 = affine_map<()[s0] -> ((s0 mod 32) floordiv 2)> +# CHECK-NEXT: #map8 = affine_map<(d0) -> (d0 * 4)> +# CHECK-NEXT: #map9 = affine_map<(d0, d1, d2) -> (d0, d2)> +# CHECK-NEXT: #map10 = affine_map<(d0, d1, d2) -> (d2, d1)> +# CHECK-NEXT: #map11 = affine_map<(d0, d1, d2) -> (d0, d1)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview_4 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_4 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %c1_1 = arith.constant 1 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c1_3 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c128, %arg10 = %c8, %arg11 = %c1_3) threads(%arg6, %arg7, %arg8) in (%arg12 = %c64, %arg13 = %c1_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_4 = arith.constant 0 : index +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %1 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %2 = affine.apply #map1(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%1, 0] [8, 512] [1, 1] : memref<1024x512xf32> to memref<8x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_6 = memref.subview %arg1[0, %2] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %arg2[%1, %2] [8, 128] [1, 1] : memref<1024x1024xf32> to memref<8x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %3 = affine.apply #map2()[%thread_id_x] +# CHECK-NEXT: %4 = affine.apply #map3(%3) +# CHECK-NEXT: %subview_8 = memref.subview %subview_6[0, %4] [512, 64] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<512x64xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_7[0, %4] [8, 64] [1, 1] : memref<8x128xf32, strided<[1024, 1], offset: ?>> to memref<8x64xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x_10 = gpu.thread_id x +# CHECK-NEXT: %thread_id_y_11 = gpu.thread_id y +# CHECK-NEXT: %thread_id_z_12 = gpu.thread_id z +# CHECK-NEXT: %5 = affine.apply #map4()[%thread_id_x_10, %c0_4, %c0_4] +# CHECK-NEXT: %6 = affine.apply #map5()[%thread_id_x_10] +# CHECK-NEXT: %7 = affine.apply #map6()[%thread_id_x_10] +# CHECK-NEXT: %8 = affine.apply #map7()[%thread_id_x_10] +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %9 = arith.cmpi ult, %6, %c32 : index +# CHECK-NEXT: scf.if %9 { +# CHECK-NEXT: %10 = affine.apply #map8(%7) +# CHECK-NEXT: %11 = affine.apply #map8(%8) +# CHECK-NEXT: %subview_13 = memref.subview %subview[%10, 0] [4, 512] [1, 1] : memref<8x512xf32, strided<[512, 1], offset: ?>> to memref<4x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_8[0, %11] [512, 4] [1, 1] : memref<512x64xf32, strided<[1024, 1], offset: ?>> to memref<512x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_9[%10, %11] [4, 4] [1, 1] : memref<8x64xf32, strided<[1024, 1], offset: ?>> to memref<4x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg15 = %c0 to %c512 step %c16 { +# CHECK-NEXT: %subview_16 = memref.subview %subview_13[0, %arg15] [4, 16] [1, 1] : memref<4x512xf32, strided<[512, 1], offset: ?>> to memref<4x16xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_17 = memref.subview %subview_14[%arg15, 0] [16, 4] [1, 1] : memref<512x4xf32, strided<[1024, 1], offset: ?>> to memref<16x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %12 = vector.transfer_read %subview_16[%c0, %c0], %0 {in_bounds = [true, true]} : memref<4x16xf32, strided<[512, 1], offset: ?>>, vector<4x16xf32> +# CHECK-NEXT: %13 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<16x4xf32, strided<[1024, 1], offset: ?>>, vector<16x4xf32> +# CHECK-NEXT: %14 = vector.transfer_read %subview_15[%c0, %c0], %0 {in_bounds = [true, true]} : memref<4x4xf32, strided<[1024, 1], offset: ?>>, vector<4x4xf32> +# CHECK-NEXT: %15 = vector.contract {indexing_maps = [#map9, #map10, #map11], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind} %12, %13, %14 : vector<4x16xf32>, vector<16x4xf32> into vector<4x4xf32> +# CHECK-NEXT: vector.transfer_write %15, %subview_15[%c0, %c0] {in_bounds = [true, true]} : vector<4x4xf32>, memref<4x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: } +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1024x512xfloat32 +# CHECK-NEXT: - %1 : 512x1024xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 1024x1024xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py b/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py index cb04ff23a..a638b6355 100644 --- a/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py +++ b/tests/filecheck/evaluation/test_matmul_pmu_counters_gpu.py @@ -22,6 +22,7 @@ sch.vectorize(["j1"]) sch.gpu_block(["i"]) sch.gpu_thread(["i1"]) +sch.interchange(["i", "j", "k", "i1", "i2", "j1"]) sched = sch.schedule() comp = impl.get_compiler( diff --git a/tests/filecheck/schedules/test_matmul_descript_gpu.py b/tests/filecheck/schedules/test_matmul_descript_gpu.py index 28c991b66..277ab9444 100644 --- a/tests/filecheck/schedules/test_matmul_descript_gpu.py +++ b/tests/filecheck/schedules/test_matmul_descript_gpu.py @@ -9,7 +9,7 @@ gpu = GPUDevice() I, J, K, dtype = 1024, 1024, 512, "float32" -a = O.tensor((I, K), dtype, name="A") +a = O.tensor((I, K), dtype, name="A", device=gpu) b = O.tensor((K, J), dtype, name="B", device=gpu) with O.graph(name="matmul") as gb: @@ -53,7 +53,7 @@ # CHECK: // -----// IR Dump Before transform //----- // # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) # CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) @@ -70,98 +70,95 @@ # CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op # CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %1 tile_sizes [128, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_3 "C/I" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 128, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_5 "C/J" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_7 "C/K" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [32, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_9 "C/I0" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_11 "C/J0" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_13 "C/I1" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_15 "C/J1" : !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "C/I" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "C/K" : !transform.any_op +# CHECK-NEXT: %tiled_op_4, %forall_op_5 = transform.structured.tile_using_forall %tiled_linalg_op_2 tile_sizes [32, 32, 0](mapping = [#gpu.thread, #gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_5 "C/I0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_op_4 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "C/I1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "C/J1" : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 4, 1] : (!transform.any_op) -> !transform.any_op # CHECK-NEXT: transform.yield # CHECK-NEXT: } # CHECK-NEXT: } # CHECK-NEXT: # CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> # CHECK-NEXT: module attributes {transform.with_named_sequence} { -# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { # CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 # CHECK-NEXT: %c0 = arith.constant 0 : index # CHECK-NEXT: %c1024 = arith.constant 1024 : index # CHECK-NEXT: %c1 = arith.constant 1 : index # CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { # CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_2 = arith.constant 0 : index -# CHECK-NEXT: %c1024_3 = arith.constant 1024 : index -# CHECK-NEXT: %c1_4 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg4 = %c0_2 to %c1024_3 step %c1_4 { -# CHECK-NEXT: %subview_5 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_5 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c1024_6 = arith.constant 1024 : index +# CHECK-NEXT: %c1_7 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_5 to %c1024_6 step %c1_7 { +# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_8 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) # CHECK-NEXT: } {"./j"} # CHECK-NEXT: } {"./i"} -# CHECK-NEXT: %c0_0 = arith.constant 0 : index -# CHECK-NEXT: %c1024_1 = arith.constant 1024 : index -# CHECK-NEXT: %c128 = arith.constant 128 : index -# CHECK-NEXT: scf.for %arg3 = %c0_0 to %c1024_1 step %c128 { -# CHECK-NEXT: %subview = memref.subview %arg0[%arg3, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_2 = memref.subview %arg1[0, 0] [512, 1024] [1, 1] : memref<512x1024xf32> to memref<512x1024xf32, strided<[1024, 1]>> -# CHECK-NEXT: %subview_3 = memref.subview %arg2[%arg3, 0] [128, 1024] [1, 1] : memref<1024x1024xf32> to memref<128x1024xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_4 = arith.constant 0 : index -# CHECK-NEXT: %c1024_5 = arith.constant 1024 : index -# CHECK-NEXT: %c128_6 = arith.constant 128 : index -# CHECK-NEXT: scf.for %arg4 = %c0_4 to %c1024_5 step %c128_6 { -# CHECK-NEXT: %subview_7 = memref.subview %subview[0, 0] [128, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x512xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_8 = memref.subview %subview_2[0, %arg4] [512, 128] [1, 1] : memref<512x1024xf32, strided<[1024, 1]>> to memref<512x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_9 = memref.subview %subview_3[0, %arg4] [128, 128] [1, 1] : memref<128x1024xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_10 = arith.constant 0 : index -# CHECK-NEXT: %c512 = arith.constant 512 : index -# CHECK-NEXT: %c1_11 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg5 = %c0_10 to %c512 step %c1_11 { -# CHECK-NEXT: %subview_12 = memref.subview %subview_7[0, %arg5] [128, 1] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_13 = memref.subview %subview_8[%arg5, 0] [1, 128] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_14 = memref.subview %subview_9[0, 0] [128, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_15 = arith.constant 0 : index -# CHECK-NEXT: %c128_16 = arith.constant 128 : index -# CHECK-NEXT: %c32 = arith.constant 32 : index -# CHECK-NEXT: scf.for %arg6 = %c0_15 to %c128_16 step %c32 { -# CHECK-NEXT: %subview_17 = memref.subview %subview_12[%arg6, 0] [32, 1] [1, 1] : memref<128x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_18 = memref.subview %subview_13[0, 0] [1, 128] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_19 = memref.subview %subview_14[%arg6, 0] [32, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x128xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_20 = arith.constant 0 : index -# CHECK-NEXT: %c128_21 = arith.constant 128 : index -# CHECK-NEXT: %c32_22 = arith.constant 32 : index -# CHECK-NEXT: scf.for %arg7 = %c0_20 to %c128_21 step %c32_22 { -# CHECK-NEXT: %subview_23 = memref.subview %subview_17[0, 0] [32, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_24 = memref.subview %subview_18[0, %arg7] [1, 32] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_25 = memref.subview %subview_19[0, %arg7] [32, 32] [1, 1] : memref<32x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_26 = arith.constant 0 : index -# CHECK-NEXT: %c32_27 = arith.constant 32 : index -# CHECK-NEXT: %c1_28 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg8 = %c0_26 to %c32_27 step %c1_28 { -# CHECK-NEXT: %subview_29 = memref.subview %subview_23[%arg8, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_30 = memref.subview %subview_24[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_31 = memref.subview %subview_25[%arg8, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %c0_32 = arith.constant 0 : index -# CHECK-NEXT: %c32_33 = arith.constant 32 : index -# CHECK-NEXT: %c1_34 = arith.constant 1 : index -# CHECK-NEXT: scf.for %arg9 = %c0_32 to %c32_33 step %c1_34 { -# CHECK-NEXT: %subview_35 = memref.subview %subview_29[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> -# CHECK-NEXT: %subview_36 = memref.subview %subview_30[0, %arg9] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: %subview_37 = memref.subview %subview_31[0, %arg9] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> -# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_35, %subview_36 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_37 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) -# CHECK-NEXT: } {"C/J1"} -# CHECK-NEXT: } {"C/I1"} -# CHECK-NEXT: } {"C/J0"} -# CHECK-NEXT: } {"C/I0"} -# CHECK-NEXT: } {"C/K"} -# CHECK-NEXT: } {"C/J"} -# CHECK-NEXT: } {"C/I"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c4_1 = arith.constant 4 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c8_3 = arith.constant 8 : index +# CHECK-NEXT: %c1_4 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c8_3, %arg11 = %c1_4) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c4_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_5 = arith.constant 0 : index +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_9 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1_10 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg15 = %c0_9 to %c512 step %c1_10 { +# CHECK-NEXT: %subview_11 = memref.subview %subview[0, %arg15] [128, 1] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_12 = memref.subview %subview_7[%arg15, 0] [1, 128] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_8[0, 0] [128, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %2 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_y) +# CHECK-NEXT: %subview_14 = memref.subview %subview_11[%2, 0] [32, 1] [1, 1] : memref<128x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_12[0, %3] [1, 32] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_13[%2, %3] [32, 32] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_17 = arith.constant 0 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c1_18 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg16 = %c0_17 to %c32 step %c1_18 { +# CHECK-NEXT: %subview_19 = memref.subview %subview_14[%arg16, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_15[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_21 = memref.subview %subview_16[%arg16, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_22 = arith.constant 0 : index +# CHECK-NEXT: %c32_23 = arith.constant 32 : index +# CHECK-NEXT: %c1_24 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg17 = %c0_22 to %c32_23 step %c1_24 { +# CHECK-NEXT: %subview_25 = memref.subview %subview_19[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_20[0, %arg17] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_27 = memref.subview %subview_21[0, %arg17] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_25, %subview_26 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_27 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"C/J1"} +# CHECK-NEXT: } {"C/I1"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: } {"C/K"} +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: } diff --git a/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py b/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py new file mode 100644 index 000000000..225920f9a --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py @@ -0,0 +1,179 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend +from xtc.schedules.descript import descript_scheduler + +from xtc.runtimes.accelerator.gpu import GPUDevice + +gpu = GPUDevice() +I, J, K, dtype = 1024, 1024, 512, "float32" +a = O.tensor((I, K), dtype, name="A", device=gpu) +b = O.tensor((K, J), dtype, name="B", device=gpu) + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +descript_scheduler( + scheduler = sch, + node_name = "C", + abstract_dims = ["I","J","K"], + spec = { + "I": {"gpu_block": 0}, + "J": {"gpu_block": 1}, + "K": {}, + "I#128": {"gpu_warp": 0}, + "J#128": {}, + "I#32": {"gpu_lane": 0}, + "J#32": {"gpu_lane": 1}, + + } +) + +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="matmul_descript_mlir_gpu", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "C/I" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "C/K" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [32, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "C/I0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "C/J0" : !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 "C/I1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "C/J1" : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_3 = arith.constant 0 : index +# CHECK-NEXT: %c1024_4 = arith.constant 1024 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_3 to %c1024_4 step %c1_5 { +# CHECK-NEXT: %subview_6 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_6 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c8_1 = arith.constant 8 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c8_1, %arg11 = %c1_2) threads(%arg6, %arg7, %arg8) in (%arg12 = %c1_0, %arg13 = %c1_0, %arg14 = %c1_0) { +# CHECK-NEXT: %c0_3 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_4 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_5 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1_7 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg15 = %c0_6 to %c512 step %c1_7 { +# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg15] [128, 1] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_4[%arg15, 0] [1, 128] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_5[0, 0] [128, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_11 = arith.constant 0 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg16 = %c0_11 to %c128 step %c32 { +# CHECK-NEXT: %subview_12 = memref.subview %subview_8[%arg16, 0] [32, 1] [1, 1] : memref<128x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_9[0, 0] [1, 128] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_10[%arg16, 0] [32, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_15 = arith.constant 0 : index +# CHECK-NEXT: %c128_16 = arith.constant 128 : index +# CHECK-NEXT: %c32_17 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg17 = %c0_15 to %c128_16 step %c32_17 { +# CHECK-NEXT: %subview_18 = memref.subview %subview_12[0, 0] [32, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, %arg17] [1, 32] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_14[0, %arg17] [32, 32] [1, 1] : memref<32x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_21 = arith.constant 0 : index +# CHECK-NEXT: %c32_22 = arith.constant 32 : index +# CHECK-NEXT: %c1_23 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_21 to %c32_22 step %c1_23 { +# CHECK-NEXT: %subview_24 = memref.subview %subview_18[%arg18, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_25 = memref.subview %subview_19[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_20[%arg18, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_27 = arith.constant 0 : index +# CHECK-NEXT: %c32_28 = arith.constant 32 : index +# CHECK-NEXT: %c1_29 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg19 = %c0_27 to %c32_28 step %c1_29 { +# CHECK-NEXT: %subview_30 = memref.subview %subview_24[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_31 = memref.subview %subview_25[0, %arg19] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_26[0, %arg19] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_30, %subview_31 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_32 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"C/J1"} +# CHECK-NEXT: } {"C/I1"} +# CHECK-NEXT: } {"C/J0"} +# CHECK-NEXT: } {"C/I0"} +# CHECK-NEXT: } {"C/K"} +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1024x512xfloat32 +# CHECK-NEXT: - %1 : 512x1024xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 1024x1024xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/search/test_conv_oo.py b/tests/filecheck/search/test_conv_oo.py index e842e3386..da68dc81b 100644 --- a/tests/filecheck/search/test_conv_oo.py +++ b/tests/filecheck/search/test_conv_oo.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 2, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 2, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 3, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 3, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 7, 1] @@ -99,4 +99,4 @@ # CHECK-NEXT: sample 76: [2, 2, 2, 8, 1, 1, 1] # CHECK-NEXT: sample 77: [2, 2, 2, 16, 1, 1, 1] # CHECK-NEXT: stats {'filtered': 78, 'all': 384} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './c1': 1, './s1': 1, './r1': 1, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './c1': 1, './s1': 1, './r1': 1, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprp.py b/tests/filecheck/search/test_conv_pprprp.py index aa0763a39..0c5234a21 100644 --- a/tests/filecheck/search/test_conv_pprprp.py +++ b/tests/filecheck/search/test_conv_pprprp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 1] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 16, 1, 1, 3] # CHECK-NEXT: stats {'filtered': 100, 'all': 202} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpv.py b/tests/filecheck/search/test_conv_pprprpv.py index afecd1a9b..0f3303958 100644 --- a/tests/filecheck/search/test_conv_pprprpv.py +++ b/tests/filecheck/search/test_conv_pprprpv.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 7, 1] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 16, 7, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 32, 1, 1, 1] # CHECK-NEXT: stats {'filtered_vec': 100, 'filtered': 1520, 'all': 4521} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpvr.py b/tests/filecheck/search/test_conv_pprprpvr.py index a9f78119c..3ecca074a 100644 --- a/tests/filecheck/search/test_conv_pprprpvr.py +++ b/tests/filecheck/search/test_conv_pprprpvr.py @@ -20,13 +20,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 7, 1] @@ -128,4 +128,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 32, 1, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 32, 1, 1, 3] # CHECK-NEXT: stats {'filtered_l2': 100, 'filtered_l1': 102, 'filtered_reg': 132, 'filtered_vec': 134, 'filtered': 1918, 'all': 3178} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 2}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 2, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 2}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 2, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpvr_rnd.py b/tests/filecheck/search/test_conv_pprprpvr_rnd.py index e6f11e898..11e3b57c4 100644 --- a/tests/filecheck/search/test_conv_pprprpvr_rnd.py +++ b/tests/filecheck/search/test_conv_pprprpvr_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 2, 16, 7, 1, 1] # CHECK-NEXT: sample 19: [1, 2, 1, 1, 2, 1, 2, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: stats {'filtered_l2': 5, 'filtered_l1': 5, 'filtered_reg': 6, 'filtered_vec': 6, 'filtered': 100} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2, './b2': 2, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2, './b2': 2, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_prp.py b/tests/filecheck/search/test_conv_prp.py index fd1577917..781149065 100644 --- a/tests/filecheck/search/test_conv_prp.py +++ b/tests/filecheck/search/test_conv_prp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 2, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 2, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 4] @@ -68,4 +68,4 @@ # CHECK-NEXT: sample 45: [2, 2, 2, 8] # CHECK-NEXT: sample 46: [2, 2, 2, 16] # CHECK-NEXT: stats {'filtered': 47, 'all': 48} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_prp_rnd.py b/tests/filecheck/search/test_conv_prp_rnd.py index e86b96833..d03a0ab7e 100644 --- a/tests/filecheck/search/test_conv_prp_rnd.py +++ b/tests/filecheck/search/test_conv_prp_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [2, 2, 1, 8] # CHECK-NEXT: sample 19: [2, 2, 1, 4] # CHECK-NEXT: stats {'filtered': 20} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 1}, './f': {'./f1': 4}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=['./b', './h', './w', './f'], unrolling={'./f1': 4, './w1': 1, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 1}, './f': {'./f1': 4}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=['./b', './h', './w', './f'], unrolling={'./f1': 4, './w1': 1, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_goto.py b/tests/filecheck/search/test_matmul_goto.py index ccf37bac2..ca283121c 100644 --- a/tests/filecheck/search/test_matmul_goto.py +++ b/tests/filecheck/search/test_matmul_goto.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 0, 0] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 0, 1] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 0] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 2, 2, 3, 1, 0] # CHECK-NEXT: sample 99: [1, 1, 1, 2, 2, 3, 1, 1] # CHECK-NEXT: stats {'filtered': 100, 'all': 108} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 2, './j2': 2}, './k': {'./k1': 2}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 2, './j2': 2}, './k': {'./k1': 2}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_goto_r.py b/tests/filecheck/search/test_matmul_goto_r.py index e76e2cdf7..2410936f7 100644 --- a/tests/filecheck/search/test_matmul_goto_r.py +++ b/tests/filecheck/search/test_matmul_goto_r.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [3, 1, 2, 1, 1, 1, 0, 0] # CHECK-NEXT: sample 1: [3, 1, 2, 1, 1, 1, 0, 1] # CHECK-NEXT: sample 2: [3, 1, 2, 1, 1, 1, 1, 0] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [3, 1, 2, 2, 6, 2, 1, 0] # CHECK-NEXT: sample 99: [3, 1, 2, 2, 6, 2, 1, 1] # CHECK-NEXT: stats {'filtered': 3256, 'all': 6620} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1}, './j': {'./j1': 4, './j2': 2}, './k': {'./k1': 6}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1}, './j': {'./j1': 4, './j2': 2}, './k': {'./k1': 6}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_oo.py b/tests/filecheck/search/test_matmul_oo.py index 0a9dfa591..eadb24e4a 100644 --- a/tests/filecheck/search/test_matmul_oo.py +++ b/tests/filecheck/search/test_matmul_oo.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [3, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 12, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 12, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 3] @@ -66,4 +66,4 @@ # CHECK-NEXT: sample 43: [7, 8, 1] # CHECK-NEXT: sample 44: [7, 16, 1] # CHECK-NEXT: stats {'filtered': 45, 'all': 144} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_p1.py b/tests/filecheck/search/test_matmul_p1.py index 43773867b..eee25c081 100644 --- a/tests/filecheck/search/test_matmul_p1.py +++ b/tests/filecheck/search/test_matmul_p1.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 0] # CHECK-NEXT: sample 1: [1, 1, 1, 1] # CHECK-NEXT: sample 2: [1, 1, 1, 2] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 32, 1, 1] # CHECK-NEXT: sample 99: [1, 32, 1, 4] # CHECK-NEXT: stats {'filtered': 100, 'all': 185} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 32}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 32, './i1': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 32}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 32, './i1': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_p1v.py b/tests/filecheck/search/test_matmul_p1v.py index 84a94b71d..b639c5f20 100644 --- a/tests/filecheck/search/test_matmul_p1v.py +++ b/tests/filecheck/search/test_matmul_p1v.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 16, 1, 1] # CHECK-NEXT: sample 1: [1, 16, 1, 4] # CHECK-NEXT: sample 2: [1, 16, 2, 1] @@ -47,4 +47,4 @@ # CHECK-NEXT: sample 24: [7, 16, 1, 1] # CHECK-NEXT: sample 25: [7, 16, 1, 4] # CHECK-NEXT: stats {'filtered': 154, 'all': 864} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprp.py b/tests/filecheck/search/test_matmul_pprprp.py index e79fd3449..c19c1be33 100644 --- a/tests/filecheck/search/test_matmul_pprprp.py +++ b/tests/filecheck/search/test_matmul_pprprp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 3] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 16, 2, 6] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 32, 1, 1] # CHECK-NEXT: stats {'filtered': 100, 'all': 121} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpv.py b/tests/filecheck/search/test_matmul_pprprpv.py index 30becfda8..2928c1141 100644 --- a/tests/filecheck/search/test_matmul_pprprpv.py +++ b/tests/filecheck/search/test_matmul_pprprpv.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 16, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 16, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 16, 3] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [3, 1, 1, 1, 1, 16, 3] # CHECK-NEXT: sample 99: [3, 1, 1, 1, 1, 16, 4] # CHECK-NEXT: stats {'filtered_vec': 100, 'filtered': 1472, 'all': 3052} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 4}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 4}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 4}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 4}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpvr.py b/tests/filecheck/search/test_matmul_pprprpvr.py index 9e11e191d..d035261f5 100644 --- a/tests/filecheck/search/test_matmul_pprprpvr.py +++ b/tests/filecheck/search/test_matmul_pprprpvr.py @@ -20,13 +20,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 16, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 16, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 16, 3] @@ -128,4 +128,4 @@ # CHECK-NEXT: sample 98: [1, 21, 1, 1, 1, 32, 3] # CHECK-NEXT: sample 99: [1, 21, 1, 1, 2, 16, 1] # CHECK-NEXT: stats {'filtered_l2': 100, 'filtered_l1': 105, 'filtered_reg': 115, 'filtered_vec': 154, 'filtered': 2126, 'all': 2749} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 21, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 21, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpvr_rnd.py b/tests/filecheck/search/test_matmul_pprprpvr_rnd.py index 5780a7eb6..72543a25d 100644 --- a/tests/filecheck/search/test_matmul_pprprpvr_rnd.py +++ b/tests/filecheck/search/test_matmul_pprprpvr_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [1, 1, 1, 1, 2, 16, 3] # CHECK-NEXT: sample 19: [7, 1, 3, 1, 1, 16, 2] # CHECK-NEXT: stats {'filtered_l2': 2, 'filtered_l1': 2, 'filtered_reg': 3, 'filtered_vec': 3, 'filtered': 70} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 2}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 2}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_prp.py b/tests/filecheck/search/test_matmul_prp.py index 199db2797..c7b19775b 100644 --- a/tests/filecheck/search/test_matmul_prp.py +++ b/tests/filecheck/search/test_matmul_prp.py @@ -12,13 +12,30 @@ utils.print_all_opt_schedules(backend, strategy) utils.print_exhaustive_samples(backend, strategy, 100) -<<<<<<< HEAD -======= ->>>>>>> 7fae6dc (Add mapping order for the gpu thread and block) -<<<<<<< HEAD -======= ->>>>>>> 7fae6dc (Add mapping order for the gpu thread and block) -# CHECK: File "/home/ruicesista/Documents/xtc/xtc/tests/filecheck/search/test_matmul_prp.py", line 17 -# CHECK-NEXT: >>>>>>> 7fae6dc (Add mapping order for the gpu thread and block) -# CHECK-NEXT: ^ -# CHECK-NEXT: SyntaxError: invalid decimal literal +# CHECK: schedule O0: [1, 1] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] +# CHECK-NEXT: schedule O1: [1, 1] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] +# CHECK-NEXT: schedule O2: [1, 16] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] +# CHECK-NEXT: schedule O3: [3, 16] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] +# CHECK-NEXT: sample 0: [1, 1] +# CHECK-NEXT: sample 1: [1, 2] +# CHECK-NEXT: sample 2: [1, 4] +# CHECK-NEXT: sample 3: [1, 8] +# CHECK-NEXT: sample 4: [1, 16] +# CHECK-NEXT: sample 5: [1, 32] +# CHECK-NEXT: sample 6: [3, 1] +# CHECK-NEXT: sample 7: [3, 2] +# CHECK-NEXT: sample 8: [3, 4] +# CHECK-NEXT: sample 9: [3, 8] +# CHECK-NEXT: sample 10: [3, 16] +# CHECK-NEXT: sample 11: [3, 32] +# CHECK-NEXT: sample 12: [7, 1] +# CHECK-NEXT: sample 13: [7, 2] +# CHECK-NEXT: sample 14: [7, 4] +# CHECK-NEXT: sample 15: [7, 8] +# CHECK-NEXT: sample 16: [7, 16] +# CHECK-NEXT: stats {'filtered': 17, 'all': 24} +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_prp_rnd.py b/tests/filecheck/search/test_matmul_prp_rnd.py index aafbc2661..d534a1703 100644 --- a/tests/filecheck/search/test_matmul_prp_rnd.py +++ b/tests/filecheck/search/test_matmul_prp_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [3, 2] # CHECK-NEXT: sample 19: [21, 8] # CHECK-NEXT: stats {'filtered': 19} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21}, './j': {'./j1': 8}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=['./i', './j'], unrolling={'./j1': 8, './i1': 21}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21}, './j': {'./j1': 8}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=['./i', './j'], unrolling={'./j1': 8, './i1': 21}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] From 96d17ee32a5bebacabfe0d723dcbe1aaaed90de6 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Tue, 8 Sep 2026 09:35:50 +0200 Subject: [PATCH 7/7] gpu tests: change dump file name --- .../backends/target_gpu/test_matmul_mlir_offload_lane.py | 2 +- .../target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py | 2 +- .../backends/target_gpu/test_matmul_mlir_with_tensor.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py index 3bcecb18e..5ee8f0e99 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py @@ -34,7 +34,7 @@ comp = impl.get_compiler( target=gpu, shared_lib=True, - dump_file="gpu_matmul_mlir_offload_tensor", + dump_file="gpu_matmul_mlir_offload_tensor_lane", print_source_ir=True, print_transformed_ir=True, ) 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 0b95988e7..41d00861a 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 @@ -36,7 +36,7 @@ comp = impl.get_compiler( target=gpu, shared_lib=True, - dump_file="gpu_matmul_mlir_offload_tensor_vectorise", + dump_file="gpu_matmul_mlir_offload_tensor_warp_lane_vectorise", print_source_ir=True, print_transformed_ir=True, ) diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py index 32976c805..16a080d5e 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py @@ -34,7 +34,7 @@ comp = impl.get_compiler( target=gpu, shared_lib=True, - dump_file="gpu_matmul_mlir_offload_tensor", + dump_file="gpu_matmul_mlir_with_tensor", print_source_ir=True, print_transformed_ir=True, print_bufferization_ir=True,