[TVMC] Add support for the MLF to 'compile' command (#8086)

* [TVMC] Add support for the MLF to 'compile' command

Add support for the Model Library Format (MLF) to 'tvmc' so users can
output compilation artifacts to a MLF archive passing the new flag
'--output-format mlf'. For instance:

$ python3 -m tvm.driver.tvmc compile ./sine_model.tflite --target="c" --output sine.tar --output-format mlf

will generate a sine.tar archive that is serialized accordingly to the
MLF.

Since the MLF is currently meant to be used only on micro targets, an
error is generated if one tries to run a MLF outside a micro context.

The micro context does not exist yet but will be later introduced as
part of the [RFC] "TVMC: Add support for µTVM".

That commit also adds 3 pytest tests to test tvmc + MLF.

Finally, it also fixes some missing periods in the 'compile' command
help sections and renames export_format to output_format so there is
no confusion with flag '--dump-code', which contains "formats to export"
in its help section.

Signed-off-by: Gustavo Romero <[email protected]>

* Fix missing importorskip in the import_package test

Fix missing importorskip() in the import_package test allowing the
test in question to be skipped when 'tflite' is not installed in the
test environment, otherwise the test will fail with:

[...]
>       archive_path = exported_tvmc_package.package_path
E       AttributeError: 'str' object has no attribute 'package_path'

python/tvm/driver/tvmc/compiler.py

@@ -40,40 +40,48 @@
 def add_compile_parser(subparsers):
     """ Include parser for 'compile' subcommand """
 
-    parser = subparsers.add_parser("compile", help="compile a model")
+    parser = subparsers.add_parser("compile", help="compile a model.")
     parser.set_defaults(func=drive_compile)
     parser.add_argument(
         "--cross-compiler",
         default="",
-        help="the cross compiler to generate target libraries, e.g. 'aarch64-linux-gnu-gcc'",
+        help="the cross compiler to generate target libraries, e.g. 'aarch64-linux-gnu-gcc'.",
     )
     parser.add_argument(
         "--cross-compiler-options",
         default="",
-        help="the cross compiler options to generate target libraries, e.g. '-mfpu=neon-vfpv4'",
+        help="the cross compiler options to generate target libraries, e.g. '-mfpu=neon-vfpv4'.",
     )
     parser.add_argument(
         "--desired-layout",
         choices=["NCHW", "NHWC"],
         default=None,
-        help="change the data layout of the whole graph",
+        help="change the data layout of the whole graph.",
     )
     parser.add_argument(
         "--dump-code",
         metavar="FORMAT",
         default="",
-        help="comma separarated list of formats to export, e.g. 'asm,ll,relay' ",
+        help="comma separated list of formats to export the input model, e.g. 'asm,ll,relay'.",
     )
     parser.add_argument(
         "--model-format",
         choices=frontends.get_frontend_names(),
-        help="specify input model format",
+        help="specify input model format.",
     )
     parser.add_argument(
         "-o",
         "--output",
         default="module.tar",
-        help="output the compiled module to an archive",
+        help="output the compiled module to a specifed archive. Defaults to 'module.tar'.",
+    )
+    parser.add_argument(
+        "-f",
+        "--output-format",
+        choices=["so", "mlf"],
+        default="so",
+        help="output format. Use 'so' for shared object or 'mlf' for Model Library Format "
+        "(only for µTVM targets). Defaults to 'so'.",
     )
     parser.add_argument(
         "--target",
@@ -85,23 +93,23 @@ def add_compile_parser(subparsers):
         metavar="PATH",
         default="",
         help="path to an auto-tuning log file by AutoTVM. If not presented, "
-        "the fallback/tophub configs will be used",
+        "the fallback/tophub configs will be used.",
     )
-    parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity")
+    parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity.")
     # TODO (@leandron) This is a path to a physical file, but
     #     can be improved in future to add integration with a modelzoo
     #     or URL, for example.
-    parser.add_argument("FILE", help="path to the input model file")
+    parser.add_argument("FILE", help="path to the input model file.")
     parser.add_argument(
         "--input-shapes",
         help="specify non-generic shapes for model to run, format is "
-        '"input_name:[dim1,dim2,...,dimn] input_name2:[dim1,dim2]"',
+        '"input_name:[dim1,dim2,...,dimn] input_name2:[dim1,dim2]".',
         type=common.parse_shape_string,
         default=None,
     )
     parser.add_argument(
         "--disabled-pass",
-        help="disable specific passes, comma-separated list of pass names",
+        help="disable specific passes, comma-separated list of pass names.",
         type=common.parse_pass_list_str,
         default="",
     )
@@ -132,6 +140,7 @@ def drive_compile(args):
         package_path=args.output,
         cross=args.cross_compiler,
         cross_options=args.cross_compiler_options,
+        output_format=args.output_format,
         dump_code=dump_code,
         target_host=None,
         desired_layout=args.desired_layout,
@@ -148,7 +157,7 @@ def compile_model(
     package_path: Optional[str] = None,
     cross: Optional[Union[str, Callable]] = None,
     cross_options: Optional[str] = None,
-    export_format: str = "so",
+    output_format: str = "so",
     dump_code: Optional[List[str]] = None,
     target_host: Optional[str] = None,
     desired_layout: Optional[str] = None,
@@ -177,7 +186,7 @@ def compile_model(
         Function that performs the actual compilation
     cross_options : str, optional
         Command line options to be passed to the cross compiler.
-    export_format : str
+    output_format : str
         What format to use when saving the function library. Must be one of "so" or "tar".
         When compiling for a remote device without a cross compiler, "tar" will likely work better.
     dump_code : list, optional
@@ -262,7 +271,11 @@ def compile_model(
 
     # Create a new tvmc model package object from the graph definition.
     package_path = tvmc_model.export_package(
-        graph_module, package_path, cross, cross_options, export_format
+        graph_module,
+        package_path,
+        cross,
+        cross_options,
+        output_format,
     )
 
     # Write dumps to file.

python/tvm/driver/tvmc/model.py

@@ -53,6 +53,7 @@
 from tvm import relay
 from tvm.contrib import utils
 from tvm.relay.backend.executor_factory import GraphExecutorFactoryModule
+from tvm.micro import export_model_library_format
 
 from .common import TVMCException
 
@@ -175,7 +176,7 @@ def default_package_path(self):
         """
         return self._tmp_dir.relpath("model_package.tar")
 
-    def export_package(
+    def export_classic_format(
         self,
         executor_factory: GraphExecutorFactoryModule,
         package_path: Optional[str] = None,
@@ -203,8 +204,6 @@ def export_package(
         package_path : str
             The path that the package was saved to.
         """
-        if lib_format not in ["so", "tar"]:
-            raise TVMCException("Only .so and .tar export formats are supported.")
         lib_name = "mod." + lib_format
         graph_name = "mod.json"
         param_name = "mod.params"
@@ -241,6 +240,50 @@ def export_package(
 
         return package_path
 
+    def export_package(
+        self,
+        executor_factory: GraphExecutorFactoryModule,
+        package_path: Optional[str] = None,
+        cross: Optional[Union[str, Callable]] = None,
+        cross_options: Optional[str] = None,
+        output_format: str = "so",
+    ):
+        """Save this TVMCModel to file.
+        Parameters
+        ----------
+        executor_factory : GraphExecutorFactoryModule
+            The factory containing compiled the compiled artifacts needed to run this model.
+        package_path : str, None
+            Where the model should be saved. Note that it will be packaged as a .tar file.
+            If not provided, the package will be saved to a generically named file in tmp.
+        cross : str or callable object, optional
+            Function that performs the actual compilation.
+        cross_options : str, optional
+            Command line options to be passed to the cross compiler.
+        output_format : str
+            How to save the modules function library. Must be one of "so" and "tar" to save
+            using the classic format or "mlf" to save using the Model Library Format.
+
+        Returns
+        -------
+        package_path : str
+            The path that the package was saved to.
+        """
+        if output_format not in ["so", "tar", "mlf"]:
+            raise TVMCException("Only 'so', 'tar', and 'mlf' output formats are supported.")
+
+        if output_format == "mlf" and cross:
+            raise TVMCException("Specifying the MLF output and a cross compiler is not supported.")
+
+        if output_format in ["so", "tar"]:
+            package_path = self.export_classic_format(
+                executor_factory, package_path, cross, cross_options, output_format
+            )
+        elif output_format == "mlf":
+            package_path = export_model_library_format(executor_factory, package_path)
+
+        return package_path
+
     def summary(self, file: TextIO = None):
         """Print the IR corressponding to this model.
 
@@ -274,25 +317,41 @@ def import_package(self, package_path: str):
         package_path : str
             The path to the saved TVMCPackage.
         """
-        lib_name_so = "mod.so"
-        lib_name_tar = "mod.tar"
-        graph_name = "mod.json"
-        param_name = "mod.params"
-
         temp = self._tmp_dir
         t = tarfile.open(package_path)
         t.extractall(temp.relpath("."))
 
-        with open(temp.relpath(param_name), "rb") as param_file:
-            self.params = bytearray(param_file.read())
-        self.graph = open(temp.relpath(graph_name)).read()
-        if os.path.exists(temp.relpath(lib_name_so)):
-            self.lib_name = lib_name_so
-        elif os.path.exists(temp.relpath(lib_name_tar)):
-            self.lib_name = lib_name_tar
+        if os.path.exists(temp.relpath("metadata.json")):
+            # Model Library Format (MLF)
+            self.lib_name = None
+            self.lib_path = None
+
+            graph = temp.relpath("runtime-config/graph/graph.json")
+            params = temp.relpath("parameters/default.params")
+
+            self.type = "mlf"
         else:
-            raise TVMCException("Couldn't find exported library in the package.")
-        self.lib_path = temp.relpath(self.lib_name)
+            # Classic format
+            lib_name_so = "mod.so"
+            lib_name_tar = "mod.tar"
+            if os.path.exists(temp.relpath(lib_name_so)):
+                self.lib_name = lib_name_so
+            elif os.path.exists(temp.relpath(lib_name_tar)):
+                self.lib_name = lib_name_tar
+            else:
+                raise TVMCException("Couldn't find exported library in the package.")
+            self.lib_path = temp.relpath(self.lib_name)
+
+            graph = temp.relpath("mod.json")
+            params = temp.relpath("mod.params")
+
+            self.type = "classic"
+
+        with open(params, "rb") as param_file:
+            self.params = bytearray(param_file.read())
+
+        with open(graph) as graph_file:
+            self.graph = graph_file.read()
 
 
 class TVMCResult(object):

python/tvm/driver/tvmc/runner.py

@@ -359,6 +359,14 @@ def run_module(
             "Try calling tvmc.compile on the model before running it."
         )
 
+    # Currently only two package formats are supported: "classic" and
+    # "mlf". The later can only be used for micro targets, i.e. with µTVM.
+    if tvmc_package.type == "mlf":
+        raise TVMCException(
+            "You're trying to run a model saved using the Model Library Format (MLF)."
+            "MLF can only be used to run micro targets (µTVM)."
+        )
+
     if hostname:
         if isinstance(port, str):
             port = int(port)

python/tvm/micro/model_library_format.py

@@ -216,6 +216,11 @@ def export_model_library_format(mod: executor_factory.ExecutorFactoryModule, fil
         The return value of tvm.relay.build, which will be exported into Model Library Format.
     file_name : str
         Path to the .tar archive to generate.
+
+    Returns
+    -------
+    file_name : str
+        The path to the generated .tar archive.
     """
     tempdir = utils.tempdir()
     is_aot = isinstance(mod, executor_factory.AOTExecutorFactoryModule)
@@ -260,3 +265,5 @@ def reset(tarinfo):
             return tarinfo
 
         tar_f.add(tempdir.temp_dir, arcname=".", filter=reset)
+
+    return file_name

tests/python/driver/tvmc/conftest.py

@@ -41,7 +41,7 @@ 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):
+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"
@@ -53,7 +53,10 @@ def get_sample_compiled_module(target_dir, package_filename):
 
     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)
+        tvmc_model,
+        target="llvm",
+        package_path=os.path.join(target_dir, package_filename),
+        output_format=output_format,
     )
 
 
@@ -182,6 +185,24 @@ def tflite_compiled_model(tmpdir_factory):
     return get_sample_compiled_module(target_dir, "mock.tar")
 
 
+@pytest.fixture(scope="session")
+def tflite_compiled_model_mlf(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", "mlf")
+
+
 @pytest.fixture(scope="session")
 def imagenet_cat(tmpdir_factory):
     tmpdir_name = tmpdir_factory.mktemp("data")