[TVMC] Compose target options from target registry (apache#9218)

* [TVMC] Split common tvmc test file into more specific files

The `test_tvmc_common.py` file was becoming a bit of a mixed bag of
tests and as we now want to extend the `Target` processing logic it made
sense to split each out into its own file to make it clearer what each
does.

`test_common.py` has also been renamed before we start using it for all the
tests instead.

* [TVMC] Compose target options from target registry

[The RFC for this is still under discussion](https://github.com/apache/tvm-rfcs/pulls), but doing this before splitting the registries makes the most sense. This enables the `tvmc` driver to re-combobulate Target options from arguments:

```
tvmc --target=llvm \
    --target-llvm-mcpu=cortex-m3
```

include/tvm/target/target_kind.h

@@ -229,6 +229,12 @@ class TargetKindRegEntry {
    * \return The entry names.
    */
   TVM_DLL static Array<String> ListTargetKinds();
+  /*!
+   * \brief Get all supported option names and types for a given Target kind.
+   * \return Map of option name to type
+   */
+  TVM_DLL static Map<String, String> ListTargetKindOptions(const TargetKind& kind);
+
   /*!
    * \brief Register or get a new entry.
    * \param target_kind_name The name of the TargetKind.

python/tvm/driver/tvmc/autotuner.py

@@ -21,7 +21,7 @@
 import logging
 import time
 from copy import deepcopy
-from typing import Optional, Dict, List, Union
+from typing import Any, Optional, Dict, List, Union
 
 from urllib.parse import urlparse
 
@@ -38,6 +38,7 @@
 from .common import TVMCException
 from .main import register_parser
 from .model import TVMCModel
+from .target import generate_target_args, reconstruct_target_args
 
 
 # pylint: disable=invalid-name
@@ -106,16 +107,14 @@ def add_tune_parser(subparsers):
         help="hostname (required) and port (optional, defaults to 9090) of the RPC tracker, "
         "e.g. '192.168.0.100:9999'",
     )
-    parser.add_argument(
-        "--target",
-        help="compilation target as plain string, inline JSON or path to a JSON file",
-        required=True,
-    )
+
+    generate_target_args(parser)
     parser.add_argument(
         "--target-host",
         help="the host compilation target, defaults to 'llvm'",
         default="llvm",
     )
+
     parser.add_argument("--timeout", type=int, default=10, help="compilation timeout, in seconds")
     parser.add_argument(
         "--trials",
@@ -286,6 +285,7 @@ def drive_tune(args):
         hardware_params=hardware_params,
         include_simple_tasks=args.include_simple_tasks,
         log_estimated_latency=args.log_estimated_latency,
+        additional_target_options=reconstruct_target_args(args),
     )
 
 
@@ -311,6 +311,7 @@ def tune_model(
     hardware_params: Optional[HardwareParams] = None,
     include_simple_tasks: bool = False,
     log_estimated_latency: bool = False,
+    additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None,
 ):
     """Use tuning to automatically optimize the functions in a model.
 
@@ -367,13 +368,15 @@ def tune_model(
         the autoscheduler.
     log_estimated_latency : bool, optional
         If using the autoscheduler, write the estimated latency at each step of tuning to file.
+    additional_target_options: Optional[Dict[str, Dict[str, Any]]]
+        Additional target options in a dictionary to combine with initial Target arguments
 
     Returns
     -------
     tuning_records : str
         The path to the produced tuning log file.
     """
-    target, extra_targets = common.target_from_cli(target)
+    target, extra_targets = common.target_from_cli(target, additional_target_options)
     target, target_host = Target.check_and_update_host_consist(target, target_host)
     # TODO(jwfromm) Remove this deepcopy once AlterOpLayout bug that mutates source
     # model is fixed. For now, creating a clone avoids the issue.

python/tvm/driver/tvmc/common.py

@@ -80,7 +80,7 @@ def convert_graph_layout(mod, desired_layout):
             )
 
 
-def validate_targets(parse_targets):
+def validate_targets(parse_targets, additional_target_options=None):
     """
     Apply a series of validations in the targets provided via CLI.
     """
@@ -104,6 +104,15 @@ def validate_targets(parse_targets):
             f"Found: {verbose_tvm_targets}."
         )
 
+    if additional_target_options is not None:
+        for target_name in additional_target_options:
+            if not any([target for target in parse_targets if target["name"] == target_name]):
+                first_option = list(additional_target_options[target_name].keys())[0]
+                raise TVMCException(
+                    f"Passed --target-{target_name}-{first_option}"
+                    f" but did not specify {target_name} target"
+                )
+
 
 def tokenize_target(target):
     """
@@ -261,7 +270,21 @@ def is_inline_json(target):
         return False
 
 
-def target_from_cli(target):
+def _combine_target_options(target, additional_target_options=None):
+    if additional_target_options is None:
+        return target
+    if target["name"] in additional_target_options:
+        target["opts"].update(additional_target_options[target["name"]])
+    return target
+
+
+def _recombobulate_target(target):
+    name = target["name"]
+    opts = " ".join([f"-{key}={value}" for key, value in target["opts"].items()])
+    return f"{name} {opts}"
+
+
+def target_from_cli(target, additional_target_options=None):
     """
     Create a tvm.target.Target instance from a
     command line interface (CLI) string.
@@ -272,6 +295,10 @@ def target_from_cli(target):
         compilation target as plain string,
         inline JSON or path to a JSON file
 
+    additional_target_options: Optional[Dict[str, Dict[str,str]]]
+        dictionary of additional target options to be
+        combined with parsed targets
+
     Returns
     -------
     tvm.target.Target
@@ -298,18 +325,22 @@ def target_from_cli(target):
         except ValueError as ex:
             raise TVMCException(f"Error parsing target string '{target}'.\nThe error was: {ex}")
 
-        validate_targets(parsed_targets)
-        tvm_targets = [t for t in parsed_targets if t["is_tvm_target"]]
+        validate_targets(parsed_targets, additional_target_options)
+        tvm_targets = [
+            _combine_target_options(t, additional_target_options)
+            for t in parsed_targets
+            if t["is_tvm_target"]
+        ]
 
         # Validated target strings have 1 or 2 tvm targets, otherwise
         # `validate_targets` above will fail.
         if len(tvm_targets) == 1:
-            target = tvm_targets[0]["raw"]
+            target = _recombobulate_target(tvm_targets[0])
             target_host = None
         else:
             assert len(tvm_targets) == 2
-            target = tvm_targets[0]["raw"]
-            target_host = tvm_targets[1]["raw"]
+            target = _recombobulate_target(tvm_targets[0])
+            target_host = _recombobulate_target(tvm_targets[1])
 
         extra_targets = [t for t in parsed_targets if not t["is_tvm_target"]]
 

python/tvm/driver/tvmc/compiler.py

@@ -19,7 +19,7 @@
 """
 import logging
 import os.path
-from typing import Optional, Dict, List, Union, Callable
+from typing import Any, Optional, Dict, List, Union, Callable
 from pathlib import Path
 
 import tvm
@@ -30,6 +30,7 @@
 from . import common, composite_target, frontends
 from .model import TVMCModel, TVMCPackage
 from .main import register_parser
+from .target import generate_target_args, reconstruct_target_args
 
 
 # pylint: disable=invalid-name
@@ -91,11 +92,7 @@ def add_compile_parser(subparsers):
         "times, each one to set one configuration value, "
         "e.g. '--pass-config relay.backend.use_auto_scheduler=0'.",
     )
-    parser.add_argument(
-        "--target",
-        help="compilation targets as comma separated string, inline JSON or path to a JSON file.",
-        required=True,
-    )
+    generate_target_args(parser)
     parser.add_argument(
         "--tuning-records",
         metavar="PATH",
@@ -154,6 +151,7 @@ def drive_compile(args):
         desired_layout=args.desired_layout,
         disabled_pass=args.disabled_pass,
         pass_context_configs=args.pass_config,
+        additional_target_options=reconstruct_target_args(args),
     )
 
     return 0
@@ -172,6 +170,7 @@ def compile_model(
     desired_layout: Optional[str] = None,
     disabled_pass: Optional[str] = None,
     pass_context_configs: Optional[List[str]] = None,
+    additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None,
 ):
     """Compile a model from a supported framework into a TVM module.
 
@@ -215,6 +214,8 @@ def compile_model(
     pass_context_configs: list[str], optional
         List of strings containing a set of configurations to be passed to the
         PassContext.
+    additional_target_options: Optional[Dict[str, Dict[str, Any]]]
+        Additional target options in a dictionary to combine with initial Target arguments
 
 
     Returns
@@ -230,7 +231,7 @@ def compile_model(
     if desired_layout:
         mod = common.convert_graph_layout(mod, desired_layout)
 
-    tvm_target, extra_targets = common.target_from_cli(target)
+    tvm_target, extra_targets = common.target_from_cli(target, additional_target_options)
     tvm_target, target_host = Target.check_and_update_host_consist(tvm_target, target_host)
 
     for codegen_from_cli in extra_targets:

python/tvm/driver/tvmc/target.py

@@ -0,0 +1,74 @@
+# 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.
+"""
+This file contains functions for processing target inputs for the TVMC CLI
+"""
+
+from tvm.target import Target
+
+# We can't tell the type inside an Array but all current options are strings so
+# it can default to that. Bool is used alongside Integer but aren't distinguished
+# between as both are represented by IntImm
+INTERNAL_TO_NATIVE_TYPE = {"runtime.String": str, "IntImm": int, "Array": str}
+INTERNAL_TO_HELP = {"runtime.String": " string", "IntImm": "", "Array": " options"}
+
+
+def _generate_target_kind_args(parser, kind):
+    target_group = parser.add_argument_group(f"target {kind.name}")
+    for target_option, target_type in kind.options.items():
+        if target_type in INTERNAL_TO_NATIVE_TYPE:
+            target_group.add_argument(
+                f"--target-{kind.name}-{target_option}",
+                type=INTERNAL_TO_NATIVE_TYPE[target_type],
+                help=f"target {kind.name} {target_option}{INTERNAL_TO_HELP[target_type]}",
+            )
+
+
+def generate_target_args(parser):
+    """Walks through the TargetKind registry and generates arguments for each Target's options"""
+    parser.add_argument(
+        "--target",
+        help="compilation target as plain string, inline JSON or path to a JSON file",
+        required=True,
+    )
+    target_kinds = Target.list_kinds()
+    for target_kind in target_kinds:
+        target = Target(target_kind)
+        _generate_target_kind_args(parser, target.kind)
+
+
+def _reconstruct_target_kind_args(args, kind):
+    kind_options = {}
+    for target_option, target_type in kind.options.items():
+        if target_type in INTERNAL_TO_NATIVE_TYPE:
+            var_name = f"target_{kind.name}_{target_option.replace('-', '_')}"
+            option_value = getattr(args, var_name)
+            if option_value is not None:
+                kind_options[target_option] = getattr(args, var_name)
+    return kind_options
+
+
+def reconstruct_target_args(args):
+    """Reconstructs the target options from the arguments"""
+    target_kinds = Target.list_kinds()
+    reconstructed = {}
+    for target_kind in target_kinds:
+        target = Target(target_kind)
+        kind_options = _reconstruct_target_kind_args(args, target.kind)
+        if kind_options:
+            reconstructed[target.kind.name] = kind_options
+    return reconstructed

python/tvm/target/target.py

@@ -31,6 +31,11 @@
 class TargetKind(Object):
     """Kind of a compilation target"""
 
+    @property
+    def options(self):
+        """Returns the dict of available option names and types"""
+        return dict(_ffi_api.ListTargetKindOptions(self))
+
 
 @tvm._ffi.register_object
 class Target(Object):

src/target/target_kind.cc

@@ -49,6 +49,14 @@ Array<String> TargetKindRegEntry::ListTargetKinds() {
   return TargetKindRegistry::Global()->ListAllNames();
 }
 
+Map<String, String> TargetKindRegEntry::ListTargetKindOptions(const TargetKind& target_kind) {
+  Map<String, String> options;
+  for (const auto& kv : target_kind->key2vtype_) {
+    options.Set(kv.first, kv.second.type_key);
+  }
+  return options;
+}
+
 TargetKindRegEntry& TargetKindRegEntry::RegisterOrGet(const String& target_kind_name) {
   return TargetKindRegistry::Global()->RegisterOrGet(target_kind_name);
 }
@@ -359,5 +367,7 @@ TVM_REGISTER_TARGET_KIND("composite", kDLCPU).add_attr_option<Array<Target>>("de
 /**********  Registry  **********/
 
 TVM_REGISTER_GLOBAL("target.ListTargetKinds").set_body_typed(TargetKindRegEntry::ListTargetKinds);
+TVM_REGISTER_GLOBAL("target.ListTargetKindOptions")
+    .set_body_typed(TargetKindRegEntry::ListTargetKindOptions);
 
 }  // namespace tvm

tests/cpp/target_test.cc

@@ -152,8 +152,18 @@ TEST(TargetCreation, DeduplicateKeys) {
   ICHECK_EQ(target->GetAttr<Bool>("link-params"), false);
 }
 
-TEST(TargetKindRegistryListTargetKinds, Basic) {
+TEST(TargetKindRegistry, ListTargetKinds) {
   Array<String> names = TargetKindRegEntry::ListTargetKinds();
   ICHECK_EQ(names.empty(), false);
   ICHECK_EQ(std::count(std::begin(names), std::end(names), "llvm"), 1);
 }
+
+TEST(TargetKindRegistry, ListTargetOptions) {
+  TargetKind llvm = TargetKind::Get("llvm").value();
+  Map<String, String> attrs = TargetKindRegEntry::ListTargetKindOptions(llvm);
+  ICHECK_EQ(attrs.empty(), false);
+
+  ICHECK_EQ(attrs["mattr"], "Array");
+  ICHECK_EQ(attrs["mcpu"], "runtime.String");
+  ICHECK_EQ(attrs["system-lib"], "IntImm");
+}

tests/python/driver/tvmc/test_compiler.py

@@ -397,7 +397,7 @@ def test_compile_tflite_module_with_external_codegen_cmsisnn(
 
     tvmc_package = tvmc.compiler.compile_model(
         tvmc_model,
-        target=f"cmsis-nn, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 --executor=aot",
+        target=f"cmsis-nn, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 -executor=aot",
         output_format="mlf",
         package_path=output_file_name,
         pass_context_configs=["tir.disable_vectorize=true"],
@@ -455,7 +455,7 @@ def test_compile_tflite_module_with_external_codegen_ethosu(
 
         tvmc_package = tvmc.compiler.compile_model(
             tvmc_model,
-            target=f"ethos-u -accelerator_config={accel_type}, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 --executor=aot",
+            target=f"ethos-u -accelerator_config={accel_type}, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 -executor=aot",
             output_format="mlf",
             package_path=output_file_name,
             pass_context_configs=["tir.disable_vectorize=true"],

tests/python/driver/tvmc/test_mlf.py

@@ -27,7 +27,7 @@
 
 
 @pytest.mark.parametrize(
-    "target,pass_configs", [["llvm", []], ["c --executor=aot", ["tir.disable_vectorize=1"]]]
+    "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")
@@ -114,7 +114,7 @@ def test_tvmc_import_package_mlf_aot(tflite_mobilenet_v1_1_quant, tflite_compile
 
     tflite_compiled_model_mlf = tflite_compile_model(
         tflite_mobilenet_v1_1_quant,
-        target="c --executor=aot",
+        target="c -executor=aot",
         output_format="mlf",
         pass_context_configs=["tir.disable_vectorize=1"],
     )