def test_nested_sessions():
"""Test entering and exiting nested session contexts."""
if not tvm.runtime.enabled("micro_dev"):
return
shape = (1024,)
dtype = "float32"
# Construct Relay add program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
ret = relay.add(x, relay.const(1.0))
add_const_func = relay.Function([x], ret)
sess_a = micro.Session(DEV_CONFIG_A)
sess_b = micro.Session(DEV_CONFIG_B)
with sess_a:
np_tensor_a = np.random.uniform(size=shape).astype(dtype)
micro_tensor_a = tvm.nd.array(np_tensor_a, tvm.micro_dev(0))
with sess_b:
np_tensor_b = np.random.uniform(size=shape).astype(dtype)
micro_tensor_b = tvm.nd.array(np_tensor_b, tvm.micro_dev(0))
add_const_mod = relay_micro_build(add_const_func, DEV_CONFIG_A)
add_const_mod.run(x=micro_tensor_a)
add_result = add_const_mod.get_output(0).asnumpy()
tvm.testing.assert_allclose(
add_result, np_tensor_a + 1.0)
def test_inactive_session_use():
"""Test the use of objects allocated in a session that is no longer active."""
if not tvm.runtime.enabled("micro_dev"):
return
shape = (1024,)
dtype = "float32"
# Construct Relay add program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
ret = relay.add(x, relay.const(1.0))
add_const_func = relay.Function([x], ret)
sess_a = micro.Session(DEV_CONFIG_A)
sess_b = micro.Session(DEV_CONFIG_B)
with sess_a:
np_tensor_a = np.random.uniform(size=shape).astype(dtype)
micro_tensor_a = tvm.nd.array(np_tensor_a, tvm.micro_dev(0))
add_const_mod = relay_micro_build(add_const_func, DEV_CONFIG_A)
with sess_b:
# These objects belong to `sess_a`.
add_const_mod.run(x=micro_tensor_a)
add_result = add_const_mod.get_output(0).asnumpy()
tvm.testing.assert_allclose(
add_result, np_tensor_a + 1.0)
def test_interleave_sessions():
"""Test closing and reopening sessions."""
if not tvm.module.enabled("micro_dev"):
return
shape = (1024, )
dtype = "float32"
# Construct Relay add program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
ret = relay.add(x, relay.const(1.0))
add_const_func = relay.Function([x], ret)
sess_a = micro.Session(DEVICE_TYPE, TOOLCHAIN_PREFIX)
sess_b = micro.Session(DEVICE_TYPE, TOOLCHAIN_PREFIX)
with sess_a:
np_tensor_a = np.random.uniform(size=shape).astype(dtype)
micro_tensor_a = tvm.nd.array(np_tensor_a, tvm.micro_dev(0))
with sess_b:
np_tensor_b = np.random.uniform(size=shape).astype(dtype)
micro_tensor_b = tvm.nd.array(np_tensor_b, tvm.micro_dev(0))
with sess_a:
add_const_mod = relay_micro_build(add_const_func, TOOLCHAIN_PREFIX)
add_const_mod.run(x=micro_tensor_a)
add_result = add_const_mod.get_output(0).asnumpy()
tvm.testing.assert_allclose(add_result, np_tensor_a + 1.0)
with sess_b:
add_const_mod = relay_micro_build(add_const_func, TOOLCHAIN_PREFIX)
add_const_mod.run(x=micro_tensor_b)
add_result = add_const_mod.get_output(0).asnumpy()
tvm.testing.assert_allclose(add_result, np_tensor_b + 1.0)
def test_workspace_add():
"""Test a module which uses a workspace to compute an intermediate value."""
if not tvm.runtime.enabled("micro_dev"):
return
shape = (1024,)
dtype = "float32"
reset_gdbinit()
# Construct TVM expression.
tvm_shape = tvm.runtime.convert(shape)
A = te.placeholder(tvm_shape, name="A", dtype=dtype)
B = te.placeholder(tvm_shape, name="B", dtype=dtype)
B = te.compute(A.shape, lambda *i: A(*i) + 1, name="B")
C = te.compute(A.shape, lambda *i: B(*i) + 1, name="C")
s = te.create_schedule(C.op)
func_name = "fadd_two_workspace"
c_mod = tvm.build(s, [A, C], target="c", name=func_name)
with micro.Session(DEV_CONFIG_A) as sess:
micro_mod = micro.create_micro_mod(c_mod, DEV_CONFIG_A)
micro_func = micro_mod[func_name]
ctx = tvm.micro_dev(0)
a_np = np.random.uniform(size=shape).astype(dtype)
a = tvm.nd.array(a_np, ctx)
c = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx)
micro_func(a, c)
# ensure input wasn't corrupted
tvm.testing.assert_allclose(
a.asnumpy(), a_np)
# ensure output is correct
tvm.testing.assert_allclose(
c.asnumpy(), a.asnumpy() + 2.0)
def test_conv2d():
if not tvm.runtime.enabled("micro_dev"):
return
from tvm.relay import create_executor
from tvm.relay import transform
dshape = (1, 4, 16, 16)
dtype = 'int8'
func_name = 'fused_nn_conv2d'
reset_gdbinit()
# Construct Relay program.
x = relay.var("x", shape=dshape, dtype=dtype)
conv_expr = relay.nn.conv2d(x,
relay.var("w"),
kernel_size=(3, 3),
padding=(1, 1),
channels=4)
func = relay.Function(relay.analysis.free_vars(conv_expr), conv_expr)
mod = tvm.IRModule.from_expr(func)
mod = transform.InferType()(mod)
x_shape = list(
map(lambda x: x.value, mod['main'].params[0].checked_type.shape))
w_shape = list(
map(lambda x: x.value, mod['main'].params[1].checked_type.shape))
out_shape = list(map(lambda x: x.value, mod['main'].ret_type.shape))
with tvm.transform.PassContext(config={"tir.disable_vectorize": True}):
graph, c_mod, params = relay.build(mod, target="c")
with micro.Session(DEV_CONFIG_A):
micro_mod = micro.create_micro_mod(c_mod, DEV_CONFIG_A)
candidate_func_name = func_name
for i in range(100):
try:
micro_func = micro_mod[candidate_func_name]
break
except tvm.TVMError as e:
candidate_func_name = f'{func_name}_{i}'
else:
assert False
ctx = tvm.micro_dev(0)
x_data = tvm.nd.array(
np.random.uniform(size=x_shape).astype(dtype), ctx)
w_data = tvm.nd.array(
np.random.uniform(size=w_shape).astype(dtype), ctx)
result = tvm.nd.array(np.zeros(shape=out_shape, dtype=dtype), ctx)
micro_func(x_data, w_data, result)
out_data = np.zeros(out_shape, dtype=dtype)
params = {'x': x_data.asnumpy(), 'w': w_data.asnumpy()}
intrp = create_executor('debug')
expected_result = intrp.evaluate(mod['main'])(x_data, w_data)
tvm.testing.assert_allclose(result.asnumpy(),
expected_result.asnumpy())
def test_multiple_modules():
"""Test loading multiple modules on the device simultaneously."""
if not tvm.module.enabled("micro_dev"):
return
shape = (1024, )
dtype = "float32"
# Construct Relay add program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
ret = relay.add(x, relay.const(1.0))
add_const_func = relay.Function([x], ret)
# Construct Relay subtract program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
ret = relay.subtract(x, relay.const(1.0))
sub_const_func = relay.Function([x], ret)
with micro.Session(DEVICE_TYPE, TOOLCHAIN_PREFIX):
add_const_mod = relay_micro_build(add_const_func, TOOLCHAIN_PREFIX)
sub_const_mod = relay_micro_build(sub_const_func, TOOLCHAIN_PREFIX)
x_in = np.random.uniform(size=shape[0]).astype(dtype)
add_const_mod.run(x=x_in)
add_result = add_const_mod.get_output(0).asnumpy()
sub_const_mod.run(x=x_in)
sub_result = sub_const_mod.get_output(0).asnumpy()
tvm.testing.assert_allclose(add_result, x_in + 1.0)
tvm.testing.assert_allclose(sub_result, x_in - 1.0)
def test_workspace_add():
"""Test a module which uses a workspace to compute an intermediate value."""
if not tvm.module.enabled("micro_dev"):
return
shape = (1024, )
dtype = "float32"
# Construct TVM expression.
tvm_shape = tvm.convert(shape)
A = tvm.placeholder(tvm_shape, name="A", dtype=dtype)
B = tvm.placeholder(tvm_shape, name="B", dtype=dtype)
B = tvm.compute(A.shape, lambda *i: A(*i) + 1, name="B")
C = tvm.compute(A.shape, lambda *i: B(*i) + 1, name="C")
s = tvm.create_schedule(C.op)
func_name = "fadd_two_workspace"
c_mod = tvm.build(s, [A, C], target="c", name=func_name)
with micro.Session(DEVICE_TYPE, TOOLCHAIN_PREFIX):
micro_mod = create_micro_mod(c_mod, TOOLCHAIN_PREFIX)
micro_func = micro_mod[func_name]
ctx = tvm.micro_dev(0)
a = tvm.nd.array(np.random.uniform(size=shape).astype(dtype), ctx)
c = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx)
micro_func(a, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 2.0)
def test_add():
"""Test a module which performs addition."""
if not tvm.module.enabled("micro_dev"):
return
shape = (1024,)
dtype = "float32"
# Construct TVM expression.
tvm_shape = tvm.convert(shape)
A = tvm.placeholder(tvm_shape, name="A", dtype=dtype)
B = tvm.placeholder(tvm_shape, name="B", dtype=dtype)
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
s = tvm.create_schedule(C.op)
func_name = "fadd"
c_mod = tvm.build(s, [A, B, C], target="c", name=func_name)
with micro.Session(DEV_CONFIG):
micro_mod = create_micro_mod(c_mod, DEV_CONFIG)
micro_func = micro_mod[func_name]
ctx = tvm.micro_dev(0)
a = tvm.nd.array(np.random.uniform(size=shape).astype(dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=shape).astype(dtype), ctx)
c = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx)
micro_func(a, b, c)
tvm.testing.assert_allclose(
c.asnumpy(), a.asnumpy() + b.asnumpy())
def spike_model(model_path, input_x, input_name):
"""Test a program which uses the graph runtime."""
if not tvm.runtime.enabled("micro_dev"):
print("not enable micro_dev")
return
input_x = input_x.astype("float32")
shape_dict = {input_name: input_x.shape}
mod, params = load_model(model_path, shape_dict)
with micro.Session(DEV_RISCV):
ctx = tvm.micro_dev(0)
disable_vectorize = tvm.target.build_config(disable_vectorize=True)
disable_fusion = relay.build_config(disabled_pass={'FuseOps'})
with disable_vectorize:
graph, c_mod, params = relay.build(mod,
target=TARGET,
params=params)
print("Part Success")
micro_mod = micro.create_micro_mod(c_mod, DEV_RISCV)
mod = graph_runtime.create(graph, micro_mod, ctx)
mod.set_input(**params)
mod.set_input(input_name, tvm.nd.array(input_x))
mod.run
tvm_output = mod.get_output(0).asnumpy()
return tvm_output
def server_start():
# pylint: disable=unused-variable
session = micro.Session(dev_config)
session._enter()
@tvm.register_func("tvm.rpc.server.shutdown", override=True)
def server_shutdown():
session._exit()
def test_alloc():
"""Test tensor allocation on the device."""
if not tvm.runtime.enabled("micro_dev"):
return
shape = (1024,)
dtype = "float32"
with micro.Session(DEV_CONFIG_A):
ctx = tvm.micro_dev(0)
np_tensor = np.random.uniform(size=shape).astype(dtype)
micro_tensor = tvm.nd.array(np_tensor, ctx)
tvm.testing.assert_allclose(np_tensor, micro_tensor.asnumpy())
def test_model():
"""Test a program which uses the graph runtime."""
if not tvm.runtime.enabled("micro_dev"):
print("not enable micro_dev")
return
model_url = 'https://people.linaro.org/~tom.gall/sine_model.tflite'
model_file = 'sine_model.tflite'
model_path = download_testdata(model_url, model_file, module='data')
tflite_model_buf = open(model_path, "rb").read()
# Using the buffer, transform into a tflite model python object
try:
import tflite
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)
# Print out the version of the model
version = tflite_model.Version()
print ("Model Version: " + str(version))
input_tensor = "dense_4_input"
input_shape = (1,)
input_dtype = "float32"
mod, params = relay.frontend.from_tflite(tflite_model,
shape_dict={input_tensor: input_shape},
dtype_dict={input_tensor: input_dtype})
with micro.Session(DEV_RISCV):
ctx = tvm.micro_dev(0)
disable_vectorize = tvm.target.build_config(disable_vectorize=True)
disable_fusion = relay.build_config(disabled_pass={'FuseOps'})
with disable_vectorize, disable_fusion:
graph, c_mod, params = relay.build(mod, target=TARGET, params=params)
micro_mod = micro.create_micro_mod(c_mod, DEV_RISCV)
mod = graph_runtime.create(graph, micro_mod, ctx)
mod.set_input(**params)
mod.set_input(input_tensor, tvm.nd.array(np.array([0.5], dtype="float32")))
tvm_output = mod.get_output(0).asnumpy()
print("result is: "+str(tvm_output))
def get_comm_overhead(dev_config, num_trials=1):
"""Get communication overhead by executing an empty kernel."""
class EmptyCMod:
def __init__(self):
pass
def export_library(self, out_obj_path, fcompile=None):
assert fcompile is not None
fcompile(out_obj_path, f'{os.path.dirname(__file__)}/empty.c')
# do multiple trials, then calc the average comm overhead
results = []
with micro.Session(dev_config) as sess:
micro_mod = create_micro_mod(EmptyCMod(), dev_config)
micro_func = micro_mod['empty']
for _ in range(num_trials):
results.append(benchmark_micro_func(sess, micro_func, [], 1, 0.0))
return sum(results) / len(results)
def test_graph_runtime():
"""Test a program which uses the graph runtime."""
if not tvm.module.enabled("micro_dev"):
return
shape = (1024, )
dtype = "float32"
# Construct Relay program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
xx = relay.multiply(x, x)
z = relay.add(xx, relay.const(1.0))
func = relay.Function([x], z)
with micro.Session(DEVICE_TYPE, TOOLCHAIN_PREFIX):
mod = relay_micro_build(func, TOOLCHAIN_PREFIX)
x_in = np.random.uniform(size=shape[0]).astype(dtype)
mod.run(x=x_in)
result = mod.get_output(0).asnumpy()
tvm.testing.assert_allclose(result, x_in * x_in + 1.0)
def test_model():
"""Test a program which uses the graph runtime."""
if not tvm.runtime.enabled("micro_dev"):
print("not enable micro_dev")
return
import onnx
model_path = "mobilenet_v2.onnx"
# now you have super_resolution.onnx on disk mobilenet_v2.onnx
onnx_model = onnx.load(model_path)
from PIL import Image
img_path = "cat.png"
img = Image.open(img_path).resize((224, 224))
x = np.array(img)[np.newaxis, :, :, :]
print(x.shape)
x = x.transpose([0,3,1,2])
print(x.shape)
input_name = "input"
shape_dict = {input_name: x.shape}
mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)
with micro.Session(DEV_RISCV):
ctx = tvm.micro_dev(0)
disable_vectorize = tvm.target.build_config(disable_vectorize=True)
disable_fusion = relay.build_config(disabled_pass={'FuseOps'})
with disable_vectorize:
graph, c_mod, params = relay.build(mod, target=TARGET, params=params)
print("I Find the wrong")
micro_mod = micro.create_micro_mod(c_mod, DEV_RISCV)
mod = graph_runtime.create(graph, micro_mod, ctx)
mod.set_input(**params)
mod.set_input(input_name, tvm.nd.array(x))
tvm_output = mod.get_output(0).asnumpy()
print("result is: "+str(tvm_output))
# Grab an example
print("[Test Input]")
image, label = get_sample_point()
input(' Check "test_input.png"...')
print()
# Begin a session
import time
print("[Initialization]")
start_time = time.time()
RISCV_TOOLCHAIN_PREFIX = "riscv64-unknown-elf-"
with micro.Session("openocd",
RISCV_TOOLCHAIN_PREFIX,
base_addr=0x10010000,
server_addr="127.0.0.1",
port=6666) as sess:
end_time = time.time()
print(f' Initialization took {end_time - start_time} seconds')
print()
input(' Press enter to continue ')
print()
# Build the function
print("[Building]")
start_time = time.time()
mod = relay_micro_build(func, RISCV_TOOLCHAIN_PREFIX, params=params)
end_time = time.time()
print(f' Build took {end_time - start_time} seconds')
print()
