Dynamic subgraph compile support

example/extensions/lib_subgraph/subgraph_lib.cc

@@ -206,7 +206,7 @@ MXReturnValue mySupportedOps(std::string json,
     }
 
     //check if op dtype is float
-    if(dtype == kFloat32) {
+    if((dtype == kFloat32 && options.count("reqFloat") > 0) || options.count("reqFloat") == 0) {
       //check if op is in whitelist
       if(std::find(op_names.begin(),op_names.end(),op.str.c_str()) != op_names.end()) {
         // found op in whitelist, set value to 1 to include op in subgraph
@@ -219,10 +219,26 @@ 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>& attrs) {
+                               std::unordered_map<std::string, std::string>& attrs,
+                               std::map<std::string, MXTensor>& args) {
   for (auto kv : options) {
     std::cout << "option: " << kv.first << " ==> " << kv.second << std::endl;
   }
+  for (auto kv : args) {
+    std::cout << "arg: " << kv.first << " ==> (";
+    for (auto s : kv.second.shape)
+      std::cout << s << ",";
+    std::cout << ") [";
+    for (int i=0; i<kv.second.size(); i++)
+      std::cout << kv.second.data<float>()[i] << ", ";
+    std::cout << "]" << std::endl;
+  }
+  if(options.count("reqArgs") > 0 && args.size() == 0) {
+    *accept = false;
+    std::cout << "rejecting subgraph since args were not provided" << std::endl;
+    return MX_SUCCESS;
+  }
+
   if(options.find("reject") != options.end() &&
      options["reject"].compare("True") == 0) {
     *accept = false;

example/extensions/lib_subgraph/test_subgraph.py

@@ -35,12 +35,17 @@
     path = os.path.abspath('libsubgraph_lib.dll')
     mx.library.load(path)
 
+# example model, ops to be partitioned do not have args (use outputs from other ops as inputs)
 a = mx.sym.var('a')
 b = mx.sym.var('b')
 c = a + b
 d = mx.sym.exp(c)
 sym = mx.sym.log(d)
 
+# example model, ops to be partitioned have args
+d2 = mx.sym.exp(a)
+sym2 = mx.sym.log(d2)
+
 #execute in MXNet
 print('-------------------------------')
 print('Testing regular MXNet execution')
@@ -74,3 +79,28 @@
 exe3 = mysym3.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2)), 'b':mx.nd.ones((3,2))})
 out3 = exe3.forward()
 print(out3)
+
+#execute in MXNet
+print('-------------------------------')
+print('Testing regular MXNet execution')
+exe4 = sym2.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2))})
+out4 = exe4.forward()
+print(out4)
+
+# with propogating shapes/types
+print('-------------------------------')
+print('Testing partitioning with shapes/types')
+arg_array = [mx.nd.ones((3,2),dtype='float32')]
+mysym5 = sym2.optimize_for("myProp", arg_array, reqArgs=True)
+print(mysym5.tojson())
+exe5 = mysym5.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2))})
+out5 = exe5.forward()
+print(out5)
+
+# without propogating shapes/types
+print('-------------------------------')
+print('Testing partitioning without shapes/types')
+mysym6 = sym2.optimize_for("myProp", reqArgs=True)
+exe6 = mysym6.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2))})
+out6 = exe6.forward()
+print(out6)

include/mxnet/lib_api.h

@@ -39,7 +39,7 @@
 #include <utility>
 #include <stdexcept>
 
