[GitHub] eric-haibin-lin commented on a change in pull request #8814: csr slice operator, gpu implementation

[GitBox <gi...@apache.org>]

eric-haibin-lin commented on a change in pull request #8814: csr slice operator, gpu implementation
URL: https://github.com/apache/incubator-mxnet/pull/8814#discussion_r153104061
 
 

 ##########
 File path: src/operator/tensor/matrix_op.cc
 ##########
 @@ -28,6 +28,104 @@
 
 namespace mxnet {
 namespace op {
+
+
+template<>
+void SliceDimOneCsrImpl<cpu>(const TShape &begin, const TShape &end, const OpContext& ctx,
+                             const NDArray &in, const NDArray &out) {
+  using namespace mshadow;
+  using namespace mxnet_op;
+  using namespace csr;
+  nnvm::dim_t begin_row = begin[0];
+  nnvm::dim_t end_row = end[0];
+  nnvm::dim_t indptr_len = end_row - begin_row + 1;
+  out.CheckAndAllocAuxData(kIndPtr, Shape1(indptr_len));
+  // assume idx indptr share the same type
+  MSHADOW_IDX_TYPE_SWITCH(in.aux_type(kIndPtr), RType, {
+    MSHADOW_IDX_TYPE_SWITCH(in.aux_type(kIdx), IType, {
+      MSHADOW_TYPE_SWITCH(in.dtype(), DType, {
+        RType* in_indptr = in.aux_data(kIndPtr).dptr<RType>();
+        RType* out_indptr = out.aux_data(kIndPtr).dptr<RType>();
+        SliceCsrIndPtrImpl<cpu, RType>(begin_row, end_row, ctx.run_ctx, in_indptr, out_indptr);
 
 Review comment:
   Is it true that only retrieving nnz has different implementation on GPU while the rest is the same? For memcpy we can use mshadow::Copy also. Shall we have a templated fucntion that performs retrieving nnz, so there's no need to copy the majority of the code body again?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services