Is block-spmm supported？

[mzy2240]

https://developer.nvidia.com/blog/accelerating-matrix-multiplication-with-block-sparse-format-and-nvidia-tensor-cores/

Aside cusparselt, just curious whether block-spmm is supported? Seems to ship with cusparse.

[maleadt]

Please don't use the issue tracker for questions.

---

In general, all library APIs are supported, the question is at what level:

• C wrappers: functions like cusparseSpMM can always be called directly, as CUDA.jl exposes the library's C APIs:

  CUDA.jl/lib/cusparse/libcusparse.jl

  Lines 6545 to 6603 in 108e75f

  	@cenum cusparseSpMMAlg_t::UInt32 begin
  	CUSPARSE_MM_ALG_DEFAULT = 0
  	CUSPARSE_COOMM_ALG1 = 1
  	CUSPARSE_COOMM_ALG2 = 2
  	CUSPARSE_COOMM_ALG3 = 3
  	CUSPARSE_CSRMM_ALG1 = 4
  	CUSPARSE_SPMM_ALG_DEFAULT = 0
  	CUSPARSE_SPMM_COO_ALG1 = 1
  	CUSPARSE_SPMM_COO_ALG2 = 2
  	CUSPARSE_SPMM_COO_ALG3 = 3
  	CUSPARSE_SPMM_COO_ALG4 = 5
  	CUSPARSE_SPMM_CSR_ALG1 = 4
  	CUSPARSE_SPMM_CSR_ALG2 = 6
  	CUSPARSE_SPMM_CSR_ALG3 = 12
  	CUSPARSE_SPMM_BLOCKED_ELL_ALG1 = 13
  	end
  	@checked function cusparseSpMM_bufferSize(handle, opA, opB, alpha, matA, matB, beta, matC,
  	computeType, alg, bufferSize)
  	initialize_context()
  	@ccall libcusparse.cusparseSpMM_bufferSize(handle::cusparseHandle_t,
  	opA::cusparseOperation_t,
  	opB::cusparseOperation_t,
  	alpha::PtrOrCuPtr{Cvoid},
  	matA::cusparseSpMatDescr_t,
  	matB::cusparseDnMatDescr_t,
  	beta::PtrOrCuPtr{Cvoid},
  	matC::cusparseDnMatDescr_t,
  	computeType::cudaDataType,
  	alg::cusparseSpMMAlg_t,
  	bufferSize::Ptr{Csize_t})::cusparseStatus_t
  	end
  	@checked function cusparseSpMM_preprocess(handle, opA, opB, alpha, matA, matB, beta, matC,
  	computeType, alg, externalBuffer)
  	initialize_context()
  	@ccall libcusparse.cusparseSpMM_preprocess(handle::cusparseHandle_t,
  	opA::cusparseOperation_t,
  	opB::cusparseOperation_t,
  	alpha::PtrOrCuPtr{Cvoid},
  	matA::cusparseSpMatDescr_t,
  	matB::cusparseDnMatDescr_t,
  	beta::PtrOrCuPtr{Cvoid},
  	matC::cusparseDnMatDescr_t,
  	computeType::cudaDataType,
  	alg::cusparseSpMMAlg_t,
  	externalBuffer::CuPtr{Cvoid})::cusparseStatus_t
  	end
  	@checked function cusparseSpMM(handle, opA, opB, alpha, matA, matB, beta, matC, computeType,
  	alg, externalBuffer)
  	initialize_context()
  	@ccall libcusparse.cusparseSpMM(handle::cusparseHandle_t, opA::cusparseOperation_t,
  	opB::cusparseOperation_t, alpha::PtrOrCuPtr{Cvoid},
  	matA::cusparseSpMatDescr_t, matB::cusparseDnMatDescr_t,
  	beta::PtrOrCuPtr{Cvoid}, matC::cusparseDnMatDescr_t,
  	computeType::cudaDataType, alg::cusparseSpMMAlg_t,
  	externalBuffer::CuPtr{Cvoid})::cusparseStatus_t
  	end

