eric-haibin-lin commented on a change in pull request #10371: [MXNET-263] [WIP] Support for dot(dns, csr) = dns and dot(dns, csr.T) = dns on GPU URL: https://github.com/apache/incubator-mxnet/pull/10371#discussion_r182912981
########## File path: src/operator/tensor/dot-inl.cuh ########## @@ -442,6 +445,105 @@ struct DotCsrRspDnsScalarKernel { } }; +/*! + * \brief GPU Kernel to scatter row id to corresponding entries + * \param tid global thread id + * \param csr_indptr indptr array of csr + * \param csr_rows array of row id of csr elements + * \param num_rows total number of rows in csr matrix + * Parallelization by output elements: 1 thread/row + */ +struct CsrRowScatterKernel { + template<typename CType> + __device__ __forceinline__ static void Map(int tid, + const CType* csr_indptr, + CType* csr_rows, + const nnvm::dim_t num_rows) { + if (tid < num_rows) { + for (CType i = csr_indptr[tid]; i < csr_indptr[tid+1]; ++i) { + csr_rows[i] = tid; + } + } + } +}; + +/*! + * \brief GPU Kernel of generation of transposed csr matrix + * \param tid global thread id Review comment: nit: remove line 472, 473 ---------------------------------------------------------------- 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: us...@infra.apache.org With regards, Apache Git Services