def test_rpc_file_exchange():
if not tvm.runtime.enabled("rpc"):
return
server = rpc.Server("localhost")
remote = rpc.connect(server.host, server.port)
blob = bytearray(np.random.randint(0, 10, size=(10)))
remote.upload(blob, "dat.bin")
rev = remote.download("dat.bin")
assert(rev == blob)
def test_rpc_return_func():
server = rpc.Server(key="x1")
client = rpc.connect("127.0.0.1", server.port, key="x1")
def check_remote():
f1 = client.get_function("rpc.test.add_to_lhs")
fadd = f1(10)
assert fadd(12) == 22
check_remote()
def test_rpc_return_func():
@tvm.register_func("rpc.test.remote_func")
def addone(x):
return lambda y: x + y
server = rpc.Server("localhost", key="x1")
client = rpc.connect(server.host, server.port, key="x1")
f1 = client.get_function("rpc.test.remote_func")
fadd = f1(10)
assert fadd(12) == 22
def test_rpc_array():
x = np.ones((3, 4))
server = rpc.Server()
remote = rpc.connect("127.0.0.1", server.port)
r_cpu = tvm.nd.array(x, remote.cpu(0))
assert str(r_cpu.device).startswith("remote")
np.testing.assert_equal(r_cpu.asnumpy(), x)
fremote = remote.get_function("rpc.test.remote_array_func")
fremote(r_cpu)
def test_rpc_array():
x = np.ones((3, 4))
server = rpc.Server("localhost")
remote = rpc.connect(server.host, server.port)
r_cpu = tvm.nd.array(x, remote.cpu(0))
assert str(r_cpu.context).startswith("remote")
np.testing.assert_equal(r_cpu.asnumpy(), x)
fremote = remote.get_function("rpc.test.remote_array_func")
fremote(r_cpu)
def verify_rpc_gpu_remove_package_params(obj_format):
if not tvm.runtime.enabled("cuda"):
print("Skip because cuda is not enabled")
return
mod, params = relay.testing.resnet.get_workload(num_layers=18)
with relay.build_config(opt_level=3):
complied_graph_lib = relay.build_module.build(mod,
"cuda",
params=params)
from tvm.contrib import util
temp = util.tempdir()
if obj_format == ".so":
file_name = "deploy_lib.so"
else:
assert obj_format == ".tar"
file_name = "deploy_lib.tar"
path_lib = temp.relpath(file_name)
complied_graph_lib_no_params = complied_graph_lib["remove_params"]()
complied_graph_lib_no_params.export_library(path_lib)
path_params = temp.relpath("deploy_param.params")
with open(path_params, "wb") as fo:
fo.write(relay.save_param_dict(complied_graph_lib.get_params()))
from tvm import rpc
server = rpc.Server("localhost", use_popen=True)
remote = rpc.connect(server.host, server.port)
remote.upload(path_lib)
loaded_lib = remote.load_module(path_lib)
data = np.random.uniform(-1, 1,
size=(1, 3, 224, 224)).astype("float32")
ctx = remote.gpu()
# raw api
gmod = loaded_lib['default'](ctx)
set_input = gmod["set_input"]
run = gmod["run"]
get_output = gmod["get_output"]
load_params = gmod["load_params"]
loaded_params = bytearray(open(path_params, "rb").read())
set_input("data", tvm.nd.array(data, ctx=ctx))
load_params(loaded_params)
run()
out = get_output(0).asnumpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
# graph runtime wrapper
gmod = graph_runtime.GraphModule(loaded_lib['default'](ctx))
loaded_params = bytearray(open(path_params, "rb").read())
gmod.set_input("data", data)
gmod.load_params(loaded_params)
gmod.run()
out = gmod.get_output(0).asnumpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
def test_estimate_peak_fma_flops_rpc():
target = "llvm -mattr=+fma,+avx2"
server = rpc.Server(key="profiling")
remote = rpc.connect("127.0.0.1", server.port, key="profiling")
dev = remote.device(target)
flops = tvm.utils.estimate_peak_fma_flops(tvm.target.Target(target),
dev,
remote=remote)
# Assume we can achieve 1 GFLOP/s per thread, which is 1 FLOP per cycle on a 1GHz cpu.
assert (flops > 10**9 and flops < 10**14
), f"FLOP/s should be between 10^9 and 10^14, but it is {flops}"
def test_rpc_runtime_string():
server = rpc.Server(key="x1")
client = rpc.connect("127.0.0.1", server.port, key="x1")
def check_remote():
func = client.get_function("rpc.test.runtime_str_concat")
x = tvm.runtime.container.String("abc")
y = tvm.runtime.container.String("def")
assert str(func(x, y)) == "abcdef"
check_remote()
def test_rpc_file_exchange():
server = rpc.Server()
remote = rpc.connect("127.0.0.1", server.port)
def check_remote():
blob = bytearray(np.random.randint(0, 10, size=(10)))
remote.upload(blob, "dat.bin")
rev = remote.download("dat.bin")
assert rev == blob
check_remote()
def test_rpc_large_array():
# testcase of large array creation
server = rpc.Server()
remote = rpc.connect("127.0.0.1", server.port)
dev = remote.cpu(0)
a_np = np.ones((5041, 720)).astype("float32")
b_np = np.ones((720, 192)).astype("float32")
a = tvm.nd.array(a_np, dev)
b = tvm.nd.array(b_np, dev)
np.testing.assert_equal(a.asnumpy(), a_np)
np.testing.assert_equal(b.asnumpy(), b_np)
def test_rpc_large_array():
# testcase of large array creation
server = rpc.Server("localhost")
remote = rpc.connect(server.host, server.port)
ctx = remote.cpu(0)
a_np = np.ones((5041, 720)).astype('float32')
b_np = np.ones((720, 192)).astype('float32')
a = tvm.nd.array(a_np, ctx)
b = tvm.nd.array(b_np, ctx)
np.testing.assert_equal(a.asnumpy(), a_np)
np.testing.assert_equal(b.asnumpy(), b_np)
def test_rpc_simple():
server = rpc.Server(key="x1")
client = rpc.connect("127.0.0.1", server.port, key="x1")
f1 = client.get_function("rpc.test.addone")
assert f1(10) == 11
f3 = client.get_function("rpc.test.except")
with pytest.raises(tvm._ffi.base.TVMError):
f3("abc")
f2 = client.get_function("rpc.test.strcat")
assert f2("abc", 11) == "abc:11"
def test_rpc_simple():
server = rpc.Server("localhost", key="x1")
client = rpc.connect(server.host, server.port, key="x1")
f1 = client.get_function("rpc.test.addone")
assert f1(10) == 11
f3 = client.get_function("rpc.test.except")
with pytest.raises(tvm.error.RPCError):
f3("abc")
f2 = client.get_function("rpc.test.strcat")
assert f2("abc", 11) == "abc:11"
def test_rpc_tracker_via_proxy():
"""
tracker
/ \
Host -- Proxy -- RPC server
"""
device_key = "test_device"
tracker_server = Tracker(port=9000, port_end=9100)
proxy_server = Proxy(
host=tracker_server.host,
port=8888,
port_end=8988,
tracker_addr=(tracker_server.host, tracker_server.port),
)
server1 = rpc.Server(
host=proxy_server.host,
port=proxy_server.port,
key=device_key,
tracker_addr=(tracker_server.host, tracker_server.port),
is_proxy=True,
)
server2 = rpc.Server(
host=proxy_server.host,
port=proxy_server.port,
key=device_key,
tracker_addr=(tracker_server.host, tracker_server.port),
is_proxy=True,
)
client = rpc.connect_tracker(tracker_server.host, tracker_server.port)
remote1 = client.request(device_key, session_timeout=30) # pylint: disable=unused-variable
remote2 = client.request(device_key, session_timeout=30) # pylint: disable=unused-variable
server2.terminate()
server1.terminate()
proxy_server.terminate()
tracker_server.terminate()
def test_estimate_peak_bandwidth_rpc():
target = "llvm -mattr=+fma,+avx2"
server = rpc.Server(key="profiling")
remote = rpc.connect("127.0.0.1", server.port, key="profiling")
dev = remote.device(target)
bandwidth = tvm.utils.estimate_peak_bandwidth(tvm.target.Target(target),
dev,
remote=remote)
# Assume we can achieve 1 GB/s. DDR2 should transfer somewhere around 6
# GB/s, so this should leave enough wiggle room.
assert (
bandwidth > 10**9 and bandwidth < 10**12
), f"Bandwidth should be between 10^9 and 10^12, but it is {bandwidth}"
def test_rpc_echo():
def check(remote):
fecho = remote.get_function("testing.echo")
assert fecho(1, 2, 3) == 1
assert fecho(100, 2, 3) == 100
assert fecho("xyz") == "xyz"
assert bytes(fecho(bytearray(b"123"))) == b"123"
with pytest.raises(RuntimeError):
raise_err = remote.get_function(
"testing.test_raise_error_callback")("RuntimeError")
raise_err()
remote.cpu().sync()
with pytest.raises(AttributeError):
f3 = remote.system_lib()["notexist"]
temp = rpc.server._server_env([])
server = rpc.Server("localhost")
client = rpc.connect(server.host, server.port)
check(rpc.LocalSession())
check(client)
# Test minrpc server.
temp = util.tempdir()
minrpc_exec = temp.relpath("minrpc")
tvm.rpc.with_minrpc(cc.create_executable)(minrpc_exec, [])
check(rpc.PopenSession(minrpc_exec))
# minrpc on the remote
server = rpc.Server("localhost")
client = rpc.connect(
server.host,
server.port,
session_constructor_args=[
"rpc.PopenSession",
open(minrpc_exec, "rb").read()
],
)
check(client)
def test_rpc_runtime_string():
if not tvm.runtime.enabled("rpc"):
return
@tvm.register_func("rpc.test.runtime_str_concat")
def strcat(x, y):
return x + y
server = rpc.Server("localhost", key="x1")
client = rpc.connect(server.host, server.port, key="x1")
func = client.get_function("rpc.test.runtime_str_concat")
x = tvm.runtime.container.String("abc")
y = tvm.runtime.container.String("def")
assert str(func(x, y)) == "abcdef"
def test_rpc_vm():
server = rpc.Server(key="profiling")
remote = rpc.connect("127.0.0.1", server.port, key="profiling")
mod, params = mlp.get_workload(1)
exe = relay.vm.compile(mod, "llvm", params=params)
temp = utils.tempdir()
path = temp.relpath("lib.tar")
exe.mod.export_library(path)
remote.upload(path)
rexec = remote.load_module("lib.tar")
vm = profiler_vm.VirtualMachineProfiler(rexec, remote.cpu())
report = vm.profile(tvm.nd.array(np.ones((1, 1, 28, 28), dtype="float32"), device=remote.cpu()))
assert len(report.calls) > 0
def test_rpc_array():
if not tvm.runtime.enabled("rpc"):
return
x = np.random.randint(0, 10, size=(3, 4))
@tvm.register_func("rpc.test.remote_array_func")
def remote_array_func(y):
np.testing.assert_equal(y.asnumpy(), x)
server = rpc.Server("localhost")
remote = rpc.connect(server.host, server.port)
r_cpu = tvm.nd.array(x, remote.cpu(0))
assert str(r_cpu.context).startswith("remote")
np.testing.assert_equal(r_cpu.asnumpy(), x)
fremote = remote.get_function("rpc.test.remote_array_func")
fremote(r_cpu)
def verify_rpc_gpu_export(obj_format):
if not tvm.testing.device_enabled("cuda"):
print("Skip because cuda is not enabled")
return
mod, params = relay.testing.synthetic.get_workload()
with relay.build_config(opt_level=3):
complied_graph_lib = relay.build_module.build(mod,
"cuda",
params=params)
from tvm.contrib import utils
temp = utils.tempdir()
if obj_format == ".so":
file_name = "deploy_lib.so"
else:
assert obj_format == ".tar"
file_name = "deploy_lib.tar"
path_lib = temp.relpath(file_name)
complied_graph_lib.export_library(path_lib)
from tvm import rpc
def check_remote(server):
remote = rpc.connect(server.host, server.port)
remote.upload(path_lib)
loaded_lib = remote.load_module(path_lib)
data = np.random.uniform(-1, 1,
size=input_shape(mod)).astype("float32")
dev = remote.cuda()
# raw api
gmod = loaded_lib["default"](dev)
set_input = gmod["set_input"]
run = gmod["run"]
get_output = gmod["get_output"]
set_input("data", tvm.nd.array(data, device=dev))
run()
out = get_output(0).numpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
# graph executor wrapper
gmod = graph_executor.GraphModule(loaded_lib["default"](dev))
gmod.set_input("data", data)
gmod.run()
out = gmod.get_output(0).numpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
check_remote(rpc.Server("127.0.0.1"))
def test_rpc_session_constructor_args():
# start server
server0 = rpc.Server(key="x0")
server1 = rpc.Server(key="x1")
def check_multi_hop():
# use server0 as proxy to connect to server1
client = rpc.connect(
"127.0.0.1",
server0.port,
key="x0",
session_constructor_args=[
"rpc.Connect", "127.0.0.1", server1.port, "x1"
],
)
fecho = client.get_function("testing.echo")
assert fecho(1, 2, 3) == 1
assert fecho(100, 2, 3) == 100
assert fecho("xyz") == "xyz"
assert bytes(fecho(bytearray(b"123"))) == b"123"
nd = tvm.nd.array([1, 2, 3], device=client.cpu(0))
assert nd.asnumpy()[1] == 2
def check_error_handling():
with pytest.raises(tvm.error.RPCError):
client = rpc.connect(
"127.0.0.1",
server0.port,
key="x0",
session_constructor_args=["rpc.NonExistingConstructor"],
)
check_multi_hop()
check_error_handling()
def test_rpc_tracker_request():
# test concurrent request
tracker = Tracker("localhost", port=9000, port_end=10000)
device_key = "test_device"
server = rpc.Server(
"localhost",
port=9000,
port_end=10000,
key=device_key,
tracker_addr=(tracker.host, tracker.port),
)
client = rpc.connect_tracker(tracker.host, tracker.port)
def target(host, port, device_key, timeout):
client = rpc.connect_tracker(host, port)
remote = client.request(device_key, session_timeout=timeout)
while True:
pass
remote.cpu()
proc1 = multiprocessing.Process(target=target,
args=(tracker.host, tracker.port,
device_key, 4))
proc2 = multiprocessing.Process(target=target,
args=(tracker.host, tracker.port,
device_key, 200))
proc1.start()
time.sleep(0.5)
proc2.start()
time.sleep(0.5)
summary = client.summary()
assert summary["queue_info"][device_key]["free"] == 0
assert summary["queue_info"][device_key]["pending"] == 1
proc1.terminate()
proc1.join()
time.sleep(0.5)
summary = client.summary()
assert summary["queue_info"][device_key]["free"] == 0
assert summary["queue_info"][device_key]["pending"] == 0
proc2.terminate()
proc2.join()
server.terminate()
tracker.terminate()
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 check_minrpc():
if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None:
return
# Test minrpc server.
temp = utils.tempdir()
minrpc_exec = temp.relpath("minrpc")
tvm.rpc.with_minrpc(cc.create_executable)(minrpc_exec, [])
check(rpc.PopenSession(minrpc_exec))
# minrpc on the remote
server = rpc.Server("localhost")
client = rpc.connect(
server.host,
server.port,
session_constructor_args=["rpc.PopenSession", open(minrpc_exec, "rb").read()],
)
check(client)
def check_remote():
mlib = tvm.build(s, [A, B], "llvm", name="myadd")
server = rpc.Server("127.0.0.1")
remote = rpc.connect(server.host, server.port)
temp = utils.tempdir()
dev = remote.cpu(0)
path_dso = temp.relpath("dev_lib.so")
mlib.export_library(path_dso)
remote.upload(path_dso)
mlib = remote.load_module("dev_lib.so")
mod = graph_executor.create(graph, mlib, remote.cpu(0))
a = np.random.uniform(size=(n, )).astype(A.dtype)
mod.run(x=tvm.nd.array(a, dev))
out = tvm.nd.empty((n, ), device=dev)
out = mod.get_output(0, out)
np.testing.assert_equal(out.asnumpy(), a + 1)
def test_remote():
if not tvm.runtime.enabled("tflite"):
print("skip because tflite runtime is not enabled...")
return
if not tvm.get_global_func("tvm.tflite_runtime.create", True):
print("skip because tflite runtime is not enabled...")
return
try:
import tensorflow as tf
except ImportError:
print("skip because tensorflow not installed...")
return
tflite_fname = "model.tflite"
tflite_model = _create_tflite_model()
temp = utils.tempdir()
tflite_model_path = temp.relpath(tflite_fname)
open(tflite_model_path, "wb").write(tflite_model)
# inference via tflite interpreter python apis
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
input_shape = input_details[0]["shape"]
tflite_input = np.array(np.random.random_sample(input_shape),
dtype=np.float32)
interpreter.set_tensor(input_details[0]["index"], tflite_input)
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]["index"])
# inference via remote tvm tflite runtime
server = rpc.Server("localhost")
remote = rpc.connect(server.host, server.port)
ctx = remote.cpu(0)
a = remote.upload(tflite_model_path)
with open(tflite_model_path, "rb") as model_fin:
runtime = tflite_runtime.create(model_fin.read(), remote.cpu(0))
runtime.set_input(0, tvm.nd.array(tflite_input, remote.cpu(0)))
runtime.invoke()
out = runtime.get_output(0)
np.testing.assert_equal(out.asnumpy(), tflite_output)
server.terminate()
def test_rpc_return_ndarray():
# start server
server = rpc.Server("localhost", key="x1")
client = rpc.connect(server.host, server.port, key="x1")
m = client.get_function("rpc.test.remote_return_nd")
get_arr = m("get_arr")
ref_count = m("ref_count")
get_elem = m("get_elem")
get_arr_elem = m("get_arr_elem")
# array test
def run_arr_test():
arr = get_arr()
assert get_elem(0) == 0.0
assert get_arr_elem(arr, 0) == 0.0
run_arr_test()
def test_rpc_graph():
server = rpc.Server(key="profiling")
remote = rpc.connect("127.0.0.1", server.port, key="profiling")
mod, params = mlp.get_workload(1)
exe = relay.build(mod, "llvm", params=params)
temp = utils.tempdir()
path = temp.relpath("lib.tar")
exe.export_library(path)
remote.upload(path)
rexec = remote.load_module("lib.tar")
gr = debug_executor.create(exe.get_graph_json(), rexec, remote.cpu())
data = np.random.rand(1, 1, 28, 28).astype("float32")
report = gr.profile(data=data)
assert len(report.calls) > 0
def verify_rpc_gpu_export(obj_format):
if not tvm.runtime.enabled("cuda"):
print("Skip because cuda is not enabled")
return
mod, params = relay.testing.synthetic.get_workload()
with relay.build_config(opt_level=3):
complied_graph_lib = relay.build_module.build(mod,
"cuda",
params=params)
from tvm.contrib import util
temp = util.tempdir()
if obj_format == ".so":
file_name = "deploy_lib.so"
else:
assert obj_format == ".tar"
file_name = "deploy_lib.tar"
path_lib = temp.relpath(file_name)
complied_graph_lib.export_library(path_lib)
from tvm import rpc
server = rpc.Server("localhost", use_popen=True, port=9094)
remote = rpc.connect(server.host, server.port)
remote.upload(path_lib)
loaded_lib = remote.load_module(path_lib)
data = np.random.uniform(-1, 1,
size=input_shape(mod)).astype("float32")
ctx = remote.gpu()
# raw api
gmod = loaded_lib['default'](ctx)
set_input = gmod["set_input"]
run = gmod["run"]
get_output = gmod["get_output"]
set_input("data", tvm.nd.array(data, ctx=ctx))
run()
out = get_output(0).asnumpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
# graph runtime wrapper
gmod = graph_runtime.GraphModule(loaded_lib['default'](ctx))
gmod.set_input("data", data)
gmod.run()
out = gmod.get_output(0).asnumpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
def test_rpc(dtype):
if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
print("skip because extern function is not available")
return
if not tvm.testing.device_enabled("rpc") or not tvm.runtime.enabled("llvm"):
return
np_ones = np.ones((512, 512), dtype=dtype)
server = rpc.Server("localhost")
remote = rpc.connect(server.host, server.port)
value = tvm.nd.empty(np_ones.shape, np_ones.dtype, remote.cpu())
random_fill = remote.get_function("tvm.contrib.random.random_fill")
random_fill(value)
assert np.count_nonzero(value.asnumpy()) == 512 * 512
# make sure arithmentic doesn't overflow too
np_values = value.asnumpy()
assert np.isfinite(np_values * np_values + np_values).any()