optimization for dot(csr.T, dense) = rsp

[ZiyueHuang]

Description

Use prefix sum to compute nnr in order to allocate the row_sparse output.

Currently dot(csr.T, dense) = rsp will allocate the dense output and then cast it to row_sparse, but not free the unused memory.

I use run_benchmark(context, lhs="csr", rhs="default", lhs_trans=True, ...) in mxnet/benchmark/python/sparse/dot.py. Please correct me if I'm wrong.

But is dot(csr.T, dense) = rsp in master slow like this? Might due to others are using my machine at the same time?

Performance of origin dot(csr.T, dense) = rsp,

[hanfeng@model-gpu00:sparse]$ python dot.py --num-omp-threads 20
========================================================
  mxnet sparse dot benchmark: dot(csr, default) = default
  (matrix multiplication: (m x k)^T * (k x n) = m x n)
========================================================
 lhs_density(%)  rhs_density(%)    context        m        k        n  t_sparse(ms)   t_dense(ms)  speedup
            1.0           100.0     cpu(0)      128  1000000      256        366.19        135.76     0.37
            1.0           100.0     cpu(0)      128  1000000     1000       1327.12        503.92     0.38
            1.0           100.0     cpu(0)      128  1000000     1000       1237.33        454.01     0.37
            1.0           100.0     cpu(0)       64  1000000     1000        868.38        345.38     0.40
            1.0           100.0     cpu(0)      128  1000000     1000       1237.09        437.32     0.35

After this PR,

[hanfeng@model-gpu00:sparse]$ python dot.py --num-omp-threads 20
========================================================
  mxnet sparse dot benchmark: dot(csr, default) = default
  (matrix multiplication: (m x k)^T * (k x n) = m x n)
========================================================
 lhs_density(%)  rhs_density(%)    context        m        k        n  t_sparse(ms)   t_dense(ms)  speedup
            1.0           100.0     cpu(0)      128  1000000      256         83.90        137.18     1.64
            1.0           100.0     cpu(0)      128  1000000     1000        410.63        448.30     1.09
            1.0           100.0     cpu(0)      128  1000000     1000        467.91        492.87     1.05
            1.0           100.0     cpu(0)       64  1000000     1000        259.99        348.32     1.34
            1.0           100.0     cpu(0)      128  1000000     1000        481.77        416.20     0.86

As a feature requested in #8168

cc @eric-haibin-lin

Checklist

Essentials

• Passed code style checking (make lint)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• For user-facing API changes, API doc string has been updated.
• To my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• unittests already exist

Comments

• If this change is a backward incompatible change, why must this change be made.
• Intersting edge cases to note here

[eric-haibin-lin]

Can we reuse this kernel and give it a better name?
https://github.com/apache/incubator-mxnet/blob/master/src/operator/tensor/util/tensor_util-inl.h#L37

[eric-haibin-lin]

Should we replace ret->shape()[0] by nnr instead?
https://github.com/apache/incubator-mxnet/blob/master/src/operator/tensor/indexing_op.h#L661

[ZiyueHuang]

I think no, this range should cover num_cols of csr, i.e. num_rows of output. So nnr is not helpful for this kernel.

[eric-haibin-lin]

Yes, the original dot(csr.T, dense) performance is not good. The same for scipy - it's not always faster than the dense dot even with very sparse data.

[eric-haibin-lin]

The documentation should be updated

[ZiyueHuang]

Add benchmark for n=2 case,

Before,

python dot.py --num-omp-threads 16
========================================================
  mxnet sparse dot benchmark: dot(csr, default) = default
  (matrix multiplication: (m x k)^T * (k x n) = m x n)
========================================================
 lhs_density(%)  rhs_density(%)    context        m        k        n  t_sparse(ms)   t_dense(ms)  speedup
            1.0           100.0     cpu(0)      256  1000000        2         41.61         36.01     0.87

After,

python dot.py --num-omp-threads 16
========================================================
  mxnet sparse dot benchmark: dot(csr, default) = default
  (matrix multiplication: (m x k)^T * (k x n) = m x n)
========================================================
 lhs_density(%)  rhs_density(%)    context        m        k        n  t_sparse(ms)   t_dense(ms)  speedup
            1.0           100.0     cpu(0)      256  1000000        2         14.44         32.46     2.25

[moveforever]

hi, @eric-haibin-lin @ZiyueHuang , can the storage type of dot(csr.T, rsp) be rsp?

[eric-haibin-lin]

It's now indeed rsp:

a = mx.nd.ones((2,2)).tostype('row_sparse')
>>> a

<RowSparseNDArray 2x2 @cpu(0)>
>>> b = mx.nd.ones((2,2)).tostype('csr')
>>> mx.nd.sparse.dot(b,a,transpose_a=True)

<RowSparseNDArray 2x2 @cpu(0)>

[moveforever]

oh, i know, thanks.