eric-haibin-lin commented on a change in pull request #10894: [MXNET-399]
Elemwise_mul between dense and csr on CPU & GPU
URL: https://github.com/apache/incubator-mxnet/pull/10894#discussion_r191573736
##########
File path: src/operator/tensor/elemwise_binary_op-inl.h
##########
@@ -495,6 +495,91 @@ void ElemwiseBinaryOp::DnsCsrDnsOp(mshadow::Stream<xpu>
*s,
});
}
+/*!
+ * \brief Kernel for performing elemwise op between dense and csr matrix
+ * \param i global thread id
+ * \param req type of request
+ * \param out output array
+ * \param dns_data data array of dense input
+ * \param csr_data data array of csr input
+ * \param csr_indices indices array of csr input
+ * \param csr_indptr indptr array of csr input
+ * \param num_rows number of rows of both inputs
+ * \param num_cols number of columns of both inputs
+ */
+template<int req, typename OP, bool reverse>
+struct ElemwiseDnsCsrCsrKernel {
+ template<typename DType, typename IType, typename CType>
+ MSHADOW_XINLINE static void Map(int i, DType* out, DType* dns_data,
+ const DType* csr_data, const IType*
csr_indices,
+ const CType* csr_indptr, const nnvm::dim_t
num_rows,
+ const nnvm::dim_t num_cols) {
+ if (i < num_rows) {
+ for (int j = csr_indptr[i]; j < csr_indptr[i+1]; ++j) {
+ KERNEL_ASSIGN(out[j], req, reverse ?
+ OP::Map(dns_data[i * num_cols +
csr_indices[j]], csr_data[j]) :
+ OP::Map(csr_data[j], dns_data[i * num_cols
+ csr_indices[j]]));
+ }
+ }
+ }
+};
+
+/*! \brief DNS -op- CSR binary operator for non-canonical NDArray */
+template<typename xpu, typename OP>
+void ElemwiseBinaryOp::DnsCsrCsrOp(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const NDArray &dns,
+ const NDArray &csr,
+ const OpReqType req,
+ const NDArray &output,
+ const bool reverse) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ using namespace csr;
+ CHECK_EQ(dns.storage_type(), kDefaultStorage);
+ CHECK_EQ(csr.storage_type(), kCSRStorage);
+ CHECK_EQ(req, kWriteTo) << "elemwise(dns, csr) = csr only supports kWriteTo";
+ if (req == kNullOp) return;
+ const bool supported_op = std::is_same<OP, mshadow_op::mul>::value;
+ CHECK(supported_op == true) << "elemwise(dns, csr) = csr only supports
mul/div";
Review comment:
remove div?
----------------------------------------------------------------
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:
[email protected]
With regards,
Apache Git Services
