lib_api.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef MXNET_LIB_API_H_
#define MXNET_LIB_API_H_

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <vector>
#include <map>
#include <unordered_set>
#include <unordered_map>
#include <string>
#include <iostream>
#include <utility>
#include <stdexcept>
#include <functional>
#include <random>
#include <sstream>

#if defined(__NVCC__)
  #include <cuda_runtime.h>
  #include <curand_kernel.h>
#endif

/* Make sure to update the version number everytime you make changes */
#define MX_LIBRARY_VERSION 11

#if defined(_WIN32) || defined(_WIN64) || defined(__WINDOWS__)
  #define PRIVATE_SYMBOL
#else
  #define PRIVATE_SYMBOL  __attribute__ ((visibility ("hidden")))
#endif

/*
 * Import from DLPack https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h
 */
#ifndef DLPACK_VERSION
#ifdef __cplusplus
#define DLPACK_EXTERN_C extern "C"
#else
#define DLPACK_EXTERN_C
#endif

#define DLPACK_VERSION 020

#ifdef _WIN32
#ifdef DLPACK_EXPORTS
#define DLPACK_DLL __declspec(dllexport)
#else
#define DLPACK_DLL __declspec(dllimport)
#endif
#else
#define DLPACK_DLL
#endif

#include <stdint.h>
#include <stddef.h>

#ifdef __cplusplus
extern "C" {
  #endif

  typedef enum {
    kDLCPU = 1,
    kDLGPU = 2,
    kDLCPUPinned = 3,
    kDLOpenCL = 4,
    kDLVulkan = 7,
    kDLMetal = 8,
    kDLVPI = 9,
    kDLROCM = 10,
    kDLExtDev = 12,
  } DLDeviceType;

  typedef struct {
    DLDeviceType device_type;
    int device_id;
  } DLContext;

  typedef enum {
    kDLInt = 0U,
    kDLUInt = 1U,
    kDLFloat = 2U,
  } DLDataTypeCode;

  typedef struct {
    uint8_t code;
    uint8_t bits;
    uint16_t lanes;
  } DLDataType;

  typedef struct {
    void* data;
    DLContext ctx;
    int ndim;
    DLDataType dtype;
    int64_t* shape;
    int64_t* strides;
    uint64_t byte_offset;
  } DLTensor;
#ifdef __cplusplus
}  // DLPACK_EXTERN_C
#endif
#endif

namespace mxnet {
namespace ext {

/* \brief Class to store error messages from extensions to pass to MXNet */
class MXerrorMsgs {
 public:
  /* \brief get singleton pointer to class */
  static MXerrorMsgs* get();

  /* \brief add a new error message */
  std::stringstream& add(const char* file, int line);

  /* \brief return number of error messages */
  int size();

  /* \brief get error message at index */
  const std::string* get(int idx);

 private:
  MXerrorMsgs() {}
  ~MXerrorMsgs();
  std::vector<std::stringstream*> messages;
};

// Add a new error message, example: MX_ERROR_MSG << "my error msg";
#define MX_ERROR_MSG mxnet::ext::MXerrorMsgs::get()->add(__FILE__, __LINE__)

enum MXDType {
  kFloat32 = 0,
  kFloat64 = 1,
  kFloat16 = 2,
  kUint8 = 3,
  kInt32 = 4,
  kInt8  = 5,
  kInt64 = 6,
  kUNSET = 100,
};

/*
 * MXTensor storage type.
 */
enum MXStorageType {
  // dense
  kDefaultStorage = 0,
  // row sparse
  kRowSparseStorage = 1,
  // csr
  kCSRStorage = 2,
};

struct MXContext {
  MXContext();
  explicit MXContext(std::string dev_type_, int dev_id_);
  explicit MXContext(const char* dev_type_, int dev_id_);
  static MXContext CPU();
  static MXContext GPU();
  static MXContext CPU(int dev_id);
  static MXContext GPU(int dev_id);

  std::string dev_type;
  int dev_id;
};

enum MXReturnValue {
  MX_FAIL = 0,
  MX_SUCCESS = 1,
};

// For sparse tensors, read/write the data from NDarray via pointers.
struct MXSparse {
  // Pointer to data.
  void *data{nullptr};
  // length of (non-zero) data.
  int64_t data_len;

  // To store aux data for sparse.
  // For CSR, indices stores the col index of non-zero elements.
  // For row sparse, indices store row index of rows which have non-zero elements.
  int64_t* indices;
  int64_t indices_len;

  // For CSR, indptr gives the start and end index of data for each row.
  // For row sparse, indptr is not used.
  int64_t* indptr = nullptr;
  int64_t indptr_len;

  void set(void *data_ptr, const int64_t* dims, int ndims, void *idx,
           int64_t num_idx, void *idx_ptr = nullptr, int64_t num_idx_ptr = 0);
};

struct MXTensor {
  MXTensor();
  MXTensor(const MXTensor& oth);
  MXTensor(void *data_ptr, std::vector<int64_t> shape, MXDType dtype,
           size_t vID, MXContext mx_ctx, MXStorageType stype = kDefaultStorage);

  void setTensor(void *dptr, MXDType type, const int64_t* dims, int ndims,
                 size_t vID, MXContext mx_ctx, MXStorageType storage_type);

  void setDLTensor();

