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_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_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")
mock_pc.assert_called_once_with(
opt_level=3,
config={"relay.ext.mock.options": {"testopt": "value"}},
disabled_pass=None,
)
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 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 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 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 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 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_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 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 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_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 test_cross_compile_aarch64_tflite_module(tflite_mobilenet_v1_1_quant):
pytest.importorskip("tflite")
graph, lib, params, dumps = tvmc.compile(
tflite_mobilenet_v1_1_quant,
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_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_save_load_model(keras_simple, tmpdir_factory):
pytest.importorskip("onnx")
tmpdir = tmpdir_factory.mktemp("data")
tvmc_model = tvmc.load(keras_simple)
# Create tuning artifacts
tvmc.tune(tvmc_model, target="llvm", trials=2)
# Create package artifacts
tvmc.compile(tvmc_model, target="llvm")
# Save the model to disk
model_path = os.path.join(tmpdir, "saved_model.tar")
tvmc_model.save(model_path)
# Load the model into a new TVMCModel
new_tvmc_model = TVMCModel(model_path=model_path)
# Check that the two models match.
assert str(new_tvmc_model.mod) == str(tvmc_model.mod)
# Check that tuning records and the compiled package are recoverable.
assert path.exists(new_tvmc_model.default_package_path())
assert path.exists(new_tvmc_model.default_tuning_records_path())
def test_compile_opencl(tflite_mobilenet_v1_0_25_128):
pytest.importorskip("tflite")
graph, lib, params, dumps = tvmc.compile(
tflite_mobilenet_v1_0_25_128,
target="opencl",
target_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_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_cross_compile_aarch64_onnx_module(onnx_resnet50):
# some CI environments wont offer onnx, so skip in case it is not present
pytest.importorskip("onnx")
graph, lib, params, dumps = tvmc.compile(
onnx_resnet50,
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_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: None
mock_fe.return_value = (None, None)
mock_ct.return_value = mock_codegen
mock_relay.return_value = mock.MagicMock()
graph, lib, params, dumps = tvmc.compile(
"no_file_needed", target="mockcodegen -testopt=value, llvm")
mock_pc.assert_called_once_with(
opt_level=3, config={"relay.ext.mock.options": {
"testopt": "value"
}})
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_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)
def test_cross_compile_options_aarch64_onnx_module(onnx_resnet50):
# some CI environments wont offer onnx, so skip in case it is not present
pytest.importorskip("onnx")
fake_sysroot_dir = utils.tempdir().relpath("")
tvmc_model = tvmc.load(onnx_resnet50)
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",
cross_options="--sysroot=" + fake_sysroot_dir,
)
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)
