diff --git a/src/xtc/backends/jir/JIRScheduler.py b/src/xtc/backends/jir/JIRScheduler.py index 4634bc9d6..226b53394 100644 --- a/src/xtc/backends/jir/JIRScheduler.py +++ b/src/xtc/backends/jir/JIRScheduler.py @@ -361,6 +361,26 @@ def distributed_buffer_at( # TODO: not implemented for now pass + @override + def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + def get_schedule_str(self) -> str: return str(JIRSchedule(scheduler=self)) diff --git a/src/xtc/backends/mlir/MlirCompilerPasses.py b/src/xtc/backends/mlir/MlirCompilerPasses.py index ae4ed40fd..3929f3d10 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 @@ -135,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 [] ) @@ -240,6 +248,13 @@ def _generate_scheduling(self) -> OpResult: ) if schedule.vectorization or self._always_vectorize: self._post_vectorize(scheduling_state, schedule) + + # GPU mapping + if schedule.gpu_blocks: + self._gpu_mapping( + schedule, + scheduling_state, + ) handle = scheduling_state.handle assert handle, "At least 1 operation should have been processed" @@ -321,7 +336,9 @@ def _generate_node_scheduling( permutation = schedule.permutation[root] if not permutation: 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 @@ -357,12 +374,113 @@ 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 = [ + 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( + loop_name=loop_name, + tiling_vector=tile_vect, + mapping_order=mapping_order, + schedule=schedule, + 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) + 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( + loop_name=loop_name, + tiling_vector=tile_vect, + mapping_order=mapping_order, + schedule=schedule, + 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, + tiling_vector=tiles_sizes_by_loops[loop_name], + mapping_order=[], + schedule=schedule, + sched_state=sched_state, + ) if loop_name in schedule.distribution: self._distribute_loop(loop_name, schedule, sched_state) # Fuse the producers @@ -494,13 +612,35 @@ def _strip_mine( self, loop_name: str, tiling_vector: list[int], + mapping_order: list[int], schedule: MlirNodeSchedule, sched_state: SchedulingState, ) -> OpResult: - if loop_name in schedule.parallelization: - tiling_command = TileUsingForallOp( - sched_state.handle, 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.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: tiling_command = TileUsingForOp(sched_state.handle, sizes=tiling_vector) # Extract the results @@ -508,6 +648,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 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) @@ -582,11 +724,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" ] @@ -704,6 +847,141 @@ def _collect_fused_producers(self, unscheduled_handles: set[str | None]): return fused_producers + def _get_lane_id(self, index: int) -> Attribute: + ctx = self._mlir_program.mlir_context + return Attribute.parse(f"#gpu.lane", context=ctx) + + def _get_warp_id(self, index: int) -> Attribute: + ctx = self._mlir_program.mlir_context + return Attribute.parse(f"#gpu.warp<{_GPU_DIM[index]}>", context=ctx) + + def _get_thread_id(self, index: int) -> Attribute: + ctx = self._mlir_program.mlir_context + return Attribute.parse(f"#gpu.thread<{_GPU_DIM[index]}>", context=ctx) + + 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 = self._generate_tiling_insns(schedule) + if schedule.gpu_blocks and not self._using_tensors: + 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 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, + ) + 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 + 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, + ) + 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 diff --git a/src/xtc/backends/mlir/MlirNodeScheduler.py b/src/xtc/backends/mlir/MlirNodeScheduler.py index a87c20cdb..1b12afe21 100644 --- a/src/xtc/backends/mlir/MlirNodeScheduler.py +++ b/src/xtc/backends/mlir/MlirNodeScheduler.py @@ -109,6 +109,26 @@ def fuse_producer_at( def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None: self._plain_sch.fuse_consumer_at(axis, root) + def map_gpu_threads(self, axes: list[str], root: str = DEFAULT_ROOT): + assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu thread" + assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu thread" + self._plain_sch.gpu_thread(axes, root) + + def map_gpu_blocks(self, axes: list[str], root: str = DEFAULT_ROOT): + assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu block" + assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu block" + self._plain_sch.gpu_block(axes, root) + + def map_gpu_lanes(self, axes: list[str], root: str = DEFAULT_ROOT): + assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu lane" + assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu lane" + self._plain_sch.gpu_lane(axes, root) + + def map_gpu_warps(self, axes: list[str], root: str = DEFAULT_ROOT): + assert len(axes) == len(set(axes)), "Duplicate in the axes for gpu warp" + assert len(axes) <= 3, "We cannot map more than 3 dimension for gpu warp" + self._plain_sch.gpu_warp(axes, root) + def get_node_schedule(self) -> MlirNodeSchedule: plain_schedule = self._plain_sch.get_plain_schedule() return MlirNodeSchedule(**asdict(plain_schedule)) diff --git a/src/xtc/backends/mlir/MlirScheduler.py b/src/xtc/backends/mlir/MlirScheduler.py index c3b25ff2a..5d7480c8a 100644 --- a/src/xtc/backends/mlir/MlirScheduler.py +++ b/src/xtc/backends/mlir/MlirScheduler.py @@ -218,6 +218,22 @@ def distributed_buffer_at( axis, input_idx, memory_axes, root=root ) + @override + def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + self._current_scheduler.map_gpu_lanes(axes, root=root) + + @override + def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + self._current_scheduler.map_gpu_warps(axes, root=root) + + @override + def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + self._current_scheduler.map_gpu_threads(axes, root=root) + + @override + def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + self._current_scheduler.map_gpu_blocks(axes, root=root) + @override def get_loop_nest(self) -> LoopNest: node_schedule = self._current_scheduler.get_node_schedule() diff --git a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py index 8173c0840..7a83f424a 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py +++ b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py @@ -442,25 +442,48 @@ def _lowering_pipeline(self, sm_arch: str, ptx_version: str) -> list[str]: "cse", "sccp", # From complex control to the soup of basic blocks - "expand-strided-metadata", - "scf-forall-to-parallel", - "canonicalize", - "cse", - "sccp", - "func.func(gpu-map-parallel-loops)", - "convert-parallel-loops-to-gpu", "convert-linalg-to-loops", "canonicalize", "cse", "sccp", - "convert-vector-to-llvm", "buffer-results-to-out-params", + # GPU to LLVM pipeline + "func.func(gpu-eliminate-barriers)", + "convert-nvgpu-to-nvvm", + "gpu-kernel-outlining", + "gpu-launch-sink-index-computations", + "convert-vector-to-scf", + "convert-vector-to-llvm{vector-contract-lowering=outerproduct}", + "nvgpu-optimize-shared-memory", + "convert-scf-to-cf", + "convert-nvvm-to-llvm", "convert-func-to-llvm{use-bare-ptr-memref-call-conv=true}", - "gpu-lower-to-nvvm-pipeline{cubin-chip=" + "expand-strided-metadata", + "nvvm-attach-target{chip=" + sm_arch - + " cubin-features=+ptx" + + " features=+ptx" + "".join(ptx_version.split(".")) - + " opt-level=3}", + + " O=3 fast=true}", + "lower-affine", + "convert-arith-to-llvm", + "convert-index-to-llvm", + "canonicalize", + "cse", + "reconcile-unrealized-casts", + "gpu.module(convert-gpu-to-nvvm{use-bare-ptr-memref-call-conv=true})", + "gpu.module(canonicalize)", + "gpu.module(cse)", + "gpu.module(reconcile-unrealized-casts)", + "gpu-to-llvm{use-bare-pointers-for-host=true use-bare-pointers-for-kernels=true}", + "gpu.module(reconcile-unrealized-casts)", + "reconcile-unrealized-casts", + "convert-math-to-llvm", + "gpu.module(reconcile-unrealized-casts)", + "reconcile-unrealized-casts", + "gpu-module-to-binary", + "canonicalize", + "cse", + "reconcile-unrealized-casts", ] def run(self, sm_arch: str, ptx_version: str) -> None: diff --git a/src/xtc/backends/tvm/TVMScheduler.py b/src/xtc/backends/tvm/TVMScheduler.py index afe7702a3..87fac7b2a 100644 --- a/src/xtc/backends/tvm/TVMScheduler.py +++ b/src/xtc/backends/tvm/TVMScheduler.py @@ -378,6 +378,26 @@ def distributed_buffer_at( # TODO: not implemented for now pass + @override + def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + + @override + def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + # TODO: not implemented for now + pass + def _get_plain_schedule(self) -> PlainNodeSchedule: return self._plain_sch.get_plain_schedule() diff --git a/src/xtc/itf/schd/scheduler.py b/src/xtc/itf/schd/scheduler.py index 4c2f0a14c..91b297493 100644 --- a/src/xtc/itf/schd/scheduler.py +++ b/src/xtc/itf/schd/scheduler.py @@ -319,6 +319,77 @@ def distributed_buffer_at( """ ... + @abstractmethod + def gpu_lane(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + """Maps iteration axes to GPU lane dimensions. + + The mapping is positional: + - axes[0] -> linear_dim_0 + - axes[1] -> linear_dim_1 + - axes[2] -> linear_dim_2 + At most 3 axes may be provided. If there are less than 3 axes, + only the corresponding lane dimensions are assigned. + By default lane are mapped on 32 id. + + Args: + axes: list of maximum 3 element that are maps to the dimension x, y, z respectively + root: the parent split (or the operator's absolute root) + """ + ... + + @abstractmethod + def gpu_warp(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + """Maps iteration axes to GPU warp dimensions. + + The mapping is positional: + - axes[0] -> warp for threadIdx.x + - axes[1] -> warp for threadIdx.y + - axes[2] -> warp for threadIdx.z + At most 3 axes may be provided. If there are less than 3 axes, + only the corresponding warp dimensions are assigned. + It is preferable to map warp only on the x axis. + The thread size of the warp should be a multiple of 32. + + Args: + axes: list of maximum 3 element that are maps to the dimension x, y, z respectively + root: the parent split (or the operator's absolute root) + """ + ... + + @abstractmethod + def gpu_thread(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + """Maps iteration axes to GPU thread dimensions. + + The mapping is positional: + - axes[0] -> threadIdx.x + - axes[1] -> threadIdx.y + - axes[2] -> threadIdx.z + At most 3 axes may be provided. If there are less than 3 axes, + only the corresponding thread dimensions are assigned. + + Args: + axes: list of maximum 3 element that are maps to the dimension x, y, z respectively + root: the parent split (or the operator's absolute root) + """ + ... + + @abstractmethod + def gpu_block(self, axes: list[str], root: str = DEFAULT_ROOT) -> None: + """ + Maps iteration axes to GPU block dimensions. + The mapping is positional: + - axes[0] -> blockIdx.x + - axes[1] -> blockIdx.y + - axes[2] -> blockIdx.z + At most 3 axes may be provided. If there are less than 3 axes, + only the corresponding block dimensions are assigned. + + Args: + axes: list of maximum 3 element that are maps to the dimension x, y, z respectively + root: the parent split (or the operator's absolute root) + """ + ... + @abstractmethod def get_loop_nest(self) -> LoopNest: """Return a LoopNest representation of the current schedule. diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 99a94c8b3..3bfc6a9a7 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -401,6 +401,18 @@ def _apply_annotations( if annotations.fuse_consumer: node.fuse_consumer_at.append(loop_name) + if annotations.gpu_lane is not None: + node.gpu_lane[loop_name] = annotations.gpu_lane + + if annotations.gpu_warp is not None: + node.gpu_warp[loop_name] = annotations.gpu_warp + + if annotations.gpu_block is not None: + node.gpu_block[loop_name] = annotations.gpu_block + + if annotations.gpu_thread is not None: + node.gpu_thread[loop_name] = annotations.gpu_thread + def _check_splitting_intervals( self, item: SplitDecl, @@ -525,7 +537,6 @@ def _apply_loop_nest(self, loop_nest: LoopNest, scheduler: Scheduler) -> None: def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None: """Recursively apply a LoopNestNode and its children to the scheduler.""" root = node.root - for d, s in node.splits.items(): scheduler.split(d, s, root=root) @@ -548,6 +559,35 @@ def _apply_node(self, node: LoopNestNode, scheduler: Scheduler) -> None: for axis in node.fuse_consumer_at: scheduler.fuse_consumer_at(axis, root=root) + + if node.gpu_lane: + sorted_keys = sorted( + (k for k, v in node.gpu_lane.items() if v is not None), + key=lambda k: node.gpu_lane[k], + ) + scheduler.gpu_lane(sorted_keys, root=root) + + if node.gpu_warp: + sorted_keys = sorted( + (k for k, v in node.gpu_warp.items() if v is not None), + key=lambda k: node.gpu_warp[k], + ) + scheduler.gpu_thread(sorted_keys, root=root) + + if node.gpu_block: + sorted_keys = sorted( + (k for k, v in node.gpu_block.items() if v is not None), + key=lambda k: node.gpu_block[k], + ) + scheduler.gpu_block(sorted_keys, root=root) + + if node.gpu_thread: + sorted_keys = sorted( + (k for k, v in node.gpu_thread.items() if v is not None), + key=lambda k: node.gpu_thread[k], + ) + scheduler.gpu_thread(sorted_keys, root=root) + # Recursively apply children for child in node.children: self._apply_node(child, scheduler) diff --git a/src/xtc/schedules/loop_nest.py b/src/xtc/schedules/loop_nest.py index bee0f67e0..8c6b9f857 100644 --- a/src/xtc/schedules/loop_nest.py +++ b/src/xtc/schedules/loop_nest.py @@ -118,6 +118,10 @@ class LoopNestNode(Node["LoopNestNode"]): pack_at: dict[str, tuple[int, str | None, bool]] = field(default_factory=dict) fuse_producer_at: dict[str, int] = field(default_factory=dict) fuse_consumer_at: list[str] = field(default_factory=list) + gpu_lane: dict[str, int] = field(default_factory=dict) + gpu_warp: dict[str, int] = field(default_factory=dict) + gpu_block: dict[str, int] = field(default_factory=dict) + gpu_thread: dict[str, int] = field(default_factory=dict) def pretty_print(self, indent: int = 0) -> str: """Return a human-readable representation of the loop nest. @@ -245,6 +249,14 @@ def _add_annotations(self, line: str, loop_name: str) -> str: annotations.append(f"fuse_producer({prod_idx})") if loop_name in self.fuse_consumer_at: annotations.append("fuse_consumer") + if loop_name in self.gpu_lane: + annotations.append(f"gpu_lane({self.gpu_lane[loop_name]})") + if loop_name in self.gpu_warp: + annotations.append(f"gpu_warp({self.gpu_warp[loop_name]})") + if loop_name in self.gpu_block: + annotations.append(f"gpu_block({self.gpu_block[loop_name]})") + if loop_name in self.gpu_thread: + annotations.append(f"gpu_thread({self.gpu_thread[loop_name]})") if annotations: line += " // " + ", ".join(annotations) return line diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index 822073fa9..3f15aa988 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -111,6 +111,12 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): fuse_producer_at: Producer fusion configuration per axis. Maps axis names to producer indices. fuse_consumer_at: List of axes where the output consumer is fused. + gpu_block: Maps loops to block id, that loop need to be parallelize + gpu_thread: Maps loops to thread id, that loop need to be parallelize + gpu_lane: Maps loops to lane id + gpu_warp: Maps loops to warp id + gpu_block: Maps loops to block id + gpu_thread: Maps loops to thread id """ root: str @@ -127,6 +133,10 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): fuse_producer_at: dict[str, int] = field(default_factory=dict) fuse_consumer_at: list[str] = field(default_factory=list) constraints: list[str] = field(default_factory=list) + gpu_lane: dict[str, int] = field(default_factory=dict) + gpu_warp: dict[str, int] = field(default_factory=dict) + gpu_block: dict[str, int] = field(default_factory=dict) + gpu_thread: dict[str, int] = field(default_factory=dict) def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: """ @@ -174,6 +184,10 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: if self.split_origin is not None else None ) + gpu_warp = self.gpu_warp + gpu_lane = self.gpu_lane + gpu_block = self.gpu_block + gpu_thread = self.gpu_thread return LoopNestNode( root=root, tiles=tiles, @@ -188,6 +202,10 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: fuse_consumer_at=fuse_consumer_at, children=children, split_origin=split_origin, + gpu_lane=gpu_lane, + gpu_warp=gpu_warp, + gpu_block=gpu_block, + gpu_thread=gpu_thread, ) def pretty_print(self, indent: int = 0) -> str: @@ -316,6 +334,14 @@ def _add_annotations(self, line: str, loop_name: str) -> str: annotations.append(f"fuse_producer({prod_idx})") if loop_name in self.fuse_consumer_at: annotations.append("fuse_consumer") + if loop_name in self.gpu_lane: + annotations.append(f"gpu_lane({self.gpu_lane[loop_name]})") + if loop_name in self.gpu_warp: + annotations.append(f"gpu_warp({self.gpu_warp[loop_name]})") + if loop_name in self.gpu_block: + annotations.append(f"gpu_block({self.gpu_block[loop_name]})") + if loop_name in self.gpu_thread: + annotations.append(f"gpu_thread({self.gpu_thread[loop_name]})") if annotations: line += " // " + ", ".join(annotations) return line diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 3fe10568d..d7ad51592 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -13,6 +13,7 @@ from .exceptions import ScheduleParseError literal = int | str +_GPU_DIM = {"x": 0, "y": 1, "z": 2} def toliteral(s: str) -> literal: @@ -54,6 +55,10 @@ 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 @dataclass(frozen=True) @@ -181,6 +186,10 @@ 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 for key, param in value.items(): match key: @@ -243,12 +252,71 @@ 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): + 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' + ) + 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): + 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' + ) + 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}") 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, @@ -262,6 +330,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations fuse_consumer=fuse_consumer, partial=partial, full=full, + gpu_block=gpu_block, + gpu_thread=gpu_thread, ) def _parse_pack_param( diff --git a/src/xtc/schedules/plain_schedule.py b/src/xtc/schedules/plain_schedule.py index 4514fd453..dff7bf1c2 100644 --- a/src/xtc/schedules/plain_schedule.py +++ b/src/xtc/schedules/plain_schedule.py @@ -33,6 +33,10 @@ class PlainNodeSchedule: distributed_buffers: dict[str, dict] fused: list[tuple[str, int]] fused_consumers: list[str] + gpu_blocks: list[str] + gpu_threads: list[str] + gpu_lanes: list[str] + gpu_warps: list[str] # Optional caller-provided vector sizes, keyed by vectorized axis name. # When an axis has a size, its dimension is vectorized with masking for # non-divisible extents; axes absent from this mapping are vectorized to @@ -116,6 +120,10 @@ def __init__( self.distributed_buffers: dict[str, dict] = {} self.fused: list[tuple[str, int]] = [] self.fused_consumers: list[str] = [] + self.gpu_blocks: list[str] = [] + self.gpu_threads: list[str] = [] + self.gpu_lanes: list[str] = [] + self.gpu_warps: list[str] = [] def get_plain_schedule(self) -> PlainNodeSchedule: return PlainNodeSchedule( @@ -137,6 +145,10 @@ def get_plain_schedule(self) -> PlainNodeSchedule: distributed_buffers=deepcopy(self.distributed_buffers), fused=deepcopy(self.fused), fused_consumers=deepcopy(self.fused_consumers), + gpu_blocks=deepcopy(self.gpu_blocks), + gpu_threads=deepcopy(self.gpu_threads), + gpu_lanes=deepcopy(self.gpu_lanes), + gpu_warps=deepcopy(self.gpu_warps), vectorization_sizes=deepcopy(self.vectorization_sizes), ) @@ -274,3 +286,23 @@ 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_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_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_1D_offload_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_1D_offload_tensor.py new file mode 100644 index 000000000..8f1169655 --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_1D_offload_tensor.py @@ -0,0 +1,212 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir.MlirGraphBackend import MlirGraphBackend as Backend + +from xtc.runtimes.accelerator.gpu import GPUDevice + +# Create device +gpu = GPUDevice() + +I, J, K, dtype = 512, 512, 512, "float32" +a = O.tensor((I, K), dtype, name="A") # A lives on the host +b = O.tensor((K, J), dtype, name="B", device=gpu) # B lives on the accelerator + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) # C must live on the accelerator + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.tile("i", {"i1": 128, "i2": 32}) +sch.tile("j", {"j1": 128, "j2": 32}) +sch.tile("k", {"k1": 64}) +sch.unroll({"i2": 2}) +sch.gpu_block(["i"]) +sch.gpu_thread(["i1"]) +sch.interchange(["i", "j", "i1", "j1","k", "k1", "i2", "j2"]) +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="gpu_matmul_mlir_1D_offload_tensor", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<512x512xf32>, memref<512x512xf32>) outs(%arg2 : memref<512x512xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 0, 0](mapping = [#gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 128, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./j" : !transform.any_op +# CHECK-NEXT: %tiled_op_4, %forall_op_5 = transform.structured.tile_using_forall %tiled_linalg_op_2 tile_sizes [32, 0, 0](mapping = [#gpu.thread]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_5 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_op_4 tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./j1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 0, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "./k" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./k1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "./i2" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "./j2" : !transform.any_op +# CHECK-NEXT: transform.loop.unroll %loops_13 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %3 = transform.gpu.map_nested_forall_to_threads %2 block_dims = [4, 1, 1] : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<512x512xf32> {llvm.noalias}, %arg1: memref<512x512xf32> {llvm.noalias, memref.on_device}, %arg2: memref<512x512xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c512 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 512] [1, 1] : memref<512x512xf32> to memref<1x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c512_7 = arith.constant 512 : index +# CHECK-NEXT: %c1_8 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_6 to %c512_7 step %c1_8 { +# CHECK-NEXT: %subview_9 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x512xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_9 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c1_1 = arith.constant 1 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: %c4_3 = arith.constant 4 : index +# CHECK-NEXT: %c1_4 = arith.constant 1 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c4_3, %arg10 = %c1_4, %arg11 = %c1_5) threads(%arg6, %arg7, %arg8) in (%arg12 = %c4, %arg13 = %c1_1, %arg14 = %c1_2) { +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c0_7 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %arg1[0, 0] [512, 512] [1, 1] : memref<512x512xf32> to memref<512x512xf32, strided<[512, 1]>> +# CHECK-NEXT: %subview_9 = memref.subview %arg2[%0, 0] [128, 512] [1, 1] : memref<512x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_10 = arith.constant 0 : index +# CHECK-NEXT: %c512_11 = arith.constant 512 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: scf.for %arg15 = %c0_10 to %c512_11 step %c128 { +# CHECK-NEXT: %subview_12 = memref.subview %subview[0, 0] [128, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_8[0, %arg15] [512, 128] [1, 1] : memref<512x512xf32, strided<[512, 1]>> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_9[0, %arg15] [128, 128] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %thread_id_x = gpu.thread_id x +# CHECK-NEXT: %thread_id_y = gpu.thread_id y +# CHECK-NEXT: %thread_id_z = gpu.thread_id z +# CHECK-NEXT: %1 = affine.apply #map1(%thread_id_x) +# CHECK-NEXT: %subview_15 = memref.subview %subview_12[%1, 0] [32, 512] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_13[0, 0] [512, 128] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_17 = memref.subview %subview_14[%1, 0] [32, 128] [1, 1] : memref<128x128xf32, strided<[512, 1], offset: ?>> to memref<32x128xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_18 = arith.constant 0 : index +# CHECK-NEXT: %c128_19 = arith.constant 128 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg16 = %c0_18 to %c128_19 step %c32 { +# CHECK-NEXT: %subview_20 = memref.subview %subview_15[0, 0] [32, 512] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_21 = memref.subview %subview_16[0, %arg16] [512, 32] [1, 1] : memref<512x128xf32, strided<[512, 1], offset: ?>> to memref<512x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_22 = memref.subview %subview_17[0, %arg16] [32, 32] [1, 1] : memref<32x128xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_23 = arith.constant 0 : index +# CHECK-NEXT: %c512_24 = arith.constant 512 : index +# CHECK-NEXT: %c64 = arith.constant 64 : index +# CHECK-NEXT: scf.for %arg17 = %c0_23 to %c512_24 step %c64 { +# CHECK-NEXT: %subview_25 = memref.subview %subview_20[0, %arg17] [32, 64] [1, 1] : memref<32x512xf32, strided<[512, 1], offset: ?>> to memref<32x64xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_21[%arg17, 0] [64, 32] [1, 1] : memref<512x32xf32, strided<[512, 1], offset: ?>> to memref<64x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_27 = memref.subview %subview_22[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_28 = arith.constant 0 : index +# CHECK-NEXT: %c64_29 = arith.constant 64 : index +# CHECK-NEXT: %c1_30 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_28 to %c64_29 step %c1_30 { +# CHECK-NEXT: %subview_31 = memref.subview %subview_25[0, %arg18] [32, 1] [1, 1] : memref<32x64xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_26[%arg18, 0] [1, 32] [1, 1] : memref<64x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_33 = memref.subview %subview_27[0, 0] [32, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<32x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_34 = arith.constant 0 : index +# CHECK-NEXT: %c32_35 = arith.constant 32 : index +# CHECK-NEXT: %c1_36 = arith.constant 1 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index +# CHECK-NEXT: scf.for %arg19 = %c0_34 to %c32_35 step %c2 { +# CHECK-NEXT: %subview_37 = memref.subview %subview_31[%arg19, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_38 = memref.subview %subview_32[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_39 = memref.subview %subview_33[%arg19, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_40 = arith.constant 0 : index +# CHECK-NEXT: %c32_41 = arith.constant 32 : index +# CHECK-NEXT: %c1_42 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg20 = %c0_40 to %c32_41 step %c1_42 { +# CHECK-NEXT: %subview_50 = memref.subview %subview_37[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_51 = memref.subview %subview_38[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_52 = memref.subview %subview_39[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_50, %subview_51 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_52 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: %c1_43 = arith.constant 1 : index +# CHECK-NEXT: %2 = arith.muli %c1_36, %c1_43 : index +# CHECK-NEXT: %3 = arith.addi %arg19, %2 : index +# CHECK-NEXT: %subview_44 = memref.subview %subview_31[%3, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_45 = memref.subview %subview_32[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_46 = memref.subview %subview_33[%3, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[512, 1], offset: ?>> to memref<1x32xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %c0_47 = arith.constant 0 : index +# CHECK-NEXT: %c32_48 = arith.constant 32 : index +# CHECK-NEXT: %c1_49 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg20 = %c0_47 to %c32_48 step %c1_49 { +# CHECK-NEXT: %subview_50 = memref.subview %subview_44[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_51 = memref.subview %subview_45[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_52 = memref.subview %subview_46[0, %arg20] [1, 1] [1, 1] : memref<1x32xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_50, %subview_51 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[512, 1], offset: ?>>) outs(%subview_52 : memref<1x1xf32, strided<[512, 1], offset: ?>>) +# CHECK-NEXT: } {"./j2"} +# CHECK-NEXT: } {"./i2"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: } {"./j1"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 512x512xfloat32 +# CHECK-NEXT: - %1 : 512x512xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 512x512xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [512x512xfloat32, 512x512xfloat32] -> [512x512xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_lane.py new file mode 100644 index 000000000..5ee8f0e99 --- /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_lane", + 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.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py index 7c1f84923..6138d3a69 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor.py @@ -9,7 +9,7 @@ # Create device gpu = GPUDevice() -I, J, K, dtype = 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,10 +22,13 @@ 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({"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() comp = impl.get_compiler( @@ -41,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<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}) { @@ -58,92 +61,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<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: %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: %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 %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: 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: %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: %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) +# CHECK-NEXT: %3 = affine.apply #map1(%thread_id_y) +# 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_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: %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: %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: %4 = arith.muli %c1_25, %c1_32 : index +# CHECK-NEXT: %5 = arith.addi %arg17, %4 : index +# 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"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: gpu.barrier +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } # CHECK-NEXT: return # CHECK-NEXT: } # CHECK-NEXT: } @@ -151,11 +184,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_offload_tensor_vectorise.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py index 0194b7262..9ce23e876 100644 --- a/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_tensor_vectorise.py @@ -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,14 @@ 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": 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"]) +sch.vectorize(["j2"]) sched = sch.schedule() comp = impl.get_compiler( @@ -42,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<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 +62,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, 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 [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: %c8 = arith.constant 8 : 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 %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: ?>> +# 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 +169,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_offload_warp_lane_vectorise.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_offload_warp_lane_vectorise.py new file mode 100644 index 000000000..41d00861a --- /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_warp_lane_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/backends/target_gpu/test_matmul_mlir_with_tensor.py b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py new file mode 100644 index 000000000..16a080d5e --- /dev/null +++ b/tests/filecheck/backends/target_gpu/test_matmul_mlir_with_tensor.py @@ -0,0 +1,287 @@ +# 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.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_with_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: 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: %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) +# 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.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: 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: %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 +# 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: %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: %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: } +# 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 ab671deab..a638b6355 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,12 @@ 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"]) +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 new file mode 100644 index 000000000..277ab9444 --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_gpu.py @@ -0,0 +1,176 @@ +# 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_thread": 0}, + "J#128": {"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<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_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, 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_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: %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: } +# 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/schedules/test_matmul_descript_gpu_warp_lane.py b/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py new file mode 100644 index 000000000..225920f9a --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_gpu_warp_lane.py @@ -0,0 +1,179 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend +from xtc.schedules.descript import descript_scheduler + +from xtc.runtimes.accelerator.gpu import GPUDevice + +gpu = GPUDevice() +I, J, K, dtype = 1024, 1024, 512, "float32" +a = O.tensor((I, K), dtype, name="A", device=gpu) +b = O.tensor((K, J), dtype, name="B", device=gpu) + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C", device=gpu) + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +descript_scheduler( + scheduler = sch, + node_name = "C", + abstract_dims = ["I","J","K"], + spec = { + "I": {"gpu_block": 0}, + "J": {"gpu_block": 1}, + "K": {}, + "I#128": {"gpu_warp": 0}, + "J#128": {}, + "I#32": {"gpu_lane": 0}, + "J#32": {"gpu_lane": 1}, + + } +) + +sched = sch.schedule() + +comp = impl.get_compiler( + target=gpu, + shared_lib=True, + dump_file="matmul_descript_mlir_gpu", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<1024x512xf32>, memref<512x1024xf32>) outs(%arg2 : memref<1024x1024xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [128, 128, 0](mapping = [#gpu.block, #gpu.block]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "C/I" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "C/K" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [32, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "C/I0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "C/J0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "C/I1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "C/J1" : !transform.any_op +# CHECK-NEXT: %2 = transform.gpu.map_forall_to_blocks %forall_op generate_gpu_launch : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 128)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<1024x512xf32> {llvm.noalias, memref.on_device}, %arg1: memref<512x1024xf32> {llvm.noalias, memref.on_device}, %arg2: memref<1024x1024xf32> {llvm.noalias, memref.on_device}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 1024] [1, 1] : memref<1024x1024xf32> to memref<1x1024xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_3 = arith.constant 0 : index +# CHECK-NEXT: %c1024_4 = arith.constant 1024 : index +# CHECK-NEXT: %c1_5 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg4 = %c0_3 to %c1024_4 step %c1_5 { +# CHECK-NEXT: %subview_6 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x1024xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%subview_6 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: %c1_0 = arith.constant 1 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c8_1 = arith.constant 8 : index +# CHECK-NEXT: %c1_2 = arith.constant 1 : index +# CHECK-NEXT: gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c8, %arg10 = %c8_1, %arg11 = %c1_2) threads(%arg6, %arg7, %arg8) in (%arg12 = %c1_0, %arg13 = %c1_0, %arg14 = %c1_0) { +# CHECK-NEXT: %c0_3 = arith.constant 0 : index +# CHECK-NEXT: %block_id_x = gpu.block_id x +# CHECK-NEXT: %block_id_y = gpu.block_id y +# CHECK-NEXT: %block_id_z = gpu.block_id z +# CHECK-NEXT: %0 = affine.apply #map(%block_id_x) +# CHECK-NEXT: %1 = affine.apply #map(%block_id_y) +# CHECK-NEXT: %subview = memref.subview %arg0[%0, 0] [128, 512] [1, 1] : memref<1024x512xf32> to memref<128x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_4 = memref.subview %arg1[0, %1] [512, 128] [1, 1] : memref<512x1024xf32> to memref<512x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_5 = memref.subview %arg2[%0, %1] [128, 128] [1, 1] : memref<1024x1024xf32> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_6 = arith.constant 0 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c1_7 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg15 = %c0_6 to %c512 step %c1_7 { +# CHECK-NEXT: %subview_8 = memref.subview %subview[0, %arg15] [128, 1] [1, 1] : memref<128x512xf32, strided<[512, 1], offset: ?>> to memref<128x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_4[%arg15, 0] [1, 128] [1, 1] : memref<512x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_5[0, 0] [128, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<128x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_11 = arith.constant 0 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg16 = %c0_11 to %c128 step %c32 { +# CHECK-NEXT: %subview_12 = memref.subview %subview_8[%arg16, 0] [32, 1] [1, 1] : memref<128x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_9[0, 0] [1, 128] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_10[%arg16, 0] [32, 128] [1, 1] : memref<128x128xf32, strided<[1024, 1], offset: ?>> to memref<32x128xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_15 = arith.constant 0 : index +# CHECK-NEXT: %c128_16 = arith.constant 128 : index +# CHECK-NEXT: %c32_17 = arith.constant 32 : index +# CHECK-NEXT: scf.for %arg17 = %c0_15 to %c128_16 step %c32_17 { +# CHECK-NEXT: %subview_18 = memref.subview %subview_12[0, 0] [32, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<32x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_19 = memref.subview %subview_13[0, %arg17] [1, 32] [1, 1] : memref<1x128xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_14[0, %arg17] [32, 32] [1, 1] : memref<32x128xf32, strided<[1024, 1], offset: ?>> to memref<32x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_21 = arith.constant 0 : index +# CHECK-NEXT: %c32_22 = arith.constant 32 : index +# CHECK-NEXT: %c1_23 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg18 = %c0_21 to %c32_22 step %c1_23 { +# CHECK-NEXT: %subview_24 = memref.subview %subview_18[%arg18, 0] [1, 1] [1, 1] : memref<32x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_25 = memref.subview %subview_19[0, 0] [1, 32] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_26 = memref.subview %subview_20[%arg18, 0] [1, 32] [1, 1] : memref<32x32xf32, strided<[1024, 1], offset: ?>> to memref<1x32xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %c0_27 = arith.constant 0 : index +# CHECK-NEXT: %c32_28 = arith.constant 32 : index +# CHECK-NEXT: %c1_29 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg19 = %c0_27 to %c32_28 step %c1_29 { +# CHECK-NEXT: %subview_30 = memref.subview %subview_24[0, 0] [1, 1] [1, 1] : memref<1x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_31 = memref.subview %subview_25[0, %arg19] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: %subview_32 = memref.subview %subview_26[0, %arg19] [1, 1] [1, 1] : memref<1x32xf32, strided<[1024, 1], offset: ?>> to memref<1x1xf32, strided<[1024, 1], offset: ?>> +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%subview_30, %subview_31 : memref<1x1xf32, strided<[512, 1], offset: ?>>, memref<1x1xf32, strided<[1024, 1], offset: ?>>) outs(%subview_32 : memref<1x1xf32, strided<[1024, 1], offset: ?>>) +# CHECK-NEXT: } {"C/J1"} +# CHECK-NEXT: } {"C/I1"} +# CHECK-NEXT: } {"C/J0"} +# CHECK-NEXT: } {"C/I0"} +# CHECK-NEXT: } {"C/K"} +# CHECK-NEXT: gpu.terminator +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1024x512xfloat32 +# CHECK-NEXT: - %1 : 512x1024xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 1024x1024xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [1024x512xfloat32, 512x1024xfloat32] -> [1024x1024xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/search/test_conv_oo.py b/tests/filecheck/search/test_conv_oo.py index 4798c0913..da68dc81b 100644 --- a/tests/filecheck/search/test_conv_oo.py +++ b/tests/filecheck/search/test_conv_oo.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 2, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 1, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 2, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 3, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './c1': 3, './s1': 1, './r1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 7, 1] @@ -99,4 +99,4 @@ # CHECK-NEXT: sample 76: [2, 2, 2, 8, 1, 1, 1] # CHECK-NEXT: sample 77: [2, 2, 2, 16, 1, 1, 1] # CHECK-NEXT: stats {'filtered': 78, 'all': 384} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './c1': 1, './s1': 1, './r1': 1, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './r', './s', './c', './h', './w', './f', './b1', './r1', './s1', './c1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './c1': 1, './s1': 1, './r1': 1, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprp.py b/tests/filecheck/search/test_conv_pprprp.py index a56b9fb53..0c5234a21 100644 --- a/tests/filecheck/search/test_conv_pprprp.py +++ b/tests/filecheck/search/test_conv_pprprp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 1] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 16, 1, 1, 3] # CHECK-NEXT: stats {'filtered': 100, 'all': 202} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpv.py b/tests/filecheck/search/test_conv_pprprpv.py index 01f3322d9..0f3303958 100644 --- a/tests/filecheck/search/test_conv_pprprpv.py +++ b/tests/filecheck/search/test_conv_pprprpv.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 7, 1] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 16, 7, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 32, 1, 1, 1] # CHECK-NEXT: stats {'filtered_vec': 100, 'filtered': 1520, 'all': 4521} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpvr.py b/tests/filecheck/search/test_conv_pprprpvr.py index 337e138aa..3ecca074a 100644 --- a/tests/filecheck/search/test_conv_pprprpvr.py +++ b/tests/filecheck/search/test_conv_pprprpvr.py @@ -20,13 +20,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 1, './w2': 1, './w3': 1}, './f': {'./f1': 1, './f2': 1, './f3': 1}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 1, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 16, 1, 1, 3] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 1, './h2': 1, './h3': 1}, './w': {'./w1': 2, './w2': 2, './w3': 2}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 2, './h3': 1, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 1, 3] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 1, 7, 1] @@ -128,4 +128,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 32, 1, 1, 1] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 32, 1, 1, 3] # CHECK-NEXT: stats {'filtered_l2': 100, 'filtered_l1': 102, 'filtered_reg': 132, 'filtered_vec': 134, 'filtered': 1918, 'all': 3178} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 2}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 2, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1, './b2': 1, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 2}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 32, './f2': 32, './f3': 32}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 3}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 32, './w3': 1, './h3': 2, './b3': 1, './c1': 3, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_pprprpvr_rnd.py b/tests/filecheck/search/test_conv_pprprpvr_rnd.py index 2ce179546..11e3b57c4 100644 --- a/tests/filecheck/search/test_conv_pprprpvr_rnd.py +++ b/tests/filecheck/search/test_conv_pprprpvr_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 2, 16, 7, 1, 1] # CHECK-NEXT: sample 19: [1, 2, 1, 1, 2, 1, 2, 1, 1, 1, 1, 16, 1, 1, 1] # CHECK-NEXT: stats {'filtered_l2': 5, 'filtered_l1': 5, 'filtered_reg': 6, 'filtered_vec': 6, 'filtered': 100} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2, './b2': 2, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2, './b2': 2, './b3': 1}, './h': {'./h1': 2, './h2': 2, './h3': 1}, './w': {'./w1': 2, './w2': 1, './w3': 1}, './f': {'./f1': 16, './f2': 16, './f3': 16}, './r': {'./r1': 1}, './s': {'./s1': 1}, './c': {'./c1': 1}}, permutation={'.': ['./b', './h', './w', './f', './b1', './h1', './w1', './f1', './r', './s', './c', './b2', './h2', './w2', './f2', './r1', './s1', './c1', './b3', './h3', './w3', './f3']}, vectorization=['./f3'], parallelization=['./b', './h', './w', './f'], unrolling={'./f3': 16, './w3': 1, './h3': 1, './b3': 1, './c1': 1, './s1': 1, './r1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_prp.py b/tests/filecheck/search/test_conv_prp.py index afc1a4511..781149065 100644 --- a/tests/filecheck/search/test_conv_prp.py +++ b/tests/filecheck/search/test_conv_prp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 1}, './f': {'./f1': 1}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 1, './w1': 1, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 2, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 2, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 1}, './h': {'./h1': 1}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 1, './b1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 4] @@ -68,4 +68,4 @@ # CHECK-NEXT: sample 45: [2, 2, 2, 8] # CHECK-NEXT: sample 46: [2, 2, 2, 16] # CHECK-NEXT: stats {'filtered': 47, 'all': 48} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 2}, './f': {'./f1': 16}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=[], unrolling={'./f1': 16, './w1': 2, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_conv_prp_rnd.py b/tests/filecheck/search/test_conv_prp_rnd.py index 622601c0b..d03a0ab7e 100644 --- a/tests/filecheck/search/test_conv_prp_rnd.py +++ b/tests/filecheck/search/test_conv_prp_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [2, 2, 1, 8] # CHECK-NEXT: sample 19: [2, 2, 1, 4] # CHECK-NEXT: stats {'filtered': 20} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 1}, './f': {'./f1': 4}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=['./b', './h', './w', './f'], unrolling={'./f1': 4, './w1': 1, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['b', 'h', 'w', 'f'], loop_stamps=[], splits={}, tiles={'./b': {}, './h': {}, './w': {}, './f': {}}, permutation={'.': ['./b', './h', './w', './f']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['b', 'h', 'w', 'f', 'r', 's', 'c'], loop_stamps=[], splits={}, tiles={'./b': {'./b1': 2}, './h': {'./h1': 2}, './w': {'./w1': 1}, './f': {'./f1': 4}}, permutation={'.': ['./b', './h', './w', './f', './r', './s', './c', './b1', './h1', './w1', './f1']}, vectorization=['./f1'], parallelization=['./b', './h', './w', './f'], unrolling={'./f1': 4, './w1': 1, './h1': 2, './b1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_goto.py b/tests/filecheck/search/test_matmul_goto.py index 615edb798..ca283121c 100644 --- a/tests/filecheck/search/test_matmul_goto.py +++ b/tests/filecheck/search/test_matmul_goto.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 0, 0] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 0, 1] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 1, 0] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 2, 2, 3, 1, 0] # CHECK-NEXT: sample 99: [1, 1, 1, 2, 2, 3, 1, 1] # CHECK-NEXT: stats {'filtered': 100, 'all': 108} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 2, './j2': 2}, './k': {'./k1': 2}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 2, './j2': 2}, './k': {'./k1': 2}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_goto_r.py b/tests/filecheck/search/test_matmul_goto_r.py index b250298df..2410936f7 100644 --- a/tests/filecheck/search/test_matmul_goto_r.py +++ b/tests/filecheck/search/test_matmul_goto_r.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1, 1, 0, 0, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1}, './j': {'./j1': 1, './j2': 1}, './k': {'./k1': 1}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 0}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [3, 1, 2, 1, 1, 1, 0, 0] # CHECK-NEXT: sample 1: [3, 1, 2, 1, 1, 1, 0, 1] # CHECK-NEXT: sample 2: [3, 1, 2, 1, 1, 1, 1, 0] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [3, 1, 2, 2, 6, 2, 1, 0] # CHECK-NEXT: sample 99: [3, 1, 2, 2, 6, 2, 1, 1] # CHECK-NEXT: stats {'filtered': 3256, 'all': 6620} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1}, './j': {'./j1': 4, './j2': 2}, './k': {'./k1': 6}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1}, './j': {'./j1': 4, './j2': 2}, './k': {'./k1': 6}}, permutation={'.': ['./j', './k', './i', './j1', './i1', './k1', './i2', './j2']}, vectorization=['./j2'], parallelization=[], unrolling={'./i2': 1, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_oo.py b/tests/filecheck/search/test_matmul_oo.py index 6fb3d9a44..eadb24e4a 100644 --- a/tests/filecheck/search/test_matmul_oo.py +++ b/tests/filecheck/search/test_matmul_oo.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [3, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 12, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 12, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 3] @@ -66,4 +66,4 @@ # CHECK-NEXT: sample 43: [7, 8, 1] # CHECK-NEXT: sample 44: [7, 16, 1] # CHECK-NEXT: stats {'filtered': 45, 'all': 144} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './k', './j', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './k1': 1, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_p1.py b/tests/filecheck/search/test_matmul_p1.py index ffa955612..eee25c081 100644 --- a/tests/filecheck/search/test_matmul_p1.py +++ b/tests/filecheck/search/test_matmul_p1.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 0] # CHECK-NEXT: sample 1: [1, 1, 1, 1] # CHECK-NEXT: sample 2: [1, 1, 1, 2] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 32, 1, 1] # CHECK-NEXT: sample 99: [1, 32, 1, 4] # CHECK-NEXT: stats {'filtered': 100, 'all': 185} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 32}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 32, './i1': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 32}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 32, './i1': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_p1v.py b/tests/filecheck/search/test_matmul_p1v.py index 92b21f24a..b639c5f20 100644 --- a/tests/filecheck/search/test_matmul_p1v.py +++ b/tests/filecheck/search/test_matmul_p1v.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 0] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1', './k1']}, vectorization=[], parallelization=[], unrolling={'./k1': 1, './j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './k1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './k1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 16, 1, 1] # CHECK-NEXT: sample 1: [1, 16, 1, 4] # CHECK-NEXT: sample 2: [1, 16, 2, 1] @@ -47,4 +47,4 @@ # CHECK-NEXT: sample 24: [7, 16, 1, 1] # CHECK-NEXT: sample 25: [7, 16, 1, 4] # CHECK-NEXT: stats {'filtered': 154, 'all': 864} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './k', './k1', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprp.py b/tests/filecheck/search/test_matmul_pprprp.py index 5a31e9712..c19c1be33 100644 --- a/tests/filecheck/search/test_matmul_pprprp.py +++ b/tests/filecheck/search/test_matmul_pprprp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 1, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 1, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 1, 3] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [1, 1, 1, 1, 16, 2, 6] # CHECK-NEXT: sample 99: [1, 1, 1, 1, 32, 1, 1] # CHECK-NEXT: stats {'filtered': 100, 'all': 121} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpv.py b/tests/filecheck/search/test_matmul_pprprpv.py index 2eb1bf600..2928c1141 100644 --- a/tests/filecheck/search/test_matmul_pprprpv.py +++ b/tests/filecheck/search/test_matmul_pprprpv.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 16, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 16, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 16, 3] @@ -121,4 +121,4 @@ # CHECK-NEXT: sample 98: [3, 1, 1, 1, 1, 16, 3] # CHECK-NEXT: sample 99: [3, 1, 1, 1, 1, 16, 4] # CHECK-NEXT: stats {'filtered_vec': 100, 'filtered': 1472, 'all': 3052} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 4}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 4}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 4}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 4}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpvr.py b/tests/filecheck/search/test_matmul_pprprpvr.py index fd0c84aaa..d035261f5 100644 --- a/tests/filecheck/search/test_matmul_pprprpvr.py +++ b/tests/filecheck/search/test_matmul_pprprpvr.py @@ -20,13 +20,13 @@ utils.print_exhaustive_samples(backend, strategy,100) # CHECK: schedule O0: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1, 1, 1, 1, 1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 1, './j2': 1, './j3': 1}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 1, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 1, 1, 1, 1, 16, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1, './i2': 1, './i3': 1}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [1, 1, 3, 1, 1, 16, 12] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 12}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 12}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1, 1, 1, 1, 16, 1] # CHECK-NEXT: sample 1: [1, 1, 1, 1, 1, 16, 2] # CHECK-NEXT: sample 2: [1, 1, 1, 1, 1, 16, 3] @@ -128,4 +128,4 @@ # CHECK-NEXT: sample 98: [1, 21, 1, 1, 1, 32, 3] # CHECK-NEXT: sample 99: [1, 21, 1, 1, 2, 16, 1] # CHECK-NEXT: stats {'filtered_l2': 100, 'filtered_l1': 105, 'filtered_reg': 115, 'filtered_vec': 154, 'filtered': 2126, 'all': 2749} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 21, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 21, './i3': 1}, './j': {'./j1': 32, './j2': 32, './j3': 16}, './k': {'./k1': 1}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 1, './k1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_pprprpvr_rnd.py b/tests/filecheck/search/test_matmul_pprprpvr_rnd.py index 6fa9effc2..72543a25d 100644 --- a/tests/filecheck/search/test_matmul_pprprpvr_rnd.py +++ b/tests/filecheck/search/test_matmul_pprprpvr_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [1, 1, 1, 1, 2, 16, 3] # CHECK-NEXT: sample 19: [7, 1, 3, 1, 1, 16, 2] # CHECK-NEXT: stats {'filtered_l2': 2, 'filtered_l1': 2, 'filtered_reg': 3, 'filtered_vec': 3, 'filtered': 70} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 2}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21, './i2': 3, './i3': 3}, './j': {'./j1': 16, './j2': 16, './j3': 16}, './k': {'./k1': 2}}, permutation={'.': ['./i', './j', './i1', './j1', './k', './i2', './j2', './k1', './i3', './j3']}, vectorization=['./j3'], parallelization=['./i', './j'], unrolling={'./j3': 16, './i3': 3, './k1': 2}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_prp.py b/tests/filecheck/search/test_matmul_prp.py index 757860aaa..c7b19775b 100644 --- a/tests/filecheck/search/test_matmul_prp.py +++ b/tests/filecheck/search/test_matmul_prp.py @@ -13,13 +13,13 @@ utils.print_exhaustive_samples(backend, strategy, 100) # CHECK: schedule O0: [1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O1: [1, 1] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 1}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 1, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O2: [1, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 1}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 1}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: schedule O3: [3, 16] -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 3}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 3}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] # CHECK-NEXT: sample 0: [1, 1] # CHECK-NEXT: sample 1: [1, 2] # CHECK-NEXT: sample 2: [1, 4] @@ -38,4 +38,4 @@ # CHECK-NEXT: sample 15: [7, 8] # CHECK-NEXT: sample 16: [7, 16] # CHECK-NEXT: stats {'filtered': 17, 'all': 24} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 7}, './j': {'./j1': 16}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=[], unrolling={'./j1': 16, './i1': 7}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})] diff --git a/tests/filecheck/search/test_matmul_prp_rnd.py b/tests/filecheck/search/test_matmul_prp_rnd.py index 47a46ed14..d534a1703 100644 --- a/tests/filecheck/search/test_matmul_prp_rnd.py +++ b/tests/filecheck/search/test_matmul_prp_rnd.py @@ -39,4 +39,4 @@ # CHECK-NEXT: sample 18: [3, 2] # CHECK-NEXT: sample 19: [21, 8] # CHECK-NEXT: stats {'filtered': 19} -# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21}, './j': {'./j1': 8}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=['./i', './j'], unrolling={'./j1': 8, './i1': 21}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], vectorization_sizes={})] +# CHECK-NEXT: [MlirNodeSchedule(node_name='%2_0', node_ident='__xtc_id_%2_0_', dims=['i', 'j'], loop_stamps=[], splits={}, tiles={'./i': {}, './j': {}}, permutation={'.': ['./i', './j']}, vectorization=[], parallelization=[], unrolling={}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={}), MlirNodeSchedule(node_name='%2', node_ident='__xtc_id_%2_', dims=['i', 'j', 'k'], loop_stamps=[], splits={}, tiles={'./i': {'./i1': 21}, './j': {'./j1': 8}}, permutation={'.': ['./i', './j', './k', './i1', './j1']}, vectorization=['./j1'], parallelization=['./i', './j'], unrolling={'./j1': 8, './i1': 21}, packed_buffers={}, write_buffers={}, memory_mesh={}, processor_mesh={}, distribution={}, distributed_buffers={}, fused=[], fused_consumers=[], gpu_blocks=[], gpu_threads=[], gpu_lanes=[], gpu_warps=[], vectorization_sizes={})]