chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,477 @@
|
||||
# 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
|
||||
"""Intrinsics for AMDGPU tensorization."""
|
||||
|
||||
from tvm.runtime import convert
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx.expr import Cast, IntImm
|
||||
|
||||
from .. import TensorIntrin
|
||||
from .dot_product_common import get_dp4a_intrin
|
||||
|
||||
lift = convert
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def sdot4(
|
||||
A: T.Buffer((4,), "int8", offset_factor=1, align=4, scope="shared"),
|
||||
B: T.Buffer((4,), "int8", offset_factor=1, align=4, scope="shared"),
|
||||
C: T.Buffer((1,), "int32", offset_factor=1, align=4, scope="local"),
|
||||
) -> None:
|
||||
with T.sblock("root"):
|
||||
T.reads(C[0], A[0:4], B[0:4])
|
||||
T.writes(C[0])
|
||||
|
||||
C[0] += T.call_llvm_pure_intrin(
|
||||
T.llvm_lookup_intrinsic_id("llvm.amdgcn.sdot4"),
|
||||
T.reinterpret(A.vload([0], "int8x4"), dtype="int32"),
|
||||
T.reinterpret(B.vload([0], "int8x4"), dtype="int32"),
|
||||
T.int32(0),
|
||||
T.bool(1),
|
||||
dtype="int32",
|
||||
)
|
||||
|
||||
|
||||
AMDGPU_SDOT4_INTRIN = "sdot4"
|
||||
|
||||
dp4a_desc, _ = get_dp4a_intrin("int8", "int8", "int32")
|
||||
TensorIntrin.register(AMDGPU_SDOT4_INTRIN, dp4a_desc, sdot4)
|
||||
|
||||
WARP_SIZE = 64
|
||||
M_DIM = 16
|
||||
N_DIM = 16
|
||||
|
||||
|
||||
def shared_16x4_to_local_64x1_layout_A(i, j):
|
||||
thread_id = j * 16 + i
|
||||
return thread_id, convert(0)
|
||||
|
||||
|
||||
def thread_id_shared_access_64x1_to_16x4_layout_A(thread_id, local_id):
|
||||
i = thread_id % 16
|
||||
j = thread_id // 16 + local_id
|
||||
return i, j
|
||||
|
||||
|
||||
def shared_4x16_to_local_64x1_layout_B(i, j):
|
||||
thread_id = i * 16 + j
|
||||
return thread_id, convert(0)
|
||||
|
||||
|
||||
def thread_id_shared_access_64x1_to_4x16_layout_B(thread_id, local_id):
|
||||
i = thread_id // 16
|
||||
j = thread_id % 16 + local_id
|
||||
return i, j
|
||||
|
||||
|
||||
def shared_16x16_to_local_64x4_layout_C(i, j):
|
||||
thread_id = j + (i // 4) * 16
|
||||
local = i % 4
|
||||
return thread_id, local
|
||||
|
||||
|
||||
def thread_id_shared_access_64x4_to_16x16_layout_A(thread_id, local_id):
|
||||
i = thread_id % 16
|
||||
j = (thread_id // 16) * 4 + local_id
|
||||
return i, j
|
||||
|
||||
|
||||
def shared_16x16_to_local_64x4_layout_A(i, j):
|
||||
thread_id = i + 16 * (j // 4)
|
||||
local = j % 4
|
||||
return thread_id, local
|
||||
|
||||
|
||||
def thread_id_shared_access_64x4_to_16x16_layout_B(thread_id, local_id):
|
||||
i = local_id + (thread_id // 16) * 4
|
||||
j = thread_id % 16
|
||||
return i, j
|
||||
|
||||
|
||||
def shared_16x16_to_local_64x4_layout_B(i, j):
|
||||
thread_id = j + (i // 4) * 16
|
||||
local = i % 4
|
||||
return thread_id, local
|
||||
|
||||
|
||||
def thread_id_shared_access_64x4_to_16x16_layout_C(thread_id, local_id):
|
||||
i = local_id + (thread_id // 16) * 4
|
||||
j = thread_id % 16
|
||||
return i, j
|
||||
|
||||
|
||||
def get_mma_fill_intrin(dtype, local_size):
|
||||
zero = IntImm("int32", 0).astype(dtype)
|
||||
|
||||
# Assume M = N = 16
|
||||
index_map = shared_16x16_to_local_64x4_layout_C
|
||||
|
||||
@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])
|
||||
tx = T.env_thread("threadIdx.x")
|
||||
T.launch_thread(tx, WARP_SIZE)
|
||||
for local_id in T.serial(0, local_size):
|
||||
C_warp[tx, local_id] = zero
|
||||
|
||||
return mma_fill_desc, mma_fill_impl
|
||||
|
||||
|
||||
def get_mfma_load_intrin(
|
||||
k_dim=4,
|
||||
dtype="float32",
|
||||
scope="shared",
|
||||
is_b=False,
|
||||
transposed=False,
|
||||
):
|
||||
local_size = (M_DIM * k_dim) // WARP_SIZE if not is_b else (N_DIM * k_dim) // WARP_SIZE
|
||||
memory_shape = (M_DIM, k_dim)
|
||||
if is_b:
|
||||
memory_shape = (N_DIM, k_dim) if transposed else (k_dim, N_DIM)
|
||||
|
||||
row_dim, col_dim = memory_shape
|
||||
|
||||
if k_dim == 4:
|
||||
index_map = shared_16x4_to_local_64x1_layout_A
|
||||
reverse_index_map = thread_id_shared_access_64x1_to_16x4_layout_A
|
||||
if is_b:
|
||||
index_map = (
|
||||
shared_16x4_to_local_64x1_layout_A
|
||||
if transposed
|
||||
else shared_4x16_to_local_64x1_layout_B
|
||||
)
|
||||
reverse_index_map = (
|
||||
thread_id_shared_access_64x1_to_16x4_layout_A
|
||||
if transposed
|
||||
else thread_id_shared_access_64x1_to_4x16_layout_B
|
||||
)
|
||||
elif k_dim == 16:
|
||||
index_map = shared_16x16_to_local_64x4_layout_A
|
||||
reverse_index_map = thread_id_shared_access_64x4_to_16x16_layout_A
|
||||
|
||||
if is_b:
|
||||
index_map = (
|
||||
shared_16x16_to_local_64x4_layout_A
|
||||
if transposed
|
||||
else shared_16x16_to_local_64x4_layout_B
|
||||
)
|
||||
reverse_index_map = (
|
||||
thread_id_shared_access_64x4_to_16x16_layout_A
|
||||
if transposed
|
||||
else thread_id_shared_access_64x4_to_16x16_layout_B
|
||||
)
|
||||
else:
|
||||
raise ValueError("k_dim must be 4 or 16 currently")
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_load_desc(reg_handle: T.handle, memory_handle: T.handle) -> None:
|
||||
memory = T.match_buffer(
|
||||
memory_handle,
|
||||
memory_shape,
|
||||
dtype,
|
||||
offset_factor=1,
|
||||
scope=scope,
|
||||
)
|
||||
reg = T.match_buffer(
|
||||
reg_handle, (WARP_SIZE, local_size), dtype, offset_factor=1, scope="warp"
|
||||
)
|
||||
|
||||
with T.sblock("root"):
|
||||
T.reads(memory[0:row_dim, 0:col_dim])
|
||||
T.writes(reg[0:WARP_SIZE, 0:local_size])
|
||||
|
||||
for ax0, ax1 in T.grid(row_dim, col_dim):
|
||||
with T.sblock("memory_reg"):
|
||||
v0, v1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(memory[v0, v1])
|
||||
|
||||
thread_id, local_id = T.meta_var(index_map(v0, v1))
|
||||
T.writes(reg[thread_id, local_id])
|
||||
reg[thread_id, local_id] = memory[v0, v1]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_load_impl(reg_handle: T.handle, memory_handle: T.handle) -> None:
|
||||
s0 = T.int32()
|
||||
s1 = T.int32()
|
||||
|
||||
memory = T.match_buffer(
|
||||
memory_handle,
|
||||
memory_shape,
|
||||
dtype,
|
||||
align=64,
|
||||
offset_factor=1,
|
||||
scope=scope,
|
||||
strides=[s0, s1],
|
||||
)
|
||||
reg = T.match_buffer(
|
||||
reg_handle, (WARP_SIZE, local_size), dtype, align=64, offset_factor=1, scope="warp"
|
||||
)
|
||||
|
||||
with T.sblock("root"):
|
||||
T.reads(memory[0:row_dim, 0:col_dim])
|
||||
T.writes(reg[0:WARP_SIZE, 0:local_size])
|
||||
tx = T.env_thread("threadIdx.x")
|
||||
for local_id in T.serial(0, local_size):
|
||||
row, col = T.meta_var(reverse_index_map(tx, local_id))
|
||||
T.launch_thread(tx, WARP_SIZE)
|
||||
reg[tx, local_id] = memory[row, col]
|
||||
|
||||
return mfma_load_desc, mfma_load_impl
|
||||
|
||||
|
||||
def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed=False):
|
||||
local_size = (M_DIM * k_dim) // WARP_SIZE
|
||||
local_size_out = (M_DIM * N_DIM) // WARP_SIZE
|
||||
if k_dim == 4:
|
||||
index_map_A = shared_16x4_to_local_64x1_layout_A
|
||||
index_map_B = shared_4x16_to_local_64x1_layout_B
|
||||
index_map_C = shared_16x16_to_local_64x4_layout_C
|
||||
elif k_dim == 16:
|
||||
index_map_A = shared_16x16_to_local_64x4_layout_A
|
||||
index_map_B = shared_16x16_to_local_64x4_layout_B
|
||||
index_map_C = shared_16x16_to_local_64x4_layout_C
|
||||
else:
|
||||
raise ValueError("k_dim must be 4 or 16 currently")
|
||||
|
||||
out_dtype_abbrv = {"float16": "f16", "float32": "f32", "int8": "i8", "int32": "i32"}[out_dtype]
|
||||
|
||||
in_dtype_abbrv = {"float16": "f16", "float32": "f32", "int8": "i8", "int32": "i32"}[in_dtype]
|
||||
|
||||
mfma_intrin = f"llvm.amdgcn.mfma.{out_dtype_abbrv}.{M_DIM}x{N_DIM}x{k_dim}{in_dtype_abbrv}"
|
||||
|
||||
def maybe_cast(v):
|
||||
if out_dtype != in_dtype:
|
||||
return Cast(out_dtype, v)
|
||||
return v
|
||||
|
||||
def maybe_swap(i, j):
|
||||
if b_transposed:
|
||||
return j, i
|
||||
return i, j
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
|
||||
B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
|
||||
C = T.match_buffer(c, (WARP_SIZE, local_size_out), out_dtype, offset_factor=1, 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])
|
||||
b_row_ind, b_col_ind = T.meta_var(maybe_swap(vk, vj))
|
||||
|
||||
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(vi, vk))
|
||||
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] += maybe_cast(
|
||||
A[thread_id_A, local_id_A]
|
||||
) * maybe_cast(B[thread_id_B, local_id_B])
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_sync_impl_float(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
|
||||
B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
|
||||
C = T.match_buffer(c, (WARP_SIZE, local_size_out), out_dtype, offset_factor=1, scope="warp")
|
||||
|
||||
with T.sblock("root"):
|
||||
T.reads(
|
||||
A[0:WARP_SIZE, 0:local_size],
|
||||
B[0:WARP_SIZE, 0:local_size],
|
||||
C[0:WARP_SIZE, 0:local_size_out],
|
||||
)
|
||||
T.writes(C[0:WARP_SIZE, 0:local_size_out])
|
||||
tx = T.env_thread("threadIdx.x")
|
||||
T.launch_thread(tx, WARP_SIZE)
|
||||
C[tx, 0:local_size_out] = T.call_llvm_pure_intrin(
|
||||
T.llvm_lookup_intrinsic_id(mfma_intrin),
|
||||
A[tx, 0:local_size],
|
||||
B[tx, 0:local_size],
|
||||
C[tx, 0:local_size_out],
|
||||
T.int32(0),
|
||||
T.int32(0),
|
||||
T.int32(0),
|
||||
dtype=f"{out_dtype}x4",
|
||||
)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_sync_impl_integer(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
|
||||
B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
|
||||
C = T.match_buffer(c, (WARP_SIZE, local_size_out), out_dtype, offset_factor=1, scope="warp")
|
||||
|
||||
with T.sblock("root"):
|
||||
T.reads(
|
||||
A[0:WARP_SIZE, 0:local_size],
|
||||
B[0:WARP_SIZE, 0:local_size],
|
||||
C[0:WARP_SIZE, 0:local_size_out],
|
||||
)
|
||||
T.writes(C[0:WARP_SIZE, 0:local_size_out])
|
||||
tx = T.env_thread("threadIdx.x")
|
||||
T.launch_thread(tx, WARP_SIZE)
|
||||
|
||||
C[tx, 0:local_size_out] = T.call_llvm_pure_intrin(
|
||||
T.llvm_lookup_intrinsic_id(mfma_intrin),
|
||||
T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size]),
|
||||
T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size]),
|
||||
C[tx, 0:local_size_out],
|
||||
T.int32(0),
|
||||
T.int32(0),
|
||||
T.int32(0),
|
||||
dtype=f"{out_dtype}x4",
|
||||
)
|
||||
|
||||
return (
|
||||
(mfma_sync_desc, mfma_sync_impl_integer)
|
||||
if in_dtype == "int8"
|
||||
else (mfma_sync_desc, mfma_sync_impl_float)
|
||||
)
|
||||
|
||||
|
||||
def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
|
||||
index_map = shared_16x16_to_local_64x4_layout_C
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_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]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mfma_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])
|
||||
tx = T.env_thread("threadIdx.x")
|
||||
T.launch_thread(tx, WARP_SIZE)
|
||||
for i in range(local_size):
|
||||
C[((tx // 16) * 4) + i, (tx % 16)] = C_warp[tx, i]
|
||||
|
||||
return mfma_store_desc, mfma_store_impl
|
||||
|
||||
|
||||
ROCM_MFMA_fill_16x16_f32_INTRIN = "ROCM_mfma_fill_16x16_f32"
|
||||
TensorIntrin.register(ROCM_MFMA_fill_16x16_f32_INTRIN, *get_mma_fill_intrin("float32", 4))
|
||||
|
||||
ROCM_MFMA_fill_16x16_i32_INTRIN = "ROCM_mfma_fill_16x16_i32"
|
||||
TensorIntrin.register(ROCM_MFMA_fill_16x16_i32_INTRIN, *get_mma_fill_intrin("int", 4))
|
||||
|
||||
ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN = "rocm_mfma_load_16x16_a_shared_s8"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN, *get_mfma_load_intrin(16, "int8", "shared")
|
||||
)
|
||||
ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN = "rocm_mfma_load_b_16x16_shared_s8"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN, *get_mfma_load_intrin(16, "int8", "shared", is_b=True)
|
||||
)
|
||||
|
||||
ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN = "rocm_mfma_load_16x16_a_shared_f16"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN, *get_mfma_load_intrin(16, "float16", "shared")
|
||||
)
|
||||
ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN = "rocm_mfma_load_b_16x16_shared_f16"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN,
|
||||
*get_mfma_load_intrin(16, "float16", "shared", is_b=True),
|
||||
)
|
||||
|
||||
ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN = "rocm_mfma_load_16x4_a_shared_f32"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN, *get_mfma_load_intrin(4, "float32", "shared")
|
||||
)
|
||||
ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN = "rocm_mfma_load_b_16x4_shared_f32"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN,
|
||||
*get_mfma_load_intrin(4, "float32", "shared", is_b=True),
|
||||
)
|
||||
|
||||
|
||||
ROCM_MFMA_f32f32f32_INTRIN = "rocm_mfma_f32f32f32"
|
||||
TensorIntrin.register(ROCM_MFMA_f32f32f32_INTRIN, *get_mfma_intrin(4, "float32", "float32"))
|
||||
|
||||
ROCM_MFMA_f16f16f32_INTRIN = "rocm_mfma_f16f16f32"
|
||||
TensorIntrin.register(ROCM_MFMA_f16f16f32_INTRIN, *get_mfma_intrin(16, "float16", "float32"))
|
||||
|
||||
ROCM_MFMA_s8s8s32_INTRIN = "rocm_mfma_s8s8s32"
|
||||
TensorIntrin.register(ROCM_MFMA_s8s8s32_INTRIN, *get_mfma_intrin(16, "int8", "int32"))
|
||||
|
||||
ROCM_MFMA_STORE_16x16_s32_INTRIN = "rocm_mfma_store_16x16_s32"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_STORE_16x16_s32_INTRIN, *get_mfma_store_intrin(4, "int32", "global")
|
||||
)
|
||||
|
||||
ROCM_MFMA_STORE_16x16_f32_INTRIN = "rocm_mfma_store_16x16_f32"
|
||||
TensorIntrin.register(
|
||||
ROCM_MFMA_STORE_16x16_f32_INTRIN, *get_mfma_store_intrin(4, "float32", "global")
|
||||
)
|
||||
Reference in New Issue
Block a user