Skip to content

fixes the kron implementation for sparse + diagonal matrix #2804

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: master
Choose a base branch
from
+119 −63
Open

fixes the kron implementation for sparse + diagonal matrix #2804

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
be66023
fix + tests
tam724 Jun 27, 2025
58123fa
use adapt(..)
tam724 Jun 30, 2025
5e51422
support adjoint(::Diagonal{<:CuVector{<:Complex}}) (e.g. mul! resulte…
tam724 Jul 14, 2025
c12cd91
Implement kron for {Adjoint, Transpose}(CuSparseMatrixCOO) and Diagon…
tam724 Jul 14, 2025
25b4648
Merge branch 'master' into fix_kron_cusparse_diag
tam724 Jul 14, 2025
fd04f1c
Reduce tests and test only Diagonal{Bool}
tam724 Jul 14, 2025
6195763
Merge branch 'master' into fix_kron_cusparse_diag
maleadt Jul 29, 2025
20608e3
Revert "support adjoint(::Diagonal{<:CuVector{<:Complex}}) (e.g. mul!…
tam724 Jul 29, 2025
e041932
change the type signature of kron to also match adjoint(Diagonal(.., …
tam724 Jul 29, 2025
f5baedb
populate coo-matrix manually for kron() tests
tam724 Jul 30, 2025
d675c5b
delete methods for kron with a cpu diagonal
tam724 Aug 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 66 additions & 58 deletions lib/cusparse/linalg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,79 +57,87 @@ function Base.reshape(A::CuSparseMatrixCOO, dims::Dims)
sparse(new_row, new_col, A.nzVal, dims[1], length(dims) == 1 ? 1 : dims[2], fmt = :coo)
end

function LinearAlgebra.kron(A::CuSparseMatrixCOO{T, Ti}, B::CuSparseMatrixCOO{T, Ti}) where {Ti, T}
mA,nA = size(A)
mB,nB = size(B)
out_shape = (mA * mB, nA * nB)
Annz = Int64(A.nnz)
Bnnz = Int64(B.nnz)

if Annz == 0 || Bnnz == 0
return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape)
# deprecate? CUSPARSE should not do this internally.
LinearAlgebra.kron(A::Union{CuSparseMatrixCOO{TvA, TiA}, Transpose{TvA, <:CuSparseMatrixCOO{TvA, TiA}}, Adjoint{TvA, <:CuSparseMatrixCOO{TvA, TiA}}}, B::Diagonal{TvB, <:Vector{TvB}}) where {TvA, TiA, TvB} = kron(A, adapt(CuArray, B))
LinearAlgebra.kron(A::Diagonal{TvA, <:Vector{TvA}}, B::Union{CuSparseMatrixCOO{TvB, TiB}, Transpose{TvB, <:CuSparseMatrixCOO{TvB, TiB}}, Adjoint{TvB, <:CuSparseMatrixCOO{TvB, TiB}}}) where {TvA, TvB, TiB} = kron(adapt(CuArray, A), B)

_kron_CuSparseMatrixCOO_components(A::CuSparseMatrixCOO) = A.rowInd, A.colInd, A.nzVal, identity, Int(A.nnz)
_kron_CuSparseMatrixCOO_components(At::Transpose{<:Number, <:CuSparseMatrixCOO}) = parent(At).colInd, parent(At).rowInd, parent(At).nzVal, transpose, Int(parent(At).nnz)
_kron_CuSparseMatrixCOO_components(Ah::Adjoint{<:Number, <:CuSparseMatrixCOO}) = parent(Ah).colInd, parent(Ah).rowInd, parent(Ah).nzVal, adjoint, Int(parent(Ah).nnz)

function LinearAlgebra.kron(
A::Union{CuSparseMatrixCOO{TvA, TiA}, Transpose{TvA, <:CuSparseMatrixCOO{TvA, TiA}}, Adjoint{TvA, <:CuSparseMatrixCOO{TvA, TiA}}},
B::Union{CuSparseMatrixCOO{TvB, TiB}, Transpose{TvB, <:CuSparseMatrixCOO{TvB, TiB}}, Adjoint{TvB, <:CuSparseMatrixCOO{TvB, TiB}}}
) where {TvA, TiA, TvB, TiB}
mA, nA = size(A)
mB, nB = size(B)
Ti = promote_type(TiA, TiB)
Tv = typeof(oneunit(TvA)*oneunit(TvB))

A_rowInd, A_colInd, A_nzVal, A_nzOp, A_nnz = _kron_CuSparseMatrixCOO_components(A)
B_rowInd, B_colInd, B_nzVal, B_nzOp, B_nnz = _kron_CuSparseMatrixCOO_components(B)

if A_nnz == 0 || B_nnz == 0
return CuSparseMatrixCOO{Tv, Ti}(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{Tv}(undef, 0), (mA * mB, nA * nB))
end

row = (A.rowInd .- 1) .* mB
row = repeat(row, inner = Bnnz)
col = (A.colInd .- 1) .* nB
col = repeat(col, inner = Bnnz)
data = repeat(A.nzVal, inner = Bnnz)
C_nnz = A_nnz * B_nnz
C_rowInd = reshape(B_rowInd .+ Ti(mB) .* (reshape(A_rowInd, (1, A_nnz)) .- one(Ti)), C_nnz)
C_colInd = reshape(B_colInd .+ Ti(nB) .* (reshape(A_colInd, (1, A_nnz)) .- one(Ti)), C_nnz)
C_nzVal = reshape(B_nzOp.(B_nzVal) .* A_nzOp.(reshape(A_nzVal, (1, A_nnz))), C_nnz)

row .+= repeat(B.rowInd .- 1, outer = Annz) .+ 1
col .+= repeat(B.colInd .- 1, outer = Annz) .+ 1

data .*= repeat(B.nzVal, outer = Annz)

sparse(row, col, data, out_shape..., fmt = :coo)
C = CuSparseMatrixCOO{Tv, Ti}(C_rowInd, C_colInd, C_nzVal, (mA * mB, nA * nB), C_nnz)
return sort_coo(C)
end

function LinearAlgebra.kron(A::CuSparseMatrixCOO{T, Ti}, B::Diagonal) where {Ti, T}
mA,nA = size(A)
mB,nB = size(B)
out_shape = (mA * mB, nA * nB)
Annz = Int64(A.nnz)
Bnnz = nB

if Annz == 0 || Bnnz == 0
return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape)
end
function LinearAlgebra.kron(
A::Union{CuSparseMatrixCOO{TvA, TiA}, Transpose{TvA, <:CuSparseMatrixCOO{TvA, TiA}}, Adjoint{TvA, <:CuSparseMatrixCOO{TvA, TiA}}},
B::Diagonal{TvB, <:Union{CuVector{TvB}, Base.ReshapedArray{TvB, 1, <:Adjoint{TvB, <:CuVector{TvB}}}}}
) where {TvA, TiA, TvB}
mA, nA = size(A)
mB, nB = size(B)
Ti = TiA
Tv = typeof(oneunit(TvA)*oneunit(TvB))

row = (A.rowInd .- 1) .* mB
row = repeat(row, inner = Bnnz)
col = (A.colInd .- 1) .* nB
col = repeat(col, inner = Bnnz)
data = repeat(A.nzVal, inner = Bnnz)
A_rowInd, A_colInd, A_nzVal, A_nzOp, A_nnz = _kron_CuSparseMatrixCOO_components(A)
B_rowInd, B_colInd, B_nzVal, B_nnz = one(Ti):Ti(nB), one(Ti):Ti(nB), B.diag, Int(nB)

row .+= CuVector(repeat(0:nB-1, outer = Annz)) .+ 1
col .+= CuVector(repeat(0:nB-1, outer = Annz)) .+ 1
if A_nnz == 0 || B_nnz == 0
return CuSparseMatrixCOO{Tv, Ti}(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{Tv}(undef, 0), (mA * mB, nA * nB))
end

data .*= repeat(CUDA.ones(T, nB), outer = Annz)
C_nnz = A_nnz * B_nnz
C_rowInd = reshape(B_rowInd .+ Ti(mB) .* (reshape(A_rowInd, (1, A_nnz)) .- one(Ti)), C_nnz)
C_colInd = reshape(B_colInd .+ Ti(nB) .* (reshape(A_colInd, (1, A_nnz)) .- one(Ti)), C_nnz)
C_nzVal = reshape(B_nzVal .* A_nzOp.(reshape(A_nzVal, (1, A_nnz))), C_nnz)

sparse(row, col, data, out_shape..., fmt = :coo)
C = CuSparseMatrixCOO{Tv, Ti}(C_rowInd, C_colInd, C_nzVal, (mA * mB, nA * nB), C_nnz)
return sort_coo(C)
end

function LinearAlgebra.kron(A::Diagonal, B::CuSparseMatrixCOO{T, Ti}) where {Ti, T}
mA,nA = size(A)
mB,nB = size(B)
out_shape = (mA * mB, nA * nB)
Annz = nA
Bnnz = Int64(B.nnz)

if Annz == 0 || Bnnz == 0
return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape)
end
function LinearAlgebra.kron(
A::Diagonal{TvA, <:Union{CuVector{TvA}, Base.ReshapedArray{TvA, 1, <:Adjoint{TvA, <:CuVector{TvA}}}}},
B::Union{CuSparseMatrixCOO{TvB, TiB}, Transpose{TvB, <:CuSparseMatrixCOO{TvB, TiB}}, Adjoint{TvB, <:CuSparseMatrixCOO{TvB, TiB}}}
) where {TvA, TvB, TiB}
mA, nA = size(A)
mB, nB = size(B)
Ti = TiB
Tv = typeof(oneunit(TvA)*oneunit(TvB))