-#define MX_LIBRARY_VERSION 3
+#define MX_LIBRARY_VERSION 4
 
 /*!
  * \brief For loading multiple custom op libraries in Linux, exporting same symbol multiple
@@ -717,7 +717,8 @@ 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>&);
+                                          std::unordered_map<std::string, std::string>&,
+                                          std::map<std::string, MXTensor>&);
 
 /*!
  * \brief An abstract class for subgraph property
@@ -920,7 +921,12 @@ typedef int (*partCallSupportedOps_t)(supportedOps_t supportedOps, const char *j
 typedef int (*partCallAcceptSubgraph_t)(acceptSubgraph_t acceptSubgraph, const char *json,
                                         int subgraph_id, int *accept, const char* const* opt_keys,
                                         const char* const* opt_vals, int num_opts,
-                                        char*** attr_keys, char*** attr_vals, int *num_attrs);
+                                        char*** attr_keys, char*** attr_vals, int *num_attrs,
+                                        const char* const* arg_names, int num_args,
+                                        void* const* arg_data, const int64_t* const* arg_shapes,
+                                        const int* arg_dims, const int* arg_types,
+                                        const size_t* arg_IDs, const char* const* arg_dev_type,
+                                        const int* arg_dev_id);
 
 #define MXLIB_INITIALIZE_STR "initialize"
 typedef int (*initialize_t)(int version);
@@ -1283,7 +1289,12 @@ extern "C" {
   _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,
-                          char*** attr_keys, char*** attr_vals, int *num_attrs) {
+                          char*** attr_keys, char*** attr_vals, int *num_attrs,
+                          const char* const* arg_names, int num_args,
+                          void* const* arg_data, const int64_t* const* arg_shapes,
+                          const int* arg_dims, const int* arg_types,
+                          const size_t* arg_IDs, const char* const* arg_dev_type,
+                          const int* arg_dev_id) {
     std::string subgraph_json(json);
     bool accept_bool = false;
     // create map of attributes from list
@@ -1292,10 +1303,24 @@ extern "C" {
       opts[std::string(opt_keys[i])] = std::string(opt_vals[i]);
     }
 
+    // create a map of named tensors for args
+    std::map<std::string, MXTensor> args;
+    for (int i = 0; i < num_args; i++) {
+      std::vector<int64_t> shapes;
+      for (int j = 0; j < arg_dims[i]; j++)
+        shapes.push_back(arg_shapes[i][j]);
+
+      MXTensor tensor(arg_data[i], shapes, (MXDType)arg_types[i],
+            arg_IDs[i], {arg_dev_type[i], arg_dev_id[i]});
+      args[arg_names[i]] = tensor;
+    }
+
+
     // 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);
+    MXReturnValue retval = acceptSubgraph(subgraph_json, subgraph_id, &accept_bool,
+                                          opts, attrs, args);
     *accept = accept_bool;
 
     if (attrs.size() > 0) {

src/c_api/c_api.cc

@@ -821,11 +821,10 @@ int MXLoadLib(const char *path) {
       LOG(INFO) << "\t\tStrategy[" << j << "] " << strategy_str
                 << " subgraphOp: '" << op_name_str << "'";
 
-      // MXNET_REGISTER_SUBGRAPH_PROPERTY(customBackend, CustomSubgraphProperty);
-      mxnet::op::SubgraphBackendRegistry::Get()->__REGISTER_CUSTOM_PROPERTY__(name_str,
-                            std::make_shared<mxnet::op::CustomSubgraphProperty>(
-                           strategy_str, callSupportedOps, supportedOps_fp,
-                           callAcceptSubgraph, acceptSubgraph_fp, callFree, op_name_str));
+      mxnet::op::SubgraphBackendRegistry::Get()->__REGISTER_CUSTOM_PROPERTY__
+        (name_str, std::make_shared<mxnet::op::CustomSubgraphProperty>
+          (strategy_str, callSupportedOps, supportedOps_fp,
+           callAcceptSubgraph, acceptSubgraph_fp, callFree, op_name_str));
     }
   }
   API_END();

src/c_api/c_api_symbolic.cc

@@ -1383,6 +1383,14 @@ int MXOptimizeForBackend(SymbolHandle sym_handle,
       common::HandleInferStorageTypeError(num_forward_inputs, indexed_graph,
                                           g.GetAttr<StorageTypeVector>("storage_type"));
     }
+    std::vector<std::string> arg_names = sym->ListInputNames(nnvm::Symbol::kReadOnlyArgs);
+    g.attrs["in_args"] = std::make_shared<nnvm::any>(in_args_ptr);
+    g.attrs["in_arg_names"] = std::make_shared<nnvm::any>(arg_names);
+  } else {
+    NDArray **in_args_ptr = static_cast<NDArray**>(nullptr);
+    std::vector<std::string> arg_names;
+    g.attrs["in_args"] = std::make_shared<nnvm::any>(in_args_ptr);
+    g.attrs["in_arg_names"] = std::make_shared<nnvm::any>(arg_names);
   }
   std::vector<std::pair<std::string, std::string>> options_map;
   for (mx_uint i = 0; i < num_options; ++i) {

src/operator/subgraph/build_subgraph.cc

@@ -561,10 +561,12 @@ void CutGraphInputs(const std::vector<nnvm::NodeEntry*> &input_entries,
     nnvm::ObjectPtr n = nnvm::CreateVariableNode(
         var_name + std::to_string(name_count_map[var_name]));
     // set attribute for subgraph input to indicate if it is from an arg/param to model
-    if (e->node->is_variable())
+    if (e->node->is_variable()) {
       n->attrs.dict["isArg"] = "True";
-    else
+      n->attrs.dict["argName"] = var_name;
+    } else {
       n->attrs.dict["isArg"] = "False";
+    }
     *e = nnvm::NodeEntry{n, 0, 0};
   }
 }

src/operator/subgraph/partitioner/custom_subgraph_property.h

@@ -99,6 +99,31 @@ class  CustomSubgraphProperty: public SubgraphProperty {
     const std::vector<std::pair<std::string, std::string>>& options_map) {
     // clear supported_nodes to remove state from previous calls
     supported_nodes.clear();
+    // get input args and arg names
+    in_arg_names = g.GetAttr<std::vector<std::string>>("in_arg_names");
+    in_args_ptr = g.GetAttr<NDArray**>("in_args");
+
+    // convert input args
+    arg_names.clear();
+    arg_data.clear();
+    arg_shapes.clear();
+    arg_dims.clear();
+    arg_types.clear();
+    arg_verIDs.clear();
+    arg_dev_type.clear();
+    arg_dev_id.clear();
+    for (size_t i=0; i < in_arg_names.size(); i++) {
+      arg_names.push_back(in_arg_names[i].c_str());
+      const auto &in_arg = *(in_args_ptr[i]);
+      arg_data.push_back(in_arg.data().dptr_);
+      arg_shapes.push_back(in_arg.shape().data());
+      arg_dims.push_back(in_arg.shape().ndim());
+      arg_types.push_back(in_arg.dtype());
+      arg_verIDs.push_back(in_arg.version());
+      const char* ctx_str = in_arg.ctx().dev_mask() == Context::kCPU ? "cpu" : "gpu";
+      arg_dev_type.push_back(ctx_str);
+      arg_dev_id.push_back(in_arg.ctx().real_dev_id());
+    }
 
     // remove all graph attrs, some cannot be saved to json
     nnvm::Graph graph = std::move(g);
@@ -162,7 +187,7 @@ class  CustomSubgraphProperty: public SubgraphProperty {
   }
   // override CreateSubgraphNode
   virtual nnvm::ObjectPtr CreateSubgraphNode(const nnvm::Symbol &sym,
-                                           const int subgraph_id = 0) const {
+                                             const int subgraph_id = 0) const {
     int accept = 1;
     int num_attr = 0;
     char** attr_keys = nullptr;
@@ -189,9 +214,14 @@ class  CustomSubgraphProperty: public SubgraphProperty {
 
       std::string subgraph_json = nnvm::pass::SaveJSON(g);
       CHECK(call_accept_subgraph_(accept_subgraph_, subgraph_json.c_str(),
-                                subgraph_id, &accept, opt_keys_.data(),
-                                opt_vals_.data(), opt_keys_.size(),
-                                &attr_keys, &attr_vals, &num_attr))
+                                  subgraph_id, &accept, opt_keys_.data(),
+                                  opt_vals_.data(), opt_keys_.size(),
+                                  &attr_keys, &attr_vals, &num_attr,
+                                  arg_names.data(), arg_names.size(),
+                                  arg_data.data(), arg_shapes.data(),
+                                  arg_dims.data(), arg_types.data(),
+                                  arg_verIDs.data(), arg_dev_type.data(),
+                                  arg_dev_id.data()))
         << "Error calling accept_subgraph for '" << subgraph_prop << "'";
     }
     if (accept) {
@@ -228,6 +258,16 @@ class  CustomSubgraphProperty: public SubgraphProperty {
   std::string subgraph_op_name;
   std::vector<std::pair<std::string, std::string>> options_map_;
   std::vector<const char*> opt_keys_, opt_vals_;
+  std::vector<std::string> in_arg_names;
+  NDArray **in_args_ptr;
+  std::vector<const char*> arg_names;
+  std::vector<void*> arg_data;
+  std::vector<const int64_t*> arg_shapes;
+  std::vector<int> arg_dims;
+  std::vector<int> arg_types;
+  std::vector<size_t> arg_verIDs;
+  std::vector<const char*> arg_dev_type;
+  std::vector<int> arg_dev_id;
 };
 }  // namespace op
 }  // namespace mxnet