# 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=invalid-name,missing-function-docstring,unused-variable # ruff: noqa: E731 """Intrinsics for tensorization on NVIDIA GPU.""" from typing import Literal from tvm_ffi import register_global_func from tvm.runtime import convert from tvm.script import tirx as T from tvm.tirx import Cast, IntImm, TensorIntrin from tvm.tirx.function import PrimFunc def shared_16x16_to_ldmatrix_32x8_layout(i, j): thread_id = 4 * (i % 8) + (j % 8) // 2 return thread_id, 4 * (j // 8) + (i // 8) * 2 + (j % 2) def shared_16x32_to_ldmatrix_32x16_layout(i, j): thread_id = 4 * (i % 8) + (j % 16) // 4 return thread_id, 8 * (j // 16) + (i // 8) * 4 + j % 4 def shared_32x16_to_ldmatrix_32x16_layout(i, j): thread_id = (i % 16) // 4 + 4 * (j % 8) return thread_id, 8 * (j // 8) + (i // 16) * 4 + i % 4 def ldmatrix_32x8_to_shared_16x16_layout(thread_id, local_id): row = 8 * (local_id % 4 // 2) + (thread_id // 4) col = 8 * (local_id // 4) + (thread_id % 4) * 2 + (local_id % 2) return row, col @register_global_func("tirx.index_map.shared_16x16_to_ldmatrix_32x8_layout") def index_map_shared_16x16_to_ldmatrix_32x8_layout(ind): i, j = ind[0], ind[1] thread_id, local_id = shared_16x16_to_ldmatrix_32x8_layout(i, j) return convert([thread_id, local_id]) lift = convert M_DIM = 16 N_DIM = 16 WARP_SIZE = 32 HALF_WARP = WARP_SIZE // 2 HALF_WARP_expr = lift(HALF_WARP) def get_ldmatrix_intrin( k_dim: int, dtype: str, matrix_name: Literal["A", "B"], transposed: bool, shared_scope: str = "shared", ): local_size = (M_DIM * k_dim) // WARP_SIZE smem_offset = None index_map = None if matrix_name == "A": transpose_in_ldmatrix = transposed # transpose_layout_for_ldmatrix_input: Every thread loads 8 bytes data. This determines # which 8 bytes every thread loads. # If transpose_layout_for_ldmatrix_input is False, the load pattern is # T0 T0 T0 T0 T16 T16 T16 T16 # T1 T1 T1 T1 T17 T17 T17 T17 # ... # T8 T8 T8 T8 T24 T24 T24 T24 # T9 T9 T9 T9 T25 T25 T25 T25 # ... # T15 T15 T15 T15 T31 T31 T31 T31 # Otherwise, the load pattern is # T0 T0 T0 T0 T8 T8 T8 T8 # T1 T1 T1 T1 T9 T9 T9 T9 # ... # T7 T7 T7 T7 T15 T15 T15 T15 # T16 T16 T16 T16 T24 T24 T24 T24 # T17 T17 T17 T17 T25 T25 T25 T25 # ... # T23 T23 T23 T23 T31 T31 T31 T31 transpose_layout_for_ldmatrix_input = transposed smem_tile_row, smem_tile_col = (M_DIM, k_dim) if not transposed else (k_dim, M_DIM) else: assert matrix_name == "B" transpose_in_ldmatrix = not transposed transpose_layout_for_ldmatrix_input = transposed smem_tile_row, smem_tile_col = (k_dim, N_DIM) if not transposed else (N_DIM, k_dim) if k_dim == 16: assert dtype == "float16" index_map = shared_16x16_to_ldmatrix_32x8_layout if transpose_layout_for_ldmatrix_input: smem_offset = lambda tx, stride: ( stride * 8 * (tx // HALF_WARP_expr) + stride * (tx % 8) + 8 * ((tx % HALF_WARP_expr) // 8) ) else: smem_offset = lambda tx, stride: ( stride * (tx % HALF_WARP_expr) + 8 * (tx // HALF_WARP_expr) ) else: # TODO(yixin): Support TN and TT matmul for int8 assert matrix_name == "B" or not transposed, ( "Now only B matrix can be transposed for int8 matmul" ) assert k_dim == 32 and ( dtype == "int8" or dtype == "float8_e4m3fn" or dtype == "float8_e5m2" ), "Only k_dim == 16 (float16) or k_dim == 32 (int8) supported for now" if matrix_name == "B" and not transposed: index_map = shared_32x16_to_ldmatrix_32x16_layout # A dummy offset, ldmatrix cannot be used for int8 + trans case. # We still use the ldmatrix intrinsic, but lower it to a manual loop in the codegen. # Only the stride information is required. smem_offset = lambda _, stride: stride elif matrix_name == "B" and transposed: index_map = shared_16x32_to_ldmatrix_32x16_layout smem_offset = lambda tx, stride: ( stride * 8 * (tx // HALF_WARP_expr) + (tx % 8) * stride + 16 * ((tx % HALF_WARP_expr) // 8) ) else: # A, not transposed index_map = shared_16x32_to_ldmatrix_32x16_layout smem_offset = lambda tx, stride: stride * (tx % 16) + 16 * (tx // 16) offset_factor = smem_tile_col @T.prim_func(s_tir=True) def ldmatrix_desc(warp_handle: T.handle, shared_handle: T.handle) -> None: shared = T.match_buffer( shared_handle, (smem_tile_row, smem_tile_col), dtype, align=64, offset_factor=offset_factor, scope=shared_scope, ) warp = T.match_buffer( warp_handle, (WARP_SIZE, local_size), dtype, align=64, offset_factor=offset_factor, scope="warp", ) with T.sblock("root"): T.reads(shared[0:smem_tile_row, 0:smem_tile_col]) T.writes(warp[0:WARP_SIZE, 0:local_size]) for ax0, ax1 in T.grid(smem_tile_row, smem_tile_col): with T.sblock("shared_warp"): v0, v1 = T.axis.remap("SS", [ax0, ax1]) T.reads(shared[v0, v1]) thread_id, local_id = T.meta_var(index_map(v0, v1)) T.writes(warp[thread_id, local_id]) warp[thread_id, local_id] = shared[v0, v1] @T.prim_func(s_tir=True) def ldmatrix_impl(warp_handle: T.handle, shared_handle: T.handle) -> None: s0 = T.int32() s1 = T.int32() shared = T.match_buffer( shared_handle, (smem_tile_row, smem_tile_col), dtype, align=64, offset_factor=offset_factor, scope=shared_scope, strides=[s0, s1], ) warp = T.match_buffer( warp_handle, (WARP_SIZE, local_size), dtype, align=64, offset_factor=offset_factor, scope="warp", ) with T.sblock("root"): T.reads(shared[0:smem_tile_row, 0:smem_tile_col]) T.writes(warp[0:WARP_SIZE, 0:local_size]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): T.evaluate( T.ptx.ldmatrix_legacy( transpose_in_ldmatrix, 4, # Always load 4 matrices ".b16", warp.data, warp.elem_offset + lift(local_size) * tx, shared.access_ptr("r"), smem_offset(tx, s0), dtype=dtype, ) ) return ldmatrix_desc, ldmatrix_impl LDMATRIX_f16_A_INTRIN = "mma_ldmatrix_f16_a" TensorIntrin.register(LDMATRIX_f16_A_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", False)) LDMATRIX_f16_B_INTRIN = "mma_ldmatrix_f16_b" TensorIntrin.register(LDMATRIX_f16_B_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", False)) LDMATRIX_f16_A_TRANS_INTRIN = "mma_ldmatrix_f16_a_trans" TensorIntrin.register(LDMATRIX_f16_A_TRANS_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", True)) LDMATRIX_f16_B_TRANS_INTRIN = "mma_ldmatrix_f16_b_trans" TensorIntrin.register(LDMATRIX_f16_B_TRANS_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", True)) LDMATRIX_f16_A_DYN_INTRIN = "mma_ldmatrix_f16_a_dyn" TensorIntrin.register( LDMATRIX_f16_A_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", False, "shared.dyn") ) LDMATRIX_f16_B_DYN_INTRIN = "mma_ldmatrix_f16_b_dyn" TensorIntrin.register( LDMATRIX_f16_B_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", False, "shared.dyn") ) LDMATRIX_f16_A_TRANS_DYN_INTRIN = "mma_ldmatrix_f16_a_trans_dyn" TensorIntrin.register( LDMATRIX_f16_A_TRANS_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", True, "shared.dyn") ) LDMATRIX_f16_B_TRANS_DYN_INTRIN = "mma_ldmatrix_f16_b_trans_dyn" TensorIntrin.register( LDMATRIX_f16_B_TRANS_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", True, "shared.dyn") ) LDMATRIX_i8_A_INTRIN = "mma_ldmatrix_i8_a" TensorIntrin.register(LDMATRIX_i8_A_INTRIN, *get_ldmatrix_intrin(32, "int8", "A", False)) LDMATRIX_i8_B_INTRIN = "mma_ldmatrix_i8_b" TensorIntrin.register(LDMATRIX_i8_B_INTRIN, *get_ldmatrix_intrin(32, "int8", "B", False)) LDMATRIX_i8_B_TRANS_INTRIN = "mma_ldmatrix_i8_b_trans" TensorIntrin.register(LDMATRIX_i8_B_TRANS_INTRIN, *get_ldmatrix_intrin(32, "int8", "B", True)) LDMATRIX_e4m3_A_INTRIN = "mma_ldmatrix_e4m3_a" TensorIntrin.register(LDMATRIX_e4m3_A_INTRIN, *get_ldmatrix_intrin(32, "float8_e4m3fn", "A", False)) LDMATRIX_e4m3_B_INTRIN = "mma_ldmatrix_e4m3_b" TensorIntrin.register(LDMATRIX_e4m3_B_INTRIN, *get_ldmatrix_intrin(32, "float8_e4m3fn", "B", False)) LDMATRIX_e4m3_B_TRANS_INTRIN = "mma_ldmatrix_e4m3_b_trans" TensorIntrin.register( LDMATRIX_e4m3_B_TRANS_INTRIN, *get_ldmatrix_intrin(32, "float8_e4m3fn", "B", True) ) LDMATRIX_e5m2_A_INTRIN = "mma_ldmatrix_e5m2_a" TensorIntrin.register(LDMATRIX_e5m2_A_INTRIN, *get_ldmatrix_intrin(32, "float8_e5m2", "A", False)) LDMATRIX_e5m2_B_INTRIN = "mma_ldmatrix_e5m2_b" TensorIntrin.register(LDMATRIX_e5m2_B_INTRIN, *get_ldmatrix_intrin(32, "float8_e5m2", "B", False)) LDMATRIX_e5m2_B_TRANS_INTRIN = "mma_ldmatrix_e5m2_b_trans" TensorIntrin.register( LDMATRIX_e5m2_B_TRANS_INTRIN, *get_ldmatrix_intrin(32, "float8_e5m2", "B", True) ) def get_mma_intrin( k_dim, a_dtype="float16", b_dtype="float16", out_dtype="float16", a_transposed=False, b_transposed=False, ): local_size = (M_DIM * k_dim) // WARP_SIZE local_size_out = (M_DIM * N_DIM) // 32 index_map_C = shared_16x16_to_ldmatrix_32x8_layout if k_dim == 16: index_map_A = shared_16x16_to_ldmatrix_32x8_layout index_map_B = shared_16x16_to_ldmatrix_32x8_layout mma_prefix = "m16n8k16" elif k_dim == 32 and b_transposed: index_map_A = index_map_B = shared_16x32_to_ldmatrix_32x16_layout mma_prefix = "m16n8k32" elif k_dim == 32 and not b_transposed: index_map_A = shared_16x32_to_ldmatrix_32x16_layout index_map_B = shared_32x16_to_ldmatrix_32x16_layout mma_prefix = "m16n8k32" else: assert False dtype_abbrv = { "float16": "fp16", "float32": "fp32", "int8": "int8", "int32": "int32", "float8_e4m3fn": "e4m3", "float8_e5m2": "e5m2", } a_dtype_abbrv = dtype_abbrv[a_dtype] b_dtype_abbrv = dtype_abbrv[b_dtype] out_dtype_abbrv = dtype_abbrv[out_dtype] def cast_to_out_dtype(v): if out_dtype in ["float32", "int32"]: return Cast(out_dtype, v) return v def swap_if_flag(i, j, flag): return (j, i) if flag else (i, j) A_offset_factor = M_DIM if a_transposed else k_dim B_offset_factor = k_dim if b_transposed else N_DIM out_offset_factor = N_DIM @T.prim_func(s_tir=True) def mma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None: A = T.match_buffer( a, (WARP_SIZE, local_size), a_dtype, align=64, offset_factor=A_offset_factor, scope="warp", ) B = T.match_buffer( b, (WARP_SIZE, local_size), b_dtype, align=64, offset_factor=B_offset_factor, scope="warp", ) C = T.match_buffer( c, (WARP_SIZE, local_size_out), out_dtype, align=64, offset_factor=out_offset_factor, scope="warp", ) with T.sblock("root"): T.reads( C[0:WARP_SIZE, 0:local_size_out], A[0:WARP_SIZE, 0:local_size], B[0:WARP_SIZE, 0:local_size], ) T.writes(C[0:WARP_SIZE, 0:local_size_out]) for i, j, k in T.grid(M_DIM, N_DIM, k_dim): with T.sblock("C"): vi, vj, vk = T.axis.remap("SSR", [i, j, k]) a_row_ind, a_col_ind = T.meta_var(swap_if_flag(vi, vk, a_transposed)) b_row_ind, b_col_ind = T.meta_var(swap_if_flag(vk, vj, b_transposed)) thread_id_C, local_id_C = T.meta_var(index_map_C(vi, vj)) thread_id_A, local_id_A = T.meta_var(index_map_A(a_row_ind, a_col_ind)) thread_id_B, local_id_B = T.meta_var(index_map_B(b_row_ind, b_col_ind)) T.reads( C[thread_id_C, local_id_C], A[thread_id_A, local_id_A], B[thread_id_B, local_id_B], ) T.writes(C[thread_id_C, local_id_C]) C[thread_id_C, local_id_C] += cast_to_out_dtype( A[thread_id_A, local_id_A] ) * cast_to_out_dtype(B[thread_id_B, local_id_B]) @T.prim_func(s_tir=True) def mma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None: A = T.match_buffer( a, (WARP_SIZE, local_size), a_dtype, align=64, offset_factor=A_offset_factor, scope="warp", ) B = T.match_buffer( b, (WARP_SIZE, local_size), b_dtype, align=64, offset_factor=B_offset_factor, scope="warp", ) C = T.match_buffer( c, (WARP_SIZE, local_size_out), out_dtype, align=64, offset_factor=out_offset_factor, scope="warp", ) with T.sblock("root"): T.reads( C[0:WARP_SIZE, 0:local_size_out], A[0:WARP_SIZE, 0:local_size], B[0:WARP_SIZE, 0:local_size], ) T.writes(C[0:WARP_SIZE, 0:local_size_out]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): T.evaluate( T.ptx.mma.legacy( mma_prefix, "row", "col", a_dtype_abbrv, b_dtype_abbrv, out_dtype_abbrv, A.data, A.elem_offset + tx * lift(local_size), B.data, B.elem_offset + tx * lift(local_size), C.data, C.elem_offset + tx * lift(local_size_out), False, dtype=out_dtype, ) ) T.evaluate( T.ptx.mma.legacy( mma_prefix, "row", "col", a_dtype_abbrv, b_dtype_abbrv, out_dtype_abbrv, A.data, A.elem_offset + tx * lift(local_size), B.data, B.elem_offset + tx * lift(local_size) + lift(local_size) // 2, C.data, C.elem_offset + tx * lift(local_size_out) + lift(local_size_out) // 2, False, dtype=out_dtype, ) ) return mma_sync_desc, mma_sync_impl MMA_f16f16f32_INTRIN = "mma_f16f16f32" TensorIntrin.register( MMA_f16f16f32_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", False, False) ) MMA_f16f16f32_TRANS_B_INTRIN = "mma_f16f16f32_trans_b" TensorIntrin.register( MMA_f16f16f32_TRANS_B_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", False, True) ) MMA_f16f16f32_TRANS_A_INTRIN = "mma_f16f16f32_trans_a" TensorIntrin.register( MMA_f16f16f32_TRANS_A_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", True, False) ) MMA_f16f16f32_TRANS_A_TRANS_B_INTRIN = "mma_f16f16f32_trans_a_trans_b" TensorIntrin.register( MMA_f16f16f32_TRANS_A_TRANS_B_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", True, True), ) MMA_f16f16f16_INTRIN = "mma_f16f16f16" TensorIntrin.register( MMA_f16f16f16_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", False, False) ) MMA_f16f16f16_TRANS_B_INTRIN = "mma_f16f16f16_trans_b" TensorIntrin.register( MMA_f16f16f16_TRANS_B_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", False, True) ) MMA_f16f16f16_TRANS_A_INTRIN = "mma_f16f16f16_trans_a" TensorIntrin.register( MMA_f16f16f16_TRANS_A_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", True, False) ) MMA_f16f16f16_TRANS_A_TRANS_B_INTRIN = "mma_f16f16f16_trans_a_trans_b" TensorIntrin.register( MMA_f16f16f16_TRANS_A_TRANS_B_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", True, True), ) MMA_i8i8i32_INTRIN = "mma_i8i8i32" TensorIntrin.register( MMA_i8i8i32_INTRIN, *get_mma_intrin(32, "int8", "int8", "int32", False, False) ) MMA_i8i8i32_TRANS_B_INTRIN = "mma_i8i8i32_trans_b" TensorIntrin.register( MMA_i8i8i32_TRANS_B_INTRIN, *get_mma_intrin(32, "int8", "int8", "int32", False, True) ) MMA_e5m2e5m2f32_INTRIN = "mma_e5m2e5m2f32" TensorIntrin.register( MMA_e5m2e5m2f32_INTRIN, *get_mma_intrin(32, "float8_e5m2", "float8_e5m2", "float32", False, False), ) MMA_e5m2e5m2f32_TRANS_B_INTRIN = "mma_e5m2e5m2f32_trans_b" TensorIntrin.register( MMA_e5m2e5m2f32_TRANS_B_INTRIN, *get_mma_intrin(32, "float8_e5m2", "float8_e5m2", "float32", False, True), ) MMA_e4m3e4m3f32_INTRIN = "mma_e4m3e4m3f32" TensorIntrin.register( MMA_e4m3e4m3f32_INTRIN, *get_mma_intrin(32, "float8_e4m3fn", "float8_e4m3fn", "float32", False, False), ) MMA_e4m3e4m3f32_TRANS_B_INTRIN = "mma_e4m3e4m3f32_trans_b" TensorIntrin.register( MMA_e4m3e4m3f32_TRANS_B_INTRIN, *get_mma_intrin(32, "float8_e4m3fn", "float8_e4m3fn", "float32", False, True), ) def get_mma_fill_intrin(dtype, local_size): zero = IntImm("int32", 0).astype(dtype) # Assume M = N = 16 index_map = shared_16x16_to_ldmatrix_32x8_layout @T.prim_func(s_tir=True) def mma_fill_desc(a: T.handle) -> None: C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp") with T.sblock("root"): T.reads() T.writes(C_warp[0:WARP_SIZE, 0:local_size]) for i0, i1 in T.grid(M_DIM, N_DIM): with T.sblock("C_warp"): i, j = T.axis.remap("SS", [i0, i1]) thread_id, local_id = T.meta_var(index_map(i, j)) T.reads() T.writes(C_warp[thread_id, local_id]) C_warp[thread_id, local_id] = zero @T.prim_func(s_tir=True) def mma_fill_impl(a: T.handle) -> None: C_warp = T.match_buffer( a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1 ) with T.sblock("root"): T.reads() T.writes(C_warp[0:WARP_SIZE, 0:local_size]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): T.evaluate( T.cuda.mma_fill_legacy(local_size, C_warp.data, C_warp.elem_offset, dtype=dtype) ) return mma_fill_desc, mma_fill_impl MMA_fill_16x16_f32_INTRIN = "mma_fill_16x16_f32" TensorIntrin.register(MMA_fill_16x16_f32_INTRIN, *get_mma_fill_intrin("float32", 8)) MMA_fill_16x16_f16_INTRIN = "mma_fill_16x16_f16" TensorIntrin.register(MMA_fill_16x16_f16_INTRIN, *get_mma_fill_intrin("float16", 8)) MMA_fill_16x16_i32_INTRIN = "mma_fill_16x16_i32" TensorIntrin.register(MMA_fill_16x16_i32_INTRIN, *get_mma_fill_intrin("int32", 8)) def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrinic=True): # Assume M = N = 16 index_map = shared_16x16_to_ldmatrix_32x8_layout index_map_rev = ldmatrix_32x8_to_shared_16x16_layout @T.prim_func(s_tir=True) def mma_store_desc(a: T.handle, c: T.handle) -> None: C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp") C = T.match_buffer(c, [M_DIM, N_DIM], dtype=dtype, scope=scope) with T.sblock("root"): T.reads(C_warp[0:WARP_SIZE, 0:local_size]) T.writes(C[0:M_DIM, 0:N_DIM]) for i0, i1 in T.grid(M_DIM, N_DIM): with T.sblock("C_warp"): v0, v1 = T.axis.remap("SS", [i0, i1]) thread_id, local_id = T.meta_var(index_map(v0, v1)) T.reads(C_warp[thread_id, local_id]) T.writes(C[v0, v1]) C[v0, v1] = C_warp[thread_id, local_id] if use_mma_store_intrinic: @T.prim_func(s_tir=True) def mma_store_impl(a: T.handle, c: T.handle) -> None: s0 = T.int32() s1 = T.int32() C_warp = T.match_buffer( a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1 ) C = T.match_buffer( c, [M_DIM, N_DIM], dtype=dtype, scope=scope, offset_factor=1, strides=[s0, s1] ) with T.sblock("root"): T.reads(C_warp[0:WARP_SIZE, 0:local_size]) T.writes(C[0:M_DIM, 0:N_DIM]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): T.evaluate( T.cuda.mma_store_legacy( M_DIM, N_DIM, C.access_ptr("w"), C_warp.data, C_warp.elem_offset, s0, dtype=dtype, ) ) else: @T.prim_func(s_tir=True) def mma_store_impl(a: T.handle, c: T.handle) -> None: s0 = T.int32() s1 = T.int32() C_warp = T.match_buffer( a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1 ) C = T.match_buffer( c, [M_DIM, N_DIM], dtype=dtype, scope=scope, offset_factor=1, strides=[s0, s1] ) with T.sblock("root"): T.reads(C_warp[0:WARP_SIZE, 0:local_size]) T.writes(C[0:M_DIM, 0:N_DIM]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): for local_id in T.serial(local_size): row, col = T.meta_var(index_map_rev(tx, local_id)) C[row, col] = C_warp[tx, local_id] return mma_store_desc, mma_store_impl MMA_store_16x16_f32_global_INTRIN = "mma_store_16x16_f32_global_" TensorIntrin.register( MMA_store_16x16_f32_global_INTRIN, *get_mma_store_intrin("float32", 8, "global", True) ) MMA_store_16x16_f32_shared_dyn_INTRIN = "mma_store_16x16_f32_shared_dyn_" TensorIntrin.register( MMA_store_16x16_f32_shared_dyn_INTRIN, *get_mma_store_intrin("float32", 8, "shared.dyn", True) ) MMA_store_16x16_f32_shared_dyn_INTRIN_SIMPLE = "mma_store_16x16_f32_shared_dyn_simple_" TensorIntrin.register( MMA_store_16x16_f32_shared_dyn_INTRIN_SIMPLE, *get_mma_store_intrin("float32", 8, "shared.dyn", False), ) MMA_store_16x16_f16_shared_dyn_INTRIN_SIMPLE = "mma_store_16x16_f16_shared_dyn_simple_" TensorIntrin.register( MMA_store_16x16_f16_shared_dyn_INTRIN_SIMPLE, *get_mma_store_intrin("float16", 8, "shared.dyn", False), ) MMA_store_16x16_f16_global_INTRIN = "mma_store_16x16_f16_global_" TensorIntrin.register( MMA_store_16x16_f16_global_INTRIN, *get_mma_store_intrin("float16", 8, "global", True) ) MMA_store_16x16_i32_global_INTRIN = "mma_store_16x16_i32_global_" TensorIntrin.register( MMA_store_16x16_i32_global_INTRIN, *get_mma_store_intrin("int32", 8, "global", True) ) def get_mma_intrin_group( load_scope: Literal["shared", "shared.dyn"], store_scope: Literal["global", "shared", "shared.dyn"], in_dtype: Literal["float16", "int8", "float8_e4m3fn", "float8_e5m2"], out_dtype: Literal["float16", "float32", "int32"], trans_a: bool, trans_b: bool, not_use_mma_store_intrinic: bool = True, store_to_smem_dtype: Literal["float16", "float32", "int32"] | None = None, ) -> dict[str, str]: """Get a group of intrinsics for mma tensor core with the given configurations Parameters ---------- load_scope : Literal["shared", "shared.dyn"] The memory scope of the input buffer. store_scope : Literal["global", "shared", "shared.dyn"] The memory scope of the result buffer. in_dtype : str The input data type. out_dtype : str The output data dtype. trans_a : bool Whether the input matrix A is transposed. trans_b : bool Whether the input matrix B is transposed. not_use_mma_store_intrinic : bool Whether to not use the mma_store intrinsic. If True, use BufferStore stmts to store the result of mma. Otherwise, use mma_store intrinsic. This is because if we use mma_store intrinsic, during swizzling shared memory visits, our rearrangement scheme will involve areas accessed by different mma_store calls. This makes swizzling quite complex. But BufferStore will not face this problem. store_to_smem_dtype : Optional[Literal["float16", "float32", "int32"]] The dtype that we use to store from register to shared memory. By default it is out_dtype. Returns ------- ret : Dict[str, str] A group of tensor intrinsics. """ assert load_scope in ["shared", "shared.dyn"] assert store_scope in ["global", "shared", "shared.dyn"] assert in_dtype in ["float16", "int8", "float8_e4m3fn", "float8_e5m2"] assert out_dtype in ["float16", "float32", "int32"] shape = "16x16" dtype_mapping = { "float16": "f16", "float32": "f32", "int8": "i8", "float8_e4m3fn": "e4m3", "float8_e5m2": "e5m2", "int32": "i32", } a_dtype = dtype_mapping[in_dtype] b_dtype = dtype_mapping[in_dtype] out_dtype = dtype_mapping[out_dtype] # e.g. mma_fill_16x16_f32 init_intrin = f"mma_fill_{shape}_{out_dtype}" # e.g. mma_ldmatrix_f16_a_trans_dyn, mma_ldmatrix_f16_b_trans_dyn trans_a = "_trans" if trans_a else "" trans_b = "_trans" if trans_b else "" load_scope = "_dyn" if load_scope == "shared.dyn" else "" load_a_intrin = f"mma_ldmatrix_{a_dtype}_a{trans_a}{load_scope}" load_b_intrin = f"mma_ldmatrix_{b_dtype}_b{trans_b}{load_scope}" # e.g. mma_f16f16f32_trans_a_trans_b trans_a_str = trans_a + "_a" if trans_a != "" else "" trans_b_str = trans_b + "_b" if trans_b != "" else "" compute_intrin = f"mma_{a_dtype}{b_dtype}{out_dtype}{trans_a_str}{trans_b_str}" # e.g. mma_store_16x16_f32_shared_dyn_simple_ store_scope = store_scope.replace(".", "_") store_to_smem_dtype = dtype_mapping[store_to_smem_dtype] if store_to_smem_dtype else out_dtype suffix = "simple_" if not_use_mma_store_intrinic else "" store_intrin = f"mma_store_{shape}_{store_to_smem_dtype}_{store_scope}_{suffix}" return { "init": init_intrin, "load_a": load_a_intrin, "load_b": load_b_intrin, "compute": compute_intrin, "store": store_intrin, } ######## WMMA intrinsics ######## def get_wmma_fragment_index(buffer, stride, m_dim, n_dim): """Compute wmma fragment index using elem_offset of the buffer""" frag_index_m = buffer.elem_offset // stride // m_dim frag_index_n = buffer.elem_offset % stride // n_dim num_fragments_per_row = stride // n_dim return frag_index_m * num_fragments_per_row + frag_index_n def get_wmma_load_intrin( m_dim: int, n_dim: int, k_dim: int, dtype: str, shared_scope: str, is_b: bool, is_col_major: bool, ) -> tuple[PrimFunc, PrimFunc]: """Generator of wmma_load intrins""" wmma_fragment_scope = f"wmma.matrix_{'b' if is_b else 'a'}" layout = "col_major" if is_col_major else "row_major" frag_m, frag_n = (k_dim, n_dim) if is_b else (m_dim, k_dim) if is_col_major: frag_m, frag_n = frag_n, frag_m offset_factor = frag_n @T.prim_func(s_tir=True) def wmma_load_desc(a: T.handle, c: T.handle) -> None: A = T.match_buffer( a, (frag_m, frag_n), dtype, align=64, offset_factor=offset_factor, scope=shared_scope ) C = T.match_buffer( c, (frag_m, frag_n), dtype, align=64, offset_factor=offset_factor, scope=wmma_fragment_scope, ) with T.sblock("root"): T.reads(A[0:frag_m, 0:frag_n]) T.writes(C[0:frag_m, 0:frag_n]) for i, j in T.grid(frag_m, frag_n): with T.sblock("load"): vii, vjj = T.axis.remap("SS", [i, j]) C[vii, vjj] = A[vii, vjj] @T.prim_func(s_tir=True) def wmma_load_impl(a: T.handle, c: T.handle) -> None: s1 = T.int32() s0 = T.int32() d1 = T.int32() d0 = T.int32() A = T.match_buffer( a, (frag_m, frag_n), dtype, align=64, offset_factor=offset_factor, scope=shared_scope, strides=[s1, s0], ) C = T.match_buffer( c, (frag_m, frag_n), dtype, align=64, offset_factor=offset_factor, scope=wmma_fragment_scope, strides=[d1, d0], ) with T.sblock("root"): T.reads(A[0:frag_m, 0:frag_n]) T.writes(C[0:frag_m, 0:frag_n]) T.evaluate( T.tvm_load_matrix_sync( C.data, m_dim, n_dim, k_dim, get_wmma_fragment_index(C, d1, frag_m, frag_n), A.access_ptr("r", ptr_type=dtype), s1, layout, dtype="void", ) ) return wmma_load_desc, wmma_load_impl def get_wmma_fill_intrin( m_dim: int, n_dim: int, k_dim: int, dtype: str ) -> tuple[PrimFunc, PrimFunc]: """Generator of wmma_fill intrins""" zero = IntImm("int32", 0).astype(dtype) offset_factor = n_dim @T.prim_func(s_tir=True) def wmma_fill_desc(c: T.handle) -> None: C = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope="wmma.accumulator", ) with T.sblock("root"): T.reads() T.writes(C[0:m_dim, 0:n_dim]) for i, j in T.grid(m_dim, n_dim): with T.sblock("init"): vii, vjj = T.axis.remap("SS", [i, j]) C[vii, vjj] = zero @T.prim_func(s_tir=True) def wmma_fill_impl(c: T.handle) -> None: d1 = T.int32() d0 = T.int32() C = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope="wmma.accumulator", strides=[d1, d0], ) with T.sblock("root"): T.reads() T.writes(C[0:m_dim, 0:n_dim]) T.evaluate( T.tvm_fill_fragment( C.data, m_dim, n_dim, k_dim, get_wmma_fragment_index(C, d1, m_dim, n_dim), T.float32(0), dtype="void", ) ) return wmma_fill_desc, wmma_fill_impl def get_wmma_store_intrin( m_dim: int, n_dim: int, k_dim: int, dtype: str, scope: str ) -> tuple[PrimFunc, PrimFunc]: """Generator of wmma_store intrins""" offset_factor = n_dim @T.prim_func(s_tir=True) def wmma_store_desc(a: T.handle, c: T.handle) -> None: A = T.match_buffer( a, (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope="wmma.accumulator", ) C = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope=scope ) with T.sblock("root"): T.reads(A[0:m_dim, 0:n_dim]) T.writes(C[0:m_dim, 0:n_dim]) for i, j in T.grid(m_dim, n_dim): with T.sblock("store"): vii, vjj = T.axis.remap("SS", [i, j]) C[vii, vjj] = A[vii, vjj] @T.prim_func(s_tir=True) def wmma_store_impl(a: T.handle, c: T.handle) -> None: s1 = T.int32() s0 = T.int32() d1 = T.int32() d0 = T.int32() A = T.match_buffer( a, (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope="wmma.accumulator", strides=[d1, d0], ) C = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=offset_factor, scope=scope, strides=[s1, s0], ) with T.sblock("root"): T.reads(A[0:m_dim, 0:n_dim]) T.writes(C[0:m_dim, 0:n_dim]) T.evaluate( T.tvm_store_matrix_sync( A.data, m_dim, n_dim, k_dim, get_wmma_fragment_index(A, d1, m_dim, n_dim), C.access_ptr("w", ptr_type=dtype), s1, "row_major", dtype="void", ) ) return wmma_store_desc, wmma_store_impl def get_wmma_sync_intrin( m_dim: int, n_dim: int, k_dim: int, in_dtype: str, out_dtype: str, b_transposed: bool ) -> tuple[PrimFunc, PrimFunc]: """Generator of wmma_sync intrins""" def maybe_cast(v): if in_dtype != out_dtype: return Cast(out_dtype, v) return v def maybe_swap(i, j): if b_transposed: return j, i return i, j b_shape_0, b_shape_1 = maybe_swap(k_dim, n_dim) A_offset_factor = k_dim B_offset_factor = b_shape_1 out_offset_factor = n_dim @T.prim_func(s_tir=True) def wmma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None: A = T.match_buffer( a, (m_dim, k_dim), in_dtype, align=64, offset_factor=A_offset_factor, scope="wmma.matrix_a", ) B = T.match_buffer( b, maybe_swap(k_dim, n_dim), in_dtype, align=64, offset_factor=B_offset_factor, scope="wmma.matrix_b", ) C = T.match_buffer( c, (m_dim, n_dim), out_dtype, align=64, offset_factor=out_offset_factor, scope="wmma.accumulator", ) with T.sblock("root"): T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:b_shape_0, 0:b_shape_1]) T.writes(C[0:m_dim, 0:n_dim]) for i, j, k in T.grid(m_dim, n_dim, k_dim): with T.sblock(""): vii, vjj, vkk = T.axis.remap("SSR", [i, j, k]) B_index_0, B_index_1 = T.meta_var(maybe_swap(vkk, vjj)) C[vii, vjj] = C[vii, vjj] + maybe_cast(A[vii, vkk]) * maybe_cast( B[B_index_0, B_index_1] ) @T.prim_func(s_tir=True) def wmma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None: a1 = T.int32() a0 = T.int32() b1 = T.int32() b0 = T.int32() c1 = T.int32() c0 = T.int32() A = T.match_buffer( a, (m_dim, k_dim), in_dtype, align=64, offset_factor=A_offset_factor, scope="wmma.matrix_a", strides=[a1, a0], ) B = T.match_buffer( b, maybe_swap(k_dim, n_dim), in_dtype, align=64, offset_factor=B_offset_factor, scope="wmma.matrix_b", strides=[b1, b0], ) C = T.match_buffer( c, (m_dim, n_dim), out_dtype, align=64, offset_factor=out_offset_factor, scope="wmma.accumulator", strides=[c1, c0], ) with T.sblock("root"): T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:b_shape_0, 0:b_shape_1]) T.writes(C[0:m_dim, 0:n_dim]) T.evaluate( T.tvm_mma_sync( C.data, get_wmma_fragment_index(C, c1, m_dim, n_dim), A.data, get_wmma_fragment_index(A, a1, m_dim, k_dim), B.data, get_wmma_fragment_index(B, b1, b_shape_0, b_shape_1), C.data, get_wmma_fragment_index(C, c1, m_dim, n_dim), dtype="void", ) ) return wmma_sync_desc, wmma_sync_impl WMMA_SYNC_16x16x16_f16f16f32_INTRIN = "wmma_sync_16x16x16_f16f16f32" TensorIntrin.register( WMMA_SYNC_16x16x16_f16f16f32_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "float16", "float32", False), ) WMMA_SYNC_16x16x16_f16f16f32_TRANS_INTRIN = "wmma_sync_16x16x16_f16f16f32_trans" TensorIntrin.register( WMMA_SYNC_16x16x16_f16f16f32_TRANS_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "float16", "float32", True), ) WMMA_SYNC_16x16x16_f16f16f16_INTRIN = "wmma_sync_16x16x16_f16f16f16" TensorIntrin.register( WMMA_SYNC_16x16x16_f16f16f16_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "float16", "float16", False), ) WMMA_SYNC_16x16x16_f16f16f16_TRANS_INTRIN = "wmma_sync_16x16x16_f16f16f16_trans" TensorIntrin.register( WMMA_SYNC_16x16x16_f16f16f16_TRANS_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "float16", "float16", True), ) WMMA_SYNC_16x16x16_s8s8s32_INTRIN = "wmma_sync_16x16x16_s8s8s32" TensorIntrin.register( WMMA_SYNC_16x16x16_s8s8s32_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "int8", "int32", False) ) WMMA_SYNC_16x16x16_s8s8s32_TRANS_INTRIN = "wmma_sync_16x16x16_s8s8s32_trans" TensorIntrin.register( WMMA_SYNC_16x16x16_s8s8s32_TRANS_INTRIN, *get_wmma_sync_intrin(16, 16, 16, "int8", "int32", True), ) WMMA_SYNC_8x8x32_s4s4s32_TRANS_INTRIN = "wmma_sync_8x8x32_s4s4s32_trans" TensorIntrin.register( WMMA_SYNC_8x8x32_s4s4s32_TRANS_INTRIN, *get_wmma_sync_intrin(8, 8, 32, "int4", "int32", True) ) WMMA_LOAD_16x16x16_F16_A_INTRIN = "wmma_load_16x16x16_f16_a_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_A_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared", False, False), ) WMMA_LOAD_16x16x16_F16_A_DYN_INTRIN = "wmma_load_16x16x16_f16_a_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_A_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", False, False), ) WMMA_LOAD_16x16x16_F16_B_INTRIN = "wmma_load_16x16x16_f16_b_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_B_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared", True, False), ) WMMA_LOAD_16x16x16_F16_B_DYN_INTRIN = "wmma_load_16x16x16_f16_b_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_B_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", True, False), ) WMMA_LOAD_16x16x16_F16_A_TRANS_INTRIN = "wmma_load_16x16x16_f16_a_trans_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_A_TRANS_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared", False, True), ) WMMA_LOAD_16x16x16_F16_A_TRANS_DYN_INTRIN = "wmma_load_16x16x16_f16_a_trans_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_A_TRANS_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", False, True), ) WMMA_LOAD_16x16x16_F16_B_TRANS_INTRIN = "wmma_load_16x16x16_f16_b_trans_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_B_TRANS_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared", True, True), ) WMMA_LOAD_16x16x16_F16_B_TRANS_DYN_INTRIN = "wmma_load_16x16x16_f16_b_trans_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_F16_B_TRANS_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "float16", "shared.dyn", True, True), ) WMMA_LOAD_16x16x16_S8_A_INTRIN = "wmma_load_16x16x16_s8_a_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_A_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared", False, False), ) WMMA_LOAD_16x16x16_S8_A_DYN_INTRIN = "wmma_load_16x16x16_s8_a_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_A_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", False, False), ) WMMA_LOAD_16x16x16_S8_B_INTRIN = "wmma_load_16x16x16_s8_b_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_B_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared", True, False) ) WMMA_LOAD_16x16x16_S8_B_DYN_INTRIN = "wmma_load_16x16x16_s8_b_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_B_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", True, False), ) WMMA_LOAD_16x16x16_S8_A_TRANS_INTRIN = "wmma_load_16x16x16_s8_a_trans_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_A_TRANS_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared", False, True), ) WMMA_LOAD_16x16x16_S8_A_TRANS_DYN_INTRIN = "wmma_load_16x16x16_s8_a_trans_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_A_TRANS_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", False, True), ) WMMA_LOAD_16x16x16_S8_B_TRANS_INTRIN = "wmma_load_16x16x16_s8_b_trans_shared" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_B_TRANS_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared", True, True), ) WMMA_LOAD_16x16x16_S8_B_TRANS_DYN_INTRIN = "wmma_load_16x16x16_s8_b_trans_shared_dyn" TensorIntrin.register( WMMA_LOAD_16x16x16_S8_B_TRANS_DYN_INTRIN, *get_wmma_load_intrin(16, 16, 16, "int8", "shared.dyn", True, True), ) WMMA_LOAD_8x8x32_S4_A_INTRIN = "wmma_load_8x8x32_s4_a_shared" TensorIntrin.register( WMMA_LOAD_8x8x32_S4_A_INTRIN, *get_wmma_load_intrin(8, 8, 32, "int4", "shared", False, False) ) WMMA_LOAD_8x8x32_S4_A_DYN_INTRIN = "wmma_load_8x8x32_s4_a_shared_dyn" TensorIntrin.register( WMMA_LOAD_8x8x32_S4_A_DYN_INTRIN, *get_wmma_load_intrin(8, 8, 32, "int4", "shared.dyn", False, False), ) WMMA_LOAD_8x8x32_S4_B_TRANS_INTRIN = "wmma_load_8x8x32_s4_b_trans_shared" TensorIntrin.register( WMMA_LOAD_8x8x32_S4_B_TRANS_INTRIN, *get_wmma_load_intrin(8, 8, 32, "int4", "shared", True, True), ) WMMA_LOAD_8x8x32_S4_B_TRANS_DYN_INTRIN = "wmma_load_8x8x32_s4_b_trans_shared_dyn" TensorIntrin.register( WMMA_LOAD_8x8x32_S4_B_TRANS_DYN_INTRIN, *get_wmma_load_intrin(8, 8, 32, "int4", "shared.dyn", True, True), ) WMMA_FILL_16x16x16_F32_INTRIN = "wmma_fill_16x16x16_f32" TensorIntrin.register(WMMA_FILL_16x16x16_F32_INTRIN, *get_wmma_fill_intrin(16, 16, 16, "float32")) WMMA_FILL_16x16x16_F16_INTRIN = "wmma_fill_16x16x16_f16" TensorIntrin.register(WMMA_FILL_16x16x16_F16_INTRIN, *get_wmma_fill_intrin(16, 16, 16, "float16")) WMMA_FILL_16x16x16_S32_INTRIN = "wmma_fill_16x16x16_s32" TensorIntrin.register(WMMA_FILL_16x16x16_S32_INTRIN, *get_wmma_fill_intrin(16, 16, 16, "int32")) WMMA_FILL_8x8x32_S32_INTRIN = "wmma_fill_8x8x32_s32" TensorIntrin.register(WMMA_FILL_8x8x32_S32_INTRIN, *get_wmma_fill_intrin(8, 8, 32, "int32")) WMMA_STORE_16x16x16_F32_SHARED_INTRIN = "wmma_store_16x16x16_f32_shared" TensorIntrin.register( WMMA_STORE_16x16x16_F32_SHARED_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float32", "shared") ) WMMA_STORE_16x16x16_F32_SHARED_DYN_INTRIN = "wmma_store_16x16x16_f32_shared_dyn" TensorIntrin.register( WMMA_STORE_16x16x16_F32_SHARED_DYN_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float32", "shared.dyn"), ) WMMA_STORE_16x16x16_F16_SHARED_INTRIN = "wmma_store_16x16x16_f16_shared" TensorIntrin.register( WMMA_STORE_16x16x16_F16_SHARED_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float16", "shared") ) WMMA_STORE_16x16x16_F16_SHARED_DYN_INTRIN = "wmma_store_16x16x16_f16_shared_dyn" TensorIntrin.register( WMMA_STORE_16x16x16_F16_SHARED_DYN_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float16", "shared.dyn"), ) WMMA_STORE_16x16x16_S32_SHARED_INTRIN = "wmma_store_16x16x16_s32_shared" TensorIntrin.register( WMMA_STORE_16x16x16_S32_SHARED_INTRIN, *get_wmma_store_intrin(16, 16, 16, "int32", "shared") ) WMMA_STORE_16x16x16_S32_SHARED_DYN_INTRIN = "wmma_store_16x16x16_s32_shared_dyn" TensorIntrin.register( WMMA_STORE_16x16x16_S32_SHARED_DYN_INTRIN, *get_wmma_store_intrin(16, 16, 16, "int32", "shared.dyn"), ) WMMA_STORE_8x8x32_S32_SHARED_INTRIN = "wmma_store_8x8x32_s32_shared" TensorIntrin.register( WMMA_STORE_8x8x32_S32_SHARED_INTRIN, *get_wmma_store_intrin(8, 8, 32, "int32", "shared") ) WMMA_STORE_8x8x32_S32_SHARED_DYN_INTRIN = "wmma_store_8x8x32_s32_shared_dyn" TensorIntrin.register( WMMA_STORE_8x8x32_S32_SHARED_DYN_INTRIN, *get_wmma_store_intrin(8, 8, 32, "int32", "shared.dyn") ) WMMA_STORE_16x16x16_F32_GLOBAL_INTRIN = "wmma_store_16x16x16_f32_global" TensorIntrin.register( WMMA_STORE_16x16x16_F32_GLOBAL_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float32", "global") ) WMMA_STORE_16x16x16_F16_GLOBAL_INTRIN = "wmma_store_16x16x16_f16_global" TensorIntrin.register( WMMA_STORE_16x16x16_F16_GLOBAL_INTRIN, *get_wmma_store_intrin(16, 16, 16, "float16", "global") ) WMMA_STORE_16x16x16_S32_GLOBAL_INTRIN = "wmma_store_16x16x16_s32_global" TensorIntrin.register( WMMA_STORE_16x16x16_S32_GLOBAL_INTRIN, *get_wmma_store_intrin(16, 16, 16, "int32", "global") ) WMMA_STORE_8x8x32_S32_GLOBAL_INTRIN = "wmma_store_8x8x32_s32_global" TensorIntrin.register( WMMA_STORE_8x8x32_S32_GLOBAL_INTRIN, *get_wmma_store_intrin(8, 8, 32, "int32", "global") ) def get_wmma_intrin_group( load_scope: Literal["shared", "shared.dyn"], store_scope: Literal["global", "shared", "shared.dyn"], in_dtype: str, out_dtype: str, trans_b: bool, ) -> dict[str, str]: """Get a group of intrinsics for wmma tensor core with the given configurations Parameters ---------- load_scope : Literal["shared", "shared.dyn"] The memory scope of the input buffer. store_scope : Literal["global", "shared", "shared.dyn"] The memory scope of the result buffer. in_dtype : str The input data type. out_dtype : str The output data dtype. trans_b : bool Whether the input matrix B is transposed. Returns ------- ret : Dict[str, str] A group of tensor intrinsics. """ assert load_scope in ["shared", "shared.dyn"] assert store_scope in ["global", "shared", "shared.dyn"] assert in_dtype in ["float16", "int8"] assert out_dtype in ["float16", "float32", "int32"] shape = "16x16x16" in_dtype = "f16" if in_dtype == "float16" else "s8" out_dtype = "f16" if out_dtype == "float16" else "f32" if out_dtype == "float32" else "s32" # convert "shared.dyn" to "shared_dyn" load_scope = load_scope.replace(".", "_") store_scope = store_scope.replace(".", "_") trans_a = "" trans_b = "_trans" if trans_b else "" # e.g. wmma_load_16x16x16_f16_a_shared load_a_intrin = f"wmma_load_{shape}_{in_dtype}_a{trans_a}_{load_scope}" # e.g. wmma_load_16x16x16_f16_b_trans_shared_dyn load_b_intrin = f"wmma_load_{shape}_{in_dtype}_b{trans_b}_{load_scope}" # e.g. wmma_sync_16x16x16_f16f16f32_trans compute_intrin = f"wmma_sync_{shape}_{in_dtype}{in_dtype}{out_dtype}{trans_b}" # e.g. wmma_fill_16x16x16_f16 init_intrin = f"wmma_fill_{shape}_{out_dtype}" # e.g. wmma_store_16x16x16_f16_shared_dyn store_intrin = f"wmma_store_{shape}_{out_dtype}_{store_scope}" return { "init": init_intrin, "load_a": load_a_intrin, "load_b": load_b_intrin, "compute": compute_intrin, "store": store_intrin, } ######## MMA intrinsics ######## def get_index_A(elem_offset, stride): i = elem_offset // stride j = elem_offset % stride stride_b = stride // 8 bi = i // 32 bj = j // 8 no = bi * stride_b + bj return no * 8 + (i % 32) // 16 * 4 def get_index_B(elem_offset, stride): i = elem_offset // stride j = elem_offset % stride stride_b = stride // 32 bi = i // 8 bj = j // 32 no = bi * stride_b + bj return no * 8 + (j % 32) // 8 * 2 def get_index_C(elem_offset, stride): i = elem_offset // stride j = elem_offset % stride stride_b = stride // 8 bi = i // 8 bj = j // 8 return ((bi // 2) * 2 * stride_b + bi % 2 + bj * 2) * 2 def get_mma_init_intrin( m_dim: int, n_dim: int, k_dim: int, dtype: str ) -> tuple[PrimFunc, PrimFunc]: """Generator of mma init intrins""" del k_dim # unused zero = IntImm("int32", 0).astype(dtype) assert m_dim % 8 == 0 and n_dim % 4 == 0, "m_dim and n_dim must be multiple of 8 and 4" assert dtype in ["float16", "float32"] assert n_dim // 4 * int(dtype[-2:]) <= 128, "n_dim vectorize failed" @T.prim_func(s_tir=True) def mma_init_desc(c: T.handle) -> None: dst = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC" ) with T.sblock("root"): T.reads() T.writes(dst[0:m_dim, 0:n_dim]) for i, j in T.grid(m_dim, n_dim): with T.sblock("init"): vi, vj = T.axis.remap("SS", [i, j]) dst[vi, vj] = zero @T.prim_func(s_tir=True) def mma_init_impl(c: T.handle) -> None: dst = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC" ) with T.sblock("root"): T.reads() T.writes(dst[0:m_dim, 0:n_dim]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): for b in range(m_dim // 8): for v in T.vectorized(n_dim // 4): dst[b * 8 + tx // 4, (tx % 4) * (n_dim // 4) + v] = zero return mma_init_desc, mma_init_impl def get_mma_load_intrin( m_dim: int, n_dim: int, k_dim: int, dtype: str, shared_scope: str, is_b: bool, is_col_major: bool, ) -> tuple[PrimFunc, PrimFunc]: """Generator of mma ldmatrix intrins""" mma_fragment_scope = f"m16n8k8.matrix{'B' if is_b else 'A'}" frag_m, frag_n = (k_dim, n_dim) if is_b else (m_dim, k_dim) trans = (not is_col_major) if is_b else is_col_major if is_col_major: frag_m, frag_n = frag_n, frag_m get_index = get_index_B if is_b else get_index_A get_tx_index = ( (lambda tx, s0: (tx % 8) * s0 + (tx // 8) * 8) if trans else (lambda tx, s0: tx * s0) ) @T.prim_func(s_tir=True) def mma_load_desc(a: T.handle, c: T.handle) -> None: src = T.match_buffer( a, (frag_m, frag_n), dtype, align=64, offset_factor=1, scope=shared_scope ) dst = T.match_buffer( c, (frag_m, frag_n), dtype, align=64, offset_factor=1, scope=mma_fragment_scope ) with T.sblock("root"): T.reads(src[0:frag_m, 0:frag_n]) T.writes(dst[0:frag_m, 0:frag_n]) for i, j in T.grid(frag_m, frag_n): with T.sblock("root"): vi, vj = T.axis.remap("SS", [i, j]) dst[vi, vj] = src[vi, vj] @T.prim_func(s_tir=True) def mma_load_impl(a: T.handle, c: T.handle) -> None: s0 = T.int32() s1 = T.int32() src = T.match_buffer( a, (frag_m, frag_n), dtype, align=64, offset_factor=1, scope=shared_scope, strides=[s0, s1], ) d0 = T.int32() d1 = T.int32() dst = T.match_buffer( c, (frag_m, frag_n), dtype, align=64, offset_factor=1, scope=mma_fragment_scope, strides=[d0, d1], ) with T.sblock("root"): T.reads(src[0:frag_m, 0:frag_n]) T.writes(dst[0:frag_m, 0:frag_n]) for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"): T.evaluate( T.ptx.ldmatrix_legacy( trans, 4, # Always load 4 matrices ".b16", dst.data, get_index(dst.elem_offset, d0), src.access_ptr("r"), get_tx_index(tx, s0), dtype=dtype, ) ) return mma_load_desc, mma_load_impl def get_mma_sync_intrin( m_dim: int, n_dim: int, k_dim: int, in_dtype: str, out_dtype: str, b_transposed: bool ) -> tuple[PrimFunc, PrimFunc]: """Generator of mma sync intrins""" def maybe_cast(v): if in_dtype != out_dtype: return Cast(out_dtype, v) return v def maybe_swap(i, j): if b_transposed: return j, i return i, j B_shape_0, B_shape_1 = maybe_swap(k_dim, n_dim) @T.prim_func(s_tir=True) def mma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None: A = T.match_buffer( a, (m_dim, k_dim), in_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixA" ) B = T.match_buffer( b, (B_shape_0, B_shape_1), in_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixB" ) C = T.match_buffer( c, (m_dim, n_dim), out_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC" ) with T.sblock("root"): T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:B_shape_0, 0:B_shape_1]) T.writes(C[0:m_dim, 0:n_dim]) for i, j, k in T.grid(m_dim, n_dim, k_dim): with T.sblock("m16n8k8_sync"): vi, vj, vk = T.axis.remap("SSR", [i, j, k]) B_index_0, B_index_1 = T.meta_var(maybe_swap(vk, vj)) C[vi, vj] = C[vi, vj] + maybe_cast(A[vi, vk]) * maybe_cast( B[B_index_0, B_index_1] ) @T.prim_func(s_tir=True) def mma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None: a0 = T.int32() a1 = T.int32() A = T.match_buffer( a, (m_dim, k_dim), in_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixA", strides=[a0, a1], ) b0 = T.int32() b1 = T.int32() B = T.match_buffer( b, (B_shape_0, B_shape_1), in_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixB", strides=[b0, b1], ) c0 = T.int32() c1 = T.int32() C = T.match_buffer( c, (m_dim, n_dim), out_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC", strides=[c0, c1], ) with T.sblock("root"): T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:B_shape_0, 0:B_shape_1]) T.writes(C[0:m_dim, 0:n_dim]) T.evaluate( T.ptx.mma.legacy( f"m{m_dim}n{n_dim}k{k_dim}", "row", "col", in_dtype, in_dtype, out_dtype, A.data, get_index_A(A.elem_offset, a0), B.data, get_index_B(B.elem_offset, b0), C.data, get_index_C(C.elem_offset, c0), False, dtype=out_dtype, ) ) return mma_sync_desc, mma_sync_impl def get_mma_store_dummy_intrin( m_dim: int, n_dim: int, k_dim: int, dtype: str ) -> tuple[PrimFunc, PrimFunc]: """Disable mma store intrin for now.""" del k_dim # unused @T.prim_func(s_tir=True) def mma_store_desc(a: T.handle, c: T.handle) -> None: src = T.match_buffer( a, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC" ) dst = T.match_buffer( c, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="shared.dyn" ) with T.sblock("root"): T.reads(src[0:m_dim, 0:n_dim]) T.writes(dst[0:m_dim, 0:n_dim]) for i, j in T.grid(m_dim, n_dim): with T.sblock("m16n8k8_store"): vi, vj = T.axis.remap("SS", [i, j]) dst[vi, vj] = src[vi, vj] return mma_store_desc, mma_store_desc TensorIntrin.register("mma_init_m16n8k8_f16", *get_mma_init_intrin(16, 8, 8, "float16")) TensorIntrin.register("mma_init_m16n8k8_f32", *get_mma_init_intrin(16, 8, 8, "float32")) TensorIntrin.register( "mma_load_m16n8k8_f16_A_shared_dyn", *get_mma_load_intrin(32, 32, 8, "float16", "shared.dyn", False, False), ) TensorIntrin.register( "mma_load_m16n8k8_f16_B_shared_dyn", *get_mma_load_intrin(32, 32, 8, "float16", "shared.dyn", True, False), ) TensorIntrin.register( "mma_sync_m16n8k8_f16f16f16", *get_mma_sync_intrin(16, 8, 8, "float16", "float16", False) ) TensorIntrin.register( "mma_sync_m16n8k8_f16f16f32", *get_mma_sync_intrin(16, 8, 8, "float16", "float32", False) ) TensorIntrin.register( "mma_store_m16n8k8_f16_global", *get_mma_store_dummy_intrin(16, 8, 8, "float16") ) TensorIntrin.register( "mma_store_m16n8k8_f32_global", *get_mma_store_dummy_intrin(16, 8, 8, "float32") ) @register_global_func("tirx.index_map_m16n8k8.matrixC") def index_map_m16n8k8_matrixC(ind): i, j = ind[0], ind[1] return convert([(i // 8) // 2, j // 8, (i // 8) % 2, (j % 8) % 2])