[tomicapretto]

Description

I incorporate some lessons from work on #2469 and this comment to improve the performance of previously existing sparse ops in numba

What I did (update)

From largest to lowest impact:

• Avoid operations at the array level, they usually introduced potentially large allocations. It's better to just write the for loop and do math with scalars.
• The previous point includes np.dot(...).
• Do not use arr.astype on entire arrays, which would allocate an entire new vector. Instead, do to_type(scalar).
• Type indices and indptr as unsigned. We are using uint32 because they're int32.
• Make sure all variables used to index arrays are unsigned.
• Use np.empty instead of np.zeros when possible, which is when we do set all values of the new array.

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Maintenance

Benchmarks

Sparse-dense dot product

Details

X: (2000, 2000), nnz=200000, density=0.05, dtype=float64
format layout width  before ms   after ms  speedup  allocations before/after
   csr      C     1      5.755      0.488   11.79x     400005/5
   csc      C     1      5.974      0.188   31.71x     400005/5
   csr      F     1      5.846      0.486   12.04x     400005/5
   csc      F     1      5.938      0.185   32.05x     400005/5
   csr      C   128     14.674      3.299    4.45x     400005/5
   csc      C   128     13.882      3.002    4.62x     400005/5
   csr      F   128     25.843     14.088    1.83x     400005/5
   csc      F   128     25.150     10.705    2.35x     400005/5
   csr      C   256     26.743      6.263    4.27x     400005/5
   csc      C   256     25.834      6.420    4.02x     400005/5
   csr      F   256     58.086     31.721    1.83x     400005/5
   csc      F   256     46.709     22.149    2.11x     400005/5
   csr      C   512     44.144     13.011    3.39x     400005/5
   csc      C   512     47.233     13.993    3.38x     400005/5
   csr      F   512    118.937     76.609    1.55x     400005/5
   csc      F   512     89.791     53.678    1.67x     400005/5

Sparse-sparse dot product

Details

NUMBA_NRT_STATS=1  python .benchmark/benchmark_sparse_dot.py
X, Y: (1000, 1000); nnz(X)=10000, nnz(Y)=10000; float64
Median of warmed-up batches; compilation excluded.
                case  before ms   after ms  speedup  allocs before/after
           csr @ csr      0.494      0.433    1.14x                11/11
           csr @ csc      0.509      0.452    1.13x                14/14
           csc @ csr      0.510      0.452    1.13x                14/14
           csc @ csc      0.496      0.435    1.14x                11/11

Sparse-dense dot gradient

Details

X: (2000, 1200), nnz=120000; width=128; density=0.05; float64
Median of warmed-up batches; compilation excluded.
                case  before ms   after ms  speedup  allocs before/after
     csr sparse grad      4.594      4.565    1.01x                  9/9
     csc sparse grad      4.590      4.323    1.06x                  9/9

Sparse-dense multiply

Details

NUMBA_NRT_STATS=1  python .benchmark/benchmark_sparse_multiply.py 
X: (2000, 2000), nnz=200000; density=0.05; float64
Median of warmed-up batches; compilation excluded.
                case  before ms   after ms  speedup  allocs before/after
        csr * vector      0.218      0.117    1.87x                  7/7
      csr * matrix C      1.345      0.773    1.74x                  7/7
      csr * matrix F      1.620      0.794    2.04x                  7/7
        csc * vector      0.127      0.062    2.06x                  7/7
      csc * matrix C      1.479      0.789    1.88x                  7/7
      csc * matrix F      1.345      0.739    1.82x                  7/7

Sparse-dense vector dot

Details

X: (2000, 2000), nnz=200000; density=0.05; float64
Median of warmed-up batches; compilation excluded.
                case  before ms   after ms  speedup  allocs before/after
          csr matvec      0.080      0.079    1.02x                  5/5
          csc matvec      0.083      0.083    1.00x                  5/5

Structured sparse-dense dot product gradient

Details

X: (2000, 1200), nnz=120000; width=128; Y=float64, G=float64
Median of warmed-up batches; JIT excluded; allocations include output.
format   layout  before ms   after ms  speedup  allocations before/after
   csr        C      5.399      1.819    2.97x          5/5
   csc        C      5.423      1.904    2.85x          5/5
   csr        F     23.268     10.597    2.20x     240005/5
   csc        F     24.205     11.112    2.18x     240005/5

Dense-sparse dot product

Details

Sparse Y: (2000, 2000), nnz=200000, float64
Median of warmed-up batches; JIT excluded; allocations include output.
format layout width  before ms   after ms  speedup  allocs before/after   avoided bytes/call
   csr      C     1      5.928      0.228   26.00x             400008/5            2,408,004
   csc      C     1      5.788      0.500   11.57x             400008/5            2,408,004
   csr      F     1      5.906      0.227   25.99x             400008/5            2,408,004
   csc      F     1      5.726      0.500   11.46x             400008/5            2,408,004
   csr      C   128     22.035     10.941    2.01x             400008/5            2,408,004
   csc      C   128     26.466     14.383    1.84x             400008/5            2,408,004
   csr      F   128     16.732      2.904    5.76x             400008/5            2,408,004
   csc      F   128     14.059      2.739    5.13x             400008/5            2,408,004
   csr      C   256     43.262     20.509    2.11x             400008/5            2,408,004
   csc      C   256     53.787     28.680    1.88x             400008/5            2,408,004
   csr      F   256     23.952      5.169    4.63x             400008/5            2,408,004
   csc      F   256     24.342      5.073    4.80x             400008/5            2,408,004
   csr      C   512     86.744     52.868    1.64x             400008/5            2,408,004
   csc      C   512    104.306     68.520    1.52x             400008/5            2,408,004
   csr      F   512     38.935     13.598    2.86x             400008/5            2,408,004
   csc      F   512     41.743     13.294    3.14x             400008/5            2,408,004

Sparse X.sum()

Details

float64; native-axis sums; warmed-up medians; JIT excluded
   rows  cols  density format  before ms   after ms  speedup allocs old allocs new
   2000   128    0.001    csr      0.004      0.004    1.16x          3          3
   2000   128    0.001    csc      0.003      0.003    1.03x          3          3
   2000   128    0.050    csr      0.010      0.007    1.43x          3          3
   2000   128    0.050    csc      0.007      0.004    1.93x          3          3
 100000   128    0.001    csr      0.141      0.045    3.12x          3          3
 100000   128    0.001    csc      0.007      0.004    1.92x          3          3
 100000   128    0.050    csr      0.797      0.758    1.05x          3          3
 100000   128    0.050    csc      0.404      0.049    8.29x          3          3

[ricardoV94]

Amazing results, I'll see if they reproduce on my machine, and check if compile time didn't get too much worse

[tomicapretto]

You'll need the benchmark scripts (codex generated but I think they're OK)

benchmark_dense_sparse_dot.py

benchmark_dot.py

benchmark_sparse_dot.py

benchmark_sparse_grad.py

benchmark_sparse_multiply.py

benchmark_sparse_sum.py

benchmark_spmv.py

benchmark_structured_dot_grad.py

[ricardoV94]

@tomicapretto My bot butler confirmed the speedups on my machine, both your scripts or using regular pytensor.function. Compile times are not meaningfully affected.

Only thing is that SpSum’s cache version was not bumped. Happy to approve after that

[tomicapretto]

Done, thanks for the review!

[ricardoV94]

It's a bit more readable if you assign cache_version=2 and then return the named variable.

[tomicapretto]

Done, thanks. Will merge once CI is green.