diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index f2c403e0..300f984f 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -61,7 +61,7 @@ PolyesterForwardDiff = "0.1.1" ReverseDiff = "1.15.1" SparseArrays = "<0.0.1,1" SparseConnectivityTracer = "0.5.0,0.6" -SparseMatrixColorings = "0.3.5" +SparseMatrixColorings = "0.4.0" Symbolics = "5.27.1, 6" Tapir = "0.2.4" Tracker = "0.2.33" diff --git a/DifferentiationInterface/docs/src/operators.md b/DifferentiationInterface/docs/src/operators.md index 97230c6f..7a740967 100644 --- a/DifferentiationInterface/docs/src/operators.md +++ b/DifferentiationInterface/docs/src/operators.md @@ -167,8 +167,7 @@ For this to work, three ingredients are needed (read [this survey](https://epubs - [`TracerSparsityDetector`](@extref SparseConnectivityTracer.TracerSparsityDetector) from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl) - [`SymbolicsSparsityDetector`](@extref Symbolics.SymbolicsSparsityDetector) from [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) - [`DenseSparsityDetector`](@ref) from DifferentiationInterface.jl (beware that this detector only gives a locally valid pattern) -3. A coloring algorithm like: - - [`GreedyColoringAlgorithm`](@extref SparseMatrixColorings.GreedyColoringAlgorithm) from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl) +3. A coloring algorithm: [`GreedyColoringAlgorithm`](@extref SparseMatrixColorings.GreedyColoringAlgorithm) from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl) is the only one we support. These ingredients can be combined within the [`AutoSparse`](@extref ADTypes.AutoSparse) wrapper, which DifferentiationInterface.jl re-exports. Note that for sparse Hessians, you need to put the `SecondOrder` backend inside `AutoSparse`, and not the other way around. diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index 00479bc6..e7d7aa07 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/DifferentiationInterface/src/DifferentiationInterface.jl @@ -12,7 +12,7 @@ module DifferentiationInterface using ADTypes: ADTypes, AbstractADType using ADTypes: mode, ForwardMode, ForwardOrReverseMode, ReverseMode, SymbolicMode using ADTypes: AutoSparse, dense_ad -using ADTypes: coloring_algorithm, column_coloring, row_coloring +using ADTypes: coloring_algorithm using ADTypes: sparsity_detector, jacobian_sparsity, hessian_sparsity using ADTypes: AutoChainRules, @@ -35,16 +35,16 @@ using LinearAlgebra: Symmetric, Transpose, dot, parent, transpose using PackageExtensionCompat: @require_extensions using SparseArrays: SparseMatrixCSC, nonzeros, nzrange, rowvals, sparse using SparseMatrixColorings: + AbstractColoringResult, + ColoringProblem, GreedyColoringAlgorithm, - color_groups, - decompress_columns, - decompress_columns!, - decompress_rows, - decompress_rows!, - decompress_symmetric, - decompress_symmetric!, - symmetric_coloring_detailed, - StarSet + coloring, + column_colors, + row_colors, + column_groups, + row_groups, + decompress, + decompress! abstract type Extras end @@ -74,7 +74,6 @@ include("second_order/hessian.jl") include("fallbacks/no_extras.jl") include("sparse/fallbacks.jl") -include("sparse/matrices.jl") include("sparse/jacobian.jl") include("sparse/hessian.jl") diff --git a/DifferentiationInterface/src/sparse/hessian.jl b/DifferentiationInterface/src/sparse/hessian.jl index e7b1c669..828ea940 100644 --- a/DifferentiationInterface/src/sparse/hessian.jl +++ b/DifferentiationInterface/src/sparse/hessian.jl @@ -1,11 +1,14 @@ struct SparseHessianExtras{ - B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,R,E2<:HVPExtras,E1<:GradientExtras + B, + C<:AbstractColoringResult{:symmetric,:column}, + M<:AbstractMatrix{<:Real}, + D, + R, + E2<:HVPExtras, + E1<:GradientExtras, } <: HessianExtras - sparsity::S - colors::Vector{Int} - star_set::StarSet - groups::Vector{Vector{Int}} - compressed::C + coloring_result::C + compressed_matrix::M batched_seeds::Vector{Batch{B,D}} batched_results::Vector{Batch{B,R}} hvp_batched_extras::E2 @@ -13,23 +16,16 @@ struct SparseHessianExtras{ end function SparseHessianExtras{B}(; - sparsity::S, - colors, - star_set, - groups, - compressed::C, + coloring_result::C, + compressed_matrix::M, batched_seeds::Vector{Batch{B,D}}, batched_results::Vector{Batch{B,R}}, hvp_batched_extras::E2, gradient_extras::E1, -) where {B,S,C,D,R,E2,E1} - @assert size(sparsity, 1) == size(sparsity, 2) == size(compressed, 1) == length(colors) - return SparseHessianExtras{B,S,C,D,R,E2,E1}( - sparsity, - colors, - star_set, - groups, - compressed, +) where {B,C,M,D,R,E2,E1} + return SparseHessianExtras{B,C,M,D,R,E2,E1}( + coloring_result, + compressed_matrix, batched_seeds, batched_results, hvp_batched_extras, @@ -41,14 +37,16 @@ end function prepare_hessian(f::F, backend::AutoSparse, x) where {F} dense_backend = dense_ad(backend) - initial_sparsity = hessian_sparsity(f, x, sparsity_detector(backend)) - sparsity = col_major(initial_sparsity) - colors, star_set = symmetric_coloring_detailed(sparsity, coloring_algorithm(backend)) - groups = color_groups(colors) + sparsity = hessian_sparsity(f, x, sparsity_detector(backend)) + problem = ColoringProblem{:symmetric,:column}() + coloring_result = coloring( + sparsity, problem, coloring_algorithm(backend); decompression_eltype=eltype(x) + ) + groups = column_groups(coloring_result) Ng = length(groups) B = pick_batchsize(maybe_outer(dense_backend), Ng) seeds = map(group -> make_seed(x, group), groups) - compressed = stack(_ -> vec(similar(x)), groups; dims=2) + compressed_matrix = stack(_ -> vec(similar(x)), groups; dims=2) batched_seeds = Batch.([ ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % Ng], Val(B)) for @@ -58,11 +56,8 @@ function prepare_hessian(f::F, backend::AutoSparse, x) where {F} hvp_batched_extras = prepare_hvp_batched(f, dense_backend, x, batched_seeds[1]) gradient_extras = prepare_gradient(f, maybe_inner(dense_backend), x) return SparseHessianExtras{B}(; - sparsity, - colors, - star_set, - groups, - compressed, + coloring_result, + compressed_matrix, batched_seeds, batched_results, hvp_batched_extras, @@ -71,11 +66,9 @@ function prepare_hessian(f::F, backend::AutoSparse, x) where {F} end function hessian(f::F, backend::AutoSparse, x, extras::SparseHessianExtras{B}) where {F,B} - @compat (; - sparsity, compressed, colors, star_set, groups, batched_seeds, hvp_batched_extras - ) = extras + @compat (; coloring_result, batched_seeds, hvp_batched_extras) = extras dense_backend = dense_ad(backend) - Ng = length(groups) + Ng = length(column_groups(coloring_result)) hvp_batched_extras_same = prepare_hvp_batched_same_point( f, dense_backend, x, batched_seeds[1], hvp_batched_extras @@ -86,28 +79,25 @@ function hessian(f::F, backend::AutoSparse, x, extras::SparseHessianExtras{B}) w stack(vec, dg_batch.elements; dims=2) end - compressed = reduce(hcat, compressed_blocks) - if Ng < size(compressed, 2) - compressed = compressed[:, 1:Ng] + compressed_matrix = reduce(hcat, compressed_blocks) + if Ng < size(compressed_matrix, 2) + compressed_matrix = compressed_matrix[:, 1:Ng] end - return decompress_symmetric(sparsity, compressed, colors, star_set) + return decompress(compressed_matrix, coloring_result) end function hessian!( f::F, hess, backend::AutoSparse, x, extras::SparseHessianExtras{B} ) where {F,B} @compat (; - sparsity, - compressed, - colors, - star_set, - groups, + coloring_result, + compressed_matrix, batched_seeds, batched_results, hvp_batched_extras, ) = extras dense_backend = dense_ad(backend) - Ng = length(groups) + Ng = length(column_groups(coloring_result)) hvp_batched_extras_same = prepare_hvp_batched_same_point( f, dense_backend, x, batched_seeds[1], hvp_batched_extras @@ -125,13 +115,13 @@ function hessian!( for b in eachindex(batched_results[a].elements) copyto!( - view(compressed, :, 1 + ((a - 1) * B + (b - 1)) % Ng), + view(compressed_matrix, :, 1 + ((a - 1) * B + (b - 1)) % Ng), vec(batched_results[a].elements[b]), ) end end - decompress_symmetric!(hess, sparsity, compressed, colors, star_set) + decompress!(hess, compressed_matrix, coloring_result) return hess end diff --git a/DifferentiationInterface/src/sparse/jacobian.jl b/DifferentiationInterface/src/sparse/jacobian.jl index 4dee60ca..b40e37ef 100644 --- a/DifferentiationInterface/src/sparse/jacobian.jl +++ b/DifferentiationInterface/src/sparse/jacobian.jl @@ -3,45 +3,45 @@ abstract type SparseJacobianExtras <: JacobianExtras end struct PushforwardSparseJacobianExtras{ - B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,R,E<:PushforwardExtras + B, + C<:AbstractColoringResult{:nonsymmetric,:column}, + M<:AbstractMatrix{<:Real}, + D, + R, + E<:PushforwardExtras, } <: SparseJacobianExtras - sparsity::S - colors::Vector{Int} - groups::Vector{Vector{Int}} - compressed::C + coloring_result::C + compressed_matrix::M batched_seeds::Vector{Batch{B,D}} batched_results::Vector{Batch{B,R}} pushforward_batched_extras::E end struct PullbackSparseJacobianExtras{ - B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,R,E<:PullbackExtras + B, + C<:AbstractColoringResult{:nonsymmetric,:row}, + M<:AbstractMatrix{<:Real}, + D, + R, + E<:PullbackExtras, } <: SparseJacobianExtras - sparsity::S - colors::Vector{Int} - groups::Vector{Vector{Int}} - compressed::C + coloring_result::C + compressed_matrix::M batched_seeds::Vector{Batch{B,D}} batched_results::Vector{Batch{B,R}} pullback_batched_extras::E end function PushforwardSparseJacobianExtras{B}(; - sparsity::S, - colors, - groups, - compressed::C, + coloring_result::C, + compressed_matrix::M, batched_seeds::Vector{Batch{B,D}}, batched_results::Vector{Batch{B,R}}, pushforward_batched_extras::E, -) where {B,S,C,D,R,E} - @assert size(sparsity, 1) == size(compressed, 1) - @assert size(sparsity, 2) == length(colors) - return PushforwardSparseJacobianExtras{B,S,C,D,R,E}( - sparsity, - colors, - groups, - compressed, +) where {B,C,M,D,R,E} + return PushforwardSparseJacobianExtras{B,C,M,D,R,E}( + coloring_result, + compressed_matrix, batched_seeds, batched_results, pushforward_batched_extras, @@ -49,21 +49,15 @@ function PushforwardSparseJacobianExtras{B}(; end function PullbackSparseJacobianExtras{B}(; - sparsity::S, - colors, - groups, - compressed::C, + coloring_result::C, + compressed_matrix::M, batched_seeds::Vector{Batch{B,D}}, batched_results::Vector{Batch{B,R}}, pullback_batched_extras::E, -) where {B,S,C,D,R,E} - @assert size(sparsity, 2) == size(compressed, 2) - @assert size(sparsity, 1) == length(colors) - return PullbackSparseJacobianExtras{B,S,C,D,R,E}( - sparsity, - colors, - groups, - compressed, +) where {B,C,M,D,R,E} + return PullbackSparseJacobianExtras{B,C,M,D,R,E}( + coloring_result, + compressed_matrix, batched_seeds, batched_results, pullback_batched_extras, @@ -87,14 +81,19 @@ function prepare_sparse_jacobian_aux( f_or_f!y::FY, backend::AutoSparse, x, y, ::PushforwardFast ) where {FY} dense_backend = dense_ad(backend) - initial_sparsity = jacobian_sparsity(f_or_f!y..., x, sparsity_detector(backend)) - sparsity = col_major(initial_sparsity) - colors = column_coloring(sparsity, coloring_algorithm(backend)) - groups = color_groups(colors) + sparsity = jacobian_sparsity(f_or_f!y..., x, sparsity_detector(backend)) + problem = ColoringProblem{:nonsymmetric,:column}() + coloring_result = coloring( + sparsity, + problem, + coloring_algorithm(backend); + decompression_eltype=promote_type(eltype(x), eltype(y)), + ) + groups = column_groups(coloring_result) Ng = length(groups) B = pick_batchsize(dense_backend, Ng) seeds = map(group -> make_seed(x, group), groups) - compressed = stack(_ -> vec(similar(y)), groups; dims=2) + compressed_matrix = stack(_ -> vec(similar(y)), groups; dims=2) batched_seeds = Batch.([ ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % Ng], Val(B)) for @@ -105,10 +104,8 @@ function prepare_sparse_jacobian_aux( f_or_f!y..., dense_backend, x, batched_seeds[1] ) return PushforwardSparseJacobianExtras{B}(; - sparsity, - colors, - groups, - compressed, + coloring_result, + compressed_matrix, batched_seeds, batched_results, pushforward_batched_extras, @@ -119,14 +116,19 @@ function prepare_sparse_jacobian_aux( f_or_f!y::FY, backend::AutoSparse, x, y, ::PushforwardSlow ) where {FY} dense_backend = dense_ad(backend) - initial_sparsity = jacobian_sparsity(f_or_f!y..., x, sparsity_detector(backend)) - sparsity = row_major(initial_sparsity) - colors = row_coloring(sparsity, coloring_algorithm(backend)) - groups = color_groups(colors) + sparsity = jacobian_sparsity(f_or_f!y..., x, sparsity_detector(backend)) + problem = ColoringProblem{:nonsymmetric,:row}() + coloring_result = coloring( + sparsity, + problem, + coloring_algorithm(backend); + decompression_eltype=promote_type(eltype(x), eltype(y)), + ) + groups = row_groups(coloring_result) Ng = length(groups) B = pick_batchsize(dense_backend, Ng) seeds = map(group -> make_seed(y, group), groups) - compressed = stack(_ -> vec(similar(x)), groups; dims=1) + compressed_matrix = stack(_ -> vec(similar(x)), groups; dims=1) batched_seeds = Batch.([ ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % Ng], Val(B)) for @@ -137,10 +139,8 @@ function prepare_sparse_jacobian_aux( f_or_f!y..., dense_backend, x, batched_seeds[1] ) return PullbackSparseJacobianExtras{B}(; - sparsity, - colors, - groups, - compressed, + coloring_result, + compressed_matrix, batched_seeds, batched_results, pullback_batched_extras, @@ -204,11 +204,9 @@ end function sparse_jacobian_aux( f_or_f!y::FY, backend::AutoSparse, x, extras::PushforwardSparseJacobianExtras{B} ) where {FY,B} - @compat (; - sparsity, compressed, colors, groups, batched_seeds, pushforward_batched_extras - ) = extras + @compat (; coloring_result, batched_seeds, pushforward_batched_extras) = extras dense_backend = dense_ad(backend) - Ng = length(groups) + Ng = length(column_groups(coloring_result)) pushforward_batched_extras_same = prepare_pushforward_batched_same_point( f_or_f!y..., dense_backend, x, batched_seeds[1], pushforward_batched_extras @@ -225,21 +223,19 @@ function sparse_jacobian_aux( stack(vec, dy_batch.elements; dims=2) end - compressed = reduce(hcat, compressed_blocks) - if Ng < size(compressed, 2) - compressed = compressed[:, 1:Ng] + compressed_matrix = reduce(hcat, compressed_blocks) + if Ng < size(compressed_matrix, 2) + compressed_matrix = compressed_matrix[:, 1:Ng] end - return decompress_columns(sparsity, compressed, colors) + return decompress(compressed_matrix, coloring_result) end function sparse_jacobian_aux( f_or_f!y::FY, backend::AutoSparse, x, extras::PullbackSparseJacobianExtras{B} ) where {FY,B} - @compat (; - sparsity, compressed, colors, groups, batched_seeds, pullback_batched_extras - ) = extras + @compat (; coloring_result, batched_seeds, pullback_batched_extras) = extras dense_backend = dense_ad(backend) - Ng = length(groups) + Ng = length(row_groups(coloring_result)) pullback_batched_extras_same = prepare_pullback_batched_same_point( f_or_f!y..., dense_backend, x, batched_seeds[1], pullback_batched_extras @@ -256,27 +252,25 @@ function sparse_jacobian_aux( stack(vec, dx_batch.elements; dims=1) end - compressed = reduce(vcat, compressed_blocks) - if Ng < size(compressed, 1) - compressed = compressed[1:Ng, :] + compressed_matrix = reduce(vcat, compressed_blocks) + if Ng < size(compressed_matrix, 1) + compressed_matrix = compressed_matrix[1:Ng, :] end - return decompress_rows(sparsity, compressed, colors) + return decompress(compressed_matrix, coloring_result) end function sparse_jacobian_aux!( f_or_f!y::FY, jac, backend::AutoSparse, x, extras::PushforwardSparseJacobianExtras{B} ) where {FY,B} @compat (; - sparsity, - compressed, - colors, - groups, + coloring_result, + compressed_matrix, batched_seeds, batched_results, pushforward_batched_extras, ) = extras dense_backend = dense_ad(backend) - Ng = length(groups) + Ng = length(column_groups(coloring_result)) pushforward_batched_extras_same = prepare_pushforward_batched_same_point( f_or_f!y..., dense_backend, x, batched_seeds[1], pushforward_batched_extras @@ -294,13 +288,13 @@ function sparse_jacobian_aux!( for b in eachindex(batched_results[a].elements) copyto!( - view(compressed, :, 1 + ((a - 1) * B + (b - 1)) % Ng), + view(compressed_matrix, :, 1 + ((a - 1) * B + (b - 1)) % Ng), vec(batched_results[a].elements[b]), ) end end - decompress_columns!(jac, sparsity, compressed, colors) + decompress!(jac, compressed_matrix, coloring_result) return jac end @@ -308,16 +302,14 @@ function sparse_jacobian_aux!( f_or_f!y::FY, jac, backend::AutoSparse, x, extras::PullbackSparseJacobianExtras{B} ) where {FY,B} @compat (; - sparsity, - compressed, - colors, - groups, + coloring_result, + compressed_matrix, batched_seeds, batched_results, pullback_batched_extras, ) = extras dense_backend = dense_ad(backend) - Ng = length(groups) + Ng = length(row_groups(coloring_result)) pullback_batched_extras_same = prepare_pullback_batched_same_point( f_or_f!y..., dense_backend, x, batched_seeds[1], pullback_batched_extras @@ -335,12 +327,12 @@ function sparse_jacobian_aux!( for b in eachindex(batched_results[a].elements) copyto!( - view(compressed, 1 + ((a - 1) * B + (b - 1)) % Ng, :), + view(compressed_matrix, 1 + ((a - 1) * B + (b - 1)) % Ng, :), vec(batched_results[a].elements[b]), ) end end - decompress_rows!(jac, sparsity, compressed, colors) + decompress!(jac, compressed_matrix, coloring_result) return jac end diff --git a/DifferentiationInterface/src/sparse/matrices.jl b/DifferentiationInterface/src/sparse/matrices.jl deleted file mode 100644 index 8b70446e..00000000 --- a/DifferentiationInterface/src/sparse/matrices.jl +++ /dev/null @@ -1,17 +0,0 @@ -## Conversion between row- and col-major - -""" - col_major(A::AbstractMatrix) - -Construct a column-major representation of the matrix `A`. -""" -col_major(A::M) where {M<:AbstractMatrix} = A -col_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = M(A) - -""" - row_major(A::AbstractMatrix) - -Construct a row-major representation of the matrix `A`. -""" -row_major(A::M) where {M<:AbstractMatrix} = transpose(M(transpose(A))) -row_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = A