[GitHub] haojin2 commented on a change in pull request #10645: [MXNET-354] Support elemwise_add/sub between dense and row sparse tensors on CPU/GPU

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

haojin2 commented on a change in pull request #10645: [MXNET-354] Support elemwise_add/sub between dense and row sparse tensors on CPU/GPU
URL: https://github.com/apache/incubator-mxnet/pull/10645#discussion_r185668098
 
 

 ##########
 File path: src/operator/tensor/elemwise_binary_op-inl.h
 ##########
 @@ -450,6 +493,91 @@ void ElemwiseBinaryOp::DnsCsrDnsOp(mshadow::Stream<cpu> *s,
   });
 }
 
+/*!
+ * \brief Kernel for performing elemwise op between dense and rsp tensor
+ * \param i            global thread id
+ * \param req          type of request
+ * \param out          output array
+ * \param dns_data     data array of dense input
+ * \param rsp_data     data array of rsp input
+ * \param rsp_indices  indices array of rsp input
+ * \param num_rows     number of rows of both inputs
+ * \param nz_rows      number of non-zero rows of rsp tensor
+ * \param num_cols     number of columns of both inputs
+ */
+template<int req, typename OP>
+struct ElemwiseDnsRspDnsKernel {
+  template<typename DType, typename IType>
+  MSHADOW_XINLINE static void Map(int i, DType* out, DType* dns_data,
+                                  const DType* rsp_data, const IType* rsp_indices,
+                                  const nnvm::dim_t num_rows, const nnvm::dim_t nz_rows,
+                                  const nnvm::dim_t num_cols) {
+    if (i < nz_rows * num_cols) {
+      const nnvm::dim_t rsp_idx = i / num_cols;
+      const nnvm::dim_t dns_row = rsp_indices[rsp_idx];
+      const nnvm::dim_t col = i % num_cols;
+      KERNEL_ASSIGN(out[dns_row * num_cols + col], req,
+                    OP::Map(dns_data[dns_row * num_cols + col],
+                            rsp_data[rsp_idx * num_cols + col]));
+    }
+  }
+};
+
+/*! \brief DNS -op- RSP binary operator for non-canonical NDArray */
+template<typename xpu, typename OP>
+void ElemwiseBinaryOp::DnsRspDnsOp(mshadow::Stream<xpu> *s,
+                                   const nnvm::NodeAttrs &attrs,
+                                   const OpContext &ctx,
+                                   const NDArray &dns,
+                                   const NDArray &rsp,
+                                   const OpReqType req,
+                                   const NDArray &output,
+                                   const bool reverse) {
+  using namespace mshadow;
+  using namespace mxnet_op;
+  CHECK_EQ(dns.storage_type(), kDefaultStorage);
+  CHECK_EQ(rsp.storage_type(), kRowSparseStorage);
+  CHECK_EQ(output.data().Size(), dns.data().Size());
+  CHECK(req != kAddTo);
+  CHECK(req != kNullOp);
 
 Review comment:
   Done.

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