address review comments.

python/tvm/contrib/pipeline_executor.py

@@ -71,7 +71,7 @@ def build(pipe_configs):
         )
 
         mconf["dev"] = "{},{}".format(dev.device_type, dev.device_id)
-        # Create a pipeline configuration, 'mod_idx' start from 0.
+        # Create a pipeline configuration.
         string_config[mod_idx] = mconf
         libs[mod_idx] = {"lib": lib, "dev": dev}
 
@@ -479,8 +479,11 @@ def get_config(self):
                     for dep in binding.bindings:
                         dep_item = {}
                         _, dname = dep.get_name()
-                        dep_item["mod_idx"] = dep.get_owner_idx()
-                        dep_item["input_name"] = dname
+                        if dep.is_pipeline_executor_interface():
+                            dep_item["global_output_index"] = int(dname)
+                        else:
+                            dep_item["mod_idx"] = dep.get_owner_idx()
+                            dep_item["input_name"] = dname
                         dep_conf.append(dep_item)
 
                 # The value of ouput_idx start from 0.
@@ -633,15 +636,12 @@ def export_library(self, directory_path):
 
             load_config.append(mconfig)
 
-        # Export the configuration file to disk.
         with open(load_config_file_name, "w") as file_handle:
             json.dump(load_config, file_handle)
 
-        # Export the pipeline configuration file to disk.
         with open(pipeline_config_file_name, "w") as file_handle:
             json.dump(self.mods_config, file_handle)
 
-        # Export the configuration file to disk.
         config_file_name = "{}/config".format(directory_path)
         with open(config_file_name, "w") as file_handle:
             json.dump(config, file_handle)

src/runtime/pipeline/pipeline_executor.h

@@ -57,8 +57,7 @@ class TVM_DLL PipelineExecutor : public ModuleNode {
   void Init(const std::vector<Module>& modules, const std::string& pipeline_json);
   /*!
    * \brief Use the information of mod_config to create a graph executor list.
-   * \param mod_config The configuration information generated by the library export library
-   *  function call.
+   * \param mod_config The configuration information generated by the library export function call.
    */
   std::vector<Module> CreateGraphModules(const ModuleConfig& mod_config);
   /*!
@@ -102,8 +101,8 @@ class TVM_DLL PipelineExecutor : public ModuleNode {
           LOG(FATAL) << "do not support key " << key;
         }
       }
-      // Check if mod_idx is read successfully, in this level reading there all the moudles
-      // are graph executor modules, hence the mod_idx should start from 0.
+      // Check if mod_idx is read successfully, All the moudles here are graph executor modules,
+      // hence the mod_idx should start from 0.
       ICHECK(mod_idx >= 0) << "Invalid mod_idx value " << mod_idx;
       // Load the lib, json, and params information.
       ICHECK(!lib_name.empty()) << "lib_name is empty.";
@@ -143,8 +142,8 @@ class TVM_DLL PipelineExecutor : public ModuleNode {
           LOG(FATAL) << "do not support key " << key;
         }
       }
-      // Check if mod_idx is read successfully, in this level reading there is no any moudle
-      // is PipelineExecutor, hence the mod_idx should start from 0.
+      // Check if mod_idx is read successfully, All moudles here are graph executor modules, hence
+      // the mod_idx should start from 0.
       ICHECK(mod_idx >= 0) << "Invalid mod_idx value " << mod_idx;
       // Check if the output is read successfully.
       ICHECK(!output.Empty()) << "Invalid output binding result.";

src/runtime/pipeline/pipeline_struct.h

@@ -26,7 +26,6 @@
 #include <string>
 #include <unordered_map>
 #include <vector>
-#define PIPELINE_EXECUTOR_INDEX -1
 /*!
  * \brief All binding information of a output interface.
  */
@@ -36,23 +35,12 @@ struct OutputBindings {
    *  of the module.
    */
   std::unordered_map<int, std::string> bindings;
-  /*!
-   * \brief If there is one PipelineExecutor binding in bindings, then the current output is
-   *  PipelineExecutor interface.
-   * \return Whether this output interface is PipelineExecutor output interface.
-   */
-  bool IsGlobalOutput() const {
-    int num_output = 0;
-    for (auto binding : bindings) {
-      /* The output is a PipelineExecutor output when the index value
-       * equal PIPELINE_EXECUTOR_INDEX.
-       */
-      num_output += (binding.first == PIPELINE_EXECUTOR_INDEX);
-    }
-    /* If this output is a global output then there is only one such output in map.*/
-    ICHECK(num_output <= 1);
-    return num_output == 1;
-  }
+  /*! Whether this output binding with global output.*/
+  bool global_binding = false;
+  /*! The index of global output that this output binding with.*/
+  int global_output_index = std::numeric_limits<int>::min();
+  /*!\brief Whether this binding bind with a PipelineExecutor output interface.*/
+  bool IsGlobalOutput() const { return global_binding; }
   /*!
    * \brief Create a module interface map from JSONReader.
    * \param reader JSON reader.
@@ -67,19 +55,31 @@ struct OutputBindings {
       while (reader->NextObjectItem(&key)) {
         if (key == "mod_idx") {
           reader->Read(&mod_idx);
-        }
-        if (key == "input_name") {
+        } else if (key == "input_name") {
           reader->Read(&input_name);
+        } else if (key == "global_output_index") {
+          reader->Read(&global_output_index);
+          // When find key with value as 'global_output_index', this output should bind with
+          // a global output.*/
+          global_binding = true;
+        } else {
+          LOG(FATAL) << "do not support key " << key;
         }
       }
