wehub-resource-sync
1758 lines
58 KiB
Python
1758 lines
58 KiB
Python
# 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])
|