[TVMC] Add composite target passes for compilation and tuning (apache…

…#7304)

* Extend --target syntax to cover multiple targets for compilation and tuning
 * Add a new composite_target module to implement custom codegen passes into TVMC
 * Provide implementation to integrate TVMC, to target Arm Ethos-N NPU and
   Compute Library for the Arm Architecture (ACL)

Change-Id: Iaee53fe22f0c14eb4e4c8ec47e72bade0c5e32cc

python/tvm/driver/tvmc/autotuner.py

@@ -29,7 +29,7 @@
 from tvm.autotvm.tuner import RandomTuner
 from tvm.autotvm.tuner import XGBTuner
 
-from . import common, frontends
+from . import common, composite_target, frontends
 from .common import TVMCException
 from .main import register_parser
 
@@ -241,9 +241,14 @@ def drive_tune(args):
                 "need to provide an RPC tracker key (--rpc-key) for remote tuning"
             )
 
-    target = common.target_from_cli(args.target)
+    target, extra_targets = common.target_from_cli(args.target)
     mod, params = frontends.load_model(args.FILE, args.model_format, shape_dict=args.input_shapes)
 
+    for codegen_from_cli in extra_targets:
+        codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"])
+        partition_function = codegen["pass_pipeline"]
+        mod = partition_function(mod, params)
+
     # min_repeat_ms should be:
     # a. the value provided by the user, if any, or
     # b. 0ms in case target is "cpu"; otherwise 1000ms

python/tvm/driver/tvmc/common.py