  template<typename data_type>
  inline data_type* data() {
    return reinterpret_cast<data_type*>(data_ptr);
  }

  int64_t size() const;

  bool isSame(const MXTensor &oth) const;

  // For dense, data_ptr points to 1D flattened tensor data
  // For sparse, data_ptr points to MXSparse
  void *data_ptr;

  // shape is in [2,3,4] format to represent high-dim tensor
  std::vector<int64_t> shape;

  // type can only be MXDType enum types
  MXDType dtype;

  // version number updated if the tensor has changed since the last use by custom op
  size_t verID;

  // context of MXTensor representing which device the tensor data is located
  MXContext ctx;

  // corresponding DLTensor repr of MXTensor
  // easy way to reuse functions taking DLTensor
  DLTensor dltensor;

  // storage type
  MXStorageType stype;
};

typedef void* (*xpu_malloc_t)(void*, int);
typedef void (*sparse_malloc_t)(void*, int, int, int, void**, int64_t**, int64_t**);
typedef void (*nd_malloc_t)(const void* _ndarray_alloc, const int64_t* shapes, int num_shapes,
                            const char* dev_str, int dev_id, int dtype, const char* name,
                            int isArg, void** data);
#if defined(__NVCC__)
  typedef cudaStream_t mx_stream_t;
  typedef curandStatePhilox4_32_10_t mx_gpu_rand_t;
#else
  typedef void* mx_stream_t;
  typedef void* mx_gpu_rand_t;
#endif
typedef std::mt19937 mx_cpu_rand_t;

/* Each thread should generate random number unique sequence out of different states */
#define MX_NUM_CPU_RANDOM_STATES 1024
#define MX_NUM_GPU_RANDOM_STATES 32768

/* \brief Class to help allocate new args/aux params in graph passes */
class PassResource {
 public:
  PassResource(std::unordered_map<std::string, MXTensor>* new_args,
               std::unordered_map<std::string, MXTensor>* new_aux,
               nd_malloc_t nd_malloc, const void* nd_alloc);

  // allocate new arg param, adds to args map, returns newly allocated tensor
  MXTensor* alloc_arg(const std::string& name, const std::vector<int64_t>& shapes,
                      const MXContext &ctx, MXDType dtype) const;

  // allocate new aux param, adds to aux map, returns newly allocated tensor
  MXTensor* alloc_aux(const std::string& name, const std::vector<int64_t>& shapes,
                      const MXContext &ctx, MXDType dtype) const;

 private:
  std::unordered_map<std::string, MXTensor>* new_args_;
  std::unordered_map<std::string, MXTensor>* new_aux_;
  nd_malloc_t nd_malloc_;
  const void* nd_alloc_;
};

class OpResource {
 public:
  OpResource(xpu_malloc_t cpu_malloc_fp, void* cpu_alloc_fp,
             xpu_malloc_t gpu_malloc_fp, void* gpu_alloc_fp, void* stream,
             sparse_malloc_t sparse_malloc_fp, void* sparse_alloc_fp,
             void* rng_cpu_states, void* rng_gpu_states);

  void* alloc_cpu(int size) const;

  void* alloc_gpu(int size) const;

  inline mx_stream_t get_cuda_stream() const {
    return static_cast<mx_stream_t>(cuda_stream);
  }

  void alloc_sparse(MXSparse* sparse, int index, int indices_len, int indptr_len = 0) const;

  /* Access each state by states[id], but this id should be <= MX_NUM_CPU_RANDOM_STATES */
  mx_cpu_rand_t* get_cpu_rand_states() const;

  /* Access each state by states[id], but this id should be <= MX_NUM_GPU_RANDOM_STATES */
  /* Note that if you are using cpu build, it will return a nullptr */
  inline mx_gpu_rand_t* get_gpu_rand_states() const {
    return static_cast<mx_gpu_rand_t*>(rand_gpu_states);
  }

 private:
  xpu_malloc_t cpu_malloc, gpu_malloc;
  void *cpu_alloc, *gpu_alloc;
  void *cuda_stream;
  sparse_malloc_t sparse_malloc;
  void *sparse_alloc;
  void *rand_cpu_states, *rand_gpu_states;
};

#define MX_STR_SUBGRAPH_SYM_JSON "subgraph_sym_json"

#define MX_STR_DTYPE "__ext_dtype__"

#define MX_STR_SHAPE "__ext_shape__"

#define MX_STR_EXTRA_INPUTS "__ext_extra_inputs__"

/* \brief get shape value from list of shapes string
 *
 * Examples:
 *
 * getShapeAt("[[1]]", 0) returns "[1]"
 * getShapeAt("[[1],[2,3]]", 1) returns "[2,3]"
 */
std::string getShapeAt(const std::string& shape, unsigned index);

/* \brief get dtype value from list of dtypes string
 *
 * Examples:
 *
 * getDtypeAt("[1]", 0) returns "1"
 * getDtypeAt("[1,2]", 1) returns "2" 
 */
std::string getDtypeAt(const std::string& dtype, unsigned index);

enum JsonType {ERR, STR, NUM, LIST, MAP};

struct JsonVal {
  JsonVal();  // default constructor
  // construct a JSON object by type
  explicit JsonVal(JsonType t);
  // construct a string JSON object
  explicit JsonVal(std::string s);
  // construct a number JSON object
  explicit JsonVal(int n);
  // complex constructor
  JsonVal(JsonType t, int n, std::string s);
  bool operator<(const JsonVal &o) const;

