Dynamic subgraph property

example/extensions/lib_subgraph/subgraph_lib.cc

@@ -28,26 +28,6 @@
 #include <algorithm>
 #include "lib_api.h"
 
-MXReturnValue parseAttrs(std::map<std::string, std::string> attrs,
-                         int* num_in, int* num_out) {
-  *num_in = 1;
-  *num_out = 1;
-  if (attrs.count(SUBGRAPH_SYM_JSON)) {
-    // example of subgraph json parsing
-    JsonParser jp;
-    JsonVal val = jp.parse_to_json(attrs[SUBGRAPH_SYM_JSON]);
-    int input = 0;
-    for (auto &item : val.map[JsonVal("nodes")].list) {
-      if (item.map[JsonVal("op")].str == "null")
-        input++;
-    }
-    int output = val.map[JsonVal("heads")].list.size();
-    *num_in = input;
-    *num_out = output;
-  }
-  return MX_SUCCESS;
-}
-
 /* function to execute log operator on floats */
 void myLog(MXTensor &in, MXTensor &out) {
   float* inp = in.data<float>();
@@ -157,7 +137,12 @@ MXReturnValue myExecutor(std::vector<MXTensor> inputs,
 
 class MyStatefulOp : public CustomStatefulOp {
  public:
-  explicit MyStatefulOp(std::string sym) : subgraph_sym(sym) {}
+  explicit MyStatefulOp(std::string sym, std::map<std::string, std::string> attrs)
+    : subgraph_sym(sym), attrs_(attrs) {
+    for (auto kv : attrs) {
+      std::cout << "subgraphOp attributes: " << kv.first << " ==> " << kv.second << std::endl;
+    }
+  }
 
   MXReturnValue Forward(std::vector<MXTensor> inputs,
                         std::vector<MXTensor> outputs,
@@ -167,6 +152,7 @@ class MyStatefulOp : public CustomStatefulOp {
 
  private:
   std::string subgraph_sym;
+  std::map<std::string, std::string> attrs_;
 };
 
 MXReturnValue createOpState(std::map<std::string, std::string> attrs,
@@ -178,13 +164,13 @@ MXReturnValue createOpState(std::map<std::string, std::string> attrs,
     // user can now parse json and run other custom ops inside subgraph
     serialized_subgraph = attrs[SUBGRAPH_SYM_JSON];
   }
-  *op_inst = new MyStatefulOp(serialized_subgraph);
+  attrs.erase(SUBGRAPH_SYM_JSON);
+  *op_inst = new MyStatefulOp(serialized_subgraph, attrs);
   std::cout << "Info: stateful operator created" << std::endl;
   return MX_SUCCESS;
 }
 
 REGISTER_OP(_custom_subgraph_op)
-.setParseAttrs(parseAttrs)
 .setIsSubgraphOp()
 .setCreateOpState(createOpState);
 
@@ -232,7 +218,8 @@ MXReturnValue mySupportedOps(std::string json,
 }
 
 MXReturnValue myAcceptSubgraph(std::string json, int subraph_id, bool* accept,
-                               std::unordered_map<std::string, std::string>& options) {
+                               std::unordered_map<std::string, std::string>& options,
+                               std::unordered_map<std::string, std::string>& attrs) {
   for (auto kv : options) {
     std::cout << "option: " << kv.first << " ==> " << kv.second << std::endl;
   }
@@ -243,6 +230,7 @@ MXReturnValue myAcceptSubgraph(std::string json, int subraph_id, bool* accept,
   } else {
     *accept = true;
     std::cout << "accepting subgraph" << std::endl;
+    attrs["myKey"] = "myVal";
   }
   return MX_SUCCESS;
 }

example/extensions/lib_subgraph/test_subgraph.py

@@ -53,6 +53,7 @@
 print('Testing partitioning with shapes/types')
 arg_array = [mx.nd.ones((3,2),dtype='float32'), mx.nd.ones((3,2),dtype='float32')]
 mysym2 = sym.optimize_for("myProp",arg_array)
+print(mysym2.tojson())
 exe2 = mysym2.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2)), 'b':mx.nd.ones((3,2))})
 out2 = exe2.forward()
 print(out2)

include/mxnet/lib_api.h

@@ -30,6 +30,7 @@
 
 #include <stdint.h>
 #include <stdlib.h>
+#include <string.h>
 #include <vector>
 #include <map>
 #include <unordered_map>
