diff --git a/Project.toml b/Project.toml index e568b052..c2d27182 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ArrayInterface" uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -version = "7.27.0" +version = "7.28.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/docs/src/sparsearrays.md b/docs/src/sparsearrays.md index 97fc9ba0..50d22058 100644 --- a/docs/src/sparsearrays.md +++ b/docs/src/sparsearrays.md @@ -11,6 +11,7 @@ These routines allow for improving sparse iteration and indexing. ArrayInterface.isstructured ArrayInterface.findstructralnz ArrayInterface.has_sparsestruct +ArrayInterface.same_sparsity_structure ``` ## Matrix Coloring diff --git a/ext/ArrayInterfaceCUDAExt.jl b/ext/ArrayInterfaceCUDAExt.jl index 7ba6216a..95b7b717 100644 --- a/ext/ArrayInterfaceCUDAExt.jl +++ b/ext/ArrayInterfaceCUDAExt.jl @@ -25,6 +25,19 @@ const CuSparseArray = Union{ } ArrayInterface.can_setindex(::Type{<:CuSparseArray}) = false +# GPU CSC/CSR do not subtype `SparseArrays.AbstractSparseMatrixCSC`, so they need their own +# structure comparison. The stored index arrays are `CuVector`s; `==` reduces on-device. +function ArrayInterface.same_sparsity_structure( + A::CUSPARSE.CuSparseMatrixCSC, B::CUSPARSE.CuSparseMatrixCSC + ) + size(A) == size(B) && A.colPtr == B.colPtr && A.rowVal == B.rowVal +end +function ArrayInterface.same_sparsity_structure( + A::CUSPARSE.CuSparseMatrixCSR, B::CUSPARSE.CuSparseMatrixCSR + ) + size(A) == size(B) && A.rowPtr == B.rowPtr && A.colVal == B.colVal +end + function ArrayInterface.promote_eltype( ::Type{<:CUDA.CuArray{T, N, M}}, ::Type{T2} ) where {T, N, M, T2} diff --git a/ext/ArrayInterfaceSparseArraysExt.jl b/ext/ArrayInterfaceSparseArraysExt.jl index 9f0e06b4..9451e2be 100644 --- a/ext/ArrayInterfaceSparseArraysExt.jl +++ b/ext/ArrayInterfaceSparseArraysExt.jl @@ -1,8 +1,9 @@ module ArrayInterfaceSparseArraysExt -import ArrayInterface: buffer, has_sparsestruct, issingular, findstructralnz, bunchkaufman_instance, DEFAULT_CHOLESKY_PIVOT, cholesky_instance, ldlt_instance, lu_instance, qr_instance +import ArrayInterface: buffer, has_sparsestruct, issingular, findstructralnz, same_sparsity_structure, bunchkaufman_instance, DEFAULT_CHOLESKY_PIVOT, cholesky_instance, ldlt_instance, lu_instance, qr_instance using ArrayInterface.LinearAlgebra using SparseArrays +using SparseArrays: AbstractSparseMatrixCSC, getcolptr, rowvals buffer(x::SparseMatrixCSC) = getfield(x, :nzval) buffer(x::SparseVector) = getfield(x, :nzval) @@ -14,6 +15,13 @@ function findstructralnz(x::SparseMatrixCSC) (rowind, colind) end +# Compares the stored index arrays directly (no `findnz` allocation), so it is cheap +# enough for per-iteration reuse checks. Covers any CPU CSC type implementing the +# `getcolptr`/`rowvals` interface. +function same_sparsity_structure(A::AbstractSparseMatrixCSC, B::AbstractSparseMatrixCSC) + Base.size(A) == Base.size(B) && getcolptr(A) == getcolptr(B) && rowvals(A) == rowvals(B) +end + function bunchkaufman_instance(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti} bunchkaufman(SparseMatrixCSC{Tv, Ti}(similar(A, 1, 1)), check = false) end diff --git a/src/ArrayInterface.jl b/src/ArrayInterface.jl index 0b2593a4..0ecfec03 100644 --- a/src/ArrayInterface.jl +++ b/src/ArrayInterface.jl @@ -341,6 +341,50 @@ function findstructralnz(x::Union{Tridiagonal, SymTridiagonal}) (rowind, colind) end +""" + same_sparsity_structure(A, B) -> Bool + +Return `true` when `A` and `B` store their nonzeros in exactly the same structural +positions, i.e. they share an identical sparsity pattern (and shape). This is the condition +under which values may be written into `B` reusing `A`'s stored layout (or a cached symbolic +factorization of `A` reused for `B`) without rebuilding the structure. + +Semantics by category: + + - Dense arrays store every position, so two of them share a structure exactly when they + have the same shape. + - Structured (`Diagonal`, `Bidiagonal`, `Tridiagonal`, `SymTridiagonal`) and sparse + (`SparseMatrixCSC`, GPU CSC/CSR) matrices compare by their pattern — shape (plus + `uplo`/format) or the stored index arrays (e.g. `colptr`/`rowval` for CSC). + - Two arrays whose base types differ (e.g. a structured matrix and a dense one, or two + different structured types) return `false`: a layout is not reused across categories, + so they are treated as "not the same structure". + - Two arrays of the *same* base type with no specialized method are **unknown**, not + `false`: rather than fabricate an answer, a `MethodError` is thrown so the type can + define its own method. (`has_sparsestruct` can be used to check applicability first.) + +Unlike `findstructralnz`, this performs no allocation on its fast paths, so it is suitable +for use in hot loops (for example, deciding per Jacobian evaluation whether a sparse +`jac_prototype`'s structure needs rebuilding). +""" +function same_sparsity_structure(A::AbstractArray, B::AbstractArray) + # Same base type but no specialized method: we genuinely do not know -- error rather + # than return a possibly-wrong `false`. Different base types cannot share a structural + # category, so those are `false`. + parameterless_type(A) === parameterless_type(B) && + throw(MethodError(same_sparsity_structure, (A, B))) + return false +end +same_sparsity_structure(A::DenseArray, B::DenseArray) = Base.size(A) == Base.size(B) +same_sparsity_structure(A::Diagonal, B::Diagonal) = Base.size(A) == Base.size(B) +function same_sparsity_structure(A::Bidiagonal, B::Bidiagonal) + Base.size(A) == Base.size(B) && A.uplo == B.uplo +end +# Base-type methods (not `A::T, B::T`) so differing type parameters, e.g. +# `Tridiagonal{Float64}` vs `Tridiagonal{Int}`, still compare as the same structure. +same_sparsity_structure(A::Tridiagonal, B::Tridiagonal) = Base.size(A) == Base.size(B) +same_sparsity_structure(A::SymTridiagonal, B::SymTridiagonal) = Base.size(A) == Base.size(B) + abstract type ColoringAlgorithm end """ @@ -1005,8 +1049,8 @@ end :cholesky_instance, :ldlt_instance, :lu_instance, :qr_instance, :svd_instance, # sparse arrays - :isstructured, :findstructralnz, :has_sparsestruct, :fast_matrix_colors, - :matrix_colors, + :isstructured, :findstructralnz, :has_sparsestruct, :same_sparsity_structure, + :fast_matrix_colors, :matrix_colors, # wrapping :is_forwarding_wrapper, :buffer, :parent_type, # tuples diff --git a/test/core.jl b/test/core.jl index 62cabf88..3ad091b5 100644 --- a/test/core.jl +++ b/test/core.jl @@ -216,6 +216,44 @@ end @test [STri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4,5,6,7,5,6,7] end +@testset "same_sparsity_structure" begin + sss = ArrayInterface.same_sparsity_structure + + # structured matrices: pattern is fixed by shape (plus uplo) + @test sss(Diagonal([1,2,3]), Diagonal([4,5,6])) + @test !sss(Diagonal([1,2,3]), Diagonal([4,5])) + @test sss(Bidiagonal([1,2,3],[7,8],:U), Bidiagonal([4,5,6],[9,1],:U)) + @test !sss(Bidiagonal([1,2,3],[7,8],:U), Bidiagonal([4,5,6],[9,1],:L)) + @test sss(Tridiagonal([1,2],[1,2,3],[4,5]), Tridiagonal([6,7],[8,9,1],[2,3])) + @test !sss(Tridiagonal([1,2],[1,2,3],[4,5]), Tridiagonal([1],[1,2],[4])) + @test sss(SymTridiagonal([1,2,3],[4,5]), SymTridiagonal([6,7,8],[9,1])) + + # same base type, differing type parameters ⇒ still the same structure + @test sss(Tridiagonal([1,2],[1,2,3],[4,5]), Tridiagonal([1.0,2],[1.0,2,3],[4.0,5])) + + # dense arrays: every position is structural, so same shape ⇒ same structure + @test sss(rand(3,3), rand(3,3)) + @test !sss(rand(3,3), rand(3,2)) + @test sss(rand(4), rand(4)) + + # different base types ⇒ false (not compared position-by-position) + @test !sss(Diagonal([1,2,3]), Bidiagonal([1,2,3],[7,8],:U)) + @test !sss(Diagonal([1,2,3]), rand(3,3)) + @test !sss(Tridiagonal([1,2],[1,2,3],[4,5]), SymTridiagonal([1,2,3],[4,5])) + + # same base type but no specialized method ⇒ unknown, error rather than fabricate false + @test_throws MethodError sss(UpperTriangular(rand(3,3)), UpperTriangular(rand(3,3))) + + # sparse CSC: compares the stored index arrays + A = sparse([1,2,3],[1,2,3],[1.0,2.0,3.0]) + B = sparse([1,2,3],[1,2,3],[4.0,5.0,6.0]) # same pattern, different values + C = sparse([1,2,3],[1,2,1],[4.0,5.0,6.0]) # different pattern + @test sss(A, B) + @test !sss(A, C) + @test sss(A, similar(A)) # `similar` preserves the pattern + @test !sss(A, sparse([1,2],[1,2],[1.0,2.0])) # different size +end + @testset "ndims_index" begin @test @inferred(ArrayInterface.ndims_index(CartesianIndices(()))) == 1 @test @inferred(ArrayInterface.ndims_index(trues(2, 2))) == 2