pydata / pydata/sparse

Faster matmul for case sparse.coo * np.ndarray

Open
#653 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
668
Forks
141
Avg merge
2d 8h
Merged PRs (30d)
4

Description

I am using this module quite heavily in a recent project and found that sparse.matmul could work faster.
I have implemented a faster version for the case matmul(sparse.coo, np.ndarray) using numba for the case of 2D arrays (cut be generalized). I also implemented a parallellized Version. The single threaded version is roughly 3 times faster than sparse.matmul in this case and scales linearly with the number of entries in the sparse matrix. The multi-threaded Version has some overhead but scales well for very large matrices (10000000+ entries).

Here is the code:

@numba.njit()
def matmul_fast(dim1:np.ndarray,dim2:np.ndarray,V:np.ndarray,Arr: np.ndarray,out: np.ndarray):
    for i in range(V.shape[0]):
        out[dim1[i]] += V[i] * Arr[dim2[i]]
    return out
@numba.njit(parallel=True,fastmath=True,boundscheck=False)
def matmul_fast_parallel(dim1:np.ndarray,dim2:np.ndarray,V:np.ndarray,Arr: np.ndarray,out: np.ndarray,n: np.int64 = 100000):
    f = V.shape[0]
    s = out.shape
    out_loc = np.zeros(s,dtype=np.float64)
    g = 0
    i = 0
    for i_0 in numba.prange(0,f//n+1):
        i = i_0 * n
        out_loc = np.zeros(s,dtype=np.float64)
        for g in range(i,i+n):
            if g>=f:
                break
            out_loc[dim1[g]] += V[g] * Arr[dim2[g]]
        out += out_loc
    return out

These functions can be driven by a non numba function that unpacks the sparse array object into coordinates, data and shape.

def matmul_sparse(Arr_sparse, Arr):
    out = np.zeros((Arr_sparse.shape[0],Arr.shape[1]))
    return matmul_fast_parallel(Arr_sparse.coords[0],Arr_sparse.coords[1],Arr_sparse.data,Arr,out)

Are there any optimizations planned that improve the performance of sparse.matmul?
If not would there be interest in introducing specific solutions like this as an option in the current implementation?

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the current sparse.matmul entry point and how it handles sparse.coo multiplied by a NumPy ndarray. Compare that path with the proposed single-threaded and parallel implementations; the work is complete only when the project has agreed on an implementation scope and verified performance improvements for the stated 2D case.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.