-      // In this level 'Load' that reading the output binding , the module can be
-      // a 'PipelineExecutor', hence the value of 'mod_idx' should start from
-      // PIPELINE_EXECUTOR_INDEX.
-      ICHECK(mod_idx >= PIPELINE_EXECUTOR_INDEX);
-      bindings[mod_idx] = input_name;
+      // When this output bind with a global interface, check whether the interface index
+      // is correct.
+      if (global_binding) {
+        ICHECK(global_output_index >= 0);
+      } else {
+        // When this output does not bind with a global interface, check whether the binding
+        // module index is correct.
+        ICHECK(mod_idx >= 0);
+        bindings[mod_idx] = input_name;
+      }
     }
   }
 };
+
 /*!
  * \brief The binding information of all outputs of a module.
  */
@@ -121,9 +121,10 @@ struct OutputMap {
       while (reader->NextObjectItem(&key)) {
         if (key == "output_idx") {
           reader->Read(&output_idx);
-        }
-        if (key == "dependent") {
+        } else if (key == "dependent") {
           reader->Read(&binding);
+        } else {
+          LOG(FATAL) << "do not support key " << key;
         }
       }
       ICHECK(output_idx >= 0);

tests/python/relay/test_pipeline_executor.py

@@ -81,7 +81,7 @@ def get_manual_conf(mods, target):
         "output": [
             {"output_idx": 0, "dependent": [{"mod_idx": 1, "input_name": "data_0"}]},
             {"output_idx": 1, "dependent": [{"mod_idx": 2, "input_name": "data_0"}]},
-            {"output_idx": 2, "dependent": [{"mod_idx": -1, "input_name": "0"}]},
+            {"output_idx": 2, "dependent": [{"global_output_index": 0}]},
         ],
     }
     mod_config[mods[0]] = {
@@ -112,7 +112,7 @@ def get_manual_conf(mods, target):
 
     pipe_config3 = {
         "mod_idx": 2,
-        "output": [{"output_idx": 0, "dependent": [{"mod_idx": -1, "input_name": "1"}]}],
+        "output": [{"output_idx": 0, "dependent": [{"global_output_index": 1}]}],
     }
     mod_config[mods[2]] = {
         "pipeline": pipe_config3,