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)