chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,220 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user