chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,578 @@
|
||||
# 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 sys
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
exec_mode = tvm.testing.parameter("bytecode", "compiled")
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class InputModule:
|
||||
@R.function
|
||||
def foo(x: R.Tensor(("m", "n"), "int64")):
|
||||
y = R.unique(x, sorted=False)
|
||||
y_sorted = R.unique(x)
|
||||
return y, y_sorted
|
||||
|
||||
|
||||
def run_cpu(mod, func_name, *args, exec_mode):
|
||||
if isinstance(mod, relax.Function):
|
||||
func = mod
|
||||
args = [func_name, *args]
|
||||
func_name = func.attrs["global_symbol"]
|
||||
mod = tvm.IRModule.from_expr(func)
|
||||
|
||||
target = tvm.target.Target("llvm")
|
||||
ex = relax.build(mod, target, exec_mode=exec_mode)
|
||||
vm = relax.VirtualMachine(ex, tvm.cpu())
|
||||
|
||||
return vm[func_name](*args)
|
||||
|
||||
|
||||
def test_unique(exec_mode):
|
||||
# TODO(prakalp): also add test for compiling and running on CUDA device.
|
||||
data_numpy = np.random.randint(0, 16, (16, 16))
|
||||
data = tvm.runtime.tensor(data_numpy)
|
||||
result, result_sorted = run_cpu(InputModule, "foo", data, exec_mode=exec_mode)
|
||||
|
||||
expected_output_sorted, indices = np.unique(data_numpy, return_index=True)
|
||||
expected_output = [data_numpy.flatten()[index] for index in sorted(indices)]
|
||||
|
||||
np.testing.assert_array_equal(expected_output_sorted, result_sorted.numpy())
|
||||
np.testing.assert_array_equal(expected_output, result.numpy())
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class PrintTest:
|
||||
@R.function(pure=False)
|
||||
def foo(x: R.Tensor((), "int32")):
|
||||
# results have to be bound, but we don't use them
|
||||
# TODO: We should allow calls whose results are not bound for side effects;
|
||||
# it would be easy syntactic sugar to add.
|
||||
p1 = R.print(x)
|
||||
p2 = R.print(x, format="Number: {}")
|
||||
t = (x, x)
|
||||
p3 = R.print(t, format="Tuple: {}")
|
||||
p4 = R.print(x, t)
|
||||
p5 = R.print(x, x, format="Custom print: {} {}")
|
||||
p6 = R.print(x, t, format="Another print: {} {}")
|
||||
return x
|
||||
|
||||
|
||||
def test_print(exec_mode):
|
||||
try:
|
||||
stdout = sys.stdout
|
||||
with tempfile.TemporaryFile(mode="w+") as test_out:
|
||||
sys.stdout = test_out
|
||||
run_cpu(
|
||||
PrintTest,
|
||||
"foo",
|
||||
tvm.runtime.tensor(np.array(1).astype("int32")),
|
||||
exec_mode=exec_mode,
|
||||
)
|
||||
test_out.seek(0)
|
||||
printed_text = str(test_out.read())
|
||||
expected = "1\nNumber: 1\nTuple: (1, 1)\n1 (1, 1)\nCustom print: 1 1\nAnother print: 1 (1, 1)\n"
|
||||
assert printed_text in expected, ("printed_text is ", printed_text)
|
||||
finally:
|
||||
sys.stdout = stdout
|
||||
|
||||
|
||||
def test_assert_passes(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(relax.const(True))
|
||||
return x
|
||||
|
||||
run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")), exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_passes_with_format_args(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(relax.const(True), x, format="You won't see me")
|
||||
return x
|
||||
|
||||
run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")), exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_fails(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(relax.const(False))
|
||||
return x
|
||||
|
||||
with pytest.raises(AssertionError, match="Assertion Failed"):
|
||||
run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")), exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_fails_with_message(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(relax.const(False), format="I failed...")
|
||||
return x
|
||||
|
||||
with pytest.raises(AssertionError, match="I failed..."):
|
||||
run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")), exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_fails_with_args(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(relax.const(False), [x, x])
|
||||
return x
|
||||
|
||||
with pytest.raises(AssertionError, match="5, 5"):
|
||||
run_cpu(func, tvm.runtime.tensor(np.array(5).astype("int32")), exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_fails_with_formatted_args(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(relax.const(False), x, format="Number: {}")
|
||||
return x
|
||||
|
||||
with pytest.raises(AssertionError, match="Number: 6"):
|
||||
run_cpu(func, tvm.runtime.tensor(np.array(6).astype("int32")), exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_on_argument_passes(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(condition: R.Tensor((), "bool"), x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(condition)
|
||||
return x
|
||||
|
||||
condition = tvm.runtime.tensor(np.array(True))
|
||||
x = tvm.runtime.tensor(np.array(5).astype("int32"))
|
||||
run_cpu(func, condition, x, exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_on_argument_fails(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(condition: R.Tensor((), "bool"), x: R.Tensor((), "int32")):
|
||||
_ = R.assert_op(condition)
|
||||
return x
|
||||
|
||||
condition = tvm.runtime.tensor(np.array(False))
|
||||
x = tvm.runtime.tensor(np.array(5).astype("int32"))
|
||||
with pytest.raises(AssertionError):
|
||||
run_cpu(func, condition, x, exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_on_symbolic_var_passes(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor(["N"], "int32")):
|
||||
N = T.int64()
|
||||
_ = R.assert_op(R.prim_value(N % 8 == 0))
|
||||
return x
|
||||
|
||||
x = tvm.runtime.tensor(np.arange(8, dtype="int32"))
|
||||
run_cpu(func, x, exec_mode=exec_mode)
|
||||
|
||||
|
||||
def test_assert_on_symbolic_var_fails(exec_mode):
|
||||
@R.function(pure=False)
|
||||
def func(x: R.Tensor(["N"], "int32")):
|
||||
N = T.int64()
|
||||
_ = R.assert_op(R.prim_value(N % 8 == 0))
|
||||
return x
|
||||
|
||||
x = tvm.runtime.tensor(np.arange(10, dtype="int32"))
|
||||
with pytest.raises(AssertionError):
|
||||
run_cpu(func, x, exec_mode=exec_mode)
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class ShapeOfTest:
|
||||
@R.function
|
||||
def get_shape(t: R.Tensor(ndim=-1, dtype="int32")) -> R.Shape(ndim=-1):
|
||||
return R.shape_of(t)
|
||||
|
||||
@R.function
|
||||
def get_constrained_shape(t: R.Tensor(ndim=1, dtype="int32")) -> R.Shape(ndim=1):
|
||||
# require the input tensor to have rank 1
|
||||
return R.shape_of(t)
|
||||
|
||||
@R.function
|
||||
def get_scalar_shape() -> R.Shape(()):
|
||||
x: R.Tensor((), "int32") = R.const(1, dtype="int32")
|
||||
return R.shape_of(x)
|
||||
|
||||
@R.function
|
||||
def get_constant_shape() -> R.Shape((2, 2)):
|
||||
x: R.Tensor((2, 2), "int32") = R.const(
|
||||
np.array([[1, 2], [3, 4]], dtype="int32"), dtype="int32"
|
||||
)
|
||||
return R.shape_of(x)
|
||||
|
||||
|
||||
def test_op_shape_of(exec_mode):
|
||||
unit_shape = run_cpu(ShapeOfTest, "get_scalar_shape", exec_mode=exec_mode)
|
||||
assert unit_shape == tvm_ffi.Shape([])
|
||||
|
||||
const_shape = run_cpu(ShapeOfTest, "get_constant_shape", exec_mode=exec_mode)
|
||||
assert const_shape == tvm_ffi.Shape([2, 2])
|
||||
|
||||
scalar_shape = run_cpu(
|
||||
ShapeOfTest,
|
||||
"get_shape",
|
||||
tvm.runtime.tensor(np.array(1, dtype="int32")),
|
||||
exec_mode=exec_mode,
|
||||
)
|
||||
assert scalar_shape == tvm_ffi.Shape([])
|
||||
|
||||
tensor_shape = run_cpu(
|
||||
ShapeOfTest,
|
||||
"get_shape",
|
||||
tvm.runtime.tensor(np.zeros((1, 2, 3)).astype("int32")),
|
||||
exec_mode=exec_mode,
|
||||
)
|
||||
assert tensor_shape == tvm_ffi.Shape([1, 2, 3])
|
||||
|
||||
constrained_shape = run_cpu(
|
||||
ShapeOfTest,
|
||||
"get_constrained_shape",
|
||||
tvm.runtime.tensor(np.zeros((1,)).astype("int32")),
|
||||
exec_mode=exec_mode,
|
||||
)
|
||||
assert constrained_shape == tvm_ffi.Shape([1])
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class ShapeToTensorTest:
|
||||
@R.function
|
||||
def const_shape(shape: R.Shape(ndim=-1)) -> R.Tensor(ndim=-1):
|
||||
return R.shape_to_tensor(shape)
|
||||
|
||||
@R.function
|
||||
def symbolic_shape(shape: R.Shape(("m", "n"))) -> R.Tensor(ndim=-1):
|
||||
m = T.int64()
|
||||
n = T.int64()
|
||||
return R.shape_to_tensor(shape)
|
||||
|
||||
|
||||
def test_op_shape_to_tensor(exec_mode):
|
||||
# Check type
|
||||
isinstance(ShapeToTensorTest["const_shape"].body.ty, tvm.relax.TensorType)
|
||||
assert ShapeToTensorTest["const_shape"].body.ty.ndim == 1
|
||||
isinstance(ShapeToTensorTest["symbolic_shape"].body.ty, tvm.relax.TensorType)
|
||||
assert ShapeToTensorTest["symbolic_shape"].body.ty.ndim == 1
|
||||
|
||||
# Check its functionality
|
||||
out2d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 2]), exec_mode=exec_mode)
|
||||
assert isinstance(out2d, tvm.runtime.Tensor)
|
||||
assert np.array_equal(out2d.numpy(), np.array([3, 2]))
|
||||
|
||||
out3d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 3, 2]), exec_mode=exec_mode)
|
||||
assert isinstance(out3d, tvm.runtime.Tensor)
|
||||
assert np.array_equal(out3d.numpy(), np.array([3, 3, 2]))
|
||||
|
||||
out4d = run_cpu(
|
||||
ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 3, 2, 2]), exec_mode=exec_mode
|
||||
)
|
||||
assert isinstance(out4d, tvm.runtime.Tensor)
|
||||
assert np.array_equal(out4d.numpy(), np.array([3, 3, 2, 2]))
|
||||
|
||||
outs = run_cpu(ShapeToTensorTest, "symbolic_shape", tvm_ffi.Shape([3, 2]), exec_mode=exec_mode)
|
||||
assert isinstance(outs, tvm.runtime.Tensor)
|
||||
assert np.array_equal(outs.numpy(), np.array([3, 2]))
|
||||
|
||||
|
||||
def test_op_call_pure_packed(exec_mode):
|
||||
@tvm.script.ir_module
|
||||
class CallPureTest:
|
||||
@R.function
|
||||
def pure_copy(x: R.Tensor((3, 4), "float32")):
|
||||
z = R.call_pure_packed(
|
||||
"vm.builtin.copy", x, ty_args=(R.Tensor((3, 4), dtype="float32"))
|
||||
)
|
||||
return z
|
||||
|
||||
np.random.seed(0) # to avoid flakiness
|
||||
arr = np.random.rand(3, 4).astype("float32")
|
||||
copy_found = run_cpu(CallPureTest, "pure_copy", tvm.runtime.tensor(arr), exec_mode=exec_mode)
|
||||
assert (copy_found.numpy() == arr).all()
|
||||
|
||||
|
||||
def test_op_call_inplace_packed(exec_mode):
|
||||
# in this case we can use the same test as above
|
||||
@tvm.script.ir_module
|
||||
class CallInplaceTest:
|
||||
@R.function
|
||||
def pure_copy(x: R.Tensor((3, 4), "float32")):
|
||||
z = R.call_inplace_packed(
|
||||
"vm.builtin.copy",
|
||||
x,
|
||||
inplace_indices=0,
|
||||
ty_args=(R.Tensor((3, 4), dtype="float32")),
|
||||
)
|
||||
return z
|
||||
|
||||
@tvm.register_global_func("test.inplace.add", override=True)
|
||||
def inplace_add(a, b):
|
||||
arr_a = a.numpy()
|
||||
arr_b = b.numpy()
|
||||
for i in range(len(arr_a)):
|
||||
for j in range(len(arr_a[i])):
|
||||
arr_a[i][j] = arr_a[i][j] + arr_b[i][j]
|
||||
a.copyfrom(arr_a)
|
||||
return a
|
||||
|
||||
@tvm.script.ir_module
|
||||
class CallInplaceAddTest:
|
||||
@R.function
|
||||
def inplace_add(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")):
|
||||
z = R.call_inplace_packed(
|
||||
"test.inplace.add",
|
||||
x,
|
||||
y,
|
||||
inplace_indices=0,
|
||||
ty_args=(R.Tensor((3, 4), dtype="float32")),
|
||||
)
|
||||
return z
|
||||
|
||||
np.random.seed(1) # to avoid flakiness
|
||||
arr_a = np.random.rand(3, 4).astype("float32")
|
||||
arr_b = np.random.rand(3, 4).astype("float32")
|
||||
sum = arr_a + arr_b
|
||||
tvm_arr_a = tvm.runtime.tensor(arr_a)
|
||||
result = run_cpu(
|
||||
CallInplaceAddTest, "inplace_add", tvm_arr_a, tvm.runtime.tensor(arr_b), exec_mode=exec_mode
|
||||
)
|
||||
assert result == tvm_arr_a
|
||||
assert (result.numpy() == sum).all()
|
||||
|
||||
@tvm.register_global_func("test.inplace.tuple_add", override=True)
|
||||
def inplace_tuple_add(a, b):
|
||||
arr_a = a.numpy()
|
||||
arr_b = b.numpy()
|
||||
c = tvm.runtime.tensor(arr_a + arr_b)
|
||||
for i in range(len(arr_a)):
|
||||
for j in range(len(arr_a[i])):
|
||||
arr_a[i][j] = arr_a[i][j] + arr_b[i][j]
|
||||
a.copyfrom(arr_a)
|
||||
return tvm.runtime.convert([a, c])
|
||||
|
||||
@tvm.script.ir_module
|
||||
class CallInplaceTuple:
|
||||
@R.function
|
||||
def inplace_tuple(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")):
|
||||
z = R.call_inplace_packed(
|
||||
"test.inplace.tuple_add",
|
||||
x,
|
||||
y,
|
||||
inplace_indices=[0, -1],
|
||||
ty_args=(R.Tensor((3, 4), dtype="float32"), R.Tensor((3, 4), dtype="float32")),
|
||||
)
|
||||
return z
|
||||
|
||||
np.random.seed(2) # to avoid flakiness
|
||||
arr_a = np.random.rand(3, 4).astype("float32")
|
||||
arr_b = np.random.rand(3, 4).astype("float32")
|
||||
sum = arr_a + arr_b
|
||||
tvm_arr_a = tvm.runtime.tensor(arr_a)
|
||||
tvm_arr_b = tvm.runtime.tensor(arr_b)
|
||||
result = run_cpu(CallInplaceTuple, "inplace_tuple", tvm_arr_a, tvm_arr_b, exec_mode=exec_mode)
|
||||
assert result[0] == tvm_arr_a
|
||||
assert (result[0].numpy() == sum).all()
|
||||
assert result[1] != tvm_arr_a and result[1] != tvm_arr_b
|
||||
assert (result[1].numpy() == sum).all()
|
||||
|
||||
|
||||
def test_op_call_py_func(exec_mode):
|
||||
"""Test R.call_py_func operator functionality."""
|
||||
import torch
|
||||
|
||||
def torch_relu(x):
|
||||
if isinstance(x, tvm.runtime.Tensor):
|
||||
x_torch = torch.from_numpy(x.numpy())
|
||||
elif hasattr(x, "asnumpy"):
|
||||
x_torch = torch.from_numpy(x.asnumpy())
|
||||
else:
|
||||
x_np = np.array(x)
|
||||
if isinstance(x_np, tvm.runtime.Tensor):
|
||||
x_torch = torch.from_numpy(x_np.numpy())
|
||||
elif len(x_np) > 0 and isinstance(x_np[0], tvm.runtime.Tensor):
|
||||
x_torch = torch.from_numpy(np.array([t.numpy() for t in x_np]))
|
||||
if x_torch.ndim > 1:
|
||||
x_torch = x_torch.flatten()
|
||||
else:
|
||||
x_torch = torch.from_numpy(x_np)
|
||||
result = torch.relu(x_torch)
|
||||
return tvm.runtime.tensor(result.numpy())
|
||||
|
||||
def torch_sigmoid(x):
|
||||
if isinstance(x, tvm.runtime.Tensor):
|
||||
x_torch = torch.from_numpy(x.numpy())
|
||||
elif hasattr(x, "asnumpy"):
|
||||
x_torch = torch.from_numpy(x.asnumpy())
|
||||
else:
|
||||
x_np = np.array(x)
|
||||
if isinstance(x_np, tvm.runtime.Tensor):
|
||||
x_torch = torch.from_numpy(x_np.numpy())
|
||||
elif len(x_np) > 0 and isinstance(x_np[0], tvm.runtime.Tensor):
|
||||
x_torch = torch.from_numpy(np.array([t.numpy() for t in x_np]))
|
||||
if x_torch.ndim > 1:
|
||||
x_torch = x_torch.flatten()
|
||||
else:
|
||||
x_torch = torch.from_numpy(x_np)
|
||||
result = torch.sigmoid(x_torch)
|
||||
return tvm.runtime.tensor(result.numpy())
|
||||
|
||||
register_func = tvm.get_global_func("vm.builtin.register_py_func")
|
||||
register_func("torch_relu", torch_relu)
|
||||
register_func("torch_sigmoid", torch_sigmoid)
|
||||
|
||||
@tvm.script.ir_module
|
||||
class CallPyFuncTest:
|
||||
@R.function
|
||||
def simple_call(x: R.Tensor((3,), "float32")):
|
||||
result = R.call_py_func(R.str("torch_relu"), (x,), out_ty=R.Tensor((3,), "float32"))
|
||||
return result
|
||||
|
||||
@R.function
|
||||
def multiple_calls(x: R.Tensor((2,), "float32")):
|
||||
y = R.call_py_func(R.str("torch_relu"), (x,), out_ty=R.Tensor((2,), "float32"))
|
||||
z = R.call_py_func(R.str("torch_sigmoid"), (y,), out_ty=R.Tensor((2,), "float32"))
|
||||
return z
|
||||
|
||||
np.random.seed(0)
|
||||
x_data = np.array([-1.0, 0.0, 1.0], dtype=np.float32)
|
||||
x_tvm = tvm.runtime.tensor(x_data)
|
||||
|
||||
result = run_cpu(CallPyFuncTest, "simple_call", x_tvm, exec_mode=exec_mode)
|
||||
expected = np.maximum(x_data, 0.0)
|
||||
assert (result.numpy() == expected).all()
|
||||
|
||||
y_data = np.array([-0.5, 0.5], dtype=np.float32)
|
||||
y_tvm = tvm.runtime.tensor(y_data)
|
||||
|
||||
result2 = run_cpu(CallPyFuncTest, "multiple_calls", y_tvm, exec_mode=exec_mode)
|
||||
expected2 = 1.0 / (1.0 + np.exp(-np.maximum(y_data, 0.0)))
|
||||
assert (result2.numpy() == expected2).all()
|
||||
|
||||
clear_func = tvm.get_global_func("vm.builtin.clear_py_func_registry")
|
||||
clear_func()
|
||||
|
||||
|
||||
def test_op_to_device(exec_mode):
|
||||
@tvm.script.ir_module
|
||||
class CallToDevice:
|
||||
@R.function
|
||||
def to_dev(x: R.Tensor((3, 4), "float32")):
|
||||
z = R.call_pure_packed(
|
||||
"vm.builtin.to_device",
|
||||
x,
|
||||
1,
|
||||
0,
|
||||
ty_args=(R.Tensor((3, 4), dtype="float32")),
|
||||
)
|
||||
return z
|
||||
|
||||
np.random.seed(0) # to avoid flakiness
|
||||
arr = np.random.rand(3, 4).astype("float32")
|
||||
copy_found = run_cpu(CallToDevice, "to_dev", tvm.runtime.tensor(arr), exec_mode=exec_mode)
|
||||
assert (copy_found.numpy() == arr).all()
|
||||
|
||||
|
||||
def test_op_to_vdevice(exec_mode):
|
||||
@tvm.script.ir_module
|
||||
class ToVDevice:
|
||||
I.module_global_infos({"vdevice": [I.vdevice("llvm")]})
|
||||
|
||||
@R.function
|
||||
def to_vdev(x: R.Tensor((3, 4), "float32")):
|
||||
dst_vdev = tvm.ir.VDevice("llvm", 0, "global")
|
||||
ret = R.to_vdevice(x, "llvm")
|
||||
return ret
|
||||
|
||||
np.random.seed(0)
|
||||
arr = np.random.rand(3, 4).astype("float32")
|
||||
copy_found = run_cpu(ToVDevice, "to_vdev", tvm.runtime.tensor(arr), exec_mode=exec_mode)
|
||||
assert (copy_found.numpy() == arr).all()
|
||||
|
||||
|
||||
def test_scalar_tensor_as_branch_condition(exec_mode):
|
||||
"""The condition of a branch may be a scalar tensor"""
|
||||
|
||||
@R.function
|
||||
def func(condition: R.Tensor((), "bool")):
|
||||
if condition:
|
||||
out = R.prim_value(5)
|
||||
else:
|
||||
out = R.prim_value(10)
|
||||
return out
|
||||
|
||||
res = run_cpu(func, tvm.runtime.tensor(np.array(True)), exec_mode=exec_mode)
|
||||
assert res == 5
|
||||
|
||||
res = run_cpu(func, tvm.runtime.tensor(np.array(False)), exec_mode=exec_mode)
|
||||
assert res == 10
|
||||
|
||||
|
||||
def test_prim_value_as_branch_condition(exec_mode):
|
||||
"""The condition may be a Expr"""
|
||||
|
||||
@R.function
|
||||
def func(condition: R.Prim("bool")):
|
||||
if condition:
|
||||
out = R.prim_value(5)
|
||||
else:
|
||||
out = R.prim_value(10)
|
||||
return out
|
||||
|
||||
res = run_cpu(func, True, exec_mode=exec_mode)
|
||||
assert res == 5
|
||||
|
||||
res = run_cpu(func, False, exec_mode=exec_mode)
|
||||
assert res == 10
|
||||
|
||||
|
||||
def test_computed_prim_value_as_branch_condition(exec_mode):
|
||||
"""The R.Prim condition may be computed within the function"""
|
||||
|
||||
@R.function
|
||||
def func(x: R.Tensor(["N"], "int64")):
|
||||
N = T.int64()
|
||||
if R.prim_value(N % 16 == 0):
|
||||
out = R.prim_value(5)
|
||||
else:
|
||||
out = R.prim_value(10)
|
||||
return out
|
||||
|
||||
res = run_cpu(func, tvm.runtime.tensor(np.arange(16)), exec_mode=exec_mode)
|
||||
assert res == 5
|
||||
|
||||
res = run_cpu(func, tvm.runtime.tensor(np.arange(20)), exec_mode=exec_mode)
|
||||
assert res == 10
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user