wehub-resource-sync
1224 lines
47 KiB
Python
1224 lines
47 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
|
|
import math
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
from tvm.tirx import Buffer
|
|
|
|
|
|
def _get_source(func: tvm.tirx.PrimFunc) -> tuple[str, tvm.IRModule]:
|
|
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 _run_tensormap_encode(shape, dtype, encode_args):
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shape, dtype=dtype, align=32)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, dtype, len(shape), A.data, *encode_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for threadIdx in T.thread_binding(1, thread="threadIdx.x"):
|
|
T.evaluate(blockIdx + threadIdx)
|
|
# fmt: on
|
|
|
|
target = tvm.target.Target("cuda")
|
|
mod = tvm.IRModule({"main": main})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
def run_and_check():
|
|
A = tvm.runtime.tensor(np.zeros(shape, dtype=dtype), device=tvm.cuda(0))
|
|
mod(A)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize("inc", [False, True])
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_ptx_setmaxnreg(inc):
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tid = T.thread_id([128])
|
|
T.ptx.setmaxnreg(inc, 32)
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
assert "setmaxnreg" in src
|
|
if inc:
|
|
assert "inc" in src
|
|
else:
|
|
assert "dec" in src
|
|
|
|
|
|
@pytest.mark.parametrize("trans", [False, True])
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_stmatrix_sync_aligned(trans):
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer((16, 16), "float16")):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([32])
|
|
A_smem = T.alloc_buffer((16, 16), "float16", scope="shared", align=16)
|
|
reg = T.alloc_buffer((8,), "float16", scope="local")
|
|
for i in range(8):
|
|
reg[i] = tx * 8 + i
|
|
T.ptx.stmatrix(
|
|
trans, 4, ".b16",
|
|
A_smem.ptr_to([tx % 16, tx // 16 * 8]),
|
|
reg.ptr_to([0]), reg.ptr_to([2]), reg.ptr_to([4]), reg.ptr_to([6]),
|
|
)
|
|
if tx == 0:
|
|
for i, j in T.grid(16, 16):
|
|
A[i, j] = A_smem[i, j]
|
|
# fmt: on
|
|
|
|
target = tvm.target.Target("cuda")
|
|
mod = tvm.IRModule({"main": func})
|
|
with target:
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
if not trans:
|
|
assert "stmatrix.sync.aligned.m8n8.x4.shared.b16" in src
|
|
else:
|
|
assert "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16" in src
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A_np = np.zeros((16, 16), dtype="float16")
|
|
A = tvm.runtime.tensor(A_np, device=dev)
|
|
mod(A)
|
|
A_ref = np.zeros((16, 16), dtype="float16")
|
|
for tx in range(32):
|
|
row = tx // 4
|
|
col = tx % 4 * 2
|
|
if not trans:
|
|
A_ref[row, col] = tx * 8
|
|
A_ref[row, col + 1] = tx * 8 + 1
|
|
A_ref[row + 8, col] = tx * 8 + 2
|
|
A_ref[row + 8, col + 1] = tx * 8 + 3
|
|
A_ref[row, col + 8] = tx * 8 + 4
|
|
A_ref[row, col + 9] = tx * 8 + 5
|
|
A_ref[row + 8, col + 8] = tx * 8 + 6
|
|
A_ref[row + 8, col + 9] = tx * 8 + 7
|
|
else:
|
|
A_ref[col, row] = tx * 8
|
|
A_ref[col + 1, row] = tx * 8 + 1
|
|
A_ref[col + 8, row] = tx * 8 + 2
|
|
A_ref[col + 9, row] = tx * 8 + 3
|
|
A_ref[col, row + 8] = tx * 8 + 4
|
|
A_ref[col + 1, row + 8] = tx * 8 + 5
|
|
A_ref[col + 8, row + 8] = tx * 8 + 6
|
|
A_ref[col + 9, row + 8] = tx * 8 + 7
|
|
np.testing.assert_allclose(A.numpy(), A_ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize("trans", [False, True])
|
|
@pytest.mark.parametrize("num", [1, 2, 4])
|
|
def test_ptx_stmatrix(trans, num):
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A: T.Buffer((16, 16), "float16")):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([32])
|
|
A_shared = T.alloc_shared([16, 16], "float16")
|
|
if tx == 0:
|
|
for i, j in T.grid(16, 16):
|
|
A_shared[i, j] = T.float16(0.0)
|
|
T.cuda.cta_sync()
|
|
A_local = T.alloc_local([8], "float16")
|
|
for i in range(8):
|
|
A_local[i] = (i // 2) * 64 + tx * 2 + i % 2
|
|
T.ptx.stmatrix(
|
|
trans, num, ".b16",
|
|
A_shared.ptr_to([tx % 16, tx // 16 * 8]),
|
|
*[A_local.ptr_to([i * 2]) for i in range(num)],
|
|
)
|
|
T.cuda.cta_sync()
|
|
if tx == 0:
|
|
for i, j in T.grid(16, 16):
|
|
A[i, j] = A_shared[i, j]
|
|
# fmt: on
|
|
|
|
target = tvm.target.Target("cuda")
|
|
mod = tvm.IRModule({"main": main})
|
|
with target:
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
A_np = np.zeros((16, 16), dtype="float16")
|
|
A_ref = np.zeros((16, 16), dtype="float16")
|
|
A_full = np.zeros((16, 16), dtype="float16")
|
|
A_full[0:8, 0:8] = np.arange(8 * 8, dtype="float16").reshape((8, 8))
|
|
A_full[8:16, 0:8] = np.arange(8 * 8, 16 * 8, dtype="float16").reshape((8, 8))
|
|
A_full[0:8, 8:16] = np.arange(16 * 8, 24 * 8, dtype="float16").reshape((8, 8))
|
|
A_full[8:16, 8:16] = np.arange(24 * 8, 32 * 8, dtype="float16").reshape((8, 8))
|
|
print(src)
|
|
|
|
if num == 1:
|
|
A_ref[0:8, 0:8] = A_full[0:8, 0:8] if not trans else A_full[0:8, 0:8].T
|
|
elif num == 2:
|
|
A_ref[0:8, 0:8] = A_full[0:8, 0:8] if not trans else A_full[0:8, 0:8].T
|
|
A_ref[8:16, 0:8] = A_full[8:16, 0:8] if not trans else A_full[8:16, 0:8].T
|
|
elif num == 4:
|
|
A_ref[0:8, 0:8] = A_full[0:8, 0:8] if not trans else A_full[0:8, 0:8].T
|
|
A_ref[0:8, 8:16] = A_full[0:8, 8:16] if not trans else A_full[0:8, 8:16].T
|
|
A_ref[8:16, 0:8] = A_full[8:16, 0:8] if not trans else A_full[8:16, 0:8].T
|
|
A_ref[8:16, 8:16] = A_full[8:16, 8:16] if not trans else A_full[8:16, 8:16].T
|
|
|
|
def run_and_check():
|
|
A = tvm.runtime.tensor(A_np, device=tvm.cuda(0))
|
|
mod(A)
|
|
np.testing.assert_allclose(A.numpy(), A_ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize("trans", [False, True])
|
|
@pytest.mark.parametrize("num", [1, 2, 4])
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_ptx_stmatrix_noncontiguous(trans, num):
|
|
"""Symmetric stmatrix API: ``num`` independent src handles.
|
|
|
|
Spaces fragments by 4 fp16 (vs the natural 2 contiguous) so per-src
|
|
pointers are non-contiguous — exercises what the old single-``local_ptr``
|
|
API couldn't express.
|
|
"""
|
|
STRIDE = 4 # 2 fp16 data + 2 fp16 gap per fragment
|
|
LOCAL_SIZE = STRIDE * num
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A: T.Buffer((16, 16), "float16")):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([32])
|
|
A_shared = T.alloc_shared([16, 16], "float16")
|
|
if tx == 0:
|
|
for i, j in T.grid(16, 16):
|
|
A_shared[i, j] = T.float16(0.0)
|
|
T.cuda.cta_sync()
|
|
A_local = T.alloc_local([LOCAL_SIZE], "float16")
|
|
for i in range(num):
|
|
A_local[i * STRIDE + 0] = T.float16(i * 64 + tx * 2 + 0)
|
|
A_local[i * STRIDE + 1] = T.float16(i * 64 + tx * 2 + 1)
|
|
T.ptx.stmatrix(
|
|
trans, num, ".b16",
|
|
A_shared.ptr_to([tx % 16, tx // 16 * 8]),
|
|
*[A_local.ptr_to([i * STRIDE]) for i in range(num)],
|
|
)
|
|
T.cuda.cta_sync()
|
|
if tx == 0:
|
|
for i, j in T.grid(16, 16):
|
|
A[i, j] = A_shared[i, j]
|
|
# fmt: on
|
|
|
|
target = tvm.target.Target("cuda")
|
|
mod = tvm.IRModule({"main": main})
|
|
with target:
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
trans_inst = ".trans" if trans else ""
|
|
assert f"stmatrix.sync.aligned.m8n8.x{num}{trans_inst}.shared.b16" in src
|
|
# num distinct src register loads in the helper body.
|
|
for i in range(num):
|
|
assert f"*(uint32_t*)src{i}" in src
|
|
|
|
A_np = np.zeros((16, 16), dtype="float16")
|
|
A_ref = np.zeros((16, 16), dtype="float16")
|
|
A_full = np.zeros((16, 16), dtype="float16")
|
|
A_full[0:8, 0:8] = np.arange(8 * 8, dtype="float16").reshape((8, 8))
|
|
A_full[8:16, 0:8] = np.arange(8 * 8, 16 * 8, dtype="float16").reshape((8, 8))
|
|
A_full[0:8, 8:16] = np.arange(16 * 8, 24 * 8, dtype="float16").reshape((8, 8))
|
|
A_full[8:16, 8:16] = np.arange(24 * 8, 32 * 8, dtype="float16").reshape((8, 8))
|
|
if num >= 1:
|
|
A_ref[0:8, 0:8] = A_full[0:8, 0:8] if not trans else A_full[0:8, 0:8].T
|
|
if num >= 2:
|
|
A_ref[8:16, 0:8] = A_full[8:16, 0:8] if not trans else A_full[8:16, 0:8].T
|
|
if num >= 4:
|
|
A_ref[0:8, 8:16] = A_full[0:8, 8:16] if not trans else A_full[0:8, 8:16].T
|
|
A_ref[8:16, 8:16] = A_full[8:16, 8:16] if not trans else A_full[8:16, 8:16].T
|
|
|
|
def run_and_check():
|
|
A = tvm.runtime.tensor(A_np, device=tvm.cuda(0))
|
|
mod(A)
|
|
np.testing.assert_allclose(A.numpy(), A_ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_bar_arrive():
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tid = T.thread_id([128])
|
|
T.ptx.bar.arrive(0, 128)
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
assert "tvm_builtin_ptx_bar_arrive(0, 128)" in src
|
|
assert 'bar.arrive %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory"' in src
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_bar_sync():
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tid = T.thread_id([128])
|
|
T.ptx.bar.sync(0, 128)
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
assert "tvm_builtin_ptx_bar_sync(0, 128)" in src
|
|
assert 'bar.sync %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory"' in src
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_fence_mbarrier_init_release_clsuter():
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tid = T.thread_id([128])
|
|
T.ptx.fence.mbarrier_init()
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
assert "fence.mbarrier_init.release.cluster" in src
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_ptx_elect_sync():
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([128])
|
|
if (T.ptx.elect_sync()):
|
|
A[tx] = tx
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
print(src)
|
|
assert "elect.sync %%rx|%%px, %2;" in src
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
@pytest.mark.parametrize("sem,scope", [("sc", "cta"), ("acq_rel", "gpu"), ("sc", "sys")])
|
|
def test_ptx_fence(sem, scope):
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tid = T.thread_id([128])
|
|
T.ptx.fence(sem, scope)
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
assert f"fence.{sem}.{scope};" in src
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_fence_proxy_async():
|
|
# fmt: off
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tid = T.thread_id([128])
|
|
T.ptx.fence.proxy_async("global")
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
|
|
# fmt: on
|
|
|
|
src, mod = _get_source(func)
|
|
assert "fence.proxy.async.global" in src
|
|
assert "fence.proxy.async.shared::cta" in src
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
@pytest.mark.parametrize("dtype", ["float16", "float32", "float8_e4m3fn", "float8_e5m2"])
|
|
@pytest.mark.parametrize(
|
|
"inputs",
|
|
[
|
|
((128,), [128, 128, 1, 0, 0, 0, 0]),
|
|
((16, 16), [16, 16, 16, 16, 16, 1, 1, 0, 0, 0, 0]),
|
|
((16, 64), [64, 16, 64, 64, 16, 1, 1, 0, 0, 0, 0]),
|
|
],
|
|
)
|
|
def test_cp_async_bulk_tensor_global_to_shared_unicast(dtype, inputs):
|
|
import ml_dtypes
|
|
|
|
def get_ir(shape, tma_args):
|
|
t_dtype = tvm.DataType(dtype)
|
|
total_bytes = math.prod(shape) * t_dtype.bits // 8
|
|
coord = [0 for _ in shape]
|
|
tma_args_copy = tma_args.copy()
|
|
for i in range(len(shape) - 1):
|
|
tma_args_copy[len(shape) + i] *= t_dtype.bits // 8
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle, B_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shape, dtype=dtype, align=16)
|
|
B = T.match_buffer(B_ptr, shape, dtype=dtype, align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, dtype, len(shape), A.data, *tma_args_copy) # noqa: E501
|
|
B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, dtype, len(shape), B.data, *tma_args_copy) # noqa: E501
|
|
|
|
T.device_entry()
|
|
for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
bar = T.shared_scalar("uint64")
|
|
phase: T.int32
|
|
A_smem = T.alloc_buffer(shape, dtype, scope="shared", align=128)
|
|
|
|
phase = 0
|
|
if threadIdx == 0:
|
|
T.ptx.mbarrier.init(T.address_of(bar), 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coord) # noqa: E501
|
|
T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes)
|
|
T.ptx.mbarrier.try_wait(T.address_of(bar), phase)
|
|
phase = phase ^ 1
|
|
|
|
T.cuda.cta_sync()
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
|
|
if threadIdx == 0:
|
|
T.ptx.cp_async.bulk.tensor.s2g(len(shape), A_smem.access_ptr("r", offset=0), T.address_of(B_map), "", *coord) # noqa: E501
|
|
T.ptx.cp_async.bulk.commit_group()
|
|
T.ptx.cp_async.bulk.wait_group(0)
|
|
# fmt: on
|
|
|
|
return main
|
|
|
|
target = tvm.target.Target("cuda")
|
|
shape, tma_args = inputs
|
|
mod = tvm.IRModule({"main": get_ir(shape, tma_args)})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
assert "const __grid_constant__ CUtensorMap" in src
|
|
|
|
A_np = np.random.randn(math.prod(shape))
|
|
|
|
def get_np_dtype(dtype):
|
|
if dtype == "float8_e4m3fn":
|
|
return ml_dtypes.float8_e4m3fn
|
|
if dtype == "float8_e5m2":
|
|
return ml_dtypes.float8_e5m2
|
|
return np.dtype(dtype)
|
|
|
|
A_np = np.array(A_np).reshape(shape).astype(get_np_dtype(dtype))
|
|
B_np = np.zeros(shape).astype(get_np_dtype(dtype))
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A = tvm.runtime.tensor(A_np, device=dev)
|
|
B = tvm.runtime.tensor(B_np, device=dev)
|
|
mod(A, B)
|
|
assert np.allclose(A.numpy().astype("float32"), B.numpy().astype("float32"))
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
@pytest.mark.parametrize(
|
|
("shape", "dtype", "encode_args", "error_msg"),
|
|
[
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[0, 16, 32, 16, 16, 1, 1, 0, 0, 0, 0],
|
|
r"globalDim\[0\] must be non-zero",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[(1 << 32) + 1, 16, 32, 16, 16, 1, 1, 0, 0, 0, 0],
|
|
r"globalDim\[0\] must be less than or equal to 2\^32",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[16, 16, 1 << 40, 16, 16, 1, 1, 0, 0, 0, 0],
|
|
r"globalStrides\[0\] must be less than 2\^40",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[16, 16, 32, 0, 16, 1, 1, 0, 0, 0, 0],
|
|
r"boxDim\[0\] must be non-zero",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[16, 16, 32, 7, 16, 1, 1, 0, 0, 0, 0],
|
|
r"boxDim\[0\] \* elementSizeInBytes\(tensorDataType\) must be a multiple of 16 bytes",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[16, 16, 32, 16, 16, 0, 1, 0, 0, 0, 0],
|
|
r"elementStrides\[0\] must be non-zero",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[16, 16, 32, 16, 16, 9, 1, 0, 0, 0, 0],
|
|
r"elementStrides\[0\] must be less than or equal to 8",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"float16",
|
|
[16, 16, 32, 16, 16, 1, 1, 2, 0, 0, 0],
|
|
r"tensorRank must be greater than or equal to 3 when interleave is not NONE",
|
|
),
|
|
(
|
|
(8, 8, 8),
|
|
"float16",
|
|
[8, 8, 8, 16, 128, 8, 8, 8, 1, 1, 1, 2, 0, 0, 0],
|
|
r"globalStrides\[0\] must be a multiple of 32",
|
|
),
|
|
(
|
|
(16, 16),
|
|
"int32",
|
|
[16, 16, 64, 4, 16, 1, 1, 0, 0, 0, 1],
|
|
(
|
|
r"CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA requires a "
|
|
r"floating-point tensorDataType"
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_tensormap_encode_tiled_runtime_validation(shape, dtype, encode_args, error_msg):
|
|
with pytest.raises(tvm.error.InternalError, match=error_msg):
|
|
_run_tensormap_encode(shape, dtype, encode_args)
|
|
|
|
|
|
@pytest.mark.parametrize("swizzle", [1, 2, 3])
|
|
@pytest.mark.parametrize("dtype", ["uint8", "float16", "float32"])
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_cp_async_bulk_tensor_global_to_shared_swizzle(swizzle, dtype):
|
|
def get_ir(swizzle, dtype):
|
|
dtype = tvm.DataType(dtype)
|
|
elem_bytes = dtype.bits // 8
|
|
|
|
shape = [16, 64]
|
|
tma_args = [16, 64, 16, 16, 64, 1, 1, 0, 0, 0, 0] # 8x16B, atom for WGMMA
|
|
shape[0] = shape[0] * (1 << swizzle) // elem_bytes
|
|
tma_args[0] = tma_args[0] * (1 << swizzle) // elem_bytes
|
|
tma_args[2] = tma_args[2] * (1 << swizzle)
|
|
tma_args[3] = tma_args[3] * (1 << swizzle) // elem_bytes
|
|
|
|
load_args = tma_args.copy()
|
|
load_args[-3] = swizzle
|
|
store_args = tma_args.copy()
|
|
|
|
shape = tuple(shape)
|
|
total_elems = math.prod(shape)
|
|
total_bytes = total_elems * elem_bytes
|
|
coord = [0 for _ in shape]
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle, B_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, total_elems, dtype=dtype, align=16)
|
|
B = T.match_buffer(B_ptr, total_elems, dtype=dtype, align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, dtype, len(shape), A.data, *load_args) # noqa: E501
|
|
B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, dtype, len(shape), B.data, *store_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
A_smem = T.alloc_buffer((total_elems,), dtype, scope="shared", align=128)
|
|
bar = T.shared_scalar("uint64")
|
|
phase: T.int32
|
|
|
|
phase = 0
|
|
if threadIdx == 0:
|
|
T.ptx.mbarrier.init(T.address_of(bar), 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coord) # noqa: E501
|
|
T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes)
|
|
T.ptx.mbarrier.try_wait(T.address_of(bar), phase)
|
|
phase = phase ^ 1
|
|
|
|
T.cuda.cta_sync()
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
|
|
if threadIdx == 0:
|
|
T.ptx.cp_async.bulk.tensor.s2g(len(shape), A_smem.access_ptr("r", offset=0), T.address_of(B_map), "", *coord) # noqa: E501
|
|
T.ptx.cp_async.bulk.commit_group()
|
|
T.ptx.cp_async.bulk.wait_group(0)
|
|
# fmt: on
|
|
|
|
return main, shape
|
|
|
|
target = tvm.target.Target("cuda")
|
|
func, shape = get_ir(swizzle, dtype)
|
|
mod = tvm.IRModule({"main": func})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
assert "const __grid_constant__ CUtensorMap" in src
|
|
|
|
total_elems = math.prod(shape)
|
|
A_np = [i for i in range(total_elems)]
|
|
A_np = np.array(A_np).astype(dtype)
|
|
B_np = np.zeros((total_elems,)).astype(dtype)
|
|
dtype = tvm.DataType(dtype)
|
|
layout = T.SwizzleLayout(
|
|
per_element=int(math.log2(128 // dtype.bits)), swizzle_len=swizzle, atom_len=3
|
|
)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A = tvm.runtime.tensor(A_np, device=dev)
|
|
B = tvm.runtime.tensor(B_np, device=dev)
|
|
mod(A, B)
|
|
B_result = B.numpy()
|
|
B_swizzle = [B_result[int(layout.apply(i)["m"])] for i in range(total_elems)]
|
|
B_swizzle = np.array(B_swizzle).astype(str(dtype))
|
|
assert np.allclose(A.numpy(), B_swizzle)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"inputs",
|
|
[
|
|
((128,), [128, 128, 1, 0, 0, 0, 0]),
|
|
((16, 16), [16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0]),
|
|
((4, 4, 4), [4, 4, 4, 16, 64, 4, 4, 4, 1, 1, 1, 0, 0, 0, 0]),
|
|
((4, 4, 4, 4), [4, 4, 4, 4, 16, 64, 256, 4, 4, 4, 4, 1, 1, 1, 1, 0, 0, 0, 0]),
|
|
(
|
|
(4, 2, 2, 2, 2),
|
|
[4, 2, 2, 2, 2, 16, 32, 64, 128, 4, 2, 2, 2, 2, 1, 1, 1, 1, 1, 0, 0, 0, 0],
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_cp_async_bulk_tensor_global_to_shared_multicast1(inputs):
|
|
# 1 CTA does the copy, and then multicast to all CTAs in the cluster
|
|
def get_ir(shape, tma_args):
|
|
total_bytes = 4 * math.prod(shape)
|
|
coord = [0 for _ in shape]
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle, B_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shape, dtype="float32", align=16)
|
|
B = T.match_buffer(B_ptr, shape, dtype="float32", align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, "float32", len(shape), A.data, *tma_args) # noqa: E501
|
|
B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, "float32", len(shape), B.data, *tma_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
for clusterCtaIdx in T.thread_binding(4, thread="clusterCtaIdx.x"):
|
|
for bx in T.thread_binding(4, thread="blockIdx.x"):
|
|
for tx in T.thread_binding(128, thread="threadIdx.x"):
|
|
bar = T.shared_scalar("uint64")
|
|
phase: T.int32
|
|
A_smem = T.alloc_buffer(shape[::-1], "float32", scope="shared", align=128)
|
|
|
|
phase = 0
|
|
if tx == 0:
|
|
# leader thread in each CTA
|
|
T.ptx.mbarrier.init(T.address_of(bar), 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes)
|
|
if clusterCtaIdx == 0:
|
|
# only the first CTA in the cluster does the copy, and then multicast # noqa: E501
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord) # noqa: E501
|
|
# wait for the copy to finish
|
|
T.ptx.mbarrier.try_wait(T.address_of(bar), phase)
|
|
phase = phase ^ 1
|
|
T.cuda.cta_sync()
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
|
|
if bx == 2:
|
|
if tx == 0:
|
|
T.ptx.cp_async.bulk.tensor.s2g(len(shape), A_smem.access_ptr("r", offset=0), T.address_of(B_map), "", *coord) # noqa: E501
|
|
T.ptx.cp_async.bulk.commit_group()
|
|
T.ptx.cp_async.bulk.wait_group(0)
|
|
# fmt: on
|
|
|
|
return main
|
|
|
|
target = tvm.target.Target("cuda")
|
|
shape, tma_args = inputs
|
|
mod = tvm.IRModule({"main": get_ir(shape, tma_args)})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
assert "const __grid_constant__ CUtensorMap" in src
|
|
|
|
A_np = [i for i in range(math.prod(shape))]
|
|
A_np = np.array(A_np, dtype="float32").reshape(shape)
|
|
B_np = np.zeros(shape, dtype="float32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A = tvm.runtime.tensor(A_np, device=dev)
|
|
B = tvm.runtime.tensor(B_np, device=dev)
|
|
mod(A, B)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"inputs",
|
|
[
|
|
((128,), [128, 32, 1, 0, 0, 0, 0]),
|
|
((16, 16), [16, 16, 64, 16, 4, 1, 1, 0, 0, 0, 0]),
|
|
((16, 16, 4), [16, 16, 4, 64, 64 * 16, 16, 16, 1, 1, 1, 1, 0, 0, 0, 0]),
|
|
],
|
|
)
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_cp_async_bulk_tensor_global_to_shared_multicast2(inputs):
|
|
# 4 CTAs in the cluster do the copy of separate chunks, and then multicast to all CTAs in the cluster # noqa: E501
|
|
def get_ir(shape, tma_args):
|
|
assert shape[0] % 4 == 0
|
|
total_bytes = 4 * math.prod(shape)
|
|
coord0 = [0 for _ in shape]
|
|
coord1 = [0 for _ in shape[:-1]] + [shape[-1] // 4]
|
|
coord2 = [0 for _ in shape[:-1]] + [shape[-1] // 2]
|
|
coord3 = [0 for _ in shape[:-1]] + [3 * shape[-1] // 4]
|
|
|
|
tma_store_args = tma_args.copy()
|
|
tma_store_args[3 * len(shape) - 2] = shape[-1]
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle, B_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shape, dtype="float32", align=16)
|
|
B = T.match_buffer(B_ptr, shape, dtype="float32", align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, "float32", len(shape), A.data, *tma_args) # noqa: E501
|
|
B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, "float32", len(shape), B.data, *tma_store_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
for clusterCtaIdx in T.thread_binding(4, thread="clusterCtaIdx.x"):
|
|
for bx in T.thread_binding(4, thread="blockIdx.x"):
|
|
for tx in T.thread_binding(128, thread="threadIdx.x"):
|
|
bar = T.shared_scalar("uint64")
|
|
phase: T.int32
|
|
A_smem = T.alloc_buffer(shape[::-1], "float32", scope="shared", align=128)
|
|
|
|
phase = 0
|
|
if tx == 0:
|
|
# leader thread in each CTA
|
|
T.ptx.mbarrier.init(T.address_of(bar), 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes)
|
|
if clusterCtaIdx == 0:
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord0[::-1])), # noqa: E501
|
|
T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord0) # noqa: E501
|
|
if clusterCtaIdx == 1:
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord1[::-1])), # noqa: E501
|
|
T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord1) # noqa: E501
|
|
if clusterCtaIdx == 2:
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord2[::-1])), # noqa: E501
|
|
T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord2) # noqa: E501
|
|
if clusterCtaIdx == 3:
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(Buffer.WRITE, offset=A_smem.elem_offset_of(coord3[::-1])), # noqa: E501
|
|
T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord3) # noqa: E501
|
|
# wait for the copy to finish
|
|
T.ptx.mbarrier.try_wait(T.address_of(bar), phase)
|
|
phase = phase ^ 1
|
|
T.cuda.cta_sync()
|
|
|
|
if bx == 1:
|
|
if tx == 0:
|
|
T.ptx.cp_async.bulk.tensor.s2g(len(shape), A_smem.access_ptr("r", offset=0), T.address_of(B_map), "", *coord0) # noqa: E501
|
|
T.ptx.cp_async.bulk.commit_group()
|
|
T.ptx.cp_async.bulk.wait_group(0)
|
|
# fmt: on
|
|
|
|
return main
|
|
|
|
target = tvm.target.Target("cuda")
|
|
shape, tma_args = inputs
|
|
mod = tvm.IRModule({"main": get_ir(shape, tma_args)})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
assert "const __grid_constant__ CUtensorMap" in src
|
|
|
|
A_np = [i for i in range(math.prod(shape))]
|
|
A_np = np.array(A_np, dtype="float32").reshape(shape)
|
|
B_np = np.zeros(shape, dtype="float32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A = tvm.runtime.tensor(A_np, device=dev)
|
|
B = tvm.runtime.tensor(B_np, device=dev)
|
|
mod(A, B)
|
|
assert np.allclose(A.numpy(), B.numpy())
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"inputs",
|
|
[
|
|
((128,), [128, 128, 1, 0, 0, 0, 0]),
|
|
((16, 16), [16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0]),
|
|
((16, 16, 4), [16, 16, 4, 64, 64 * 16, 16, 16, 4, 1, 1, 1, 0, 0, 0, 0]),
|
|
],
|
|
)
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_cp_async_bulk_tensor_shared_to_global(inputs):
|
|
def get_ir(shape, tma_args):
|
|
assert shape[0] % 4 == 0
|
|
elems = math.prod(shape)
|
|
coord = [0 for _ in shape]
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shape, dtype="float32", align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, "float32", len(shape), A.data, *tma_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([128])
|
|
|
|
A_smem = T.alloc_buffer(elems, "float32", scope="shared", align=128)
|
|
|
|
if tx == 0:
|
|
for i in T.serial(0, elems):
|
|
A_smem[i] = i
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.cuda.cta_sync()
|
|
|
|
if tx == 0:
|
|
T.ptx.cp_async.bulk.tensor.s2g(len(shape), A_smem.access_ptr("r", offset=0), T.address_of(A_map), "", *coord) # noqa: E501
|
|
T.ptx.cp_async.bulk.commit_group()
|
|
T.ptx.cp_async.bulk.wait_group(0)
|
|
# fmt: on
|
|
|
|
return main
|
|
|
|
target = tvm.target.Target("cuda")
|
|
shape, tma_args = inputs
|
|
mod = tvm.IRModule({"main": get_ir(shape, tma_args)})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
src = mod.mod.imports[0].inspect_source()
|
|
assert "const __grid_constant__ CUtensorMap" in src
|
|
|
|
A_np = np.zeros(shape, dtype="float32")
|
|
A_ref = [i for i in range(math.prod(shape))]
|
|
A_ref = np.array(A_ref, dtype="float32").reshape(shape)
|
|
|
|
def run_and_check():
|
|
A = tvm.runtime.tensor(A_np, device=tvm.cuda(0))
|
|
mod(A)
|
|
np.testing.assert_allclose(A.numpy(), A_ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9, exact=True), reason="need cuda compute == 9.0")
|
|
def test_wgmma_ss_nt():
|
|
def get_ir(
|
|
shapeA,
|
|
shapeB,
|
|
shapeC,
|
|
A_tma_args,
|
|
B_tma_args,
|
|
in_dtype,
|
|
out_dtype,
|
|
A_encode_args,
|
|
B_encode_args,
|
|
):
|
|
coordA = [0 for _ in shapeA]
|
|
coordB = [0 for _ in shapeB]
|
|
A_bytes = tvm.DataType(in_dtype).bits // 8 * math.prod(shapeA)
|
|
B_bytes = tvm.DataType(in_dtype).bits // 8 * math.prod(shapeB)
|
|
|
|
C_elems = math.prod(shapeC) // 128
|
|
|
|
M, K = shapeA if not transA else shapeA[::-1]
|
|
N, _ = shapeB if not transB else shapeB[::-1]
|
|
|
|
def get_init_value(dtype):
|
|
if dtype == "float32":
|
|
return T.float32(0.0)
|
|
assert False, f"Unsupported dtype {dtype}"
|
|
|
|
def get_accum_list(C, C_elems):
|
|
return [C[i] for i in range(C_elems)]
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shapeA, dtype=in_dtype, align=16)
|
|
B = T.match_buffer(B_ptr, shapeB, dtype=in_dtype, align=16)
|
|
C = T.match_buffer(C_ptr, shapeC, dtype=out_dtype, align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, in_dtype, len(shapeA), A.data, *A_tma_args) # noqa: E501
|
|
B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, in_dtype, len(shapeB), B.data, *B_tma_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([128]) # A warpgroup is 128 threads
|
|
|
|
A_smem = T.alloc_buffer(shapeA, in_dtype, scope="shared", align=1024)
|
|
B_smem = T.alloc_buffer(shapeB, in_dtype, scope="shared", align=1024)
|
|
bar = T.shared_scalar("uint64")
|
|
phase: T.int32
|
|
|
|
descA: T.uint64
|
|
descB: T.uint64
|
|
C_local = T.alloc_buffer((C_elems,), out_dtype, scope="local")
|
|
|
|
# init phase and bar
|
|
phase = 0
|
|
if tx == 0:
|
|
T.ptx.mbarrier.init(T.address_of(bar), 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.cuda.cta_sync()
|
|
# load A and B to smem
|
|
if tx == 0:
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shapeA), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coordA) # noqa: E501
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shapeB), B_smem.data, T.address_of(bar), T.address_of(B_map), 0, 1, "", *coordB) # noqa: E501
|
|
T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), A_bytes + B_bytes)
|
|
T.ptx.mbarrier.try_wait(T.address_of(bar), phase)
|
|
phase = phase ^ 1
|
|
T.cuda.cta_sync()
|
|
|
|
# init C_local
|
|
for i in T.serial(0, C_elems):
|
|
C_local[i] = T.Cast(out_dtype, get_init_value(out_dtype))
|
|
T.ptx.wgmma.noop_barrier(C_local[i])
|
|
|
|
# do wgmma
|
|
T.ptx.wgmma.encode_matrix_descriptor(T.address_of(descA), A_smem.data, *A_encode_args) # noqa: F821
|
|
T.ptx.wgmma.encode_matrix_descriptor(T.address_of(descB), B_smem.data, *B_encode_args) # noqa: F821
|
|
T.ptx.wgmma.fence()
|
|
T.ptx.wgmma.mma_async.ss(descA, descB, *get_accum_list(C_local, C_elems), # noqa: F821
|
|
M=M, N=N, K=K, in_dtype=in_dtype, out_dtype=out_dtype, transA=transA, transB=transB, scaleA=1.0, scaleB=1.0, scaleD=False) # noqa: E501
|
|
T.ptx.wgmma.commit_group()
|
|
T.ptx.wgmma.wait_group(0)
|
|
|
|
for i in T.serial(0, C_elems):
|
|
T.ptx.wgmma.noop_barrier(C_local[i])
|
|
|
|
# store C_local to C
|
|
for i in T.serial(0, C_elems // 4):
|
|
row = T.meta_var((tx % 32) // 4 + (tx // 32) * 16)
|
|
col = T.meta_var(i * 8 + tx % 4 * 2)
|
|
C[row, col] = C_local[i * 4]
|
|
C[row, col + 1] = C_local[i * 4 + 1]
|
|
C[row + 8, col] = C_local[i * 4 + 2]
|
|
C[row + 8, col + 1] = C_local[i * 4 + 3]
|
|
# fmt: on
|
|
|
|
return main
|
|
|
|
in_dtype = "float16"
|
|
out_dtype = "float32"
|
|
transA = transB = True
|
|
swizzleA = swizzleB = 3
|
|
|
|
t_in_dtype = tvm.DataType(in_dtype)
|
|
elem_bytes = t_in_dtype.bits // 8
|
|
|
|
target = tvm.target.Target("cuda")
|
|
M = 64
|
|
N = 64
|
|
K = 256 // t_in_dtype.bits
|
|
shapeA = (M, K) if not transA else (K, M)
|
|
shapeB = (N, K) if not transB else (K, N)
|
|
shapeC = (M, N)
|
|
|
|
# A tma args
|
|
A_outer, A_inner = shapeA
|
|
A_tma_args = [A_inner, A_outer, A_inner * elem_bytes, A_inner, A_outer, 1, 1, 0, swizzleA, 0, 0]
|
|
# B tma args
|
|
B_outer, B_inner = shapeB
|
|
B_tma_args = [B_inner, B_outer, B_inner * elem_bytes, B_inner, B_outer, 1, 1, 0, swizzleB, 0, 0]
|
|
# A encode args
|
|
A_encode_args = [1, 64, swizzleA]
|
|
B_encode_args = [1, 64, swizzleB]
|
|
|
|
func = get_ir(
|
|
shapeA,
|
|
shapeB,
|
|
shapeC,
|
|
A_tma_args,
|
|
B_tma_args,
|
|
in_dtype,
|
|
out_dtype,
|
|
A_encode_args,
|
|
B_encode_args,
|
|
)
|
|
mod = tvm.IRModule({"main": func})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
|
|
np.random.seed(0)
|
|
A_np = np.random.randn(*shapeA).astype(in_dtype)
|
|
B_np = np.random.randn(*shapeB).astype(in_dtype)
|
|
C_np = np.zeros(shapeC).astype(out_dtype)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A_tvm = tvm.runtime.tensor(A_np, device=dev)
|
|
B_tvm = tvm.runtime.tensor(B_np, device=dev)
|
|
C_tvm = tvm.runtime.tensor(C_np, device=dev)
|
|
mod(A_tvm, B_tvm, C_tvm)
|
|
C_ref = np.dot(A_np.T, B_np).astype(out_dtype)
|
|
tvm.testing.assert_allclose(C_tvm.numpy(), C_ref, rtol=1e-3, atol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9, exact=True), reason="need cuda compute == 9.0")
|
|
def test_wgmma_rs_nt():
|
|
def get_ir(
|
|
shapeA, shapeB, shapeC, B_tma_args, in_dtype, in_dtype_bits, out_dtype, B_encode_args
|
|
):
|
|
coordB = [0 for _ in shapeB]
|
|
B_bytes = tvm.DataType(in_dtype).bits // 8 * math.prod(shapeB)
|
|
|
|
A_elems = math.prod(shapeA) // 128
|
|
C_elems = math.prod(shapeC) // 128
|
|
|
|
M, K = shapeA if not transA else shapeA[::-1]
|
|
N, _ = shapeB if not transB else shapeB[::-1]
|
|
|
|
def get_init_value(dtype):
|
|
if dtype == "float32":
|
|
return T.float32(0.0)
|
|
assert False, f"Unsupported dtype {dtype}"
|
|
|
|
def get_A_list(A_local, A_elems):
|
|
return [A_local[i] for i in range(A_elems)]
|
|
|
|
def get_accum_list(C, C_elems):
|
|
return [C[i] for i in range(C_elems)]
|
|
|
|
# fmt: off
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, shapeA, dtype=in_dtype, align=16)
|
|
B = T.match_buffer(B_ptr, shapeB, dtype=in_dtype, align=16)
|
|
C = T.match_buffer(C_ptr, shapeC, dtype=out_dtype, align=16)
|
|
|
|
B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, in_dtype, len(shapeB), B.data, *B_tma_args) # noqa: E501
|
|
|
|
T.device_entry()
|
|
cta_id = T.cta_id([1])
|
|
tx = T.thread_id([128]) # A warpgroup is 128 threads
|
|
|
|
B_smem = T.alloc_buffer(shapeB, in_dtype, scope="shared", align=1024)
|
|
# bar = T.alloc_buffer((1,), "uint64", scope="shared", align=8)
|
|
bar = T.shared_scalar("uint64")
|
|
|
|
# descB = T.alloc_buffer((1,), "uint64", scope="local")
|
|
descB: T.uint64
|
|
A_local = T.alloc_buffer((A_elems,), in_dtype, scope="local")
|
|
C_local = T.alloc_buffer((C_elems,), out_dtype, scope="local")
|
|
|
|
A_elems_b32 = T.meta_var(A_elems // (32 // in_dtype_bits))
|
|
A_local_b32 = T.decl_buffer((A_elems_b32,), "uint32", data=A_local.data)
|
|
|
|
# load A to regs
|
|
for i in T.serial(0, A_elems // 4):
|
|
row = T.meta_var((tx % 32) // 4 + (tx // 32) * 16)
|
|
col = T.meta_var(i * 8 + tx % 4 * 2)
|
|
A_local[i * 4] = A[row, col]
|
|
A_local[i * 4 + 1] = A[row, col + 1]
|
|
A_local[i * 4 + 2] = A[row + 8, col]
|
|
A_local[i * 4 + 3] = A[row + 8, col + 1]
|
|
# init bar, and make sure it's visible to all threads and async proxy
|
|
if tx == 0:
|
|
T.ptx.mbarrier.init(T.address_of(bar), 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.cuda.cta_sync()
|
|
# load B to smem
|
|
if tx == 0:
|
|
T.ptx.cp_async.bulk.tensor.g2c(len(shapeB), B_smem.data, T.address_of(bar), T.address_of(B_map), 0, 1, "", *coordB) # noqa: E501
|
|
T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), B_bytes)
|
|
T.ptx.mbarrier.try_wait(T.address_of(bar), 0)
|
|
T.cuda.cta_sync()
|
|
|
|
# init C_local
|
|
for i in T.serial(0, C_elems):
|
|
C_local[i] = T.Cast(out_dtype, get_init_value(out_dtype))
|
|
|
|
# fence A_local and C_local
|
|
for i in T.serial(0, A_elems_b32):
|
|
T.ptx.wgmma.noop_barrier(A_local_b32[i])
|
|
for i in T.serial(0, C_elems):
|
|
T.ptx.wgmma.noop_barrier(C_local[i])
|
|
# do wgmma
|
|
T.ptx.wgmma.encode_matrix_descriptor(T.address_of(descB), B_smem.data, *B_encode_args) # noqa: F821
|
|
T.ptx.wgmma.fence()
|
|
T.ptx.wgmma.mma_async.rs(descB, *(get_A_list(A_local_b32, A_elems_b32) + get_accum_list(C_local, C_elems)), # noqa: E501, F821
|
|
M=M, N=N, K=K, in_dtype=in_dtype, out_dtype=out_dtype, transA=transA, transB=transB, scaleA=1.0, scaleB=1.0, scaleD=False) # noqa: E501
|
|
T.ptx.wgmma.commit_group()
|
|
T.ptx.wgmma.wait_group(0)
|
|
|
|
# fence A_local
|
|
for i in T.serial(0, A_elems_b32):
|
|
T.ptx.wgmma.noop_barrier(A_local_b32[i])
|
|
# fence C_local
|
|
for i in T.serial(0, C_elems):
|
|
T.ptx.wgmma.noop_barrier(C_local[i])
|
|
|
|
# store C_local to C
|
|
for i in T.serial(0, C_elems // 4):
|
|
row = T.meta_var((tx % 32) // 4 + (tx // 32) * 16)
|
|
col = T.meta_var(i * 8 + tx % 4 * 2)
|
|
C[row, col] = C_local[i * 4]
|
|
C[row, col + 1] = C_local[i * 4 + 1]
|
|
C[row + 8, col] = C_local[i * 4 + 2]
|
|
C[row + 8, col + 1] = C_local[i * 4 + 3]
|
|
# fmt: on
|
|
|
|
return main
|
|
|
|
in_dtype = "float16"
|
|
in_dtype_bits = 16
|
|
out_dtype = "float32"
|
|
transA = False
|
|
transB = True
|
|
swizzleB = 3
|
|
|
|
t_in_dtype = tvm.DataType(in_dtype)
|
|
elem_bytes = t_in_dtype.bits // 8
|
|
|
|
target = tvm.target.Target("cuda")
|
|
M = 64
|
|
N = 64
|
|
K = 256 // t_in_dtype.bits
|
|
shapeA = (M, K) if not transA else (K, M)
|
|
shapeB = (N, K) if not transB else (K, N)
|
|
shapeC = (M, N)
|
|
|
|
# B tma args
|
|
B_outer, B_inner = shapeB
|
|
B_tma_args = [B_inner, B_outer, B_inner * elem_bytes, B_inner, B_outer, 1, 1, 0, swizzleB, 0, 0]
|
|
# B encode args
|
|
B_encode_args = [1, 64, swizzleB]
|
|
|
|
func = get_ir(
|
|
shapeA, shapeB, shapeC, B_tma_args, in_dtype, in_dtype_bits, out_dtype, B_encode_args
|
|
)
|
|
mod = tvm.IRModule({"main": func})
|
|
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
|
|
|
|
np.random.seed(0)
|
|
A_np = np.random.randn(*shapeA).astype(in_dtype)
|
|
B_np = np.random.randn(*shapeB).astype(in_dtype)
|
|
C_np = np.zeros(shapeC).astype(out_dtype)
|
|
|
|
np.printoptions(threshold=np.inf)
|
|
np.printoptions(linewidth=np.inf)
|
|
np.printoptions(precision=2)
|
|
|
|
C_ref = np.dot(A_np, B_np).astype(out_dtype)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
A_tvm = tvm.runtime.tensor(A_np, device=dev)
|
|
B_tvm = tvm.runtime.tensor(B_np, device=dev)
|
|
C_tvm = tvm.runtime.tensor(C_np, device=dev)
|
|
mod(A_tvm, B_tvm, C_tvm)
|
|
tvm.testing.assert_allclose(C_tvm.numpy(), C_ref, rtol=1e-3, atol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_ptx_map_shared_rank():
|
|
@T.prim_func
|
|
def func(A: T.Buffer(1)):
|
|
T.device_entry()
|
|
cbx = T.cta_id_in_cluster([2])
|
|
cta_id = T.cta_id([2])
|
|
tx = T.thread_id([128])
|
|
A_smem = T.alloc_buffer([1], "uint32", scope="shared")
|
|
if cbx == 0 and tx == 0:
|
|
T.ptx.map_shared_rank(A_smem.data, cbx)
|
|
|
|
src, mod = _get_source(func)
|
|
print(src)
|
|
assert "tvm_builtin_ptx_mapa_u64(A_smem" in src
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|