Generate correct output tensor names in C Interface API

Change-Id: I90d86aef7d8b93fb937391a87f79ae9362207058

Author: grant-arm

apps/microtvm/ethosu/src/demo_bare_metal.c

@@ -40,7 +40,7 @@ int main(int argc, char** argv) {
 
   printf("Running inference\n");
   struct tvmgen_default_outputs outputs = {
-      .output = output,
+      .MobilenetV2_Predictions_Reshape_11 = output,
   };
   struct tvmgen_default_inputs inputs = {
       .tfl_quantize = input,

python/tvm/micro/model_library_format.py

@@ -270,9 +270,13 @@ def _get_inputs_and_outputs_from_module(mod):
     main_func = _get_main_relay_func(mod)
     inputs = [argument.name_hint for argument in main_func.params]
 
-    outputs = ["output"]
-    if isinstance(main_func.ret_type, TupleType):
-        outputs = _convert_tuple_to_outputs(main_func.ret_type)
+    if "output_tensor_names" in main_func.attrs:
+        outputs = main_func.attrs["output_tensor_names"]
+    else:
+        if isinstance(main_func.ret_type, TupleType):
+            outputs = _convert_tuple_to_outputs(main_func.ret_type)
+        else:
+            outputs = ["output"]
 
     return inputs, outputs
 

python/tvm/relay/frontend/tflite.py

@@ -18,6 +18,7 @@
 """Tensorflow lite frontend."""
 import itertools
 import math
+import re
 
 import numpy as np
 import tvm
@@ -3769,6 +3770,15 @@ def from_tflite(model, shape_dict=None, dtype_dict=None, op_converter=OperatorCo
     params = {k: _nd.array(np.array(v)) for k, v in exp_tab.params.items()}
     outputs = [exp_tab.get_expr(get_tensor_name(subgraph, i)) for i in model_outputs]
     outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
-    func = _function.Function(analysis.free_vars(outputs), outputs)
+    attrs = tvm.ir.make_node(
+        "DictAttrs",
+        **{
+            "output_tensor_names": [
+                re.sub(r"\W", "_", get_tensor_name(subgraph, model_output))
+                for model_output in model_outputs
+            ]
+        },
+    )
+    func = _function.Function(analysis.free_vars(outputs), outputs, attrs=attrs)
     mod = IRModule.from_expr(func)
     return mod, params

src/relay/backend/aot_executor_codegen.cc

@@ -957,7 +957,27 @@ class AOTExecutorCodegen : public MixedModeVisitor {
         Array<tir::Var>(tir_main_func->params.begin(),
                         tir_main_func->params.begin() + tir_main_func->params.size() -
                             return_sid_.size() - pool_vars.size() - devices.size());
-    ret.metadata = ExecutorCodegenMetadata(inputs, pool_vars, devices, return_sid_.size(),
+
+    std::vector<String> output_var_names;
+    if (auto opt = func->GetAttr<Array<String>>("output_tensor_names")) {
+      Array<String> output_tensor_names = opt.value();
+      for (size_t i = 0; i < output_tensor_names.size(); ++i) {
+        output_var_names.push_back(output_tensor_names[i]);
+      }
+    }
+
+    // If output names have not been specified then generate default output names
+    if (output_var_names.size() == 0) {
+      if (return_sid_.size() == 1) {
+        output_var_names.push_back(String("output"));
+      } else {
+        for (size_t i = 0; i < return_sid_.size(); ++i) {
+          output_var_names.push_back(String("output" + std::to_string(i)));
+        }
+      }
+    }
+
+    ret.metadata = ExecutorCodegenMetadata(inputs, pool_vars, devices, output_var_names,
                                            runtime::kTvmExecutorAot, mod_name, interface_api,
                                            use_unpacked_api_, pool_var_info);
     return ret;

src/relay/backend/utils.cc

@@ -179,14 +179,14 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
     });
 
 ExecutorCodegenMetadata::ExecutorCodegenMetadata(
-    Array<tir::Var> inputs, Array<tir::Var> pools, Array<String> devices, Integer num_outputs,
+    Array<tir::Var> inputs, Array<tir::Var> pools, Array<String> devices, Array<String> outputs,
     String executor, String mod_name, String interface_api, bool unpacked_api,
     Map<tir::Var, tir::usmp::AllocatedPoolInfo> pool_inputs) {
   auto n = make_object<ExecutorCodegenMetadataNode>();
   n->inputs = inputs;
   n->pools = pools;
   n->devices = devices;
-  n->num_outputs = num_outputs;
+  n->outputs = outputs;
   n->executor = executor;
   n->interface_api = interface_api;
   n->unpacked_api = unpacked_api;

src/relay/backend/utils.h

@@ -61,10 +61,10 @@ class ExecutorCodegenMetadataNode : public Object {
  public:
   /*! \brief input information for the main function */
   Array<tir::Var> inputs;
+  /*! \brief output information for the main function */
+  Array<String> outputs;
   /*! \brief pool information for the main function */
   Array<tir::Var> pools;
-  /*! \brief number of outputs of the main function */
-  Integer num_outputs = 1;
   /*! \brief device contexts information for the main function */
   Array<String> devices;
   /*! \brief the executor to be used to run the model */
@@ -81,7 +81,7 @@ class ExecutorCodegenMetadataNode : public Object {
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("inputs", &inputs);
     v->Visit("pools", &pools);
-    v->Visit("num_outputs", &num_outputs);
+    v->Visit("outputs", &outputs);
     v->Visit("devices", &devices);
     v->Visit("executor", &executor);
     v->Visit("unpacked_api", &unpacked_api);
@@ -98,7 +98,7 @@ class ExecutorCodegenMetadataNode : public Object {
 class ExecutorCodegenMetadata : public ObjectRef {
  public:
   TVM_DLL ExecutorCodegenMetadata(Array<tir::Var> inputs, Array<tir::Var> pools,
-                                  Array<String> devices, Integer num_outputs, String executor,
+                                  Array<String> devices, Array<String> outputs, String executor,
                                   String mod_name, String interface_api = "packed",
                                   bool unpacked_api = false,
                                   Map<tir::Var, tir::usmp::AllocatedPoolInfo> pool_inputs =

src/target/source/source_module.cc

@@ -273,7 +273,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
         }
         call_args_ss << " " << input_var->name_hint << ",";
       }
-      for (int i = 0; i < metadata_->num_outputs->value; ++i) {
+      for (unsigned int i = 0; i < metadata_->outputs.size(); ++i) {
         call_args_ss << "void* output" << i << ",";
       }
       for (const tir::Var& pool_var : metadata_->pools) {
@@ -300,7 +300,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
       for (unsigned int i = 0; i < metadata_->inputs.size(); ++i) {
         call_args_ss << "((DLTensor*)(((TVMValue*)args)[" << i << "].v_handle))[0].data,";
       }
-      for (int i = 0; i < metadata_->num_outputs->value; ++i) {
+      for (unsigned int i = 0; i < metadata_->outputs.size(); ++i) {
         int j = metadata_->inputs.size() + i;
         call_args_ss << "((DLTensor*)(((TVMValue*)args)[" << j << "].v_handle))[0].data,";
       }
@@ -328,7 +328,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
       entrypoint_arg_count++;
       run_func_arg_count++;
     }
-    for (int i = 0; i < metadata_->num_outputs->value; i++) {
+    for (unsigned int i = 0; i < metadata_->outputs.size(); i++) {
       run_func_to_entry_point_args[run_func_arg_count] = Integer(entrypoint_arg_count);
       entrypoint_arg_count++;
       run_func_arg_count++;
@@ -356,7 +356,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
 
     // We are creating a copy of the set of pointers
     size_t number_of_io_tensors =
-        metadata_->inputs.size() + metadata_->num_outputs->value + metadata_->pools.size();
+        metadata_->inputs.size() + metadata_->outputs.size() + metadata_->pools.size();
     code_ << "TVMValue tensors[" << number_of_io_tensors << "];\n";
 
     std::unordered_map<int, ObjectRef> run_func_to_entry_point_args =
@@ -395,7 +395,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
         }
         call_args_ss << " " << relay::backend::SanitizeName(input_var->name_hint) << ",";
       }
-      for (int i = 0; i < metadata_->num_outputs->value; ++i) {
+      for (unsigned int i = 0; i < metadata_->outputs.size(); ++i) {
         call_args_ss << "void* output" << i << ",";
       }
       for (const tir::Var& pool_var : metadata_->pools) {
@@ -449,13 +449,11 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
       for (const auto& input : metadata_->inputs) {
         call_args_ss << "inputs->" << relay::backend::SanitizeName(input->name_hint) << ",";
       }
-      if (metadata_->num_outputs->value == 1) {
-        call_args_ss << "outputs->output,";
-      } else {
-        for (int i = 0; i < metadata_->num_outputs->value; ++i) {
-          call_args_ss << "outputs->output" << i << ",";
-        }
+      for (const auto& output : metadata_->outputs) {
+        call_args_ss << "outputs->" << relay::backend::SanitizeName(output);
+        call_args_ss << ",";
       }
+
       for (const tir::Var& pool_var : metadata_->pools) {
         String pool_name = metadata_->pool_inputs.value()[pool_var]->pool_info->pool_name;
         if (IsInternalWorkspaceBuffer(pool_var)) {

tests/python/contrib/test_ethosu/infra.py

@@ -306,9 +306,10 @@ def generate_ref_data_tflite(model):
         interpreter.set_tensor(index, value)
     interpreter.invoke()
 
-    expected_output_data = [
-        interpreter.get_tensor(output_detail["index"]) for output_detail in output_details
-    ]
+    expected_output_data = {
+        output_detail["name"]: interpreter.get_tensor(output_detail["index"])
+        for output_detail in output_details
+    }
 
     return input_data, expected_output_data
 

tests/python/contrib/test_ethosu/test_codegen.py

@@ -754,7 +754,7 @@ def rounding_right_shift(lhs, rhs):
         "ifm": lhs,
         "ifm2": rhs,
     }
-    output_data = generate_output_data(input_data)
+    output_data = {"output": generate_output_data(input_data)[0]}
     ethosu_mod = _create_ethosu_partition(cpu_mod)
 
     _compare_ethosu_with_reference(ethosu_mod, input_data, output_data, accel_type)
@@ -781,7 +781,7 @@ def generate_output_data(input_data):
 
     cpu_mod = create_model()
     input_data = {"ifm": np.random.randint(-120, high=120, size=ifm_shape, dtype="int8")}
-    output_data = generate_output_data(input_data)
+    output_data = {"output": generate_output_data(input_data)[0]}
     ethosu_mod = _create_ethosu_partition(cpu_mod)
 
     _compare_ethosu_with_reference(
@@ -910,7 +910,7 @@ def clz_comp(n):
 
     cpu_mod = create_model()
     input_data = {"ifm": np.random.randint(-500000, high=500000, size=ifm_shape, dtype="int32")}
-    output_data = generate_output_data(input_data)
+    output_data = {"output": generate_output_data(input_data)[0]}
     ethosu_mod = _create_ethosu_partition(cpu_mod)
 
     _compare_ethosu_with_reference(ethosu_mod, input_data, output_data, accel_type)

tests/python/contrib/test_ethosu/test_lookup_table.py

@@ -154,7 +154,7 @@ def test_random_lut(accel_type):
     compiled_models = infra.build_source(
         mod,
         {"ifm": in_data},
-        out_data,
+        {"output": out_data},
         accel_type,
     )