Merge pull request #281 from avik-pal/ap/polyester-forward-diff

Add PolyesterForwardDiff Support

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
-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"

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

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

test/1.9specific/Project.toml

@@ -0,0 +1,3 @@
+[deps]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"

test/allocs/Project.toml

@@ -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

test/runtests.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