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

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")