row = (0:nA-1) .* mB
row = CuVector(repeat(row, inner = Bnnz))
col = (0:nA-1) .* nB
col = CuVector(repeat(col, inner = Bnnz))
data = repeat(CUDA.ones(T, nA), inner = Bnnz)
A_rowInd, A_colInd, A_nzVal, A_nnz = one(Ti):Ti(nA), one(Ti):Ti(nA), A.diag, Int(nA)
B_rowInd, B_colInd, B_nzVal, B_nzOp, B_nnz = _kron_CuSparseMatrixCOO_components(B)

row .+= repeat(B.rowInd .- 1, outer = Annz) .+ 1
col .+= repeat(B.colInd .- 1, outer = Annz) .+ 1
if A_nnz == 0 || B_nnz == 0
return CuSparseMatrixCOO{Tv, Ti}(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{Tv}(undef, 0), (mA * mB, nA * nB))
end

data .*= repeat(B.nzVal, outer = Annz)
C_nnz = A_nnz * B_nnz
C_rowInd = reshape(B_rowInd .+ Ti(mB) .* (reshape(A_rowInd, (1, A_nnz)) .- one(Ti)), C_nnz)
C_colInd = reshape(B_colInd .+ Ti(nB) .* (reshape(A_colInd, (1, A_nnz)) .- one(Ti)), C_nnz)
C_nzVal = reshape(B_nzOp.(B_nzVal) .* reshape(A_nzVal, (1, A_nnz)), C_nnz)

