def test_compile_check_configs_composite_target(mock_pkg, mock_pc, mock_fe,
mock_ct, mock_relay):
mock_codegen = {}
mock_codegen["config_key"] = "relay.ext.mock.options"
mock_codegen["pass_pipeline"] = lambda *args, **kwargs: None
mock_fe.return_value = mock.MagicMock()
mock_ct.return_value = mock_codegen
mock_relay.return_value = mock.MagicMock()
tvmc_model = tvmc.load("no_file_needed")
tvmc.compile(tvmc_model, target="mockcodegen -testopt=value, llvm")
assert mock_pc.call_count == 2
codegen_partition_context = mock.call(
config={"relay.ext.mock.options": {
"testopt": "value"
}}, )
codegen_compile_context = mock.call(
config={"relay.ext.mock.options": {
"testopt": "value"
}},
opt_level=3,
disabled_pass=None,
)
mock_pc.assert_has_calls([
codegen_partition_context,
codegen_partition_context.__enter__(),
codegen_partition_context.__exit__(None, None, None),
codegen_compile_context,
codegen_compile_context.__enter__(),
codegen_compile_context.__exit__(None, None, None),
])
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_package = tvmc.compiler.compile_model(
tvmc_model,
target=
f"cmsis-nn, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 --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_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_package = tvmc.compiler.compile_model(
tvmc_model,
target=
f"ethos-u -accelerator_config={accel_type}, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 --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) == 17
def test_load_model__paddle(paddle_resnet50):
# some CI environments wont offer Paddle, so skip in case it is not present
pytest.importorskip("paddle")
tvmc_model = tvmc.load(paddle_resnet50, model_format="paddle")
assert type(tvmc_model) is TVMCModel
assert type(tvmc_model.mod) is IRModule
assert type(tvmc_model.params) is dict
def test_load_model__pb(pb_mobilenet_v1_1_quant):
# some CI environments wont offer TensorFlow, so skip in case it is not present
pytest.importorskip("tensorflow")
mod, params = tvmc.load(pb_mobilenet_v1_1_quant)
assert type(mod) is IRModule
assert type(params) is dict
# check whether one known value is part of the params dict
assert "MobilenetV1/Conv2d_0/weights" in params.keys()
def test_load_model__tflite(tflite_mobilenet_v1_1_quant):
# some CI environments wont offer TFLite, so skip in case it is not present
pytest.importorskip("tflite")
mod, params = tvmc.load(tflite_mobilenet_v1_1_quant)
assert type(mod) is IRModule
assert type(params) is dict
# check whether one known value is part of the params dict
assert "_param_1" in params.keys()
def verify_compile_onnx_module(model, shape_dict=None, use_vm=False):
# some CI environments wont offer onnx, so skip in case it is not present
pytest.importorskip("onnx")
tvmc_model = tvmc.load(model, shape_dict=shape_dict)
tvmc_package = tvmc.compile(tvmc_model,
target="llvm",
dump_code="ll",
use_vm=use_vm)
dumps_path = tvmc_package.package_path + ".ll"
verify_tvmc_package(tvmc_package, dumps_path, use_vm=use_vm)
def verify_compile_tflite_module(model, shape_dict=None, use_vm=False):
pytest.importorskip("tflite")
tvmc_model = tvmc.load(model, shape_dict=shape_dict)
tvmc_package = tvmc.compile(tvmc_model,
target="llvm",
dump_code="ll",
desired_layout="NCHW",
use_vm=use_vm)
dumps_path = tvmc_package.package_path + ".ll"
verify_tvmc_package(tvmc_package, dumps_path, use_vm=use_vm)
def test_load_model__pth(pytorch_resnet18):
# some CI environments wont offer torch, so skip in case it is not present
pytest.importorskip("torch")
pytest.importorskip("torchvision")
mod, params = tvmc.load(pytorch_resnet18,
shape_dict={"input": [1, 3, 224, 224]})
assert type(mod) is IRModule
assert type(params) is dict
# check whether one known value is part of the params dict
assert "layer1.0.conv1.weight" in params.keys()
def verify_compile_onnx_module(model, shape_dict=None):
# some CI environments wont offer onnx, so skip in case it is not present
pytest.importorskip("onnx")
mod, params = tvmc.load(model, shape_dict=shape_dict)
graph, lib, params, dumps = tvmc.compile(mod, params, target="llvm", dump_code="ll")
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
assert "ll" in dumps.keys()
def test_compile_tflite_module_with_external_codegen(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
mod, params = tvmc.load(tflite_mobilenet_v1_1_quant)
graph, lib, params, dumps = tvmc.compile(
mod, params, target="ethos-n77, llvm", dump_code="relay"
)
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
def verify_compile_tflite_module(model, shape_dict=None):
pytest.importorskip("tflite")
mod, params = tvmc.load(model, shape_dict=shape_dict)
graph, lib, params, dumps = tvmc.compile(
mod, params, target="llvm", dump_code="ll", alter_layout="NCHW"
)
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
def verify_compile_tflite_module(model, shape_dict=None):
pytest.importorskip("tflite")
tvmc_model = tvmc.load(model, shape_dict=shape_dict)
tvmc_package = tvmc.compile(tvmc_model, target="llvm", dump_code="ll", desired_layout="NCHW")
dumps_path = tvmc_package.package_path + ".ll"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)
def test_compile_tflite_module_with_external_codegen(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
tvmc_package = tvmc.compile(tvmc_model, target="ethos-n77, llvm", dump_code="relay")
dumps_path = tvmc_package.package_path + ".relay"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)
def test_load_quantized_model__pth(pytorch_mobilenetv2_quantized):
# some CI environments wont offer torch, so skip in case it is not present
pytest.importorskip("torch")
pytest.importorskip("torchvision")
tvmc_model = tvmc.load(pytorch_mobilenetv2_quantized, shape_dict={"input": [1, 3, 224, 224]})
assert type(tvmc_model) is TVMCModel
assert type(tvmc_model.mod) is IRModule
assert type(tvmc_model.params) is dict
# checking weights remain quantized and are not float32
for p in tvmc_model.params.values():
assert p.dtype in ["int8", "uint8", "int32"] # int32 for bias
def test_compile_keras__save_module(keras_resnet50, tmpdir_factory):
# some CI environments wont offer tensorflow/Keras, so skip in case it is not present
pytest.importorskip("tensorflow")
mod, params = tvmc.load(keras_resnet50)
graph, lib, params, dumps = tvmc.compile(mod, params, target="llvm", dump_code="ll")
expected_temp_dir = tmpdir_factory.mktemp("saved_output")
expected_file_name = "saved.tar"
module_file = os.path.join(expected_temp_dir, expected_file_name)
tvmc.compiler.save_module(module_file, graph, lib, params)
assert os.path.exists(module_file), "output file {0} should exist".format(module_file)
def verify_compile_onnx_module(model, shape_dict=None):
# some CI environments wont offer onnx, so skip in case it is not present
pytest.importorskip("onnx")
tvmc_model = tvmc.load(model, shape_dict=shape_dict)
tvmc_package = tvmc.compile(tvmc_model, target="llvm", dump_code="ll")
dumps_path = tvmc_package.package_path + ".ll"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)
def test_tvmc_workflow(keras_simple):
pytest.importorskip("tensorflow")
tvmc_model = tvmc.load(keras_simple)
tuning_records = tvmc.tune(tvmc_model, target="llvm", enable_autoscheduler=True, trials=2)
tvmc_package = tvmc.compile(tvmc_model, tuning_records=tuning_records, target="llvm")
result = tvmc.run(tvmc_package, device="cpu")
assert type(tvmc_model) is TVMCModel
assert type(tvmc_package) is TVMCPackage
assert type(result) is TVMCResult
assert path.exists(tuning_records)
assert type(result.outputs) is dict
assert type(result.times) is BenchmarkResult
assert "output_0" in result.outputs.keys()
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_compile_check_workspace_pools(mock_pkg, mock_fe, mock_relay):
mock_fe.return_value = mock.MagicMock()
mock_relay.return_value = mock.MagicMock()
memory_pools = WorkspaceMemoryPools(
[WorkspacePoolInfo(pool_name="sram", targets=[Target("llvm")])])
tvmc_model = tvmc.load("no_file_needed")
tvmc.compile(
tvmc_model,
target="llvm,c",
workspace_pools=memory_pools,
)
assert mock_relay.call_count == 1
assert mock_relay.call_args_list[0][1][
"workspace_memory_pools"] == memory_pools
def test_compile_opencl(tflite_mobilenet_v1_0_25_128):
pytest.importorskip("tflite")
mod, params = tvmc.load(tflite_mobilenet_v1_0_25_128)
graph, lib, params, dumps = tvmc.compile(
mod,
params,
target="opencl --host=llvm",
alter_layout="NCHW",
)
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
def test_cross_compile_aarch64_tflite_module(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
mod, params = tvmc.load(tflite_mobilenet_v1_1_quant)
graph, lib, params, dumps = tvmc.compile(
mod,
params,
target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr='+neon'",
dump_code="asm",
)
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
def test_compile_tflite_module_with_external_codegen_vitis_ai(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
mod, params = tvmc.load(tflite_mobilenet_v1_1_quant)
graph, lib, params, dumps = tvmc.compiler.compile_model(
mod,
params,
target="vitis-ai -dpu=DPUCZDX8G-zcu104 -export_runtime_module=vitis_ai.rtmod, llvm",
dump_code="relay",
)
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
def test_compile_opencl(tflite_mobilenet_v1_0_25_128):
pytest.importorskip("tflite")
tvmc_model = tvmc.load(tflite_mobilenet_v1_0_25_128)
tvmc_package = tvmc.compile(
tvmc_model,
target="opencl --host=llvm",
desired_layout="NCHW",
)
dumps_path = tvmc_package.package_path + ".asm"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)
def test_compile_tflite_module_with_external_codegen_vitis_ai(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
tvmc_package = tvmc.compiler.compile_model(
tvmc_model,
target="vitis-ai -dpu=DPUCZDX8G-zcu104 -export_runtime_module=vitis_ai.rtmod, llvm",
dump_code="relay",
)
dumps_path = tvmc_package.package_path + ".relay"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)
def test_compile_check_configs_composite_target(mock_pc, mock_fe, mock_ct, mock_relay):
mock_codegen = {}
mock_codegen["config_key"] = "relay.ext.mock.options"
mock_codegen["pass_pipeline"] = lambda *args, **kwargs: None
mock_fe.return_value = (None, None)
mock_ct.return_value = mock_codegen
mock_relay.return_value = mock.MagicMock()
mod, params = tvmc.load("no_file_needed")
graph, lib, params, dumps = tvmc.compile(mod, params, target="mockcodegen -testopt=value, llvm")
mock_pc.assert_called_once_with(
opt_level=3,
config={"relay.ext.mock.options": {"testopt": "value"}},
disabled_pass=None,
)
def test_cross_compile_aarch64_onnx_module(onnx_resnet50):
# some CI environments wont offer onnx, so skip in case it is not present
pytest.importorskip("onnx")
mod, params = tvmc.load(onnx_resnet50)
graph, lib, params, dumps = tvmc.compile(
mod,
params,
target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr=+neon",
dump_code="asm",
)
# check for output types
assert type(graph) is str
assert type(lib) is tvm.runtime.module.Module
assert type(params) is dict
assert type(dumps) is dict
assert "asm" in dumps.keys()
def test_cross_compile_aarch64_tflite_module(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
tvmc_package = tvmc.compile(
tvmc_model,
target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr='+neon'",
dump_code="asm",
cross="aarch64-linux-gnu-gcc",
)
dumps_path = tvmc_package.package_path + ".asm"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)
def test_compile_keras__save_module(keras_resnet50, tmpdir_factory):
# some CI environments wont offer tensorflow/Keras, so skip in case it is not present
pytest.importorskip("tensorflow")
expected_temp_dir = tmpdir_factory.mktemp("saved_output")
expected_file_name = "saved.tar"
module_file = os.path.join(expected_temp_dir, expected_file_name)
tvmc_model = tvmc.load(keras_resnet50)
tvmc.compile(tvmc_model, target="llvm", dump_code="ll", package_path=module_file)
assert os.path.exists(module_file), "output file {0} should exist".format(module_file)
# Test that we can load back in a module.
tvmc_package = TVMCPackage(package_path=module_file)
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.graph) is str
assert type(tvmc_package.params) is bytearray
def test_cross_compile_aarch64_paddle_module(paddle_resnet50):
# some CI environments wont offer paddle, so skip in case it is not present
pytest.importorskip("paddle")
tvmc_model = tvmc.load(paddle_resnet50, "paddle")
tvmc_package = tvmc.compile(
tvmc_model,
target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr=+neon",
dump_code="asm",
cross="aarch64-linux-gnu-gcc",
)
dumps_path = tvmc_package.package_path + ".asm"
# check for output types
assert type(tvmc_package) is TVMCPackage
assert type(tvmc_package.graph) is str
assert type(tvmc_package.lib_path) is str
assert type(tvmc_package.params) is bytearray
assert os.path.exists(dumps_path)