From 1d22b78f933517005f71fe2b86fef39225e67c3d Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 11 May 2026 11:37:14 +0200 Subject: [PATCH 1/8] 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 4634bc9d..037caec4 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 69b9a583..9ab4c1a6 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 @@ -326,7 +333,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 @@ -362,12 +370,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 @@ -384,6 +419,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_consumers_into_loops( @@ -550,9 +589,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 @@ -819,6 +875,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 a87c20cd..7bd7110a 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 e421c935..38e36fb0 100644 --- a/src/xtc/backends/mlir/MlirScheduler.py +++ b/src/xtc/backends/mlir/MlirScheduler.py @@ -217,6 +217,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 8173c084..99651f85 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 61f46ee3..c78f8380 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -332,6 +332,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 99a94c8b..51fc208d 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 14fdeb15..f5570536 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/src/xtc/schedules/loop_nest.py @@ -119,6 +119,8 @@ class LoopNestNode(Node["LoopNestNode"]): fuse_producer_at: dict[str, int] = field(default_factory=dict) fuse_consumer_at: list[str] = field(default_factory=list) external_at: dict[str, str] = field(default_factory=dict) + gpu_block: dict[str, int] = field(default_factory=dict) + gpu_thread: dict[str, int] = field(default_factory=dict) def pretty_print(self, indent: int = 0) -> str: """Return a human-readable representation of the loop nest. @@ -248,6 +250,10 @@ def _add_annotations(self, line: str, loop_name: str) -> str: annotations.append("fuse_consumer") if loop_name in self.external_at: annotations.append(f"external({self.external_at[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 3fe10568..d6cb1219 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 9f55df68..efc827bb 100644 --- a/src/xtc/schedules/plain_schedule.py +++ b/src/xtc/schedules/plain_schedule.py @@ -34,6 +34,8 @@ class PlainNodeSchedule: fused_producers: list[tuple[str, int]] fused_consumers: list[str] externals: dict[str, 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 @@ -118,6 +120,8 @@ def __init__( self.fused_producers: list[tuple[str, int]] = [] self.fused_consumers: list[str] = [] self.externals: dict[str, str] = {} + self.gpu_blocks: list[str] = [] + self.gpu_threads: list[str] = [] def get_plain_schedule(self) -> PlainNodeSchedule: return PlainNodeSchedule( @@ -140,6 +144,8 @@ def get_plain_schedule(self) -> PlainNodeSchedule: fused_producers=deepcopy(self.fused_producers), fused_consumers=deepcopy(self.fused_consumers), externals=deepcopy(self.externals), + gpu_blocks=deepcopy(self.gpu_blocks), + gpu_threads=deepcopy(self.gpu_threads), vectorization_sizes=deepcopy(self.vectorization_sizes), ) @@ -282,3 +288,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 7c1f8492..9cbaeb9a 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 00000000..7a1080d5 --- /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 62cd7f9b..504b7bf2 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 0b4a6249..8a2dd2c5 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 05dff96a..d12fa684 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 282114d1..4f3e05fb 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 2297f093..6e5e9055 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 4c9c9661..e40d12cf 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 4e7926c0..79f458f0 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 3d187493..2e36d0b2 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 e677ad9e..af3a7866 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 a2bfe8ed..633e7740 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 8f185ded..30a244fa 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 e6bfc557..0820924d 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 5e6da758..63e37476 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 7fb68ceb..bc39254f 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 c00a4cca..5f473fdd 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 3fa48a9c..1bf68cd2 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 1b9007bf..aac0e3d5 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 8c7499df..e2d99a3c 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, gpu_blocks=[], gpu_threads=[], vectorization_sizes={})] From b8084a33d7410ac1a5345b75c434694575bd0186 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 18 May 2026 12:31:48 +0200 Subject: [PATCH 2/8] 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 9ab4c1a6..ce4da01a 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -372,26 +372,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, ) @@ -400,6 +414,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, ) @@ -585,6 +600,7 @@ def _strip_mine( self, loop_name: str, tiling_vector: list[int], + mapping_order: list[int], schedule: MlirNodeSchedule, sched_state: SchedulingState, ) -> OpResult: From adfda4272495c4c532638c7ebf22724bdf6449c3 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 18 May 2026 14:38:04 +0200 Subject: [PATCH 3/8] 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 ++++++++++-------- 7 files changed, 188 insertions(+), 144 deletions(-) diff --git a/src/xtc/backends/jir/JIRScheduler.py b/src/xtc/backends/jir/JIRScheduler.py index 037caec4..5627ad0e 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 ce4da01a..0c4854c0 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 if schedule.fused_consumers: @@ -381,11 +388,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, ) @@ -401,11 +415,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, ) @@ -434,10 +454,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_consumers_into_loops( @@ -608,22 +624,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) @@ -867,30 +874,6 @@ def find_consumer_handles(module: Module, root_handle: str) -> list[str | None]: consumer_handles.append(attr) return consumer_handles - -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) @@ -903,8 +886,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( ( @@ -915,27 +898,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 c78f8380..c705eb4c 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -333,7 +333,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. @@ -342,7 +342,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 51fc208d..8ac2f799 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 822073fa..b4e2f9e9 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 d6cb1219..4e2bac4b 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 9cbaeb9a..2a02e131 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 From 115aa0fb1373277f0ab69b4222b5250dcd85cd70 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Fri, 19 Jun 2026 16:51:43 +0200 Subject: [PATCH 4/8] 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 0c4854c0..cb9beb7b 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 [] ) @@ -630,6 +631,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: @@ -639,6 +641,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) @@ -888,7 +892,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] @@ -916,6 +920,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 00000000..0858d657 --- /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 dbf827a40db7e1bda028f9fee20ded85673d0bd8 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 22 Jun 2026 13:55:53 +0200 Subject: [PATCH 5/8] 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 cb9beb7b..05e8cd3e 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -621,18 +621,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) @@ -641,7 +643,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) @@ -717,11 +719,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 7bd7110a..e633537d 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 99651f85..419d0f8f 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 c705eb4c..b02ed77a 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -334,19 +334,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 00000000..8f116965 --- /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 2a02e131..6138d3a6 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 0194b726..2fcd2d32 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 0858d657..32976c80 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 ab671dea..cb04ff23 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 7a1080d5..28c991b6 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 fecc1d312966b7038f75cc6f9dea7c0408e4e508 Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Mon, 29 Jun 2026 16:40:45 +0200 Subject: [PATCH 6/8] 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 | 153 +++++++++++-- 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/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 | 10 +- tests/filecheck/search/test_matmul_prp_rnd.py | 2 +- 35 files changed, 1077 insertions(+), 213 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 5627ad0e..226b5339 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 05e8cd3e..7a0c59b8 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -343,6 +343,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 @@ -405,6 +406,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) @@ -431,6 +455,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, @@ -634,6 +681,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: @@ -858,28 +915,13 @@ def _collect_fused_producers(self, unscheduled_handles: set[str | None]): return fused_producer_handles + def _get_lane_id(self, index: int) -> Attribute: + ctx = self._mlir_program.mlir_context + return Attribute.parse(f"#gpu.lane", context=ctx) -def find_consumer_handles(module: Module, root_handle: str) -> list[str | None]: - # returns the handles for each consumer op of the operation specified by root_handle - consumer_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 consumer_handles - - for use in root_op.results[0].uses: - consumer_op = use.owner - for attr in consumer_op.attributes: - if attr.startswith("__xtc_id_"): - consumer_handles.append(attr) - return consumer_handles + 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 @@ -910,6 +952,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]) @@ -918,6 +965,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, @@ -948,6 +1015,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]) @@ -956,6 +1024,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, @@ -963,6 +1051,28 @@ def _gpu_mapping( ) +def find_consumer_handles(module: Module, root_handle: str) -> list[str | None]: + # returns the handles for each consumer op of the operation specified by root_handle + consumer_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 consumer_handles + + for use in root_op.results[0].uses: + consumer_op = use.owner + for attr in consumer_op.attributes: + if attr.startswith("__xtc_id_"): + consumer_handles.append(attr) + return consumer_handles + 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] = [] @@ -987,7 +1097,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 e633537d..1b12afe2 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 38e36fb0..3ae70a23 100644 --- a/src/xtc/backends/mlir/MlirScheduler.py +++ b/src/xtc/backends/mlir/MlirScheduler.py @@ -217,6 +217,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 419d0f8f..7a83f424 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 926fd754..3a281326 100644 --- a/src/xtc/backends/tvm/TVMScheduler.py +++ b/src/xtc/backends/tvm/TVMScheduler.py @@ -425,6 +425,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 b02ed77a..6ba62eb9 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -332,6 +332,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. @@ -341,7 +378,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 8ac2f799..3bfc6a9a 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/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index b4e2f9e9..3f15aa98 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 4e2bac4b..d7ad5159 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 efc827bb..6c1cd7fc 100644 --- a/src/xtc/schedules/plain_schedule.py +++ b/src/xtc/schedules/plain_schedule.py @@ -36,6 +36,8 @@ class PlainNodeSchedule: externals: dict[str, 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 @@ -122,6 +124,8 @@ def __init__( self.externals: dict[str, 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( @@ -146,6 +150,8 @@ def get_plain_schedule(self) -> PlainNodeSchedule: externals=deepcopy(self.externals), 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), ) @@ -289,12 +295,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 00000000..3bcecb18 --- /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 2fcd2d32..9ce23e87 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 00000000..0b95988e --- /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 cb04ff23..a638b635 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 28c991b6..277ab944 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 00000000..225920f9 --- /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 504b7bf2..b000de9f 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 8a2dd2c5..95cfba4f 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 d12fa684..e70518d8 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 4f3e05fb..186a734b 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 6e5e9055..38e92cb5 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 e40d12cf..f56a2529 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 79f458f0..af499191 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 2e36d0b2..13aa90e0 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 af3a7866..80798ad6 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 633e7740..5adbfd66 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 30a244fa..e5686009 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 0820924d..24187cc7 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 63e37476..c1ca2d9d 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 bc39254f..430e4fe1 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 5f473fdd..afa616e5 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 1bf68cd2..85ede46b 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 aac0e3d5..044431c8 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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] @@ -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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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 e2d99a3c..5d534c54 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, 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_producers=[], fused_consumers=[], externals={}, gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] From b6973f798b809821e5452d0e6dc6fcf849d188fa Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Tue, 8 Sep 2026 09:35:50 +0200 Subject: [PATCH 7/8] Refactor some code + fix parsing for gpu_lane and gpu_warp Fix some tests on gpu + add schedule to nvgpu test --- Makefile | 4 +- src/xtc/backends/mlir/MlirCompilerPasses.py | 232 ++++++------------ src/xtc/backends/mlir/MlirNodeScheduler.py | 8 - src/xtc/schedules/descript.py | 2 +- src/xtc/schedules/loop_nest.py | 10 + src/xtc/schedules/parameter_loop_nest.py | 2 - src/xtc/schedules/parsing.py | 20 +- .../test_matmul_mlir_offload_lane.py | 2 +- ...matmul_mlir_offload_warp_lane_vectorise.py | 2 +- .../test_matmul_mlir_with_tensor.py | 2 +- .../schedules/test_descript_pretty_print.py | 39 ++- .../test_matmul_descript_gpu_warp_lane.py | 140 ++++++----- 12 files changed, 228 insertions(+), 235 deletions(-) diff --git a/Makefile b/Makefile index 81f00bf2..766364e2 100644 --- a/Makefile +++ b/Makefile @@ -63,7 +63,7 @@ check-lit-c: env XTC_MLIR_TARGET=c lit -v tests/filecheck/backends tests/filecheck/mlir_loop check-lit-nvgpu: - [ `uname -s` = Darwin ] || env XTC_MLIR_TARGET=nvgpu lit -v tests/filecheck/backends tests/filecheck/mlir_loop tests/filecheck/evaluation + [ `uname -s` = Darwin ] || env XTC_MLIR_TARGET=nvgpu lit -v tests/filecheck/backends tests/filecheck/mlir_loop tests/filecheck/evaluation tests/filecheck/schedules check-lit-mppa: [ `uname -s` = Darwin ] || env XTC_MLIR_TARGET=mppa lit -v -j 1 tests/filecheck/backends/target_mppa tests/filecheck/evaluation/test_matmul_pmu_counters_mppa.py @@ -105,5 +105,5 @@ claude: run-tutorial: marimo run docs/tutorials/xtc_101.py -.PHONY: help test check check-lit-all check-lit check-lit-c check-lit-nvpgu check-pytest check-type check-pyright check-mypy check-format check-format-ruff check-license check-banwords format format-ruff format-license agents claude check-tutorials run-tutorial check-dependencies dependencies wheel pages +.PHONY: help test check check-lit-all check-lit check-lit-c check-lit-nvgpu check-pytest check-type check-pyright check-mypy check-format check-format-ruff check-license check-banwords format format-ruff format-license agents claude check-tutorials run-tutorial check-dependencies dependencies wheel pages .SUFFIXES: diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index 7a0c59b8..8970dbcf 100644 --- a/src/xtc/backends/mlir/MlirCompilerPasses.py +++ b/src/xtc/backends/mlir/MlirCompilerPasses.py @@ -380,102 +380,47 @@ def _generate_node_scheduling( break elif loop_name in tiles_sizes_by_loops: if loop_name in schedule.gpu_blocks: - tile_vect = [ - sum(values) - for values in zip( - *[ - tiles_sizes_by_loops[loop] - for loop in schedule.gpu_blocks - ] - ) - ] - tile_vect = tile_vect + [0] * (3 - len(tile_vect)) - # TODO: Do not work with splitting - position_index = [ - permutation.index(loop) for loop in schedule.gpu_blocks - ] - mapping_order = sorted( - range(len(position_index)), key=lambda i: position_index[i] - ) if gpu_material: - self._strip_mine( + self._gpu_strip_mine( loop_name=loop_name, - tiling_vector=tile_vect, - mapping_order=mapping_order, schedule=schedule, sched_state=sched_state, + gpu_list=schedule.gpu_blocks, + permutation=permutation, + tiles_sizes_by_loops=tiles_sizes_by_loops, ) 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( + self._gpu_strip_mine( loop_name=loop_name, - tiling_vector=tile_vect, - mapping_order=mapping_order, schedule=schedule, sched_state=sched_state, + gpu_list=schedule.gpu_warps, + permutation=permutation, + tiles_sizes_by_loops=tiles_sizes_by_loops, ) gpu_warp_thread = False elif loop_name in schedule.gpu_threads: - tile_vect = [ - sum(values) - for values in zip( - *[ - tiles_sizes_by_loops[loop] - for loop in schedule.gpu_threads - ] - ) - ] - tile_vect = tile_vect + [0] * (3 - len(tile_vect)) - position_index = [ - permutation.index(loop) for loop in schedule.gpu_threads - ] - mapping_order = sorted( - range(len(position_index)), key=lambda i: position_index[i] - ) if gpu_mat_thread: - self._strip_mine( + self._gpu_strip_mine( loop_name=loop_name, - tiling_vector=tile_vect, - mapping_order=mapping_order, schedule=schedule, sched_state=sched_state, + gpu_list=schedule.gpu_threads, + permutation=permutation, + tiles_sizes_by_loops=tiles_sizes_by_loops, ) 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( + self._gpu_strip_mine( loop_name=loop_name, - tiling_vector=tile_vect, - mapping_order=mapping_order, schedule=schedule, sched_state=sched_state, + gpu_list=schedule.gpu_lanes, + permutation=permutation, + tiles_sizes_by_loops=tiles_sizes_by_loops, ) gpu_mat_thread = False else: @@ -936,7 +881,6 @@ def _gpu_mapping( schedule: MlirNodeSchedule, sched_state: SchedulingState, ): - tiles_sizes_by_loops = self._generate_tiling_insns(schedule) if schedule.gpu_blocks and not self._using_tensors: new_loop = next( ( @@ -946,50 +890,7 @@ def _gpu_mapping( ), None, ) - # Since we know there only 1 non zero number - # TODO Find a way to put thread number instead of putting tile size - new_loop = MapForallToBlocks( - new_loop, - generate_gpu_launch=True, - ).result - # 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]) - // max(tiles_sizes_by_loops[loop_name]) - for loop_name, loop_name_block in zip( - 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, - block_dims=block_dims, - ) + self._gpu_mapping_helper(schedule, new_loop) elif ( schedule.gpu_blocks and self._using_tensors @@ -1009,46 +910,68 @@ def _gpu_mapping( "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 + self._gpu_mapping_helper(schedule, gpu_block_handle) + + def _gpu_mapping_helper(self, schedule: MlirNodeSchedule, handle: OpResult): + tiles_sizes_by_loops = self._generate_tiling_insns(schedule) + new_loop = MapForallToBlocks( + handle, + generate_gpu_launch=True, + ).result + block_dims: list[int] = [] + for curType, gpu_list in enumerate( + [schedule.gpu_threads, schedule.gpu_lanes, schedule.gpu_warps] + ): + if not gpu_list: + continue + # If there is a something in gpu warp multiply it by 32 + thread_size = 1 + if curType == 2: + thread_size = 32 block_dims = [] - 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 - ) - ] - 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, - block_dims=block_dims, + for loop_name in gpu_list: + tile_size = schedule.size_of_tile(loop_name) + + if tile_size is None: + block_dims.append(1) + else: + block_dims.append( + thread_size + * (tile_size // max(tiles_sizes_by_loops[loop_name])) ) + if block_dims: + block_dims = block_dims + [1] * (3 - len(block_dims)) + MapNestedForallToThreads( + new_loop, + block_dims=block_dims, + ) + + def _gpu_strip_mine( + self, + loop_name: str, + schedule: MlirNodeSchedule, + sched_state: SchedulingState, + gpu_list: list[str], + permutation: list[str], + tiles_sizes_by_loops: dict[str, list[int]], + ): + tile_vect = [ + sum(values) + for values in zip(*[tiles_sizes_by_loops[loop] for loop in gpu_list]) + ] + tile_vect = tile_vect + [0] * (3 - len(tile_vect)) + # TODO: Make it work with splitting + position_index = [permutation.index(loop) for loop in gpu_list] + mapping_order = sorted( + range(len(position_index)), key=lambda i: position_index[i] + ) + self._strip_mine( + loop_name=loop_name, + tiling_vector=tile_vect, + mapping_order=mapping_order, + schedule=schedule, + sched_state=sched_state, + ) def find_consumer_handles(module: Module, root_handle: str) -> list[str | None]: @@ -1073,6 +996,7 @@ def find_consumer_handles(module: Module, root_handle: str) -> list[str | None]: consumer_handles.append(attr) return consumer_handles + 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] = [] diff --git a/src/xtc/backends/mlir/MlirNodeScheduler.py b/src/xtc/backends/mlir/MlirNodeScheduler.py index 1b12afe2..8cb42652 100644 --- a/src/xtc/backends/mlir/MlirNodeScheduler.py +++ b/src/xtc/backends/mlir/MlirNodeScheduler.py @@ -110,23 +110,15 @@ def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None: self._plain_sch.fuse_consumer_at(axis, root) def map_gpu_threads(self, axes: list[str], root: str = DEFAULT_ROOT): - assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu thread" - assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu thread" self._plain_sch.gpu_thread(axes, root) def map_gpu_blocks(self, axes: list[str], root: str = DEFAULT_ROOT): - assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu block" - assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu block" self._plain_sch.gpu_block(axes, root) def map_gpu_lanes(self, axes: list[str], root: str = DEFAULT_ROOT): - assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu lane" - assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu lane" self._plain_sch.gpu_lane(axes, root) def map_gpu_warps(self, axes: list[str], root: str = DEFAULT_ROOT): - assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu warp" - assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu warp" self._plain_sch.gpu_warp(axes, root) def get_node_schedule(self) -> MlirNodeSchedule: diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 3bfc6a9a..c89c8422 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -572,7 +572,7 @@ def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None: (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) + scheduler.gpu_warp(sorted_keys, root=root) if node.gpu_block: sorted_keys = sorted( diff --git a/src/xtc/schedules/loop_nest.py b/src/xtc/schedules/loop_nest.py index f5570536..c3c3b6e7 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/src/xtc/schedules/loop_nest.py @@ -105,6 +105,10 @@ class LoopNestNode(Node["LoopNestNode"]): 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: List of loops to map to gpu block + gpu_thread: List of loops to map to gpu thread + gpu_warp: List of loops to map to gpu warp + gpu_lane: List of loops to map to gpu lane """ root: str @@ -121,6 +125,8 @@ class LoopNestNode(Node["LoopNestNode"]): external_at: dict[str, str] = field(default_factory=dict) gpu_block: dict[str, int] = field(default_factory=dict) gpu_thread: dict[str, int] = field(default_factory=dict) + gpu_warp: dict[str, int] = field(default_factory=dict) + gpu_lane: dict[str, int] = field(default_factory=dict) def pretty_print(self, indent: int = 0) -> str: """Return a human-readable representation of the loop nest. @@ -254,6 +260,10 @@ def _add_annotations(self, line: str, loop_name: str) -> str: 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 loop_name in self.gpu_warp: + annotations.append(f"gpu_warp({self.gpu_warp[loop_name]})") + if loop_name in self.gpu_lane: + annotations.append(f"gpu_lane({self.gpu_lane[loop_name]})") if annotations: line += " // " + ", ".join(annotations) return line diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index 3f15aa98..23ff548e 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -111,8 +111,6 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): fuse_producer_at: Producer fusion configuration per axis. Maps axis names to producer indices. fuse_consumer_at: List of axes where the output consumer is fused. - gpu_block: Maps loops to block id, that loop need to be parallelize - gpu_thread: Maps loops to thread id, that loop need to be parallelize gpu_lane: Maps loops to lane id gpu_warp: Maps loops to warp id gpu_block: Maps loops to block id diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index d7ad5159..8f1b5df1 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -257,10 +257,14 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations 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' + f'`{{"gpu_lane" = {param}}}`: gpu_lane parameter should be a string or int' ) elif isinstance(param, int): gpu_lane = param + if gpu_lane is None: + raise ScheduleParseError( + f'`{{"gpu_lane" = {param}}}`: gpu_lane string parameter should x, y or z' + ) else: raise ScheduleParseError( f'`{{"gpu_lane" = {param}}}`: gpu_block parameter should be a string or int' @@ -274,6 +278,10 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations ) elif isinstance(param, int): gpu_warp = param + if gpu_warp is None: + raise ScheduleParseError( + f'`{{"gpu_warp" = {param}}}`: gpu_warp string parameter should x, y or z' + ) else: raise ScheduleParseError( f'`{{"gpu_warp" = {param}}}`: gpu_warp parameter should be a string or int' @@ -287,6 +295,10 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations ) elif isinstance(param, int): gpu_block = param + if gpu_block is None: + raise ScheduleParseError( + f'`{{"gpu_block" = {param}}}`: gpu_block string parameter should x, y or z' + ) else: raise ScheduleParseError( f'`{{"gpu_block" = {param}}}`: gpu_block parameter should be a string or int' @@ -300,7 +312,7 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations ) elif isinstance(param, int): gpu_thread = param - if param < 0 and param < 3: + if param < 0 or param >= 3: raise ScheduleParseError( f'`{{"gpu_thread" = {param}}}`: gpu_thread int parameter should 0, 1 or 2' ) @@ -313,10 +325,12 @@ 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( unroll_factor=unroll_factor, unroll_specified=unroll_specified, @@ -332,6 +346,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations full=full, gpu_block=gpu_block, gpu_thread=gpu_thread, + gpu_lane=gpu_lane, + gpu_warp=gpu_warp, ) def _parse_pack_param( 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 3bcecb18..5ee8f0e9 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 0b95988e..41d00861 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 32976c80..16a080d5 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, diff --git a/tests/filecheck/schedules/test_descript_pretty_print.py b/tests/filecheck/schedules/test_descript_pretty_print.py index 52309940..50c4c021 100644 --- a/tests/filecheck/schedules/test_descript_pretty_print.py +++ b/tests/filecheck/schedules/test_descript_pretty_print.py @@ -7,6 +7,8 @@ # RUN: python %s --pack 2>&1 | filecheck %s --check-prefix=CHECK-PACK # RUN: python %s --fuse-producer 2>&1 | filecheck %s --check-prefix=CHECK-FUSE-PRODUCER # RUN: python %s --fuse-consumer 2>&1 | filecheck %s --check-prefix=CHECK-FUSE-CONSUMER +# RUN: python %s --gpu 2>&1 | filecheck %s --check-prefix=CHECK-GPU +# RUN: python %s --gpu-warp-lane 2>&1 | filecheck %s --check-prefix=CHECK-GPU-WARP-LANE import sys from xtc.schedules.parsing import ScheduleParser @@ -89,7 +91,28 @@ ast = parser.parse(spec) loop_nest = interpreter.interpret(ast, root="C") print(loop_nest.root_node.pretty_print()) - +elif "--gpu" in sys.argv: + spec = { + "i": {}, + "i#32": {"gpu_block": 0}, + "i#4": {"gpu_thread": 0}, + "j": {}, + "k": {} + } + ast = parser.parse(spec) + loop_nest = interpreter.interpret(ast, root="C") + print(loop_nest.root_node.pretty_print()) +elif "--gpu-warp-lane" in sys.argv: + spec = { + "i": {"gpu_block": 0}, + "i#32": {"gpu_warp": 0}, + "i#4": {"gpu_lane": 0}, + "j": {}, + "k": {} + } + ast = parser.parse(spec) + loop_nest = interpreter.interpret(ast, root="C") + print(loop_nest.root_node.pretty_print()) # CHECK-SIMPLE: loop i # CHECK-SIMPLE-NEXT: loop k # CHECK-SIMPLE-NEXT: loop j @@ -147,3 +170,17 @@ # CHECK-FUSE-CONSUMER-NEXT: loop k # CHECK-FUSE-CONSUMER-NEXT: loop j // fuse_consumer # CHECK-FUSE-CONSUMER-NEXT: ... + +# CHECK-GPU: loop i +# CHECK-GPU-NEXT: tile(i, 32) // gpu_block(0) +# CHECK-GPU-NEXT: tile(i, 4) // gpu_thread(0) +# CHECK-GPU-NEXT: loop j +# CHECK-GPU-NEXT: loop k +# CHECK-GPU-NEXT: ... + +# CHECK-GPU-WARP-LANE: loop i // gpu_block(0) +# CHECK-GPU-WARP-LANE-NEXT: tile(i, 32) // gpu_warp(0) +# CHECK-GPU-WARP-LANE-NEXT: tile(i, 4) // gpu_lane(0) +# CHECK-GPU-WARP-LANE-NEXT: loop j +# CHECK-GPU-WARP-LANE-NEXT: loop k +# CHECK-GPU-WARP-LANE-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 index 225920f9..6fc3445e 100644 --- a/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py +++ b/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py @@ -30,9 +30,9 @@ "J": {"gpu_block": 1}, "K": {}, "I#128": {"gpu_warp": 0}, - "J#128": {}, + "J#32": {}, "I#32": {"gpu_lane": 0}, - "J#32": {"gpu_lane": 1}, + "J#4": {}, } ) @@ -42,7 +42,7 @@ comp = impl.get_compiler( target=gpu, shared_lib=True, - dump_file="matmul_descript_mlir_gpu", + dump_file="matmul_descript_mlir_gpu_warp_lane", print_source_ir=True, print_transformed_ir=True, ) @@ -70,25 +70,30 @@ # 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: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 32, 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: %tiled_op_4, %forall_op_5 = transform.structured.tile_using_forall %tiled_linalg_op_2 tile_sizes [32, 0, 0](mapping = [#gpu.warp]) : (!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 [0, 4, 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: %tiled_op_8, %forall_op_9 = transform.structured.tile_using_forall %tiled_linalg_op_6 tile_sizes [1, 0, 0](mapping = [#gpu.lane]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_9 "C/I1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_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: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [128, 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: #map2 = affine_map<()[s0] -> (s0 floordiv 32)> +# CHECK-NEXT: #map3 = affine_map<()[s0, s1, s2] -> (s0 + s1 * 128 + s2 * 128)> +# CHECK-NEXT: #map4 = affine_map<()[s0] -> (s0 mod 32)> # 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 @@ -97,68 +102,79 @@ # 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: %c0_4 = arith.constant 0 : index +# CHECK-NEXT: %c1024_5 = arith.constant 1024 : index +# CHECK-NEXT: %c1_6 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_4 to %c1024_5 step %c1_6 { +# CHECK-NEXT: %subview_7 = 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_7 : 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: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c1_1 = arith.constant 1 : 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: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %c1_3 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c32, %arg11 = %c1_3) threads(%arg6, %arg7, %arg8) in (%arg12 = %c128, %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: %0 = affine.apply #map(%block_id_x) -# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %1 = affine.apply #map1(%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: %subview_6 = memref.subview %arg1[0, %1] [512, 32] [1, 1] : memref<512x1024xf32> to memref<512x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %arg2[%0, %1] [128, 32] [1, 1] : memref<1024x1024xf32> to memref<128x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_8 = 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: %c1_9 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg15 = %c0_8 to %c512 step %c1_9 { +# CHECK-NEXT: %subview_10 = 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_11 = memref.subview %subview_6[%arg15, 0] [1, 32] [1, 1] : memref<512x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_12 = memref.subview %subview_7[0, 0] [128, 32] [1, 1] : memref<128x32xf32, strided<[1024, 1], offset: ?>> to memref<128x32xf32, 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 #map2()[%thread_id_x] +# CHECK-NEXT: %3 = affine.apply #map1(%2) +# CHECK-NEXT: %subview_13 = memref.subview %subview_10[%3, 0] [32, 1] [1, 1] : memref<128x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_11[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_12[%3, 0] [32, 32] [1, 1] : memref<128x32xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_16 = arith.constant 0 : index +# CHECK-NEXT: %c32_17 = arith.constant 32 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: scf.for %arg16 = %c0_16 to %c32_17 step %c4 { +# CHECK-NEXT: %subview_18 = memref.subview %subview_13[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_14[0, %arg16] [1, 4] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_15[0, %arg16] [32, 4] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<32x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x_21 = gpu.thread_id x +# CHECK-NEXT: %thread_id_y_22 = gpu.thread_id y +# CHECK-NEXT: %thread_id_z_23 = gpu.thread_id z +# CHECK-NEXT: %4 = affine.apply #map3()[%thread_id_x_21, %c0_4, %c0_4] +# CHECK-NEXT: %5 = affine.apply #map4()[%thread_id_x_21] +# CHECK-NEXT: %6 = affine.apply #map4()[%thread_id_x_21] +# CHECK-NEXT: %c32_24 = arith.constant 32 : index +# CHECK-NEXT: %7 = arith.cmpi ult, %5, %c32_24 : index +# CHECK-NEXT: scf.if %7 { +# CHECK-NEXT: %subview_25 = memref.subview %subview_18[%6, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_19[0, 0] [1, 4] [1, 1] : memref<1x4xf32, strided<[1024, 1], offset: ?>> to memref<1x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_27 = memref.subview %subview_20[%6, 0] [1, 4] [1, 1] : memref<32x4xf32, strided<[1024, 1], offset: ?>> to memref<1x4xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_28 = arith.constant 0 : index +# CHECK-NEXT: %c4_29 = arith.constant 4 : index +# CHECK-NEXT: %c1_30 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg17 = %c0_28 to %c4_29 step %c1_30 { +# CHECK-NEXT: %subview_31 = memref.subview %subview_25[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_26[0, %arg17] [1, 1] [1, 1] : memref<1x4xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_33 = memref.subview %subview_27[0, %arg17] [1, 1] [1, 1] : memref<1x4xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_31, %subview_32 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_33 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"C/J1"} +# CHECK-NEXT: } +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: } {"C/J0"} +# CHECK-NEXT: gpu.barrier # CHECK-NEXT: } {"C/K"} # CHECK-NEXT: gpu.terminator # CHECK-NEXT: } From c00df99befe5622a33aad2ea531d5b91830c539a Mon Sep 17 00:00:00 2001 From: Rui Cesista Date: Wed, 23 Sep 2026 11:58:06 +0200 Subject: [PATCH 8/8] mlir: Add conv2d test for nvgpu and add check in loop nest to verify that we don't put loop to several gpu_primitive --- src/xtc/schedules/loop_nest.py | 29 +++ src/xtc/schedules/parameter_loop_nest.py | 29 +++ .../target_gpu/test_conv2d_r181_gpu_mlir.py | 230 ++++++++++++++++++ 3 files changed, 288 insertions(+) create mode 100644 tests/filecheck/backends/target_gpu/test_conv2d_r181_gpu_mlir.py diff --git a/src/xtc/schedules/loop_nest.py b/src/xtc/schedules/loop_nest.py index c3c3b6e7..b13d448c 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/src/xtc/schedules/loop_nest.py @@ -6,6 +6,7 @@ from typing import Generic, TypeVar from dataclasses import dataclass, field +from itertools import combinations from .exceptions import ScheduleValidationError @@ -385,6 +386,7 @@ def check(self): self._check_external_consistency() self._check_tiling_consistency(info) self._check_sizes(info) + self._check_gpu_consistency() def _check_use_defined_dims(self, info: LoopInfo): for dim in self.abstract_dims: @@ -505,6 +507,33 @@ def _check_sizes(self, info: LoopInfo): f'`{{"unroll" = {unroll_factor}}}`: unroll factor should be smaller than {loop_size}.' ) + def _check_gpu_consistency(self) -> None: + for sched in self.nodes: + gpu_sets = { + "gpu_block": set(sched.gpu_block.keys()), + "gpu_thread": set(sched.gpu_thread.keys()), + "gpu_lane": set(sched.gpu_lane.keys()), + "gpu_warp": set(sched.gpu_warp.keys()), + } + + primitive_names = list(gpu_sets.keys()) + for prim1, prim2 in combinations(primitive_names, 2): + overlap = gpu_sets[prim1] & gpu_sets[prim2] + if overlap: + loops_str = ", ".join(sorted(overlap)) + raise ScheduleValidationError( + f"Loops {loops_str} appear in both {prim1} and {prim2}." + ) + + has_block = bool(sched.gpu_block) + has_thread_or_lane_or_warp = ( + bool(sched.gpu_thread) or bool(sched.gpu_lane) or bool(sched.gpu_warp) + ) + if has_block and not has_thread_or_lane_or_warp: + raise ScheduleValidationError( + "gpu_block requires either gpu_thread or gpu_lane or gpu_warp to be specified." + ) + @staticmethod def _must_be_smaller_routine( new_size: int, current_sizes: dict[str, int | None], loop_name: str, axis: str diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index 23ff548e..85a65799 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from typing import Generic, TypeVar, Any +from itertools import combinations from xtc.schedules.loop_nest import LoopNest, LoopNestNode, SplitOrigin @@ -472,6 +473,7 @@ def check(self): self._check_vectorization_consistency() self._check_tiling_consistency(info) self._check_sizes(info) + self._check_gpu_consistency() def apply_sample(self, sample: dict[str, int]) -> LoopNest: """ @@ -586,6 +588,33 @@ def _check_sizes(self, info: ParameterLoopInfo): f'`{{"unroll" = {unroll_factor}}}`: unroll factor should be smaller than {loop_size}.' ) + def _check_gpu_consistency(self) -> None: + for sched in self.nodes: + gpu_sets = { + "gpu_block": set(sched.gpu_block.keys()), + "gpu_thread": set(sched.gpu_thread.keys()), + "gpu_lane": set(sched.gpu_lane.keys()), + "gpu_warp": set(sched.gpu_warp.keys()), + } + + primitive_names = list(gpu_sets.keys()) + for prim1, prim2 in combinations(primitive_names, 2): + overlap = gpu_sets[prim1] & gpu_sets[prim2] + if overlap: + loops_str = ", ".join(sorted(overlap)) + raise ScheduleValidationError( + f"Loops {loops_str} appear in both {prim1} and {prim2}." + ) + + has_block = bool(sched.gpu_block) + has_thread_or_lane_or_warp = ( + bool(sched.gpu_thread) or bool(sched.gpu_lane) or bool(sched.gpu_warp) + ) + if has_block and not has_thread_or_lane_or_warp: + raise ScheduleValidationError( + "gpu_block requires either gpu_thread or gpu_lane or gpu_warp to be specified." + ) + @staticmethod def _must_be_smaller_routine( new_size: literal, diff --git a/tests/filecheck/backends/target_gpu/test_conv2d_r181_gpu_mlir.py b/tests/filecheck/backends/target_gpu/test_conv2d_r181_gpu_mlir.py new file mode 100644 index 00000000..458e26c8 --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_conv2d_r181_gpu_mlir.py @@ -0,0 +1,230 @@ +# 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.artifacts import get_operation +from xtc.artifacts import get_operation + +from xtc.runtimes.accelerator.gpu import GPUDevice + +# Create device +gpu = GPUDevice() +op = get_operation("conv2d", "ResNet18_01") +N, H, W, F, R, S, C = [op["dims"][k] for k in ["n", "h", "w", "f", "r", "s", "c"]] +SH, SW = [op["params"][k] for k in ["SH", "SW"]] +dtype = "float32" + +a = O.tensor((N, H + R - 1, W + S - 1, C), dtype, device=gpu) +b = O.tensor((R, S, C, F), dtype, device=gpu) + +with O.graph(name="conv2d_nhwc_r181") as gb: + O.conv2d(a, b, stride=(SH, SW), name="O", device=gpu) + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.tile("w", {"w1": 4}) +sch.tile("f", {"f1": 16}) +sch.gpu_block(["b"]) +sch.gpu_thread(["w1", "f1"]) +sch.interchange(["b", "h", "w", "f", "r", "s", "c", "w1", "f1"]) +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="conv2d_nhwc_r181_gpu_mlir", + 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: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d6)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_r181(%arg0: memref<1x230x230x3xf32> {llvm.noalias, memref.on_device}, %arg1: memref<7x7x3x64xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1x112x112x64xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_O_0_} ins(%cst : f32) outs(%arg2 : memref<1x112x112x64xf32>) +# CHECK-NEXT: linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %arg1 : memref<1x230x230x3xf32>, memref<7x7x3x64xf32>) outs(%arg2 : memref<1x112x112x64xf32>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_0: f32, %out: f32): +# CHECK-NEXT: %0 = arith.mulf %in, %in_0 fastmath : f32 +# CHECK-NEXT: %1 = arith.addf %out, %0 fastmath : f32 +# CHECK-NEXT: linalg.yield %1 : f32 +# CHECK-NEXT: } +# 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_O_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, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./b" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./h" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_linalg_op_0 tile_sizes [0, 0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./w" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./f" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_O_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [1, 0, 0, 0, 0, 0, 0](mapping = [#gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./b" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 1, 0, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./h" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 0, 4, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "./w" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 0, 0, 16, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./f" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [0, 0, 0, 0, 1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "./r" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 0, 0, 0, 0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "./s" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_16, %loops_17 = transform.structured.tile_using_for %tiled_linalg_op_14 tile_sizes [0, 0, 0, 0, 0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_17 "./c" : !transform.any_op +# CHECK-NEXT: %tiled_op_18, %forall_op_19 = transform.structured.tile_using_forall %tiled_linalg_op_16 tile_sizes [0, 0, 1, 1, 0, 0, 0](mapping = [#gpu.thread, #gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_19 "./w1" : !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, 16, 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: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d6)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_r181(%arg0: memref<1x230x230x3xf32> {llvm.noalias, memref.on_device}, %arg1: memref<7x7x3x64xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1x112x112x64xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1 step %c1_0 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0, 0, 0] [1, 112, 112, 64] [1, 1, 1, 1] : memref<1x112x112x64xf32> to memref<1x112x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c112 = arith.constant 112 : index +# CHECK-NEXT: %c1_7 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_6 to %c112 step %c1_7 { +# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg4, 0, 0] [1, 1, 112, 64] [1, 1, 1, 1] : memref<1x112x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_9 = arith.constant 0 : index +# CHECK-NEXT: %c112_10 = arith.constant 112 : index +# CHECK-NEXT: %c1_11 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg5 = %c0_9 to %c112_10 step %c1_11 { +# CHECK-NEXT: %subview_12 = memref.subview %subview_8[0, 0, %arg5, 0] [1, 1, 1, 64] [1, 1, 1, 1] : memref<1x1x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x1x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_13 = arith.constant 0 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %c1_14 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg6 = %c0_13 to %c64 step %c1_14 { +# CHECK-NEXT: %subview_15 = memref.subview %subview_12[0, 0, 0, %arg6] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_O_0_} ins(%cst : f32) outs(%subview_15 : memref<1x1x1x1xf32, strided<[802816, 7168, 64, 1], offset: ?>>) +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: } {"./b"} +# CHECK-NEXT: %c1_1 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c1_3 = arith.constant 1 : 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 = %c1_3, %arg10 = %c1_4, %arg11 = %c1_5) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c16, %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: %subview = memref.subview %arg0[%block_id_x, 0, 0, 0] [1, 229, 229, 3] [1, 1, 1, 1] : memref<1x230x230x3xf32> to memref<1x229x229x3xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg1[0, 0, 0, 0] [7, 7, 3, 64] [1, 1, 1, 1] : memref<7x7x3x64xf32> to memref<7x7x3x64xf32, strided<[1344, 192, 64, 1]>> +# CHECK-NEXT: %subview_9 = memref.subview %arg2[%block_id_x, 0, 0, 0] [1, 112, 112, 64] [1, 1, 1, 1] : memref<1x112x112x64xf32> to memref<1x112x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_10 = arith.constant 0 : index +# CHECK-NEXT: %c112 = arith.constant 112 : index +# CHECK-NEXT: %c1_11 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg15 = %c0_10 to %c112 step %c1_11 { +# CHECK-NEXT: %0 = affine.apply #map(%arg15) +# CHECK-NEXT: %subview_12 = memref.subview %subview[0, %0, 0, 0] [1, 7, 229, 3] [1, 1, 1, 1] : memref<1x229x229x3xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x7x229x3xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_8[0, 0, 0, 0] [7, 7, 3, 64] [1, 1, 1, 1] : memref<7x7x3x64xf32, strided<[1344, 192, 64, 1]>> to memref<7x7x3x64xf32, strided<[1344, 192, 64, 1]>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_9[0, %arg15, 0, 0] [1, 1, 112, 64] [1, 1, 1, 1] : memref<1x112x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_15 = arith.constant 0 : index +# CHECK-NEXT: %c112_16 = arith.constant 112 : index +# CHECK-NEXT: %c4_17 = arith.constant 4 : index +# CHECK-NEXT: scf.for %arg16 = %c0_15 to %c112_16 step %c4_17 { +# CHECK-NEXT: %1 = affine.apply #map(%arg16) +# CHECK-NEXT: %subview_18 = memref.subview %subview_12[0, 0, %1, 0] [1, 7, 13, 3] [1, 1, 1, 1] : memref<1x7x229x3xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x7x13x3xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, 0, 0, 0] [7, 7, 3, 64] [1, 1, 1, 1] : memref<7x7x3x64xf32, strided<[1344, 192, 64, 1]>> to memref<7x7x3x64xf32, strided<[1344, 192, 64, 1]>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_14[0, 0, %arg16, 0] [1, 1, 4, 64] [1, 1, 1, 1] : memref<1x1x112x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x4x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_21 = arith.constant 0 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: %c16_22 = arith.constant 16 : index +# CHECK-NEXT: scf.for %arg17 = %c0_21 to %c64 step %c16_22 { +# CHECK-NEXT: %subview_23 = memref.subview %subview_18[0, 0, 0, 0] [1, 7, 13, 3] [1, 1, 1, 1] : memref<1x7x13x3xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x7x13x3xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_24 = memref.subview %subview_19[0, 0, 0, %arg17] [7, 7, 3, 16] [1, 1, 1, 1] : memref<7x7x3x64xf32, strided<[1344, 192, 64, 1]>> to memref<7x7x3x16xf32, strided<[1344, 192, 64, 1], offset: ?>> +# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, 0, 0, %arg17] [1, 1, 4, 16] [1, 1, 1, 1] : memref<1x1x4x64xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_26 = arith.constant 0 : index +# CHECK-NEXT: %c7 = arith.constant 7 : index +# CHECK-NEXT: %c1_27 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_26 to %c7 step %c1_27 { +# CHECK-NEXT: %subview_28 = memref.subview %subview_23[0, %arg18, 0, 0] [1, 1, 13, 3] [1, 1, 1, 1] : memref<1x7x13x3xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x1x13x3xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_29 = memref.subview %subview_24[%arg18, 0, 0, 0] [1, 7, 3, 16] [1, 1, 1, 1] : memref<7x7x3x16xf32, strided<[1344, 192, 64, 1], offset: ?>> to memref<1x7x3x16xf32, strided<[1344, 192, 64, 1], offset: ?>> +# CHECK-NEXT: %subview_30 = memref.subview %subview_25[0, 0, 0, 0] [1, 1, 4, 16] [1, 1, 1, 1] : memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_31 = arith.constant 0 : index +# CHECK-NEXT: %c7_32 = arith.constant 7 : index +# CHECK-NEXT: %c1_33 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg19 = %c0_31 to %c7_32 step %c1_33 { +# CHECK-NEXT: %subview_34 = memref.subview %subview_28[0, 0, %arg19, 0] [1, 1, 7, 3] [1, 1, 1, 1] : memref<1x1x13x3xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x1x7x3xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_35 = memref.subview %subview_29[0, %arg19, 0, 0] [1, 1, 3, 16] [1, 1, 1, 1] : memref<1x7x3x16xf32, strided<[1344, 192, 64, 1], offset: ?>> to memref<1x1x3x16xf32, strided<[1344, 192, 64, 1], offset: ?>> +# CHECK-NEXT: %subview_36 = memref.subview %subview_30[0, 0, 0, 0] [1, 1, 4, 16] [1, 1, 1, 1] : memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: %c0_37 = arith.constant 0 : index +# CHECK-NEXT: %c3 = arith.constant 3 : index +# CHECK-NEXT: %c1_38 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg20 = %c0_37 to %c3 step %c1_38 { +# CHECK-NEXT: %subview_39 = memref.subview %subview_34[0, 0, 0, %arg20] [1, 1, 7, 1] [1, 1, 1, 1] : memref<1x1x7x3xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x1x7x1xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_40 = memref.subview %subview_35[0, 0, %arg20, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x3x16xf32, strided<[1344, 192, 64, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1344, 192, 64, 1], offset: ?>> +# CHECK-NEXT: %subview_41 = memref.subview %subview_36[0, 0, 0, 0] [1, 1, 4, 16] [1, 1, 1, 1] : memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x4x16xf32, strided<[802816, 7168, 64, 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 #map(%thread_id_x) +# CHECK-NEXT: %subview_42 = memref.subview %subview_39[0, 0, %2, 0] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x7x1xf32, strided<[158700, 690, 3, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[158700, 690, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_43 = memref.subview %subview_40[0, 0, 0, %thread_id_y] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1344, 192, 64, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1344, 192, 64, 1], offset: ?>> +# CHECK-NEXT: %subview_44 = memref.subview %subview_41[0, 0, %thread_id_x, %thread_id_y] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x4x16xf32, strided<[802816, 7168, 64, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[802816, 7168, 64, 1], offset: ?>> +# CHECK-NEXT: linalg.generic {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%subview_42, %subview_43 : memref<1x1x1x1xf32, strided<[158700, 690, 3, 1], offset: ?>>, memref<1x1x1x1xf32, strided<[1344, 192, 64, 1], offset: ?>>) outs(%subview_44 : memref<1x1x1x1xf32, strided<[802816, 7168, 64, 1], offset: ?>>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_45: f32, %out: f32): +# CHECK-NEXT: %3 = arith.mulf %in, %in_45 fastmath : f32 +# CHECK-NEXT: %4 = arith.addf %out, %3 fastmath : f32 +# CHECK-NEXT: linalg.yield %4 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: } {"./c"} +# CHECK-NEXT: } {"./s"} +# CHECK-NEXT: } {"./r"} +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: conv2d_nhwc_r181 +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1x230x230x3xfloat32 +# CHECK-NEXT: - %1 : 7x7x3x64xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 1x112x112x64xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: conv2d(%0, %1, stride=(2, 2)) {name = 'O'} : [1x230x230x3xfloat32, 7x7x3x64xfloat32] -> [1x112x112x64xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0