diff --git a/Project.toml b/Project.toml index 1055d5cf..ac09c34a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" authors = ["Pankaj Mishra ", "Chris Rackauckas "] -version = "2.15.1" +version = "2.16.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -27,17 +27,19 @@ VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f" [weakdeps] Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] SparseDiffToolsEnzymeExt = "Enzyme" +SparseDiffToolsPolyesterForwardDiffExt = "PolyesterForwardDiff" SparseDiffToolsSymbolicsExt = "Symbolics" SparseDiffToolsZygoteExt = "Zygote" [compat] -ADTypes = "0.2.1" -Adapt = "3.0, 4" +ADTypes = "0.2.6" +Adapt = "3, 4" ArrayInterface = "7.4.2" Compat = "4" DataStructures = "0.18" @@ -47,7 +49,8 @@ ForwardDiff = "0.10" Graphs = "1" LinearAlgebra = "<0.0.1, 1" PackageExtensionCompat = "1" -Random = "<0.0.1, 1" +PolyesterForwardDiff = "0.1.1" +Random = "1.6" Reexport = "1" SciMLOperators = "0.3.7" Setfield = "1" @@ -67,6 +70,7 @@ BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl new file mode 100644 index 00000000..18f704dc --- /dev/null +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -0,0 +1,77 @@ +module SparseDiffToolsPolyesterForwardDiffExt + +using ADTypes, SparseDiffTools, PolyesterForwardDiff +import ForwardDiff +import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection, + ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache, sparse_jacobian!, + sparse_jacobian_static_array, __standard_tag, __chunksize + +struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <: + AbstractMaybeSparseJacobianCache + coloring::CO + cache::CA + jac_prototype::J + fx::FX + x::X +end + +function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff, + AutoPolyesterForwardDiff}, sd::AbstractMaybeSparsityDetection, f::F, x; + fx = nothing) where {F} + coloring_result = sd(ad, f, x) + fx = fx === nothing ? similar(f(x)) : fx + if coloring_result isa NoMatrixColoring + cache = __chunksize(ad, x) + jac_prototype = nothing + else + @warn """Currently PolyesterForwardDiff does not support sparsity detection + natively. Falling back to using ForwardDiff.jl""" maxlog=1 + tag = __standard_tag(nothing, x) + # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof` + cache = ForwardColorJacCache(f, x, __chunksize(ad); coloring_result.colorvec, + dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag)) + jac_prototype = coloring_result.jacobian_sparsity + end + return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) +end + +function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff, + AutoPolyesterForwardDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, + x) where {F} + coloring_result = sd(ad, f!, fx, x) + if coloring_result isa NoMatrixColoring + cache = __chunksize(ad, x) + jac_prototype = nothing + else + @warn """Currently PolyesterForwardDiff does not support sparsity detection + natively. Falling back to using ForwardDiff.jl""" maxlog=1 + tag = __standard_tag(nothing, x) + # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof` + cache = ForwardColorJacCache(f!, x, __chunksize(ad); coloring_result.colorvec, + dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag)) + jac_prototype = coloring_result.jacobian_sparsity + end + return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) +end + +function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacobianCache, + f::F, x) where {F} + if cache.cache isa ForwardColorJacCache + forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff + else + PolyesterForwardDiff.threaded_jacobian!(f, J, x, cache.cache) # Don't try to exploit sparsity + end + return J +end + +function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacobianCache, + f!::F, fx, x) where {F} + if cache.cache isa ForwardColorJacCache + forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff + else + PolyesterForwardDiff.threaded_jacobian!(f!, fx, J, x, cache.cache) # Don't try to exploit sparsity + end + return J +end + +end diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 1f609fed..82f1cd7d 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -269,7 +269,8 @@ function init_jacobian end const __init_𝒥 = init_jacobian # Misc Functions -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) where {C} +function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, + AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}, x) where {C} C isa ForwardDiff.Chunk && return C return __chunksize(Val(C), x) end @@ -285,7 +286,8 @@ end __chunksize(x) = ForwardDiff.Chunk(x) __chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}() -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C} +function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, + AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C} C === nothing && return nothing C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C) return nothing @@ -347,4 +349,5 @@ end @inline __backend(::Union{AutoEnzyme, AutoSparseEnzyme}) = :Enzyme @inline __backend(::Union{AutoZygote, AutoSparseZygote}) = :Zygote @inline __backend(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = :ForwardDiff +@inline __backend(::Union{AutoPolyesterForwardDiff, AutoSparsePolyesterForwardDiff}) = :PolyesterForwardDiff @inline __backend(::Union{AutoFiniteDiff, AutoSparseFiniteDiff}) = :FiniteDiff diff --git a/test/1.9specific/Project.toml b/test/1.9specific/Project.toml new file mode 100644 index 00000000..c151568d --- /dev/null +++ b/test/1.9specific/Project.toml @@ -0,0 +1,3 @@ +[deps] +AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" +PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" diff --git a/test/allocs/Project.toml b/test/allocs/Project.toml deleted file mode 100644 index 2ce39fac..00000000 --- a/test/allocs/Project.toml +++ /dev/null @@ -1,2 +0,0 @@ -[deps] -AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" diff --git a/test/runtests.jl b/test/runtests.jl index 52b53f92..38085bb0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -42,7 +42,7 @@ if GROUP == "Core" || GROUP == "All" end if GROUP == "InterfaceI" || GROUP == "All" - VERSION ≥ v"1.9" && activate_env("allocs") + VERSION ≥ v"1.9" && activate_env("1.9specific") @time @safetestset "Jac Vecs and Hes Vecs" begin include("test_jaches_products.jl") end diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index d6d4273a..395604c5 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -2,6 +2,22 @@ using SparseDiffTools, Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays +@static if VERSION ≥ v"1.9" + using PolyesterForwardDiff +end + +function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, + AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C} + return C +end + +function __isinferrable(difftype) + return !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff || + difftype isa AutoSparsePolyesterForwardDiff || + difftype isa AutoPolyesterForwardDiff) || + (__chunksize(difftype) isa Int && __chunksize(difftype) > 0) +end + @views function fdiff(y, x) # in-place L = length(x) y[2:(L - 1)] .= x[1:(L - 2)] .- 2 .* x[2:(L - 1)] .+ x[3:L] @@ -38,11 +54,22 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Sparsity Detection: $(nameof(typeof(sd)))" @info "Out of Place Function" - @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in (AutoSparseZygote(), - AutoZygote(), AutoSparseForwardDiff(), AutoForwardDiff(), - AutoSparseForwardDiff(; chunksize = 0), AutoForwardDiff(; chunksize = 0), - AutoSparseForwardDiff(; chunksize = 4), AutoForwardDiff(; chunksize = 4), - AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme()) + DIFFTYPES = [AutoSparseZygote(), AutoZygote(), AutoSparseForwardDiff(), + AutoForwardDiff(), AutoSparseForwardDiff(; chunksize = 0), + AutoForwardDiff(; chunksize = 0), AutoSparseForwardDiff(; chunksize = 4), + AutoForwardDiff(; chunksize = 4), AutoSparseFiniteDiff(), AutoFiniteDiff(), + AutoEnzyme(), AutoSparseEnzyme()] + + if VERSION ≥ v"1.9" + append!(DIFFTYPES, + [AutoSparsePolyesterForwardDiff(), AutoPolyesterForwardDiff(), + AutoSparsePolyesterForwardDiff(; chunksize = 0), + AutoPolyesterForwardDiff(; chunksize = 0), + AutoSparsePolyesterForwardDiff(; chunksize = 4), + AutoPolyesterForwardDiff(; chunksize = 4)]) + end + + @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in DIFFTYPES @testset "Cache & Reuse" begin cache = sparse_jacobian_cache(difftype, sd, fdiff, x) J = init_jacobian(cache) @@ -59,7 +86,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, cache, fdiff, x) end @@ -71,7 +98,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa J = sparse_jacobian(difftype, sd, fdiff, x) @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, sd, fdiff, x) end @@ -114,7 +141,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa J = sparse_jacobian(difftype, cache, fdiff, y, x) @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, cache, fdiff, y, x) end @@ -126,7 +153,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa J = sparse_jacobian(difftype, sd, fdiff, y, x) @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, sd, fdiff, y, x) end