wehub-resource-sync
1202 lines
45 KiB
Python
1202 lines
45 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# ruff: noqa: E501, E741, F401, F841
|
|
|
|
# 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.
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.support.nvcc
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.support.nvcc import have_bf16, have_fp16, have_int8
|
|
from tvm.testing import env
|
|
|
|
|
|
@pytest.fixture(autouse=True, params=["nvcc", "nvrtc"])
|
|
def setup_cuda_compile_mode(request):
|
|
mode = request.param
|
|
if mode == "nvrtc":
|
|
try:
|
|
from cuda.bindings import nvrtc
|
|
except ImportError:
|
|
pytest.skip("cuda-python not available, skipping nvrtc tests")
|
|
|
|
orig_func = tvm.support.nvcc.tvm_callback_cuda_compile
|
|
|
|
def compile_mode_wrapper(code):
|
|
if mode == "nvcc":
|
|
return tvm.support.nvcc.compile_cuda(code, target_format="fatbin", compiler="nvcc")
|
|
elif mode == "nvrtc":
|
|
return tvm.support.nvcc.compile_cuda(code, target_format="cubin", compiler="nvrtc")
|
|
else:
|
|
raise ValueError(f"Unknown mode: {mode}")
|
|
|
|
tvm.register_global_func("tvm_callback_cuda_compile", compile_mode_wrapper, override=True)
|
|
# yield back to the original function so that each test runs twice
|
|
yield
|
|
tvm.register_global_func("tvm_callback_cuda_compile", orig_func, override=True)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_vectorize_add():
|
|
num_thread = 8
|
|
|
|
def check_cuda(dtype, n, lanes):
|
|
if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
|
|
print("Skip because gpu does not have fp16 support")
|
|
return
|
|
if dtype == "int8" and not have_int8(tvm.cuda(0).compute_version):
|
|
print("skip because gpu does not support int8")
|
|
return
|
|
vec_dtype = f"{dtype}x{lanes}"
|
|
one = tvm.tirx.const(1, vec_dtype)
|
|
num_blocks = (n + num_thread - 1) // num_thread
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
|
|
T.where(i_0 * num_thread + i_1 < n)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i] + one
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np.random.uniform(size=(n, lanes)))
|
|
c = tvm.runtime.empty((n,), vec_dtype, dev)
|
|
fun(a, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_cuda("float32", 64, 2)
|
|
check_cuda("float32", 64, 3)
|
|
check_cuda("float32", 64, 4)
|
|
check_cuda("int8", 64, 2)
|
|
check_cuda("int8", 64, 3)
|
|
check_cuda("int8", 64, 4)
|
|
check_cuda("uint8", 64, 2)
|
|
check_cuda("uint8", 64, 3)
|
|
check_cuda("uint8", 64, 4)
|
|
check_cuda("float16", 64, 2)
|
|
check_cuda("float16", 64, 4)
|
|
check_cuda("float16", 64, 6)
|
|
check_cuda("float16", 64, 8)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_bf16_vectorize_add():
|
|
if not have_bf16(tvm.cuda(0).compute_version):
|
|
print("skip because gpu does not support bf16")
|
|
return
|
|
num_thread = 8
|
|
|
|
def np_float2np_bf16(arr):
|
|
"""Convert a numpy array of float to a numpy array
|
|
of bf16 in uint16"""
|
|
orig = arr.view("<u4")
|
|
bias = np.bitwise_and(np.right_shift(orig, 16), 1) + 0x7FFF
|
|
return np.right_shift(orig + bias, 16).astype("uint16")
|
|
|
|
def np_bf162np_float(arr):
|
|
"""Convert a numpy array of bf16 (uint16) to a numpy array
|
|
of float"""
|
|
u32 = np.left_shift(arr.astype("uint32"), 16)
|
|
return u32.view("<f4")
|
|
|
|
def check_cuda(n, lanes):
|
|
vec_dtype = f"bfloat16x{lanes}"
|
|
num_blocks = n // num_thread
|
|
one = tvm.tirx.Broadcast(tvm.tirx.const(1, "bfloat16"), lanes)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i] + one
|
|
|
|
with tvm.transform.PassContext(
|
|
disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination", "tirx.BF16TypeLowering"]
|
|
):
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
np_a = np.random.uniform(size=(n, lanes)).astype("float32")
|
|
np_a = np_bf162np_float(np_float2np_bf16(np_a))
|
|
a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_float2np_bf16(np_a))
|
|
c = tvm.runtime.empty((n,), vec_dtype, dev)
|
|
fun(a, c)
|
|
c = tvm.runtime.empty((n, lanes), "uint16", dev).copyfrom(c)
|
|
tvm.testing.assert_allclose(c.numpy(), np_float2np_bf16(np_a + 1))
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_cuda(64, 2)
|
|
check_cuda(64, 4)
|
|
check_cuda(64, 6)
|
|
check_cuda(64, 8)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_multiply_add():
|
|
num_thread = 8
|
|
|
|
def check_cuda(dtype, n, lanes):
|
|
if dtype == "int8" and not have_int8(tvm.cuda(0).compute_version):
|
|
print("skip because gpu does not support int8")
|
|
return
|
|
vec_dtype = f"{dtype}x{lanes}"
|
|
num_blocks = n // num_thread
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((n,), vec_dtype),
|
|
B: T.Buffer((n,), vec_dtype),
|
|
C: T.Buffer((n,), "int32"),
|
|
D: T.Buffer((n,), "int32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
|
|
with T.sblock("D"):
|
|
v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
|
|
T.reads(A[v_i], B[v_i], C[v_i])
|
|
T.writes(D[v_i])
|
|
D[v_i] = T.call_pure_extern("int32", "__dp4a", A[v_i], B[v_i], C[v_i])
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
|
|
np_b = np.random.randint(low=-128, high=127, size=(n, lanes))
|
|
np_c = np.random.randint(low=0, high=127, size=(n,))
|
|
np_d = [sum(x * y) + z for x, y, z in zip(np_a, np_b, np_c)]
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_a)
|
|
b = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_b)
|
|
c = tvm.runtime.empty((n,), "int32", dev).copyfrom(np_c)
|
|
d = tvm.runtime.empty((n,), "int32", dev)
|
|
fun(a, b, c, d)
|
|
tvm.testing.assert_allclose(d.numpy(), np_d)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_cuda("int8", 64, 4)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_vectorize_load():
|
|
num_thread = 8
|
|
|
|
def check_cuda(dtype, n, lanes):
|
|
vec_dtype = f"{dtype}x{lanes}"
|
|
num_blocks = n // num_thread
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i]
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
|
|
a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_a)
|
|
b = tvm.runtime.empty((n,), vec_dtype, dev)
|
|
fun(a, b)
|
|
tvm.testing.assert_allclose(a.numpy(), b.numpy())
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_cuda("int8", 64, 2)
|
|
check_cuda("int8", 64, 3)
|
|
check_cuda("int8", 64, 4)
|
|
check_cuda("int8", 64, 8)
|
|
check_cuda("int8", 64, 16)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_make_int8():
|
|
def check_cuda(n, value, lanes):
|
|
dtype = "int8"
|
|
const_value = tvm.tirx.const(value, dtype=dtype)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n, lanes), dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in T.thread_binding(n, thread="blockIdx.x"):
|
|
for j in T.vectorized(lanes):
|
|
with T.sblock("A"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads()
|
|
T.writes(A[v_i, v_j])
|
|
A[v_i, v_j] = const_value
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
np_a = np.full((n, lanes), value, dtype=dtype)
|
|
a = tvm.runtime.empty(np_a.shape, dtype, dev)
|
|
fun(a)
|
|
np.testing.assert_equal(a.numpy(), np_a)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_cuda(64, np.uint8(0xAB).view(np.int8), 4)
|
|
check_cuda(64, 0, 4)
|
|
check_cuda(64, -3, 4)
|
|
check_cuda(64, np.uint8(0xAB).view(np.int8), 3)
|
|
check_cuda(64, 0, 3)
|
|
check_cuda(64, -3, 3)
|
|
check_cuda(64, np.uint8(0xAB).view(np.int8), 2)
|
|
check_cuda(64, 0, 2)
|
|
check_cuda(64, -3, 2)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_inf_nan():
|
|
def check_inf_nan(n, value, dtype):
|
|
inf_value = tvm.tirx.const(value, dtype=dtype)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), dtype), C: T.Buffer((n,), dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(1, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(8, thread="threadIdx.x"):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(n, i_0 * 8 + i_1)
|
|
T.where(i_0 * 8 + i_1 < n)
|
|
T.reads()
|
|
T.writes(C[v_i])
|
|
C[v_i] = inf_value
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.device("cuda", 0)
|
|
a = tvm.runtime.empty((n,), dtype, dev)
|
|
c = tvm.runtime.empty((n,), dtype, dev)
|
|
fun(a, c)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_inf_nan(1, -float("inf"), "float32")
|
|
check_inf_nan(1, -float("inf"), "float64")
|
|
check_inf_nan(1, float("inf"), "float32")
|
|
check_inf_nan(1, float("inf"), "float64")
|
|
check_inf_nan(1, float("nan"), "float32")
|
|
check_inf_nan(1, float("nan"), "float64")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"target",
|
|
[
|
|
pytest.param("cuda", marks=pytest.mark.gpu),
|
|
pytest.param("rocm", marks=pytest.mark.gpu),
|
|
],
|
|
)
|
|
def test_crossthread_reduction1(target):
|
|
if not tvm.testing.device_enabled(target):
|
|
pytest.skip(f"{target} not enabled")
|
|
|
|
def sched(nthd):
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n, m = T.int32(), T.int32()
|
|
A = T.match_buffer(var_A, (n, m))
|
|
B = T.match_buffer(var_B, (n,))
|
|
for i in T.thread_binding(n, thread="blockIdx.x"):
|
|
for m_0 in T.thread_binding(nthd, thread="threadIdx.x"):
|
|
for m_1 in range((m + nthd - 1) // nthd):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i)
|
|
v_m = T.axis.reduce(m, m_0 * ((m + nthd - 1) // nthd) + m_1)
|
|
T.where(m_0 * ((m + nthd - 1) // nthd) + m_1 < m)
|
|
T.reads(A[v_i, v_m])
|
|
T.writes(B[v_i])
|
|
with T.init():
|
|
B[v_i] = T.float32(0.0)
|
|
B[v_i] = B[v_i] + A[v_i, v_m]
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
return fun
|
|
|
|
def verify(nthd):
|
|
func = sched(nthd)
|
|
nn = 3
|
|
# checks three typical cases
|
|
vals = [nthd - 1, nthd, nthd + 1]
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target)
|
|
for kk in vals:
|
|
size = (nn, kk)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
|
|
func(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=1), rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
verify(16)
|
|
verify(32)
|
|
verify(64)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"target",
|
|
[
|
|
pytest.param("cuda", marks=pytest.mark.gpu),
|
|
pytest.param("rocm", marks=pytest.mark.gpu),
|
|
],
|
|
)
|
|
def test_crossthread_reduction2(target):
|
|
if not tvm.testing.device_enabled(target):
|
|
pytest.skip(f"{target} not enabled")
|
|
|
|
def sched(nthdx, nthdy):
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n, k0, k1 = T.int32(), T.int32(), T.int32()
|
|
A = T.match_buffer(var_A, (n, k0, k1))
|
|
B = T.match_buffer(var_B, (n,))
|
|
for i in T.thread_binding(n, thread="blockIdx.x"):
|
|
for k0_0 in T.thread_binding(nthdx, thread="threadIdx.x"):
|
|
for k0_1 in range((k0 + nthdx - 1) // nthdx):
|
|
for k1_0 in T.thread_binding(nthdy, thread="threadIdx.y"):
|
|
for k1_1 in range((k1 + nthdy - 1) // nthdy):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i)
|
|
v_k0 = T.axis.reduce(
|
|
k0, k0_0 * ((k0 + nthdx - 1) // nthdx) + k0_1
|
|
)
|
|
v_k1 = T.axis.reduce(
|
|
k1, k1_0 * ((k1 + nthdy - 1) // nthdy) + k1_1
|
|
)
|
|
T.where(
|
|
k0_0 * ((k0 + nthdx - 1) // nthdx) + k0_1 < k0
|
|
and k1_0 * ((k1 + nthdy - 1) // nthdy) + k1_1 < k1
|
|
)
|
|
T.reads(A[v_i, v_k0, v_k1])
|
|
T.writes(B[v_i])
|
|
with T.init():
|
|
B[v_i] = T.float32(0.0)
|
|
B[v_i] = B[v_i] + A[v_i, v_k0, v_k1]
|
|
|
|
func = tvm.compile(Module, target="cuda")
|
|
return func
|
|
|
|
def verify(nthdx, nthdy):
|
|
func = sched(nthdx, nthdy)
|
|
nn = 3
|
|
# checks three typical cases
|
|
vx = [nthdx - 1, nthdx, nthdx + 1]
|
|
vy = [nthdy - 1, nthdy, nthdy + 1]
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target)
|
|
for kk0, kk1 in [(x, y) for x in vx for y in vy]:
|
|
size = (nn, kk0, kk1)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
|
|
func(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=(1, 2)), rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
verify(16, 16)
|
|
verify(32, 32)
|
|
verify(16, 32)
|
|
verify(32, 16)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_reduction_binding():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((96, 32), "float32"), B: T.Buffer((96,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for k in range(32):
|
|
for m_0 in T.thread_binding(3, thread="blockIdx.x"):
|
|
for m_1 in range(32):
|
|
with T.sblock("B"):
|
|
v_m = T.axis.spatial(96, m_0 * 32 + m_1)
|
|
v_k = T.axis.reduce(32, k)
|
|
T.reads(A[v_m, v_k])
|
|
T.writes(B[v_m])
|
|
with T.init():
|
|
B[v_m] = T.float32(0.0)
|
|
B[v_m] = B[v_m] + A[v_m, v_k]
|
|
|
|
func = tvm.compile(Module, target="cuda")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_const_float_to_half():
|
|
# This import is required to use nvcc to perform code gen;
|
|
# otherwise it is found that the code gen is done by nvrtc.
|
|
|
|
half_const = tvm.tirx.const(0.5, dtype="float16")
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.Buffer((2, 3, 4), "float16"), C: T.Buffer((2, 3, 4), "bool")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_j_k_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
|
for i_j_k_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(2, (i_j_k_fused_0 * 64 + i_j_k_fused_1) // 12)
|
|
v_j = T.axis.spatial(3, (i_j_k_fused_0 * 64 + i_j_k_fused_1) % 12 // 4)
|
|
v_k = T.axis.spatial(4, (i_j_k_fused_0 * 64 + i_j_k_fused_1) % 4)
|
|
T.where(i_j_k_fused_0 * 64 + i_j_k_fused_1 < 24)
|
|
T.reads(a[v_i, v_j, v_k])
|
|
T.writes(C[v_i, v_j, v_k])
|
|
C[v_i, v_j, v_k] = half_const < a[v_i, v_j, v_k]
|
|
|
|
func = tvm.compile(Module, target="cuda")
|
|
|
|
shape = (2, 3, 4)
|
|
a_np = np.random.uniform(size=shape).astype("float16")
|
|
c_np = np.zeros(shape=shape, dtype="bool")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
c = tvm.runtime.tensor(c_np, dev)
|
|
func(a, c)
|
|
np.testing.assert_equal(c.numpy(), a_np > 0.5)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_floordiv_with_vectorization():
|
|
with tvm.target.Target("cuda"):
|
|
# B[i] = A[floordiv(i, k)]
|
|
n = 256
|
|
k = 37
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(1, thread="blockIdx.x"):
|
|
for i_1_0 in T.thread_binding(64, thread="threadIdx.x"):
|
|
for i_1_1 in T.vectorized(4):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(256, i_0 * 256 + i_1_0 * 4 + i_1_1)
|
|
T.reads(A[v_i // 37])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i // 37]
|
|
|
|
func = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a_np = np.random.uniform(size=(n,)).astype("float32")
|
|
b_np = np.array([a_np[i // k] for i in range(0, n)])
|
|
a_nd = tvm.runtime.tensor(a_np, dev)
|
|
b_nd = tvm.runtime.tensor(np.zeros(b_np.shape, dtype=b_np.dtype), dev)
|
|
func(a_nd, b_nd)
|
|
tvm.testing.assert_allclose(b_nd.numpy(), b_np, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_floormod_with_vectorization():
|
|
with tvm.target.Target("cuda"):
|
|
# B[i] = A[floormod(i, k)]
|
|
n = 256
|
|
k = 37
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(1, thread="blockIdx.x"):
|
|
for i_1_0 in T.thread_binding(64, thread="threadIdx.x"):
|
|
for i_1_1 in T.vectorized(4):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(256, i_0 * 256 + i_1_0 * 4 + i_1_1)
|
|
T.reads(A[v_i % 37])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i % 37]
|
|
|
|
func = tvm.compile(Module, target="cuda")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a_np = np.random.uniform(size=(n,)).astype("float32")
|
|
b_np = np.array([a_np[i % k] for i in range(0, n)])
|
|
a_nd = tvm.runtime.tensor(a_np, dev)
|
|
b_nd = tvm.runtime.tensor(np.zeros(b_np.shape, dtype=b_np.dtype), dev)
|
|
func(a_nd, b_nd)
|
|
tvm.testing.assert_allclose(b_nd.numpy(), b_np, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_vectorized_casts():
|
|
def check(t0, t1, factor):
|
|
if (t0 == "float16" or t1 == "float16") and not have_fp16(tvm.cuda(0).compute_version):
|
|
print("Skip because gpu does not have fp16 support")
|
|
return
|
|
|
|
n = 128
|
|
num_thread = n // factor
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), t0), B: T.Buffer((n,), t1), C: T.Buffer((n,), t0)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(num_thread, thread="threadIdx.x"):
|
|
for i_1 in T.vectorized(factor):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(n, i_0 * factor + i_1)
|
|
T.reads(A[v_i], B[v_i])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[v_i] + T.Cast(t0, B[v_i])
|
|
|
|
func = tvm.compile(Module, target="cuda")
|
|
|
|
# correctness
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
low, high = (0, 20) if t0.startswith("u") or t1.startswith("u") else (-10, 10)
|
|
a_np = np.random.randint(low, high, size=n).astype(t0)
|
|
b_np = np.random.randint(low, high, size=n).astype(t1)
|
|
c_np = (a_np + b_np).astype(t0)
|
|
a_nd = tvm.runtime.tensor(a_np, dev)
|
|
b_nd = tvm.runtime.tensor(b_np, dev)
|
|
c_nd = tvm.runtime.tensor(np.zeros(c_np.shape, dtype=c_np.dtype), dev)
|
|
func(a_nd, b_nd, c_nd)
|
|
tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
def skip(t0, t1):
|
|
if t0 == t1:
|
|
return True
|
|
# CUDA does support cast between {u}int8 and fp16.
|
|
skip_set = {"float16", "uint8", "int8"}
|
|
if t0 in skip_set and t1 in skip_set:
|
|
return True
|
|
return False
|
|
|
|
types_4 = [
|
|
"float16",
|
|
"float32",
|
|
"int8",
|
|
"uint8",
|
|
"int16",
|
|
"uint16",
|
|
"int32",
|
|
"uint32",
|
|
"float64",
|
|
"int64",
|
|
"uint64",
|
|
]
|
|
types_8 = ["float16", "float32", "int8", "uint8", "int16", "uint16", "int32", "uint32"]
|
|
for t0, t1 in [(x, y) for x in types_4 for y in types_4 if not skip(x, y)]:
|
|
check(t0, t1, 4)
|
|
for t0, t1 in [(x, y) for x in types_8 for y in types_8 if not skip(x, y)]:
|
|
check(t0, t1, 8)
|
|
check("int8", "uint8", 16)
|
|
check("uint8", "int8", 16)
|
|
|
|
|
|
def sched(compute_fn, dtype, n=128):
|
|
"""Create a vectorized CUDA module with the given compute function.
|
|
|
|
The schedule structure is: split [1, None] -> split [32, None] -> split [None, 4]
|
|
then vectorize innermost, bind blockIdx.x and threadIdx.x.
|
|
For n=128 this gives: blockIdx.x=1, threadIdx.x=32, serial=1, vectorized=4.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), dtype), B: T.Buffer((n,), dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0_0 in T.thread_binding(1, thread="blockIdx.x"):
|
|
for i0_1_0 in T.thread_binding(32, thread="threadIdx.x"):
|
|
for i0_1_1_0 in range(1):
|
|
for i0_1_1_1 in T.vectorized(4):
|
|
with T.sblock("B"):
|
|
v_i0 = T.axis.spatial(n, i0_1_0 * 4 + i0_1_1_0 * 4 + i0_1_1_1)
|
|
T.reads(A[v_i0])
|
|
T.writes(B[v_i0])
|
|
B[v_i0] = compute_fn(A[v_i0])
|
|
|
|
return tvm.compile(Module, target="cuda")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_vectorized_intrin1():
|
|
test_funcs = [
|
|
(tvm.tirx.floor, lambda x: np.floor(x)),
|
|
(tvm.tirx.ceil, lambda x: np.ceil(x)),
|
|
(tvm.tirx.trunc, lambda x: np.trunc(x)),
|
|
(tvm.tirx.abs, lambda x: np.fabs(x)),
|
|
(tvm.tirx.round, lambda x: np.round(x)),
|
|
(tvm.tirx.exp, lambda x: np.exp(x)),
|
|
(tvm.tirx.exp2, lambda x: np.exp2(x)),
|
|
(tvm.tirx.exp10, lambda x: np.power(10, x)),
|
|
(tvm.tirx.log, lambda x: np.log(x)),
|
|
(tvm.tirx.log2, lambda x: np.log2(x)),
|
|
(tvm.tirx.log10, lambda x: np.log10(x)),
|
|
(tvm.tirx.tan, lambda x: np.tan(x)),
|
|
(tvm.tirx.cos, lambda x: np.cos(x)),
|
|
(tvm.tirx.cosh, lambda x: np.cosh(x)),
|
|
(tvm.tirx.sin, lambda x: np.sin(x)),
|
|
(tvm.tirx.sinh, lambda x: np.sinh(x)),
|
|
(tvm.tirx.atan, lambda x: np.arctan(x)),
|
|
(tvm.tirx.tanh, lambda x: np.tanh(x)),
|
|
(tvm.tirx.sqrt, lambda x: np.sqrt(x)),
|
|
]
|
|
|
|
def run_test(tvm_intrin, np_func, dtype):
|
|
if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
|
|
print("Skip because gpu does not have fp16 support")
|
|
return
|
|
# set of intrinsics does not support fp16 yet.
|
|
skip_set = {
|
|
tvm.tirx.abs,
|
|
tvm.tirx.round,
|
|
tvm.tirx.tan,
|
|
tvm.tirx.atan,
|
|
tvm.tirx.tanh,
|
|
tvm.tirx.cosh,
|
|
tvm.tirx.sinh,
|
|
}
|
|
if dtype == "float16" and tvm_intrin in skip_set:
|
|
print(f"Skip because '{tvm_intrin.__name__}' does not support fp16 yet")
|
|
return
|
|
|
|
n = 128
|
|
f = sched(tvm_intrin, dtype, n)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
|
|
f(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
for func in test_funcs:
|
|
run_test(*func, "float32")
|
|
run_test(*func, "float16")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_vectorized_intrin2(dtype="float32"):
|
|
c2 = tvm.tirx.const(2, dtype=dtype)
|
|
test_funcs = [
|
|
(tvm.tirx.power, lambda x: np.power(x, 2.0)),
|
|
(tvm.tirx.fmod, lambda x: np.fmod(x, 2.0)),
|
|
]
|
|
|
|
def run_test(tvm_intrin, np_func):
|
|
n = 128
|
|
f = sched(lambda x: tvm_intrin(x, c2), dtype, n)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
|
|
f(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
for func in test_funcs:
|
|
run_test(*func)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_vectorized_popcount():
|
|
def ref_popcount(x):
|
|
cnt = 0
|
|
while x:
|
|
x -= x & -x
|
|
cnt += 1
|
|
return cnt
|
|
|
|
def run_test(dtype):
|
|
n = 128
|
|
f = sched(lambda x: tvm.tirx.popcount(x), dtype, n)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.tensor(np.random.randint(0, 100000, size=n).astype(dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
|
|
f(a, b)
|
|
ref = np.vectorize(ref_popcount)(a.numpy())
|
|
tvm.testing.assert_allclose(b.numpy(), ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
run_test("uint32")
|
|
run_test("uint64")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_vectorize_load_permute_pad():
|
|
def check_cuda(dtype, n, l, padding, lanes):
|
|
if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
|
|
print("Skip because gpu does not have fp16 support")
|
|
return
|
|
|
|
zero = tvm.tirx.const(0, dtype)
|
|
dim0 = n // lanes
|
|
dim1 = l + 2 * padding
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n, l), dtype), B: T.Buffer((dim0, dim1, lanes), dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in T.thread_binding(dim0, thread="blockIdx.x"):
|
|
for j in T.thread_binding(dim1, thread="threadIdx.x"):
|
|
for k in T.vectorized(lanes):
|
|
with T.sblock("B"):
|
|
v_i, v_j, v_k = T.axis.remap("SSS", [i, j, k])
|
|
T.reads(A[v_i * lanes + v_k, v_j - padding])
|
|
T.writes(B[v_i, v_j, v_k])
|
|
B[v_i, v_j, v_k] = T.if_then_else(
|
|
v_j < padding or l + padding <= v_j,
|
|
zero,
|
|
A[v_i * lanes + v_k, v_j - padding],
|
|
)
|
|
|
|
fun = tvm.compile(Module, target="cuda")
|
|
|
|
np_a = np.random.randint(low=-128, high=127, size=(n, l)).astype(dtype)
|
|
np_a_reshape = np_a.reshape(n // lanes, lanes, l).transpose(0, 2, 1)
|
|
ref = np.pad(
|
|
np_a_reshape, ((0, 0), (padding, padding), (0, 0)), mode="constant", constant_values=0
|
|
)
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.empty((n, l), dtype, dev).copyfrom(np_a)
|
|
b = tvm.runtime.empty((dim0, dim1, lanes), dtype, dev)
|
|
fun(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_cuda("int8", 64, 16, 3, 2)
|
|
check_cuda("uint8", 64, 16, 3, 2)
|
|
check_cuda("int8", 64, 16, 3, 4)
|
|
check_cuda("uint8", 64, 16, 3, 4)
|
|
check_cuda("int32", 64, 16, 3, 4)
|
|
check_cuda("float16", 64, 16, 3, 4)
|
|
check_cuda("float32", 64, 16, 3, 4)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_try_unaligned_vector_load():
|
|
def build(N, C_N, offset):
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((N,), "float16"), C: T.Buffer((C_N,), "float16")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(C_N // 2, thread="threadIdx.x"):
|
|
for i_1 in T.vectorized(2):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(C_N, i_0 * 2 + i_1)
|
|
T.reads(A[v_i + offset])
|
|
T.writes(C[v_i])
|
|
C[v_i] = A[v_i + offset]
|
|
|
|
f = tvm.tirx.build(Module, target="cuda")
|
|
|
|
kernel_source = f.imports[0].inspect_source()
|
|
a_data = np.arange(0, N).astype("float16")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda()
|
|
a = tvm.runtime.tensor(a_data, dev)
|
|
c = tvm.runtime.tensor(np.zeros(C_N, dtype="float16"), dev)
|
|
f(a, c)
|
|
return c.numpy()
|
|
|
|
c = tvm.testing.run_with_gpu_lock(run_and_check)
|
|
return a_data, c, kernel_source
|
|
|
|
# Unaligned case: N=3, C_N=2, offset=1
|
|
a_data, c, kernel_source = build(3, 2, 1)
|
|
# (uint1*)(A + (1)) is invalid
|
|
assert "A_ptr + (1)" not in kernel_source
|
|
|
|
expected = a_data[1 : 2 + 1]
|
|
assert np.allclose(c, expected), f"expected={expected}\nactual={c}"
|
|
|
|
# Aligned case: N=4, C_N=2, offset=2
|
|
a_data, c, kernel_source = build(4, 2, 2)
|
|
# (uint1*)(A + (2)) is a valid vector load
|
|
assert "A_ptr + 2" in kernel_source
|
|
|
|
expected = a_data[2 : 2 + 2]
|
|
assert np.allclose(c, expected), f"expected={expected}\nactual={c}"
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_thread_sync_inside_condition():
|
|
@T.prim_func(s_tir=True)
|
|
def func1(A: T.Buffer((4, 4), "float32")) -> None:
|
|
A_shared = T.sblock_alloc_buffer((4, 4), "float32", scope="shared")
|
|
for bx in T.thread_binding(1, "blockIdx.x"):
|
|
for tx in T.thread_binding(32, "threadIdx.x"):
|
|
if A[0, 0] > 1.0:
|
|
for i, j in T.grid(4, 4):
|
|
A_shared[i, j] = A[i, j]
|
|
for i, j in T.grid(4, 4):
|
|
A[i, j] = A_shared[i, j] + 1.0
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func2(A: T.Buffer((4, 4), "float32")) -> None:
|
|
A_shared = T.sblock_alloc_buffer((4, 4), "float32", scope="shared")
|
|
for bx in T.thread_binding(1, "blockIdx.x"):
|
|
for tx in T.thread_binding(32, "threadIdx.x"):
|
|
if T.tvm_thread_invariant(A[0, 0] > 1.0):
|
|
for i, j in T.grid(4, 4):
|
|
A_shared[i, j] = A[i, j]
|
|
for i, j in T.grid(4, 4):
|
|
A[i, j] = A_shared[i, j] + 1.0
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func3(A: T.Buffer((4, 4), "float32")) -> None:
|
|
A_shared = T.sblock_alloc_buffer((4, 4), "float32", scope="shared")
|
|
for bx in T.thread_binding(1, "blockIdx.x"):
|
|
for tx in T.thread_binding(32, "threadIdx.x"):
|
|
while T.tvm_thread_invariant(A[0, 0] > 1.0):
|
|
for i, j in T.grid(4, 4):
|
|
A_shared[i, j] = A[i, j]
|
|
for i, j in T.grid(4, 4):
|
|
A[i, j] = A_shared[i, j] + 1.0
|
|
|
|
mod = tvm.IRModule({"main": func1})
|
|
with pytest.raises(tvm.error.InternalError):
|
|
tvm.compile(mod, target="cuda")
|
|
|
|
mod = tvm.IRModule({"main": func2})
|
|
tvm.compile(mod, target="cuda")
|
|
|
|
mod = tvm.IRModule({"main": func3})
|
|
tvm.compile(mod, target="cuda")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_invalid_reinterpret():
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((4,), "uint32"), B: T.Buffer((4,), "uint8")) -> None:
|
|
for tx in T.thread_binding(4, "threadIdx.x"):
|
|
B[tx] = T.call_intrin("uint8", "tirx.reinterpret", A[tx])
|
|
|
|
with pytest.raises(RuntimeError):
|
|
tvm.compile(func, target="cuda")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
|
def test_cuda_tensormap():
|
|
# fmt: off
|
|
@T.prim_func(s_tir=True)
|
|
def main(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapInit", A_map, "float32", 2, A.data,
|
|
16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0)
|
|
|
|
for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
if threadIdx == 0:
|
|
A[0, 0] = T.reinterpret("float64", A_map)
|
|
# fmt: on
|
|
|
|
mod = tvm.IRModule({"main": main})
|
|
mod = tvm.compile(mod, target="cuda")
|
|
assert (
|
|
"""
|
|
extern "C" __global__ void __launch_bounds__(128) main_kernel(const __grid_constant__ CUtensorMap A_map, float* __restrict__ A_ptr) {
|
|
if (((int)threadIdx.x) == 0) {
|
|
A_ptr[0] = ((float)(*(double *)(&(A_map))));
|
|
}
|
|
}""".strip()
|
|
in mod.mod.imports[0].inspect_source()
|
|
)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_device_func_call():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(a: T.float32, b: T.float32) -> T.float32:
|
|
return a + b
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((1024, 1024), "float32"),
|
|
B: T.Buffer((1024, 1024), "float32"),
|
|
C: T.Buffer((1024, 1024), "float32"),
|
|
):
|
|
for bx in T.thread_binding(1024, "blockIdx.x"):
|
|
for tx in T.thread_binding(1024, "threadIdx.x"):
|
|
C[bx, tx] = Module.add(A[bx, tx], B[bx, tx])
|
|
|
|
lib = tvm.compile(Module, target="cuda")
|
|
cuda_code = lib.mod.imports[0].inspect_source()
|
|
assert 'extern "C" __device__ float add(float a, float b) {\n return (a + b);\n}' in cuda_code
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_float_const_hex_format():
|
|
"""Test that float constants are emitted in hexadecimal format for precision"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((1024, 1024), "float32"),
|
|
):
|
|
for bx in T.thread_binding(1024, "blockIdx.x"):
|
|
for tx in T.thread_binding(1024, "threadIdx.x"):
|
|
A[bx, tx] = T.float32(1 / 27)
|
|
|
|
lib = tvm.compile(Module, target="cuda")
|
|
cuda_code = lib.mod.imports[0].inspect_source()
|
|
assert "0x1.2f684bda12f68p-5f" in cuda_code
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_device_host_call_same_func():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(a: T.int32, b: T.int32) -> T.int32:
|
|
return a + b
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
A: T.Buffer((128, 128), "int32"),
|
|
B: T.Buffer((128, 128), "int32"),
|
|
C: T.Buffer((128, 128), "int32"),
|
|
):
|
|
length: T.let[T.int32] = Module.add(64, 64) # Call from host
|
|
for bx in T.thread_binding(length, "blockIdx.x"):
|
|
for tx in T.thread_binding(length, "threadIdx.x"):
|
|
C[bx, tx] = Module.add(A[bx, tx], B[bx, tx]) # Call from device
|
|
|
|
# 1. If we set host to llvm, it will raise an error of
|
|
# "the tirx.ret should be transformed to return zero before the llvm code generation."
|
|
# Need to revisit this.
|
|
# 2. We set a dummy mcpu value for testing purpose,
|
|
# in order to avoid checking a function is host or device based on the "cpu" substring.
|
|
target = tvm.target.Target({"kind": "cuda", "mcpu": "dummy_mcpu"}, host="c")
|
|
lib = tvm.compile(Module, target=target)
|
|
cuda_code = lib.mod.imports[0].inspect_source()
|
|
assert 'extern "C" __device__ int add(int a, int b) {\n return (a + b);\n}' in cuda_code
|
|
|
|
# Run a simple test
|
|
a_np = np.random.randint(0, 10, (128, 128), dtype="int32")
|
|
b_np = np.random.randint(0, 10, (128, 128), dtype="int32")
|
|
|
|
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.empty((128, 128), dtype="int32", device=dev)
|
|
lib["main"](a_tvm, b_tvm, c_tvm)
|
|
tvm.testing.assert_allclose(c_tvm.numpy(), a_np + b_np)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_thread_return():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")):
|
|
for bx in T.thread_binding(32, "blockIdx.x"):
|
|
for tx in T.thread_binding(32, "threadIdx.x"):
|
|
if bx >= 16 or tx >= 16:
|
|
T.thread_return()
|
|
B[bx, tx] = A[bx, tx]
|
|
|
|
lib = tvm.compile(Module, target="cuda")
|
|
cuda_code = lib.mod.imports[0].inspect_source()
|
|
assert "return;" in cuda_code
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_cuda_loop_step():
|
|
@T.prim_func(s_tir=True)
|
|
def cuda_loop_step(
|
|
A: T.Buffer((1024,), "float32"),
|
|
B: T.Buffer((1024,), "float32"),
|
|
C: T.Buffer((1024,), "float32"),
|
|
):
|
|
# Each thread computes a strided subset of the i loop: start = tx*3, step = 96 (3 * 32 threads)
|
|
for bx in T.thread_binding(1, "blockIdx.x"):
|
|
for tx in T.thread_binding(96, "threadIdx.x"):
|
|
for i in T.serial(tx, 1024, step=96):
|
|
C[i] = A[i] + B[i]
|
|
|
|
target = tvm.target.Target({"kind": "cuda"})
|
|
with tvm.transform.PassContext(disabled_pass=["s_tir.CanonicalizeLoop"]):
|
|
lib = tvm.compile(cuda_loop_step, target=target)
|
|
|
|
cuda_src = lib.mod.imports[0].inspect_source()
|
|
assert "i += 96" in cuda_src
|
|
a_np = np.random.uniform(1, 100, (1024,)).astype("float32")
|
|
b_np = np.random.uniform(1, 100, (1024,)).astype("float32")
|
|
c_np = np.zeros((1024,), dtype="float32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a_nd = tvm.runtime.tensor(a_np, dev)
|
|
b_nd = tvm.runtime.tensor(b_np, dev)
|
|
c_nd = tvm.runtime.tensor(c_np, dev)
|
|
lib["main"](a_nd, b_nd, c_nd)
|
|
tvm.testing.assert_allclose(c_nd.numpy(), a_np + b_np)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
|
def test_export_load_with_fallback(monkeypatch, tmp_path):
|
|
"""Force the codegen wrapper into the fallback branch, then export+load+run."""
|
|
n = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(32, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i_0 * 32 + i_1)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i] + 1.0
|
|
|
|
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
|
|
host_lib = tvm.compile(Module, target="cuda")
|
|
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
|
|
|
|
lib_path = str(tmp_path / "lib.so")
|
|
host_lib.export_library(lib_path)
|
|
reloaded = tvm.runtime.load_module(lib_path)
|
|
|
|
a_np = np.random.uniform(size=(n,)).astype("float32")
|
|
b_np = np.zeros((n,), dtype="float32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(b_np, dev)
|
|
reloaded["main"](a, b)
|
|
np.testing.assert_allclose(b.numpy(), a_np + 1.0, rtol=1e-5)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|