implementation for Csr slice on cpu

python/mxnet/sparse_ndarray.py

@@ -251,13 +251,13 @@ def _slice(self, start, stop):
         >>> a[1:2].asnumpy()
         array([[0, 0, 3]])
 
-        >>> a[1:2].indptr.asnumpy()
+        >>> a[1:2]._indptr.asnumpy()
         array([[2, 3]])
 
-        >>> a[1:2].indicies.asnumpy()
+        >>> a[1:2]._indicies.asnumpy()
         array([0, 2, 2, 0, 1, 2])
 
-        >>> a[1:2].values.asnumpy()
+        >>> a[1:2]._values.asnumpy()
         array([1, 2, 3, 4, 5, 6])
 
         """
@@ -293,11 +293,27 @@ def _aux_type(self, i):
         return _DTYPE_MX_TO_NP[aux_type.value]
 
     @property
-    def values(self):
+    def _values(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.
+
+        Returns
+        -------
+        NDArray
+            This SparseNDArray's values array.
+        """
         return self._data(0)
 
     @property
-    def indices(self):
+    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.
+
+        Returns
+        -------
+        NDArray
+            This SparseNDArray's indices array.
+        """
         stype = self.storage_type
         if stype == 'row_sparse':
             return self._aux_data(0)
@@ -306,7 +322,16 @@ def indices(self):
         raise Exception("unknown storage type " + stype)
 
     @property
-    def indptr(self):
+    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.
+
+        Returns
+        -------
+        NDArray
+            This SparseNDArray's indptr array.
+        """
         stype = self.storage_type
         if stype == 'csr':
             return self._aux_data(0)
@@ -383,6 +408,7 @@ def _aux_data(self, i, writable=False):
         SparseNDArray is not yet compacted, the returned result may include invalid values.
 
         """
+        self.wait_to_read()
         hdl = NDArrayHandle()
         check_call(_LIB.MXNDArrayGetAuxNDArray(self.handle, i, ctypes.byref(hdl)))
         return NDArray(hdl, writable)
@@ -392,6 +418,7 @@ def _data(self, writable=False):
         SparseNDArray is not yet compacted, the returned result may include invalid values.
 
         """
+        self.wait_to_read()
         hdl = NDArrayHandle()
         check_call(_LIB.MXNDArrayGetDataNDArray(self.handle, ctypes.byref(hdl)))
         return NDArray(hdl, writable)

src/operator/tensor/indexing_op.h

