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_multiple_relay_modules_aot_graph():
mod = get_conv2d_relay_module()
executor1 = Executor("graph")
executor2 = 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=executor1,
mod_name="mod1")
factory2 = tvm.relay.build(mod,
target,
runtime=runtime,
executor=executor2,
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", "mod1_lib2.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_mod2.h"))
with open(os.path.join(extract_dir, "metadata.json")) as f:
metadata = json.load(f)
assert metadata["modules"]["mod1"]["executors"] == ["graph"]
assert metadata["modules"]["mod2"]["executors"] == ["aot"]
assert metadata["version"] == _GENERATED_VERSION
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 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 test_create_executor_attr_type_incorrect():
with pytest.raises(
TVMError,
match='Attribute "interface-api" should have type "runtime.String"'
' but instead found "IntImm"',
):
Executor("aot", {"interface-api": True})
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_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_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 _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_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_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_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 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 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_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 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_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_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 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
def _get_ethosu_workspace_size(mod, params, accel_type, pool_size,
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,
}
workspace_memory_pools = WorkspaceMemoryPools([
PoolInfo(
"SRAM",
{
target: PoolInfo.READ_WRITE_ACCESS,
ethosu_target: PoolInfo.READ_WRITE_ACCESS
},
size_hint_bytes=pool_size,
read_bandwidth_bytes_per_cycle=16,
write_bandwidth_bytes_per_cycle=16,
target_burst_bytes={ethosu_target: 1},
),
])
with tvm.transform.PassContext(opt_level=3, config=pass_config):
lib = tvm.relay.build(
mod,
target,
executor=executor,
runtime=runtime,
workspace_memory_pools=workspace_memory_pools,
params=params,
)
mlf_memory_map = mlf._build_function_memory_map(lib.function_metadata)
return mlf_memory_map["main"][0]["workspace_size_bytes"]
def test_onnx(board, arduino_cli_cmd, tvm_debug, workspace_dir):
"""Testing a simple ONNX model."""
model = test_utils.ARDUINO_BOARDS[board]
build_config = {"debug": tvm_debug}
# Load test images.
this_dir = pathlib.Path(__file__).parent
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)
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(model, board, arduino_cli_cmd, workspace_dir, 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()
print(result)
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
def test_aot_executor():
"""Test use of the AOT 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})
executor = Executor("aot")
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
factory = tvm.relay.build(relay_mod,
target=TARGET,
runtime=runtime,
executor=executor)
def do_test():
aot_executor = tvm.runtime.executor.aot_executor.AotModule(
sess._rpc.get_function("tvm.aot_executor.create")(
sess.get_system_lib(), sess.device, "default"))
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(temp_dir, factory) as sess:
do_test()
def test_compile_tflite_module_with_mod_name_and_ethosu(
tmpdir_factory, tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
pytest.importorskip("ethosu.vela")
output_dir = tmpdir_factory.mktemp("mlf")
tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
output_file_name = f"{output_dir}/file.tar"
tvmc.compiler.compile_model(
tvmc_model,
target=f"ethos-u -accelerator_config=ethos-u55-256, c -mcpu=cortex-m55",
runtime=Runtime("crt"),
executor=Executor("aot", {"unpacked-api": True}),
output_format="mlf",
package_path=output_file_name,
pass_context_configs=["tir.disable_vectorize=true"],
mod_name="classify",
)
# check that an MLF package was created
assert os.path.exists(output_file_name)
with tarfile.open(output_file_name) as mlf_package:
# check that the C source files have been named classify_lib*.c
c_source_files = [
name for name in mlf_package.getnames()
if re.match(r"\./codegen/host/src/classify_lib\d+\.c", name)
]
assert len(c_source_files) > 0
# check that "default" doesn't occur in any of the C source files
# check that function names are of the form "tvmgen_classify_*"
for file_name in c_source_files:
with mlf_package.extractfile(file_name) as f:
content = f.read()
assert b"default" not in content
assert b"tvmgen_classify_" in content
# check that tvmgen_classify_run() function exists
with mlf_package.extractfile(
"./codegen/host/src/classify_lib0.c") as f:
content = f.read()
assert b"tvmgen_classify_run(" in content
# check that microNPU function names are of the form "tvmgen_classify_ethos_u_main_*"
with mlf_package.extractfile(
"./codegen/host/src/classify_lib2.c") as f:
content = f.read()
assert b"tvmgen_classify_ethos_u_main_" in content
def test_qemu_make_fail(workspace_dir, board, west_cmd, microtvm_debug):
"""Testing QEMU make fail."""
if board not in ["qemu_x86", "mps2_an521", "mps3_an547"]:
pytest.skip(msg="Only for QEMU targets.")
model = test_utils.ZEPHYR_BOARDS[board]
build_config = {"debug": microtvm_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)
sample = np.zeros(shape=shape, dtype=dtype)
project, project_dir = test_utils.generate_project(
workspace_dir,
board,
west_cmd,
lowered,
build_config,
sample,
shape,
dtype,
load_cmsis=False,
)
file_path = pathlib.Path(project_dir) / "build" / "build.ninja"
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_pass_link_params():
"""
This test checks ensures that proper executor is passed to interpreter instance
The test will fail if FoldConstant does not override the executor due to "int8"
is not supported in ScheduleBuilder
"""
def expr():
z = relay.const(10, dtype="int8")
return relay.cast(z, dtype="int32")
mod = tvm.IRModule.from_expr(expr())
mod = tvm.relay.transform.InferType()(mod)
# Add executor with link-params
mod = mod.with_attr("executor", Executor("aot", {"link-params": True}))
mod = tvm.relay.transform.FoldConstant()(mod)
def test_aot_uses_anf():
"""Checks that A-Normal Form is being used in the AOT lowering pipeline."""
x = relay.var("x", shape=(1, 10, 10, 10))
y = relay.var("y", shape=(1, 10, 10, 10))
z = relay.add(x, y)
func = relay.Function([x, y], z)
@pass_instrument
class CheckANFRuns:
def __init__(self):
self.did_run_anf = False
def run_before_pass(self, _, info):
if info.name == "ToANormalForm":
self.did_run_anf = True
if info.name == "LowerTE":
assert self.did_run_anf, "ToANormalForm pass should run before LowerTE."
check_run_anf = CheckANFRuns()
model = AOTTestModel(module=IRModule.from_expr(func),
inputs=None,
outputs=None)
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": 8,
"interface-api": "c",
"unpacked-api": True,
},
)
config = {"tir.disable_vectorize": True}
with tvm.transform.PassContext(opt_level=3,
config=config,
instruments=[check_run_anf]):
tvm.relay.build(
model.module,
tvm.target.Target("c"),
executor=executor,
runtime=runtime,
workspace_memory_pools=None,
params=model.params,
mod_name=model.name,
)
assert check_run_anf.did_run_anf, "Expected ToANormalForm pass to have run."
def test_workspace_calculation_cmsis_nn():
"""This tests cmsis_nn codegen for workspace calculation.
This is tested specially because cmsis-nn codegen creates
multiple PrimFuncs per offloaded relay function in a non
-hierarchical manner."""
pytest.importorskip("tflite")
# pylint: disable=import-outside-toplevel
from tvm.relay.op.contrib import cmsisnn
from tvm.contrib.download import download_testdata
# pylint: enable=import-outside-toplevel
target = "c"
runtime = Runtime("crt")
executor = Executor(
"aot",
{
"workspace-byte-alignment": 16,
"interface-api": "c",
"unpacked-api": True,
},
)
base_url = ("https://github.com/ARM-software/ML-zoo/raw/"
"48a22ee22325d15d2371a6df24eb7d67e21dcc97"
"/models/keyword_spotting/cnn_small/tflite_int8")
file_to_download = "cnn_s_quantized.tflite"
file_saved = "cnn_s_quantized_15Dec2021.tflite"
model_file = download_testdata("{}/{}".format(base_url, file_to_download),
file_saved)
mod, _, params = create_relay_module_and_inputs_from_tflite_file(
model_file)
mod = cmsisnn.partition_for_cmsisnn(mod, params)
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"] == 14384
def test_mobilenet(hexagon_session: Session):
"""Test mobilenet with graph executor"""
dtype = "float32"
onnx_model = get_mobilenet()
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_kws_project(board, arduino_cli_cmd, tvm_debug, workspace_dir):
this_dir = pathlib.Path(__file__).parent
model = ARDUINO_BOARDS[board]
build_config = {"debug": tvm_debug}
with open(this_dir.parent / "testdata" / "kws" / "yes_no.tflite",
"rb") as f:
tflite_model_buf = f.read()
# TFLite.Model.Model has changed to TFLite.Model from 1.14 to 2.1
try:
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
mod, params = relay.frontend.from_tflite(tflite_model)
target = tvm.target.target.micro(model)
runtime = Runtime("crt")
executor = Executor("aot", {"unpacked-api": True})
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
mod = relay.build(mod,
target,
runtime=runtime,
executor=executor,
params=params)
return tvm.micro.generate_project(
str(TEMPLATE_PROJECT_DIR),
mod,
workspace_dir / "project",
{
"arduino_board": board,
"arduino_cli_cmd": arduino_cli_cmd,
"project_type": "example_project",
"verbose": bool(build_config.get("debug")),
},
)
def test_compile_tflite_module_with_external_codegen_ethosu(
tmpdir_factory, tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
pytest.importorskip("ethosu.vela")
ACCEL_TYPES = [
"ethos-u55-256", "ethos-u55-128", "ethos-u55-64", "ethos-u55-32"
]
output_dir = tmpdir_factory.mktemp("mlf")
tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
for accel_type in ACCEL_TYPES:
output_file_name = f"{output_dir}/file_{accel_type}.tar"
tvmc.compiler.compile_model(
tvmc_model,
target=
f"ethos-u -accelerator_config={accel_type}, c -mcpu=cortex-m55",
runtime=Runtime("crt"),
executor=Executor("aot", {"unpacked-api": True}),
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)
]
# The number of c_source_files depends on the number of fused subgraphs that
# get offloaded to the NPU, e.g. conv2d->depthwise_conv2d->conv2d gets offloaded
# as a single subgraph if both of these operators are supported by the NPU.
# Currently there are two source files for CPU execution and one offload graph
assert len(c_source_files) == 3
