wehub-resource-sync
299 lines
9.5 KiB
Python
299 lines
9.5 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.
|
|
"""Lowest level testing VM. Test execbuilder and execution."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.relax.testing.vm import check_saved_func
|
|
from tvm.script import relax as R
|
|
|
|
|
|
def test_vm_execute():
|
|
ib = relax.ExecBuilder()
|
|
with ib.function("func0", num_inputs=2):
|
|
ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
|
|
ib.emit_ret(ib.r(2))
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
a = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
b = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
|
|
add_res = check_saved_func(vm, "func0", a, b)
|
|
tvm.testing.assert_allclose(add_res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
|
|
|
|
|
|
def test_vm_multiple_func():
|
|
ib = relax.ExecBuilder()
|
|
with ib.function("func0", num_inputs=2):
|
|
ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
|
|
ib.emit_ret(ib.r(2))
|
|
with ib.function("func1", num_inputs=2):
|
|
ib.emit_call("test.vm.mul", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
|
|
ib.emit_ret(ib.r(2))
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
a = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
b = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
mul_res = check_saved_func(vm, "func1", a, b)
|
|
add_res = check_saved_func(vm, "func0", a, b)
|
|
tvm.testing.assert_allclose(add_res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
|
|
tvm.testing.assert_allclose(mul_res.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)
|
|
|
|
|
|
def test_vm_checker():
|
|
ib = relax.ExecBuilder()
|
|
with pytest.raises(RuntimeError):
|
|
with ib.function("func0", num_inputs=2):
|
|
ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(2)], dst=ib.r(2))
|
|
ib.emit_ret(ib.r(2))
|
|
ib.get()
|
|
|
|
|
|
def test_neg_imm():
|
|
ib = relax.ExecBuilder()
|
|
|
|
with ib.function("func0", num_inputs=1):
|
|
ib.emit_call("test.vm.add_scalar", args=[ib.imm(-3), ib.r(0)], dst=ib.r(1))
|
|
ib.emit_ret(ib.r(1))
|
|
ib.get()
|
|
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
assert vm["func0"](1) == -2
|
|
assert vm["func0"](-3) == -6
|
|
|
|
|
|
def test_emit_cache():
|
|
ib = relax.ExecBuilder()
|
|
|
|
with ib.function("func0", num_inputs=1):
|
|
x0 = ib.convert_constant("str0")
|
|
x1 = ib.convert_constant("str0")
|
|
# cache constant str
|
|
assert x0 == x1
|
|
s0 = ib.convert_constant(tvm_ffi.Shape([1, 2]))
|
|
s1 = ib.convert_constant(tvm_ffi.Shape([1, 2]))
|
|
s2 = ib.convert_constant(tvm_ffi.Shape([1, 3]))
|
|
assert s0 == s1
|
|
assert s1 != s2
|
|
y0 = ib.convert_constant(tvm.runtime.tensor(np.array([1, 2, 3]).astype("int32")))
|
|
y1 = ib.convert_constant(tvm.runtime.tensor(np.array([1, 2, 3]).astype("int32")))
|
|
assert y0 == y1
|
|
ib.emit_ret(ib.r(0))
|
|
|
|
|
|
def test_vm_formalize():
|
|
ib0 = relax.ExecBuilder()
|
|
ib1 = relax.ExecBuilder()
|
|
with ib0.function("func0", num_inputs=2):
|
|
ib0.emit_call("test.vm.add", args=[ib0.r(0), ib0.r(1)], dst=ib0.r(100))
|
|
ib0.emit_call("test.vm.mul", args=[ib0.r(1), ib0.r(100)], dst=ib0.r(50))
|
|
ib0.emit_ret(ib0.r(50))
|
|
with ib1.function("func0", num_inputs=2):
|
|
ib1.emit_call("test.vm.add", args=[ib1.r(0), ib1.r(1)], dst=ib1.r(2))
|
|
ib1.emit_call("test.vm.mul", args=[ib1.r(1), ib1.r(2)], dst=ib1.r(3))
|
|
ib1.emit_ret(ib1.r(3))
|
|
exec0 = ib0.get()
|
|
exec1 = ib1.get()
|
|
assert exec0.as_text() == exec1.as_text()
|
|
|
|
|
|
def test_vm_operand():
|
|
ib0 = relax.ExecBuilder()
|
|
with ib0.function("func0", num_inputs=2):
|
|
ib0.emit_call("test.vm.add_scalar", args=[ib0.r(0), ib0.r(1)], dst=ib0.r(2))
|
|
ib0.emit_ret(ib0.r(2))
|
|
exec0 = ib0.get()
|
|
vm = relax.VirtualMachine(exec0, tvm.cpu())
|
|
res = vm["func0"](2, 3)
|
|
assert res == 5
|
|
|
|
ib1 = relax.ExecBuilder()
|
|
with ib1.function("func1", num_inputs=1):
|
|
ib1.emit_call("test.vm.get_device_id", args=[ib1.r(0)], dst=ib1.r(1))
|
|
ib1.emit_ret(ib1.r(1))
|
|
exec1 = ib1.get()
|
|
vm = relax.VirtualMachine(exec1, tvm.cpu())
|
|
res = vm["func1"](tvm.cpu(3))
|
|
assert res == 3
|
|
|
|
|
|
def test_vm_shapeof():
|
|
ib = relax.ExecBuilder()
|
|
shape = (32, 16)
|
|
arr = tvm.runtime.tensor(np.random.rand(*shape))
|
|
with ib.function("main", num_inputs=0):
|
|
ib.emit_call("vm.builtin.shape_of", args=[arr], dst=ib.r(0))
|
|
ib.emit_ret(ib.r(0))
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
res = vm["main"]()
|
|
for i, s in enumerate(res):
|
|
assert s == shape[i]
|
|
|
|
|
|
def test_vm_storage():
|
|
dtype = tvm.DataType("float32")
|
|
shape = (4, 6)
|
|
ib = relax.ExecBuilder()
|
|
with ib.function("main", num_inputs=0):
|
|
ib.emit_call(
|
|
"vm.builtin.alloc_storage",
|
|
args=[
|
|
ib.vm_state(),
|
|
(24,),
|
|
ib.convert_constant(0),
|
|
dtype,
|
|
ib.convert_constant("global"),
|
|
],
|
|
dst=ib.r(1),
|
|
)
|
|
ib.emit_call(
|
|
"vm.builtin.alloc_tensor", args=[ib.r(1), ib.imm(0), shape, dtype], dst=ib.r(2)
|
|
)
|
|
ib.emit_ret(ib.r(2))
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
res = vm["main"]()
|
|
assert res.device == tvm.cpu()
|
|
assert res.shape == shape
|
|
|
|
|
|
def test_vm_goto():
|
|
ib = relax.ExecBuilder()
|
|
with ib.function("main", num_inputs=2):
|
|
ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
|
|
ib.emit_goto(2)
|
|
ib.emit_call("test.vm.mul", args=[ib.r(2), ib.r(1)], dst=ib.r(2))
|
|
ib.emit_ret(ib.r(2))
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
a = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
b = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
res = check_saved_func(vm, "main", a, b)
|
|
tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
|
|
|
|
|
|
def test_vm_if():
|
|
ib = relax.ExecBuilder()
|
|
with ib.function("main", num_inputs=3):
|
|
ib.emit_if(ib.r(0), 3)
|
|
ib.emit_call("test.vm.add", args=[ib.r(1), ib.r(2)], dst=ib.r(3))
|
|
ib.emit_goto(2)
|
|
ib.emit_call("test.vm.mul", args=[ib.r(1), ib.r(2)], dst=ib.r(3))
|
|
ib.emit_ret(ib.r(3))
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
a = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
b = tvm.runtime.tensor(
|
|
np.random.rand(
|
|
4,
|
|
)
|
|
)
|
|
res = vm["main"](0, a, b)
|
|
tvm.testing.assert_allclose(res.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)
|
|
res = vm["main"](1, a, b)
|
|
tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)
|
|
|
|
|
|
def test_vm_invoke_closure():
|
|
ib = relax.ExecBuilder()
|
|
with ib.function("lifted_func_1", num_inputs=4):
|
|
ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(4))
|
|
ib.emit_call("test.vm.add", args=[ib.r(2), ib.r(4)], dst=ib.r(5))
|
|
ib.emit_call("test.vm.add", args=[ib.r(3), ib.r(5)], dst=ib.r(6))
|
|
ib.emit_ret(ib.r(6))
|
|
with ib.function("main", num_inputs=2):
|
|
ib.emit_call(
|
|
"vm.builtin.make_closure", args=[ib.f("lifted_func_1"), ib.r(0), ib.r(1)], dst=ib.r(2)
|
|
)
|
|
ib.emit_ret(ib.r(2))
|
|
|
|
ex = ib.get()
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
w_inp = tvm.runtime.tensor(np.random.rand(2, 3))
|
|
x_inp = tvm.runtime.tensor(np.random.rand(2, 3))
|
|
y_inp = tvm.runtime.tensor([[3.1, 4.0, 5.0], [6.0, 7.1, 9.0]])
|
|
z_inp = tvm.runtime.tensor(np.random.rand(2, 3))
|
|
clo = vm["main"](w_inp, x_inp)
|
|
res = vm.invoke_closure(clo, y_inp, z_inp)
|
|
tvm.testing.assert_allclose(
|
|
res.numpy(), w_inp.numpy() + x_inp.numpy() + y_inp.numpy() + z_inp.numpy()
|
|
)
|
|
|
|
|
|
def test_vm_stack_restore_after_failure():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(inp: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.multiply(inp, R.const(2, "float32"))
|
|
gv: R.Tensor((10, 10), dtype="float32") = lv
|
|
R.output(gv)
|
|
return gv
|
|
|
|
ex = tvm.compile(Module, "llvm")
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
|
|
correct_input = tvm.runtime.tensor(np.random.normal(size=(10, 10)).astype("float32"))
|
|
incorrect_input = tvm.runtime.tensor(np.random.normal(size=(12, 10)).astype("float32"))
|
|
|
|
try:
|
|
vm["main"](incorrect_input)
|
|
except RuntimeError:
|
|
pass
|
|
|
|
# VM should executes correctly after encountered incorrect shape in previous invocation
|
|
vm["main"](correct_input)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|