API to trigger partitioning

apache · Aug 14, 2019 · 478c60a · 478c60a
1 parent 843c3ab
commit 478c60a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
@@ -1987,6 +1987,15 @@ MXNET_DLL int MXSetCalibTableToQuantizedSymbol(SymbolHandle qsym_handle,
                                                const float* high_quantiles,
                                                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);
 /*!
  * \brief Run subgraph pass based on the backend provided
  * \param sym_handle symbol to be converted

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,