[TVMC][microNPU] tvmc option for printing which operators are offloaded to Ethos-U

python/tvm/driver/tvmc/compiler.py

@@ -20,8 +20,12 @@
 """
 import logging
 import os.path
+import re
+import itertools
+from copy import deepcopy
 from typing import Any, Optional, Dict, List, Union, Callable, Sequence
 from pathlib import Path
+from collections import defaultdict
 
 import tvm
 from tvm import autotvm, auto_scheduler
@@ -31,6 +35,8 @@
 from tvm.ir.memory_pools import WorkspaceMemoryPools
 from tvm.target import Target
 from tvm.relay.backend import Executor, Runtime
+from tvm.relay.analysis.operations_distribution import analyze_operations_distribution
+from tvm.relay.transform.suffixes import tag_suffixes
 
 from . import composite_target, frontends, TVMCException
 from .model import TVMCModel, TVMCPackage
@@ -69,6 +75,16 @@ def add_compile_parser(subparsers, _, json_params):
         default="",
         help="comma separated list of formats to export the input model, e.g. 'asm,ll,tir,relay'.",
     )
+    parser.add_argument(
+        "--dump-offloads",
+        default="",
+        help="output a mapping of which operations of the initial Relay "
+        "will be transferred to which backend, indicating the composite "
+        "that includes those operations, "
+        "e.g. '--dump-offloads -' to dump to the console, "
+        "e.g. '--dump-offloads <path_to_file>' to dump to the file. "
+        "If not presented, no output is done. ",
+    )
     parser.add_argument(
         "--model-format",
         choices=frontends.get_frontend_names(),
@@ -171,6 +187,8 @@ def drive_compile(args):
 
     dump_code = [x.strip() for x in args.dump_code.split(",")] if args.dump_code else None
 
+    dump_offloads = args.dump_offloads if args.dump_offloads else ""
+
     additional_targets = reconstruct_target_args(args)
     workspace_pools_target, extra_targets = target_from_cli(args.target, additional_targets)
     transform_args = parse_graph_transform_args(args)
@@ -187,6 +205,7 @@ def drive_compile(args):
         cross_options=args.cross_compiler_options,
         output_format=args.output_format,
         dump_code=dump_code,
+        dump_offloads=dump_offloads,
         target_host=None,
         disabled_pass=args.disabled_pass,
         pass_context_configs=args.pass_config,
@@ -213,6 +232,7 @@ def compile_model(
     cross_options: Optional[str] = None,
     output_format: str = "so",
     dump_code: Optional[List[str]] = None,
+    dump_offloads: str = "",
     target_host: Optional[str] = None,
     disabled_pass: Optional[str] = None,
     pass_context_configs: Optional[List[str]] = None,
@@ -259,6 +279,10 @@ def compile_model(
     dump_code : list[str], optional
         Dump the generated code for the specified source types, on
         the requested target. Choose from: ["asm", "ll", "tir", "relay"].
+    dump_offloads : str
+        Dump the information about the partition of input model's layers by external codegen.
+        Can be '' to not dump at all, '-' to dump to the console
+        or '<path_to_file>' to dump to the specified file.
     target_host : str, optional
         The target of the host machine if host-side code
         needs to be generated.
@@ -313,6 +337,13 @@ def compile_model(
     if "tir" in dump_code:
         config, dumps = add_tir_to_dumps(config, dumps)
 
+    initial_relay = None
+    if dump_offloads != "":
+        # add suffixes to the span field for calls in Relay
+        mod = tag_suffixes(mod)
+        # remember initial Relay
+        initial_relay = deepcopy(mod)
+
     tvm_target, extra_targets = target_from_cli(target, additional_target_options)
     tvm_target, target_host = Target.canon_target_and_host(tvm_target, target_host)
 
@@ -337,6 +368,10 @@ def compile_model(
         for partition_function, opts in zip(partition_functions, partition_opts):
             mod = partition_function(mod, params, mod_name=mod_name, **opts)
 
+        if initial_relay:
+            # dump which operations are offloaded to which backend
+            dump_operation_offloads(mod, initial_relay, dump_offloads)
+
         if tuning_records and os.path.exists(tuning_records):
             logger.debug("tuning records file provided: %s", tuning_records)
 
@@ -496,3 +531,141 @@ def save_dumps(module_name: str, dumps: Dict[str, str], dump_root: str = "."):
         dump_name = module_name + "." + dump_format
         with open(Path(dump_root, dump_name), "w") as f:
             f.write(dumps[dump_format])
+
+
+def dump_operation_offloads(mod: tvm.ir.IRModule, initial_mod: tvm.ir.IRModule, dump_path: str):
+    """This helper function forms a line-by-line output of the initial Relay lines,
+    indicating which operations are ported to which target,
+    and indicating the composite that includes those operations;
+    the 'generic' target refers to operations uploaded to the host, e.g
+    'target1        <-     target1.qnn_conv2d'
+    'target1        <-          %0 = qnn.conv2d(%tfl.quantize, %v_param_1, ...'
+    'target1        <-          %1 = nn.bias_add(%0, %v_param_2, axis=3);'
+    'target1        <-          %2 = qnn.requantize(%1, meta[relay.Constant]...'
+    'target2        <-     target2.reshape'
+    'target2        <-          %3 = reshape(%2, newshape=[1, 1001]);'
+    'generic        <-     %4 = nn.pad(%3, -128f, pad_width=[[0, 0], [1, 1]...'
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The partitioned IRModule with external global functions.
+    initial_mod : tvm.ir.IRModule
+        The initial IRModule that gets generated from a relay frontend.
+    dump_path: str
+        Value of the "dump_offloads" compiler atribute.
+        Could be dash ("-") or file path or empty string for
+        printing to console, file or doing nothing respectively.
+    """
+    print_to_console = dump_path == "-"
+    save_to_file = all([dump_path != "-", dump_path != ""])
+
+    if print_to_console or save_to_file:
+
+        operations_distribution = analyze_operations_distribution(mod)
+
+        def annotate_f(x):
+            ret = ""
+            if isinstance(x, relay.Call):
+                # if there is no x.span.source_name.name in operations_distribution,
+                # this could mean that the span was not copied during the application of passes
+                # to the Relay, in which case we can not associate the initial Relay string
+                # with the resulting Relay call
+                source_name = x.span.source_name.name
+                if source_name in operations_distribution:
+                    compiler_name, op_name, func_id = operations_distribution[source_name]
+                    ret = (
+                        f", compiler_name: {compiler_name}, op_name: {op_name}, "
+                        f"func_id: {func_id}"
+                    )
+                else:
+                    ret = ", compiler_name: unknown, op_name: unknown, func_id: unknown"
+            return ret
+
+        initial_relay_astext = initial_mod.astext(show_meta_data=False, annotate=annotate_f).split(
+            "\n"
+        )
+
+        # funcs_list is a list of internal composite/function IDs
+        # generated by analyze_operations_distribution().
+        # funcs_list helps keep the order of lines from the initial Relay.
+        funcs_list = []
+
+        # target_statistic is a mapping of the target name to the
+        # number of initial Relay calls offloaded on the target
+        target_statistic = defaultdict(int)
+
+        # funcs_dict is a mapping of the generated analyze_operations_distribution
+        # internal composite/function IDs to a list, where:
+        # 1st element is
+        #   (1a): target name - it could be "generic" or "unknown" or
+        #   (1b): specific target name, like "ethos-u" or "cmsis-nn"
+        # 2nd element is
+        #   (2a): corresponding initial Relay line for the case (1a) or
+        #   (2b): the name of the target composite functon in the other case (1b)
+        # 3rd element or subsequent ones are presented only for the case (2b)
+        # and are the initial Relay lines included in the corresponding
+        # target composite functon
+        funcs_dict = {}
+
+        # Here we group together initial Relay lines from the one composite
+        counter = itertools.count()
+        for s in initial_relay_astext:
+            result = re.search(
+                r"(compiler_name: )(.*)(, op_name: )(.*)(, func_id: )((.*)(?=;)|(.*))", s
+            )
+            if result:
+                target_name = result.group(2)
+                op_name = result.group(4)
+                func_id = result.group(6)
+                s = re.sub(r", compiler_name: (.*)", "", s).lstrip()
+                target_statistic[target_name] += 1
+
+                # create an identifier for each "unknown" case to keep the lines order
+                if func_id == "unknown":
+                    func_id = str(next(counter) * -1)
+
+                if func_id not in funcs_dict:
+                    funcs_list.append(func_id)
+                    funcs_dict[func_id] = [target_name]
+                    if target_name not in ["unknown", "generic"]:
+                        funcs_dict[func_id].append(op_name)
+
+                funcs_dict[func_id].append(s)
+
+        # Here we prepare the output for printing.
+        # The output in most cases keeps the original order of the Relay lines
+        # but some lines are moved to be in the corresponding composite group
+        output = []
+        total = 0
+        output.append("Total number of operators and distribution by targets")
+        output.append("Total:")
+        for target, statistic in target_statistic.items():
+            total += statistic
+            output.append(f"{target}: {statistic}")
+        output[1] += f" {total}"
+        output[len(target_statistic) + 1] += "\n"
+
+        for func_id in funcs_list:
+            _list = funcs_dict[func_id]
+            output.append(f"{_list[0]:10}     <-     {_list[1]}")
+            if _list[0] == "unknown":
+                output.append(
+                    "Warning: The above line means that some pass(es) \
+                              in Relay partitioning"
+                )
+                output.append("do not copy the span when the call is recreated")
+                output.append(
+                    "and a line from initial Relay could not be associated \
+                              with the resulting Relay"
+                )
+            for el in _list[2:]:
+                output.append(f"{_list[0]:10}     <-          {el}")
+
+        if print_to_console:
+            print("\n" + "\n".join(output))
+        if save_to_file:
+            file_path = os.path.abspath(dump_path)
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            with open(file_path, "w") as f:
+                f.write("\n".join(output))

python/tvm/relay/analysis/operations_distribution.py

@@ -0,0 +1,102 @@
+# 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.
+"""Utilities that enable analyze Relay and get mappings for
+the unique identifier of the Relay line to the tuple of
+compiler name, composite name and composite/function identifier."""
+import tvm
+from tvm import relay
+from tvm.relay.expr_functor import ExprVisitor
+
+
+class AnalyzeOperationsDistribution(ExprVisitor):
+    """A visitor pass that maintains the dictionary unique_op_ids where
+    the tuple (compiler name, composite name, composite/function identifier)
+    corresponds to the unique identifier of the Relay line.
+    TVMC compiler adds a unique Relay line identifier as a suffix
+    to the call span field using the tag_suffixes pass
+    if the --dump-offloads option is specified.
+
+    Attributes
+    ----------
+    unique_op_ids : Dict[str, str, int]
+        Mapping the unique identifier of the Relay line obtained from
+        the "span" field of the Call and the tuple of compiler name,
+        composite name and internal composite/function identifier.
+    func_name : str
+        The name of the composite name in the partitioned Relay or
+        'generic' in case the Call has not been included in any composite.
+    func_id : int
+        Internal(inside unique_op_ids) composite/function identifier.
+    compiler_name : str
+        A name of the compiler (e.g. 'ethos-u' or 'cmsis-nn') or 'generic'
+        in case the Call has not been included in any composite.
+    """
+
+    def __init__(self):
+        self.unique_op_ids = {}
+        self.func_name = ""
+        self.func_id = 1
+        self.compiler_name = ""
+        super().__init__()
+
+    def extract(self, call: relay.Call):
+        self.compiler_name = "generic"
+        self.func_name = "generic"
+        if "Compiler" in call.attrs:
+            self.compiler_name = call.attrs["Compiler"]
+        self.visit(call)
+
+    def visit_call(self, call: relay.Call):
+        if isinstance(call.op, tvm.ir.Op):
+            if call.span:
+                src = call.span.source_name.name
+                self.unique_op_ids[src] = [self.compiler_name, self.func_name, self.func_id]
+                if self.func_name == "generic":
+                    self.func_id += 1
+        if isinstance(call.op, relay.Function):
+            self.func_name = call.op.attrs["Composite"]
+            self.func_id += 1
+        super().visit_call(call)
+
+
+def analyze_operations_distribution(mod):
+    """Traverses the partitioned graph to get the unique identifier
+    of the Relay line from the Call's span field.
+    The result is maintained in the dictionary unique_op_ids where
+    the unique indicator obtained from the op's span corresponds to
+    the tuple (compiler name, composite name, composite/function identifier).
+    With this information we can annotate the textual representation
+    of the initial Relay by indicating into which target composite
+    and function the operators are converted
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The partitioned Relay graph usually obtained with
+        partition_for_<target> function
+
+    Returns
+    -------
+    unique_op_ids : Dict[str, str, int]
+        Mapping from the unique identifier of the Relay line to the tuple of
+        compiler name, composite name, internal composite/function
+        identifier.
+    """
+    analyze = AnalyzeOperationsDistribution()
+    for _, func in mod.functions.items():
+        analyze.extract(func)
+    return analyze.unique_op_ids