wehub-resource-sync
120 lines
4.6 KiB
Python
120 lines
4.6 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.
|
|
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_gpu(), reason="need gpu")
|
|
def test_large_uint_imm():
|
|
value = (1 << 63) + 123
|
|
value_const = tvm.tirx.const(value, "uint64")
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((12,), "uint64")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0_0 in T.thread_binding(6, thread="blockIdx.x"):
|
|
for i0_1 in T.thread_binding(2, thread="threadIdx.x"):
|
|
with T.sblock("A"):
|
|
v_i0 = T.axis.spatial(12, i0_0 * 2 + i0_1)
|
|
T.reads()
|
|
T.writes(A[v_i0])
|
|
A[v_i0] = value_const + T.uint64(3)
|
|
|
|
def check_target(target):
|
|
target_kind = target["kind"] if isinstance(target, dict) else target
|
|
if not tvm.testing.device_enabled(target_kind):
|
|
return
|
|
f = tvm.compile(Module, target=target)
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(target_kind, 0)
|
|
a = tvm.runtime.empty((12,), dtype="uint64", device=dev)
|
|
f(a)
|
|
assert a.numpy()[0] == value + 3
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_target("cuda")
|
|
check_target({"kind": "vulkan", "from_device": 0})
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_add_pipeline():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, B: T.Buffer((), "float32"), var_D: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n,))
|
|
D = T.match_buffer(var_D, (n,))
|
|
C = T.sblock_alloc_buffer((n,))
|
|
for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
|
|
for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
|
|
with T.sblock("C"):
|
|
v_i0 = T.axis.spatial(n, i0_0 * 256 + i0_1)
|
|
T.where(i0_0 * 256 + i0_1 < n)
|
|
T.reads(A[v_i0], B[()])
|
|
T.writes(C[v_i0])
|
|
C[v_i0] = A[v_i0] + B[()]
|
|
for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
|
|
for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
|
|
with T.sblock("D"):
|
|
v_i0 = T.axis.spatial(n, i0_0 * 256 + i0_1)
|
|
T.where(i0_0 * 256 + i0_1 < n)
|
|
T.reads(C[v_i0])
|
|
T.writes(D[v_i0])
|
|
D[v_i0] = C[v_i0] + T.float32(1.0)
|
|
|
|
def check_target(device, host):
|
|
if not tvm.testing.device_enabled(device) or not tvm.testing.device_enabled(host):
|
|
return
|
|
target = tvm.target.Target(device, host)
|
|
mhost = tvm.tirx.build(Module, target=target)
|
|
f = mhost.main
|
|
n = 1027
|
|
|
|
def run_and_check():
|
|
dev = tvm.device(device, 0)
|
|
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
|
|
b = tvm.runtime.tensor(np.random.uniform(size=()).astype("float32"), dev)
|
|
d = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
|
|
f(a, b, d)
|
|
tvm.testing.assert_allclose(d.numpy(), a.numpy() + b.numpy() + 1)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
check_target("cuda", host="llvm")
|
|
# check_target("nvptx", host="llvm") # nvptx kernel entry-point lookup not wired here
|
|
check_target("vulkan", host="llvm")
|
|
check_target("rocm", host="llvm")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_large_uint_imm()
|
|
test_add_pipeline()
|