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 relay_micro_build(func, dev_config, params=None):
"""Create a graph runtime module with a micro device context from a Relay function.
Parameters
----------
func : relay.Function
function to compile
dev_config : Dict[str, Any]
MicroTVM config dict for the target device
params : dict
input parameters that do not change during inference
Return
------
mod : tvm.runtime.Module
graph runtime module for the target device
"""
with tvm.transform.PassContext(disabled_pass={'FuseOps'}, config={
"tir.disable_vectorize": True
}):
graph, c_mod, params = relay.build(func, target=TARGET, params=params)
micro_mod = micro.create_micro_mod(c_mod, dev_config)
ctx = tvm.micro_dev(0)
mod = graph_runtime.create(graph, micro_mod, ctx)
mod.set_input(**params)
return mod
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 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 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 relay_micro_build(func,
dev_config,
target,
params=None,
lib_headers=None,
lib_include_paths=None):
"""Create a graph runtime module with a micro device context from a Relay function.
Parameters
----------
func : relay.Function
function to compile
dev_config : TODO
TODO
target : TODO
TODO
params : dict
input parameters that do not change during inference
lib_headers : TODO
TODO
lib_include_paths : TODO
TODO
Return
------
mod : tvm.module.Module
graph runtime module for the target device
"""
with tvm.target.build_config(opt_level=3, disable_vectorize=True):
graph, c_mod, params = relay.build(func, target=target, params=params)
micro_mod = micro.create_micro_mod(c_mod,
dev_config,
lib_headers=lib_headers,
lib_include_paths=lib_include_paths)
ctx = tvm.micro_dev(0)
if DEBUG_MODE:
dump_root = f'{get_repo_root()}/debug/micro'
mod = debug_runtime.create(graph, micro_mod, ctx, dump_root=dump_root)
else:
mod = graph_runtime.create(graph, micro_mod, ctx)
mod.set_input(**params)
return mod
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_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))