• Julia wrappers: to make functions more user-friendly (e.g., making them type-generic) we often provide slightly high level wrappers. In the case of cusparseSpMM, that's the mm! function:

  CUDA.jl/lib/cusparse/generic.jl

  Lines 177 to 294 in 108e75f

  	function mm!(transa::SparseChar, transb::SparseChar, alpha::Number, A::Union{CuSparseMatrixCSC{T},CuSparseMatrixCSR{T},CuSparseMatrixCOO{T}},
  	B::DenseCuMatrix{T}, beta::Number, C::DenseCuMatrix{T}, index::SparseChar, algo::cusparseSpMMAlg_t=CUSPARSE_MM_ALG_DEFAULT) where {T}
  	# Support transa = 'C' and `transb = 'C' for real matrices
  	transa = T <: Real && transa == 'C' ? 'T' : transa
  	transb = T <: Real && transb == 'C' ? 'T' : transb
  	# Issue # 1610
  	if isa(A, CuSparseMatrixCSC) && transa == 'C' && T <: Complex
  	throw(ArgumentError("Matrix-matrix multiplication with the adjoint of a complex CSC matrix" *
  	" is not supported. Use a CSR or COO matrix instead."))
  	end
  	if isa(A, CuSparseMatrixCSC)
  	# cusparseSpMM doesn't support CSC matrices with CUSPARSE.version() < v"11.6.1".
  	# cusparseSpMM supports the CSC matrices with CUSPARSE.version() ≥ v"11.6.1"
  	# but it doesn't work for complex numbers when transa == 'C'.
  	# We use Aᵀ to model them as CSR matrices for now.
  	# It should be fixed with CUDA toolkit v"12.0".
  	descA = CuSparseMatrixDescriptor(A, index, transposed=true)
  	k,m = size(A)
  	transa = transa == 'N' ? 'T' : 'N'
  	else
  	descA = CuSparseMatrixDescriptor(A, index)
  	m,k = size(A)
  	end
  	n = size(C)[2]
  	if transa == 'N' && transb == 'N'
  	chkmmdims(B,C,k,n,m,n)
  	elseif transa == 'N' && transb != 'N'
  	chkmmdims(B,C,n,k,m,n)
  	elseif transa != 'N' && transb == 'N'
  	chkmmdims(B,C,m,n,k,n)
  	elseif transa != 'N' && transb != 'N'
  	chkmmdims(B,C,n,m,k,n)
  	end
  	descB = CuDenseMatrixDescriptor(B)
  	descC = CuDenseMatrixDescriptor(C)
  	function bufferSize()
  	out = Ref{Csize_t}()
  	cusparseSpMM_bufferSize(
  	handle(), transa, transb, Ref{T}(alpha), descA, descB, Ref{T}(beta),
  	descC, T, algo, out)
  	return out[]
  	end
  	with_workspace(bufferSize) do buffer
  	# Uncomment if we find a way to reuse the buffer (issue #1362)
  	# cusparseSpMM_preprocess(
  	# handle(), transa, transb, Ref{T}(alpha), descA, descB, Ref{T}(beta),
  	# descC, T, algo, buffer)
  	# end
  	cusparseSpMM(
  	handle(), transa, transb, Ref{T}(alpha), descA, descB, Ref{T}(beta),
  	descC, T, algo, buffer)
  	end
  	return C
  	end
  	function mm!(transa::SparseChar, transb::SparseChar, alpha::Number, A::DenseCuMatrix{T},
  	B::Union{CuSparseMatrixCSC{T},CuSparseMatrixCSR{T},CuSparseMatrixCOO{T}},
  	beta::Number, C::DenseCuMatrix{T}, index::SparseChar, algo::cusparseSpMMAlg_t=CUSPARSE_MM_ALG_DEFAULT) where {T}
  	CUSPARSE.version() < v"11.7.4" && throw(ErrorException("This operation is not supported by the current CUDA version."))
  	# Support transa = 'C' and `transb = 'C' for real matrices
  	transa = T <: Real && transa == 'C' ? 'T' : transa
  	transb = T <: Real && transb == 'C' ? 'T' : transb
  	# cusparseSpMM can be also used to perform the multiplication of a dense matrix and a sparse matrix by switching the dense matrices layout:
  	# Cc = α * Ac * B + β * Cc → α * Bᵀ * Ar + β * Cr
  	# Cc = α * Ac * Bᵀ + β * Cc → α * B * Ar + β * Cr
  	# Cc = α * Ac * Bᴴ + β * Cc → α * B̅ * Ar + β * Cr
  	# where B is a sparse matrix, Ac and Cc indicate column-major layout, while Ar and Cr refer to row-major layout.
  	# Issue # 1610
  	if isa(B, CuSparseMatrixCSR) && transb == 'C' && T <: Complex
  	throw(ArgumentError("Matrix-matrix multiplication with the adjoint of a complex CSR matrix is not supported. Use a CSC or COO matrix instead."))
  	end
  	m,k = size(A)
  	n = size(C)[2]
  	if transa == 'N' && transb == 'N'
  	chkmmdims(B,C,k,n,m,n)
  	elseif transa == 'N' && transb != 'N'
  	chkmmdims(B,C,n,k,m,n)
  	elseif transa != 'N' && transb == 'N'
  	chkmmdims(B,C,m,n,k,n)
  	elseif transa != 'N' && transb != 'N'
  	chkmmdims(B,C,n,m,k,n)
  	end
  	descA = CuDenseMatrixDescriptor(A, transposed=true)
  	descB = CuSparseMatrixDescriptor(B, index, transposed=true)
  	descC = CuDenseMatrixDescriptor(C, transposed=true)
  	function bufferSize()
  	out = Ref{Csize_t}()
  	cusparseSpMM_bufferSize(
  	handle(), transb, transa, Ref{T}(alpha), descB, descA, Ref{T}(beta),
  	descC, T, algo, out)
  	return out[]
  	end
  	with_workspace(bufferSize) do buffer
  	# Uncomment if we find a way to reuse the buffer (issue #1362)
  	# cusparseSpMM_preprocess(
  	# handle(), transb, transa, Ref{T}(alpha), descB, descA, Ref{T}(beta),
  	# descC, T, algo, buffer)
  	# end
  	cusparseSpMM(
  	handle(), transb, transa, Ref{T}(alpha), descB, descA, Ref{T}(beta),
  	descC, T, algo, buffer)
  	end
  	return C
  	end

