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 6 commits into
base: master
Choose a base branch
from
Open

fixes the kron implementation for sparse + diagonal matrix #2804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 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
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
2 changes: 2 additions & 0 deletions lib/cublas/linalg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,8 @@ function LinearAlgebra.inv(D::Diagonal{T, <:CuArray{T}}) where {T}
Diagonal(Di)
end

LinearAlgebra.adjoint(D::Diagonal{T, <:CuVector{T}}) where T <: Complex = Diagonal(map(adjoint, D.diag))

LinearAlgebra.rdiv!(A::CuArray, D::Diagonal) = _rdiv!(A, A, D)

Base.:/(A::CuArray, D::Diagonal) = _rdiv!(similar(A, typeof(oneunit(eltype(A)) / oneunit(eltype(D)))), A, D)
Expand Down
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, <: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_nzOp, B_nnz = one(Ti):Ti(nB), one(Ti):Ti(nB), B.diag, identity, 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_nzOp.(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, <: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_nzOp, A_nnz = one(Ti):Ti(nA), one(Ti):Ti(nA), A.diag, identity, 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) .* 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.dot(y::CuVector{T}, A::CuSparseMatrixCSC{T}, x::CuVector{T}) where {T<:Union{BlasInt, BlasFloat}}
Expand Down
12 changes: 12 additions & 0 deletions test/libraries/cublas/extensions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -538,11 +538,23 @@ k = 13
mul!(d_XA, d_X, d_A)
Array(d_XA) ≈ Diagonal(x) * A

XA = rand(elty,m,n)
d_XA = CuArray(XA)
d_X = Diagonal(d_x)
mul!(d_XA, d_X', d_A)
Array(d_XA) ≈ Diagonal(x)' * A

AY = rand(elty,m,n)
d_AY = CuArray(AY)
d_Y = Diagonal(d_y)
mul!(d_AY, d_A, d_Y)
Array(d_AY) ≈ A * Diagonal(y)

AY = rand(elty,m,n)
d_AY = CuArray(AY)
d_Y = Diagonal(d_y)
mul!(d_AY, d_A, d_Y')
Array(d_AY) ≈ A * Diagonal(y)'

YA = rand(elty,n,m)
d_YA = CuArray(YA)
Expand Down
49 changes: 48 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,51 @@ 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)

dA = CuSparseMatrixCOO{T}(A)
dB = CuSparseMatrixCOO{T}(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)
dB = CuSparseMatrixCOO{TvB}(B)

@testset "kronecker (diagonal ⊗ COO) opa = $opa, opb = $opb" for opa in (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, transpose, adjoint), opb in (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