[GitHub] cjolivier01 commented on a change in pull request #8826: [WIP] improve memory usage in sparse.dot

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

cjolivier01 commented on a change in pull request #8826: [WIP] improve memory usage in sparse.dot
URL: https://github.com/apache/incubator-mxnet/pull/8826#discussion_r153243691
 
 

 ##########
 File path: src/operator/tensor/dot-inl.h
 ##########
 @@ -613,15 +613,15 @@ inline void DotCsrDnsRspImpl(const OpContext& ctx,
       MSHADOW_IDX_TYPE_SWITCH(col_idx_l.type_flag_, CType, {  // col 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));
+          size_t workspace_size = 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<MarkRowFlgKernel, cpu>::Launch(s, lhs.aux_shape(csr::kIdx)[0], row_flg,
+          // prefix sum array re-uses the row_flg array temp space
+          dim_t* prefix_sum = row_flg;
+          Kernel<op_with_req<set_to_int<0>, kWriteTo>, cpu>::Launch(s, num_rows, row_flg);
 
 Review comment:
   set_zero is just a 'using' of set_to_int<0>, so it's the same thing.

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