def test_multi_targets():
# Build an IRModule.
n = 10
x = relay.var("x", shape=(n,))
y = relay.var("y", shape=(n,))
z = relay.var("z", shape=(n,))
f = relay.Function([x, y, z], x + relay.op.annotation.on_device(y + z, tvm.cpu()))
mod = IRModule.from_expr(f)
# Compile to VMExecutable.
with tvm.transform.PassContext(
opt_level=3, config={"relay.fallback_device_type": tvm.cuda().device_type}
):
exe = relay.vm.compile(
mod, target={"cpu": tvm.target.Target("llvm"), "cuda": tvm.target.Target("cuda")}
)
# Run
vm = runtime.vm.VirtualMachine(exe, [tvm.cuda(), tvm.cpu()])
x_data = np.random.rand(
n,
).astype("float32")
y_data = np.random.rand(
n,
).astype("float32")
z_data = np.random.rand(
n,
).astype("float32")
actual_result = vm.invoke("main", x_data, y_data, z_data)
# Test
expected_result = x_data + y_data + z_data
tvm.testing.assert_allclose(actual_result.numpy(), expected_result)
def veval(f, *args, device=tvm.cpu(), target="llvm"):
if isinstance(f, relay.Expr):
mod = tvm.IRModule()
mod["main"] = f
else:
assert isinstance(f, tvm.IRModule), "expected expression or module"
mod = f
exe = relay.vm.compile(mod, target)
vm = runtime.vm.VirtualMachine(exe, device)
return vm.invoke("main", *args)