CuSparseArrayCSR (N dim array) with batched matmatmul (bmm) - #1944

Interesting! cc @amontoison

The only restriction I'm aware of is that the number of non-zeros of each matrix slices (batch-element) must be the same. The benefit of this representation is that we can more easily make use of CUDA's batched sparse mat muls, ex. Ci = Ai * Bi, which i've implemented for sparse * dense batched matmul in lib/cusparse/generic.jl (see bmm!).

If the representation isn't a fully generic 3D SparseArray, and the main use case is batched operations, why not a hypothetical Batched{CuSparseMatrix}? I'm concerned that users may think they can do more with a CuSparseArray type than they actually can.

If the representation isn't a fully generic 3D SparseArray, and the main use case is batched operations, why not a hypothetical Batched{CuSparseMatrix}? I'm concerned that users may think they can do more with a CuSparseArray type than they actually can.

Sure, that makes sense to me. Let me know if I'm understanding you correctly with this below,

mutable struct Batched{<: CuSparseMatrixCSR, Tv, Ti} <: AbstractCuSparseArray{Tv, Ti, 3}
    rowPtr::CuMatrix{Ti}
    colVal::CuMatrix{Ti}
    nzVal::CuMatrix{Tv}
    dims::NTuple{3,Int}
    nnz::Ti

    function Batched{A, Tv, Ti}(rowPtr::CuMatrix{<:Integer}, colVal::CuMatrix{<:Integer}, nzVal::CuMatrix, dims::NTuple{3,<:Integer}) where {Tv, Ti<:Integer}
        new{A, Tv, Ti}(rowPtr, colVal, nzVal, dims, length(nzVal))
    end
end

@nikopj Nice!
I like you proposition with the Batched structure.
We should be able to use it with different sparse matrices.

Getting this going again.

I've made things more general by allowing the "batch" dimension of the CuSparseArrayCSR{Tv,Ti,N} to be several dimensions (N-2 batch dims). I'm motiavted to do this by some deep-learning sparse attention stuff, where we might have different sparse attention matrices per mini-batch element, and each mini-batch element might want to make use of several separate attention matrices (ex. multi-head self attention).

So N=3 is the same case as before, but N>3 is now also possible. I've made it such that batched sparse-dense matmul also works for N>3 if the sizes make sense. See the end of test/libraries/cusparse/bmm.jl for an example size (m,n,2,3) CuSparseArrayCSR.

Because of the extended dimensions, I'm in favor of keeping the original naming convention. It also makes sense to keep things specific to the sparsity type as that dictates the fields of the struct. If we end up implementing a similar type for CSC, COO, etc, we could do a "Batched" union type for all of them.

The printing / showing / indexing of CuSparseArrayCSR is also working better, making testing on the REPL easier.

Base.cat and Base.reshape are also working for sensible arguments.

CI failures look related.

@maleadt

CI failures look related.

It's failing on CUDA 11.4 and 11.5, + the julia nightly build (unrelated to cuSparse).
I'm not sure where the error lies in these versions as they API says it can handle batched sparse-dense mul.

A quick fix would be to only define bmm! for versions 11.6 and higher.

A quick fix would be to only define bmm! for versions 11.6 and higher.

At the very least, yes. You could make the tests check for the cuSPARSE version (see the start of the CI logs), and maybe even add an error to the function itself.

The nightly issue is unrelated indeed.

According to Aqua, this added a couple of ambiguities. Not sure why CI didn't spot those...

julia> ambs = Aqua.detect_ambiguities(CUDA; recursive=true)
 (kwcall(::NamedTuple, ::typeof(cat), As::CUDA.CUSPARSE.CuSparseMatrixCSR...) @ CUDA.CUSPARSE ~/Julia/pkg/CUDA/lib/cusparse/batched.jl:1, kwcall(::NamedTuple, ::typeof(cat), As::CUDA.CUSPARSE.CuSparseArrayCSR...) @ CUDA.CUSPARSE ~/Julia/pkg/CUDA/lib/cusparse/batched.jl:14)
 (cat(As::CUDA.CUSPARSE.CuSparseMatrixCSR...; dims) @ CUDA.CUSPARSE ~/Julia/pkg/CUDA/lib/cusparse/batched.jl:1, cat(As::CUDA.CUSPARSE.CuSparseArrayCSR...; dims) @ CUDA.CUSPARSE ~/Julia/pkg/CUDA/lib/cusparse/batched.jl:14)

Could you fix those?

Ok I'll take a look.