API to trigger partitioning

apache · Aug 15, 2019 · 3860dee · 3860dee
1 parent 843c3ab
commit 3860dee
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
@@ -2002,6 +2002,15 @@ MXNET_DLL int MXGenBackendSubgraph(SymbolHandle sym_handle, const char *backend,
  * \param ret_sym_handle returned atomic symbol
  */
 MXNET_DLL int MXGenAtomicSymbolFromSymbol(SymbolHandle sym_handle, SymbolHandle *ret_sym_handle);
+/*!
+ * \brief Partitions symbol for given backend, potentially creating subgraphs
+ * \param sym_handle symbol to be partitioned
+ * \param backend backend name
+ * \param ret_sym_handle partitioned symbol returned
+ */
+MXNET_DLL int MXOptimizeForBackend(SymbolHandle sym_handle,
+                                   const char* backend,
+                                   SymbolHandle* ret_sym_handle);
 
 
 //--------------------------------------------

diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
@@ -1437,6 +1437,12 @@ def _gen_atomic_symbol(self):
         return Symbol(handle)
 
 
+    def optimizeFor(self, backend):
+        """Partition symbol and optimize it for a given backend"""
+        out = SymbolHandle()
+        check_call(_LIB.MXOptimizeForBackend(self.handle, c_str(backend), ctypes.byref(out)))
+
+
     # pylint: disable=too-many-locals
     def simple_bind(self, ctx, grad_req='write', type_dict=None, stype_dict=None,
                     group2ctx=None, shared_arg_names=None, shared_exec=None,

diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc
@@ -1181,3 +1181,31 @@ int MXShallowCopySymbol(SymbolHandle src, SymbolHandle* out) {
   *out = out_sym;
   API_END_HANDLE_ERROR(delete out_sym);
 }
+
+inline nnvm::Symbol MXOptimizeForBackend(SymbolHandle sym_handle,
+                                         const char* backend_name,
+                                         SymbolHandle* ret_sym_handle) {
+  nnvm::Symbol *s = new nnvm::Symbol();
+  API_BEGIN();
+  nnvm::Symbol *sym = static_cast<nnvm::Symbol *>(sym_handle);
+  *s = sym->Copy();
+  nnvm::Graph orig_g = Symbol2Graph(*s);
+  /* TODO
+  if (infer_shapes_flag) {
+    orig_g = InferForwardAttrs(orig_g, arg_shapes, arg_dtypes, arg_stypes,
+                               default_ctx, ctx_map, in_arg_ctxes, aux_state_ctxes, true);
+  } */
+  const auto backend = mxnet::op::SubgraphBackendRegistry::Get()->GetSubgraphBackend(backend_name);
+  const auto& subgraph_prop_list = backend->GetSubgraphProperties();
+  for (auto property : subgraph_prop_list) {
+    nnvm::Graph g = orig_g;
+    property->SetAttr("graph", g);
+    g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(property);
+    g = ApplyPass(std::move(g), "BuildSubgraph");
+    property->RemoveAttr("graph");
+    g.attrs.erase("subgraph_property");
+    s->outputs = g.outputs;
+  }
+  *ret_sym_handle = s;
+  API_END_HANDLE_ERROR(delete s);
+}