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

Author: sergio-grovety

python/tvm/driver/tvmc/compiler.py

@@ -19,8 +19,10 @@
 """
 import logging
 import os.path
+from copy import deepcopy
 from typing import Any, Optional, Dict, List, Union, Callable, Sequence
 from pathlib import Path
+import re
 
 import tvm
 from tvm import autotvm, auto_scheduler
@@ -30,6 +32,7 @@
 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 . import composite_target, frontends, TVMCException
 from .model import TVMCModel, TVMCPackage
@@ -73,6 +76,16 @@ def add_compile_parser(subparsers, _, json_params):
         default="",
         help="comma separated list of formats to export the input model, e.g. 'asm,ll,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(),
@@ -175,6 +188,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)
 
@@ -190,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,
         desired_layout=args.desired_layout,
         disabled_pass=args.disabled_pass,
@@ -216,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,
     desired_layout: Optional[str] = None,
     disabled_pass: Optional[str] = None,
@@ -257,6 +274,10 @@ def compile_model(
     dump_code : list, optional
         Dump the generated code for the specified source types, on
         the requested target.
+    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.
@@ -292,6 +313,11 @@ def compile_model(
 
     config = parse_configs(pass_context_configs)
 
+    initial_relay = None
+    if dump_offloads != "":
+        # remember initial relay
+        initial_relay = deepcopy(tvmc_model.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)
 
@@ -316,6 +342,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)
 
@@ -459,3 +489,79 @@ 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        <-     target2.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 IRModule that gets generated from a relay frontend.
+    initial_relay_astext : list
+        List of input model IR strings.
+    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):
+                assert (
+                    x.span.source_name.name in operations_distribution
+                ), f"Initial relay op {x.span.source_name.name} is not present \
+                    in relay after partition"
+                compiler_name, op_name, func_id = operations_distribution[x.span.source_name.name]
+                ret = f", func_id: {func_id}, compiler_name: {compiler_name}, op_name: {op_name}"
+            return ret
+
+        initial_relay_astext = initial_mod.astext(show_meta_data=False, annotate=annotate_f).split(
+            "\n"
+        )
+
+        output = []
+        prev_func_id = -1
+        for s in initial_relay_astext:
+            result = re.search(
+                r"(func_id: )(.*)(, compiler_name: )(.*)(, op_name: )((.*)(?=;)|(.*))", s
+            )
+            if result:
+                func_id = result.group(2)
+                target_name = result.group(4)
+                op_name = result.group(6)
+                if prev_func_id != func_id:
+                    # assume that operations of one composite go one after another in the relay
+                    prev_func_id = func_id
+                    # add the name of a composite
+                    if op_name != "generic":
+                        output.append(f"{target_name:10}     <-     {op_name}")
+                sub = re.sub(r", compiler_name: (.*)", "", s).lstrip()
+                if target_name == "generic":
+                    output.append(f"{target_name:10}     <-     {sub}")
+                else:
+                    output.append(f"{target_name:10}     <-          {sub}")
+        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,87 @@
+# 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 unique
+input module layer name to the tuple of compiler and operation name"""
+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, compiler operation name, ) corresponds
+    to unique input module layer name
+
+    Attributes
+    ----------
+    unique_op_ids : Dict[str, str, int]
+        Mapping from unique input module layer name to the tuple of compiler,
+        operation name and function id.
+    func_name : str
+        Function id in partitioned module.
+    func_id : int
+        Function identifier in partitioned module.
+    compiler_name : str
+        A name of the compiler (e.g. 'ethos-u' or 'cmsis-nn').
+    """
+
+    def __init__(self):
+        self.unique_op_ids = {}
+        self.func_name = ""
+        self.func_id = 0
+        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 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 source name field from the op's span.
+    The result is maintained in the dictionary unique_op_ids where the tuple
+    (compiler name, compiler composite name, function identifier in the partitioned graph)
+    corresponds to unique input module layer name
+
+    Parameters
+    ----------
+    tvm.ir.IRModule
+        The IRModule that gets generated from a relay frontend.
+
+    Returns
+    -------
+    unique_op_ids : Dict[str, str, int]
+        Mapping from unique input module layer name to the tuple of compiler,
+        operation name and function id.
+    """
+    analyze = AnalyzeOperationsDistribution()
+    for _, func in mod.functions.items():
+        analyze.extract(func)
+    return analyze.unique_op_ids

python/tvm/relay/frontend/common.py

@@ -1067,6 +1067,20 @@ def __init__(self, span):
             self._span = tvm.relay.Span(tvm.relay.SourceName(span.decode("utf-8")), 0, 0, 0, 0)
         else:
             assert False, f"unsupported span type: {type(span)}"
+        self.suffix_str = "_PART_"
+        self.counter = 0
+        self.distance_from_leaf = -1
+
+    def _create_span(self):
+        """Adds suffix_str + counter value to _span.source_name.name,
+        to create a unique source_name for the Relay layer
+        """
+        if self.distance_from_leaf == 0:
+            return tvm.relay.Span(tvm.relay.SourceName(self._span), 0, 0, 0, 0)
+        self.distance_from_leaf -= 1
+        span_str = "{}{}{}".format(self._span.source_name.name, self.suffix_str, str(self.counter))
+        self.counter += 1
+        return tvm.relay.Span(tvm.relay.SourceName(span_str), 0, 0, 0, 0)
 
     def visit(self, expr):
         if hasattr(expr, "span") and expr.span:
@@ -1093,7 +1107,7 @@ def visit_call(self, call):
         # ExprMutator will return directly if subject belongs to Op type
         new_op = self.visit(call.op)
         return _expr.CallWithFields(
-            call, new_op, new_args, call.attrs, call.type_args, None, self._span
+            call, new_op, new_args, call.attrs, call.type_args, None, self._create_span()
         )
 
     def visit_var(self, var):

src/relay/transforms/annotate_target.cc

@@ -258,6 +258,7 @@ class AnnotateTargetRewriter : public ExprRewriter {
     Array<Expr> compiler_begins = std::get<1>(target_n_args);
     Call new_call = Call(post_call->op, compiler_begins, post_call->attrs);
     new_call->checked_type_ = pre->checked_type_;
+    new_call->span = pre->span;
 
     // Update the target map.
     op_expr_to_target_[new_call] = target;

tests/python/contrib/test_ethosu/infra.py

@@ -285,7 +285,7 @@ def get_tflite_model(model_url):
     return tflite_model_buf
 
 
-def get_tflite_graph(tf_func, shapes, ranges=None):
+def get_tflite_graph(tf_func, shapes, ranges=None, partitioned=True):
     tensor_specs = [tf.TensorSpec(shape, dtype=tf.float32) for shape in shapes]
     if not ranges:
         ranges = [(0, 1) for _ in shapes]
@@ -314,8 +314,9 @@ def representative_dataset():
     tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
 
     relay_module, params = relay.frontend.from_tflite(tflite_model)
-    mod = partition_for_ethosu(relay_module, params)
-    return mod, tflite_graph
+    if partitioned:
+        relay_module = partition_for_ethosu(relay_module, params)
+    return relay_module, tflite_graph
 
 
 def compare_ethosu_with_reference(