add separate stateful compute and pass state_op ptr

include/mxnet/lib_api.h

@@ -495,6 +495,11 @@ typedef int (*opCallMutateInputs_t)(mutateInputs_t, const char* const*, const ch
 typedef int (*opCallCreateOpState_t)(createOpState_t, const char* const*, const char* const*, int,
                                      void**);
 
+#define MXLIB_OPCALLFSTATEFULCOMP_STR "_opCallFStatefulCompute"
+typedef int (*opCallFStatefulComp_t)(bool, void*, const int64_t**, int*, void**, int*, int,
+                                     const int64_t**, int*, void**, int*, int,
+                                     xpu_malloc_t, void*);
+
 #define MXLIB_INITIALIZE_STR "initialize"
 typedef int (*initialize_t)(int);
 
@@ -654,7 +659,7 @@ extern "C" {
     return retval;
   }
 
-  /*! \brief returns status of calling Forward function for operator from library */
+  /*! \brief returns status of calling Forward/Backward function for operator from library */
 #if defined(_WIN32) || defined(_WIN64) || defined(__WINDOWS__)
   __declspec(dllexport) int __cdecl
 #else
@@ -750,6 +755,45 @@ extern "C" {
     return create_op(attrs, op_ptr);
   }
 
+  /*! \brief returns status of calling Stateful Forward/Backward for operator from library */
+#if defined(_WIN32) || defined(_WIN64) || defined(__WINDOWS__)
+  __declspec(dllexport) int __cdecl
+#else
+  int
+#endif
+  _opCallFStatefulCompute(bool is_forward, void* state_op,
+                          const int64_t** inshapes, int* indims,
+                          void** indata, int* intypes, int num_in,
+                          const int64_t** outshapes, int* outdims,
+                          void** outdata, int* outtypes, int num_out,
+                          xpu_malloc_t cpu_malloc, void* cpu_alloc) {
+    // create a vector of tensors for inputs
+    std::vector<MXTensor> inputs(num_in);
+    for (int i = 0; i < num_in; i++) {
+      inputs[i].data_ptr = indata[i];
+      inputs[i].dtype = (MXDType)intypes[i];
+      for (int j = 0; j < indims[i]; j++) {
+        inputs[i].shape.push_back(inshapes[i][j]);
+      }
+    }
+
+    // create a vector of tensors for outputs
+    std::vector<MXTensor> outputs(num_out);
+    for (int i = 0; i < num_out; i++) {
+      outputs[i].data_ptr = outdata[i];
+      outputs[i].dtype = (MXDType) outtypes[i];
+      for (int j = 0; j < outdims[i]; j++) {
+        outputs[i].shape.push_back(outshapes[i][j]);
+      }
+    }
+    OpResource res(cpu_malloc, cpu_alloc);
+    CustomStatefulOp* op_ptr = reinterpret_cast<CustomStatefulOp*>(state_op);
+    if (is_forward) {
+      return op_ptr->Forward(inputs, outputs, res);
+    }
+    return op_ptr->Backward(inputs, outputs, res);
+  }
+
   /*!
    * \brief Checks if the MXNet version is supported by the library.
    * If supported, initializes the library.

src/c_api/c_api.cc

@@ -138,13 +138,16 @@ int MXLoadLib(const char *path) {
   opCallCreateOpState_t callCreateOpState =
     get_func<opCallCreateOpState_t>(lib, const_cast<char*>(MXLIB_OPCALLCREATEOPSTATE_STR));
 
+  opCallFStatefulComp_t callFStatefulComp =
+    get_func<opCallFStatefulComp_t>(lib, const_cast<char*>(MXLIB_OPCALLFSTATEFULCOMP_STR));
+
   // get number of operators registered in the library
   opRegSize_t opRegSize = get_func<opRegSize_t>(lib, const_cast<char*>(MXLIB_OPREGSIZE_STR));
   int numOps = opRegSize();
   LOG(INFO) << "Found " << numOps << " operators in library";
 
   /*
-   * The library has custom operators implementation
+   * Get all custom operators implementation from custom library
    * loop and register each operator in the library to NNVM
    */
   opRegGet_t opRegGet = get_func<opRegGet_t>(lib, const_cast<char*>(MXLIB_OPREGGET_STR));
@@ -368,11 +371,11 @@ int MXLoadLib(const char *path) {
 
     // lambda function to convert from external fcompute to internal MXNet types
     auto fcomp_lambda = [=](fcomp_t fcomp_fp,
-                     const nnvm::NodeAttrs& attrs,
-                     const OpContext& ctx,
-                     const std::vector<NDArray>& inputs,
-                     const std::vector<OpReqType>& req,
-                     const std::vector<NDArray>& outputs) {
+                            const nnvm::NodeAttrs& attrs,
+                            const OpContext& ctx,
+                            const std::vector<NDArray>& inputs,
+                            const std::vector<OpReqType>& req,
+                            const std::vector<NDArray>& outputs) {
       // convert attributes to vector of char*
       std::vector<const char*> attr_keys, attr_vals;
       for (auto kv : attrs.dict) {
@@ -406,6 +409,7 @@ int MXLoadLib(const char *path) {
       mshadow::Stream<mxnet::cpu> *cpu_stream = ctx.get_stream<mxnet::cpu>();
 
       // create lambda that captures stream & resource objects
+      // the memory pointer returned will eventually return to user
       auto cpu_alloc = [&](int size) {
         mshadow::Tensor<mxnet::cpu, 1, char> data =
         resource.get_space_typed<mxnet::cpu, 1, char>(mshadow::Shape1(size), cpu_stream);
@@ -418,6 +422,7 @@ int MXLoadLib(const char *path) {
       auto cpu_malloc = [](void* _cpu_alloc, int size) {
         // cast the void* argument to the type for the cpu_alloc lambda function
         alloc_type* cpualloc = static_cast<alloc_type*>(_cpu_alloc);
+        // call cpu_alloc to actually allocate memory and get the pointer
         void* ptr = (*cpualloc)(size);
         return ptr;
       };
@@ -547,37 +552,39 @@ int MXLoadLib(const char *path) {
     };
 
     // stateful forward and backward
-    auto fstateful_lambda = [=](bool forward,
+    auto fstateful_lambda = [=](bool is_forward,
                                 const OpStatePtr& state_ptr,
                                 const OpContext& ctx,
                                 const std::vector<NDArray>& inputs,
                                 const std::vector<OpReqType>& req,
                                 const std::vector<NDArray>& outputs) {
-      // create a vector of tensors for inputs
-      std::vector<MXTensor> c_inputs(inputs.size());
+      std::vector<void*> in_data, out_data;
+      std::vector<const int64_t *> in_shapes, out_shapes;
+      std::vector<int> in_dims, out_dims;
+      std::vector<int> in_types, out_types;
+
+      // convert input tensors to constituent parts
       for (size_t i = 0; i < inputs.size(); i++) {
-        c_inputs[i].data_ptr = inputs[i].data().dptr_;
-        c_inputs[i].dtype = (MXDType)inputs[i].dtype();
-        for (int_least16_t j = 0; j < inputs[i].shape().ndim(); j++) {
-          c_inputs[i].shape.push_back(inputs[i].shape().data()[j]);
-        }
+        in_data.push_back(inputs[i].data().dptr_);
+        in_shapes.push_back(inputs[i].shape().data());
+        in_dims.push_back(inputs[i].shape().ndim());
+        in_types.push_back(inputs[i].dtype());
       }
 
-      // create a vector of tensors for outputs
-      std::vector<MXTensor> c_outputs(outputs.size());
+      // convert output tensors to constituent parts
       for (size_t i = 0; i < outputs.size(); i++) {
-        c_outputs[i].data_ptr = outputs[i].data().dptr_;
-        c_outputs[i].dtype = (MXDType)outputs[i].dtype();
-        for (int j = 0; j < outputs[i].shape().ndim(); j++) {
-          c_outputs[i].shape.push_back(outputs[i].shape().data()[j]);
-        }
+        out_data.push_back(outputs[i].data().dptr_);
+        out_shapes.push_back(outputs[i].shape().data());
+        out_dims.push_back(outputs[i].shape().ndim());
+        out_types.push_back(outputs[i].dtype());
       }
 
       // get memory resource
       const Resource &resource = ctx.requested[0];
       mshadow::Stream<mxnet::cpu> *cpu_stream = ctx.get_stream<mxnet::cpu>();
 
       // create lambda that captures stream & resource objects
+      // the memory pointer returned will eventually return to user
       auto cpu_alloc = [&](int size) {
         mshadow::Tensor<mxnet::cpu, 1, char> data =
         resource.get_space_typed<mxnet::cpu, 1, char>(mshadow::Shape1(size), cpu_stream);
@@ -590,26 +597,23 @@ int MXLoadLib(const char *path) {
       auto cpu_malloc = [](void* _cpu_alloc, int size) {
         // cast the void* argument to the type for the cpu_alloc lambda function
         alloc_type* cpualloc = static_cast<alloc_type*>(_cpu_alloc);
+        // call cpu_alloc to actually allocate memory and get the pointer
         void* ptr = (*cpualloc)(size);
         return ptr;
       };
 
-      OpResource op_res(cpu_malloc, &cpu_alloc);
-
       // retrieve op state object created from CreateOpState
       CustomStatefulOpWrapper& op = state_ptr.get_state<CustomStatefulOpWrapper>();
       CustomStatefulOp* state_op_inst = op.get_instance();
       CHECK(state_op_inst != nullptr)
       << "Error MXNet cannot load custom stateful operator'" << name_str << "'";
 
-      if (forward) {
-        CHECK(state_op_inst->Forward(c_inputs, c_outputs, op_res))
-        << "Error calling ForwardStateful for custom operator '" << name_str << "'";
-      } else {
-        CHECK(state_op_inst->Backward(c_inputs, c_outputs, op_res))
-        << "Error calling BackwardStateful for custom operator '" << name_str << "'";
-      }
-      // return type void
+      // call fcompute function
+      CHECK(callFStatefulComp(is_forward, state_op_inst, in_shapes.data(), in_dims.data(),
+                              in_data.data(), in_types.data(), in_data.size(),
+                              out_shapes.data(), out_dims.data(), out_data.data(),
+                              out_types.data(), out_data.size(), cpu_malloc, &cpu_alloc))
+      << "Error calling FStatefulCompute for custom operator '" << name_str << "'";
     };
 
     auto fstateful_forward = [=](const OpStatePtr& state_ptr,