@@ -18,6 +18,7 @@
 Common utility functions shared by TVMC modules.
 """
 import re
+import json
 import logging
 import os.path
 import argparse
@@ -78,6 +79,168 @@ def convert_graph_layout(mod, desired_layout):
             )
 
 
+def validate_targets(parse_targets):
+    """
+    Apply a series of validations in the targets provided via CLI.
+    """
+    tvm_target_kinds = tvm.target.Target.list_kinds()
+    targets = [t["name"] for t in parse_targets]
+
+    if len(targets) > len(set(targets)):
+        raise TVMCException("Duplicate target definitions are not allowed")
+
+    if targets[-1] not in tvm_target_kinds:
+        tvm_target_names = ", ".join(tvm_target_kinds)
+        raise TVMCException(
+            f"The last target needs to be a TVM target. Choices: {tvm_target_names}"
+        )
+
+    tvm_targets = [t for t in targets if t in tvm_target_kinds]
+    if len(tvm_targets) > 1:
+        verbose_tvm_targets = ", ".join(tvm_targets)
+        raise TVMCException(
+            f"Only one of the following targets can be used at a time. "
+            "Found: {verbose_tvm_targets}."
+        )
+
+
+def tokenize_target(target):
+    """
+    Extract a list of tokens from a target specification text.
+
+    It covers some corner-cases that are not covered by the built-in
+    module 'shlex', such as the use of "+" as a punctuation character.
+
+
+    Example
+    -------
+
+    For the input `foo -op1=v1 -op2="v ,2", bar -op3=v-4` we
+    should obtain:
+
+        ["foo", "-op1=v1", "-op2="v ,2"", ",", "bar", "-op3=v-4"]
+
+    Parameters
+    ----------
+    target : str
+        Target options sent via CLI arguments
+
+    Returns
+    -------
+    list of str
+        a list of parsed tokens extracted from the target string
+    """
+
+    target_pattern = (
+        r"(\-{0,2}[\w\-]+\=?"
+        r"(?:[\w\+\-]+(?:,[\w\+\-])*|[\'][\w\+\-,\s]+[\']|[\"][\w\+\-,\s]+[\"])*|,)"
+    )
+
+    return re.findall(target_pattern, target)
+
+
+def parse_target(target):
+    """
+    Parse a plain string of targets provided via a command-line
+    argument.
+
+    To send more than one codegen, a comma-separated list
+    is expected. Options start with -<option_name>=<value>.
+
+    We use python standard library 'shlex' to parse the argument in
+    a POSIX compatible way, so that if options are defined as
+    strings with spaces or commas, for example, this is considered
+    and parsed accordingly.
+
+
+    Example
+    -------
+
+    For the input `--target="foo -op1=v1 -op2="v ,2", bar -op3=v-4"` we
+    should obtain:
+
+      [
+        {
+            name: "foo",
+            opts: {"op1":"v1", "op2":"v ,2"},
+            raw: 'foo -op1=v1 -op2="v ,2"'
+        },
+        {
+            name: "bar",
+            opts: {"op3":"v-4"},
+            raw: 'bar -op3=v-4'
+        }
+      ]
+
+    Parameters
+    ----------
+    target : str
+        Target options sent via CLI arguments
+
+    Returns
+    -------
+    codegens : list of dict
+        This list preserves the order in which codegens were
+        provided via command line. Each Dict contains three keys:
+        'name', containing the name of the codegen; 'opts' containing
+        a key-value for all options passed via CLI; 'raw',
+        containing the plain string for this codegen
+    """
+    codegens = []
+
+    parsed_tokens = tokenize_target(target)
+
+    split_codegens = []
+    current_codegen = []
+    split_codegens.append(current_codegen)
+    for token in parsed_tokens:
+        # every time there is a comma separating
+        # two codegen definitions, prepare for
+        # a new codegen
+        if token == ",":
+            current_codegen = []
+            split_codegens.append(current_codegen)
+        else:
+            # collect a new token for the current
+            # codegen being parsed
+            current_codegen.append(token)
+
+    # at this point we have a list of lists,
+    # each item on the first list is a codegen definition
+    # in the comma-separated values
+    for codegen_def in split_codegens:
+        # the first is expected to be the name
+        name = codegen_def[0]
+        raw_target = " ".join(codegen_def)
+        all_opts = codegen_def[1:] if len(codegen_def) > 1 else []
+        opts = {}
+        for opt in all_opts:
+            try:
+                # deal with -- prefixed flags
+                if opt.startswith("--"):
+                    opt_name = opt[2:]
+                    opt_value = True
+                else:
+                    opt = opt[1:] if opt.startswith("-") else opt
+                    opt_name, opt_value = opt.split("=", maxsplit=1)
+            except ValueError:
+                raise ValueError(f"Error when parsing '{opt}'")
+
+            opts[opt_name] = opt_value
+
+        codegens.append({"name": name, "opts": opts, "raw": raw_target})
+
+    return codegens
+
+
+def is_inline_json(target):
+    try:
+        json.loads(target)
+        return True
+    except json.decoder.JSONDecodeError:
+        return False
+
+
 def target_from_cli(target):
     """
     Create a tvm.target.Target instance from a
@@ -93,18 +256,33 @@ def target_from_cli(target):
     -------
     tvm.target.Target
         an instance of target device information
+    extra_targets : list of dict
+        This list preserves the order in which extra targets were
+        provided via command line. Each Dict contains three keys:
+        'name', containing the name of the codegen; 'opts' containing
+        a key-value for all options passed via CLI; 'raw',
+        containing the plain string for this codegen
     """
+    extra_targets = []
 
     if os.path.exists(target):
         with open(target) as target_file:
-            logger.info("using target input from file: %s", target)
+            logger.debug("target input is a path: %s", target)
             target = "".join(target_file.readlines())
+    elif is_inline_json(target):
+        logger.debug("target input is inline JSON: %s", target)
+    else:
+        logger.debug("target input is plain text: %s", target)
+        try:
+            parsed_targets = parse_target(target)
+        except ValueError as ex:
+            raise TVMCException(f"Error parsing target string '{target}'.\nThe error was: {ex}")
 
-    # TODO(@leandron) We don't have an API to collect a list of supported
-    #       targets yet
-    logger.debug("creating target from input: %s", target)
+        validate_targets(parsed_targets)
+        target = parsed_targets[-1]["raw"]
+        extra_targets = parsed_targets[:-1] if len(parsed_targets) > 1 else []
 
-    return tvm.target.Target(target)
+    return tvm.target.Target(target), extra_targets
 
 
 def tracker_host_port_from_cli(rpc_tracker_str):

python/tvm/driver/tvmc/compiler.py

@@ -28,7 +28,7 @@
 from tvm.contrib import cc
 from tvm.contrib import utils
 
-from . import common, frontends
+from . import common, composite_target, frontends
 from .main import register_parser
 
 
@@ -72,7 +72,7 @@ def add_compile_parser(subparsers):
     )
     parser.add_argument(
         "--target",
-        help="compilation target as plain string, inline JSON or path to a JSON file",
+        help="compilation targets as comma separated string, inline JSON or path to a JSON file.",
         required=True,
     )
     parser.add_argument(
@@ -185,13 +185,21 @@ def compile_model(
     """
     dump_code = [x.strip() for x in dump_code.split(",")] if dump_code else None
     mod, params = frontends.load_model(path, model_format, shape_dict)
