[GitHub] [incubator-mxnet] HahTK commented on a change in pull request #18779: [WIP] Support extra inputs for subgraph ops

[GitBox <gi...@apache.org>]

HahTK commented on a change in pull request #18779:
URL: https://github.com/apache/incubator-mxnet/pull/18779#discussion_r460655282



##########
File path: example/extensions/lib_subgraph/subgraph_lib.cc
##########
@@ -331,6 +330,46 @@ REGISTER_PARTITIONER(mySelect)
 .setCreateSelector("strategy1", createSelector)
 .setReviewSubgraph("strategy1", myReviewSubgraph);
 
+/* \brief a basic pass that adds a new input for subgraph ops */
+MXReturnValue addInputPass(const std::string& in_graph, const std::string** out_graph,

Review comment:
       The new input has unknown and arbitrary tensor shape.
   Do we handle that somewhere ?
   
   Also is this added to both the 
   a) subgraph_op and
   b) the subgraph attr of the subgraph_op ?
   
   We probably do not want that in the latter.
   
   
   Also how is shape handled ? 
   The input tensor can be an take arbitrary and unknown shape at this point in the flow.
   Once known, this shape is unknown to user and needs to be remembered.
   Therefore, this op likely needs the shape attr with a byte dtype.
   Do these already work with this update ? or is it upto the user to add ?

##########
File path: example/extensions/lib_subgraph/subgraph_lib.cc
##########
@@ -331,6 +332,46 @@ REGISTER_PARTITIONER(mySelect)
 .setCreateSelector("strategy1", createSelector)
 .setReviewSubgraph("strategy1", myReviewSubgraph);
 
+/* \brief a basic pass that adds a new input for subgraph ops */
+MXReturnValue addInputPass(const std::string& in_graph, const std::string** out_graph,
+			   const std::unordered_map<std::string, std::string>& options,
+			   const std::unordered_map<std::string, MXTensor>& args,
+			   const std::unordered_map<std::string, MXTensor>& aux,
+			   const PassResource& res) {
+  // convert graph from JSON string to Graph/Node data structure
+  Graph *g = Graph::fromString(in_graph);
+  //find node with '_custom_subgraph_op' op type
+  for(Node* n : g->nodes) {
+    if(n->op.compare("_custom_subgraph_op") == 0) {
+      //set extra input
+      n->attrs[MX_STR_EXTRA_INPUTS] = std::to_string(1);
+      
+      //create a new input Node
+      Node* input = new Node();
+      std::string input_name = n->name + "_input";
+      input->name = input_name;
+      input->op = "null";
+      //add a new node in graph
+      g->nodes.push_back(input);
+      g->inputs.push_back(input);
+      //connect new input to node
+      input->outputs.push_back({n,(int)(n->inputs.size())});
+      //connect node to new input
+      n->inputs.push_back({input,0});
+      // add a corresponding tensor for this input
+      MXTensor* arg_ = res.alloc_arg(input_name,{1},MXContext::CPU(0),kFloat32);
+    }
+  }
+
+  //convert back to JSON string from Graph/Node
+  *out_graph = new std::string(g->toString());
+  return MX_SUCCESS;
+}
+
+REGISTER_PASS(addInputPass)

Review comment:
       I propose, not overloading the optime_for() path to apply custom passes.
   Instead we should just have an API to apply_pass() or apply_pass_for().
   
   This makes the intent clear and unambiguous.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org