wehub-resource-sync
221 lines
8.2 KiB
Python
221 lines
8.2 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=missing-function-docstring
|
|
"""Codegen tests for Ampere (sm_80) warp-level ``mma.sync`` tensor cores.
|
|
|
|
These exercise the ``T.ptx.mma`` intrinsic directly (not via the gemm
|
|
dispatch). ``ptx.mma`` takes one pointer per 32-bit register for each operand
|
|
(``d_ptrs`` / ``a_ptrs`` / ``b_ptrs`` / ``c_ptrs``), enumerated in the fixed
|
|
PTX register order, so the b32 registers may be scattered in the register file
|
|
while the two packed fp16/bf16 within a b32 stay contiguous. For m16n8k{8,16}
|
|
with f32 accumulation the per-lane register counts are:
|
|
|
|
A: 2 inputs per b32 -> k16: 4 b32 (regs 0,2,4,6); k8: 2 b32 (regs 0,2)
|
|
B: 2 inputs per b32 -> k16: 2 b32 (regs 0,2); k8: 1 b32 (reg 0)
|
|
D/C: 4 f32 accumulator registers (0,1,2,3)
|
|
"""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
def _get_source(func: tvm.tirx.PrimFunc):
|
|
target = tvm.target.Target("cuda")
|
|
mod = tvm.IRModule({"main": func})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
return src, mod
|
|
|
|
|
|
def _np_in(dtype):
|
|
if dtype == "bfloat16":
|
|
return __import__("ml_dtypes").bfloat16
|
|
return np.float16
|
|
|
|
|
|
def _run_mma(mod, K, no_c_ptr, np_in):
|
|
"""Run an m16n8kK mma kernel and check D == A @ B (+ C) against numpy."""
|
|
np.random.seed(0)
|
|
A_np = np.random.randn(16, K).astype(np_in)
|
|
B_np = np.random.randn(K, 8).astype(np_in)
|
|
C_np = np.random.randn(16, 8).astype(np.float32)
|
|
ref = A_np.astype(np.float32) @ B_np.astype(np.float32)
|
|
if not no_c_ptr:
|
|
ref = ref + C_np
|
|
|
|
def run_and_check():
|
|
dev = tvm.device("cuda")
|
|
D = tvm.runtime.tensor(np.zeros((16, 8), np.float32), device=dev)
|
|
A = tvm.runtime.tensor(A_np, device=dev)
|
|
B = tvm.runtime.tensor(B_np, device=dev)
|
|
C = tvm.runtime.tensor(C_np, device=dev)
|
|
mod(D, A, B, C)
|
|
np.testing.assert_allclose(D.numpy(), ref, atol=1e-2, rtol=1e-2)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
@pytest.mark.parametrize("a_type", ["float16", "bfloat16"])
|
|
@pytest.mark.parametrize("no_c_ptr", [False, True])
|
|
def test_ptx_mma_m16n8k16(a_type, no_c_ptr):
|
|
"""m16n8k16 row.col mma, f32 accumulate: A is 16x16 (4 b32/lane), B is 16x8
|
|
as [K, N] (2 b32/lane), D/C is 16x8 (4 f32/lane)."""
|
|
if a_type == "bfloat16":
|
|
pytest.importorskip("ml_dtypes")
|
|
b_type = a_type
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(
|
|
D: T.Buffer((16, 8), "float32"),
|
|
A: T.Buffer((16, 16), a_type),
|
|
B: T.Buffer((16, 8), b_type),
|
|
C: T.Buffer((16, 8), "float32"),
|
|
):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([32])
|
|
D_local = T.alloc_local([4], "float32")
|
|
A_local = T.alloc_local([8], a_type)
|
|
B_local = T.alloc_local([4], b_type)
|
|
C_local = T.alloc_local([4], "float32")
|
|
|
|
@T.inline
|
|
def G2L(buf_local, buf_global, block_8x8, mode="row"):
|
|
if mode == "row":
|
|
for i in range(block_8x8):
|
|
row = T.meta_var(i % 2 * 8 + tx // 4)
|
|
col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
|
|
for j in range(2):
|
|
buf_local[i * 2 + j] = buf_global[row, col + j]
|
|
elif mode == "col":
|
|
for i in range(block_8x8):
|
|
row = T.meta_var(i % 2 * 8 + (tx % 4) * 2)
|
|
col = T.meta_var(i // 2 * 8 + tx // 4)
|
|
for j in range(2):
|
|
buf_local[i * 2 + j] = buf_global[row + j, col]
|
|
|
|
G2L(D_local, D, 2)
|
|
G2L(A_local, A, 4)
|
|
G2L(B_local, B, 2, "col")
|
|
G2L(C_local, C, 2)
|
|
|
|
# One pointer per b32 register, in PTX order: A=4, B=2, D/C=4.
|
|
d_ptrs = [D_local.ptr_to([i]) for i in range(4)]
|
|
a_ptrs = [A_local.ptr_to([2 * i]) for i in range(4)]
|
|
b_ptrs = [B_local.ptr_to([2 * i]) for i in range(2)]
|
|
if no_c_ptr:
|
|
T.ptx.mma("m16n8k16", "row", "col", "float32", a_type, b_type, "float32",
|
|
d_ptrs, a_ptrs, b_ptrs)
|
|
else:
|
|
c_ptrs = [C_local.ptr_to([i]) for i in range(4)]
|
|
T.ptx.mma("m16n8k16", "row", "col", "float32", a_type, b_type, "float32",
|
|
d_ptrs, a_ptrs, b_ptrs, c_ptrs)
|
|
|
|
for i in range(2):
|
|
row = T.meta_var(i % 2 * 8 + tx // 4)
|
|
col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
|
|
for j in range(2):
|
|
D[row, col + j] = D_local[i * 2 + j]
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(main)
|
|
assert "mma.sync.aligned.m16n8k16.row.col" in src
|
|
_run_mma(mod, 16, no_c_ptr, _np_in(a_type))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
@pytest.mark.parametrize("a_type", ["float16", "bfloat16"])
|
|
@pytest.mark.parametrize("no_c_ptr", [False, True])
|
|
def test_ptx_mma_m16n8k8(a_type, no_c_ptr):
|
|
"""m16n8k8 row.col mma, f32 accumulate: A is 16x8 (2 b32/lane), B is 8x8
|
|
as [K, N] (1 b32/lane), D/C is 16x8 (4 f32/lane)."""
|
|
if a_type == "bfloat16":
|
|
pytest.importorskip("ml_dtypes")
|
|
b_type = a_type
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(
|
|
D: T.Buffer((16, 8), "float32"),
|
|
A: T.Buffer((16, 8), a_type),
|
|
B: T.Buffer((8, 8), b_type),
|
|
C: T.Buffer((16, 8), "float32"),
|
|
):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([32])
|
|
D_local = T.alloc_local([4], "float32")
|
|
A_local = T.alloc_local([4], a_type)
|
|
B_local = T.alloc_local([2], b_type)
|
|
C_local = T.alloc_local([4], "float32")
|
|
|
|
@T.inline
|
|
def G2L(buf_local, buf_global, block_8x8, mode="row"):
|
|
if mode == "row":
|
|
for i in range(block_8x8):
|
|
row = T.meta_var(i % 2 * 8 + tx // 4)
|
|
col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
|
|
for j in range(2):
|
|
buf_local[i * 2 + j] = buf_global[row, col + j]
|
|
elif mode == "col":
|
|
for i in range(block_8x8):
|
|
row = T.meta_var(i % 2 * 8 + (tx % 4) * 2)
|
|
col = T.meta_var(i // 2 * 8 + tx // 4)
|
|
for j in range(2):
|
|
buf_local[i * 2 + j] = buf_global[row + j, col]
|
|
|
|
G2L(D_local, D, 2)
|
|
G2L(A_local, A, 2)
|
|
G2L(B_local, B, 1, "col")
|
|
G2L(C_local, C, 2)
|
|
|
|
# One pointer per b32 register, in PTX order: A=2, B=1, D/C=4.
|
|
d_ptrs = [D_local.ptr_to([i]) for i in range(4)]
|
|
a_ptrs = [A_local.ptr_to([2 * i]) for i in range(2)]
|
|
b_ptrs = [B_local.ptr_to([0])]
|
|
if no_c_ptr:
|
|
T.ptx.mma("m16n8k8", "row", "col", "float32", a_type, b_type, "float32",
|
|
d_ptrs, a_ptrs, b_ptrs)
|
|
else:
|
|
c_ptrs = [C_local.ptr_to([i]) for i in range(4)]
|
|
T.ptx.mma("m16n8k8", "row", "col", "float32", a_type, b_type, "float32",
|
|
d_ptrs, a_ptrs, b_ptrs, c_ptrs)
|
|
|
|
for i in range(2):
|
|
row = T.meta_var(i % 2 * 8 + tx // 4)
|
|
col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
|
|
for j in range(2):
|
|
D[row, col + j] = D_local[i * 2 + j]
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(main)
|
|
assert "mma.sync.aligned.m16n8k8.row.col" in src
|
|
_run_mma(mod, 8, no_c_ptr, _np_in(a_type))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|