• interface integrations: at the highest level of abstraction, CUDA.jl tries to integrate with libraries like LinearAlgebra.jl. For mm! that's LinearAlgebra.mul!, which in certain cases should dispatch to cusparseSpMM:

  CUDA.jl/lib/cusparse/interfaces.jl

  Lines 32 to 103 in 7e58f28

  	function mv_wrapper(transa::SparseChar, alpha::Number, A::CuSparseMatrix, X::DenseCuVector{T},
  	beta::Number, Y::CuVector{T}) where {T}
  	mv!(transa, alpha, A, X, beta, Y, 'O')
  	end
  	function mm_wrapper(transa::SparseChar, transb::SparseChar, alpha::Number,
  	A::CuSparseMatrix{T}, B::CuMatrix{T}, beta::Number, C::CuMatrix{T}) where {T}
  	n_A, m_A = (transa != 'N') ? reverse(size(A)) : size(A)
  	n_B, m_B = (transb != 'N') ? reverse(size(B)) : size(B)
  	n_C, m_C = size(C)
  	m_A == n_B || throw(DimensionMismatch())
  	n_A == n_C || throw(DimensionMismatch())
  	m_B == m_C || throw(DimensionMismatch())
  	isempty(B) && return CUDA.zeros(eltype(B), size(A, 1), 0)
  	if version() <= v"10.3.1"
  	# Generic mm! doesn't support transposed B on CUDA10
  	return mm2!(transa, transb, alpha, A, B, beta, C, 'O')
  	end
  	mm!(transa, transb, alpha, A, B, beta, C, 'O')
  	end
  	LinearAlgebra.dot(x::CuSparseVector{T}, y::DenseCuVector{T}) where {T <: BlasReal} = vv!('N', x, y, 'O')
  	LinearAlgebra.dot(x::DenseCuVector{T}, y::CuSparseVector{T}) where {T <: BlasReal} = dot(y, x)
  	LinearAlgebra.dot(x::CuSparseVector{T}, y::DenseCuVector{T}) where {T <: BlasComplex} = vv!('C', x, y, 'O')
  	LinearAlgebra.dot(x::DenseCuVector{T}, y::CuSparseVector{T}) where {T <: BlasComplex} = conj(dot(y,x))
  	tag_wrappers = ((identity, identity),
  	(T -> :(HermOrSym{T, <:$T}), A -> :(parent($A))))
  	adjtrans_wrappers = ((identity, identity),
  	(M -> :(Transpose{T, <:$M}), M -> :(_sptranspose(parent($M)))),
  	(M -> :(Adjoint{T, <:$M}), M -> :(_spadjoint(parent($M)))))
  	op_wrappers = ((identity, T -> 'N', identity),
  	(T -> :(Transpose{<:T, <:$T}), T -> 'T', A -> :(parent($A))),
  	(T -> :(Adjoint{<:T, <:$T}), T -> T <: Real ? 'T' : 'C', A -> :(parent($A))))
  	for (taga, untaga) in tag_wrappers, (wrapa, transa, unwrapa) in op_wrappers
  	TypeA = wrapa(taga(:(CuSparseMatrix{T})))
  	@eval begin
  	function LinearAlgebra.mul!(C::CuVector{T}, A::$TypeA, B::DenseCuVector{T},
  	alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}}
  	mv_wrapper($transa(T), alpha, $(untaga(unwrapa(:A))), B, beta, C)
  	end
  	function LinearAlgebra.mul!(C::CuVector{T}, A::$TypeA, B::CuSparseVector{T},
  	alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}}
  	mv_wrapper($transa(T), alpha, $(untaga(unwrapa(:A))), CuVector{T}(B), beta, C)
  	end
  	function LinearAlgebra.:(*)(A::$TypeA, x::CuSparseVector{T}) where {T <: Union{Float16, ComplexF16, BlasFloat}}
  	m, n = size(A)
  	length(x) == n || throw(DimensionMismatch())
  	y = CuVector{T}(undef, m)
  	mul!(y, A, x)
  	end
  	end
  	for (tagb, untagb) in tag_wrappers, (wrapb, transb, unwrapb) in op_wrappers
  	TypeB = wrapb(tagb(:(DenseCuMatrix{T})))
  	@eval begin
  	function LinearAlgebra.mul!(C::CuMatrix{T}, A::$TypeA, B::$TypeB,
  	alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}}
  	mm_wrapper($transa(T), $transb(T), alpha, $(untaga(unwrapa(:A))), $(untagb(unwrapb(:B))), beta, C)
  	end
  	end
  	end
  	end

What's required for block sparse gemm to work is the new BlockedEll storage format. That isn't supported by any of the high level wrappers, so you'd need to use the low-level cusparseCreateBlockedEll API, or provide your own integrations with the higher-level types (mimicking the support for CSR/CSC/COO/BSR storage types).