wehub-resource-sync
700 lines
30 KiB
Python
700 lines
30 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: E501, F841
|
|
|
|
import re
|
|
|
|
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.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import ir as I_builder
|
|
from tvm.script.ir_builder import tirx as T_builder
|
|
from tvm.testing import env
|
|
|
|
dtype = tvm.testing.parameter("float32", "int32", "float16", "int8")
|
|
fuzz_seed = tvm.testing.parameter(range(25))
|
|
|
|
|
|
# Explicitly specify a target, as this test is looking at the
|
|
# generated shader code, and is not running on an actual device.
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled(
|
|
{
|
|
"kind": "vulkan",
|
|
"supports_int8": 1,
|
|
"supports_8bit_buffer": 1,
|
|
"supports_storage_buffer_storage_class": 1,
|
|
"supports_float16": 1,
|
|
"supports_16bit_buffer": 1,
|
|
}
|
|
),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vector_comparison(dtype):
|
|
target = {
|
|
"kind": "vulkan",
|
|
"supports_int8": 1,
|
|
"supports_8bit_buffer": 1,
|
|
"supports_storage_buffer_storage_class": 1,
|
|
"supports_float16": 1,
|
|
"supports_16bit_buffer": 1,
|
|
}
|
|
target = tvm.target.Target(target)
|
|
zero = tvm.tirx.const(0, dtype)
|
|
one = tvm.tirx.const(1, dtype)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((1024,), dtype), B: T.Buffer((1024,), dtype)):
|
|
for i_0 in T.thread_binding(8, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(32, thread="threadIdx.x"):
|
|
for i_2 in T.vectorized(4):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(1024, i_0 * 128 + i_1 * 4 + i_2)
|
|
B[v_i] = T.Select(A[v_i] >= zero, A[v_i] + one, zero)
|
|
|
|
# Build
|
|
f = tvm.tirx.build(Module, target=target)
|
|
|
|
# Verify we generate the boolx4 type declaration and the OpSelect
|
|
# v4{float,half,int} instruction
|
|
assembly = f.imports[0].inspect_source()
|
|
matches = re.findall("%v4bool = OpTypeVector %bool 4", assembly)
|
|
assert len(matches) == 1
|
|
matches = re.findall("OpSelect %v4.*", assembly)
|
|
assert len(matches) == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"target",
|
|
[
|
|
"llvm",
|
|
pytest.param("cuda", marks=pytest.mark.gpu),
|
|
pytest.param("rocm", marks=pytest.mark.gpu),
|
|
pytest.param("vulkan", marks=pytest.mark.gpu),
|
|
pytest.param("metal", marks=pytest.mark.gpu),
|
|
pytest.param("opencl", marks=pytest.mark.gpu),
|
|
pytest.param("nvptx", marks=pytest.mark.gpu),
|
|
],
|
|
)
|
|
def test_array_copy(target, dtype, fuzz_seed):
|
|
if not tvm.testing.device_enabled(target):
|
|
pytest.skip(f"{target} not enabled")
|
|
np.random.seed(fuzz_seed)
|
|
|
|
log_arr_size = np.random.uniform(low=np.log(1), high=np.log(32768))
|
|
arr_size = np.exp(log_arr_size).astype(int)
|
|
a_np = np.random.uniform(size=(arr_size,)).astype(dtype)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target)
|
|
a = tvm.runtime.empty((arr_size,), dtype, dev).copyfrom(a_np)
|
|
tvm.testing.assert_allclose(a_np, a.numpy())
|
|
|
|
if target == "llvm":
|
|
run_and_check()
|
|
else:
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_array_vectorize_add(dtype):
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
target = tvm.target.Target(target)
|
|
arr_size = 64
|
|
lanes = 2
|
|
|
|
vec_dtype = f"{dtype}x{lanes}"
|
|
one = tvm.tirx.const(1, vec_dtype)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((64,), vec_dtype), B: T.Buffer((64,), vec_dtype)):
|
|
for i_0 in T.thread_binding(16, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(4, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(64, i_0 * 4 + i_1)
|
|
B[v_i] = A[v_i] + one
|
|
|
|
f = tvm.compile(Module, target=target)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target.kind.name)
|
|
a = tvm.runtime.empty((arr_size,), vec_dtype, dev).copyfrom(
|
|
np.random.uniform(size=(arr_size, lanes))
|
|
)
|
|
c = tvm.runtime.empty((arr_size,), vec_dtype, dev)
|
|
f(a, c)
|
|
tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vulkan_bool_load():
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
target = tvm.target.Target(target)
|
|
arr_size = 1024
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((1024,), "bool"), B: T.Buffer((1024,), "int32")):
|
|
for i_0 in T.thread_binding(8, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(128, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(1024, i_0 * 128 + i_1)
|
|
B[v_i] = T.Cast("int32", A[v_i])
|
|
|
|
f = tvm.compile(Module, target=target)
|
|
|
|
a_np = np.random.uniform(size=arr_size) > 0.5
|
|
b_np = np.zeros((arr_size,), dtype="int32")
|
|
ref = a_np.astype(np.int32)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target.kind.name)
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(b_np, dev)
|
|
f(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), ref)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
vulkan_parameter_impl = tvm.testing.parameter("push_constants", "ubo")
|
|
vulkan_parameter_dtype = tvm.testing.parameter("int32", "float32", "int64")
|
|
|
|
|
|
# Only run on vulkan because extremely large numbers of input
|
|
# parameters can crash cuda/llvm compiler.
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vulkan_constant_passing(vulkan_parameter_impl, vulkan_parameter_dtype):
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
target = tvm.target.Target(target)
|
|
dtype = vulkan_parameter_dtype
|
|
|
|
if not target.attrs.get("supports_int64", False):
|
|
pytest.xfail("Vulkan target does not support Int64 variables")
|
|
|
|
# f_add has 3+num_int_params scalar parameters. The other three
|
|
# are length_n, stride1, and stride2.
|
|
if vulkan_parameter_impl == "push_constants":
|
|
# 4 params, 32 bytes. Within 128-byte spec-guaranteed size of
|
|
# push constants. Uses push constants.
|
|
num_int_params = 1
|
|
else:
|
|
# 24 params, 192 bytes. May be above spec-guaranteed size of 128
|
|
# bytes for push constants. Uses either push constants or UBO,
|
|
# depending on the device.
|
|
max_push_constants_size = int(target.attrs.get("max_push_constants_size", 128))
|
|
max_int_params_in_push = max_push_constants_size // 8 - 3
|
|
num_int_params = max_int_params_in_push + 1
|
|
|
|
# Build IRModule programmatically since num_int_params is dynamic
|
|
with IRBuilder() as ib:
|
|
with I_builder.ir_module():
|
|
with T_builder.prim_func():
|
|
T_builder.func_name("main")
|
|
scalar_vars = []
|
|
for i in range(num_int_params):
|
|
v = T_builder.arg(f"scale{i}", tvm.tirx.Var("", dtype))
|
|
scalar_vars.append(v)
|
|
var_A = T_builder.arg("var_A", T_builder.handle())
|
|
var_B = T_builder.arg("var_B", T_builder.handle())
|
|
T_builder.func_attr({"tirx.noalias": True})
|
|
n_var = T_builder.int32()
|
|
A = T_builder.match_buffer(var_A, (n_var,), dtype)
|
|
B = T_builder.match_buffer(var_B, (n_var,), dtype)
|
|
scalar_sum = scalar_vars[0]
|
|
for s in scalar_vars[1:]:
|
|
scalar_sum = scalar_sum + s
|
|
with T_builder.thread_binding(
|
|
tvm.tirx.ceildiv(n_var, 64), thread="blockIdx.x"
|
|
) as i_0:
|
|
with T_builder.thread_binding(64, thread="threadIdx.x") as i_1:
|
|
with T_builder.sblock("B"):
|
|
v_i = T_builder.axis.spatial(n_var, i_0 * 64 + i_1)
|
|
T_builder.where(i_0 * 64 + i_1 < n_var)
|
|
T_builder.reads(A[v_i])
|
|
T_builder.writes(B[v_i])
|
|
T_builder.buffer_store(B, scalar_sum + A[v_i], [v_i])
|
|
mod = ib.get()
|
|
f_add = tvm.compile(mod, target=target)
|
|
|
|
n = 1024
|
|
scalars = np.array([1 for _ in range(num_int_params)]).astype(dtype)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target.kind.name)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=n).astype(dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
|
|
f_add(*scalars, a, b)
|
|
tvm.testing.assert_allclose(a.numpy() + sum(scalars), b.numpy())
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vulkan_while_if():
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
target = tvm.target.Target(target)
|
|
n = 1
|
|
dtype = "int32"
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def while_if_gpu(A: T.Buffer((1,), "int32"), B: T.Buffer((1,), "int32")):
|
|
for bx in T.thread_binding(1, thread="blockIdx.x"):
|
|
iterations = T.decl_buffer((1,), "int32", scope="local")
|
|
iterations[0] = 0
|
|
B[0] = 0
|
|
while iterations[0] < T.if_then_else(A[0] > 0, 10, 20):
|
|
iterations[0] = iterations[0] + 1
|
|
B[0] = B[0] + iterations[0]
|
|
|
|
mod = tvm.IRModule.from_expr(while_if_gpu.with_attr("target", target))
|
|
compiled_func = tvm.compile(mod, target=target)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target.kind.name)
|
|
for input_value, expected in [(5, [55]), (-5, [210])]:
|
|
a = tvm.runtime.tensor(np.array([input_value], dtype=dtype), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
|
|
compiled_func(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), expected)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vulkan_local_threadidx():
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
target = tvm.target.Target(target)
|
|
n = 32
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def local_threadidx_func(A: T.Buffer((32,), "int32"), B: T.Buffer((32,), "int32")):
|
|
# First block with thread extent 16
|
|
for _ in range(1):
|
|
for tx in T.thread_binding(16, thread="threadIdx.x"):
|
|
B[tx + 0] = A[tx + 0]
|
|
# Second block with thread extent 16
|
|
for _ in range(1):
|
|
for tx in T.thread_binding(16, thread="threadIdx.x"):
|
|
B[tx + 16] = A[tx + 16]
|
|
|
|
mod = tvm.IRModule.from_expr(local_threadidx_func)
|
|
func = tvm.compile(mod, target=target)
|
|
|
|
a_np = np.arange(n).astype(dtype="int32")
|
|
b_np = np.zeros((n,), dtype="int32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target.kind.name)
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(b_np, dev)
|
|
func(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), a_np)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vectorized_index_ramp():
|
|
"""Test vectorized copy with ramp indices (load N values, write to N locations)"""
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
n = 4
|
|
ramp_index = tvm.tirx.Ramp(0, 1, 4)
|
|
|
|
@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})
|
|
A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
|
|
B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
|
|
with T.sblock("compute"):
|
|
T.reads()
|
|
T.writes()
|
|
bx = T.launch_thread("blockIdx.x", 1)
|
|
B[ramp_index] = A[ramp_index]
|
|
|
|
f = tvm.compile(Module, target=target)
|
|
|
|
a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
|
|
b_np = np.zeros(n, dtype="int32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target["kind"])
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(b_np, dev)
|
|
f(a, b)
|
|
tvm.testing.assert_allclose(b.numpy(), a_np)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_vectorized_index_broadcast():
|
|
"""Test broadcast index (load 1 value, write to N locations)"""
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
n = 4
|
|
broadcast_index = tvm.tirx.Broadcast(0, 4)
|
|
ramp_index = tvm.tirx.Ramp(0, 1, 4)
|
|
|
|
@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})
|
|
A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
|
|
B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
|
|
with T.sblock("compute"):
|
|
T.reads()
|
|
T.writes()
|
|
bx = T.launch_thread("blockIdx.x", 1)
|
|
# Load from broadcast index (single element), store to ramp index
|
|
B[ramp_index] = A[broadcast_index]
|
|
|
|
f = tvm.compile(Module, target=target)
|
|
|
|
a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
|
|
b_np = np.zeros(n, dtype="int32")
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target["kind"])
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(b_np, dev)
|
|
f(a, b)
|
|
# All elements of b should be a[0] (broadcast load)
|
|
tvm.testing.assert_allclose(b.numpy(), np.full(n, a_np[0]))
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(
|
|
not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
|
|
reason="vulkan not enabled",
|
|
)
|
|
def test_negative_operand_divmod():
|
|
"""Test handling of negative offsets to floormod/floordiv
|
|
|
|
Even though the SPIR-V spec states that OpSRem and OpSMod can give
|
|
the signed modulo, the Vulkan spec states that any use of negative
|
|
operands is undefined behavior. This test starts with negative
|
|
operands to floordiv, validating that they are simplified into the
|
|
corresponding positive operands, such that the final TIR can be
|
|
expressed using only positive operands.
|
|
|
|
SPIR-V: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpSRem
|
|
Vulkan: https://registry.khronos.org/vulkan/specs/1.3/html/chap37.html#spirvenv-op-prec
|
|
"""
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
|
|
N = 32
|
|
offset = 16
|
|
divisor = 5
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((N, 2), "int32")):
|
|
for i in T.thread_binding(N, thread="threadIdx.x"):
|
|
with T.sblock("A"):
|
|
v_i = T.axis.spatial(N, i)
|
|
A[v_i, 0] = T.floordiv(v_i - offset, divisor)
|
|
A[v_i, 1] = T.floormod(v_i - offset, divisor)
|
|
|
|
built = tvm.compile(func, target=target)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target["kind"])
|
|
a_dev = tvm.runtime.empty([N, 2], "int32", dev)
|
|
built(a_dev)
|
|
a = a_dev.numpy()
|
|
np.testing.assert_array_equal(a[:, 0], (np.arange(N) - offset) // divisor)
|
|
np.testing.assert_array_equal(a[:, 1], (np.arange(N) - offset) % divisor)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.parametrize("out_dtype", ["float32", "float16"])
|
|
def test_cooperative_matrix(out_dtype):
|
|
M, N, K = 16, 16, 32
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(X: T.Buffer((16, 32), "float16"), W: T.Buffer((32, 16), "float16"), compute: T.Buffer((16, 16), out_dtype)):
|
|
T.func_attr({"tirx.noalias": True})
|
|
X_shared = T.sblock_alloc_buffer((16, 32), "float16", scope="shared")
|
|
W_shared = T.sblock_alloc_buffer((32, 16), "float16", scope="shared")
|
|
X_shared_wmma_matrix_a = T.sblock_alloc_buffer((16, 32), "float16", scope="wmma.matrix_a")
|
|
W_shared_wmma_matrix_b = T.sblock_alloc_buffer((32, 16), "float16", scope="wmma.matrix_b")
|
|
compute_wmma_accumulator = T.sblock_alloc_buffer((16, 16), out_dtype, scope="wmma.accumulator")
|
|
for i_0_j_0_fused in T.thread_binding(1, thread="blockIdx.x"):
|
|
with T.sblock("compute_init_o"):
|
|
v_i_o = T.axis.spatial(1, 0)
|
|
v_j_o = T.axis.spatial(1, 0)
|
|
T.reads()
|
|
T.writes(compute_wmma_accumulator[0:16, 0:16])
|
|
C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
|
T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0.0))
|
|
for k_0 in range(2):
|
|
for ax0_ax1_fused_0 in range(2):
|
|
for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
|
|
for ax0_ax1_fused_2 in T.vectorized(4):
|
|
with T.sblock("X_shared"):
|
|
v0 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
|
|
v1 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
|
|
T.reads(X[v0, v1])
|
|
T.writes(X_shared[v0, v1])
|
|
X_shared[v0, v1] = X[v0, v1]
|
|
for ax0_ax1_fused_0 in range(2):
|
|
for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
|
|
for ax0_ax1_fused_2 in T.vectorized(4):
|
|
with T.sblock("W_shared"):
|
|
v0 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
|
|
v1 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
|
|
T.reads(W[v0, v1])
|
|
T.writes(W_shared[v0, v1])
|
|
W_shared[v0, v1] = W[v0, v1]
|
|
for ax0_0 in T.unroll(1):
|
|
for ax1_0 in T.unroll(1):
|
|
with T.sblock("X_shared_wmma.matrix_a_o"):
|
|
v0_o = T.axis.spatial(1, ax0_0)
|
|
v1_o = T.axis.spatial(2, k_0 + ax1_0)
|
|
T.reads(X_shared[0:16, v1_o * 16:v1_o * 16 + 16])
|
|
T.writes(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16])
|
|
A = T.match_buffer(X_shared[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
|
|
C = T.match_buffer(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
|
|
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
|
|
for ax0_0 in T.unroll(1):
|
|
for ax1_0 in T.unroll(1):
|
|
with T.sblock("W_shared_wmma.matrix_b_o"):
|
|
v0_o = T.axis.spatial(2, k_0 + ax0_0)
|
|
v1_o = T.axis.spatial(1, ax1_0)
|
|
T.reads(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16])
|
|
T.writes(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16])
|
|
A = T.match_buffer(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
|
|
C = T.match_buffer(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
|
|
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
|
|
with T.sblock("compute_update_o"):
|
|
v_i_o = T.axis.spatial(1, 0)
|
|
v_j_o = T.axis.spatial(1, 0)
|
|
v_k_o = T.axis.reduce(2, k_0)
|
|
T.reads(compute_wmma_accumulator[0:16, 0:16], X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16])
|
|
T.writes(compute_wmma_accumulator[0:16, 0:16])
|
|
A = T.match_buffer(X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
|
|
B = T.match_buffer(W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16], (16, 16), "float16", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
|
|
C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
|
T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
|
|
with T.sblock("compute_wmma.accumulator_o"):
|
|
v0_o = T.axis.spatial(1, 0)
|
|
v1_o = T.axis.spatial(1, 0)
|
|
T.reads(compute_wmma_accumulator[0:16, 0:16])
|
|
T.writes(compute[0:16, 0:16])
|
|
A = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
|
|
C = T.match_buffer(compute[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), offset_factor=16)
|
|
T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(out_dtype), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
|
|
# fmt: on
|
|
|
|
target = {"kind": "vulkan", "from_device": 0}
|
|
tgt_attrs = tvm.target.Target(target).attrs
|
|
|
|
if tgt_attrs.get("supports_cooperative_matrix"):
|
|
f = tvm.compile(Module, target=target)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device("vulkan", 0)
|
|
A = tvm.runtime.tensor(np.random.randn(M, K).astype("float16"), dev)
|
|
B = tvm.runtime.tensor(np.random.randn(K, N).astype("float16"), dev)
|
|
C = tvm.runtime.tensor(np.random.randn(M, N).astype(out_dtype), dev)
|
|
f(A, B, C)
|
|
A_np = A.numpy()
|
|
B_np = B.numpy()
|
|
ref = np.dot(A_np.astype("float32"), B_np.astype("float32"))
|
|
tvm.testing.assert_allclose(C.numpy(), ref, rtol=1e-2, atol=1e-2)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
|
|
def test_codegen_decl_buffer():
|
|
"""The codegen should accept DeclBuffer nodes in its input"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def kernel():
|
|
T.func_attr({"calling_conv": 2, "global_symbol": "kernel", "tirx.noalias": True})
|
|
A = T.alloc_buffer((256,), dtype="float32", scope="local")
|
|
A_buf = T.decl_buffer([256], dtype="float32", scope="local", data=A.data)
|
|
|
|
target = tvm.target.Target("vulkan")
|
|
vulkan_codegen = tvm.get_global_func("target.build.vulkan")
|
|
vulkan_codegen(Module, target)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
|
|
def test_codegen_static_shared_memory():
|
|
"""The codegen should accept static shared/workgroup allocations."""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
|
|
A_shared = T.alloc_buffer((128,), dtype="float32", scope="shared")
|
|
|
|
for bx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for tx in T.thread_binding(128, thread="threadIdx.x"):
|
|
A_shared[tx] = A[tx]
|
|
B[tx] = A_shared[tx]
|
|
|
|
tvm.compile(Module, target="vulkan")
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
|
|
def test_unary():
|
|
test_funcs = [
|
|
(tvm.tirx.sin, lambda x: np.sin(x)),
|
|
(tvm.tirx.cos, lambda x: np.cos(x)),
|
|
(tvm.tirx.tan, lambda x: np.tan(x)),
|
|
(tvm.tirx.sinh, lambda x: np.sinh(x)),
|
|
(tvm.tirx.cosh, lambda x: np.cosh(x)),
|
|
(tvm.tirx.tanh, lambda x: np.tanh(x)),
|
|
(tvm.tirx.asin, lambda x: np.arcsin(x)),
|
|
(tvm.tirx.acos, lambda x: np.arccos(x)),
|
|
(tvm.tirx.atan, lambda x: np.arctan(x)),
|
|
(tvm.tirx.asinh, lambda x: np.arcsinh(x)),
|
|
(tvm.tirx.acosh, lambda x: np.arccosh(x)),
|
|
(tvm.tirx.atanh, lambda x: np.arctanh(x)),
|
|
]
|
|
|
|
def run_test(tvm_intrin, np_func):
|
|
n = 16
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle):
|
|
m = T.int32()
|
|
A = T.match_buffer(var_A, (m,), "float32")
|
|
B = T.match_buffer(var_B, (m,), "float32")
|
|
for i_0 in T.thread_binding((m + 63) // 64, thread="blockIdx.x"):
|
|
for i_1 in T.thread_binding(64, thread="threadIdx.x"):
|
|
with T.sblock("B"):
|
|
v_i = T.axis.spatial(m, i_0 * 64 + i_1)
|
|
T.where(i_0 * 64 + i_1 < m)
|
|
T.reads(A[v_i])
|
|
T.writes(B[v_i])
|
|
B[v_i] = tvm_intrin(A[v_i])
|
|
|
|
target = tvm.target.Target("vulkan")
|
|
func = tvm.compile(Module, target=target)
|
|
|
|
if tvm_intrin in [tvm.tirx.asin, tvm.tirx.acos]:
|
|
data = np.random.uniform(-1.0, 1.0, size=n)
|
|
elif tvm_intrin == tvm.tirx.atanh:
|
|
data = np.random.uniform(-0.999, 0.999, size=n)
|
|
elif tvm_intrin == tvm.tirx.acosh:
|
|
data = np.random.uniform(1.0, 5.0, size=n)
|
|
else:
|
|
data = np.random.uniform(0.1, 0.9, size=n)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target.kind.name, 0)
|
|
a = tvm.runtime.tensor(data.astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
func(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_vulkan(), reason="need vulkan")
|
|
def test_export_load_with_fallback(monkeypatch, tmp_path):
|
|
"""Force the codegen wrapper into the fallback branch, then export."""
|
|
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="vulkan")
|
|
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
|
|
|
|
lib_path = str(tmp_path / "lib.so")
|
|
host_lib.export_library(lib_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|