  // convert JSON object back to JSON-compatible string
  std::string dump() const;

  // convert JSON-compatible string to JSON object
  static JsonVal parse(const std::string& json);

  // parse a string JSON object
  static JsonVal parse_string(const std::string& json, unsigned int* idx);

  // parse a number JSON object
  static JsonVal parse_num(const std::string& json, unsigned int* idx);

  // parse a list of JSON objects
  static JsonVal parse_list(const std::string& json, unsigned int* idx);

  // parse a map of JSON objects
  static JsonVal parse_map(const std::string& json, unsigned int* idx);

  // generic parse function
  static JsonVal parse(const std::string& json, unsigned int *idx);

  // debug function to convert data structure to a debugstring
  std::string toString() const;

  JsonType type;
  int num;
  std::string str;
  std::vector<JsonVal> list;
  std::map<JsonVal, JsonVal> map;
};

class Node;
class Graph;

// Representation of an input/output to a node
struct NodeEntry {
  Node* node;  // other node thats producing/consuming inputs/outputs
  int entry;  // entry index from other node (ie. output index from producing node)
};

// Representation of a node in the graph
class Node {
 public:
  Node();

  // internally set passResource to enable tensor allocation for graph passes
  void _setPassResource(PassResource* res_);

  /* \brief allocate an arg tensor for this node */
  void alloc_arg(const std::vector<int64_t>& shapes,
                 const MXContext &ctx, MXDType dtype);

  /* \brief allocate an aux tensor for this node */
  void alloc_aux(const std::vector<int64_t>& shapes,
                 const MXContext &ctx, MXDType dtype);

  std::string op;  // operator name (ie. Convolution)
  std::string name;  // unique node name (ie. conv_0 or conv_1)
  MXTensor* tensor;  // tensor data for input nodes
  std::vector<NodeEntry> inputs;  // set of inputs to the node
  std::vector<NodeEntry> outputs;  // set of outputs from the node
  std::vector<Graph*> subgraphs;  // set of subgraphs within this node
  std::unordered_map<std::string, std::string> attrs;  // node attributes

 private:
  PassResource* res;
};

// Representation of the graph
class Graph {
 public:
  Graph();

  /* \brief deleted nodes when deleting the graph */
  ~Graph();

  /* \brief create a graph object from an unparsed string */
  static Graph* fromString(const std::string& json);

  /* \brief create a graph object from a parsed JSON object */
  static Graph* fromJson(JsonVal val);

  /* \brief convert graph object back to JSON object */
  JsonVal toJson() const;

  /* \brief convert graph object to JSON string */
  std::string toString() const;

  /* \brief visits a node "n" */
  void _dfs_util(Node* n, std::unordered_set<Node*>* to_visit,
                 std::function<void(Node*)> handler) const;

  /* \brief post-order DFS graph traversal */
  void DFS(std::function<void(Node*)> handler) const;

  /* \brief sort graph nodes in topological order */
  std::vector<Node*> topological_sort() const;

  /* \brief print out graph details */
  void print(int indent = 0) const;

  /* \brief add a new node to this graph */
  Node* addNode(const std::string& name, const std::string& op);

  /* \brief get node at index in graph */
  Node* getNode(size_t idx);

  /* \brief get const node at index in const graph */
  const Node* getNode(size_t idx) const;

  /* \brief get attribute on graph */
  const JsonVal& getAttr(const std::string& key) const;

  /* \brief get number of nodes in the graph */
  size_t size() const;

  // internally set passResource to enable tensor allocation for graph passes
  void _setPassResource(PassResource* res_);

  // internally set arg/aux params when available
  void _setParams(std::unordered_map<std::string, mxnet::ext::MXTensor>* args,
                  std::unordered_map<std::string, mxnet::ext::MXTensor>* aux);

  std::vector<Node*> inputs;
  std::vector<NodeEntry> outputs;
  std::map<std::string, JsonVal> attrs;

 private:
  std::vector<Node*> nodes;
  PassResource* res;
};

/* \brief An abstract class for library authors creating custom
 * partitioners. Optional, can just implement supportedOps instead
 */
class CustomOpSelector {
 public:
  /* \brief Select a node to include in subgraph, return true to include node
   * nodeID - index of node in graph
   */
  virtual bool Select(int nodeID) = 0;
  /* \brief Select an input node from current node to include in subgraph
   * return true to include node
   * nodeID - index of node in graph
   * input_nodeID - index of input node in graph
   */
  virtual bool SelectInput(int nodeID, int input_nodeID) = 0;
  /* \brief Select an output node from current node to include in subgraph
   * return true to include node
   * nodeID - index of node in graph
   * output_nodeID - index of output node in graph
   */
  virtual bool SelectOutput(int nodeID, int output_nodeID) = 0;
  /* \brief Review nodes to include in subgraph
   * return set of candidate nodes to keep in subgraph
   * candidates - indices of nodes to include in subgraph
   * keep - indices of nodes to keep in subgraph
   */
  virtual void Filter(const std::vector<int>& candidates,
                      std::vector<int>* keep) {
    keep->insert(keep->end(), candidates.begin(), candidates.end());
  }
  /* \brief Reset any selector state, called after growing subgraph, before filter
   * Called after finished calling SelectInput/SelectOutput and growing subgraph
   */
  virtual void Reset() {}
};

class CustomStatefulOp {
 public:
  CustomStatefulOp();
  virtual ~CustomStatefulOp();

