From dff0bf9e949acb5370d447425a2abacd37f34f46 Mon Sep 17 00:00:00 2001 From: Christopher Sidebottom Date: Thu, 1 Jul 2021 10:37:07 +0100 Subject: [PATCH] Allow tvmc to compile models with AOT executor in MLF (#8331) * Allow tvmc to compile models with AOT executor The tflite_compiled_model fixture was getting duplicated a few times so I've added a parameterized fixture tflite_tvmc_compiler which combines tmpdir_factory setup with compile_model Nested targets broke a basic string split, so in cases where we use nested targets I replaced the string split with shlex split * Clarify that graph JSON is required only for graph executor Plus other clean ups * Change parametrize fixture to use string instead of list --- python/tvm/driver/tvmc/compiler.py | 6 +-- python/tvm/driver/tvmc/model.py | 13 ++++-- tests/python/driver/tvmc/conftest.py | 60 ++++--------------------- tests/python/driver/tvmc/test_mlf.py | 47 ++++++++++++++++--- tests/python/driver/tvmc/test_runner.py | 5 ++- 5 files changed, 66 insertions(+), 65 deletions(-) diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index e79a07f1de2e..9eb85a4934cb 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -171,7 +171,7 @@ def compile_model( target_host: Optional[str] = None, desired_layout: Optional[str] = None, disabled_pass: Optional[str] = None, - pass_context_configs: Optional[str] = None, + pass_context_configs: Optional[List[str]] = None, ): """Compile a model from a supported framework into a TVM module. @@ -212,8 +212,8 @@ def compile_model( disabled_pass: str, optional Comma-separated list of passes which needs to be disabled during compilation - pass_context_configs: str, optional - String containing a set of configurations to be passed to the + pass_context_configs: list[str], optional + List of strings containing a set of configurations to be passed to the PassContext. diff --git a/python/tvm/driver/tvmc/model.py b/python/tvm/driver/tvmc/model.py index c0ebb842b994..8c8828ddd49b 100644 --- a/python/tvm/driver/tvmc/model.py +++ b/python/tvm/driver/tvmc/model.py @@ -45,6 +45,7 @@ """ import os import tarfile +import json from typing import Optional, Union, List, Dict, Callable, TextIO import numpy as np @@ -332,8 +333,11 @@ def import_package(self, package_path: str): # Model Library Format (MLF) self.lib_name = None self.lib_path = None + with open(temp.relpath("metadata.json")) as metadata_json: + metadata = json.load(metadata_json) - graph = temp.relpath("runtime-config/graph/graph.json") + is_graph_runtime = "graph" in metadata["runtimes"] + graph = temp.relpath("runtime-config/graph/graph.json") if is_graph_runtime else None params = temp.relpath("parameters/default.params") self.type = "mlf" @@ -357,8 +361,11 @@ def import_package(self, package_path: str): with open(params, "rb") as param_file: self.params = bytearray(param_file.read()) - with open(graph) as graph_file: - self.graph = graph_file.read() + if graph is not None: + with open(graph) as graph_file: + self.graph = graph_file.read() + else: + self.graph = None class TVMCResult(object): diff --git a/tests/python/driver/tvmc/conftest.py b/tests/python/driver/tvmc/conftest.py index 9c0d8fa8911e..209c371a296a 100644 --- a/tests/python/driver/tvmc/conftest.py +++ b/tests/python/driver/tvmc/conftest.py @@ -41,25 +41,6 @@ def download_and_untar(model_url, model_sub_path, temp_dir): return os.path.join(temp_dir, model_sub_path) -def get_sample_compiled_module(target_dir, package_filename, output_format="so"): - """Support function that returns a TFLite compiled module""" - base_url = "https://storage.googleapis.com/download.tensorflow.org/models" - model_url = "mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz" - model_file = download_and_untar( - "{}/{}".format(base_url, model_url), - "mobilenet_v1_1.0_224_quant.tflite", - temp_dir=target_dir, - ) - - tvmc_model = tvmc.frontends.load_model(model_file) - return tvmc.compiler.compile_model( - tvmc_model, - target="llvm", - package_path=os.path.join(target_dir, package_filename), - output_format=output_format, - ) - - # PyTest fixtures @@ -167,40 +148,17 @@ def onnx_mnist(): return model_file -@pytest.fixture(scope="session") -def tflite_compiled_model(tmpdir_factory): - - # Not all CI environments will have TFLite installed - # so we need to safely skip this fixture that will - # crash the tests that rely on it. - # As this is a pytest.fixture, we cannot take advantage - # of pytest.importorskip. Using the block below instead. - try: - import tflite - except ImportError: - print("Cannot import tflite, which is required by tflite_compiled_module_as_tarfile.") - return "" - - target_dir = tmpdir_factory.mktemp("data") - return get_sample_compiled_module(target_dir, "mock.tar") - - -@pytest.fixture(scope="session") -def tflite_compiled_model_mlf(tmpdir_factory): +@pytest.fixture +def tflite_compile_model(tmpdir_factory): + """Support function that returns a TFLite compiled module""" - # Not all CI environments will have TFLite installed - # so we need to safely skip this fixture that will - # crash the tests that rely on it. - # As this is a pytest.fixture, we cannot take advantage - # of pytest.importorskip. Using the block below instead. - try: - import tflite - except ImportError: - print("Cannot import tflite, which is required by tflite_compiled_module_as_tarfile.") - return "" + def model_compiler(model_file, **overrides): + package_path = tmpdir_factory.mktemp("data").join("mock.tar") + tvmc_model = tvmc.frontends.load_model(model_file) + args = {"target": "llvm", **overrides} + return tvmc.compiler.compile_model(tvmc_model, package_path=package_path, **args) - target_dir = tmpdir_factory.mktemp("data") - return get_sample_compiled_module(target_dir, "mock.tar", "mlf") + return model_compiler @pytest.fixture(scope="session") diff --git a/tests/python/driver/tvmc/test_mlf.py b/tests/python/driver/tvmc/test_mlf.py index 48be5a810bc5..4669fab916a6 100644 --- a/tests/python/driver/tvmc/test_mlf.py +++ b/tests/python/driver/tvmc/test_mlf.py @@ -17,6 +17,7 @@ import pytest import os +import shlex import tvm from tvm.driver import tvmc @@ -24,7 +25,10 @@ from tvm.driver.tvmc.model import TVMCPackage, TVMCException -def test_tvmc_cl_compile_run_mlf(tflite_mobilenet_v1_1_quant, tmpdir_factory): +@pytest.mark.parametrize( + "target,pass_configs", [["llvm", []], ["c --executor=aot", ["tir.disable_vectorize=1"]]] +) +def test_tvmc_cl_compile_run_mlf(tflite_mobilenet_v1_1_quant, tmpdir_factory, target, pass_configs): pytest.importorskip("tflite") output_dir = tmpdir_factory.mktemp("mlf") @@ -32,10 +36,9 @@ def test_tvmc_cl_compile_run_mlf(tflite_mobilenet_v1_1_quant, tmpdir_factory): output_file = os.path.join(output_dir, "mock.tar") # Compile the input model and generate a Model Library Format (MLF) archive. - tvmc_cmd = ( - f"tvmc compile {input_model} --target='llvm' --output {output_file} --output-format mlf" - ) - tvmc_args = tvmc_cmd.split(" ")[1:] + pass_config_args = " ".join([f"--pass-config {pass_config}" for pass_config in pass_configs]) + tvmc_cmd = f"tvmc compile {input_model} --target='{target}' {pass_config_args} --output {output_file} --output-format mlf" + tvmc_args = shlex.split(tvmc_cmd)[1:] _main(tvmc_args) assert os.path.exists(output_file), "Could not find the exported MLF archive." @@ -82,9 +85,39 @@ def test_tvmc_export_package_mlf(tflite_mobilenet_v1_1_quant, tmpdir_factory): assert str(exp.value) == expected_reason, on_error -def test_tvmc_import_package_mlf(tflite_compiled_model_mlf): +def test_tvmc_import_package_mlf_graph(tflite_mobilenet_v1_1_quant, tflite_compile_model): + pytest.importorskip("tflite") + + tflite_compiled_model_mlf = tflite_compile_model( + tflite_mobilenet_v1_1_quant, output_format="mlf" + ) + + # Compile and export a model to a MLF archive so it can be imported. + exported_tvmc_package = tflite_compiled_model_mlf + archive_path = exported_tvmc_package.package_path + + # Import the MLF archive. TVMCPackage constructor will call import_package method. + tvmc_package = TVMCPackage(archive_path) + + assert tvmc_package.lib_name is None, ".lib_name must not be set in the MLF archive." + assert tvmc_package.lib_path is None, ".lib_path must not be set in the MLF archive." + assert ( + tvmc_package.graph is not None + ), ".graph must be set in the MLF archive for Graph executor." + assert tvmc_package.params is not None, ".params must be set in the MLF archive." + assert tvmc_package.type == "mlf", ".type must be set to 'mlf' in the MLF format." + + +def test_tvmc_import_package_mlf_aot(tflite_mobilenet_v1_1_quant, tflite_compile_model): pytest.importorskip("tflite") + tflite_compiled_model_mlf = tflite_compile_model( + tflite_mobilenet_v1_1_quant, + target="c --executor=aot", + output_format="mlf", + pass_context_configs=["tir.disable_vectorize=1"], + ) + # Compile and export a model to a MLF archive so it can be imported. exported_tvmc_package = tflite_compiled_model_mlf archive_path = exported_tvmc_package.package_path @@ -94,6 +127,6 @@ def test_tvmc_import_package_mlf(tflite_compiled_model_mlf): assert tvmc_package.lib_name is None, ".lib_name must not be set in the MLF archive." assert tvmc_package.lib_path is None, ".lib_path must not be set in the MLF archive." - assert tvmc_package.graph is not None, ".graph must be set in the MLF archive." + assert tvmc_package.graph is None, ".graph must not be set in the MLF archive for AOT executor." assert tvmc_package.params is not None, ".params must be set in the MLF archive." assert tvmc_package.type == "mlf", ".type must be set to 'mlf' in the MLF format." diff --git a/tests/python/driver/tvmc/test_runner.py b/tests/python/driver/tvmc/test_runner.py index 5277a790b43b..7acb376baba6 100644 --- a/tests/python/driver/tvmc/test_runner.py +++ b/tests/python/driver/tvmc/test_runner.py @@ -71,12 +71,15 @@ def test_get_top_results_keep_results(): assert len(sut[1]) == expected_number_of_results_per_line -def test_run_tflite_module__with_profile__valid_input(tflite_compiled_model, imagenet_cat): +def test_run_tflite_module__with_profile__valid_input( + tflite_mobilenet_v1_1_quant, tflite_compile_model, imagenet_cat +): # some CI environments wont offer TFLite, so skip in case it is not present pytest.importorskip("tflite") inputs = np.load(imagenet_cat) + tflite_compiled_model = tflite_compile_model(tflite_mobilenet_v1_1_quant) result = tvmc.run( tflite_compiled_model, inputs=inputs,