[TVMC] Add --config argument for config files (apache#11012)

* [TVMC] Add `--config` argument for config files

Collecting common configurations for users of TVM and exposing them gracefully in tvmc using a `--config` option
as defined in https://github.com/apache/tvm-rfcs/blob/main/rfcs/0030-tvmc-comand-line-configuration-files.md

Co-authored-by: Shai Maor <[email protected]>

* Add correct test guards

Co-authored-by: Shai Maor <[email protected]>

configs/host/default.json

@@ -0,0 +1,7 @@
+{
+  "targets": [
+    {
+      "kind": "llvm"
+    }
+  ]
+}

configs/test/compile_config_test.json

@@ -0,0 +1,9 @@
+{
+    "targets": [
+        {"kind": "cmsis-nn", "from_device": "1"},
+	{"kind": "c", "mcpu": "cortex-m55"}
+    ],
+    "executor": { "kind": "aot"},
+    "runtime": { "kind": "crt"},
+    "pass-config": { "tir.disable_vectorize": "1"}
+}

configs/test/tune_config_test.json

@@ -0,0 +1,6 @@
+{
+    "targets": [
+        { "kind": "llvm" }
+    ],
+    "trials": "2"
+}

python/tvm/driver/tvmc/autotuner.py

@@ -47,7 +47,7 @@
 
 
 @register_parser
-def add_tune_parser(subparsers, _):
+def add_tune_parser(subparsers, _, json_params):
     """Include parser for 'tune' subcommand"""
 
     parser = subparsers.add_parser("tune", help="auto-tune a model")
@@ -224,6 +224,9 @@ def add_tune_parser(subparsers, _):
         type=parse_shape_string,
     )
 
+    for one_entry in json_params:
+        parser.set_defaults(**one_entry)
+
 
 def drive_tune(args):
     """Invoke auto-tuning with command line arguments

python/tvm/driver/tvmc/compiler.py

@@ -43,7 +43,7 @@
 
 
 @register_parser
-def add_compile_parser(subparsers, _):
+def add_compile_parser(subparsers, _, json_params):
     """Include parser for 'compile' subcommand"""
 
     parser = subparsers.add_parser("compile", help="compile a model.")
@@ -143,6 +143,9 @@ def add_compile_parser(subparsers, _):
         help="The output module name. Defaults to 'default'.",
     )
 
+    for one_entry in json_params:
+        parser.set_defaults(**one_entry)
+
 
 def drive_compile(args):
     """Invoke tvmc.compiler module with command line arguments

python/tvm/driver/tvmc/config_options.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+manipulate json config file to work with TVMC
+"""
+import os
+import json
+from tvm.driver.tvmc import TVMCException
+
+
+def find_json_file(name, path):
+    """search for json file given file name a path
+
+    Parameters
+    ----------
+    name: string
+        the file name need to be searched
+    path: string
+        path to search at
+
+    Returns
+    -------
+    string
+        the full path to that file
+
+    """
+    match = ""
+    for root, _dirs, files in os.walk(path):
+        if name in files:
+            match = os.path.join(root, name)
+            break
+
+    return match
+
+
+def read_and_convert_json_into_dict(config_args):
+    """Read json configuration file and return a dictionary with all parameters
+
+    Parameters
+    ----------
+    args: argparse.Namespace
+        Arguments from command line parser holding the json file path.
+
+    Returns
+    -------
+    dictionary
+        dictionary with all the json arguments keys and values
+
+    """
+    try:
+        if ".json" not in config_args.config:
+            config_args.config = config_args.config.strip() + ".json"
+        if os.path.isfile(config_args.config):
+            json_config_file = config_args.config
+        else:
+            config_dir = os.path.abspath(
+                os.path.join(os.path.realpath(__file__), "..", "..", "..", "..", "..", "configs")
+            )
+            json_config_file = find_json_file(config_args.config, config_dir)
+        return json.load(open(json_config_file, "rb"))
+
+    except FileNotFoundError:
+        raise TVMCException(
+            f"File {config_args.config} does not exist at {config_dir} or is wrong format."
+        )
+
+
+def parse_target_from_json(one_target, command_line_list):
+    """parse the targets out of the json file struct
+
+    Parameters
+    ----------
+    one_target: dict
+        dictionary with all target's details
+    command_line_list: list
+        list to update with target parameters
+    """
+    target_kind, *sub_type = [
+        one_target[key] if key == "kind" else (key, one_target[key]) for key in one_target
+    ]
+
+    internal_dict = {}
+    if sub_type:
+        sub_target_type = sub_type[0][0]
+        target_value = sub_type[0][1]
+        internal_dict[f"target_{target_kind}_{sub_target_type}"] = target_value
+        command_line_list.append(internal_dict)
+
+    return target_kind
+
+
+def convert_config_json_to_cli(json_params):
+    """convert all configuration keys & values from dictionary to cli format
+
+    Parameters
+    ----------
+    args: dictionary
+        dictionary with all configuration keys & values.
+
+    Returns
+    -------
+    int
+        list of configuration values in cli format
+
+    """
+    command_line_list = []
+    for param_key in json_params:
+        if param_key == "targets":
+            target_list = [
+                parse_target_from_json(one_target, command_line_list)
+                for one_target in json_params[param_key]
+            ]
+
+            internal_dict = {}
+            internal_dict["target"] = ", ".join(map(str, target_list))
+            command_line_list.append(internal_dict)
+
+        elif param_key in ("executor", "runtime"):
+            for key, value in json_params[param_key].items():
+                if key == "kind":
+                    kind = f"{value}_"
+                    new_dict_key = param_key
+                else:
+                    new_dict_key = f"{param_key}_{kind}{key}"
+
+                internal_dict = {}
+                internal_dict[new_dict_key.replace("-", "_")] = value
+                command_line_list.append(internal_dict)
+
+        elif isinstance(json_params[param_key], dict):
+            internal_dict = {}
+            modify_param_key = param_key.replace("-", "_")
+            internal_dict[modify_param_key] = []
+            for key, value in json_params[param_key].items():
+                internal_dict[modify_param_key].append(f"{key}={value}")
+            command_line_list.append(internal_dict)
+
+        else:
+            internal_dict = {}
+            internal_dict[param_key.replace("-", "_")] = json_params[param_key]
+            command_line_list.append(internal_dict)
+
+    return command_line_list

python/tvm/driver/tvmc/main.py

@@ -26,7 +26,10 @@
 import tvm
 
 from tvm.driver.tvmc import TVMCException, TVMCImportError
-
+from tvm.driver.tvmc.config_options import (
+    read_and_convert_json_into_dict,
+    convert_config_json_to_cli,
+)
 
 REGISTERED_PARSER = []
 
@@ -64,12 +67,19 @@ def _main(argv):
         # so it doesn't interfere with the creation of the dynamic subparsers.
         add_help=False,
     )
+
+    parser.add_argument("--config", default="default", help="configuration json file")
+    config_arg, argv = parser.parse_known_args(argv)
+
+    json_param_dict = read_and_convert_json_into_dict(config_arg)
+    json_config_values = convert_config_json_to_cli(json_param_dict)
+
     parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity")
     parser.add_argument("--version", action="store_true", help="print the version and exit")
 
     subparser = parser.add_subparsers(title="commands")
     for make_subparser in REGISTERED_PARSER:
-        make_subparser(subparser, parser)
+        make_subparser(subparser, parser, json_config_values)
 
     # Finally, add help for the main parser.
     parser.add_argument("-h", "--help", action="help", help="show this help message and exit.")

python/tvm/driver/tvmc/micro.py

@@ -45,7 +45,7 @@
 
 
 @register_parser
-def add_micro_parser(subparsers, main_parser):
+def add_micro_parser(subparsers, main_parser, json_params):
     """Includes parser for 'micro' context and associated subcommands:
     create-project (create), build, and flash.
     """
@@ -231,6 +231,9 @@ def _add_parser(parser, platform):
         help="show this help message which includes platform-specific options and exit.",
     )
 
+    for one_entry in json_params:
+        micro.set_defaults(**one_entry)
+
 
 def drive_micro(args):
     # Call proper handler based on subcommand parsed.

python/tvm/driver/tvmc/runner.py

@@ -60,7 +60,7 @@
 
 
 @register_parser
-def add_run_parser(subparsers, main_parser):
+def add_run_parser(subparsers, main_parser, json_params):
     """Include parser for 'run' subcommand"""
 
     # Use conflict_handler='resolve' to allow '--list-options' option to be properly overriden when
@@ -191,6 +191,9 @@ def add_run_parser(subparsers, main_parser):
         help="show this help message with platform-specific options and exit.",
     )
 
+    for one_entry in json_params:
+        parser.set_defaults(**one_entry)
+
 
 def drive_run(args):
     """Invoke runner module with command line arguments

python/tvm/driver/tvmc/target.py

@@ -81,7 +81,7 @@ def generate_target_args(parser):
     parser.add_argument(
         "--target",
         help="compilation target as plain string, inline JSON or path to a JSON file",
-        required=True,
+        required=False,
     )
     for target_kind in _valid_target_kinds():
         _generate_target_kind_args(parser, target_kind)

tests/python/driver/tvmc/test_command_line.py

@@ -56,3 +56,39 @@ def test_tvmc_cl_workflow(keras_simple, tmpdir_factory):
     run_args = run_str.split(" ")[1:]
     _main(run_args)
     assert os.path.exists(output_path)
+
+
+@pytest.mark.skipif(
+    platform.machine() == "aarch64",
+    reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
+)
+def test_tvmc_cl_workflow_json_config(keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
+    tune_config_file = "tune_config_test"
+    tmpdir = tmpdir_factory.mktemp("data")
+
+    # Test model tuning
+    log_path = os.path.join(tmpdir, "keras-autotuner_records.json")
+    tuning_str = (
+        f"tvmc tune --config {tune_config_file} --output {log_path} "
+        f"--enable-autoscheduler {keras_simple}"
+    )
+    tuning_args = tuning_str.split(" ")[1:]
+    _main(tuning_args)
+    assert os.path.exists(log_path)
+
+    # Test model compilation
+    package_path = os.path.join(tmpdir, "keras-tvm.tar")
+    compile_str = (
+        f"tvmc compile --tuning-records {log_path} " f"--output {package_path} {keras_simple}"
+    )
+    compile_args = compile_str.split(" ")[1:]
+    _main(compile_args)
+    assert os.path.exists(package_path)
+
+    # Test running the model
+    output_path = os.path.join(tmpdir, "predictions.npz")
+    run_str = f"tvmc run --outputs {output_path} {package_path}"
+    run_args = run_str.split(" ")[1:]
+    _main(run_args)
+    assert os.path.exists(output_path)