Add dedup flag to master from #19112 (#19246)

* initial commit

* update build_subgraph

* added test

* calling test

Co-authored-by: Ubuntu <[email protected]>

src/c_api/c_api_symbolic.cc

@@ -1410,6 +1410,11 @@ int MXOptimizeForBackend(SymbolHandle sym_handle,
   for (mx_uint i = 0; i < num_options; ++i)
     options_map.emplace(keys[i], vals[i]);
 
+  // set dedup option as attribute on graph to enable dedup during partitioning
+  if (options_map.count("dedup_subgraph") > 0 &&
+      options_map.at("dedup_subgraph").compare("True") == 0)
+    g.attrs["dedup_subgraph"] = std::make_shared<nnvm::any>(std::string("True"));
+
   if (mxnet::op::SubgraphBackendRegistry::Get()->backend_map_.count(backend_name) > 0) {
     // use subgraph backend
     const auto backend = mxnet::op::SubgraphBackendRegistry

src/operator/subgraph/build_subgraph.cc

@@ -539,10 +539,13 @@ void CutGraphInputs(const std::vector<nnvm::NodeEntry*> &input_entries,
                     std::vector<nnvm::NodeEntry> *orig_entries,
                     std::vector<nnvm::NodeEntry> *unique_orig_entries,
                     std::vector<nnvm::NodeEntry*> *unique_input_entries,
-                    const bool skip_var = false) {
+                    const bool skip_var = false,
+                    const bool dedup = false) {
   orig_entries->resize(input_entries.size());
   // map for creating unique var nodes for deduplicating entries from the same node
-  std::unordered_map<std::string, nnvm::NodeEntry> name_count_map;
+  std::unordered_map<std::string, nnvm::NodeEntry> name_map;
+  std::unordered_map<std::string, int> name_count_map;
+
   for (size_t i = 0; i < input_entries.size(); ++i) {
     nnvm::NodeEntry *e = input_entries[i];
     // If the node is a variable itself, we may want to skip the node.
@@ -558,18 +561,24 @@ void CutGraphInputs(const std::vector<nnvm::NodeEntry*> &input_entries,
     CHECK_EQ(output_names.size(), 1U);
     const std::string& var_name = output_names[0];
     // check if this entry is a duplicate
-    auto it = name_count_map.find(var_name);
-    if (name_count_map.end() == it) {
+    if (name_count_map.count(var_name) == 0) {
       // first use of this node as input to subgraph
+      name_count_map.emplace(var_name, 0);
       unique_orig_entries->push_back(*e);
       unique_input_entries->push_back(e);
       nnvm::ObjectPtr n = nnvm::CreateVariableNode(var_name + std::to_string(0));
-      *e = nnvm::NodeEntry{n, 0, 0};
-      // store node for re-use
-      name_count_map.emplace(var_name, *e);
+      name_map.emplace(var_name, nnvm::NodeEntry{n, 0, 0});
     } else {
       // other use of same node as input to subgraph
-      *e = it->second;
+      name_count_map[var_name]++;
+    }
+
+    if (dedup) {
+      *e = name_map[var_name];
+    } else {
+      nnvm::ObjectPtr n = nnvm::CreateVariableNode(
+        var_name + std::to_string(name_count_map[var_name]));
+      *e = nnvm::NodeEntry{n, 0, 0};
     }
   }
 }
@@ -600,13 +609,14 @@ void CreateSubgraphNode(nnvm::Graph* g,
 #if DEBUG_SUBGRAPH
   LOG(INFO) << "Searching for input entries...";
 #endif
+  bool dedup_subgraph = g->HasAttr("dedup_subgraph");
   std::vector<nnvm::NodeEntry*> input_entries;  // nodes that produce inputs to subgraph nodes
   FindInputEntries(*g, simple_nodes, subgraph_nodes, *entry_top_order_map, &input_entries);
   std::vector<nnvm::NodeEntry> orig_input_entries;  // original input entries (dupes)
   std::vector<nnvm::NodeEntry> unique_orig_entries;  // unique original input entries
   std::vector<nnvm::NodeEntry*> unique_input_entries;  // unique modified subgraph inputs
   CutGraphInputs(input_entries, &orig_input_entries, &unique_orig_entries,
-                 &unique_input_entries, false);
+                 &unique_input_entries, false, dedup_subgraph);
 #if DEBUG_SUBGRAPH
   PrintNodeEntries(input_entries);
   LOG(INFO) << "Searching for output entries...";
@@ -621,25 +631,36 @@ void CreateSubgraphNode(nnvm::Graph* g,
   nnvm::NodeEntryEqual node_equal;
   sym.outputs.resize(output_entries.size());
   for (size_t i = 0; i < output_entries.size(); ++i) {
-    if (i == 0) {  // add first entry
-      sym.outputs[idx] = *output_entries[i];
-    } else if (!node_equal(sym.outputs[idx], *output_entries[i])) {  // compare to see if diff
-      // add new entries
-      idx++;
-      sym.outputs[idx] = *output_entries[i];
-    }  // else skip over dupe entries
+    if (dedup_subgraph) {
+      if (i == 0) {  // add first entry
+        sym.outputs[idx] = *output_entries[i];
+      } else if (!node_equal(sym.outputs[idx], *output_entries[i])) {  // compare to see if diff
+        // add new entries
+        idx++;
+        sym.outputs[idx] = *output_entries[i];
+      }  // else skip over dupe entries
+    } else {
+      sym.outputs[i] = *output_entries[i];
+    }
   }
-  sym.outputs.resize(idx+1);
+  if (dedup_subgraph)
+    sym.outputs.resize(idx+1);
 
   const SubgraphPropertyPtr& subg_prop = g->GetAttr<SubgraphPropertyPtr>("subgraph_property");
-  subg_prop->InitSubgraphInputs(&unique_input_entries, &unique_orig_entries);
+  if (dedup_subgraph)
+    subg_prop->InitSubgraphInputs(&unique_input_entries, &unique_orig_entries);
+  else
+    subg_prop->InitSubgraphInputs(&input_entries, &orig_input_entries);
   nnvm::ObjectPtr n = subg_prop->CreateSubgraphNode(sym, subgraph_selector, subgraph_id);
   // CreateSubgraphNode returns NULL if subgraph property determines that subgraph is sub-optimal
   // In that case, subgraph node is not created and graph is not modified
   if (n) {
     // Connect the external nodes to the subgraph node.
     subg_prop->ConnectSubgraphOutputs(n, &output_entries);
-    subg_prop->ConnectSubgraphInputs(n, &unique_input_entries, &unique_orig_entries);
+    if (dedup_subgraph)
+      subg_prop->ConnectSubgraphInputs(n, &unique_input_entries, &unique_orig_entries);
+    else
+      subg_prop->ConnectSubgraphInputs(n, &input_entries, &orig_input_entries);
 
     const auto& indexed_graph = g->indexed_graph();
     for (size_t i = 0; i < n->inputs.size(); ++i) {

src/operator/subgraph/partitioner/custom_subgraph_property.h

@@ -319,9 +319,8 @@ class  CustomSubgraphProperty: public SubgraphProperty {
     opt_vals_.clear();
     options_map_.clear();
     // store options in map in subgraph property to re-use later for reviewSubgraph
-    for (auto& kv : options_map) {
-      options_map_.push_back(kv);
-    }
+    options_map_.insert(options_map.begin(), options_map.end());
+
     // convert options_map_ to char* to pass to backend library
     for (auto& kv : options_map_) {
       opt_keys_.push_back(kv.first.c_str());
@@ -526,7 +525,7 @@ class  CustomSubgraphProperty: public SubgraphProperty {
   mxnet::ext::opCallFree_t call_free_;
   std::unordered_map<std::string, int> supported_nodes;
   std::string subgraph_op_name;
-  std::vector<std::pair<std::string, std::string>> options_map_;
+  std::unordered_map<std::string, std::string> options_map_;
   std::vector<const char*> opt_keys_, opt_vals_;
   std::vector<std::string> in_arg_names, in_aux_names;
   NDArray **in_args_ptr;

src/operator/subgraph/subgraph_property.h

@@ -258,7 +258,7 @@ class SubgraphProperty {
     kAdjust,
   };
 
-  explicit SubgraphProperty(SgPropertyType type = kCreate) : type_(type) {}
+  explicit SubgraphProperty(SgPropertyType type = kCreate) : type_(type), dedup_subgraph(true) {}
 
   /*!
    * \brief The criteria of selecting the subgraph nodes.
@@ -269,7 +269,14 @@ class SubgraphProperty {
   }
 
   virtual void PrePartition(const nnvm::Graph& g,
-    const std::unordered_map<std::string, std::string>& options_map) {}
+    const std::unordered_map<std::string, std::string>& options_map) {
+    if (options_map.count("dedup_subgraph") > 0 &&
+        options_map.at("dedup_subgraph").compare("True") == 0) {
+      dedup_subgraph = true;
+    } else {
+      dedup_subgraph = false;
+    }
+  }
 
   virtual void PostPartition(const nnvm::Graph& g) {}
 
@@ -348,14 +355,19 @@ class SubgraphProperty {
     nnvm::NodeEntry prevNodeEntry;
     uint32_t idx = 0;
     for (size_t i = 0; i < output_entries->size(); ++i) {
-      // increment the output idx for each unique output of the subgraph
-      if (i != 0 && !node_equal(prevNodeEntry, *output_entries->at(i)))
-        idx++;
-      prevNodeEntry = *output_entries->at(i);  // make a copy so we can compare before modifying
-      // change output entry to point to subgraph instead of original node
-      *output_entries->at(i) = nnvm::NodeEntry{subgraph_node, idx, 0};
+      if (dedup_subgraph) {
+        // increment the output idx for each unique output of the subgraph
+        if (i != 0 && !node_equal(prevNodeEntry, *output_entries->at(i)))
+          idx++;
+        prevNodeEntry = *output_entries->at(i);  // make a copy so we can compare before modifying
+        // change output entry to point to subgraph instead of original node
+        *output_entries->at(i) = nnvm::NodeEntry{subgraph_node, idx, 0};
+      } else {
+        *output_entries->at(i) = nnvm::NodeEntry{subgraph_node, static_cast<uint32_t>(i), 0};
+      }
     }
   }
+
   /*!
    * \brief Connect subgraph internal input with external input entries.
    * By default, each input entry will connect in top sorted order.
@@ -417,6 +429,7 @@ class SubgraphProperty {
  protected:
   SgPropertyType type_;
   std::unordered_map<std::string, std::shared_ptr<nnvm::any>> attrs_;
+  bool dedup_subgraph;
 };
 
 using SubgraphPropertyPtr = std::shared_ptr<SubgraphProperty>;

tests/python/unittest/test_subgraph_op.py

@@ -382,6 +382,37 @@ def test_subgraph_exe8(sym, subgraph_backend, op_names):
     for i in range(len(outputs1)):
         assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))
 
+@pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2'])
+@pytest.mark.parametrize('sym,op_names', get_graphs())
+def test_subgraph_exe9(sym, subgraph_backend, op_names):
+    """Call optimize_for to infer shapes, types and dtypes followed by graph partitioning and
+    dedup subgraph, then bind and compare results of the partitioned sym and the original sym."""
+    # bind
+    sym, _, _ = sym
+    arg_shapes, _, aux_shapes = sym.infer_shape()
+    arg_names = sym.list_arguments()
+    aux_names = sym.list_auxiliary_states()
+    arg_dict = {name:mx.nd.random.uniform(shape=shape) for name,shape in zip(arg_names,arg_shapes)}
+    aux_dict = {name:mx.nd.random.uniform(shape=shape) for name,shape in zip(aux_names,aux_shapes)}
+    exe1 = sym._bind(ctx=mx.current_context(), args=arg_dict, aux_states=aux_dict, grad_req='null')
+    exe1.forward()
+
+    # infer shape/type before partition before bind
+    check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)),
+                                                   c_str_array(op_names)))
+    part_sym = sym.optimize_for(subgraph_backend, arg_dict, aux_dict, dedup_subgraph=True)
+    check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend)))
+
+    exe2 = part_sym._bind(ctx=mx.current_context(), args=arg_dict, aux_states=aux_dict, grad_req='null')
+    exe2.forward()
+
+    # compare outputs
+    outputs1 = exe1.outputs
+    outputs2 = exe2.outputs
+    assert len(outputs1) == len(outputs2)
+    for i in range(len(outputs1)):
+        assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))
+
 @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2'])
 @pytest.mark.parametrize('sym,op_names', get_graphs())
 def test_subgraph_backend_gluon(sym, subgraph_backend, op_names, tmpdir):
@@ -505,6 +536,7 @@ def hybrid_forward(self, F, x):
             test_subgraph_exe6(sym, subgraph_backend, op_names)
             test_subgraph_exe7(sym, subgraph_backend, op_names)
             test_subgraph_exe8(sym, subgraph_backend, op_names)
+            test_subgraph_exe9(sym, subgraph_backend, op_names)
             test_subgraph_backend_gluon(sym, subgraph_backend, op_names, tmpdir)
             test_subgraph_backend_gluon_ext1(tmpdir)
             test_subgraph_backend_gluon_ext2(tmpdir)