# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=missing-docstring, invalid-name # ruff: noqa: E722 """A GEMM schedule rule for GPU operators.""" from dataclasses import dataclass from enum import Enum from tvm import s_tir, tirx from tvm.ir import Range from tvm.s_tir.schedule.schedule import SBlockRV from tvm.script import tirx as T from tvm.target import Target from tvm.tirx import Expr, IterVar, Var from tvm.tirx.analysis import undefined_vars from ..analysis import IterInfo, SBlockInfo, get_root_block from .base import GPUScheduleRule def _collect_producers(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV): result = [] for producer in sch.get_producers(block): result.append(producer) result.extend(_collect_producers(sch, producer)) return result def _collect_consumers(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV): result = [] for consumer in sch.get_consumers(block): result.append(consumer) result.extend(_collect_consumers(sch, consumer)) return result def auto_inline_producers( sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV, ): while True: inlined_cnt = 0 producers = _collect_producers(sch, block) for producer in producers: try: sch.compute_inline(producer) inlined_cnt += 1 except: # pylint: disable=bare-except continue if inlined_cnt == 0: return def auto_inline_consumers( sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV, ): while True: inlined_cnt = 0 consumers = _collect_consumers(sch, block) for consumer in consumers: try: sch.compute_inline(consumer) inlined_cnt += 1 except: # pylint: disable=bare-except continue for consumer in consumers: try: sch.reverse_compute_inline(consumer) inlined_cnt += 1 except: # pylint: disable=bare-except continue if inlined_cnt == 0: return def auto_inline_consumer_chain( sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV, ): auto_inline_consumers(sch, block) remaining_consumers = sch.get_consumers(block) if len(remaining_consumers) != 0: # Some blocks have failed to be inlined to the producer cache-write stage. # This could be due to another producer block that has not been scheduled. for c in remaining_consumers: for p in sch.get_producers(c): if sch.get(p) != sch.get(block): auto_inline_producers(sch, p) sch.compute_inline(p) # Try inlining into the cache-write stage again, this time it should succeed. auto_inline_consumers(sch, block) msg = "There are some consumers of the cache-write stage that are not properly inlined." assert len(sch.get_consumers(block)) == 0, msg class IterKind(Enum): """Iter kinds for GEMM-liked programs. We can simplify the computation to C[S, I, J] += A[S, I, K] * B[S, J, K], where `I, J, K` are fundamental axes for gemm and `S` represents all other spatial axes (e.g. batches) kIter_S: spatial axes kIter_I: I axes kIter_J: J axes kIter_K: K axes kIter_T: trivial axes (i.e. with extent 1) """ kIter_S = 0 kIter_I = 1 kIter_J = 2 kIter_K = 3 kIter_T = 4 @dataclass class IterTrait: kind: IterKind extent: Expr def _is_one(x: Expr) -> bool: return isinstance(x, tirx.IntImm) and x.value == 1 def make_iter_fusion_index_map( traits: list[IterTrait], kind_order: list[IterKind], ) -> tirx.IndexMap: fused_iters: dict[IterKind, Expr] = {} input_iters: list[tirx.Var] = [] for i, trait in enumerate(traits): v_i = tirx.Var(f"i{i}", trait.extent.ty) input_iters.append(v_i) if trait.kind == IterKind.kIter_T: continue if trait.kind not in kind_order: raise ValueError(f"Unknown iter kind {trait.kind}") if trait.kind in fused_iters: fused_iters[trait.kind] = fused_iters[trait.kind] * trait.extent + v_i else: fused_iters[trait.kind] = v_i final_indices: list[tirx.Expr] = [ fused_iters.get(kind, tirx.IntImm(traits[0].extent.ty, 0)) for kind in kind_order ] return tirx.IndexMap(input_iters, final_indices, None) def detect_iter_traits(block: tirx.SBlock) -> tuple[list[IterTrait]] | None: """Detect iter traits based on the pattern C[S, I, J] += A[S, I, K] * B[S, J, K] Parameters ---------- block : tirx.SBlock The block to be analyzed Returns ------- traits : Optional[Tuple[List[IterTrait]]] The detected iter traits for axes in A, B and C. None if the block does not match the pattern. """ if len(block.reads) != 2 or len(block.writes) != 1: return None def get_access_axes(region: list[Range]) -> set[Var]: axes: set[Var] = set() for r in region: if not _is_one(r.extent): raise ValueError("Expect elemwise block access") axes = axes.union(set(undefined_vars(r.min))) return axes try: A_axes = get_access_axes(block.reads[0].region) B_axes = get_access_axes(block.reads[1].region) C_axes = get_access_axes(block.writes[0].region) except ValueError: return None traits: dict[Var, IterTrait] = {} for iter_var in block.iter_vars: var = iter_var.var kind: IterKind if _is_one(iter_var.dom.extent): kind = IterKind.kIter_T elif iter_var.iter_type == iter_var.DataPar: if var in A_axes and var in B_axes and var in C_axes: kind = IterKind.kIter_S elif var in A_axes and var in C_axes: kind = IterKind.kIter_I elif var in B_axes and var in C_axes: kind = IterKind.kIter_J else: return None elif iter_var.iter_type == tirx.IterVar.CommReduce: if var in A_axes and var in B_axes and var not in C_axes: kind = IterKind.kIter_K else: return None else: return None traits[var] = IterTrait(kind, iter_var.dom.extent) # A Gemm-kernel requires have I, J and K axes gemm_traits = {IterKind.kIter_I, IterKind.kIter_J, IterKind.kIter_K} if {x.kind for x in traits.values()}.intersection(gemm_traits) != gemm_traits: return None A_traits = [traits[iter_var.var] for iter_var in block.iter_vars if iter_var.var in A_axes] B_traits = [traits[iter_var.var] for iter_var in block.iter_vars if iter_var.var in B_axes] C_traits = [traits[iter_var.var] for iter_var in block.iter_vars if iter_var.var in C_axes] block_traits = [traits[i.var] for i in block.iter_vars] return A_traits, B_traits, C_traits, block_traits def get_index_map(block: tirx.SBlock) -> tuple[tirx.IndexMap, ...] | None: """Get index maps for the block Parameters ---------- block : tirx.SBlock The block to be analyzed Returns ------- index_maps : Optional[Tuple[tirx.IndexMap]] The index maps for the block, or None if the block is not a gemm-liked kernel """ traits = detect_iter_traits(block) if traits is None: return None A_traits, B_traits, C_traits, block_traits = traits A_index_map = make_iter_fusion_index_map( A_traits, [IterKind.kIter_S, IterKind.kIter_I, IterKind.kIter_K] ) B_index_map = make_iter_fusion_index_map( B_traits, [IterKind.kIter_S, IterKind.kIter_J, IterKind.kIter_K] ) C_index_map = make_iter_fusion_index_map( C_traits, [IterKind.kIter_S, IterKind.kIter_I, IterKind.kIter_J] ) matmul_index_map = make_iter_fusion_index_map( block_traits, [IterKind.kIter_S, IterKind.kIter_I, IterKind.kIter_J, IterKind.kIter_K] ) return ( matmul_index_map, A_index_map, B_index_map, C_index_map, ) def get_sblock_info(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV) -> SBlockInfo: def _iter_kind(loop: tirx.IterVar) -> str: return {tirx.IterVar.DataPar: "S", tirx.IterVar.CommReduce: "R"}.get(loop.iter_type, "O") def _is_reduction_block(block: s_tir.schedule.SBlockRV): for iter_var in sch.get(block).iter_vars: if _iter_kind(iter_var) == "R": return True return False return SBlockInfo( name=sch.get(block).name_hint, iters=[ IterInfo( kind=_iter_kind(iter_var), var=iter_var.var, dom=iter_var.dom.extent, loop_rv=loop_rv, ) for loop_rv, iter_var in zip(sch.get_loops(block), sch.get(block).iter_vars) ], block_rv=block, reduction_block=_is_reduction_block(block), ) def get_reduction_blocks(sch, blocks) -> bool: # Get the main computation block def is_reduction(block: SBlockRV) -> bool: block_stmt = sch.get(block) iter_types = {iter_var.iter_type for iter_var in block_stmt.iter_vars} return iter_types == {IterVar.CommReduce, IterVar.DataPar} def is_spatial(block: SBlockRV) -> bool: block_stmt = sch.get(block) iter_types = {iter_var.iter_type for iter_var in block_stmt.iter_vars} return iter_types == {IterVar.DataPar} # NOTE: We assume there is only one reduction block in the function # all blocks are required to be spatial or reduction if not all([is_reduction(block) or is_spatial(block) for block in blocks]): return None # There is only one reduction block reduction_blocks = [block for block in blocks if is_reduction(block)] if len(reduction_blocks) != 1: return None return reduction_blocks def get_in_out_dtypes(block: tirx.SBlock) -> tuple[str]: """ Detect In/Out data types for the given block based on the analysis if read/write buffers. """ assert len(block.reads) > 0 and len(block.writes) > 0 in_dtype = block.reads[0].buffer.dtype out_dtype = block.writes[0].buffer.dtype return (in_dtype, out_dtype) def check_sm_version(arch: str) -> int: sm_version = arch.replace("sm_", "") return int(sm_version) if sm_version.isdigit() else -1 class MetalMatmul(GPUScheduleRule): """ The schedule rule for Metal matmul computation. """ def apply( # pylint: disable=too-many-locals,missing-docstring self, func: tirx.PrimFunc, target: Target, _: bool, ) -> s_tir.Schedule | None: from tvm.s_tir.tensor_intrin.metal import ( # pylint: disable=import-outside-toplevel get_simdgroup_intrin_group, ) if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target): return None sch = s_tir.Schedule(func) root_block = get_root_block(sch) blocks = sch.get_child_blocks(root_block) reduction_blocks = get_reduction_blocks(sch, blocks) if reduction_blocks is None: return None # Step 0. Configs block_size_x: int = 16 block_size_y: int = 16 block_size_k: int = 32 micro_size: int = 8 warp_size: int = 32 ty_len: int = 1 tz_len: int = 4 vector_size: int = 4 # Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K] # Reindex first and than analyze the index map main_block = reduction_blocks[0] reindex_a = sch.reindex(main_block, ("read", 0)) reindex_b = sch.reindex(main_block, ("read", 1)) reindex_c = sch.reindex(main_block, ("write", 0)) index_maps = get_index_map(sch.get(main_block)) assert index_maps is not None matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps sch.transform_layout(reindex_a, ("write", 0), a_index_map) sch.transform_layout(reindex_b, ("write", 0), b_index_map) sch.transform_layout(reindex_c, ("read", 0), c_index_map) sch.transform_block_layout(main_block, matmul_index_map) # Step 2. Padding for dynamic shape kernels sch.pad_einsum( main_block, [ 1, ty_len * block_size_x, tz_len * block_size_y, block_size_k, ], ) # Step 3. Schedule matmul to use simdgroup intrinsics batch, i, j, k = sch.get_loops(main_block) bx, ty, i0, i1 = sch.split(i, [None, ty_len, block_size_x // micro_size, micro_size]) by, tz, j0, j1 = sch.split(j, [None, tz_len, block_size_y // micro_size, micro_size]) k0, k1, k2 = sch.split(k, [None, block_size_k // micro_size, micro_size]) sch.reorder(bx, by, ty, tz, k0, k1, i0, j0, i1, j1, k2) sch.bind(bx, "blockIdx.x") sch.bind(by, "blockIdx.y") sch.bind(batch, "blockIdx.z") sch.bind(ty, "threadIdx.y") sch.bind(tz, "threadIdx.z") def fetch_to_shared(block, idx): block_read = sch.cache_read(block, idx, "shared") sch.compute_at(block_read, k0, preserve_unit_loops=True) fused = sch.fuse(*sch.get_loops(block_read)[-2:]) _, _tz, _ty, _tx, vec = sch.split(fused, [None, tz_len, ty_len, warp_size, vector_size]) sch.bind(_tz, "threadIdx.z") sch.bind(_ty, "threadIdx.y") sch.bind(_tx, "threadIdx.x") sch.vectorize(vec) return block_read a_g2s = fetch_to_shared(main_block, 0) b_g2s = fetch_to_shared(main_block, 1) auto_inline_producers(sch, a_g2s) auto_inline_producers(sch, b_g2s) # create read cache to load matrix from shared memory to wmma fragments A_simdgroup = sch.cache_read(main_block, 0, "metal.simdgroup") B_simdgroup = sch.cache_read(main_block, 1, "metal.simdgroup") sch.compute_at(A_simdgroup, k1) sch.compute_at(B_simdgroup, k1) C_simd2s = sch.cache_write(main_block, 0, "metal.simdgroup") C_s2g = sch.cache_write(C_simd2s, 0, "shared") sch.reverse_compute_at(C_simd2s, tz, preserve_unit_loops=True) sch.reverse_compute_at(C_s2g, by, preserve_unit_loops=True) intrin_group = get_simdgroup_intrin_group( load_scope="shared", store_scope="shared", dtype="float16", trans_a=False, trans_b=True, ) sch.transform_layout(B_simdgroup, ("write", 0), lambda s, i, j: (s, j, i)) def tensorize_block(block: s_tir.schedule.SBlockRV, intrin: str): *_, i, j = sch.get_loops(block) io, ii = sch.split(i, [None, micro_size]) jo, ji = sch.split(j, [None, micro_size]) sch.reorder(io, jo, ii, ji) sch.tensorize(ii, intrin) C_init = sch.decompose_reduction(main_block, k0) tensorize_block(A_simdgroup, intrin_group["load_a"]) tensorize_block(B_simdgroup, intrin_group["load_b"]) tensorize_block(C_simd2s, intrin_group["store"]) tensorize_block(C_init, intrin_group["init"]) *_, i, j, k = sch.get_loops(main_block) sch.tensorize(i, intrin_group["compute"]) auto_inline_consumer_chain(sch, C_s2g) fused = sch.fuse(*sch.get_loops(C_s2g)[-2:]) _, _tz, _ty, _tx, vec = sch.split(fused, [None, tz_len, ty_len, warp_size, vector_size]) sch.bind(_tz, "threadIdx.z") sch.bind(_ty, "threadIdx.y") sch.bind(_tx, "threadIdx.x") sch.vectorize(vec) return sch class MatmulTensorization(GPUScheduleRule): """ The schedule rule for float16 tensor core matmul computation. func with attr 'dlight.do_not_tensorize' will not be tensorized. """ def apply( # pylint: disable=too-many-locals,missing-docstring self, func: tirx.PrimFunc, target: Target, _: bool, ) -> s_tir.Schedule | None: from tvm.s_tir.tensor_intrin.cuda import ( # pylint: disable=import-outside-toplevel get_wmma_intrin_group, ) if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target): return None sch = s_tir.Schedule(func) root_block = get_root_block(sch) blocks = sch.get_child_blocks(root_block) if "dlight.do_not_tensorize" in func.attrs.keys(): return None reduction_blocks = get_reduction_blocks(sch, blocks) if reduction_blocks is None: return None # Start Schedule # Step 0. Get schedule config. # NOTE: we can analyze the config by the hardware spec in the future # tensor core intrinsic size micro_size_x = 16 micro_size_y = 16 micro_size_k = 16 warp_size = 32 vector_size = 4 i_factors, j_factors, k_factors = ( [None, 1, 4, 2], [1, None, 4, 2], [None, 4], ) num_ty = i_factors[2] * j_factors[2] x_pad_factor = i_factors[2] * i_factors[3] y_pad_factor = j_factors[2] * j_factors[3] k_pad_factor = k_factors[1] # Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K] # Reindex first and than analyze the index map main_block = reduction_blocks[0] reindex_a = sch.reindex(main_block, ("read", 0)) reindex_b = sch.reindex(main_block, ("read", 1)) reindex_c = sch.reindex(main_block, ("write", 0)) index_maps = get_index_map(sch.get(main_block)) assert index_maps is not None matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps sch.transform_layout(reindex_a, ("write", 0), a_index_map) sch.transform_layout(reindex_b, ("write", 0), b_index_map) sch.transform_layout(reindex_c, ("read", 0), c_index_map) sch.transform_block_layout(main_block, matmul_index_map) # Step 2. Padding for dynamic shape kernels sch.pad_einsum( main_block, [ 1, micro_size_x * x_pad_factor, micro_size_y * y_pad_factor, micro_size_k * k_pad_factor, ], ) # Step 3. Schedule matmul to use tensor core block = main_block batch, i, j, k = sch.get_loops(block) # inner loops for tensor core computation i, i_inner = sch.split(i, factors=[None, micro_size_x]) j, j_inner = sch.split(j, factors=[None, micro_size_y]) k, k_inner = sch.split(k, factors=[None, micro_size_k]) sch.reorder(i, j, k, i_inner, j_inner, k_inner) block_inner = block block_outer = sch.blockize(i_inner) i0, i1, i2, i3 = sch.split(i, factors=i_factors) j0, j1, j2, j3 = sch.split(j, factors=j_factors) k0, k1 = sch.split(k, k_factors) sch.annotate(k0, "software_pipeline_order", [0, 3, 1, 4, 5, 2, 6]) sch.annotate(k0, "software_pipeline_stage", [0, 0, 0, 0, 0, 1, 1]) sch.annotate(k1, "software_pipeline_order", [0, 1, 2]) sch.annotate(k1, "software_pipeline_stage", [0, 0, 1]) sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3) block_idx = sch.fuse(i0, j0) block_idy = sch.fuse(i1, j1) thread_idy = sch.fuse(j2, i2) sch.bind(batch, "blockIdx.z") sch.bind(block_idx, "blockIdx.x") sch.bind(block_idy, "blockIdx.y") sch.bind(thread_idy, "threadIdx.y") def fetch_to_shared(block, idx, ndim): block_read = sch.cache_read(block, idx, "shared.dyn") sch.compute_at(block_read, k0) fused = sch.fuse(*sch.get_loops(block_read)[-ndim:]) _, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size]) sch.bind(f_2, "threadIdx.x") sch.bind(f_1, "threadIdx.y") sch.vectorize(f_3) sch.storage_align(block_read, 0, axis=-2, factor=16, offset=8) sch.annotate(block_read, "tirx.manifest_shared_memory_local_stage", 1) sch.annotate(block_read, "double_buffer_scope", 0) return block_read a_g2s = fetch_to_shared(block_outer, 0, 2) b_g2s = fetch_to_shared(block_outer, 1, 2) auto_inline_producers(sch, a_g2s) auto_inline_producers(sch, b_g2s) # create read cache to load matrix from shared memory to wmma fragments A_mat = sch.cache_read(block_outer, 0, "wmma.matrix_a") B_mat = sch.cache_read(block_outer, 1, "wmma.matrix_b") sch.compute_at(A_mat, k1) sch.compute_at(B_mat, k1) # create write cache to store matrix from wmma fragments to shared memory and global memory accumulator_shared_to_global = sch.cache_write(block_outer, 0, "shared.dyn") sch.storage_align(accumulator_shared_to_global, 0, -2, 16, 4) store = sch.cache_write(block_outer, 0, "wmma.accumulator") sch.reverse_compute_at(store, thread_idy) sch.reverse_compute_at(accumulator_shared_to_global, thread_idy) # split the store loop to match hardware intrinsic pattern i, j = sch.get_loops(store)[-2:] i0, i1 = sch.split(i, factors=[None, 16]) j0, j1 = sch.split(j, factors=[None, 16]) sch.reorder(i0, j0, i1, j1) block_init_c = sch.decompose_reduction(block_outer, k0) block_init_c_inner = sch.get_child_blocks(block_init_c)[0] # Tensorization by hardware intrinsics intrin_group = get_wmma_intrin_group( load_scope="shared.dyn", store_scope="shared.dyn", in_dtype="float16", out_dtype="float32", trans_b=True, ) try: i, j = sch.get_loops(A_mat)[-2:] i0, i1 = sch.split(i, factors=[None, 16]) j0, j1 = sch.split(j, factors=[None, 16]) sch.reorder(i0, j0, i1, j1) sch.unroll(i0) sch.unroll(j0) sch.tensorize(i1, intrin_group["load_a"]) i, j = sch.get_loops(B_mat)[-2:] i0, i1 = sch.split(i, factors=[None, 16]) j0, j1 = sch.split(j, factors=[None, 16]) sch.reorder(i0, j0, i1, j1) sch.unroll(i0) sch.unroll(j0) sch.tensorize(i1, intrin_group["load_b"]) except: # pylint: disable=bare-except return None # Try to tensorize the init, store and compute block with f16 or f32 intrinsics tensorize_success: bool = False def tensorize_init_store_compute(): sch.tensorize(sch.get_loops(block_init_c_inner)[-2], intrin_group["init"]) sch.tensorize(sch.get_loops(store)[-2], intrin_group["store"]) sch.tensorize(sch.get_loops(block_inner)[-3], intrin_group["compute"]) try: tensorize_init_store_compute() tensorize_success = True except: # pylint: disable=bare-except intrin_group = get_wmma_intrin_group( load_scope="shared.dyn", store_scope="shared.dyn", in_dtype="float16", out_dtype="float16", trans_b=True, ) if not tensorize_success: try: tensorize_init_store_compute() tensorize_success = True except: # pylint: disable=bare-except return None auto_inline_consumer_chain(sch, accumulator_shared_to_global) fused = sch.fuse(*sch.get_loops(accumulator_shared_to_global)[-2:]) _, f1, f2 = sch.split(fused, factors=[None, warp_size, vector_size]) sch.bind(f1, "threadIdx.x") sch.vectorize(f2) return sch if tensorize_success else None class MatmulInt8Tensorization(GPUScheduleRule): """ The schedule rule for int8 tensor core matmul computation. func with attr 'dlight.do_not_tensorize' will not be tensorized. """ def apply( # pylint: disable=too-many-locals,missing-docstring self, func: tirx.PrimFunc, target: Target, _: bool, ) -> s_tir.Schedule | None: from tvm.s_tir.tensor_intrin.cuda import ( # pylint: disable=import-outside-toplevel get_wmma_intrin_group, ) if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target): return None sch = s_tir.Schedule(func) root_block = get_root_block(sch) blocks = sch.get_child_blocks(root_block) if "dlight.do_not_tensorize" in func.attrs.keys(): return None reduction_blocks = get_reduction_blocks(sch, blocks) if reduction_blocks is None: return None # Start Schedule # Step 0. Get schedule config. # NOTE: we can analyze the config by the hardware spec in the future # tensor core intrinsic size micro_size_x = 16 micro_size_y = 16 micro_size_k = 16 warp_size = 32 vector_size = 4 i_factors, j_factors, k_factors = ( [None, 1, 4, 2], [1, None, 4, 2], [None, 1], ) num_ty = i_factors[2] * j_factors[2] x_pad_factor = i_factors[2] * i_factors[3] y_pad_factor = j_factors[2] * j_factors[3] k_pad_factor = k_factors[1] # Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K] # Reindex first and than analyze the index map main_block = reduction_blocks[0] reindex_a = sch.reindex(main_block, ("read", 0)) reindex_b = sch.reindex(main_block, ("read", 1)) reindex_c = sch.reindex(main_block, ("write", 0)) index_maps = get_index_map(sch.get(main_block)) assert index_maps is not None matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps sch.transform_layout(reindex_a, ("write", 0), a_index_map) sch.transform_layout(reindex_b, ("write", 0), b_index_map) sch.transform_layout(reindex_c, ("read", 0), c_index_map) sch.transform_block_layout(main_block, matmul_index_map) # Step 2. Padding for dynamic shape kernels sch.pad_einsum( main_block, [ 1, micro_size_x * x_pad_factor, micro_size_y * y_pad_factor, micro_size_k * k_pad_factor, ], ) # Step 3. Schedule matmul to use tensor core block = main_block batch, i, j, k = sch.get_loops(block) # inner loops for tensor core computation i, i_inner = sch.split(i, factors=[None, micro_size_x]) j, j_inner = sch.split(j, factors=[None, micro_size_y]) k, k_inner = sch.split(k, factors=[None, micro_size_k]) sch.reorder(i, j, k, i_inner, j_inner, k_inner) block_inner = block block_outer = sch.blockize(i_inner) i0, i1, i2, i3 = sch.split(i, factors=i_factors) j0, j1, j2, j3 = sch.split(j, factors=j_factors) k0, k1 = sch.split(k, k_factors) sch.annotate(k0, "software_pipeline_order", [0, 3, 1, 4, 5, 2, 6]) sch.annotate(k0, "software_pipeline_stage", [0, 0, 0, 0, 0, 1, 1]) sch.annotate(k1, "software_pipeline_order", [0, 1, 2]) sch.annotate(k1, "software_pipeline_stage", [0, 0, 1]) sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3) block_idx = sch.fuse(i0, j0) block_idy = sch.fuse(i1, j1) thread_idy = sch.fuse(j2, i2) sch.bind(batch, "blockIdx.z") sch.bind(block_idx, "blockIdx.x") sch.bind(block_idy, "blockIdx.y") sch.bind(thread_idy, "threadIdx.y") def fetch_to_shared(block, idx, ndim): block_read = sch.cache_read(block, idx, "shared.dyn") sch.compute_at(block_read, k0) fused = sch.fuse(*sch.get_loops(block_read)[-ndim:]) _, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size]) sch.bind(f_2, "threadIdx.x") sch.bind(f_1, "threadIdx.y") sch.vectorize(f_3) sch.storage_align(block_read, 0, axis=-2, factor=32, offset=16) sch.annotate(block_read, "tirx.manifest_shared_memory_local_stage", 1) sch.annotate(block_read, "double_buffer_scope", 0) return block_read a_g2s = fetch_to_shared(block_outer, 0, 2) b_g2s = fetch_to_shared(block_outer, 1, 2) auto_inline_producers(sch, a_g2s) auto_inline_producers(sch, b_g2s) # create read cache to load matrix from shared memory to wmma fragments A_mat = sch.cache_read(block_outer, 0, "wmma.matrix_a") B_mat = sch.cache_read(block_outer, 1, "wmma.matrix_b") sch.compute_at(A_mat, k1) sch.compute_at(B_mat, k1) # create write cache to store matrix from wmma fragments to shared memory and global memory accumulator_shared_to_global = sch.cache_write(block_outer, 0, "shared.dyn") sch.storage_align(accumulator_shared_to_global, 0, -2, 16, 4) store = sch.cache_write(block_outer, 0, "wmma.accumulator") sch.reverse_compute_at(store, thread_idy) sch.reverse_compute_at(accumulator_shared_to_global, thread_idy) # split the store loop to match hardware intrinsic pattern i, j = sch.get_loops(store)[-2:] i0, i1 = sch.split(i, factors=[None, 16]) j0, j1 = sch.split(j, factors=[None, 16]) sch.reorder(i0, j0, i1, j1) block_init_c = sch.decompose_reduction(block_outer, k0) block_init_c_inner = sch.get_child_blocks(block_init_c)[0] # Tensorization by hardware intrinsics intrin_group = get_wmma_intrin_group( load_scope="shared.dyn", store_scope="shared.dyn", in_dtype="int8", out_dtype="int32", trans_b=True, ) try: i, j = sch.get_loops(A_mat)[-2:] i0, i1 = sch.split(i, factors=[None, 16]) j0, j1 = sch.split(j, factors=[None, 16]) sch.reorder(i0, j0, i1, j1) sch.unroll(i0) sch.unroll(j0) sch.tensorize(i1, intrin_group["load_a"]) i, j = sch.get_loops(B_mat)[-2:] i0, i1 = sch.split(i, factors=[None, 16]) j0, j1 = sch.split(j, factors=[None, 16]) sch.reorder(i0, j0, i1, j1) sch.unroll(i0) sch.unroll(j0) sch.tensorize(i1, intrin_group["load_b"]) except: # pylint: disable=bare-except return None def tensorize_init_store_compute(): sch.tensorize(sch.get_loops(block_init_c_inner)[-2], intrin_group["init"]) sch.tensorize(sch.get_loops(store)[-2], intrin_group["store"]) sch.tensorize(sch.get_loops(block_inner)[-3], intrin_group["compute"]) try: tensorize_init_store_compute() except: # pylint: disable=bare-except return None auto_inline_consumer_chain(sch, accumulator_shared_to_global) fused = sch.fuse(*sch.get_loops(accumulator_shared_to_global)[-2:]) _, f1, f2 = sch.split(fused, factors=[None, warp_size, vector_size]) sch.bind(f1, "threadIdx.x") sch.vectorize(f2) return sch class Matmul(GPUScheduleRule): """The schedule rule for matmul-like computation""" @dataclass class Config: block_size_x: int = 8 block_size_y: int = 8 vthread_x: int = 1 vthread_y: int = 1 micro_size_x: int = 4 micro_size_y: int = 4 micro_size_k: int = 8 vector_size: int = 1 unroll: int = 256 # 0 means no unroll use_shared: bool = True storage_align: bool = False inner_x: bool = False def get_configs(self, target: Target) -> Config: """Get the schedule config for the target""" if target.kind.name == "cuda" or target.kind.name == "rocm": return Matmul.Config( block_size_x=8, block_size_y=16, vthread_x=1, vthread_y=1, micro_size_x=4, micro_size_y=4, micro_size_k=16, vector_size=2, unroll=256, use_shared=True, storage_align=True, inner_x=False, ) elif target.kind.name == "opencl" and ( ("android" in str(target.host)) or ("adreno" in str(target.attrs)) ): return Matmul.Config( block_size_x=32, block_size_y=4, vthread_x=1, vthread_y=1, micro_size_x=8, micro_size_y=2, micro_size_k=16, vector_size=8, unroll=16, use_shared=False, storage_align=False, inner_x=True, ) else: return Matmul.Config() def apply( # pylint: disable=too-many-locals,missing-docstring self, func: tirx.PrimFunc, target: Target, _: bool, ) -> s_tir.Schedule | None: if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target): return None sch = s_tir.Schedule(func) config = self.get_configs(target) root_block = get_root_block(sch) blocks = sch.get_child_blocks(root_block) reduction_blocks = get_reduction_blocks(sch, blocks) if reduction_blocks is None: return None main_block = reduction_blocks[0] block_stmt = sch.get(main_block) main_block_info = get_sblock_info(sch, main_block) iter_infos = main_block_info.iters if not get_index_map(block_stmt): return None # Checks if it's a inner reduction by getting the last matrix's inner Index def is_inner_reduction(block_stmt, iter_infos): end_it = block_stmt.reads[-1].region[-1].min return {it.var: it.kind for it in iter_infos}.get(end_it, "O") == "R" if ( target.kind.name == "opencl" and (("android" in str(target.host)) or ("adreno" in str(target.attrs))) ) and not is_inner_reduction(block_stmt, iter_infos): ret = self.sch_outer_reduction(sch, config, main_block, blocks) if ret is not None: return ret # Step 0. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K] # Reindex first and than analyze the index map reindex_a = sch.reindex(main_block, ("read", 0)) reindex_b = sch.reindex(main_block, ("read", 1)) reindex_c = sch.reindex(main_block, ("write", 0)) index_maps = get_index_map(sch.get(main_block)) assert index_maps is not None matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps sch.transform_layout(reindex_a, ("write", 0), a_index_map) sch.transform_layout(reindex_b, ("write", 0), b_index_map) sch.transform_layout(reindex_c, ("read", 0), c_index_map) sch.transform_block_layout(main_block, matmul_index_map) # Step 1. Check Tensor Core support # Tensorization config: # If any value of I, J, K is fixed and less than this threshold, # tensorization rule will not be applied. minimal_tensorize_threshold = 64 block_stmt = sch.get(main_block) if target.kind.name == "cuda" and check_sm_version(str(target.attrs.get("arch", ""))) >= 70: apply_tensorization: bool = True # the batch dimension is not taken into consideration. for item_var in block_stmt.iter_vars[1:]: extent = item_var.dom.extent if isinstance(extent, tirx.expr.IntImm): if extent.value <= minimal_tensorize_threshold: apply_tensorization = False if apply_tensorization: # Analyze read/write buffers and choose correct tensorizer: int8 or fp16. in_dtype, out_dtype = get_in_out_dtypes(block_stmt) tensorize_sch = None if in_dtype == "int8" and out_dtype == "int32": tensorize_sch = MatmulInt8Tensorization().apply(func, target, _) elif in_dtype == "float16" and out_dtype in ["float16", "float32"]: tensorize_sch = MatmulTensorization().apply(func, target, _) if tensorize_sch is not None: return tensorize_sch elif target.kind.name == "metal": try: return MetalMatmul().apply(func, target, _) except: # pylint: disable=bare-except pass # Step 2. Schedule matmul y_kernel_size = config.vthread_y * config.block_size_y * config.micro_size_y x_kernel_size = config.vthread_x * config.block_size_x * config.micro_size_x if config.inner_x: sch.pad_einsum( main_block, [1, y_kernel_size, x_kernel_size, config.micro_size_k], ) batch, y, x, k = sch.get_loops(main_block) else: sch.pad_einsum( main_block, [1, x_kernel_size, y_kernel_size, config.micro_size_k], ) batch, x, y, k = sch.get_loops(main_block) by, vy, ty, yi = sch.split( y, [None, config.vthread_y, config.block_size_y, config.micro_size_y] ) bx, vx, tx, xi = sch.split( x, [None, config.vthread_x, config.block_size_x, config.micro_size_x] ) ko, ki = sch.split(k, factors=[None, config.micro_size_k]) reordered_loops = [by, bx, vy, vx, ty, tx, ko, ki] + ( [yi, xi] if config.inner_x else [xi, yi] ) sch.reorder(*reordered_loops) by = sch.fuse(batch, by) sch.bind(bx, "blockIdx.x") sch.bind(by, "blockIdx.y") sch.bind(vy, "vthread.y") sch.bind(vx, "vthread.x") sch.bind(ty, "threadIdx.y") sch.bind(tx, "threadIdx.x") inner_loop = config.micro_size_x if config.inner_x else config.micro_size_y if inner_loop % config.vector_size == 0: _, v = sch.split(reordered_loops[-1], [None, config.vector_size]) sch.vectorize(v) if config.unroll > 0: sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=config.unroll) sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1) l2g = sch.cache_write(main_block, 0, "local") sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True) if config.micro_size_x % config.vector_size == 0: _, v = sch.split(sch.get_loops(l2g)[-1], [None, config.vector_size]) sch.vectorize(v) if config.use_shared: def _cooperative_fetch(index, vec_len): block = sch.cache_read(main_block, index, "shared") num_loops = len(sch.get_loops(block)) sch.compute_at(block, ko, preserve_unit_loops=True) loops = sch.get_loops(block)[-num_loops:] ty, tx, _, vec = sch.split( sch.fuse(*loops), factors=[config.block_size_y, config.block_size_x, None, vec_len], ) sch.vectorize(vec) sch.bind(ty, "threadIdx.y") sch.bind(tx, "threadIdx.x") if config.storage_align: sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len) return block a_g2s = _cooperative_fetch(0, vec_len=config.vector_size) b_g2s = _cooperative_fetch(1, vec_len=config.vector_size) auto_inline_producers(sch, a_g2s) auto_inline_producers(sch, b_g2s) else: auto_inline_producers(sch, main_block) auto_inline_consumer_chain(sch, l2g) sch.decompose_reduction(main_block, ko) return sch def sch_outer_reduction( self, sch: s_tir.Schedule, config: Config, reduction_block: s_tir.schedule.SBlockRV, blocks: list[s_tir.schedule.SBlockRV], ) -> s_tir.Schedule | None: """Get vectorization factor""" def get_max_factor(n, factors): factors = sorted(factors, reverse=True) for factor in factors: if n % factor == 0: return factor return 1 reduction_loops = sch.get_loops(reduction_block) if not len(reduction_loops) == 4: return None mb, ms, n, k = reduction_loops if not ( isinstance(sch.get(n).extent, tirx.IntImm) and isinstance(sch.get(mb).extent, tirx.IntImm) and not isinstance(sch.get(ms).extent, tirx.IntImm) ): return None Threads_X, Threads_Y, VecSize, Unroll_M = ( config.block_size_x, config.block_size_y, config.vector_size, config.unroll, ) VecSize = min(get_max_factor(sch.get(n).extent // Threads_X, [1, 2, 4, 8]), VecSize) dequant_block = None matmul_block = reduction_block epilogue_block = None if blocks[-1] is not matmul_block: epilogue_block = blocks[-1] for blk in blocks[:-1]: if "dequantize" in sch.get(blk).name_hint: dequant_block = blk elif blk is not matmul_block: sch.compute_inline(blk) block = sch.reindex(reduction_block, ("read", 0)) sch.pad_einsum(reduction_block, [1, Unroll_M, 1, 1]) sch.compute_inline(block) trans_block, matmul_reindex = ( sch.get_producers(matmul_block)[0], sch.get_consumers(matmul_block)[0], ) if epilogue_block is not None: sch.compute_inline(matmul_reindex) matmul_reindex = epilogue_block sch.transform_layout( trans_block, ("write", 0), T.index_map(lambda i0, i1, i2: (i0, i1 // Unroll_M, i2, i1 % Unroll_M)), ) # transpose block schedules # sch.set_scope(trans_block, 0, "global.texture-1d") tb, tn, tk = sch.get_loops(trans_block) tbx, ttx = sch.split(tk, [None, Threads_X]) tby, tty, tc = sch.split(tn, [None, Threads_Y, Unroll_M]) sch.bind(tb, "blockIdx.z") sch.bind(tby, "blockIdx.y") sch.bind(tbx, "blockIdx.x") sch.bind(tty, "threadIdx.y") sch.bind(ttx, "threadIdx.x") sch.reorder(tb, tby, tbx, tty, ttx, tc) sch.vectorize(tc) mb, ms, n, k = sch.get_loops(matmul_block) m = sch.fuse(mb, ms) bx, tx, vec = sch.split(n, [None, Threads_X, VecSize]) by, ty, unr = sch.split(m, [None, Threads_Y, Unroll_M]) k1, k2, k3 = sch.split(k, [None, 4, 8]) sch.reorder(bx, by, tx, ty, k1, k2, k3, unr, vec) sch.set_scope(matmul_block, 0, "local") if dequant_block is not None: sch.compute_at(dequant_block, k3) sch.set_scope(dequant_block, 0, "local") sch.bind(by, "blockIdx.y") sch.bind(bx, "blockIdx.x") sch.bind(ty, "threadIdx.y") sch.bind(tx, "threadIdx.x") sch.vectorize(vec) inp = sch.cache_read(matmul_block, read_buffer_index=0, storage_scope="local") sch.compute_at(inp, k3, preserve_unit_loops=True) sch.vectorize(sch.get_loops(inp)[-1]) sch.unroll(unr) sch.unroll(k3) if dequant_block is not None: Aq_local = sch.cache_read(dequant_block, read_buffer_index=0, storage_scope="local") sch.compute_at(Aq_local, k2, preserve_unit_loops=True) sch.vectorize(sch.get_loops(Aq_local)[-1]) As_local = sch.cache_read(dequant_block, read_buffer_index=1, storage_scope="local") sch.compute_at(As_local, k1, preserve_unit_loops=True) sch.vectorize(sch.get_loops(As_local)[-1]) sch.vectorize(sch.get_loops(dequant_block)[-1]) sch.reverse_compute_at(matmul_reindex, ty) o_ur, o_vec = sch.get_loops(matmul_reindex)[-2:] sch.vectorize(o_vec) sch.unroll(o_ur) sch.decompose_reduction(matmul_block, k1) return sch