def test_vm_save_and_load_without_designating_late_bound_consts():
"""Check that a VM can be saved and loaded without late-bound consts in play.
Specifically, this test ensures that the machinery behind late-bound const
loading does not assume the need to load late-bound consts (and cause an error)
when the user did not choose to designate any consts as such.
"""
target = tvm.target.Target("llvm")
dev = tvm.cpu()
const_data = np.random.rand(1).astype("float64")
x = relay.var("x", shape=(1,), dtype="float64")
const = relay.const(const_data, dtype="float64")
func = relay.Function([x], relay.op.add(x, const))
mod = tvm.IRModule.from_expr(func)
vm_exec = vm.compile(mod, target=target)
code, lib = vm_exec.save()
exe = runtime.vm.Executable.load_exec(code, lib)
x_data = np.random.rand(1).astype("float64")
loaded_vm = runtime.vm.VirtualMachine(exe, dev)
actual = loaded_vm.invoke("main", x_data)
expected = x_data + const_data
tvm.testing.assert_allclose(expected, actual.numpy())
def test_multiple_set_input(target, dev):
dtype = "float32"
in_shape = [1, 2, 3, 3]
in_data_name_0 = "d0"
in_data_name_1 = "d1"
mod = get_multiple_input_relay_mod(dtype, in_shape, in_data_name_0, in_data_name_1)
# Compile to VMExecutable.
vm_exec = vm.compile(mod, target=target)
exe = runtime.vm.VirtualMachine(vm_exec, dev)
data0_core = np.random.uniform(size=in_shape).astype(dtype)
data0 = tvm.nd.array(data0_core)
data1_core = np.random.uniform(size=in_shape).astype(dtype)
data1 = tvm.nd.array(data1_core)
ref_res_core = data0_core + data1_core
ref_res = tvm.nd.array(ref_res_core)
exe.set_input("main", data0, data1)
output = exe.invoke("main")
assert output.dtype == ref_res.dtype
tvm.testing.assert_allclose(ref_res_core, output.numpy())
data_dict = {in_data_name_1: data1, in_data_name_0: data0}
exe.set_input("main", **data_dict)
output = exe.invoke("main")
assert output.dtype == ref_res.dtype
tvm.testing.assert_allclose(ref_res_core, output.numpy())
def test_benchmark_end_to_end(target, dev):
mod, params = mlp.get_workload(1)
lib = vm.compile(mod, target=target, params=params)
exe = runtime.vm.VirtualMachine(lib, dev)
data = tvm.nd.array(np.random.rand(1, 1, 28, 28).astype("float32"), device=dev)
result = exe.benchmark(dev, data, func_name="main", repeat=2, number=1, end_to_end=True)
assert result.mean > 0
def test_load_late_bound_consts_with_no_late_bound_consts():
"""Check that load_late_bound_consts handles a model with no late bound consts."""
target = tvm.target.Target("llvm")
dev = tvm.cpu()
const_data = np.random.rand(1).astype("float64")
x = relay.var("x", shape=(1,), dtype="float64")
const = relay.const(const_data, dtype="float64")
func = relay.Function([x], relay.op.add(x, const))
mod = tvm.IRModule.from_expr(func)
vm_exec = vm.compile(mod, target=target)
temp = utils.tempdir()
path_consts = temp.relpath("consts")
path_dso = temp.relpath("lib.so")
# Ensure const_data is below the byte threshold for a late-bound const.
byte_limit = len(const_data.tobytes()) + 1
vm_exec.move_late_bound_consts(path_consts, byte_limit=byte_limit)
vm_exec.mod.export_library(path_dso)
mod = runtime.load_module(path_dso)
mod["load_late_bound_consts"](path_consts)
x_data = np.random.rand(1).astype("float64")
loaded_vm = runtime.vm.VirtualMachine(mod, dev)
actual = loaded_vm.invoke("main", x_data)
expected = x_data + const_data
tvm.testing.assert_allclose(expected, actual.numpy())
def test_benchmark(target, dev):
mod, params = mlp.get_workload(1)
lib = vm.compile(mod, target=target, params=params)
exe = runtime.vm.VirtualMachine(lib, tvm.cpu())
data = tvm.nd.array(np.random.rand(1, 1, 28, 28).astype("float32"))
result = exe.benchmark(tvm.cpu(),
data,
func_name="main",
repeat=2,
number=1)
assert result.mean == result.median
assert result.mean > 0
assert len(result.results) == 2
with patch.object(
tvm.runtime.module.Module,
"time_evaluator",
return_value=lambda x: tvm.runtime.module.BenchmarkResult(
[1, 2, 2, 5]),
) as method:
result = exe.benchmark(dev, data, func_name="main", repeat=2, number=1)
assert result.mean == 2.5
assert result.median == 2.0
assert result.max == 5
assert result.min == 1
assert result.std == 1.5
def test_large_constants():
"""Large constants can be serialized outside of executable"""
target = tvm.target.Target("llvm")
dev = tvm.cpu()
# fn(x) { add(x, <large constant>) }
x = relay.var("x", shape=(1000, 1000))
const_data = np.random.rand(1000, 1000).astype("float32")
const = relay.const(const_data, dtype="float32")
func = relay.Function([x], relay.op.add(x, const))
mod = tvm.IRModule.from_expr(func)
# Compile to executable.
vm_exec = vm.compile(mod, target=target)
# Save to constants and library files
temp = utils.tempdir()
path_consts = temp.relpath("consts")
vm_exec.move_late_bound_consts(path_consts, byte_limit=256)
path_dso = temp.relpath("lib.so")
vm_exec.mod.export_library(path_dso)
# Load library files and constants
mod = runtime.load_module(path_dso)
mod["load_late_bound_consts"](path_consts)
# Test main
x_data = np.random.rand(1000, 1000).astype("float32")
the_vm = runtime.vm.VirtualMachine(mod, dev)
actual = the_vm.invoke("main", x_data)
expected = x_data + const_data
tvm.testing.assert_allclose(expected, actual.numpy())
def test_get_input_index(target, dev):
# Build a IRModule.
data_0, data_1 = ["d1", "d2"]
x, y = [relay.var(c, shape=(10,)) for c in [data_0, data_1]]
f = relay.Function([x, y], x + y)
mod = IRModule.from_expr(f)
# Compile to VMExecutable.
vm_exec = vm.compile(mod, target=target)
vm_factory = runtime.vm.VirtualMachine(vm_exec, dev)
assert vm_factory.get_input_index(data_1) == 1
assert vm_factory.get_input_index(data_0) == 0
assert vm_factory.get_input_index("invalid") == -1
def test_get_output_multiple(target, dev):
# Build a IRModule.
x = relay.var("x", shape=(10,))
f = relay.Function([x], relay.Tuple([x + x, x]))
mod = IRModule.from_expr(f)
# Compile to VMExecutable.
vm_exec = vm.compile(mod, target=target)
vm_factory = runtime.vm.VirtualMachine(vm_exec, dev)
inp = np.ones(10, dtype="float32")
vm_factory.invoke_stateful("main", inp)
outputs = vm_factory.get_outputs()
assert len(outputs) == 2
np.testing.assert_allclose(outputs[0].numpy(), inp + inp)
np.testing.assert_allclose(outputs[1].numpy(), inp)
def test_vm_rpc():
"""
This test checks to make sure you can export a VMExecutable,
upload it to a remote machine using RPC and then execute it
on the other machine.
"""
target = tvm.target.Target("llvm --host=llvm")
# Build a IRModule.
x = relay.var("x", shape=(10, 1))
f = relay.Function([x], x + x)
mod = IRModule.from_expr(f)
# Compile to VMExecutable.
vm_exec = vm.compile(mod, target=target)
# Export to Disk
temp = utils.tempdir()
path = temp.relpath("vm_library.so")
vm_exec.mod.export_library(path)
# Use local rpc server for testing.
# Server must use popen so it doesn't inherit the current process state. It
# will crash otherwise.
server = rpc.Server("localhost", port=9120, use_popen=True)
time.sleep(2)
remote = rpc.connect(server.host, server.port, session_timeout=10)
# Upload the serialized Executable.
remote.upload(path)
# Get a handle to remote Executable.
rexec = remote.load_module("vm_library.so")
ctx = remote.cpu()
# Build a VM out of the executable and context.
vm_factory = runtime.vm.VirtualMachine(rexec, ctx)
np_input = np.random.uniform(size=(10, 1)).astype("float32")
input_tensor = tvm.nd.array(np_input, ctx)
# Invoke its "main" function.
out = vm_factory.invoke("main", input_tensor)
# Check the result.
np.testing.assert_allclose(out.asnumpy(), np_input + np_input)
# delete tensors before the server shuts down so we don't throw errors.
del input_tensor
del out
server.terminate()
def test_get_output_single():
target = tvm.target.Target("llvm")
# Build a IRModule.
x = relay.var("x", shape=(10, ))
f = relay.Function([x], x + x)
mod = IRModule.from_expr(f)
# Compile to VMExecutable.
vm_exec = vm.compile(mod, target=target)
vm_factory = runtime.vm.VirtualMachine(vm_exec, tvm.cpu())
inp = np.ones(10, dtype="float32")
vm_factory.invoke_stateful("main", inp)
outputs = vm_factory.get_outputs()
assert len(outputs) == 1
np.testing.assert_allclose(outputs[0].numpy(), inp + inp)
def test_benchmark_end_to_end_rpc():
server = rpc.Server("127.0.0.1")
remote = rpc.connect(server.host, server.port)
mod, params = mlp.get_workload(1)
lib = vm.compile(mod, target="llvm", params=params)
temp = utils.tempdir()
path = temp.relpath("vm_library.so")
lib.mod.export_library(path)
remote.upload(path)
rlib = remote.load_module("vm_library.so")
exe = runtime.vm.VirtualMachine(rlib, remote.cpu())
data = tvm.nd.array(np.random.rand(1, 1, 28, 28).astype("float32"), device=remote.cpu())
result = exe.benchmark(
remote.cpu(), data=data, func_name="main", repeat=2, number=1, end_to_end=True
)
assert result.mean > 0
def test_vm_rpc():
"""
This test checks to make sure you can export a VMExecutable,
upload it to a remote machine using RPC and then execute it
on the other machine.
"""
target = tvm.target.Target("llvm --host=llvm")
# Build a IRModule.
x = relay.var("x", shape=(10, 1))
f = relay.Function([x], x + x)
mod = IRModule.from_expr(f)
# Compile to VMExecutable.
vm_exec = vm.compile(mod, target=target)
# Export to Disk
temp = utils.tempdir()
path = temp.relpath("vm_library.so")
vm_exec.mod.export_library(path)
# Use LocalRPC for testing.
remote = rpc.LocalSession()
# Upload the serialized Executable.
remote.upload(path)
# Get a handle to remote Executable.
rexec = remote.load_module("vm_library.so")
ctx = remote.cpu()
# Build a VM out of the executable and context.
vm_factory = runtime.vm.VirtualMachine(rexec, ctx)
np_input = np.random.uniform(size=(10, 1)).astype("float32")
input_tensor = tvm.nd.array(np_input, ctx)
# Invoke its "main" function.
out = vm_factory.invoke("main", [input_tensor])
# Check the result.
np.testing.assert_allclose(out.asnumpy(), np_input + np_input)