  template<class A, typename ...Ts>
  static CustomStatefulOp* create(Ts...args) {
    CustomStatefulOp* op = new A(args...);
    op->created = true;
    return op;
  }

  bool wasCreated() { return created; }

  virtual MXReturnValue Forward(std::vector<MXTensor>* inputs,
                                std::vector<MXTensor>* outputs,
                                const OpResource& op_res) = 0;
  virtual MXReturnValue Backward(std::vector<MXTensor>* inputs,
                                 std::vector<MXTensor>* outputs,
                                 const OpResource& op_res) {
    MX_ERROR_MSG << "Error! Operator does not support backward" << std::endl;
    return MX_FAIL;
  }

  bool ignore_warn;

 private:
  bool created;
};

typedef MXReturnValue (*fcomp_t)(const std::unordered_map<std::string,
                                 std::string>& attributes,
                                 std::vector<MXTensor>* inputs,
                                 std::vector<MXTensor>* outputs,
                                 const OpResource& res);
typedef MXReturnValue (*parseAttrs_t)(const std::unordered_map<std::string,
                                      std::string>& attributes,
                                      int* num_inputs, int* num_outputs);
typedef MXReturnValue (*inferType_t)(const std::unordered_map<std::string,
                                     std::string>& attributes,
                                     std::vector<int>* in_types,
                                     std::vector<int>* out_types);
typedef MXReturnValue (*inferSType_t)(const std::unordered_map<std::string,
                                      std::string>& attributes,
                                      std::vector<int>* in_storage_types,
                                      std::vector<int>* out_storage_types);
typedef MXReturnValue (*inferShape_t)(const std::unordered_map<std::string,
                                      std::string>& attributes,
                                      std::vector<std::vector<unsigned int> >* in_shapes,
                                      std::vector<std::vector<unsigned int> >* out_shapes);
typedef MXReturnValue (*mutateInputs_t)(const std::unordered_map<std::string,
                                        std::string>& attributes,
                                        std::vector<int>* input_indices);
typedef MXReturnValue (*createOpState_t)(const std::unordered_map<std::string,
                                         std::string>& attributes,
                                         const MXContext& ctx,
                                         const std::vector<std::vector<unsigned int> >& in_shapes,
                                         const std::vector<int> in_types,
                                         CustomStatefulOp**);

class CustomOp {
 public:
  explicit CustomOp(const char* op_name);

  CustomOp& setForward(fcomp_t fcomp, const char* ctx);

  CustomOp& setBackward(fcomp_t fgrad, const char* ctx);

  CustomOp& setParseAttrs(parseAttrs_t func);

  CustomOp& setInferType(inferType_t func);

  CustomOp& setInferSType(inferSType_t func);

  CustomOp& setInferShape(inferShape_t func);

  CustomOp& setMutateInputs(mutateInputs_t func);

  CustomOp& setCreateOpState(createOpState_t func, const char* ctx);

  CustomOp& setIsSubgraphOp();

  void mapToVector();

  const char* name;

  parseAttrs_t parse_attrs;
  inferType_t infer_type;
  inferSType_t infer_storage_type;
  inferShape_t infer_shape;
  mutateInputs_t mutate_inputs;
  bool isSGop;

  std::vector<const char*> forward_ctx_cstr, backward_ctx_cstr, create_op_ctx_cstr;
  std::vector<fcomp_t> forward_fp, backward_fp;
  std::vector<createOpState_t> create_op_fp;

 private:
  void raiseDuplicateContextError();

  std::unordered_map<const char*, fcomp_t> forward_ctx_map, backward_ctx_map;
  std::unordered_map<const char*, createOpState_t> create_op_ctx_map;
};

typedef MXReturnValue (*graphPass_t)(mxnet::ext::Graph* graph,
                                     const std::unordered_map<std::string, std::string>& options);

class CustomPass {
 public:
  CustomPass();
  explicit CustomPass(const char* pass_name);

  CustomPass& setBody(graphPass_t fn);

  const char* name;
  graphPass_t pass;
};

typedef MXReturnValue (*supportedOps_t)(const mxnet::ext::Graph *graph, std::vector<int>* ids,
                                        const std::unordered_map<std::string,
                                                                 std::string>& options);
typedef MXReturnValue (*createSelector_t)(const mxnet::ext::Graph *graph,
                                          CustomOpSelector** sel_inst,
                                          const std::unordered_map<std::string,
                                                                   std::string>& options);
typedef MXReturnValue (*reviewSubgraph_t)(const mxnet::ext::Graph *subgraph, int subgraph_id,
                                          bool* accept,
                                          const std::unordered_map<std::string,
                                                                   std::string>& options,
                                          std::unordered_map<std::string,
                                                             std::string>* attrs);

class CustomPartitioner {
 public:
  CustomPartitioner();

  explicit CustomPartitioner(const char* backend_name);

  CustomPartitioner& addStrategy(const char* prop_name,
                                 const char* sg_name);

