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