Compiler option "--target-ethos-u-dump_npu_functions_coverage" has been replaced by more generic "--dump-offloads" with the same meaning

Author: sergio-grovety

python/tvm/driver/tvmc/compiler.py

@@ -18,9 +18,11 @@
 Provides support to compile networks both AOT and JIT.
 """
 import logging
+import os
 import os.path
 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 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,15 @@ def compile_model(
 
     config = parse_configs(pass_context_configs)
 
+    # remember initial relay
+    initial_relay_astext = (
+        mod.astext(
+            show_meta_data=False,
+        ).split("\n")
+        if dump_offloads != ""
+        else ""
+    )
+
     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 +346,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_astext != "":
+            # dump which operations are offloaded to which backend
+            dump_operation_offloads(mod, initial_relay_astext, dump_offloads)
+
         if tuning_records and os.path.exists(tuning_records):
             logger.debug("tuning records file provided: %s", tuning_records)
 
@@ -459,3 +493,65 @@ 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_relay_astext: list, dump_path: str):
+    """This helper function forms a line-by-line output of the initial Relay lines,
+    indicating which operations are ported to which backend,
+    indicating the composite that includes those operations e.g
+    'device1    <- device2.qnn_conv2d'
+    'device1    <-        %0 = qnn.conv2d(%tfl.quantize, %v_param_1, ...'
+    'device1    <-        %1 = nn.bias_add(%0, %v_param_2, axis=3);'
+    'device1    <-        %2 = qnn.requantize(%1, meta[relay.Constant]...'
+    'device2    <- device2.reshape'
+    'device2    <-        %3 = reshape(%206, newshape=[1, 1001]);'
+
+    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)
+        output = []
+        prev_op_name = ""
+        for s in initial_relay_astext:
+            result = re.search(r"(output_name: )(.*)(:0:0)(.*)\*/", s)
+            if result:
+                op = result.group(2)
+                hardw_id = ""
+                if op in operations_distribution:
+                    hardw_id = operations_distribution[op][0]
+                    result = re.search(r"(.*)(_PART_)", op)
+                    op_name = result.group(1) if result else op
+                    if prev_op_name != op_name:
+                        # assume that operations of one composite go one after another in the relay
+                        prev_op_name = op_name
+                        # add the name of a composite
+                        output.append(
+                            f"{hardw_id:10}     <-     {(operations_distribution[op][1])}"
+                        )
+                if "span" in s:
+                    sub = re.sub(r" /\* span(.*)\*/", "", s).lstrip()
+                    if hardw_id == "":
+                        hardw_id = "generic"
+                        output.append(f"{hardw_id:10}     <-     {sub}")
+                    else:
+                        output.append(f"{hardw_id: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,86 @@
+# 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 maintai the dictionary unic_op_ids where
+    the tuple (compiler, compiler operation name) corresponds
+    to unique input module layer name
+
+    Attributes
+    ----------
+    unic_op_ids : Dict[str, Tuple(str, str)]
+        Mapping from unique input module layer name to the tuple of compiler and operation name.
+    func_name : str
+        Function name.
+    compiler_name : str
+        A name of the compiler (e.g. 'ethos-u' or 'cmsis-nn').
+    """
+
+    def __init__(self):
+        self.unic_op_ids = {}
+        self.func_name = ""
+        self.compiler_name = ""
+        super().__init__()
+
+    def extract(self, call: relay.Call):
+        self.compiler_name = call.attrs["Compiler"]
+        self.visit(call)
+
+    def get_src_op_from_span(self, span):
+        span_elements = span.split(",")
+        for s in span_elements:
+            if "output_name" in s:
+                src_op = s.split("output_name: ")[1]
+                return src_op
+        return ""
+
+    def visit_call(self, call: relay.Call):
+        if isinstance(call.op, tvm.ir.Op):
+            if call.span:
+                src = self.get_src_op_from_span(call.span.source_name.name)
+                self.unic_op_ids[src] = [self.compiler_name, self.func_name]
+        if isinstance(call.op, relay.Function):
+            self.func_name = call.op.attrs["Composite"]
+        super().visit_call(call)
+
+
+def analyze_operations_distribution(mod):
+    """Traverses the graph to get "output_name" field from the op's span. The
+    result is maintained in the dictionary unic_op_ids where
+    the tuple (compiler, compiler operation name) corresponds to unique input module layer name
+
+    Parameters
+    ----------
+    tvm.ir.IRModule
+        The IRModule that gets generated from a relay frontend.
+
+    Returns
+    -------
+    unic_op_ids : Dict[str, Tuple(str, str)]
+        Mapping from unique input module layer name to the tuple of compiler and operation name.
+    """
+    analyze = AnalyzeOperationsDistribution()
+    for _, func in mod.functions.items():
+        if "Compiler" in func.attrs:
+            analyze.extract(func)
+    return analyze.unic_op_ids

python/tvm/relay/expr.py

@@ -316,10 +316,13 @@ class TupleGetItem(ExprWithOp):
 
     index: int
         The index.
+
+    span: Optional[tvm.relay.Span]
+        Span that points to original source code
     """
 
-    def __init__(self, tuple_value, index):
-        self.__init_handle_by_constructor__(_ffi_api.TupleGetItem, tuple_value, index)
+    def __init__(self, tuple_value, index, span=None):
+        self.__init_handle_by_constructor__(_ffi_api.TupleGetItem, tuple_value, index, span)
 
 
 @tvm._ffi.register_object("relay.RefCreate")

python/tvm/relay/frontend/common.py

@@ -30,6 +30,7 @@
 from .. import op as _op
 from .. import ty as _ty
 from .. import analysis
+from ..expr_functor import ExprMutator
 
 
 class DuplicateFilter:
@@ -997,3 +998,57 @@ def try_resolve_var_to_const(x, graph_params):
         return _op.const(value, dtype)
 
     return x
+
+
+# Span filling to Relay for TFLite taken from
+# commit https://github.com/apache/tvm/pull/9723
+def set_span(sym, node_name):
+    """Set up the span of relay expression(s) while converting OP"""
+
+    class SpanFiller(ExprMutator):
+        """SpanFiller"""
+
+        def __init__(self, node_name, suffix_str="_PART_"):
+            ExprMutator.__init__(self)
+            self.node_name = node_name
+            self.suffix_str = suffix_str
+            self.counter = 0
+            self.distance_from_leaf = -1
+
+        def _create_span(self):
+            if self.distance_from_leaf == 0:
+                return tvm.relay.Span(tvm.relay.SourceName(self.node_name), 0, 0, 0, 0)
+            self.distance_from_leaf -= 1
+            span_str = "{}{}{}".format(self.node_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_call(self, call):
+            if call.span is None:
+                self.distance_from_leaf += 1
+                new_args = [self.visit(arg) for arg in call.args]
+                return _expr.Call(
+                    call.op, new_args, call.attrs, call.type_args, self._create_span()
+                )
+            return call
+
+        def visit_tuple(self, tup):
+            if tup.span is None:
+                self.distance_from_leaf += 1
+                return _expr.Tuple([self.visit(field) for field in tup.fields], self._create_span())
+            return tup
+
+        def visit_tuple_getitem(self, op):
+            if op.span is None:
+                self.distance_from_leaf += 1
+                return _expr.TupleGetItem(self.visit(op.tuple_value), op.index, self._create_span())
+            return op
+
+        def fill(self, sym):
+            if isinstance(sym, _expr.TupleWrapper):
+                return _expr.TupleWrapper(self.visit(sym.tuple_value), sym.size)
+            if isinstance(sym, _expr.RelayExpr):
+                return self.visit(sym)
+            return sym
+
+    return SpanFiller(node_name).fill(sym)

python/tvm/relay/frontend/tflite.py

@@ -34,6 +34,7 @@
 from .common import ExprTable
 from .common import infer_shape as _infer_shape
 from .common import lstm_cell, to_int_list, shape_of, try_infer_value
+from .common import set_span
 from .tflite_flexbuffer import FlexBufferDecoder
 
 __all__ = ["from_tflite"]
@@ -275,14 +276,21 @@ def convert_op_to_relay(self):
             if ret is None:
                 continue
 
+            # Span filling to Relay for TFLite taken from
+            # commit https://github.com/apache/tvm/pull/9723
             if len(output_tensors) == 1:
                 tensor_idx = output_tensors[0].tensor_idx
-                self.exp_tab.set_expr(get_tensor_name(self.subgraph, tensor_idx), ret)
+                curr_output = get_tensor_name(self.subgraph, tensor_idx)
+                ret = set_span(ret, "location: {}, output_name: {}".format(op_idx, curr_output))
+                self.exp_tab.set_expr(curr_output, ret)
             else:
-                for idx, output_tensor in enumerate(output_tensors):
-                    self.exp_tab.set_expr(
-                        get_tensor_name(self.subgraph, output_tensor.tensor_idx), ret[idx]
-                    )
+                out_names = []
+                for output_tensor in output_tensors:
+                    out_names.append(get_tensor_name(self.subgraph, output_tensor.tensor_idx))
+                curr_output = ", ".join(out_names)
+                ret = set_span(ret, "location: {}, output_name: {}".format(op_idx, curr_output))
+                for idx, out_name in enumerate(out_names):
+                    self.exp_tab.set_expr(out_name, ret[idx])
 
     def get_op_code_str(self, op):
         """Get TFLite ops string representation"""

src/printer/relay_text_printer.cc

@@ -375,7 +375,11 @@ Doc RelayTextPrinter::VisitExpr_(const TupleNode* op) {
 
 Doc RelayTextPrinter::VisitExpr_(const TupleGetItemNode* op) {
   Doc doc;
-  return doc << Print(op->tuple) << "." << op->index;
+  doc << Print(op->tuple) << "." << op->index;
+  if (op->span.defined()) {
+    doc << " /* " << PrintSpan(op->span) << " */";
+  }
+  return doc;
 }
 
 Doc RelayTextPrinter::VisitExpr_(const IfNode* op) {

src/relay/ir/expr.cc

@@ -410,8 +410,8 @@ TupleGetItem WithFields(TupleGetItem tuple_get_item, Optional<Expr> opt_tuple,
 
 TVM_REGISTER_NODE_TYPE(TupleGetItemNode);
 
-TVM_REGISTER_GLOBAL("relay.ir.TupleGetItem").set_body_typed([](Expr tuple, int index) {
-  return TupleGetItem(tuple, index);
+TVM_REGISTER_GLOBAL("relay.ir.TupleGetItem").set_body_typed([](Expr tuple, int index, Span span) {
+  return TupleGetItem(tuple, index, span);
 });
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)