Add support for cast storage on same stypes (#10400)

* Add cast storage support for same stypes

* Add imports

* Fix cast

* Fix doc for cast_storage

* Fix

src/operator/tensor/cast_storage-inl.h

@@ -30,6 +30,7 @@
 #include <algorithm>
 #include "../mxnet_op.h"
 #include "../operator_common.h"
+#include "../../src/operator/tensor/init_op.h"
 #ifdef __CUDACC__
 #include "./cast_storage-inl.cuh"
 #endif  // __CUDACC__
@@ -328,6 +329,50 @@ void CastStorageCsrDnsImpl(const OpContext& ctx,
   });
 }
 
+/*!
+ * \brief Casts a csr matrix to another csr.
+ */
+template <typename xpu>
+void CastStorageCsrCsrImpl(const OpContext& ctx, const NDArray& csr,
+                           NDArray* output) {
+  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
+  if (!csr.storage_initialized()) {
+    FillZerosCsrImpl(s, *output);
+    return;
+  }
+  std::vector<TShape> aux_shapes({csr.aux_shape(csr::kIndPtr), csr.aux_shape(csr::kIdx)});
+  output->CheckAndAlloc(aux_shapes);
+  const TBlob& val = output->data();
+  const TBlob& indptr = output->aux_data(csr::kIndPtr);
+  const TBlob& idx = output->aux_data(csr::kIdx);
+  mxnet_op::copy(s, val, csr.data());
+  mxnet_op::copy(s, indptr, csr.aux_data(csr::kIndPtr));
+  mxnet_op::copy(s, idx, csr.aux_data(csr::kIdx));
+}
+
+/*!
+ * \brief Casts a rsp matrix to another rsp.
+ */
+template <typename xpu>
+void CastStorageRspRspImpl(const OpContext& ctx, const NDArray& rsp,
+                           NDArray* output) {
+  CHECK_EQ(rsp.storage_type(), output->storage_type())
+      << "Copying with different storage type";
+  mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
+  if (!rsp.storage_initialized()) {
+    FillZerosRspImpl(s, *output);
+    return;
+  }
+  auto aux_shape = rsp.aux_shape(rowsparse::kIdx);
+  output->CheckAndAlloc({aux_shape});
+  const TBlob& val = output->data();
+  const TBlob& idx = output->aux_data(rowsparse::kIdx);
+  const TBlob& from_val = rsp.data();
+  const TBlob& from_idx = rsp.aux_data(rowsparse::kIdx);
+  mxnet_op::copy(s, val, from_val);
+  mxnet_op::copy(s, idx, from_idx);
+}
+
 template<typename xpu>
 void CastStorageComputeImpl(const OpContext& ctx,
                             const NDArray& input,
@@ -346,6 +391,12 @@ void CastStorageComputeImpl(const OpContext& ctx,
   } else if (src_stype == kCSRStorage && dst_stype == kDefaultStorage) {
     TBlob ret = output.data();
     CastStorageCsrDnsImpl<xpu>(ctx, input, &ret);
+  } else if (src_stype == kCSRStorage && dst_stype == kCSRStorage) {
+    NDArray ret = output;
+    CastStorageCsrCsrImpl<xpu>(ctx, input, &ret);
+  } else if (src_stype == kRowSparseStorage && dst_stype == kRowSparseStorage) {
+    NDArray ret = output;
+    CastStorageRspRspImpl<xpu>(ctx, input, &ret);
 #if MXNET_USE_MKLDNN == 1
   } else if (src_stype == kDefaultStorage && dst_stype == kDefaultStorage) {
     CHECK_EQ(output.ctx().dev_type, input.ctx().dev_type);

src/operator/tensor/cast_storage.cc

@@ -46,6 +46,8 @@ The storage type of ``cast_storage`` output depends on stype parameter:
 - cast_storage(row_sparse, 'default') = default
 - cast_storage(default, 'csr') = csr
 - cast_storage(default, 'row_sparse') = row_sparse
+- cast_storage(csr, 'csr') = csr
+- cast_storage(row_sparse, 'row_sparse') = row_sparse
 
 Example::
 

tests/python/unittest/test_sparse_operator.py

@@ -1177,10 +1177,13 @@ def check_cast_storage(shape, density, from_stype, to_stype, check_numeric_grad=
         shape_3d = rand_shape_3d()
         check_cast_storage(shape_2d, d, 'csr', 'default')
         check_cast_storage(shape_2d, d, 'default', 'csr')
+        check_cast_storage(shape_2d, d, 'csr', 'csr')
         check_cast_storage(shape_2d, d, 'row_sparse', 'default')
         check_cast_storage(shape_2d, d, 'default', 'row_sparse')
+        check_cast_storage(shape_2d, d, 'row_sparse', 'row_sparse')
         check_cast_storage(shape_3d, d, 'row_sparse', 'default')
         check_cast_storage(shape_3d, d, 'default', 'row_sparse')
+        check_cast_storage(shape_3d, d, 'row_sparse', 'row_sparse')
         for i in range(4, 6):
             shape = rand_shape_nd(i, 5)
             check_cast_storage(shape, d, 'default', 'row_sparse')