eric-haibin-lin commented on a change in pull request #8611: optimization for
dot(csr.T, dense) = rsp
URL: https://github.com/apache/incubator-mxnet/pull/8611#discussion_r150351688
##########
File path: src/operator/tensor/dot-inl.h
##########
@@ -598,40 +614,52 @@ inline void DotCsrDnsRspImpl(const OpContext& ctx,
const TBlob col_idx_l = lhs.aux_data(csr::kIdx);
const TBlob& data_r = rhs;
- // pre-allocate spaces for ret using the dense dimension size
- ret->CheckAndAlloc({mshadow::Shape1(lhs.shape()[1])});
- const TBlob data_out = ret->data();
- const TBlob row_idx_out = ret->aux_data(rowsparse::kIdx);
-
MSHADOW_SGL_DBL_TYPE_SWITCH(data_l.type_flag_, DType, { // data type
MSHADOW_IDX_TYPE_SWITCH(indptr_l.type_flag_, IType, { // indptr type
MSHADOW_IDX_TYPE_SWITCH(col_idx_l.type_flag_, CType, { // col idx type
- MSHADOW_IDX_TYPE_SWITCH(row_idx_out.type_flag_, RType, { // row idx
type
+ MSHADOW_IDX_TYPE_SWITCH(ret->aux_type(rowsparse::kIdx), RType, { //
row idx type
+ const dim_t num_rows = lhs.shape()[1];
+ size_t workspace_size = 2 * (num_rows * sizeof(dim_t));
+ mshadow::Tensor<cpu, 1, char> workspace =
+ ctx.requested[0].get_space_typed<cpu, 1, char>(
+ mshadow::Shape1(workspace_size), s);
+ dim_t* row_flg = reinterpret_cast<dim_t*>(workspace.dptr_);
+ dim_t* prefix_sum = row_flg + num_rows;
+
+ Fill<false>(s, TBlob(row_flg, mshadow::Shape1(num_rows),
cpu::kDevMask), kWriteTo, 0);
+ mxnet_op::Kernel<MarkCsrColKernel, cpu>::Launch(s,
lhs.aux_shape(csr::kIdx)[0], row_flg,
+ col_idx_l.dptr<CType>());
+
+ prefix_sum[0] = row_flg[0];
+ for (nnvm::dim_t i = 1; i < num_rows; i++) {
+ prefix_sum[i] = prefix_sum[i - 1] + row_flg[i];
+ }
+ dim_t nnr = prefix_sum[num_rows - 1];
+
+ if (nnr == 0) {
+ FillZerosRspImpl(s, *ret);
+ return;
+ }
+
+ ret->CheckAndAlloc({mshadow::Shape1(nnr)});
+ const TBlob& data_out = ret->data();
+ const TBlob& row_idx = ret->aux_data(rowsparse::kIdx);
+
dim_t num_threads = data_out.Size();
mxnet_op::Kernel<set_zero, cpu>::Launch(s, num_threads,
data_out.dptr<DType>());
- RType* row_idx = row_idx_out.dptr<RType>();
- num_threads = row_idx_out.Size();
- mxnet_op::Kernel<set_zero, cpu>::Launch(s, num_threads, row_idx);
- num_threads = mxnet_op::get_num_threads<cpu>(data_out.shape_[0]);
- dim_t seg_len = (data_out.shape_[0] + num_threads - 1) / num_threads;
+ RType* row_idx_out = row_idx.dptr<RType>();
+
+ mxnet_op::Kernel<FillRspRowIdxKernel, cpu>::Launch(s, num_rows,
+ row_idx_out, prefix_sum, num_rows);
+
+ num_threads = mxnet_op::get_num_threads<cpu>(ret->shape()[0]);
+ dim_t seg_len = (ret->shape()[0] + num_threads - 1) / num_threads;
Review comment:
Should we replace `ret->shape()[0] ` by `nnr` instead?
https://github.com/apache/incubator-mxnet/blob/master/src/operator/tensor/indexing_op.h#L661
----------------------------------------------------------------
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
