[GitHub] eric-haibin-lin commented on a change in pull request #8331: [sparse] slice_axis for csr, cpu implementation

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

eric-haibin-lin commented on a change in pull request #8331: [sparse] slice_axis for csr, cpu implementation
URL: https://github.com/apache/incubator-mxnet/pull/8331#discussion_r146137664
 
 

 ##########
 File path: src/operator/tensor/matrix_op-inl.h
 ##########
 @@ -927,6 +927,140 @@ void SliceAxis(const nnvm::NodeAttrs& attrs,
   }
 }
 
+inline bool SliceAxisForwardInferStorageType(const nnvm::NodeAttrs& attrs,
+                                             const int dev_mask,
+                                             DispatchMode* dispatch_mode,
+                                             std::vector<int>* in_attrs,
+                                             std::vector<int>* out_attrs) {
+  CHECK_EQ(in_attrs->size(), 1);
+  CHECK_EQ(out_attrs->size(), 1);
+  const SliceAxisParam& param = nnvm::get<SliceAxisParam>(attrs.parsed);
+  const auto& in_stype = in_attrs->at(0);
+  auto& out_stype = out_attrs->at(0);
+  bool dispatched = false;
+  const bool invalid_ctx = dev_mask != mshadow::cpu::kDevMask;
+  const auto dispatch_ex = invalid_ctx ? DispatchMode::kFComputeFallback :
+                                         DispatchMode::kFComputeEx;
+  if (!dispatched && in_stype == kDefaultStorage) {
+    dispatched = storage_type_assign(&out_stype, kDefaultStorage,
+                                     dispatch_mode, DispatchMode::kFCompute);
+  }
+  if (!dispatched && in_stype == kCSRStorage && param.axis <= 1) {
+    dispatched = storage_type_assign(&out_stype, kCSRStorage,
+                                     dispatch_mode, dispatch_ex);
+  }
+  if (!dispatched) {
+    dispatch_fallback(out_attrs, dispatch_mode);
+  }
+  if (*dispatch_mode == DispatchMode::kFComputeFallback) {
+    LogStorageFallback(attrs, dev_mask, in_attrs, out_attrs);
+  }
+
+  return true;
+}
+
+struct SliceAxisOneCsrAssign {
+  template<typename IType, typename RType, typename DType>
+  MSHADOW_XINLINE static void Map(int i, IType* out_idx, DType* out_data,
+                                  const RType* out_indptr, const IType* in_idx,
+                                  const DType* in_data, const RType* in_indptr,
+                                  int begin, int end) {
+    RType ind = out_indptr[i];
+    for (int j=in_indptr[i]; j < in_indptr[i+1]; j++) {
+      if (in_idx[j] >= begin && in_idx[j] < end) {
+        out_idx[ind] = in_idx[j] - static_cast<IType>(begin);
+        out_data[ind] = in_data[j];
+        ind++;
+      }
+    }
+  }
+};
+
+template<typename xpu>
+void SliceAxisOneCsrImpl(const SliceAxisParam &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)) << "SliceAxis for CSR input only implemented for CPU";
+  if (req == kNullOp) return;
+  CHECK_NE(req, kAddTo) << "kAddTo for SliceAxis on CSR input is not supported";
+  CHECK_NE(req, kWriteInplace) << "kWriteInplace for SliceAxis on CSR input is not supported";
+  int axis, begin, end;
+  GetSliceAxisParams(param, in.shape(), &axis, &begin, &end);
+  int indptr_len = in.shape()[0] + 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>();
+        IType *in_idx = in.aux_data(kIdx).dptr<IType>();
+        DType *in_data = in.data().dptr<DType>();
+        // retrieve nnz (CPU implementation)
+        RType *out_indptr = out.aux_data(kIndPtr).dptr<RType>();
+        int nnz = 0;
+        out_indptr[0] = 0;
+        for (int i=0; i < indptr_len - 1; i++) {
+          out_indptr[i+1] = out_indptr[i];
+          for (int j=in_indptr[i]; j < in_indptr[i+1]; j++) {
+            if (in_idx[j] >= begin && in_idx[j] < end) {
+              out_indptr[i+1]++;
+              nnz++;
+            }
+          }
+        }
+        out.CheckAndAllocAuxData(kIdx, Shape1(nnz));
+        out.CheckAndAllocData(Shape1(nnz));
+        IType *out_idx = out.aux_data(kIdx).dptr<IType>();
+        DType *out_data = out.data().dptr<DType>();
+
+        Stream<xpu> *s = ctx.get_stream<xpu>();
+        Kernel<SliceAxisOneCsrAssign, xpu>::Launch(s, indptr_len-1, out_idx, out_data,
+                                                   out_indptr, in_idx, in_data, in_indptr,
+                                                   begin, end);
+      });
+    });
+  });
+}
+
+template<typename xpu>
+void SliceAxisZeroCsrImpl(const SliceAxisParam &param, const OpContext& ctx,
+                  const NDArray &in, OpReqType req, const NDArray &out) {
+  int axis, begin, end;
+  GetSliceAxisParams(param, in.shape(), &axis, &begin, &end);
+  SliceParam slice_param;
+  slice_param.begin[0] = begin;
+  slice_param.end[0] = end;
+  SliceCsrImpl<xpu>(slice_param, ctx, in, req, out);
+}
+
+template<typename xpu>
+void SliceAxisEx(const nnvm::NodeAttrs& attrs,
+          const OpContext& ctx,
 
 Review comment:
   nit: indentation

----------------------------------------------------------------
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