+    config = {}
 
     if alter_layout:
         mod = common.convert_graph_layout(mod, alter_layout)
 
-    tvm_target = common.target_from_cli(target)
+    tvm_target, extra_targets = common.target_from_cli(target)
     target_host = tvm_target if not target_host else target_host
 
+    for codegen_from_cli in extra_targets:
+        codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"])
+        partition_function = codegen["pass_pipeline"]
+        mod = partition_function(mod, params)
+        if codegen["config_key"] is not None:
+            config[codegen["config_key"]] = codegen_from_cli["opts"]
+
     if tuning_records and os.path.exists(tuning_records):
         logger.debug("tuning records file provided: %s", tuning_records)
 
@@ -203,22 +211,21 @@ def compile_model(
 
         if use_autoscheduler:
             with auto_scheduler.ApplyHistoryBest(tuning_records):
-                with tvm.transform.PassContext(
-                    opt_level=3, config={"relay.backend.use_auto_scheduler": True}
-                ):
+                config["relay.backend.use_auto_scheduler"] = True
+                with tvm.transform.PassContext(opt_level=3, config=config):
                     logger.debug("building relay graph with autoscheduler")
                     graph_module = relay.build(
                         mod, target=target, params=params, target_host=target_host
                     )
         else:
             with autotvm.apply_history_best(tuning_records):
-                with tvm.transform.PassContext(opt_level=3):
+                with tvm.transform.PassContext(opt_level=3, config=config):
                     logger.debug("building relay graph with tuning records")
                     graph_module = relay.build(
                         mod, tvm_target, params=params, target_host=target_host
                     )
     else:
-        with tvm.transform.PassContext(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3, config=config):
             logger.debug("building relay graph (no tuning records provided)")
             graph_module = relay.build(mod, tvm_target, params=params, target_host=target_host)
 

python/tvm/driver/tvmc/composite_target.py

@@ -0,0 +1,68 @@
+# 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.
+"""
+Provides support to composite target on TVMC.
+"""
+import logging
+
+from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib
+from tvm.relay.op.contrib.ethosn import partition_for_ethosn
+
+from .common import TVMCException
+
+
+# pylint: disable=invalid-name
+logger = logging.getLogger("TVMC")
+
+# Global dictionary to map targets with the configuration key
+# to be used in the PassContext (if any), and a function
+# responsible for partitioning to that target.
+REGISTERED_CODEGEN = {
+    "acl": {
+        "config_key": None,
+        "pass_pipeline": partition_for_arm_compute_lib,
+    },
+    "ethos-n77": {
+        "config_key": "relay.ext.ethos-n.options",
+        "pass_pipeline": partition_for_ethosn,
+    },
+}
+
+
+def get_codegen_names():
+    """Return a list of all registered codegens.
+
+    Returns
+    -------
+    list of str
+        all registered targets
+    """
+    return list(REGISTERED_CODEGEN.keys())
+
+
+def get_codegen_by_target(name):
+    """Return a codegen entry by name.
+
+    Returns
+    -------
+    dict
+        requested target information
+    """
+    try:
+        return REGISTERED_CODEGEN[name]
+    except KeyError:
+        raise TVMCException("Composite target %s is not defined in TVMC." % name)