Fix ndarray aux data issue

include/mxnet/c_api.h

@@ -352,6 +352,17 @@ MXNET_DLL int MXNDArraySyncCopyFromCPU(NDArrayHandle handle,
 MXNET_DLL int MXNDArraySyncCopyToCPU(NDArrayHandle handle,
                                      void *data,
                                      size_t size);
+/*!
+ * \brief Copy src.data() to dst.data() if i = -1, else dst.aux_data(i) if i >= 0
+ * This function blocks. Do not use it in performance critical code.
+ * \param handle_dst handle of a dst ndarray whose data/aux_data has been allocated
+ * \param handle_src handle of a src ndarray which has default storage type
+ * \param i dst data blob indicator
+ */
+MXNET_DLL int MXNDArraySyncCopyFromNDArray(NDArrayHandle handle_dst,
+                                           const NDArrayHandle handle_src,
+                                           const int i);
+
 /*!
  * \brief Wait until all the pending writes with respect NDArray are finished.
  *  Always call this before read data out synchronizely.
@@ -458,12 +469,20 @@ MXNET_DLL int MXNDArrayGetAuxType(NDArrayHandle handle,
                                   mx_uint i,
                                   int *out_type);
 
-// Get the ith aux data blob wrapped in an NDArray
+/*!
+ * \brief Get a deep copy of the ith aux data blob
+ * in the form of an NDArray of default storage type.
+ * This function blocks. Do not use it in performance critical code.
+ */
 MXNET_DLL int MXNDArrayGetAuxNDArray(NDArrayHandle handle,
                                      mx_uint i,
                                      NDArrayHandle *out);
 
-// Get the data blob wrapped in an NDArray
+/*!
+ * \brief Get a deep copy of the data blob
+ * in the form of an NDArray of default storage type.
+ * This function blocks. Do not use it in performance critical code.
+ */
 MXNET_DLL int MXNDArrayGetDataNDArray(NDArrayHandle handle,
                                       NDArrayHandle *out);
 /*!

include/mxnet/ndarray.h

@@ -146,11 +146,9 @@ class NDArray {
    *  make sure the memory region is available through out the life of NDArray
    * \param data the memory content of static data
    * \param dev_id the device id this tensor sits at
-   * \param shared_var the same var handle shared with others.
-            It will not be deleted during destruction.
    */
-  NDArray(const TBlob &data, int dev_id, Engine::VarHandle shared_var = nullptr)
-      : ptr_(std::make_shared<Chunk>(data, dev_id, shared_var)), shape_(data.shape_),
+  NDArray(const TBlob &data, int dev_id)
+      : ptr_(std::make_shared<Chunk>(data, dev_id)), shape_(data.shape_),
         dtype_(data.type_flag_), entry_({nullptr, 0, 0}) {
 #if MKL_EXPERIMENTAL == 1
     Mkl_mem_ = std::make_shared<MKLMemHolder>();
@@ -166,8 +164,6 @@ class NDArray {
    * \param data the memory content of static data
    * \param aux_data the memory content of static aux data
    * \param dev_id the device id this tensor sits at
-   * \param shared_var the same var handle shared with others.
-            It will not be deleted during destruction.
    */
   NDArray(const NDArrayStorageType stype, const TShape &shape,
           const TBlob &data, const std::vector<TBlob> &aux_data, int dev_id)
@@ -423,6 +419,12 @@ class NDArray {
    * \param size the size of the source array, in sizeof(DType) not raw btyes.
    */
   void SyncCopyFromCPU(const void *data, size_t size) const;
+
+  /*!
+   * \brief Copy from src.data()/aux_data(i) to this->data()/aux_data(j)
+   */
+  void SyncCopyFromNDArray(const NDArray &src, int i = -1, int j = -1);
+
   /*!
    * \brief Do a synchronize copy to a continugous CPU memory region.
    *
@@ -448,19 +450,19 @@ class NDArray {
    * \return idx-th sub array NDArray
    */
   NDArray At(index_t idx) const;
-  // Wrap the tblob of aux data into an NDArray which shares the same variable with the
-  // current one.
-  inline const NDArray aux_ndarray(size_t i) const {
-    CHECK_NE(storage_type(), kDefaultStorage);
-    CHECK(i < ptr_->aux_shapes.size());
-    return NDArray(aux_data(i), ctx().dev_id, var());
-  }
-  // Wrap the tblob of data into an NDArray which shares the same variable with the
-  // current one.
-  inline const NDArray data_ndarray() const {
-    CHECK_NE(storage_type(), kDefaultStorage);
-    return NDArray(data(), ctx().dev_id, var());
-  }
+
+  /*!
+   * \brief Generate a deep copy of aux_data(i) returned as
+   * a default storage type NDArray
+   */
+  NDArray aux_ndarray(size_t i) const;
+
+  /*!
+   * \brief Generate a deep copy of data() returned as a
+   * default storage type NDArray
+   */
+  NDArray data_ndarray() const;
+
   /*!
    * \brief Create a NDArray that shares memory with current one
    *  The new array must have smaller memory size than the current array.
@@ -506,6 +508,23 @@ class NDArray {
     CHECK_EQ(storage_type(), kDefaultStorage);
     ptr_->CheckAndAlloc();
   }
+
+  /*!
+   * \brief Allocate the space if the allocation has been delayed
+   * or the requested size is bigger than the available one.
+   * This function can only be called by ndarray of default
+   * storage type and effectively changes the ndarray's shape_.
+   * Note: This function is named as this to avoid overload conflict
+   * with CheckAndAlloc(const std::vector<TShape> &aux_shapes), since
+   * TShape tmp = some_shape is equivalent to TShape tmp = {some_shape}.
+   */
+  void ReshapeAndAlloc(const TShape& shape) {
+    CHECK_EQ(storage_type(), kDefaultStorage);
+    CHECK(!is_none());
+    shape_ = shape;
+    ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
+  }
+
   /* !
    * \brief Alloc memory for non-default storage
    * aux_shape is only known at run time
@@ -581,8 +600,6 @@ class NDArray {
     // The shape of aux data. The default value for the shape depends on the type of storage.
     // If aux_shapes[i].Size() is zero, aux data i is empty.
     std::vector<TShape> aux_shapes;
-    // \brief skip the deletion of var handle. Usually set when shared_var is present.
-    bool skip_delete_var = false;
 
     /*! \brief default cosntructor */
     Chunk() : static_data(true), delay_alloc(false) {}
@@ -598,17 +615,10 @@ class NDArray {
       if (!delay_alloc_) this->CheckAndAlloc();
     }
 
-    Chunk(const TBlob &data, int dev_id, Engine::VarHandle shared_var)
+    Chunk(const TBlob &data, int dev_id)
         : static_data(true), delay_alloc(false) {
       CHECK(storage_type == kDefaultStorage);
-      // init var
-      if (shared_var == nullptr) {
-        var = Engine::Get()->NewVariable();
-      } else {
-        skip_delete_var = true;
-        var = shared_var;
-      }
-      // init ctx
+      var = Engine::Get()->NewVariable();
       if (data.dev_mask() == cpu::kDevMask) {
         ctx = Context::CPU();
       } else {
@@ -633,6 +643,9 @@ class NDArray {
       // aux_handles always reflect the correct number of aux data
       for (size_t i = 0; i < aux_shapes.size(); i++) {
         CheckAndAllocAuxData(i, aux_shapes[i]);
+        // this line is needed in case when aux_shapes[i].Size() = 0
+        // aux_handles[i] will not be updated and take only default value.
+        aux_handles[i].ctx = ctx;
       }
       if (!delay_alloc) {
         CheckAndAllocData(storage_shape, dtype);
@@ -677,6 +690,22 @@ class NDArray {
         delay_alloc = false;
       }
     }
+
+    /*! \brief Check and alloc memory for a dense ndarray */
+    // size is the number of bytes
+    void CheckAndAlloc(uint64_t dbytes) {
+      CHECK_EQ(kDefaultStorage, storage_type);
+      if (delay_alloc) {
+        shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
+        delay_alloc = false;
+      } else if (shandle.size < dbytes) {
+        // free storage if necessary and alloc again
+        if (shandle.size > 0) Storage::Get()->Free(shandle);
+        // init storage
+        shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
+      }
+    }
+
     inline void CheckAndAlloc(const TShape &shape, const std::vector<TShape> &aux_shapes,
                               int dtype) {
       // calculate size, perform allocation
@@ -740,7 +769,6 @@ class NDArray {
     }
     /*! \brief destructor */
     ~Chunk() {
-      if (skip_delete_var) return;
       bool skip_free = static_data || delay_alloc;
       Storage::Handle h = this->shandle;
       std::vector<Storage::Handle> aux_h = this->aux_handles;

python/mxnet/ndarray/sparse_ndarray.py

@@ -23,7 +23,7 @@
 from . import ndarray
 from .ndarray import _DTYPE_NP_TO_MX, _DTYPE_MX_TO_NP
 from .ndarray import _STORAGE_TYPE_STR_TO_ID
-from .ndarray import NDArray, _storage_type, _zeros_ndarray
+from .ndarray import NDArray, _storage_type, _zeros_ndarray, array
 from . import cast_storage
 from . import slice as nd_slice
 
@@ -220,20 +220,19 @@ def _aux_type(self, i):
 
     @property
     def data(self):
-        """The values array of the SparseNDArray. This is a read-only view of the values array.
-        They reveal internal implementation details and should be used with care.
+        """Get a deep copy of the values array of the SparseNDArray.
 
         Returns
         -------
         NDArray
-            This SparseNDArray's values array.
+            A deep copy of the SparseNDArray's values array.
         """
         return self._data()
 
     @property
     def _num_aux(self):
-        ''' The number of aux data used to help store the sparse ndarray.
-        '''
+        """The number of aux data used to help store the sparse ndarray.
+        """
         return len(_STORAGE_AUX_TYPES[self.stype])
 
     @property
@@ -253,7 +252,6 @@ def _aux_types(self):
 
     def asnumpy(self):
         """Return a dense ``numpy.ndarray`` object with value copied from this array
-
         """
         return self.todense().asnumpy()
 
@@ -311,21 +309,23 @@ def copyto(self, other):
     def todense(self):
         return todense(self)
 
-    def _aux_data(self, i, writable=False):
-        """ Get an NDArray referencing the ith aux data array associated with the SparseNDArray.
+    def _aux_data(self, i):
+        """ Get a deep copy NDArray of the i-th aux data array associated with the SparseNDArray.
+        This function blocks. Do not use it in performance critical code.
         """
         self.wait_to_read()
         hdl = NDArrayHandle()
         check_call(_LIB.MXNDArrayGetAuxNDArray(self.handle, i, ctypes.byref(hdl)))
-        return NDArray(hdl, writable)
+        return NDArray(hdl)
 
-    def _data(self, writable=False):
-        """ Get an NDArray referencing the value array associated with the SparseNDArray.
+    def _data(self):
+        """ Get a deep copy NDArray of the value array associated with the SparseNDArray.
+        This function blocks. Do not use it in performance critical code.
         """
         self.wait_to_read()
         hdl = NDArrayHandle()
         check_call(_LIB.MXNDArrayGetDataNDArray(self.handle, ctypes.byref(hdl)))
-        return NDArray(hdl, writable)
+        return NDArray(hdl)
 
 # pylint: disable=abstract-method
 class CSRNDArray(SparseNDArray):
@@ -351,8 +351,8 @@ def __reduce__(self):
 
     @property
     def indices(self):
-        """The indices array of the SparseNDArray. This is a read-only view of the indices array.
-        They reveal internal implementation details and should be used with care.
+        """The indices array of the SparseNDArray with `csr` storage type.
+        This generates a deep copy of the column indices of the current `csr` matrix.
 
         Returns
         -------
@@ -364,8 +364,7 @@ def indices(self):
     @property
     def indptr(self):
         """The indptr array of the SparseNDArray with `csr` storage type.
-        This is a read-only view of the indptr array.
-        They reveal internal implementation details and should be used with care.
+        This generates a deep copy of the `indptr` of the current `csr` matrix.
 
         Returns
         -------
@@ -405,8 +404,8 @@ def __reduce__(self):
 
     @property
     def indices(self):
-        """The indices array of the SparseNDArray. This is a read-only view of the indices array.
-        They reveal internal implementation details and should be used with care.
+        """The indices array of the SparseNDArray with `row_sparse` storage type.
+        This generates a deep copy of the row indices of the current row-sparse matrix.
 
         Returns
         -------
@@ -490,31 +489,36 @@ def csr(data, indptr, indices, shape, ctx=None, dtype=None, indptr_type=None, in
     assert(len(shape) == 2)
     result = CSRNDArray(_new_alloc_handle(storage_type, shape, ctx, False, dtype,
                                           [indptr_type, indices_type], aux_shapes))
-    # assign indptr, indices and data
-    data_ref = result._data(True)
-    indptr_ref = result._aux_data(0, True)
-    indices_ref = result._aux_data(1, True)
-    data_ref[:] = data
-    indptr_ref[:] = indptr
-    indices_ref[:] = indices
+    # TODO(junwu): Convert data, indptr, and indices to mxnet NDArrays
+    # if they are not for now. In the future, we should provide a c-api
+    # to accept np.ndarray types to copy from to result.data and aux_data
+    if not isinstance(data, NDArray):
+        data = array(data, ctx, dtype)
+    if not isinstance(indptr, NDArray):
+        indptr = array(indptr, ctx, indptr_type)
+    if not isinstance(indices, NDArray):
+        indices = array(indices, ctx, indices_type)
+    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, data.handle, ctypes.c_int(-1)))
+    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, indptr.handle, ctypes.c_int(0)))
+    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, indices.handle, ctypes.c_int(1)))
     return result
 
 
-def row_sparse(values, indices, shape, ctx=None, dtype=None, indices_type=None):
+def row_sparse(data, indices, shape, ctx=None, dtype=None, indices_type=None):
     """Creates a row sparse array with a set of tensor slices at given indices.
 
     Parameters
     ----------
-    values: array_like
+    data: array_like
         An object exposing the array interface, with shape [D0, D1, .. Dn], where D0 is
         the number of rows with non-zeros entries.
     indices: array_like
         An object exposing the array interface, with shape [D0].
     ctx : Context, optional
         Device context (default is the current default context).
     dtype : str or numpy.dtype, optional
-        The data type of the output array. The default dtype is ``values.dtype``
-        if `values` is an `NDArray`, `float32` otherwise.
+        The data type of the output array. The default dtype is ``data.dtype``
+        if `data` is an `NDArray`, `float32` otherwise.
     indices_type: str or numpy.dtype, optional
         The data type of the indices array. The default dtype is ``indices.dtype``
         if `indicies` is an `NDArray`, `int32` otherwise.
@@ -540,21 +544,26 @@ def row_sparse(values, indices, shape, ctx=None, dtype=None, indices_type=None):
     if ctx is None:
         ctx = Context.default_ctx
     # prepare src array and types
-    values, dtype = _prepare_src_array(values, dtype, mx_real_t)
+    data, dtype = _prepare_src_array(data, dtype, mx_real_t)
     indices, indices_type = _prepare_src_array(indices, indices_type,
                                                _STORAGE_AUX_TYPES[storage_type][0])
     # verify types
     assert('int64' in str(indices_type)), "expected int64 for indices"
     # verify shapes
-    assert(values.ndim == len(shape))
+    assert(data.ndim == len(shape))
     assert(indices.ndim == 1)
     result = RowSparseNDArray(_new_alloc_handle(storage_type, shape, ctx, False, dtype,
                                                 [indices_type], [indices.shape]))
-    # assign indices and values
-    values_ref = result._data(True)
-    indices_ref = result._aux_data(0, True)
-    values_ref[:] = values
-    indices_ref[:] = indices
+
+    # TODO(junwu): Convert data, indptr, and indices to mxnet NDArrays
+    # if they are not for now. In the future, we should provide a c-api
+    # to accept np.ndarray types to copy from to result.data and aux_data
+    if not isinstance(data, NDArray):
+        data = array(data, ctx, dtype)
+    if not isinstance(indices, NDArray):
+        indices = array(indices, ctx, indices_type)
+    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, data.handle, ctypes.c_int(-1)))
+    check_call(_LIB.MXNDArraySyncCopyFromNDArray(result.handle, indices.handle, ctypes.c_int(0)))
     return result