Adding sparse support to MXTensor for custom operators

include/mxnet/lib_api.h

@@ -214,6 +214,18 @@ enum MXDType {
   kUNSET = 100,
 };
 
+/*
+ * MXTensor storage type.
+ */
+enum MXStorageType {
+  // dense
+  kDefaultStorage = 0,
+  // row sparse 
+  kRowSparseStorage = 1,
+  // csr
+  kCSRStorage = 2,
+};
+
 /*!
  * \brief Context info passing from MXNet OpContext
  * dev_type is string repr of supported context, currently only "cpu" and "gpu"
@@ -229,20 +241,54 @@ enum MXReturnValue {
   MX_SUCCESS = 1,
 };
 
+struct ChunkDense {
+  // Pointer to data.
+  void *data{nullptr};
+  // Size of data in bytes.
+  size_t dataSize{0};
+  // shape of data.
+  std::vector<int64_t> shape;
+  // Context of data.
+  // MXContext ctx;
+};
+
+struct ChunkSparse {
+  // Pointer to data.
+  void *data{nullptr};
+  // Size of data in bytes.
+  size_t dataSize{0};
+  // length  of data.
+  int64_t data_lens;
+
+  // To store aux data for sparse.
+  // for row_sparse, aux_data[0] = indices
+  // for csr, aux_data[0] = indptr, aux_data[1] = indices
+  std::vector<std::vector<int64_t>> aux_data;
+
+  // Lens of the aux_data.
+  // for row_sparse, aux_lens[0] = len(indices)
+  // for csr, aux_lens[0] = len(indptr), aux_lens[1] = len(indices)
+  std::vector<int64_t> aux_lens;
+  // Context of data.
+  // MXContext ctx;
+};
+
 /*!
  * \brief Tensor data structure used by custom operator
  */
 struct MXTensor {
-  MXTensor() : data_ptr(NULL), dtype(kUNSET), verID(0) {}
+  MXTensor() : data_ptr(nullptr), dtype(kUNSET), verID(0), stype(kDefaultStorage) {}
 
+  // Construtor for dense.
   MXTensor(void *data_ptr, const std::vector<int64_t> &shape, MXDType dtype,
-           size_t vID, MXContext mx_ctx)
-  : data_ptr(data_ptr), shape(shape), dtype(dtype), verID(vID), ctx(mx_ctx) {}
+           size_t vID, MXContext mx_ctx, MXStorageType stype = kDefaultStorage)
+  : data_ptr(data_ptr), shape(shape), dtype(dtype), verID(vID), ctx(mx_ctx), stype(stype) {}
 
   /*! \brief populate internal tensor fields */
+  // To do: solve for CSR and row sparse.
   void setTensor(void *dptr, MXDType type, const int64_t* dims, int ndims,
-                 size_t vID, MXContext mx_ctx) {
-    data_ptr = dptr; dtype = type; verID = vID; ctx = mx_ctx;
+                 size_t vID, MXContext mx_ctx, MXStorageType stype = kDefaultStorage) {
+    data_ptr = dptr; dtype = type; verID = vID; ctx = mx_ctx; stype = stype;
     shape.clear();
     for (int j = 0; j < ndims; j++) {
       shape.push_back(dims[j]);
@@ -335,11 +381,15 @@ struct MXTensor {
            verID == oth.verID &&
            ctx.dev_type == oth.ctx.dev_type &&
            ctx.dev_id == oth.ctx.dev_id &&
-           shape == oth.shape;
+           shape == oth.shape && 
+           stype == oth.stype;
   }
 
-  // data is flatten 1D repr of tensor, elements are in continuous memory
-  // user can access each element using the shape of tensor
+  /*! \brief get MXTensors storage type*/
+  inline MXStorageType getStorageType() { return stype; } 
+
+  // For dense, data_ptr points to ChunkDense.
+  // For sparse, data_ptr points to ChunkSparse.
   void *data_ptr;
 
   // shape is in [2,3,4] format to represent high-dim tensor
@@ -357,6 +407,9 @@ struct MXTensor {
   // corresponding DLTensor repr of MXTensor
   // easy way to reuse functions taking DLTensor
   DLTensor dltensor;
+
+  // storage type
+  MXStorageType stype;
 };
 
 /*! \brief resource malloc function to allocate memory inside Forward/Backward functions */