  CustomPartitioner& setSupportedOps(const char* prop_name, supportedOps_t fn);

  CustomPartitioner& setCreateSelector(const char* prop_name, createSelector_t fn);

  CustomPartitioner& setReviewSubgraph(const char* prop_name, reviewSubgraph_t fn);

  supportedOps_t getSupportedOps(int stg_id);

  createSelector_t getCreateSelector(int stg_id);

  reviewSubgraph_t getReviewSubgraph(int stg_id);

  const char* name;
  std::map<std::string, supportedOps_t> supported_map;
  std::map<std::string, createSelector_t> selector_map;
  std::map<std::string, reviewSubgraph_t> review_map;
  std::vector<const char*> strategies;
  std::vector<const char*> op_names;
};

template <class T>
class Registry {
 public:
  static Registry* get() PRIVATE_SYMBOL {
    static Registry inst;
    return &inst;
  }
  T& add(const char* name) {
    T *entry = new T(name);
    entries.push_back(entry);
    return *entry;
  }
  int size() {
    return entries.size();
  }
  T& get(int idx) {
    return *(entries.at(idx));
  }

 private:
  Registry() {}
  ~Registry() {}
  std::vector<T*> entries;
};

#define MX_STR_CONCAT_(__a, __b) __a ## __b
#define MX_STR_CONCAT(__a, __b) MX_STR_CONCAT_(__a, __b)

#define MX_STRINGIFY(x) #x
#define MX_TOSTRING(x) MX_STRINGIFY(x)

#define MX_REGISTER_NAME_(Name) MXNet ## _CustomOp ## _ ## Name
#define MX_REGISTER_DEF_(Name) mxnet::ext::CustomOp MX_REGISTER_NAME_(Name)

#define MX_REGISTER_PROP_NAME_(Name) MXNet ## _CustomSubProp ## _ ## Name
#define MX_REGISTER_PROP_DEF_(Name) mxnet::ext::CustomPartitioner MX_REGISTER_PROP_NAME_(Name)

#define MX_REGISTER_PASS_NAME_(Name) MXNet ## _CustomPass ## _ ## Name
#define MX_REGISTER_PASS_DEF_(Name) mxnet::ext::CustomPass MX_REGISTER_PASS_NAME_(Name)

#define REGISTER_OP(Name) MX_STR_CONCAT(MX_REGISTER_DEF_(Name), __COUNTER__) = \
    mxnet::ext::Registry<mxnet::ext::CustomOp>::get()->add(MX_TOSTRING(Name))

#define REGISTER_PARTITIONER(Name) \
  MX_STR_CONCAT(MX_REGISTER_PROP_DEF_(Name), __COUNTER__) = \
    mxnet::ext::Registry<mxnet::ext::CustomPartitioner>::get()->add(MX_TOSTRING(Name))

#define REGISTER_PASS(Name) \
  MX_STR_CONCAT(MX_REGISTER_PASS_DEF_(Name), __COUNTER__) = \
    mxnet::ext::Registry<mxnet::ext::CustomPass>::get()->add(MX_TOSTRING(Name))

/* -------------- BELOW ARE CTYPE FUNCTIONS PROTOTYPES --------------- */

#define MXLIB_OPREGSIZE_STR "_opRegSize"
typedef int (*opRegSize_t)(void);

#define MXLIB_OPREGGET_STR "_opRegGet"
typedef int (*opRegGet_t)(int idx, const char** name, int *isSGop,
                          const char*** forward_ctx, mxnet::ext::fcomp_t** forward_fp,
                          int* forward_count, const char*** backward_ctx,
                          mxnet::ext::fcomp_t** backward_fp, int* backward_count,
                          const char*** create_op_ctx, mxnet::ext::createOpState_t** create_op_fp,
                          int* create_op_count, mxnet::ext::parseAttrs_t* parse,
                          mxnet::ext::inferType_t* type, mxnet::ext::inferSType_t* stype,
                          mxnet::ext::inferShape_t* shape, mxnet::ext::mutateInputs_t* mutate);

#define MXLIB_OPCALLFREE_STR "_opCallFree"
typedef int (*opCallFree_t)(void* ptr);

#define MXLIB_OPCALLPARSEATTRS_STR "_opCallParseAttrs"
typedef int (*opCallParseAttrs_t)(parseAttrs_t parseAttrs, const char* const* keys,
                                  const char* const* vals, int num,
                                  int* num_in, int* num_out);

#define MXLIB_OPCALLINFERSHAPE_STR "_opCallInferShape"
typedef int (*opCallInferShape_t)(inferShape_t inferShape, const char* const* keys,
                                  const char* const* vals, int num,
                                  unsigned int** inshapes, int* indims, int num_in,
                                  unsigned int*** mod_inshapes, int** mod_indims,
                                  unsigned int*** outshapes, int** outdims, int num_out);

#define MXLIB_OPCALLINFERTYPE_STR "_opCallInferType"
typedef int (*opCallInferType_t)(inferType_t inferType, const char* const* keys,
                                 const char* const* vals, int num,
                                 int* intypes, int num_in, int* outtypes, int num_out);

#define MXLIB_OPCALLINFERSTYPE_STR "_opCallInferSType"
typedef int (*opCallInferSType_t)(inferSType_t inferSType, const char* const* keys,
                                 const char* const* vals, int num,
                                 int* intypes, int num_in, int* outtypes, int num_out);

#define MXLIB_OPCALLFCOMP_STR "_opCallFCompute"
typedef int (*opCallFComp_t)(fcomp_t fcomp, const char* const* keys,
                             const char* const* vals, int num,
                             const int64_t** inshapes, int* indims,
                             void** indata, int* intypes,
                             size_t* inIDs, const char** indev_type,
                             int* indev_id, int num_in,
                             const int64_t** outshapes, int* outdims,
                             void** outdata, int* outtypes,
                             size_t* outIDs, const char** outdev_type,
                             int* outdev_id, int num_out,
                             xpu_malloc_t cpu_malloc, void* cpu_alloc,
                             xpu_malloc_t gpu_malloc, void* gpu_alloc, void* cuda_stream,
                             sparse_malloc_t sparse_malloc, void* sparse_alloc,
                             int* instypes, int* outstypes,
                             void** in_indices, void** out_indices,
                             void** in_indptr, void** out_indptr,
                             int64_t* in_indices_shapes, int64_t* out_indices_shapes,
                             int64_t* in_indptr_shapes, int64_t* out_indptr_shapes,
                             void* rng_cpu_states, void* rng_gpu_states);

#define MXLIB_OPCALLMUTATEINPUTS_STR "_opCallMutateInputs"
typedef int (*opCallMutateInputs_t)(mutateInputs_t mutate, const char* const* keys,
                                    const char* const* vals, int num,
                                    int** mutate_indices, int* indices_size);

#define MXLIB_OPCALLCREATEOPSTATE_STR "_opCallCreateOpState"
typedef int (*opCallCreateOpState_t)(createOpState_t create_op, const char* const* keys,
                                     const char* const* vals, int num, const char* dev_type,
                                     int dev_id, unsigned int** inshapes, int* indims,
                                     int num_in, const int* intypes, void** state_op);

#define MXLIB_OPCALLDESTROYOPSTATE_STR "_opCallDestroyOpState"
typedef int (*opCallDestroyOpState_t)(void* state_op);

#define MXLIB_OPCALLFSTATEFULCOMP_STR "_opCallFStatefulCompute"
typedef int (*opCallFStatefulComp_t)(int is_forward, void* state_op,
                                     const int64_t** inshapes, int* indims,
                                     void** indata, int* intypes,
                                     size_t* inIDs, const char** indev_type,
                                     int* indev_id, int num_in,
                                     const int64_t** outshapes, int* outdims,
                                     void** outdata, int* outtypes,
                                     size_t* outIDs, const char** outdev_type,
                                     int* outdev_id, int num_out,
                                     xpu_malloc_t cpu_malloc, void* cpu_alloc,
                                     xpu_malloc_t gpu_malloc, void* gpu_alloc, void* stream,
                                     sparse_malloc_t sparse_malloc, void* sparse_alloc,
                                     int* instypes, int* outstypes,
                                     void** in_indices, void** out_indices,
                                     void** in_indptr, void** out_indptr,
                                     int64_t* in_indices_shapes, int64_t* out_indices_shapes,
                                     int64_t* in_indptr_shapes, int64_t* out_indptr_shapes,
                                     void* rng_cpu_states, void* rng_gpu_states);

#define MXLIB_PARTREGSIZE_STR "_partRegSize"
typedef int (*partRegSize_t)(void);

#define MXLIB_PARTREGGETCOUNT_STR "_partRegGetCount"
typedef int (*partRegGetCount_t)(int idx, const char** name);

#define MXLIB_PARTREGGET_STR "_partRegGet"
typedef void (*partRegGet_t)(int part_idx, int stg_idx, const char** strategy,
                             supportedOps_t* supportedOps, createSelector_t* createSelector,
                             reviewSubgraph_t* reviewSubgraph, const char** op_name);

#define MXLIB_PARTCALLSUPPORTEDOPS_STR "_partCallSupportedOps"
typedef int (*partCallSupportedOps_t)(supportedOps_t supportedOps, const char *json,
                                      int num_ids, int *ids, const char* const* opt_keys,
                                      const char* const* opt_vals, int num_opts);

#define MXLIB_PARTCALLCREATESELECTOR_STR "_partCallCreateSelector"
typedef int (*partCallCreateSelector_t)(createSelector_t createSelector, const char *json,
                                        void** selector, const char* const* opt_keys,
                                        const char* const* opt_vals, int num_opts);

#define MXLIB_PARTCALLSELECT_STR "_partCallSelect"
typedef void (*partCallSelect_t)(void* sel_inst, int nodeID, int* selected);

#define MXLIB_PARTCALLSELECTINPUT_STR "_partCallSelectInput"
typedef void (*partCallSelectInput_t)(void* sel_inst, int nodeID, int input_nodeID,
                                  int* selected);

#define MXLIB_PARTCALLSELECTOUTPUT_STR "_partCallSelectOutput"
typedef void (*partCallSelectOutput_t)(void* sel_inst, int nodeID, int output_nodeID,
                                   int* selected);

#define MXLIB_PARTCALLFILTER_STR "_partCallFilter"
typedef void (*partCallFilter_t)(void* sel_inst, int* candidates, int num_candidates,
                             int** keep, int* num_keep);

#define MXLIB_PARTCALLRESET_STR "_partCallReset"
typedef void (*partCallReset_t)(void* sel_inst);

#define MXLIB_PARTCALLREVIEWSUBGRAPH_STR "_partCallReviewSubgraph"
typedef int (*partCallReviewSubgraph_t)(reviewSubgraph_t reviewSubgraph, 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,
                                        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,
                                        const char* const* aux_names, int num_aux,
                                        void* const* aux_data, const int64_t* const* aux_shapes,
                                        const int* aux_dims, const int* aux_types,
                                        const size_t* aux_IDs, const char* const* aux_dev_type,
                                        const int* aux_dev_id);

#define MXLIB_PASSREGSIZE_STR "_passRegSize"
typedef int (*passRegSize_t)(void);

#define MXLIB_PASSREGGET_STR "_passRegGet"
typedef void (*passRegGet_t)(int pass_idx, graphPass_t* graphPass, const char** pass_name);

#define MXLIB_PASSCALLGRAPHPASS_STR "_passCallGraphPass"
typedef int (*passCallGraphPass_t)(graphPass_t graphPass, const char *in_graph,
                                   char** out_graph, const char* const* opt_keys,
                                   const char* const* opt_vals, int num_opts,
                                   const char* pass_name, 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,
                                   const char* const* aux_names, int num_aux,
                                   void* const* aux_data, const int64_t* const* aux_shapes,
                                   const int* aux_dims, const int* aux_types,
                                   const size_t* aux_IDs, const char* const* aux_dev_type,
                                   const int* aux_dev_id, nd_malloc_t nd_malloc,
                                   const void* nd_alloc);

#define MXLIB_INITIALIZE_STR "initialize"
typedef int (*initialize_t)(int version);

#define MXLIB_OPVERSION_STR "_opVersion"
typedef int (*opVersion_t)();

#define MXLIB_MSGSIZE_STR "_msgSize"
typedef int (*msgSize_t)(void);

#define MXLIB_MSGGET_STR "_msgGet"
typedef int (*msgGet_t)(int idx, const char** msg);

class CustomStatefulOpWrapper {
 public:
  ~CustomStatefulOpWrapper();
  explicit CustomStatefulOpWrapper(CustomStatefulOp* inst, opCallDestroyOpState_t destroy)
    : instance(inst), destroy_(destroy) {}
  CustomStatefulOp* get_instance() { return instance; }
 private:
  CustomStatefulOp* instance;
  opCallDestroyOpState_t destroy_;
};

#if defined(_WIN32) || defined(_WIN64) || defined(__WINDOWS__)
#define MX_INT_RET  __declspec(dllexport) int __cdecl
#define MX_VOID_RET __declspec(dllexport) void __cdecl
#else
#define MX_INT_RET  int
#define MX_VOID_RET void
#endif

}  // namespace ext
}  // namespace mxnet