sparse(row, col, data, out_shape..., fmt = :coo)
C = CuSparseMatrixCOO{Tv, Ti}(C_rowInd, C_colInd, C_nzVal, (mA * mB, nA * nB), C_nnz)
return sort_coo(C)
end

function LinearAlgebra.dot(y::CuVector{T}, A::CuSparseMatrixCSC{T}, x::CuVector{T}) where {T<:Union{BlasInt, BlasFloat}}
Expand Down
53 changes: 52 additions & 1 deletion test/libraries/cusparse/linalg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,18 @@
using CUDA.CUSPARSE
using LinearAlgebra, SparseArrays
using LinearAlgebra, SparseArrays, Adapt

m = 10
@testset "T = $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
A = sprand(T, m, m, 0.2)
B = sprand(T, m, m, 0.3)
ZA = spzeros(T, m, m)
C = I(div(m, 2))
D = Diagonal(rand(T, m))
@testset "type = $typ" for typ in [CuSparseMatrixCSR, CuSparseMatrixCSC]
dA = typ(A)
dB = typ(B)
dZA = typ(ZA)
dD = adapt(CuArray, D)
@testset "opnorm and norm" begin
@test opnorm(A, Inf) ≈ opnorm(dA, Inf)
@test opnorm(A, 1) ≈ opnorm(dA, 1)
Expand Down Expand Up @@ -42,6 +44,55 @@ m = 10
@test collect(kron(opa(dZA), C)) ≈ kron(opa(ZA), C)
@test collect(kron(C, opa(dZA))) ≈ kron(C, opa(ZA))
end
@testset "kronecker product with Diagonal opa = $opa" for opa in (identity, transpose, adjoint)
@test collect(kron(opa(dA), dD)) ≈ kron(opa(A), D)
@test collect(kron(dD, opa(dA))) ≈ kron(D, opa(A))
@test collect(kron(opa(dZA), dD)) ≈ kron(opa(ZA), D)
@test collect(kron(dD, opa(dZA))) ≈ kron(D, opa(ZA))
end
end
end