@@ -328,7 +328,7 @@ struct EmbeddingBackwardRsp {
     size_t segment_end = (i + 1) * segment_len;
     for (size_t y = 0; y < num_idx; y++) {
       size_t j = idx[y];
-      if (j > num_rows) j = num_rows - 1;
+      if (j >= num_rows) j = num_rows - 1;
       if (j < segment_start || j >= segment_end) continue;
       dst_idx[j] = j;
       for (size_t k = 0; k < width; k++) {

src/operator/tensor/matrix_op-inl.h

@@ -7,6 +7,7 @@
 #define MXNET_OPERATOR_TENSOR_MATRIX_OP_INL_H_
 
 #include <mxnet/operator_util.h>
+#include <dmlc/omp.h>
 #include <vector>
 #include <algorithm>
 #include <utility>
@@ -1217,6 +1218,84 @@ void Slice(const nnvm::NodeAttrs& attrs,
   });
 }
 
+// slice the indptr of a csr
+struct SliceCsrIndPtr {
+  template<typename IType>
+  MSHADOW_XINLINE static void Map(int i, IType* out, const IType* in, const IType* base) {
+    KERNEL_ASSIGN(out[i], kWriteTo, in[i] - *base);
+  }
+};
+
+/*
+ * Slice a CSR NDArray
+ * Only implemented for CPU
+ */
+template<typename xpu>
+void SliceCsrImpl(const SliceParam &param, const OpContext& ctx,
+                  const NDArray &in, OpReqType req, const NDArray &out) {
+  using namespace mshadow;
+  using namespace mxnet_op;
+  using namespace csr;
+  CHECK((std::is_same<xpu, cpu>::value)) << "Slice for CSR input only implemented for CPU";
+  if (req == kNullOp) return;
+  CHECK_NE(req, kAddTo) << "kAddTo for Slice on CSR input is not supported";
+  CHECK_NE(req, kWriteInplace) << "kWriteInplace for Slice on CSR input is not supported";
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+  int begin = *param.begin[0];
+  int end = *param.end[0];
+  int indptr_len = end - begin + 1;
+  out.CheckAndAllocAuxData(kIndPtr, Shape1(indptr_len));
+  if (!in.storage_initialized()) {
+    out.SetAuxShape(kIndPtr, Shape1(0));
+    return;
+  }
+  CHECK_EQ(in.aux_type(kIndPtr), in.aux_type(kIdx))
+           << "The type for indptr and indices are different. This is not implemented yet.";
+  // assume idx indptr share the same type
+  NDARRAY_IDX_TYPE_SWITCH(in.aux_type(kIndPtr), IType, {
+    MSHADOW_TYPE_SWITCH(in.dtype(), DType, {
+      auto in_indptr = in.aux_data(kIndPtr).dptr<IType>();
+      auto out_indptr = out.aux_data(kIndPtr).dptr<IType>();
+      int num_threads = omp_get_num_threads();
+      int segment_len = (indptr_len + num_threads - 1) / num_threads;
+      Kernel<SliceCsrIndPtr, xpu>::Launch(s, indptr_len, out_indptr, in_indptr + begin,
+                                          in_indptr + begin);
+      // retrieve nnz (CPU implementation)
+      int nnz = out_indptr[indptr_len - 1] - out_indptr[0];
+      // copy indices and values
+      out.CheckAndAllocAuxData(kIdx, Shape1(nnz));
+      out.CheckAndAllocData(Shape1(nnz));
+      auto in_idx = in.aux_data(kIdx).dptr<IType>();
+      auto out_idx = out.aux_data(kIdx).dptr<IType>();
+      auto in_data = in.data().dptr<DType>();
+      auto out_data = out.data().dptr<DType>();
+      int offset = in_indptr[begin];
+      // this is also a CPU-only implementation
+      memcpy(out_idx, in_idx + offset, nnz * sizeof(IType));
+      memcpy(out_data, in_data + offset, nnz * sizeof(DType));
+    });
+  });
+}
+
+template<typename xpu>
+void SliceEx(const nnvm::NodeAttrs& attrs,
+          const OpContext& ctx,
+          const std::vector<NDArray>& inputs,
+          const std::vector<OpReqType>& req,
+          const std::vector<NDArray>& outputs) {
+  CHECK_EQ(inputs.size(), 1);
+  CHECK_EQ(outputs.size(), 1);
+  const SliceParam& param = nnvm::get<SliceParam>(attrs.parsed);
+  auto in_stype = inputs[0].storage_type();
+  CHECK_NE(in_stype, kDefaultStorage)
+           << "SliceEx is not expected to execute for input with default storage type";
+  if (in_stype == kCSRStorage) {
+    SliceCsrImpl<xpu>(param, ctx, inputs[0], req[0], outputs[0]);
+  } else {
+    LOG(FATAL) << "Slice not implemented for storage type" << in_stype;
+  }
+}
+
 inline bool SliceAssignShape(const nnvm::NodeAttrs& attrs,
                              std::vector<TShape> *in_attrs,
                              std::vector<TShape> *out_attrs) {

src/operator/tensor/matrix_op.cc

@@ -232,6 +232,9 @@ and ``end=(e_1, e_2, ... e_n)`` indices will result in an array with the shape
 The resulting array's *k*-th dimension contains elements
  from the *k*-th dimension of the input array with the open range ``[b_k, e_k)``.
 
+For an input array of non-default storage type(e.g. `csr` or `row_sparse`), it only supports
+slicing on the first dimension.
+
 Example::
 
   x = [[  1.,   2.,   3.,   4.],
@@ -245,8 +248,10 @@ Example::
 .set_attr_parser(ParamParser<SliceParam>)
 .set_attr<nnvm::FInferShape>("FInferShape", SliceShape)
 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
+.set_attr<nnvm::FInferStorageType>("FInferStorageType", ElemwiseStorageType<1, 1>)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_slice"})
 .set_attr<FCompute>("FCompute<cpu>", Slice<cpu>)
+.set_attr<FComputeEx>(FCOMP_EX_CPU, SliceEx<cpu>)
 .add_argument("data", "NDArray-or-Symbol", "Source input")
 .add_arguments(SliceParam::__FIELDS__());
 

tests/python/unittest/test_sparse_ndarray.py

@@ -75,10 +75,10 @@ def check_sparse_nd_prop_rsp():
     shape = (rnd.randint(1, 2), rnd.randint(1, 2))
     nd, (v, idx) = rand_sparse_ndarray(shape, storage_type)
     assert(nd._num_aux == 1)
-    assert(nd.indices.dtype == np.int32)
+    assert(nd._indices.dtype == np.int32)
     assert(nd.storage_type == 'row_sparse')
     assert_almost_equal(nd._data().asnumpy(), v)
-    assert_almost_equal(nd._aux_data(0).asnumpy(), idx)
+    assert_almost_equal(nd._indices.asnumpy(), idx)
 
 def test_sparse_nd_basic():
     def check_rsp_creation(values, indices, shape):
@@ -88,13 +88,13 @@ def check_rsp_creation(values, indices, shape):
         dns[3] = mx.nd.array(values[1])
         assert_almost_equal(rsp.asnumpy(), dns.asnumpy())
         indices = mx.nd.array(indices).asnumpy()
-        assert_almost_equal(rsp.indices.asnumpy(), indices)
+        assert_almost_equal(rsp._indices.asnumpy(), indices)
 
     def check_csr_creation(shape):
         csr, (indptr, indices, values) = rand_sparse_ndarray(shape, 'csr')
-        assert_almost_equal(csr.indptr.asnumpy(), indptr)
-        assert_almost_equal(csr.indices.asnumpy(), indices)
-        assert_almost_equal(csr.values.asnumpy(), values)
+        assert_almost_equal(csr._indptr.asnumpy(), indptr)
+        assert_almost_equal(csr._indices.asnumpy(), indices)
+        assert_almost_equal(csr._values.asnumpy(), values)
 
     shape = (4,2)
     values = np.random.rand(2,2)

tests/python/unittest/test_sparse_operator.py

@@ -170,11 +170,27 @@ def test_sparse_embedding():
     grad = mx.nd.zeros(np_grad.shape)
     grad[:] = np_grad
     exe_test.backward([grad])
-    assert_almost_equal(grad_map["embed_weight"].asnumpy(), np.dot(np_onehot.T, np_grad))
+    assert_almost_equal(grad_map["embed_weight"].asnumpy(), np.dot(np_onehot.T, np_grad), atol=1e-5)
+
+def test_sparse_slice():
+    def check_csr_slice(shape, sliced_input):
+        storage_type = 'csr'
+        A, _ = rand_sparse_ndarray(shape, storage_type)
+        A = A._slice(1, shape[0] - 1) if sliced_input else A
+        A2 = A.asnumpy()
+        begin = rnd.randint(0, A.shape[0] - 1)
+        end = rnd.randint(begin + 1, A.shape[0])
+        A_slice = mx.nd.crop(A, begin=begin, end=end)
+        assert same(A_slice.asnumpy(), A2[begin:end]), (A_slice.asnumpy(), A2[begin:end])
+
+    shape = (rnd.randint(7, 15), rnd.randint(1, 10))
+    check_csr_slice(shape, True)
+    check_csr_slice(shape, False)
 
 if __name__ == '__main__':
     test_elemwise_add_ex()
     test_elemwise_add_ex_multiple_stages()
     test_cast_storage_ex()
     test_sparse_dot()
     test_sparse_embedding()
+    test_sparse_slice()