extern "C" {
  MX_INT_RET _opVersion();

  MX_INT_RET _opRegSize();

  MX_VOID_RET _opRegGet(int idx, const char** name, int *isSGop,
                        const char*** forward_ctx, mxnet::ext::fcomp_t** forward_fp,
                        int* forward_count, const char*** backward_ctx,
                        mxnet::ext::fcomp_t** backward_fp, int* backward_count,
                        const char*** create_op_ctx, mxnet::ext::createOpState_t** create_op_fp,
                        int* create_op_count, mxnet::ext::parseAttrs_t* parse,
                        mxnet::ext::inferType_t* type, mxnet::ext::inferSType_t* stype,
                        mxnet::ext::inferShape_t* shape, mxnet::ext::mutateInputs_t* mutate);

  MX_VOID_RET _opCallFree(void* ptr);

  MX_INT_RET _opCallParseAttrs(mxnet::ext::parseAttrs_t parseAttrs, const char* const* keys,
                               const char* const* vals, int num,
                               int* num_in, int* num_out);

  MX_INT_RET _opCallInferShape(mxnet::ext::inferShape_t inferShape, const char* const* keys,
                               const char* const* vals, int num,
                               unsigned int** inshapes, int* indims, int num_in,
                               unsigned int*** mod_inshapes, int** mod_indims,
                               unsigned int*** outshapes, int** outdims, int num_out);

  MX_INT_RET _opCallInferType(mxnet::ext::inferType_t inferType, const char* const* keys,
                              const char* const* vals, int num,
                              int* intypes, int num_in, int* outtypes, int num_out);

  MX_INT_RET _opCallInferSType(mxnet::ext::inferSType_t inferSType, const char* const* keys,
                               const char* const* vals, int num,
                               int* instypes, int num_in, int* outstypes, int num_out);

  MX_INT_RET _opCallFCompute(mxnet::ext::fcomp_t fcomp, const char* const* keys,
                             const char* const* vals,
                             int num, const int64_t** inshapes, int* indims, void** indata,
                             int* intypes, size_t* inIDs, const char** indev_type, int* indev_id,
                             int num_in, const int64_t** outshapes, int* outdims, void** outdata,
                             int* outtypes, size_t* outIDs, const char** outdev_type,
                             int* outdev_id, int num_out, mxnet::ext::xpu_malloc_t cpu_malloc,
                             void* cpu_alloc,
                             mxnet::ext::xpu_malloc_t gpu_malloc, void* gpu_alloc,
                             void* cuda_stream,
                             mxnet::ext::sparse_malloc_t sparse_malloc, void* sparse_alloc,
                             int* instypes, int* outstypes, void** in_indices, void** out_indices,
                             void** in_indptr, void** out_indptr,
                             int64_t* in_indices_shapes, int64_t* out_indices_shapes,
                             int64_t* in_indptr_shapes, int64_t* out_indptr_shapes,
                             void* rng_cpu_states, void* rng_gpu_states);

  MX_INT_RET _opCallMutateInputs(mxnet::ext::mutateInputs_t mutate, const char* const* keys,
                                 const char* const* vals, int num,
                                 int** mutate_indices, int* indices_size);

  MX_INT_RET _opCallCreateOpState(mxnet::ext::createOpState_t create_op, const char* const* keys,
                                  const char* const* vals, int num, const char* dev_type,
                                  int dev_id, unsigned int** inshapes, int* indims,
                                  int num_in, const int* intypes, void** state_op);

  MX_VOID_RET _opCallDestroyOpState(void* state_op);

  MX_INT_RET _opCallFStatefulCompute(int is_forward, void* state_op, const int64_t** inshapes,
                                     int* indims, void** indata, int* intypes, size_t* inIDs,
                                     const char** indev_type, int* indev_id, int num_in,
                                     const int64_t** outshapes, int* outdims, void** outdata,
                                     int* outtypes, size_t* outIDs, const char** outdev_type,
                                     int* outdev_id, int num_out,
                                     mxnet::ext::xpu_malloc_t cpu_malloc,
                                     void* cpu_alloc, mxnet::ext::xpu_malloc_t gpu_malloc,
                                     void* gpu_alloc,
                                     void* stream, mxnet::ext::sparse_malloc_t sparse_malloc,
                                     void* sparse_alloc, int* instypes, int* outstypes,
                                     void** in_indices, void** out_indices, void** in_indptr,
                                     void** out_indptr, int64_t* in_indices_shapes,
                                     int64_t* out_indices_shapes, int64_t* in_indptr_shapes,
                                     int64_t* out_indptr_shapes,
                                     void* rng_cpu_states, void* rng_gpu_states);

  MX_INT_RET _partRegSize();

  /* returns number of strategies registered for partitioner
   * at specified index */
  MX_INT_RET _partRegGetCount(int idx, const char** name);

  MX_VOID_RET _partRegGet(int part_idx, int stg_idx, const char** strategy,
                          mxnet::ext::supportedOps_t* supportedOps,
                          mxnet::ext::createSelector_t* createSelector,
                          mxnet::ext::reviewSubgraph_t* reviewSubgraph, const char** op_name);

  MX_INT_RET _partCallSupportedOps(mxnet::ext::supportedOps_t supportedOps, const char *json,
                                   int num_ids, int *ids, const char* const* opt_keys,
                                   const char* const* opt_vals, int num_opts);

  MX_INT_RET _partCallCreateSelector(mxnet::ext::createSelector_t createSelector, const char *json,
                                     void** selector, const char* const* opt_keys,
                                     const char* const* opt_vals, int num_opts);

  MX_VOID_RET _partCallSelect(void* sel_inst, int nodeID, int* selected);

  MX_VOID_RET _partCallSelectInput(void* sel_inst, int nodeID,
                                   int input_nodeID, int* selected);

  MX_VOID_RET _partCallSelectOutput(void* sel_inst, int nodeID,
                                    int output_nodeID, int* selected);

  MX_VOID_RET _partCallFilter(void* sel_inst, int* candidates, int num_candidates,
                              int** keep, int* num_keep);

  MX_VOID_RET _partCallReset(void* sel_inst);

  MX_INT_RET _partCallReviewSubgraph(mxnet::ext::reviewSubgraph_t reviewSubgraph, 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,
                                     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,
                                     const char* const* aux_names, int num_aux,
                                     void* const* aux_data, const int64_t* const* aux_shapes,
                                     const int* aux_dims, const int* aux_types,
                                     const size_t* aux_IDs, const char* const* aux_dev_type,
                                     const int* aux_dev_id);

  MX_INT_RET _passRegSize();

  MX_VOID_RET _passRegGet(int pass_idx, mxnet::ext::graphPass_t* graphPass,
                          const char** pass_name);

  MX_INT_RET _passCallGraphPass(mxnet::ext::graphPass_t graphPass, const char *json,
                                char** out_graph, const char* const* opt_keys,
                                const char* const* opt_vals, int num_opts,
                                const char* pass_name, 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, const char* const* aux_names, int num_aux,
                                void* const* aux_data, const int64_t* const* aux_shapes,
                                const int* aux_dims, const int* aux_types,
                                const size_t* aux_IDs, const char* const* aux_dev_type,
                                const int* aux_dev_id, mxnet::ext::nd_malloc_t nd_malloc,
                                const void* nd_alloc);

#if defined(_WIN32) || defined(_WIN64) || defined(__WINDOWS__)
  __declspec(dllexport) mxnet::ext::MXReturnValue __cdecl
#else
  mxnet::ext::MXReturnValue
#endif
  initialize(int version);

  MX_INT_RET _msgSize();

  MX_VOID_RET _msgGet(int idx, const char** msg);
}  // extern "C"

#endif  // MXNET_LIB_API_H_