@testset "T = $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
mat_sizes = [(2, 3), (2, 0)]
@testset "size(A) = ($(mA), $(nA)), size(B) = ($(mB), $(nB))" for (mA, nA) in mat_sizes, (mB, nB) in mat_sizes
A = sprand(T, mA, nA, 0.5)
B = sprand(T, mB, nB, 0.5)

A_I, A_J, A_V = findnz(A)
dA = CuSparseMatrixCOO{T, Cint}(adapt(CuVector{Cint}, A_I), adapt(CuVector{Cint}, A_J), adapt(CuVector{T}, A_V), size(A))
B_I, B_J, B_V = findnz(B)
dB = CuSparseMatrixCOO{T, Cint}(adapt(CuVector{Cint}, B_I), adapt(CuVector{Cint}, B_J), adapt(CuVector{T}, B_V), size(B))

@testset "kronecker (COO ⊗ COO) opa = $opa, opb = $opb" for opa in (identity, transpose, adjoint), opb in (identity, transpose, adjoint)
dC = kron(opa(dA), opb(dB))
@test collect(dC) ≈ kron(opa(A), opb(B))
@test eltype(dC) == typeof(oneunit(T) * oneunit(T))
@test dC isa CuSparseMatrixCOO
end
end
end

@testset "TA = $TA, TvB = $TvB" for TvB in [Float32, Float64, ComplexF32, ComplexF64], TA in [Bool, TvB]
A = Diagonal(rand(TA, 2))
B = sprand(TvB, 3, 4, 0.5)
dA = adapt(CuArray, A)

B_I, B_J, B_V = findnz(B)
dB = CuSparseMatrixCOO{TvB, Cint}(adapt(CuVector{Cint}, B_I), adapt(CuVector{Cint}, B_J), adapt(CuVector{TvB}, B_V), size(B))

@testset "kronecker (diagonal ⊗ COO) opa = $opa, opb = $opb" for opa in (identity, adjoint), opb in (identity, transpose, adjoint)
dC = kron(opa(dA), opb(dB))
@test collect(dC) ≈ kron(opa(A), opb(B))
@test eltype(dC) == typeof(oneunit(TA) * oneunit(TvB))
@test dC isa CuSparseMatrixCOO
end

@testset "kronecker (COO ⊗ diagonal) opa = $opa, opb = $opb" for opa in (identity, adjoint), opb in (identity, transpose, adjoint)
dC = kron(opb(dB), opa(dA))
@test collect(dC) ≈ kron(opb(B), opa(A))
@test eltype(dC) == typeof(oneunit(TvB) * oneunit(TA))
@test dC isa CuSparseMatrixCOO
end
end

Expand Down
Loading