From 9c54dd6e5f42eafd05551320bd8f45c9a9e8586f Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sun, 1 Feb 2026 00:38:45 +0100 Subject: [PATCH 1/2] Add methods of kron with `Diagonal` --- Project.toml | 5 +- .../GenericSparseArraysFillArraysExt.jl | 16 ++ .../kernels.jl | 61 ++++ .../kron_coo.jl | 113 ++++++++ .../kron_csc.jl | 46 +++ .../kron_csr.jl | 46 +++ src/GenericSparseArrays.jl | 2 +- src/conversions/conversions.jl | 70 +++++ src/matrix_coo/matrix_coo.jl | 266 +++++++++++------- src/matrix_coo/matrix_coo_kernels.jl | 73 +++++ src/matrix_csc/matrix_csc.jl | 114 +++----- src/matrix_csr/matrix_csr.jl | 118 +++----- test/Project.toml | 2 + test/runtests.jl | 7 +- test/shared/matrix_coo.jl | 43 ++- test/shared/matrix_csc.jl | 47 +++- test/shared/matrix_csr.jl | 47 +++- 17 files changed, 801 insertions(+), 275 deletions(-) create mode 100644 ext/GenericSparseArraysFillArraysExt/GenericSparseArraysFillArraysExt.jl create mode 100644 ext/GenericSparseArraysFillArraysExt/kernels.jl create mode 100644 ext/GenericSparseArraysFillArraysExt/kron_coo.jl create mode 100644 ext/GenericSparseArraysFillArraysExt/kron_csc.jl create mode 100644 ext/GenericSparseArraysFillArraysExt/kron_csr.jl diff --git a/Project.toml b/Project.toml index 2854d2e..0ea5b98 100644 --- a/Project.toml +++ b/Project.toml @@ -1,11 +1,12 @@ name = "GenericSparseArrays" uuid = "da3fe0eb-88a8-4d14-ae1a-857c283e9c70" -authors = ["Alberto Mercurio and contributors"] version = "0.1.0" +authors = ["Alberto Mercurio and contributors"] [deps] AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -15,11 +16,13 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" [extensions] +GenericSparseArraysFillArraysExt = "FillArrays" GenericSparseArraysJLArraysExt = "JLArrays" [compat] AcceleratedKernels = "0.4" Adapt = "4" +FillArrays = "1" GPUArraysCore = "0.2.0" JLArrays = "0.3" KernelAbstractions = "0.9" diff --git a/ext/GenericSparseArraysFillArraysExt/GenericSparseArraysFillArraysExt.jl b/ext/GenericSparseArraysFillArraysExt/GenericSparseArraysFillArraysExt.jl new file mode 100644 index 0000000..a75394f --- /dev/null +++ b/ext/GenericSparseArraysFillArraysExt/GenericSparseArraysFillArraysExt.jl @@ -0,0 +1,16 @@ +module GenericSparseArraysFillArraysExt + +using GenericSparseArrays: GenericSparseMatrixCOO, GenericSparseMatrixCSC, GenericSparseMatrixCSR +using LinearAlgebra: Diagonal, Transpose, Adjoint, kron + +import FillArrays: AbstractFill, getindex_value +import KernelAbstractions: @kernel, @index, get_backend, synchronize +import LinearAlgebra: kron +import SparseArrays: nnz + +include("kernels.jl") +include("kron_coo.jl") +include("kron_csc.jl") +include("kron_csr.jl") + +end # module diff --git a/ext/GenericSparseArraysFillArraysExt/kernels.jl b/ext/GenericSparseArraysFillArraysExt/kernels.jl new file mode 100644 index 0000000..fabc6ef --- /dev/null +++ b/ext/GenericSparseArraysFillArraysExt/kernels.jl @@ -0,0 +1,61 @@ +# Optimized kernels for kron with Diagonal{AbstractFill} + +# Optimized kernel for kron(D, B) where D has uniform diagonal value (e.g., scaled identity) +@kernel inbounds = true function kernel_kron_scaled_identity_coo!( + rowind_C, + colind_C, + nzval_C, + @Const(scale_val), # The uniform diagonal value + @Const(n_D::Int), # Size of the diagonal matrix + @Const(rowind_B), + @Const(colind_B), + @Const(nzval_B), + @Const(m_B::Int), + @Const(n_B::Int), + ) + idx = @index(Global, Linear) + + nnz_B = length(nzval_B) + + if idx <= n_D * nnz_B + idx_D = div(idx - 1, nnz_B) + 1 + idx_B = mod(idx - 1, nnz_B) + 1 + + i_B = rowind_B[idx_B] + j_B = colind_B[idx_B] + val_B = nzval_B[idx_B] + + rowind_C[idx] = (idx_D - 1) * m_B + i_B + colind_C[idx] = (idx_D - 1) * n_B + j_B + nzval_C[idx] = scale_val * val_B + end +end + +# Optimized kernel for kron(A, D) where D has uniform diagonal value (e.g., scaled identity) +@kernel inbounds = true function kernel_kron_coo_scaled_identity!( + rowind_C, + colind_C, + nzval_C, + @Const(rowind_A), + @Const(colind_A), + @Const(nzval_A), + @Const(scale_val), # The uniform diagonal value + @Const(p::Int), # Size of the diagonal matrix + ) + idx = @index(Global, Linear) + + nnz_A = length(nzval_A) + + if idx <= nnz_A * p + idx_A = div(idx - 1, p) + 1 + idx_D = mod(idx - 1, p) + 1 + + i_A = rowind_A[idx_A] + j_A = colind_A[idx_A] + val_A = nzval_A[idx_A] + + rowind_C[idx] = (i_A - 1) * p + idx_D + colind_C[idx] = (j_A - 1) * p + idx_D + nzval_C[idx] = val_A * scale_val + end +end diff --git a/ext/GenericSparseArraysFillArraysExt/kron_coo.jl b/ext/GenericSparseArraysFillArraysExt/kron_coo.jl new file mode 100644 index 0000000..9bf9fc3 --- /dev/null +++ b/ext/GenericSparseArraysFillArraysExt/kron_coo.jl @@ -0,0 +1,113 @@ +# Kronecker product with Diagonal{AbstractFill} for COO format + +using GenericSparseArrays: trans_adj_wrappers + +function kron( + D::Diagonal{Tv1, <:AbstractFill{Tv1}}, + B::GenericSparseMatrixCOO{Tv2, Ti}, + ) where {Tv1, Tv2, Ti} + n_D = size(D, 1) + m_B, n_B = size(B) + + # Result dimensions + m_C = n_D * m_B + n_C = n_D * n_B + nnz_C = n_D * nnz(B) + + # Determine result types + Tv = promote_type(Tv1, Tv2) + + backend = get_backend(B) + + # Allocate output arrays + rowind_C = similar(B.rowind, Ti, nnz_C) + colind_C = similar(B.colind, Ti, nnz_C) + nzval_C = similar(B.nzval, Tv, nnz_C) + + # Get the uniform fill value + fill_value = getindex_value(D.diag) + + # Launch optimized kernel + kernel! = kernel_kron_scaled_identity_coo!(backend) + kernel!( + rowind_C, + colind_C, + nzval_C, + fill_value, + n_D, + B.rowind, + B.colind, + B.nzval, + m_B, + n_B; + ndrange = nnz_C, + ) + + return GenericSparseMatrixCOO(m_C, n_C, rowind_C, colind_C, nzval_C) +end + +function kron( + A::GenericSparseMatrixCOO{Tv1, Ti}, + D::Diagonal{Tv2, <:AbstractFill{Tv2}}, + ) where {Tv1, Ti, Tv2} + m_A, n_A = size(A) + p = size(D, 1) # D is p×p + + # Result dimensions + m_C = m_A * p + n_C = n_A * p + nnz_C = nnz(A) * p + + # Determine result types + Tv = promote_type(Tv1, Tv2) + + backend = get_backend(A) + + # Allocate output arrays + rowind_C = similar(A.rowind, Ti, nnz_C) + colind_C = similar(A.colind, Ti, nnz_C) + nzval_C = similar(A.nzval, Tv, nnz_C) + + # Get the uniform fill value + fill_value = getindex_value(D.diag) + + # Launch optimized kernel + kernel! = kernel_kron_coo_scaled_identity!(backend) + kernel!( + rowind_C, + colind_C, + nzval_C, + A.rowind, + A.colind, + A.nzval, + fill_value, + p; + ndrange = nnz_C, + ) + + return GenericSparseMatrixCOO(m_C, n_C, rowind_C, colind_C, nzval_C) +end + +# kron with Diagonal{AbstractFill} and transpose/adjoint wrappers for COO +for (wrap, trans, conj, unwrap, whereT) in trans_adj_wrappers(:GenericSparseMatrixCOO) + trans == false && continue + + TypeB = wrap(:(T)) + + @eval function kron( + D::Diagonal{Tv1, <:AbstractFill{Tv1}}, + B::$TypeB, + ) where {Tv1, $(whereT(:T))} + B_coo = GenericSparseMatrixCOO(B) + return kron(D, B_coo) + end + + TypeA = wrap(:(T)) + @eval function kron( + A::$TypeA, + D::Diagonal{Tv2, <:AbstractFill{Tv2}}, + ) where {$(whereT(:T)), Tv2} + A_coo = GenericSparseMatrixCOO(A) + return kron(A_coo, D) + end +end diff --git a/ext/GenericSparseArraysFillArraysExt/kron_csc.jl b/ext/GenericSparseArraysFillArraysExt/kron_csc.jl new file mode 100644 index 0000000..d5ca065 --- /dev/null +++ b/ext/GenericSparseArraysFillArraysExt/kron_csc.jl @@ -0,0 +1,46 @@ +# Kronecker product with Diagonal{AbstractFill} for CSC format + +using GenericSparseArrays: trans_adj_wrappers + +function kron( + D::Diagonal{Tv, <:AbstractFill{Tv}}, + B::GenericSparseMatrixCSC, + ) where {Tv} + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(D, B_coo) + return GenericSparseMatrixCSC(C_coo) +end + +function kron( + A::GenericSparseMatrixCSC, + D::Diagonal{Tv, <:AbstractFill{Tv}}, + ) where {Tv} + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSC(C_coo) +end + +for (wrap, trans, conj, unwrap, whereT) in trans_adj_wrappers(:GenericSparseMatrixCSC) + trans == false && continue + + TypeB = wrap(:(T)) + + @eval function kron( + D::Diagonal{Tv1, <:AbstractFill{Tv1}}, + B::$TypeB, + ) where {Tv1, $(whereT(:T))} + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(D, B_coo) + return GenericSparseMatrixCSC(C_coo) + end + + TypeA = wrap(:(T)) + @eval function kron( + A::$TypeA, + D::Diagonal{Tv2, <:AbstractFill{Tv2}}, + ) where {$(whereT(:T)), Tv2} + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSC(C_coo) + end +end diff --git a/ext/GenericSparseArraysFillArraysExt/kron_csr.jl b/ext/GenericSparseArraysFillArraysExt/kron_csr.jl new file mode 100644 index 0000000..e7ca9f6 --- /dev/null +++ b/ext/GenericSparseArraysFillArraysExt/kron_csr.jl @@ -0,0 +1,46 @@ +# Kronecker product with Diagonal{AbstractFill} for CSR format + +using GenericSparseArrays: trans_adj_wrappers + +function kron( + D::Diagonal{Tv, <:AbstractFill{Tv}}, + B::GenericSparseMatrixCSR, + ) where {Tv} + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(D, B_coo) + return GenericSparseMatrixCSR(C_coo) +end + +function kron( + A::GenericSparseMatrixCSR, + D::Diagonal{Tv, <:AbstractFill{Tv}}, + ) where {Tv} + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSR(C_coo) +end + +for (wrap, trans, conj, unwrap, whereT) in trans_adj_wrappers(:GenericSparseMatrixCSR) + trans == false && continue + + TypeB = wrap(:(T)) + + @eval function kron( + D::Diagonal{Tv1, <:AbstractFill{Tv1}}, + B::$TypeB, + ) where {Tv1, $(whereT(:T))} + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(D, B_coo) + return GenericSparseMatrixCSR(C_coo) + end + + TypeA = wrap(:(T)) + @eval function kron( + A::$TypeA, + D::Diagonal{Tv2, <:AbstractFill{Tv2}}, + ) where {$(whereT(:T)), Tv2} + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSR(C_coo) + end +end diff --git a/src/GenericSparseArrays.jl b/src/GenericSparseArrays.jl index 08b6a1d..2d8bf2c 100644 --- a/src/GenericSparseArrays.jl +++ b/src/GenericSparseArrays.jl @@ -1,7 +1,7 @@ module GenericSparseArrays using LinearAlgebra -import LinearAlgebra: wrap, copymutable_oftype, __normalize!, kron +import LinearAlgebra: wrap, copymutable_oftype, __normalize!, kron, Diagonal using SparseArrays import SparseArrays: SparseVector, SparseMatrixCSC import SparseArrays: getcolptr, getrowval, getnzval, nonzeroinds diff --git a/src/conversions/conversions.jl b/src/conversions/conversions.jl index f9465d7..54b0096 100644 --- a/src/conversions/conversions.jl +++ b/src/conversions/conversions.jl @@ -122,6 +122,30 @@ function GenericSparseMatrixCSR( ) end +# ============================================================================ +# Transpose and Adjoint conversions for COO +# ============================================================================ + +GenericSparseMatrixCOO(A::GenericSparseMatrixCOO) = A + +GenericSparseMatrixCOO(A::Transpose{Tv, <:GenericSparseMatrixCOO}) where {Tv} = + GenericSparseMatrixCOO( + size(A, 1), + size(A, 2), + A.parent.colind, + A.parent.rowind, + A.parent.nzval, +) + +GenericSparseMatrixCOO(A::Adjoint{Tv, <:GenericSparseMatrixCOO}) where {Tv} = + GenericSparseMatrixCOO( + size(A, 1), + size(A, 2), + A.parent.colind, + A.parent.rowind, + conj.(A.parent.nzval), +) + # ============================================================================ # CSC ↔ COO Conversions # ============================================================================ @@ -201,6 +225,29 @@ GenericSparseMatrixCSC(A::Adjoint{Tv, <:GenericSparseMatrixCOO}) where {Tv} = ) ) +# Transpose and Adjoint conversions for CSC to COO +function GenericSparseMatrixCOO(A::Transpose{Tv, <:GenericSparseMatrixCSC}) where {Tv} + parent_coo = GenericSparseMatrixCOO(A.parent) + return GenericSparseMatrixCOO( + size(A, 1), + size(A, 2), + parent_coo.colind, + parent_coo.rowind, + parent_coo.nzval, + ) +end + +function GenericSparseMatrixCOO(A::Adjoint{Tv, <:GenericSparseMatrixCSC}) where {Tv} + parent_coo = GenericSparseMatrixCOO(A.parent) + return GenericSparseMatrixCOO( + size(A, 1), + size(A, 2), + parent_coo.colind, + parent_coo.rowind, + conj.(parent_coo.nzval), + ) +end + # ============================================================================ # CSR ↔ COO Conversions # ============================================================================ @@ -256,3 +303,26 @@ function GenericSparseMatrixCSR(A::GenericSparseMatrixCOO{Tv, Ti}) where {Tv, Ti return GenericSparseMatrixCSR(m, n, rowptr, colind_sorted, nzval_sorted) end + +# Transpose and Adjoint conversions for CSR to COO +function GenericSparseMatrixCOO(A::Transpose{Tv, <:GenericSparseMatrixCSR}) where {Tv} + parent_coo = GenericSparseMatrixCOO(A.parent) + return GenericSparseMatrixCOO( + size(A, 1), + size(A, 2), + parent_coo.colind, + parent_coo.rowind, + parent_coo.nzval, + ) +end + +function GenericSparseMatrixCOO(A::Adjoint{Tv, <:GenericSparseMatrixCSR}) where {Tv} + parent_coo = GenericSparseMatrixCOO(A.parent) + return GenericSparseMatrixCOO( + size(A, 1), + size(A, 2), + parent_coo.colind, + parent_coo.rowind, + conj.(parent_coo.nzval), + ) +end diff --git a/src/matrix_coo/matrix_coo.jl b/src/matrix_coo/matrix_coo.jl index 9aedaf2..0ff2afb 100644 --- a/src/matrix_coo/matrix_coo.jl +++ b/src/matrix_coo/matrix_coo.jl @@ -309,31 +309,6 @@ function _add_sparse_to_dense!(C::DenseMatrix, A::GenericSparseMatrixCOO) return C end -""" - +(A::GenericSparseMatrixCOO, B::GenericSparseMatrixCOO) - -Add two sparse matrices in COO format. Both matrices must have the same dimensions -and be on the same backend (device). - -The result is a COO matrix with entries from both A and B properly merged, -with duplicate entries (same row and column) combined by summing their values. - -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays - -julia> A = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [1.0, 2.0], 2, 2)); - -julia> B = GenericSparseMatrixCOO(sparse([1, 2], [2, 1], [3.0, 4.0], 2, 2)); - -julia> C = A + B; - -julia> collect(C) -2×2 Matrix{Float64}: - 1.0 3.0 - 4.0 2.0 -``` -""" function Base.:+(A::GenericSparseMatrixCOO, B::GenericSparseMatrixCOO) size(A) == size(B) || throw( DimensionMismatch( @@ -538,32 +513,63 @@ for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSpars end end -""" - kron(A::GenericSparseMatrixCOO, B::GenericSparseMatrixCOO) +function Base.:(*)(A::GenericSparseMatrixCOO, B::GenericSparseMatrixCOO) + size(A, 2) == size(B, 1) || throw( + DimensionMismatch( + "second dimension of A, $(size(A, 2)), does not match first dimension of B, $(size(B, 1))", + ), + ) -Compute the Kronecker product of two sparse matrices in COO format. + backend_A = get_backend(A) + backend_B = get_backend(B) + backend_A == backend_B || + throw(ArgumentError("Both matrices must have the same backend")) -The Kronecker product of two matrices `A` (size m×n) and `B` (size p×q) -is an (m*p)×(n*q) matrix formed by multiplying each element of `A` by the -entire matrix `B`. + # Convert to CSC, multiply, convert back to COO + # This is acceptable as COO doesn't have an efficient direct multiplication algorithm + # and CSC provides the sorted structure needed for efficient SpGEMM + A_csc = GenericSparseMatrixCSC(A) + B_csc = GenericSparseMatrixCSC(B) + C_csc = A_csc * B_csc + return GenericSparseMatrixCOO(C_csc) +end -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays +# Multiplication with transpose/adjoint support - all cases use the same approach +for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSparseMatrixCOO) + for (wrapb, transb, conjb, unwrapb, whereT2) in + trans_adj_wrappers(:GenericSparseMatrixCOO) + # Skip the case where both are not transposed (already handled above) + (transa == false && transb == false) && continue -julia> A = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [1.0, 2.0], 2, 2)); + TypeA = wrapa(:(T1)) + TypeB = wrapb(:(T2)) -julia> B = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [3.0, 4.0], 2, 2)); + @eval function Base.:(*)( + A::$TypeA, + B::$TypeB, + ) where {$(whereT1(:T1)), $(whereT2(:T2))} + size(A, 2) == size(B, 1) || throw( + DimensionMismatch( + "second dimension of A, $(size(A, 2)), does not match first dimension of B, $(size(B, 1))", + ), + ) -julia> C = kron(A, B); + backend_A = get_backend($(unwrapa(:A))) + backend_B = get_backend($(unwrapb(:B))) + backend_A == backend_B || + throw(ArgumentError("Both matrices must have the same backend")) -julia> size(C) -(4, 4) + # Convert to CSC (handles transpose/adjoint), multiply, convert back to COO + # Same approach as the base case since COO doesn't have an efficient + # direct multiplication algorithm + A_csc = GenericSparseMatrixCSC(A) + B_csc = GenericSparseMatrixCSC(B) + C_csc = A_csc * B_csc + return GenericSparseMatrixCOO(C_csc) + end + end +end -julia> nnz(C) -4 -``` -""" function LinearAlgebra.kron( A::GenericSparseMatrixCOO{Tv1, Ti1}, B::GenericSparseMatrixCOO{Tv2, Ti2}, @@ -607,85 +613,129 @@ function LinearAlgebra.kron( return GenericSparseMatrixCOO(m_C, n_C, rowind_C, colind_C, nzval_C) end -""" - *(A::GenericSparseMatrixCOO, B::GenericSparseMatrixCOO) +for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSparseMatrixCOO) + for (wrapb, transb, conjb, unwrapb, whereT2) in trans_adj_wrappers(:GenericSparseMatrixCOO) + # Skip the case where both are not transposed (already handled above) + (transa == false && transb == false) && continue -Multiply two sparse matrices in COO format. Both matrices must have compatible dimensions -(number of columns of A equals number of rows of B) and be on the same backend (device). + TypeA = wrapa(:(T1)) + TypeB = wrapb(:(T2)) -The multiplication converts to CSC format, performs the multiplication with GPU-compatible -kernels, and converts back to COO format. This approach is used for all cases including -transpose/adjoint since COO doesn't have an efficient direct multiplication algorithm. + @eval function LinearAlgebra.kron( + A::$TypeA, + B::$TypeB, + ) where {$(whereT1(:T1)), $(whereT2(:T2))} + return kron(GenericSparseMatrixCOO(A), GenericSparseMatrixCOO(B)) + end + end +end -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays +# ============================================================================= +# Kronecker product with Diagonal matrices +# ============================================================================= -julia> A = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [2.0, 3.0], 2, 2)); +function LinearAlgebra.kron( + D::Diagonal{Tv1}, + B::GenericSparseMatrixCOO{Tv2, Ti}, + ) where {Tv1, Tv2, Ti} + n_D = size(D, 1) + m_B, n_B = size(B) -julia> B = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [4.0, 5.0], 2, 2)); + # Result dimensions + m_C = n_D * m_B + n_C = n_D * n_B + nnz_C = n_D * nnz(B) -julia> C = A * B; + # Determine result types + Tv = promote_type(Tv1, Tv2) -julia> collect(C) -2×2 Matrix{Float64}: - 8.0 0.0 - 0.0 15.0 -``` -""" -function Base.:(*)(A::GenericSparseMatrixCOO, B::GenericSparseMatrixCOO) - size(A, 2) == size(B, 1) || throw( - DimensionMismatch( - "second dimension of A, $(size(A, 2)), does not match first dimension of B, $(size(B, 1))", - ), - ) + backend = get_backend(B) - backend_A = get_backend(A) - backend_B = get_backend(B) - backend_A == backend_B || - throw(ArgumentError("Both matrices must have the same backend")) + # Allocate output arrays + rowind_C = similar(B.rowind, Ti, nnz_C) + colind_C = similar(B.colind, Ti, nnz_C) + nzval_C = similar(B.nzval, Tv, nnz_C) - # Convert to CSC, multiply, convert back to COO - # This is acceptable as COO doesn't have an efficient direct multiplication algorithm - # and CSC provides the sorted structure needed for efficient SpGEMM - A_csc = GenericSparseMatrixCSC(A) - B_csc = GenericSparseMatrixCSC(B) - C_csc = A_csc * B_csc - return GenericSparseMatrixCOO(C_csc) + # Launch kernel + kernel! = kernel_kron_diag_coo!(backend) + kernel!( + rowind_C, + colind_C, + nzval_C, + D.diag, + B.rowind, + B.colind, + B.nzval, + m_B, + n_B; + ndrange = nnz_C, + ) + + return GenericSparseMatrixCOO(m_C, n_C, rowind_C, colind_C, nzval_C) end -# Multiplication with transpose/adjoint support - all cases use the same approach -for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSparseMatrixCOO) - for (wrapb, transb, conjb, unwrapb, whereT2) in - trans_adj_wrappers(:GenericSparseMatrixCOO) - # Skip the case where both are not transposed (already handled above) - (transa == false && transb == false) && continue +function LinearAlgebra.kron( + A::GenericSparseMatrixCOO{Tv1, Ti}, + D::Diagonal{Tv2}, + ) where {Tv1, Ti, Tv2} + m_A, n_A = size(A) + p = size(D, 1) # D is p×p - TypeA = wrapa(:(T1)) - TypeB = wrapb(:(T2)) + # Result dimensions + m_C = m_A * p + n_C = n_A * p + nnz_C = nnz(A) * p - @eval function Base.:(*)( - A::$TypeA, - B::$TypeB, - ) where {$(whereT1(:T1)), $(whereT2(:T2))} - size(A, 2) == size(B, 1) || throw( - DimensionMismatch( - "second dimension of A, $(size(A, 2)), does not match first dimension of B, $(size(B, 1))", - ), - ) + # Determine result types + Tv = promote_type(Tv1, Tv2) - backend_A = get_backend($(unwrapa(:A))) - backend_B = get_backend($(unwrapb(:B))) - backend_A == backend_B || - throw(ArgumentError("Both matrices must have the same backend")) + backend = get_backend(A) - # Convert to CSC (handles transpose/adjoint), multiply, convert back to COO - # Same approach as the base case since COO doesn't have an efficient - # direct multiplication algorithm - A_csc = GenericSparseMatrixCSC(A) - B_csc = GenericSparseMatrixCSC(B) - C_csc = A_csc * B_csc - return GenericSparseMatrixCOO(C_csc) - end + # Allocate output arrays + rowind_C = similar(A.rowind, Ti, nnz_C) + colind_C = similar(A.colind, Ti, nnz_C) + nzval_C = similar(A.nzval, Tv, nnz_C) + + # Launch kernel + kernel! = kernel_kron_coo_diag!(backend) + kernel!( + rowind_C, + colind_C, + nzval_C, + A.rowind, + A.colind, + A.nzval, + D.diag, + p; + ndrange = nnz_C, + ) + + return GenericSparseMatrixCOO(m_C, n_C, rowind_C, colind_C, nzval_C) +end + +# kron with Diagonal and transpose/adjoint wrappers +for (wrap, trans, conj, unwrap, whereT) in trans_adj_wrappers(:GenericSparseMatrixCOO) + # Skip identity case (already handled above) + trans == false && continue + + TypeB = wrap(:(T)) + + # kron(D, op(B)) + @eval function LinearAlgebra.kron( + D::Diagonal{Tv1}, + B::$TypeB, + ) where {Tv1, $(whereT(:T))} + B_coo = GenericSparseMatrixCOO(B) + return kron(D, B_coo) + end + + # kron(op(A), D) + TypeA = wrap(:(T)) + @eval function LinearAlgebra.kron( + A::$TypeA, + D::Diagonal{Tv2}, + ) where {$(whereT(:T)), Tv2} + A_coo = GenericSparseMatrixCOO(A) + return kron(A_coo, D) end end diff --git a/src/matrix_coo/matrix_coo_kernels.jl b/src/matrix_coo/matrix_coo_kernels.jl index 11f85a0..6d8e08a 100644 --- a/src/matrix_coo/matrix_coo_kernels.jl +++ b/src/matrix_coo/matrix_coo_kernels.jl @@ -182,6 +182,79 @@ end end end +# Kernel for computing Kronecker product of Diagonal (left) and COO (right) +# kron(D, B) where D is n×n diagonal creates n diagonal blocks of B scaled by D[i,i] +@kernel inbounds = true function kernel_kron_diag_coo!( + rowind_C, + colind_C, + nzval_C, + @Const(diag_vals), # Diagonal values of D + @Const(rowind_B), + @Const(colind_B), + @Const(nzval_B), + @Const(m_B::Int), + @Const(n_B::Int), + ) + idx = @index(Global, Linear) + + n_D = length(diag_vals) + nnz_B = length(nzval_B) + + if idx <= n_D * nnz_B + # Compute which diagonal element and which B element we're combining + idx_D = div(idx - 1, nnz_B) + 1 + idx_B = mod(idx - 1, nnz_B) + 1 + + d_val = diag_vals[idx_D] + + i_B = rowind_B[idx_B] + j_B = colind_B[idx_B] + val_B = nzval_B[idx_B] + + # Result position: block (idx_D, idx_D) contains d_val * B + # Row: (idx_D - 1) * m_B + i_B + # Col: (idx_D - 1) * n_B + j_B + rowind_C[idx] = (idx_D - 1) * m_B + i_B + colind_C[idx] = (idx_D - 1) * n_B + j_B + nzval_C[idx] = d_val * val_B + end +end + +# Kernel for computing Kronecker product of COO (left) and Diagonal (right) +# kron(A, D) where D is p×p diagonal: each A[i,j] creates a p×p diagonal block scaled by A[i,j] +@kernel inbounds = true function kernel_kron_coo_diag!( + rowind_C, + colind_C, + nzval_C, + @Const(rowind_A), + @Const(colind_A), + @Const(nzval_A), + @Const(diag_vals), # Diagonal values of D + @Const(p::Int), # Size of D (p×p) + ) + idx = @index(Global, Linear) + + nnz_A = length(nzval_A) + + if idx <= nnz_A * p + # Compute which A element and which diagonal element we're combining + idx_A = div(idx - 1, p) + 1 + idx_D = mod(idx - 1, p) + 1 + + i_A = rowind_A[idx_A] + j_A = colind_A[idx_A] + val_A = nzval_A[idx_A] + d_val = diag_vals[idx_D] + + # Result position: A[i_A, j_A] creates block at (i_A, j_A) with diagonal entry at (idx_D, idx_D) + # Row: (i_A - 1) * p + idx_D + # Col: (j_A - 1) * p + idx_D + rowind_C[idx] = (i_A - 1) * p + idx_D + colind_C[idx] = (j_A - 1) * p + idx_D + nzval_C[idx] = val_A * d_val + end +end + # Kernel for marking duplicate entries in sorted COO format # Returns a mask where mask[i] = true if entry i should be kept (first occurrence or sum) @kernel inbounds = true function kernel_mark_unique_coo!( diff --git a/src/matrix_csc/matrix_csc.jl b/src/matrix_csc/matrix_csc.jl index 8ecd273..07bb76e 100644 --- a/src/matrix_csc/matrix_csc.jl +++ b/src/matrix_csc/matrix_csc.jl @@ -301,28 +301,6 @@ function _add_sparse_to_dense!(C::DenseMatrix, A::GenericSparseMatrixCSC) return C end -""" - +(A::GenericSparseMatrixCSC, B::GenericSparseMatrixCSC) - -Add two sparse matrices in CSC format. Both matrices must have the same dimensions -and be on the same backend (device). - -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays - -julia> A = GenericSparseMatrixCSC(sparse([1, 2], [1, 2], [1.0, 2.0], 2, 2)); - -julia> B = GenericSparseMatrixCSC(sparse([1, 2], [2, 1], [3.0, 4.0], 2, 2)); - -julia> C = A + B; - -julia> collect(C) -2×2 Matrix{Float64}: - 1.0 3.0 - 4.0 2.0 -``` -""" function Base.:+(A::GenericSparseMatrixCSC, B::GenericSparseMatrixCSC) size(A) == size(B) || throw( DimensionMismatch( @@ -419,64 +397,62 @@ for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSpars end end -""" - kron(A::GenericSparseMatrixCSC, B::GenericSparseMatrixCSC) - -Compute the Kronecker product of two sparse matrices in CSC format. - -The Kronecker product is computed by converting to COO format, computing the -product, and converting back to CSC format. - -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays - -julia> A = GenericSparseMatrixCSC(sparse([1, 2], [1, 2], [1.0, 2.0], 2, 2)); - -julia> B = GenericSparseMatrixCSC(sparse([1, 2], [1, 2], [3.0, 4.0], 2, 2)); - -julia> C = kron(A, B); +for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSparseMatrixCSC) + for (wrapb, transb, conjb, unwrapb, whereT2) in trans_adj_wrappers(:GenericSparseMatrixCSC) + TypeA = wrapa(:(T1)) + TypeB = wrapb(:(T2)) -julia> size(C) -(4, 4) + @eval function LinearAlgebra.kron( + A::$TypeA, + B::$TypeB, + ) where {$(whereT1(:T1)), $(whereT2(:T2))} + # Convert to COO, compute kron, convert back to CSC + A_coo = GenericSparseMatrixCOO(A) + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(A_coo, B_coo) + return GenericSparseMatrixCSC(C_coo) + end + end +end -julia> nnz(C) -4 -``` -""" -function LinearAlgebra.kron(A::GenericSparseMatrixCSC, B::GenericSparseMatrixCSC) - # Convert to COO, compute kron, convert back to CSC - A_coo = GenericSparseMatrixCOO(A) +function LinearAlgebra.kron(D::Diagonal, B::GenericSparseMatrixCSC) B_coo = GenericSparseMatrixCOO(B) - C_coo = kron(A_coo, B_coo) + C_coo = kron(D, B_coo) return GenericSparseMatrixCSC(C_coo) end -""" - *(A::GenericSparseMatrixCSC, B::GenericSparseMatrixCSC) - -Multiply two sparse matrices in CSC format. Both matrices must have compatible dimensions -(number of columns of A equals number of rows of B) and be on the same backend (device). - -The multiplication uses GPU-compatible kernels for efficient sparse-sparse matrix -multiplication (SpGEMM). +function LinearAlgebra.kron(A::GenericSparseMatrixCSC, D::Diagonal) + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSC(C_coo) +end -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays +# kron with Diagonal and transpose/adjoint wrappers for CSC +for (wrap, trans, conj, unwrap, whereT) in trans_adj_wrappers(:GenericSparseMatrixCSC) + trans == false && continue -julia> A = GenericSparseMatrixCSC(sparse([1, 2], [1, 2], [2.0, 3.0], 2, 2)); + TypeB = wrap(:(T)) -julia> B = GenericSparseMatrixCSC(sparse([1, 2], [1, 2], [4.0, 5.0], 2, 2)); + @eval function LinearAlgebra.kron( + D::Diagonal{Tv1}, + B::$TypeB, + ) where {Tv1, $(whereT(:T))} + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(D, B_coo) + return GenericSparseMatrixCSC(C_coo) + end -julia> C = A * B; + TypeA = wrap(:(T)) + @eval function LinearAlgebra.kron( + A::$TypeA, + D::Diagonal{Tv2}, + ) where {$(whereT(:T)), Tv2} + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSC(C_coo) + end +end -julia> collect(C) -2×2 Matrix{Float64}: - 8.0 0.0 - 0.0 15.0 -``` -""" function Base.:(*)(A::GenericSparseMatrixCSC, B::GenericSparseMatrixCSC) size(A, 2) == size(B, 1) || throw( DimensionMismatch( diff --git a/src/matrix_csr/matrix_csr.jl b/src/matrix_csr/matrix_csr.jl index dbf6934..aaa5107 100644 --- a/src/matrix_csr/matrix_csr.jl +++ b/src/matrix_csr/matrix_csr.jl @@ -299,28 +299,6 @@ function _add_sparse_to_dense!(C::DenseMatrix, A::GenericSparseMatrixCSR) return C end -""" - +(A::GenericSparseMatrixCSR, B::GenericSparseMatrixCSR) - -Add two sparse matrices in CSR format. Both matrices must have the same dimensions -and be on the same backend (device). - -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays - -julia> A = GenericSparseMatrixCSR(sparse([1, 2], [1, 2], [1.0, 2.0], 2, 2)); - -julia> B = GenericSparseMatrixCSR(sparse([1, 2], [2, 1], [3.0, 4.0], 2, 2)); - -julia> C = A + B; - -julia> collect(C) -2×2 Matrix{Float64}: - 1.0 3.0 - 4.0 2.0 -``` -""" function Base.:+(A::GenericSparseMatrixCSR, B::GenericSparseMatrixCSR) size(A) == size(B) || throw( DimensionMismatch( @@ -417,68 +395,62 @@ for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSpars end end -""" - kron(A::GenericSparseMatrixCSR, B::GenericSparseMatrixCSR) - -Compute the Kronecker product of two sparse matrices in CSR format. - -The Kronecker product is computed by converting to COO format, computing the -product, and converting back to CSR format. - -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays - -julia> A_coo = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [1.0, 2.0], 2, 2)); - -julia> B_coo = GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [3.0, 4.0], 2, 2)); - -julia> A = GenericSparseMatrixCSR(A_coo); - -julia> B = GenericSparseMatrixCSR(B_coo); - -julia> C = kron(A, B); +for (wrapa, transa, conja, unwrapa, whereT1) in trans_adj_wrappers(:GenericSparseMatrixCSR) + for (wrapb, transb, conjb, unwrapb, whereT2) in trans_adj_wrappers(:GenericSparseMatrixCSR) + TypeA = wrapa(:(T1)) + TypeB = wrapb(:(T2)) -julia> size(C) -(4, 4) + @eval function LinearAlgebra.kron( + A::$TypeA, + B::$TypeB, + ) where {$(whereT1(:T1)), $(whereT2(:T2))} + # Convert to COO, compute kron, convert back to CSR + A_coo = GenericSparseMatrixCOO(A) + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(A_coo, B_coo) + return GenericSparseMatrixCSR(C_coo) + end + end +end -julia> nnz(C) -4 -``` -""" -function LinearAlgebra.kron(A::GenericSparseMatrixCSR, B::GenericSparseMatrixCSR) - # Convert to COO, compute kron, convert back to CSR - A_coo = GenericSparseMatrixCOO(A) +function LinearAlgebra.kron(D::Diagonal, B::GenericSparseMatrixCSR) B_coo = GenericSparseMatrixCOO(B) - C_coo = kron(A_coo, B_coo) + C_coo = kron(D, B_coo) return GenericSparseMatrixCSR(C_coo) end -""" - *(A::GenericSparseMatrixCSR, B::GenericSparseMatrixCSR) - -Multiply two sparse matrices in CSR format. Both matrices must have compatible dimensions -(number of columns of A equals number of rows of B) and be on the same backend (device). - -The multiplication uses GPU-compatible kernels for efficient sparse-sparse matrix -multiplication (SpGEMM). +function LinearAlgebra.kron(A::GenericSparseMatrixCSR, D::Diagonal) + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSR(C_coo) +end -# Examples -```jldoctest -julia> using GenericSparseArrays, SparseArrays +# kron with Diagonal and transpose/adjoint wrappers for CSR +for (wrap, trans, conj, unwrap, whereT) in trans_adj_wrappers(:GenericSparseMatrixCSR) + trans == false && continue -julia> A = GenericSparseMatrixCSR(GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [2.0, 3.0], 2, 2))); + TypeB = wrap(:(T)) -julia> B = GenericSparseMatrixCSR(GenericSparseMatrixCOO(sparse([1, 2], [1, 2], [4.0, 5.0], 2, 2))); + @eval function LinearAlgebra.kron( + D::Diagonal{Tv1}, + B::$TypeB, + ) where {Tv1, $(whereT(:T))} + B_coo = GenericSparseMatrixCOO(B) + C_coo = kron(D, B_coo) + return GenericSparseMatrixCSR(C_coo) + end -julia> C = A * B; + TypeA = wrap(:(T)) + @eval function LinearAlgebra.kron( + A::$TypeA, + D::Diagonal{Tv2}, + ) where {$(whereT(:T)), Tv2} + A_coo = GenericSparseMatrixCOO(A) + C_coo = kron(A_coo, D) + return GenericSparseMatrixCSR(C_coo) + end +end -julia> collect(C) -2×2 Matrix{Float64}: - 8.0 0.0 - 0.0 15.0 -``` -""" function Base.:(*)(A::GenericSparseMatrixCSR, B::GenericSparseMatrixCSR) size(A, 2) == size(B, 1) || throw( DimensionMismatch( diff --git a/test/Project.toml b/test/Project.toml index 69013e3..e8cf42b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -9,6 +10,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] +FillArrays = "1" LinearAlgebra = "1" Pkg = "1" SparseArrays = "1" diff --git a/test/runtests.jl b/test/runtests.jl index 76a5cec..18d4b87 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,6 +4,7 @@ using JET using GenericSparseArrays using JLArrays using Adapt +using FillArrays using LinearAlgebra using SparseArrays @@ -103,7 +104,11 @@ if GROUP in ("All", "Code-Quality") include(joinpath(@__DIR__, "shared", "code_quality.jl")) @testset "Code quality (Aqua.jl)" begin ambiguities = true # VERSION > v"1.11" - Aqua.test_all(GenericSparseArrays; ambiguities = ambiguities) + Aqua.test_all( + GenericSparseArrays; + ambiguities = ambiguities, + # stale_deps = (; ignore = [:FillArrays]), # FillArrays is a weakdep for extension + ) end @testset "Code linting (JET.jl)" verbose = true begin diff --git a/test/shared/matrix_coo.jl b/test/shared/matrix_coo.jl index 2cfd701..e63b22d 100644 --- a/test/shared/matrix_coo.jl +++ b/test/shared/matrix_coo.jl @@ -355,17 +355,48 @@ function shared_test_linearalgebra_matrix_coo( # Test with rectangular matrices A_sparse = sprand(T, 30, 25, 0.1) B_sparse = sprand(T, 20, 15, 0.1) + D_diag = Diagonal(rand(T, 4)) A = adapt(op, GenericSparseMatrixCOO(A_sparse)) B = adapt(op, GenericSparseMatrixCOO(B_sparse)) + D1 = adapt(op, D_diag) + D2 = Diagonal(FillArrays.Fill(T(2), 4)) - C = kron(A, B) - C_expected = kron(A_sparse, B_sparse) + for (op_A, op_B) in Iterators.product( + (identity, transpose, adjoint), + (identity, transpose, adjoint), + ) - @test size(C) == size(C_expected) - @test nnz(C) == nnz(C_expected) - @test Matrix(SparseMatrixCSC(C)) ≈ Matrix(C_expected) - @test C isa GenericSparseMatrixCOO + C = kron(op_A(A), op_B(B)) + C_expected = kron(op_A(A_sparse), op_B(B_sparse)) + + @test size(C) == size(C_expected) + @test nnz(C) == nnz(C_expected) + @test Matrix(SparseMatrixCSC(C)) ≈ Matrix(C_expected) + @test C isa GenericSparseMatrixCOO + + # Test with Diagonal matrix + C = kron(D1, op_B(B)) + C_expected = kron(D_diag, op_B(B_sparse)) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCOO + + C = kron(op_A(A), D1) + C_expected = kron(op_A(A_sparse), D_diag) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCOO + + # Test with FillArrays Diagonal matrix + C = kron(D2, op_B(B)) + C_expected = kron(D2, op_B(B_sparse)) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCOO + + C = kron(op_A(A), D2) + C_expected = kron(op_A(A_sparse), D2) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCOO + end end end end diff --git a/test/shared/matrix_csc.jl b/test/shared/matrix_csc.jl index 1a16778..21c6c21 100644 --- a/test/shared/matrix_csc.jl +++ b/test/shared/matrix_csc.jl @@ -354,17 +354,48 @@ function shared_test_linearalgebra_matrix_csc( # Test with rectangular matrices A_sparse = SparseMatrixCSC{T, int_types[end]}(sprand(T, 30, 25, 0.1)) B_sparse = SparseMatrixCSC{T, int_types[end]}(sprand(T, 20, 15, 0.1)) + D_diag = Diagonal(rand(T, 4)) A = adapt(op, GenericSparseMatrixCSC(A_sparse)) B = adapt(op, GenericSparseMatrixCSC(B_sparse)) - - C = kron(A, B) - C_expected = kron(A_sparse, B_sparse) - - @test size(C) == size(C_expected) - @test nnz(C) == nnz(C_expected) - @test Matrix(SparseMatrixCSC(C)) ≈ Matrix(C_expected) - @test C isa GenericSparseMatrixCSC + D1 = adapt(op, D_diag) + D2 = Diagonal(FillArrays.Fill(T(2), 4)) + + for (op_A, op_B) in Iterators.product( + (identity, transpose, adjoint), + (identity, transpose, adjoint), + ) + + C = kron(op_A(A), op_B(B)) + C_expected = kron(op_A(A_sparse), op_B(B_sparse)) + + @test size(C) == size(C_expected) + @test nnz(C) == nnz(C_expected) + @test Matrix(SparseMatrixCSC(C)) ≈ Matrix(C_expected) + @test C isa GenericSparseMatrixCSC + + # Test with Diagonal matrix + C = kron(D1, op_B(B)) + C_expected = kron(D_diag, op_B(B_sparse)) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSC + + C = kron(op_A(A), D1) + C_expected = kron(op_A(A_sparse), D_diag) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSC + + # Test with FillArrays Diagonal matrix + C = kron(D2, op_B(B)) + C_expected = kron(D2, op_B(B_sparse)) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSC + + C = kron(op_A(A), D2) + C_expected = kron(op_A(A_sparse), D2) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSC + end end end end diff --git a/test/shared/matrix_csr.jl b/test/shared/matrix_csr.jl index 7577931..5ebc8d5 100644 --- a/test/shared/matrix_csr.jl +++ b/test/shared/matrix_csr.jl @@ -353,17 +353,48 @@ function shared_test_linearalgebra_matrix_csr( # Test with rectangular matrices A_sparse = SparseMatrixCSC{T, int_types[end]}(sprand(T, 30, 25, 0.1)) B_sparse = SparseMatrixCSC{T, int_types[end]}(sprand(T, 20, 15, 0.1)) + D_diag = Diagonal(rand(T, 4)) A = adapt(op, GenericSparseMatrixCSR(A_sparse)) B = adapt(op, GenericSparseMatrixCSR(B_sparse)) - - C = kron(A, B) - C_expected = kron(A_sparse, B_sparse) - - @test size(C) == size(C_expected) - @test nnz(C) == nnz(C_expected) - @test Matrix(SparseMatrixCSC(C)) ≈ Matrix(C_expected) - @test C isa GenericSparseMatrixCSR + D1 = adapt(op, D_diag) + D2 = Diagonal(FillArrays.Fill(T(2), 4)) + + for (op_A, op_B) in Iterators.product( + (identity, transpose, adjoint), + (identity, transpose, adjoint), + ) + + C = kron(op_A(A), op_B(B)) + C_expected = kron(op_A(A_sparse), op_B(B_sparse)) + + @test size(C) == size(C_expected) + @test nnz(C) == nnz(C_expected) + @test Matrix(SparseMatrixCSC(C)) ≈ Matrix(C_expected) + @test C isa GenericSparseMatrixCSR + + # Test with Diagonal matrix + C = kron(D1, op_B(B)) + C_expected = kron(D_diag, op_B(B_sparse)) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSR + + C = kron(op_A(A), D1) + C_expected = kron(op_A(A_sparse), D_diag) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSR + + # Test with FillArrays Diagonal matrix + C = kron(D2, op_B(B)) + C_expected = kron(D2, op_B(B_sparse)) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSR + + C = kron(op_A(A), D2) + C_expected = kron(op_A(A_sparse), D2) + @test SparseMatrixCSC(C) ≈ C_expected + @test C isa GenericSparseMatrixCSR + end end end end From c1d46c5b0b91953dd5a10806008d17f162f66559 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sun, 1 Feb 2026 00:46:50 +0100 Subject: [PATCH 2/2] Add FillArrays in staledeps --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 18d4b87..2ca9daf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -107,7 +107,7 @@ if GROUP in ("All", "Code-Quality") Aqua.test_all( GenericSparseArrays; ambiguities = ambiguities, - # stale_deps = (; ignore = [:FillArrays]), # FillArrays is a weakdep for extension + stale_deps = (; ignore = [:FillArrays]), # FillArrays is a weakdep for extension ) end