wehub-resource-sync
222 lines
7.9 KiB
Python
222 lines
7.9 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.
|
|
# ruff: noqa: F841
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
|
|
def test_rocm_inf_nan():
|
|
def check_inf_nan(n, value, dtype):
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), 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(128, thread="threadIdx.x"):
|
|
with T.sblock("C"):
|
|
v_i = T.axis.spatial(1, i_0 * 128 + i_1)
|
|
T.where(i_0 * 128 + i_1 < 1)
|
|
T.reads()
|
|
T.writes(C[v_i])
|
|
C[v_i] = T.Cast(dtype, value)
|
|
|
|
fun = tvm.compile(Module, "rocm")
|
|
|
|
def run_and_check():
|
|
dev = tvm.rocm(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.gpu
|
|
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
|
|
def test_rocm_copy():
|
|
def check_rocm(dtype, n):
|
|
def run_and_check():
|
|
dev = tvm.rocm(0)
|
|
a_np = np.random.uniform(size=(n,)).astype(dtype)
|
|
a = tvm.runtime.empty((n,), dtype, dev).copyfrom(a_np)
|
|
b_np = a.numpy()
|
|
tvm.testing.assert_allclose(a_np, b_np)
|
|
tvm.testing.assert_allclose(a_np, a.numpy())
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
for _ in range(100):
|
|
dtype = np.random.choice(["float32", "float16", "int8", "int32"])
|
|
logN = np.random.randint(1, 15)
|
|
peturb = np.random.uniform(low=0.5, high=1.5)
|
|
check_rocm(dtype, int(peturb * (2**logN)))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
|
|
def test_rocm_vectorize_add():
|
|
def check_rocm(dtype, n, lanes):
|
|
vec_dtype = f"{dtype}x{lanes}"
|
|
num_blocks = n // 4
|
|
|
|
@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(4, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(n, i_0 * 4 + i_1)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = A[v_i] + T.Broadcast(T.Cast(dtype, 1), lanes)
|
|
|
|
fun = tvm.compile(Module, target="rocm")
|
|
|
|
def run_and_check():
|
|
dev = tvm.rocm(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_rocm("float32", 64, 2)
|
|
check_rocm("float16", 64, 2)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
|
|
def test_rocm_warp_shuffle():
|
|
@T.prim_func(s_tir=True)
|
|
def func(
|
|
A_handle: T.handle,
|
|
):
|
|
A = T.match_buffer(A_handle, (32,), dtype="float32")
|
|
|
|
for bx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for tx in T.thread_binding(32, thread="threadIdx.x"):
|
|
with T.sblock("test"):
|
|
A_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
|
|
mask = T.sblock_alloc_buffer((1,), "uint32", scope="local")
|
|
t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
|
|
|
|
A_local[0] = A[tx]
|
|
A_local[0] = T.tvm_warp_shuffle(mask[0], A_local[0], 0, 32, 32)
|
|
A[tx] = A_local[0]
|
|
|
|
mod = tvm.compile(func, target="rocm")
|
|
|
|
def run_and_check():
|
|
dev = tvm.rocm(0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=(32,)).astype("float32"), dev)
|
|
mod(a)
|
|
tvm.testing.assert_allclose(a.numpy(), np.ones((32,)) * a.numpy()[0])
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
|
|
def test_rocm_vectorized_exp():
|
|
@T.prim_func(s_tir=True)
|
|
def func(
|
|
A_handle: T.handle,
|
|
B_handle: T.handle,
|
|
):
|
|
A = T.match_buffer(A_handle, (4,), dtype="float32")
|
|
B = T.match_buffer(B_handle, (4,), dtype="float32")
|
|
|
|
for bx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for tx in T.thread_binding(1, thread="threadIdx.x"):
|
|
with T.sblock("test"):
|
|
for i in T.vectorized(0, 4):
|
|
B[i] = T.exp2(A[i])
|
|
|
|
mod = tvm.compile(func, target="rocm")
|
|
|
|
def run_and_check():
|
|
dev = tvm.rocm(0)
|
|
a = tvm.runtime.tensor(np.ones((4,)).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros((4,)).astype("float32"), dev)
|
|
mod(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), np.exp2(a.numpy()))
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
|
|
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="rocm")
|
|
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.rocm(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()
|