def test_crt_link_params():
from tvm import micro
for dtype in LINKABLE_DTYPES:
mod, param_init = _make_mod_and_params(dtype)
rand_input = _make_random_tensor(dtype, INPUT_SHAPE)
target = "c"
runtime = Runtime("crt", {"system-lib": True})
executor = Executor("graph", {"link-params": True})
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
factory = tvm.relay.build(mod,
target,
runtime=runtime,
executor=executor,
params=param_init)
assert set(factory.get_params().keys()) == {"p0", "p1"
} # NOTE: op folded
temp_dir = tvm.contrib.utils.tempdir()
template_project_dir = os.path.join(
tvm.micro.get_standalone_crt_dir(), "template", "host")
project = tvm.micro.generate_project(template_project_dir, factory,
temp_dir / "project",
{"verbose": 1})
project.build()
project.flash()
with tvm.micro.Session(project.transport()) as sess:
graph_rt = tvm.micro.session.create_local_graph_executor(
factory.get_graph_json(), sess.get_system_lib(),
sess.device)
# NOTE: not setting params here.
graph_rt.set_input("rand_input", rand_input)
graph_rt.run()
linked_output = graph_rt.get_output(0).numpy()
runtime = Runtime("cpp", {"system-lib": True})
with tvm.transform.PassContext(opt_level=3):
lib = tvm.relay.build(mod,
"llvm",
runtime=runtime,
params=param_init)
def _run_unlinked(lib):
graph_json, mod, lowered_params = lib
graph_rt = tvm.contrib.graph_executor.create(
graph_json, mod, tvm.cpu(0))
graph_rt.set_input("rand_input", rand_input, **lowered_params)
graph_rt.run()
return graph_rt.get_output(0).numpy()
unlinked_output = _run_unlinked(lib)
if "int" in dtype:
np.testing.assert_equal(unlinked_output, linked_output)
else:
np.testing.assert_allclose(unlinked_output, linked_output)
def test_mobilenet_aot(hexagon_session: Session, aot_host_target, aot_target,
enable_usmp):
if hexagon_session._launcher._serial_number == "simulator":
pytest.skip(msg="Skip on simulator due to long runtime.")
dtype = "float32"
onnx_model = get_mobilenet()
data_in = np.random.rand(1, 3, 224, 224).astype(dtype=dtype)
input_name = "input"
shape_dict = {input_name: data_in.shape}
relay_mod, params = relay.frontend.from_onnx(onnx_model,
shape_dict,
freeze_params=True)
inputs = {input_name: data_in}
target_llvm = tvm.target.Target("llvm")
config = {"tir.usmp.enable": enable_usmp}
with tvm.transform.PassContext(opt_level=3, config=config):
hexagon_lowered = tvm.relay.build(
relay_mod,
tvm.target.Target(aot_target, host=aot_host_target),
runtime=Runtime("cpp"),
executor=Executor("aot", {
"unpacked-api": False,
"interface-api": "packed"
}),
params=params,
)
aot_mod = hexagon_session.get_executor_from_factory(hexagon_lowered)
aot_mod.set_input(**inputs)
aot_mod.run()
hexagon_output = aot_mod.get_output(0).numpy()
with tvm.transform.PassContext(opt_level=3):
llvm_lowered = tvm.relay.build(
relay_mod,
tvm.target.Target(target_llvm, host=target_llvm),
runtime=Runtime("cpp"),
executor=Executor("graph", {"link-params": True}),
params=params,
)
llvm_graph_mod = tvm.contrib.graph_executor.GraphModule(
llvm_lowered["default"](tvm.cpu(0)))
llvm_graph_mod.set_input(**inputs)
llvm_graph_mod.run()
expected_output = llvm_graph_mod.get_output(0).numpy()
tvm.testing.assert_allclose(hexagon_output,
expected_output,
rtol=1e-4,
atol=1e-5)
def test_qemu_make_fail(temp_dir, board, west_cmd, tvm_debug):
"""Testing QEMU make fail."""
if board not in ["qemu_x86", "mps2_an521"]:
pytest.skip(msg="Only for QEMU targets.")
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": tvm_debug}
shape = (10,)
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(np.ones(shape=shape, dtype=dtype)))
func = relay.Function([x], z)
ir_mod = tvm.IRModule.from_expr(func)
target = tvm.target.target.micro(model)
executor = Executor("aot")
runtime = Runtime("crt")
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
lowered = relay.build(ir_mod, target, executor=executor, runtime=runtime)
# Generate input/output header files
with tempfile.NamedTemporaryFile() as tar_temp_file:
with tarfile.open(tar_temp_file.name, "w:gz") as tf:
with tempfile.TemporaryDirectory() as tar_temp_dir:
model_files_path = os.path.join(tar_temp_dir, "include")
os.mkdir(model_files_path)
header_path = generate_c_interface_header(
lowered.libmod_name, ["input_1"], ["output"], [], 0, model_files_path
)
tf.add(header_path, arcname=os.path.relpath(header_path, tar_temp_dir))
test_utils.create_header_file(
"input_data", np.zeros(shape=shape, dtype=dtype), "include", tf
)
test_utils.create_header_file(
"output_data", np.zeros(shape=shape, dtype=dtype), "include", tf
)
project, project_dir = test_utils.build_project(
temp_dir,
board,
west_cmd,
lowered,
build_config,
extra_files_tar=tar_temp_file.name,
)
file_path = (
pathlib.Path(project_dir) / "build" / "zephyr" / "CMakeFiles" / "run.dir" / "build.make"
)
assert file_path.is_file(), f"[{file_path}] does not exist."
# Remove a file to create make failure.
os.remove(file_path)
project.flash()
with pytest.raises(server.JSONRPCError) as excinfo:
project.transport().open()
assert "QEMU setup failed" in str(excinfo.value)
def test_create_runtime_attr_type_incorrect():
with pytest.raises(
TVMError,
match='Attribute "system-lib" should have type "IntImm"'
' but instead found "runtime.String"',
):
Runtime("c", {"system-lib": "woof"})
def test_graph_executor():
"""Test use of the graph executor with microTVM."""
ws_root = pathlib.Path(os.path.dirname(__file__) + "/micro-workspace")
if ws_root.exists():
shutil.rmtree(ws_root)
temp_dir = tvm.contrib.utils.tempdir(ws_root.resolve())
relay_mod = tvm.parser.fromtext("""
#[version = "0.0.5"]
def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), uint8]) {
%0 = %a + %b;
%0
}""")
runtime = Runtime("crt", {"system-lib": True})
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
factory = tvm.relay.build(relay_mod, target=TARGET, runtime=runtime)
with _make_session(temp_dir, factory) as sess:
graph_mod = tvm.micro.create_local_graph_executor(
factory.get_graph_json(), sess.get_system_lib(), sess.device)
A_data = tvm.nd.array(np.array([2, 3], dtype="uint8"),
device=sess.device)
assert (A_data.numpy() == np.array([2, 3])).all()
B_data = tvm.nd.array(np.array([4, 7], dtype="uint8"),
device=sess.device)
assert (B_data.numpy() == np.array([4, 7])).all()
graph_mod.run(a=A_data, b=B_data)
out = graph_mod.get_output(0)
assert (out.numpy() == np.array([6, 10])).all()
def test_multiple_relay_modules_c():
mod = get_conv2d_relay_module()
executor = Executor("aot", {"unpacked-api": True, "interface-api": "c"})
runtime = Runtime("crt")
target = tvm.target.target.micro("host")
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
factory1 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod1")
factory2 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod2")
temp_dir = utils.tempdir()
mlf_tar_path = temp_dir.relpath("lib.tar")
micro.export_model_library_format([factory1, factory2], mlf_tar_path)
tf = tarfile.open(mlf_tar_path)
extract_dir = temp_dir.relpath("extract")
os.mkdir(extract_dir)
tf.extractall(extract_dir)
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib0.c"))
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib1.c"))
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod2_lib0.c"))
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod2_lib1.c"))
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "include", "tvmgen_mod1.h"))
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "include", "tvmgen_mod2.h"))
# check CRT runtime directory
assert os.path.exists(os.path.join(extract_dir, "runtime"))
def _get_compilation_config(accel_type, enable_cascader, enable_striping):
enable_usmp = True
target = tvm.target.Target("c")
ethosu_target = tvm.target.Target("ethos-u")
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": 16,
"interface-api": "c",
"unpacked-api": True,
},
)
pass_config = {
"tir.disable_vectorize": True,
"relay.ext.ethos-u.options": {
"accelerator_config": accel_type,
"enable_cascader": enable_cascader,
"enable_striping": enable_striping,
},
"tir.usmp.enable": enable_usmp,
"tir.usmp.algorithm": "hill_climb",
"tir.disable_storage_rewrite": enable_usmp,
}
return target, ethosu_target, runtime, executor, pass_config
def test_incompatible_interface_api_errors():
"""Ensures an error is thrown if not using the C interface API"""
mod, params = tvm.relay.testing.synthetic.get_workload()
target = "c"
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"interface-api": "packed",
},
)
with pytest.raises(
tvm.TVMError,
match=re.escape(
"tir.usmp.use_workspace_io option is only compatible with interface_api c.\n"
"Please use interface_api c to be able to enable tir.usmp.use_workspace_io"
),
):
with tvm.transform.PassContext(
opt_level=3,
config={
"tir.usmp.enable": True,
"tir.usmp.use_workspace_io": True
},
):
tvm.relay.build(mod,
target,
executor=executor,
runtime=runtime,
params=params)
def test_workspace_calculation(workspace_byte_alignment, main_workspace_size):
mod, params = tvm.relay.testing.synthetic.get_workload()
target = "c"
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": workspace_byte_alignment,
},
)
with tvm.transform.PassContext(
opt_level=3,
config={
"tir.disable_vectorize": True,
},
):
lib = tvm.relay.build(mod,
target,
executor=executor,
runtime=runtime,
params=params)
mlf_memory_map = mlf._build_function_memory_map(lib.function_metadata)
assert mlf_memory_map["main"][0][
"workspace_size_bytes"] == main_workspace_size
def test_multiple_relay_modules_same_module_name():
mod = get_conv2d_relay_module()
executor = Executor("graph")
runtime = Runtime("crt")
target = tvm.target.target.micro("host")
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
factory1 = tvm.relay.build(mod,
target,
runtime=runtime,
executor=executor,
mod_name="mod")
factory2 = tvm.relay.build(mod,
target,
runtime=runtime,
executor=executor,
mod_name="mod")
temp_dir = utils.tempdir()
mlf_tar_path = temp_dir.relpath("lib.tar")
with pytest.raises(AssertionError,
match="Multiple modules should have unique names"):
micro.export_model_library_format([factory1, factory2], mlf_tar_path)
def check_result(temp_dir, relay_mod, model, zephyr_board, west_cmd,
map_inputs, out_shape, result, build_config):
"""Helper function to verify results"""
TOL = 1e-5
runtime = Runtime("crt", {"system-lib": True})
target = tvm.target.target.micro(model)
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
mod = tvm.relay.build(relay_mod, target=target, runtime=runtime)
with _make_session(temp_dir, zephyr_board, west_cmd, mod,
build_config) as session:
rt_mod = tvm.micro.create_local_graph_executor(
mod.get_graph_json(), session.get_system_lib(), session.device)
rt_mod.set_input(**mod.get_params())
for name, data in map_inputs.items():
rt_mod.set_input(name, data)
rt_mod.set_input(**mod.get_params())
rt_mod.run()
out_shapes = out_shape if isinstance(out_shape, list) else [out_shape]
results = result if isinstance(result, list) else [result]
for idx, shape in enumerate(out_shapes):
out = tvm.nd.empty(shape, device=session.device)
out = rt_mod.get_output(idx, out)
tvm.testing.assert_allclose(out.numpy(),
results[idx],
rtol=TOL,
atol=TOL)
def test_schedule_build_with_cmsis_dependency(temp_dir, board, west_cmd,
tvm_debug):
"""Test Relay schedule with CMSIS dependency. This test shows if microTVM Auto tuning
with Zephyr breaks if CMSIS dependency was required for a schedule.
"""
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": tvm_debug}
target = tvm.target.target.micro(model, options=["-keys=arm_cpu,cpu"])
isa = arm_isa.IsaAnalyzer(target)
if not isa.has_dsp_support:
pytest.skip(f"ISA does not support DSP. target: {target}")
# Create a Relay conv2d
data_shape = (1, 16, 16, 3)
weight_shape = (5, 5, 8, 3)
data = relay.var("data", relay.TensorType(data_shape, "int8"))
weight = relay.var("weight", relay.TensorType(weight_shape, "int8"))
y = relay.nn.conv2d(
data,
weight,
padding=(2, 2),
kernel_size=(5, 5),
data_layout="NHWC",
kernel_layout="HWOI",
out_dtype="int32",
)
func = relay.Function([data, weight], y)
ir_mod = tvm.IRModule.from_expr(func)
runtime = Runtime("crt", {"system-lib": True})
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
mod = tvm.relay.build(ir_mod, target=target, runtime=runtime)
project_options = {
"project_type": "host_driven",
"west_cmd": west_cmd,
"verbose": bool(build_config.get("debug")),
"zephyr_board": board,
"cmsis_path": os.getenv("CMSIS_PATH"),
}
project_dir = temp_dir / "project"
project = tvm.micro.generate_project(
str(test_utils.TEMPLATE_PROJECT_DIR),
mod,
project_dir,
project_options,
)
project.build()
with open(project_dir / "CMakeLists.txt", "r") as cmake_f:
cmake_content = cmake_f.read()
assert "CMSIS/DSP/Include" in cmake_content
assert "CMSIS/DSP/Include/dsp" in cmake_content
assert "CMSIS/DSP/Include" in cmake_content
assert "CMSIS/NN/Include" in cmake_content
def test_relay(temp_dir, board, west_cmd, tvm_debug):
"""Testing a simple relay graph"""
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": tvm_debug}
shape = (10, )
dtype = "int8"
# Construct Relay program.
x = relay.var("x", relay.TensorType(shape=shape, dtype=dtype))
xx = relay.multiply(x, x)
z = relay.add(xx, relay.const(np.ones(shape=shape, dtype=dtype)))
func = relay.Function([x], z)
ir_mod = tvm.IRModule.from_expr(func)
runtime = Runtime("crt", {"system-lib": True})
target = tvm.target.target.micro(model)
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
mod = tvm.relay.build(ir_mod, target=target, runtime=runtime)
with _make_session(temp_dir, board, west_cmd, mod,
build_config) as session:
graph_mod = tvm.micro.create_local_graph_executor(
mod.get_graph_json(), session.get_system_lib(), session.device)
graph_mod.set_input(**mod.get_params())
x_in = np.random.randint(10, size=shape[0], dtype=dtype)
graph_mod.run(x=x_in)
result = graph_mod.get_output(0).numpy()
tvm.testing.assert_allclose(graph_mod.get_input(0).numpy(), x_in)
tvm.testing.assert_allclose(result, x_in * x_in + 1)
def test_compile_tflite_module_with_external_codegen_cmsisnn(
tmpdir_factory, tflite_cnn_s_quantized):
pytest.importorskip("tflite")
output_dir = tmpdir_factory.mktemp("mlf")
tvmc_model = tvmc.load(tflite_cnn_s_quantized)
output_file_name = f"{output_dir}/file.tar"
tvmc.compiler.compile_model(
tvmc_model,
target=f"cmsis-nn, c -mcpu=cortex-m55",
runtime=Runtime("crt", {"system-lib": True}),
executor=Executor("aot"),
output_format="mlf",
package_path=output_file_name,
pass_context_configs=["tir.disable_vectorize=true"],
)
# check whether an MLF package was created
assert os.path.exists(output_file_name)
# check whether the expected number of C sources are in the tarfile
with tarfile.open(output_file_name) as mlf_package:
c_source_files = [
name for name in mlf_package.getnames()
if re.match(r"\./codegen/host/src/\D+\d+\.c", name)
]
assert len(c_source_files) == 3
def test_memory_planning(workspace_byte_alignment, main_workspace_size):
mod, params = tvm.relay.testing.synthetic.get_workload()
target = "c"
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": workspace_byte_alignment,
},
)
with tvm.transform.PassContext(
opt_level=3,
config={
"tir.disable_vectorize": True,
"tir.disable_storage_rewrite": True,
"tir.usmp.enable": True,
"tir.usmp.algorithm": "greedy_by_conflicts",
},
):
lib = tvm.relay.build(mod,
target,
executor=executor,
runtime=runtime,
params=params)
assert (sum(lib.function_metadata["__tvm_main__"].workspace_sizes.values())
== main_workspace_size)
def check_device(device):
dev = tvm.device(device, 0)
if not tvm.testing.device_enabled(device):
print("Skip because %s is not enabled" % device)
return
temp = utils.tempdir()
name = "myadd_%s" % device
if sys.platform == "darwin" or sys.platform.startswith("linux"):
runtime = Runtime("cpp", {"system-lib": True})
f = tvm.build(s, [A, B], device, "llvm", runtime=runtime, name=name)
elif sys.platform == "win32":
f = tvm.build(s, [A, B], device, "llvm", name=name)
else:
raise ValueError("Unsupported platform")
path_dso = temp.relpath("dev_lib.so")
# test cross compiler function
f.export_library(path_dso, cc.cross_compiler("g++"))
f1 = tvm.runtime.load_module(path_dso)
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), dev)
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
f1(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
if sys.platform != "win32":
f2 = tvm.runtime.system_lib()
f2[name](a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
def check_minrpc():
if tvm.get_global_func("rpc.CreatePipeClient",
allow_missing=True) is None:
return
# export to minrpc
temp = utils.tempdir()
runtime = Runtime("cpp", {"system-lib": True})
f = tvm.build(s, [A, B], "llvm", name="myadd", runtime=runtime)
path_minrpc = temp.relpath("dev_lib.minrpc")
f.export_library(path_minrpc, rpc.with_minrpc(cc.create_executable))
with pytest.raises(RuntimeError):
rpc.PopenSession("filenotexist")
# statrt the minrpc session.
remote = tvm.rpc.PopenSession(path_minrpc)
dev = remote.cpu(0)
f1 = remote.system_lib()
a = tvm.nd.array(np.random.uniform(size=102).astype(A.dtype), dev)
b = tvm.nd.array(np.zeros(102, dtype=A.dtype), dev)
time_f = f1.time_evaluator("myadd", remote.cpu(0), number=1)
cost = time_f(a, b).mean
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
# change to not executable
os.chmod(path_minrpc, stat.S_IRUSR)
with pytest.raises(RuntimeError):
rpc.PopenSession(path_minrpc)
def test_memory_planning(workspace_byte_alignment, main_workspace_size):
"""Checks calculated workspace against known values"""
mod, params = tvm.relay.testing.synthetic.get_workload()
target = "c"
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": workspace_byte_alignment,
},
)
with tvm.transform.PassContext(
opt_level=3,
config={
"tir.disable_vectorize": True,
"tir.disable_storage_rewrite": True,
"tir.usmp.enable": True,
"tir.usmp.algorithm": "greedy_by_conflicts",
},
):
lib = tvm.relay.build(mod,
target,
executor=executor,
runtime=runtime,
params=params)
# The workspace_size dictionary will have an entry for both the 'primitive' and 'host'
# targets, though both are identical.
for size in lib.function_metadata["__tvm_main__"].workspace_sizes.values():
assert size == main_workspace_size
def test_multiple_relay_modules_graph():
mod = get_conv2d_relay_module()
executor = Executor("graph")
runtime = Runtime("crt")
target = tvm.target.target.micro("host")
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
factory1 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod1")
factory2 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod2")
temp_dir = utils.tempdir()
mlf_tar_path = temp_dir.relpath("lib.tar")
micro.export_model_library_format([factory1, factory2], mlf_tar_path)
with tarfile.open(mlf_tar_path, "r:*") as tf:
tar_members = [ti.name for ti in tf.getmembers()]
print("tar members", tar_members)
assert "./metadata.json" in tar_members
assert "./codegen/host/src/mod1_lib0.c" in tar_members
assert "./codegen/host/src/mod2_lib0.c" in tar_members
with tf.extractfile("./metadata.json") as f:
metadata = json.load(f)
mod2_main_md = metadata["modules"]["mod2"]["memory"]["functions"]["main"]
assert mod2_main_md == [
{
"constants_size_bytes": 0,
"device": 1,
"io_size_bytes": 143960,
"workspace_size_bytes": 158088,
}
]
assert metadata["modules"]["mod1"]["model_name"] == "mod1"
assert metadata["modules"]["mod2"]["model_name"] == "mod2"
def check_system_lib():
dev = tvm.cpu(0)
if not tvm.testing.device_enabled("llvm"):
print("Skip because llvm is not enabled")
return
temp = utils.tempdir()
runtime = Runtime("cpp", {"system-lib": True})
fadd1 = tvm.build(s, [A, B], "llvm", runtime=runtime, name="myadd1")
fadd2 = tvm.build(s, [A, B], "llvm", runtime=runtime, name="myadd2")
path1 = temp.relpath("myadd1.o")
path2 = temp.relpath("myadd2.o")
path_dso = temp.relpath("mylib.so")
fadd1.save(path1)
fadd2.save(path2)
cc.create_shared(path_dso, [path1, path2])
# Load dll, will trigger system library registration
ctypes.CDLL(path_dso)
# Load the system wide library
mm = tvm.runtime.system_lib()
a = tvm.nd.array(np.random.uniform(size=nn).astype(A.dtype), dev)
b = tvm.nd.array(np.zeros(nn, dtype=A.dtype), dev)
mm["myadd1"](a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
mm["myadd2"](a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
def test_memory_planning(workspace_byte_alignment, main_workspace_size,
sum_workspace_size):
mod, params = tvm.relay.testing.synthetic.get_workload()
target = "c"
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": workspace_byte_alignment,
},
)
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
lib = tvm.relay.build(mod,
target,
executor=executor,
runtime=runtime,
params=params)
assert (sum(lib.function_metadata["__tvm_main__"].workspace_sizes.values())
== main_workspace_size)
assert (sum([
size for metadata in lib.function_metadata.values()
for size in metadata.workspace_sizes.values()
]) == sum_workspace_size)
def test_export_operator_model_library_format():
import tvm.micro as micro
target = tvm.target.target.micro("host")
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
A = tvm.te.placeholder((2,), dtype="int8")
B = tvm.te.placeholder((1,), dtype="int8")
C = tvm.te.compute(A.shape, lambda i: A[i] + B[0], name="C")
sched = tvm.te.create_schedule(C.op)
mod = tvm.build(
sched,
[A, B, C],
tvm.target.Target(target, target),
runtime=Runtime("crt", {"system-lib": True}),
name="add",
)
temp_dir = utils.tempdir()
mlf_tar_path = temp_dir.relpath("lib.tar")
micro.export_model_library_format(mod, mlf_tar_path)
tf = tarfile.open(mlf_tar_path)
extract_dir = temp_dir.relpath("extract")
os.mkdir(extract_dir)
tf.extractall(extract_dir)
with open(os.path.join(extract_dir, "metadata.json")) as json_f:
metadata = json.load(json_f)
assert metadata["version"] == 5
assert metadata["model_name"] == "add"
export_datetime = datetime.datetime.strptime(
metadata["export_datetime"], "%Y-%m-%d %H:%M:%SZ"
)
assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5)
assert metadata["target"] == {"1": str(target)}
assert metadata["memory"]["add"][0]["dtype"] == "int8"
assert metadata["memory"]["add"][0]["shape"] == [2]
assert metadata["memory"]["add"][0]["size_bytes"] == 2
assert metadata["memory"]["add"][1]["dtype"] == "int8"
assert metadata["memory"]["add"][1]["shape"] == [1]
assert metadata["memory"]["add"][1]["size_bytes"] == 1
assert metadata["memory"]["add"][2]["dtype"] == "int8"
assert metadata["memory"]["add"][2]["shape"] == [2]
assert metadata["memory"]["add"][2]["size_bytes"] == 2
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "lib0.c"))
assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "lib1.c"))
assert (
len(mod.ir_module_by_target) == 1
), f"expect 1 ir_model_by_target: {mod.ir_module_by_target!r}"
for target, ir_mod in mod.ir_module_by_target.items():
assert int(tvm.runtime.ndarray.device(str(target)).device_type) == 1
with open(os.path.join(extract_dir, "src", "tir-1.txt")) as tir_f:
assert tir_f.read() == str(ir_mod)
def test_graph_executor(hexagon_session: Session):
"""Test graph executor"""
dtype = "float32"
data = relay.var("data", relay.TensorType((1, 64, 64, 3), dtype))
weight = relay.var("weight", relay.TensorType((5, 5, 3, 8), dtype))
conv2d_op = relay.nn.conv2d(
data,
weight,
padding=(2, 2),
kernel_size=(5, 5),
data_layout="NHWC",
kernel_layout="HWIO",
out_dtype="float32",
)
f = relay.Function([data, weight], conv2d_op)
relay_mod = tvm.IRModule.from_expr(f)
relay_mod = relay.transform.InferType()(relay_mod)
target_hexagon = tvm.target.hexagon("v68")
runtime = Runtime("cpp")
executor = Executor("graph")
weight_in = np.random.rand(5, 5, 3, 8).astype(dtype=dtype)
data_in = np.random.rand(1, 64, 64, 3).astype(dtype=dtype)
params = {"weight": weight_in}
inputs = {"data": data_in}
with tvm.transform.PassContext(opt_level=3):
lowered = tvm.relay.build(
relay_mod,
tvm.target.Target(target_hexagon, host=target_hexagon),
runtime=runtime,
executor=executor,
)
graph_mod = hexagon_session.get_executor_from_factory(lowered)
graph_mod.set_input(**params)
graph_mod.run(**inputs)
hexagon_output = graph_mod.get_output(0).numpy()
target_llvm = tvm.target.Target("llvm")
with tvm.transform.PassContext(opt_level=3):
llvm_lowered = tvm.relay.build(
relay_mod,
tvm.target.Target(target_llvm, host=target_llvm),
runtime=runtime,
executor=executor,
)
llvm_graph_mod = tvm.contrib.graph_executor.GraphModule(
llvm_lowered["default"](tvm.cpu(0)))
llvm_graph_mod.set_input(**params)
llvm_graph_mod.run(**inputs)
expected_output = llvm_graph_mod.get_output(0).numpy()
tvm.testing.assert_allclose(hexagon_output,
expected_output,
rtol=1e-4,
atol=1e-5)
def _make_sess_from_op(
model, arduino_board, arduino_cli_cmd, workspace_dir, op_name, sched, arg_bufs, build_config
):
target = tvm.target.target.micro(model)
runtime = Runtime("crt", {"system-lib": True})
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
mod = tvm.build(sched, arg_bufs, target=target, runtime=runtime, name=op_name)
return _make_session(model, arduino_board, arduino_cli_cmd, workspace_dir, mod, build_config)
def test_aot_executor(workspace_dir, board, west_cmd, microtvm_debug, use_fvp):
"""Test use of the AOT executor with microTVM."""
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": microtvm_debug}
shape = (10, )
dtype = "int8"
print("test_relay: construct relay program\n")
# Construct Relay program.
relay_mod = tvm.parser.fromtext("""
#[version = "0.0.5"]
def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), uint8]) {
%0 = %a + %b;
%0
}""")
runtime = Runtime("crt", {"system-lib": True})
executor = Executor("aot")
target = tvm.target.target.micro(model)
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
mod = tvm.relay.build(relay_mod,
target=target,
runtime=runtime,
executor=executor)
def do_test():
aot_executor = tvm.runtime.executor.aot_executor.AotModule(
session.create_aot_executor())
assert aot_executor.get_input_index("a") == 0
assert aot_executor.get_input_index("b") == 1
assert aot_executor.get_num_inputs() == 2
assert aot_executor.get_num_outputs() == 1
A_np = np.array([[2, 3]], dtype="uint8")
B_np = np.array([[4, 7]], dtype="uint8")
A_data = aot_executor.get_input("a").copyfrom(A_np)
B_data = aot_executor.get_input("b").copyfrom(B_np)
aot_executor.run()
out = aot_executor.get_output(0)
assert (out.numpy() == np.array([6, 10])).all()
B_np_new = np.array([[5, 8]])
aot_executor.set_input("b", B_np_new)
assert (B_data.numpy() == B_np_new).all()
with _make_session(workspace_dir, board, west_cmd, mod, build_config,
use_fvp) as session:
do_test()
def test_rpc():
if not tvm.runtime.enabled("rpc"):
return
# generate the wasm library
runtime = Runtime("cpp", {"system-lib": True})
target = "llvm -mtriple=wasm32-unknown-unknown-wasm"
if not tvm.runtime.enabled(target):
raise RuntimeError("Target %s is not enbaled" % target)
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
s = te.create_schedule(B.op)
fadd = tvm.build(s, [A, B], target, runtime=runtime, name="addone")
temp = utils.tempdir()
wasm_path = temp.relpath("addone.wasm")
fadd.export_library(wasm_path, emcc.create_tvmjs_wasm)
wasm_binary = open(wasm_path, "rb").read()
remote = rpc.connect(
proxy_host,
proxy_port,
key="wasm",
session_constructor_args=["rpc.WasmSession", wasm_binary],
)
def check(remote):
# basic function checks.
faddone = remote.get_function("testing.asyncAddOne")
fecho = remote.get_function("testing.echo")
assert faddone(100) == 101
assert fecho(1, 2, 3) == 1
assert fecho(1, 2, 3) == 1
assert fecho(100, 2, 3) == 100
assert fecho("xyz") == "xyz"
assert bytes(fecho(bytearray(b"123"))) == b"123"
# run the generated library.
f1 = remote.system_lib()
dev = remote.cpu(0)
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), dev)
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
# invoke the function
addone = f1.get_function("addone")
addone(a, b)
# time evaluator
time_f = f1.time_evaluator("addone", dev, number=100, repeat=10)
time_f(a, b)
cost = time_f(a, b).mean
print("%g secs/op" % cost)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
check(remote)
def test_mobilenet(hexagon_session):
import onnx
dtype = "float32"
model_url = "https://github.com/onnx/models/raw/main/vision/classification/mobilenet/model/mobilenetv2-7.onnx"
model_path = tvm.contrib.download.download_testdata(model_url,
"mobilenetv2-7.onnx",
module="onnx")
onnx_model = onnx.load(model_path)
target_hexagon = tvm.target.hexagon("v68")
target_llvm = tvm.target.Target("llvm")
runtime = Runtime("cpp")
executor = Executor("graph", {"link-params": True})
data_in = np.random.rand(1, 3, 224, 224).astype(dtype=dtype)
input_name = "input"
shape_dict = {input_name: data_in.shape}
relay_mod, params = relay.frontend.from_onnx(onnx_model,
shape_dict,
freeze_params=True)
inputs = {input_name: data_in}
with tvm.transform.PassContext(opt_level=3):
hexagon_lowered = tvm.relay.build(
relay_mod,
tvm.target.Target(target_hexagon, host=target_hexagon),
runtime=runtime,
executor=executor,
params=params,
)
llvm_lowered = tvm.relay.build(
relay_mod,
tvm.target.Target(target_llvm, host=target_llvm),
runtime=runtime,
executor=executor,
params=params,
)
graph_mod = hexagon_session.get_executor_from_factory(hexagon_lowered)
graph_mod.set_input(**inputs)
graph_mod.run()
hexagon_output = graph_mod.get_output(0).numpy()
llvm_graph_mod = tvm.contrib.graph_executor.GraphModule(
llvm_lowered["default"](tvm.cpu(0)))
llvm_graph_mod.set_input(**inputs)
llvm_graph_mod.run()
expected_output = llvm_graph_mod.get_output(0).numpy()
tvm.testing.assert_allclose(hexagon_output,
expected_output,
rtol=1e-4,
atol=1e-5)
def _make_sess_from_op(
temp_dir, model, zephyr_board, west_cmd, op_name, sched, arg_bufs, build_config, use_fvp
):
runtime = Runtime("crt", {"system-lib": True})
target = tvm.target.target.micro(model)
target = tvm.target.Target(target=target, host=target)
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
mod = tvm.build(sched, arg_bufs, target=target, runtime=runtime, name=op_name)
return _make_session(temp_dir, zephyr_board, west_cmd, mod, build_config, use_fvp)
def main():
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
s = tvm.te.create_schedule(C.op)
s[C].parallel(s[C].op.axis[0])
runtime = Runtime("cpp", {"system-lib": True})
print(tvm.lower(s, [A, B, C], simple_mode=True))
tvm.build(s, [A, B, C], "llvm", runtime=runtime).save(osp.join(sys.argv[1], "test.o"))
def test_onnx(temp_dir, board, west_cmd, tvm_debug):
"""Testing a simple ONNX model."""
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": tvm_debug}
this_dir = pathlib.Path(os.path.dirname(__file__))
mnist_testdata = this_dir.parent / "testdata" / "mnist"
digit_2 = Image.open(mnist_testdata / "digit-2.jpg").resize((28, 28))
digit_2 = np.asarray(digit_2).astype("float32")
digit_2 = np.expand_dims(digit_2, axis=0)
digit_9 = Image.open(mnist_testdata / "digit-9.jpg").resize((28, 28))
digit_9 = np.asarray(digit_9).astype("float32")
digit_9 = np.expand_dims(digit_9, axis=0)
# Load ONNX model and convert to Relay.
onnx_model = onnx.load(mnist_testdata / "mnist-8.onnx")
shape = {"Input3": (1, 1, 28, 28)}
relay_mod, params = relay.frontend.from_onnx(onnx_model,
shape=shape,
freeze_params=True)
relay_mod = relay.transform.DynamicToStatic()(relay_mod)
# We add the link-params=True option to ensure the model parameters are compiled in.
# There is currently a bug preventing the host_driven environment from receiving
# the model weights when set using graph_mod.set_input().
# See: https://github.com/apache/tvm/issues/7567
target = tvm.target.target.micro(model)
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
executor = Executor("graph", {"link-params": True})
runtime = Runtime("crt", {"system-lib": True})
lowered = relay.build(relay_mod,
target,
params=params,
executor=executor,
runtime=runtime)
graph = lowered.get_graph_json()
with _make_session(temp_dir, board, west_cmd, lowered,
build_config) as session:
graph_mod = tvm.micro.create_local_graph_executor(
graph, session.get_system_lib(), session.device)
# Send the digit-2 image and confirm that the correct result is returned.
graph_mod.set_input("Input3", tvm.nd.array(digit_2))
graph_mod.run()
result = graph_mod.get_output(0).numpy()
assert np.argmax(result) == 2
# Send the digit-9 image and confirm that the correct result is returned.
graph_mod.set_input("Input3", tvm.nd.array(digit_9))
graph_mod.run()
result = graph_mod.get_output(0).numpy()
assert np.argmax(result) == 9