@@ -595,6 +596,7 @@ class CustomOp {
 typedef MXReturnValue (*supportedOps_t)(std::string, int, int*,
                                         std::unordered_map<std::string, std::string>&);
 typedef MXReturnValue (*acceptSubgraph_t)(std::string, int, bool*,
+                                          std::unordered_map<std::string, std::string>&,
                                           std::unordered_map<std::string, std::string>&);
 
 /*!
@@ -777,7 +779,8 @@ typedef int (*partCallSupportedOps_t)(supportedOps_t, const char*, int, int *,
 typedef int (*partCallAcceptSubgraph_t)(acceptSubgraph_t acceptSubgraph,
                                         const char *json, int subgraph_id,
                                         int *accept, const char* const*,
-                                        const char* const*, int);
+                                        const char* const*, int,
+                                        char***, char***, int*);
 
 #define MXLIB_INITIALIZE_STR "initialize"
 typedef int (*initialize_t)(int);
@@ -1142,16 +1145,39 @@ extern "C" {
 #endif
   _partCallAcceptSubgraph(acceptSubgraph_t acceptSubgraph, const char *json,
                           int subgraph_id, int *accept, const char* const* opt_keys,
-                          const char* const* opt_vals, int num_opts) {
+                          const char* const* opt_vals, int num_opts,
+                          char*** attr_keys, char*** attr_vals, int *num_attrs) {
     std::string subgraph_json(json);
     bool accept_bool = false;
     // create map of attributes from list
     std::unordered_map<std::string, std::string> opts;
     for (int i = 0; i < num_opts; i++) {
       opts[std::string(opt_keys[i])] = std::string(opt_vals[i]);
     }
-    MXReturnValue retval = acceptSubgraph(subgraph_json, subgraph_id, &accept_bool, opts);
+
+    //attributes to set on subgraph node
+    std::unordered_map<std::string, std::string> attrs;
+
+    MXReturnValue retval = acceptSubgraph(subgraph_json, subgraph_id, &accept_bool, opts, attrs);
     *accept = accept_bool;
+
+    if (attrs.size() > 0) {
+      *num_attrs = attrs.size();
+      // allocate space for attributes
+      *attr_keys = static_cast<char**>(malloc (attrs.size() * sizeof(char*)));
+      *attr_vals = static_cast<char**>(malloc (attrs.size() * sizeof(char*)));
+
+      // copy attributes
+      int i=0;
+      for (auto kv : attrs) {
+        (*attr_keys)[i] = static_cast<char*>(malloc ((kv.first.size()+1) * sizeof(char)));
+        (*attr_vals)[i] = static_cast<char*>(malloc ((kv.second.size()+1) * sizeof(char)));
+        strcpy((*attr_keys)[i], kv.first.c_str());
+        strcpy((*attr_vals)[i], kv.second.c_str());
+        i++;
+      }
+    }
+
     return retval;
   }
 

src/c_api/c_api.cc

@@ -181,10 +181,10 @@ int MXLoadLib(const char *path) {
     // set bool, dont pass bool across ABI boundary
     isSubgraphOp = _isSubgraphOp;
 
-    CHECK(parse_fp != nullptr) << "Error loading '" << name
-                               << "' custom op, ParseAttrs function was not set.";
     if (!isSubgraphOp) {
       // validate custom operator functions from the dynamic library
+      CHECK(parse_fp != nullptr) << "Error loading '" << name
+                                 << "' custom op, ParseAttrs function was not set.";
       CHECK(fcomp_fp != nullptr || create_opstate_fp != nullptr) << "Error loading '" << name
                             << "' custom op, Forward or CreateOpState function was not set.";
       CHECK(type_fp  != nullptr) << "Error loading '" << name
@@ -655,9 +655,6 @@ int MXLoadLib(const char *path) {
     // check if operator is already registered
     const nnvm::Op *regOpPtr = dmlc::Registry<nnvm::Op>::Get()->Find(name);
     nnvm::Op &regOp = dmlc::Registry<nnvm::Op>::Get()->__REGISTER_OR_GET__(name);
-    regOp.set_attr_parser(attr_parser);
-    regOp.set_num_inputs(num_inputs);
-    regOp.set_num_outputs(num_outputs);
     int plevel = 10;
     if (regOpPtr != nullptr) {
       // overwrite registration of existing op with custom op
@@ -667,6 +664,9 @@ int MXLoadLib(const char *path) {
       plevel++;
     }
     if (!isSubgraphOp) {
+      regOp.set_attr_parser(attr_parser);
+      regOp.set_num_inputs(num_inputs);
+      regOp.set_num_outputs(num_outputs);
       regOp.set_attr<nnvm::FInferType>("FInferType", infer_type, plevel);
       regOp.set_attr<mxnet::FInferShape>("FInferShape", infer_shape, plevel);
       regOp.set_attr<FInferStorageType>("FInferStorageType", infer_storage_type, plevel);
@@ -676,6 +676,8 @@ int MXLoadLib(const char *path) {
         regOp.set_attr<nnvm::FMutateInputs>("FMutateInputs", mutate_inputs, plevel);
     } else {
       using namespace mxnet::op;
+      regOp.set_num_inputs(DefaultSubgraphOpNumInputs);
+      regOp.set_num_outputs(DefaultSubgraphOpNumOutputs);
       regOp.set_attr<nnvm::FInferType>("FInferType",
                                        DefaultSubgraphOpType, plevel);
       regOp.set_attr<mxnet::FInferShape>("FInferShape",
@@ -765,7 +767,7 @@ int MXLoadLib(const char *path) {
       mxnet::op::SubgraphBackendRegistry::Get()->__REGISTER_CUSTOM_PROPERTY__(name_str,
                             std::make_shared<mxnet::op::CustomSubgraphProperty>(
                            strategy_str, callSupportedOps, supportedOps_fp,
-                           callAcceptSubgraph, acceptSubgraph_fp, op_name_str));
+                           callAcceptSubgraph, acceptSubgraph_fp, callFree, op_name_str));
     }
   }
   API_END();

src/operator/subgraph/partitioner/custom_subgraph_property.h

@@ -80,12 +80,14 @@ class  CustomSubgraphProperty: public SubgraphProperty {
                          supportedOps_t supportedOps,
                          partCallAcceptSubgraph_t callAcceptSubgraph,
                          acceptSubgraph_t acceptSubgraph,
+                         opCallFree_t callFree,
                          std::string op_name) :
     subgraphProp(subgraphProp_name),
       callSupportedOps_(callSupportedOps),
       supportedOps_(supportedOps),
       callAcceptSubgraph_(callAcceptSubgraph),
       acceptSubgraph_(acceptSubgraph),
+      callFree_(callFree),
       subgraph_op_name(op_name) {}
 
   // create custom subgraph property
@@ -162,6 +164,9 @@ class  CustomSubgraphProperty: public SubgraphProperty {
   virtual nnvm::NodePtr CreateSubgraphNode(const nnvm::Symbol &sym,
                                            const int subgraph_id = 0) const {
     int accept = 1;
+    int num_attr = 0;
+    char** attr_keys = nullptr;
+    char** attr_vals = nullptr;
     if (acceptSubgraph_) {
       nnvm::Graph g;
       g.outputs = sym.outputs;
@@ -181,18 +186,28 @@ class  CustomSubgraphProperty: public SubgraphProperty {
             node->attrs.dict["isAux"] = "False";
         }
       }
-
+      
       std::string subgraph_json = nnvm::pass::SaveJSON(g);
       CHECK(callAcceptSubgraph_(acceptSubgraph_, subgraph_json.c_str(),
                                 subgraph_id, &accept, opt_keys_.data(),
-                                opt_vals_.data(), opt_keys_.size()))
+                                opt_vals_.data(), opt_keys_.size(),
+                                &attr_keys, &attr_vals, &num_attr))
         << "Error calling acceptSubgraph for '" << subgraphProp << "'";
     }
     if (accept) {
       nnvm::NodePtr n = nnvm::Node::Create();
       n->attrs.op = Op::Get(subgraph_op_name);
       n->attrs.name = "_op" + std::to_string(subgraph_id);
       n->attrs.subgraphs.push_back(std::make_shared<nnvm::Symbol>(sym));
+      // set user specified attributes
+      for (int i=0; i < num_attr; i++) {
+        n->attrs.dict[attr_keys[i]] = attr_vals[i];
+        callFree_(attr_vals[i]);
+        callFree_(attr_keys[i]);
+      }
+      // free memory used by custom op to allocate attributes
+      callFree_(attr_vals);
+      callFree_(attr_keys);
       return n;
     } else {
       return NULL;
@@ -208,6 +223,7 @@ class  CustomSubgraphProperty: public SubgraphProperty {
   supportedOps_t supportedOps_;
   partCallAcceptSubgraph_t callAcceptSubgraph_;
   acceptSubgraph_t acceptSubgraph_;
+  opCallFree_t callFree_;
   std::unordered_set<std::string> supportedNodes;
   std::string subgraph_op_name;
   std::vector<std::pair<std::string, std::string>> options_map_;