Merge pull request JuliaDiff#285 from avik-pal/ap/speed

More extensive use of Polyester

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "2.16.0"
+version = "2.17.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -27,12 +27,14 @@ VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [weakdeps]
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SparseDiffToolsEnzymeExt = "Enzyme"
+SparseDiffToolsPolyesterExt = "Polyester"
 SparseDiffToolsPolyesterForwardDiffExt = "PolyesterForwardDiff"
 SparseDiffToolsSymbolicsExt = "Symbolics"
 SparseDiffToolsZygoteExt = "Zygote"
@@ -49,6 +51,7 @@ ForwardDiff = "0.10"
 Graphs = "1"
 LinearAlgebra = "<0.0.1, 1"
 PackageExtensionCompat = "1"
+Polyester = "0.7.9"
 PolyesterForwardDiff = "0.1.1"
 Random = "1.6"
 Reexport = "1"
@@ -70,6 +73,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"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"

ext/SparseDiffToolsPolyesterExt.jl

@@ -0,0 +1,70 @@
+module SparseDiffToolsPolyesterExt
+
+using Adapt, ArrayInterface, ForwardDiff, FiniteDiff, Polyester, SparseDiffTools,
+      SparseArrays
+import SparseDiffTools: polyesterforwarddiff_color_jacobian, ForwardColorJacCache,
+                        __parameterless_type
+
+function cld_fast(a::A, b::B) where {A, B}
+    T = promote_type(A, B)
+    return cld_fast(a % T, b % T)
+end
+function cld_fast(n::T, d::T) where {T}
+    x = Base.udiv_int(n, d)
+    x += n != d * x
+    return x
+end
+
+function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
+        x::AbstractArray{<:Number}, jac_cache::ForwardColorJacCache) where {F}
+    t = jac_cache.t
+    dx = jac_cache.dx
+    p = jac_cache.p
+    colorvec = jac_cache.colorvec
+    sparsity = jac_cache.sparsity
+    chunksize = jac_cache.chunksize
+    maxcolor = maximum(colorvec)
+
+    vecx = vec(x)
+
+    nrows, ncols = size(J)
+
+    if !(sparsity isa Nothing)
+        rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+        rows_index = [rows_index[i] for i in 1:length(rows_index)]
+        cols_index = [cols_index[i] for i in 1:length(cols_index)]
+    else
+        rows_index = 1:nrows
+        cols_index = 1:ncols
+    end
+
+    if J isa AbstractSparseMatrix
+        fill!(nonzeros(J), zero(eltype(J)))
+    else
+        fill!(J, zero(eltype(J)))
+    end
+
+    batch((length(p), min(length(p), Threads.nthreads()))) do _, start, stop
+        color_i = (start - 1) * chunksize + 1
+        for i in start:stop
+            partial_i = p[i]
+            t_ = reshape(eltype(t).(vecx, ForwardDiff.Partials.(partial_i)), size(t))
+            fx = f(t_)
+            for j in 1:chunksize
+                dx = vec(ForwardDiff.partials.(fx, j))
+                pick_inds = [idx
+                             for idx in 1:length(rows_index)
+                             if colorvec[cols_index[idx]] == color_i]
+                rows_index_c = rows_index[pick_inds]
+                cols_index_c = cols_index[pick_inds]
+                @inbounds @simd for i in 1:length(rows_index_c)
+                    J[rows_index_c[i], cols_index_c[i]] = dx[rows_index_c[i]]
+                end
+                color_i += 1
+            end
+        end
+    end
+    return J
+end
+
+end

ext/SparseDiffToolsPolyesterForwardDiffExt.jl

@@ -1,11 +1,13 @@
 module SparseDiffToolsPolyesterForwardDiffExt
 
-using ADTypes, SparseDiffTools, PolyesterForwardDiff
+using ADTypes, SparseDiffTools, PolyesterForwardDiff, UnPack, Random, SparseArrays
 import ForwardDiff
 import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection,
                         ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache,
                         sparse_jacobian!,
-                        sparse_jacobian_static_array, __standard_tag, __chunksize
+                        sparse_jacobian_static_array, __standard_tag, __chunksize,
+                        polyesterforwarddiff_color_jacobian,
+                        polyesterforwarddiff_color_jacobian!
 
 struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <:
        AbstractMaybeSparseJacobianCache
@@ -17,18 +19,14 @@ struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <:
 end
 
 function sparse_jacobian_cache(
-        ad::Union{AutoSparsePolyesterForwardDiff,
-            AutoPolyesterForwardDiff},
-        sd::AbstractMaybeSparsityDetection, f::F, x;
-        fx = nothing) where {F}
+        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,
@@ -39,17 +37,16 @@ function sparse_jacobian_cache(
 end
 
 function sparse_jacobian_cache(
-        ad::Union{AutoSparsePolyesterForwardDiff,
-            AutoPolyesterForwardDiff},
-        sd::AbstractMaybeSparsityDetection, f!::F, fx,
-        x) where {F}
+        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
+                 natively for inplace functions. 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,
@@ -62,7 +59,7 @@ 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
+        polyesterforwarddiff_color_jacobian(J, f, x, cache.cache)
     else
         PolyesterForwardDiff.threaded_jacobian!(f, J, x, cache.cache) # Don't try to exploit sparsity
     end
@@ -72,11 +69,46 @@ 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
+        forwarddiff_color_jacobian!(J, f!, x, cache.cache)
     else
         PolyesterForwardDiff.threaded_jacobian!(f!, fx, J, x, cache.cache) # Don't try to exploit sparsity
     end
     return J
 end
 
+## Approximate Sparsity Detection
+function (alg::ApproximateJacobianSparsity)(
+        ad::AutoSparsePolyesterForwardDiff, f::F, x; fx = nothing, kwargs...) where {F}
+    @unpack ntrials, rng = alg
+    fx = fx === nothing ? f(x) : fx
+    ck = __chunksize(ad, x)
+    J = fill!(similar(fx, length(fx), length(x)), 0)
+    J_cache = similar(J)
+    x_ = similar(x)
+    for _ in 1:ntrials
+        randn!(rng, x_)
+        PolyesterForwardDiff.threaded_jacobian!(f, J_cache, x_, ck)
+        @. J += abs(J_cache)
+    end
+    return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f, x;
+        fx, kwargs...)
+end
+
+function (alg::ApproximateJacobianSparsity)(
+        ad::AutoSparsePolyesterForwardDiff, f::F, fx, x;
+        kwargs...) where {F}
+    @unpack ntrials, rng = alg
+    ck = __chunksize(ad, x)
+    J = fill!(similar(fx, length(fx), length(x)), 0)
+    J_cache = similar(J)
+    x_ = similar(x)
+    for _ in 1:ntrials
+        randn!(rng, x_)
+        PolyesterForwardDiff.threaded_jacobian!(f, fx, J_cache, x_, ck)
+        @. J += abs(J_cache)
+    end
+    return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f, x;
+        fx, kwargs...)
+end
+
 end

src/differentiation/compute_jacobian_ad.jl

@@ -173,6 +173,9 @@ function forwarddiff_color_jacobian(f::F, x::AbstractArray{<:Number},
     end
 end
 
+# Defined in extension. Polyester version of `forwarddiff_color_jacobian`
+function polyesterforwarddiff_color_jacobian end
+
 # When J is mutable, this version of forwarddiff_color_jacobian will mutate J to avoid allocations
 function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
         x::AbstractArray{<:Number},
@@ -249,9 +252,8 @@ function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
 end
 
 # When J is immutable, this version of forwarddiff_color_jacobian will avoid mutating J
-function forwarddiff_color_jacobian_immutable(f, x::AbstractArray{<:Number},
-        jac_cache::ForwardColorJacCache,
-        jac_prototype = nothing)
+function forwarddiff_color_jacobian_immutable(f::F, x::AbstractArray{<:Number},
+        jac_cache::ForwardColorJacCache, jac_prototype = nothing) where {F}
     t = jac_cache.t
     dx = jac_cache.dx
     p = jac_cache.p
@@ -315,16 +317,16 @@ function forwarddiff_color_jacobian_immutable(f, x::AbstractArray{<:Number},
     return J
 end
 
-function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f,
+function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f::F,
         x::AbstractArray{<:Number}; dx = similar(x, size(J, 1)), colorvec = 1:length(x),
-        sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing)
+        sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) where {F}
     forwarddiff_color_jacobian!(J, f, x, ForwardColorJacCache(f, x; dx, colorvec, sparsity))
 end
 
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
-        f,
+        f::F,
         x::AbstractArray{<:Number},
-        jac_cache::ForwardColorJacCache)
+        jac_cache::ForwardColorJacCache) where {F}
     t = jac_cache.t
     fx = jac_cache.fx
     dx = jac_cache.dx

src/highlevel/coloring.jl

@@ -36,7 +36,11 @@ function (alg::ApproximateJacobianSparsity)(
         ad::AbstractSparseADType, f::F, x; fx = nothing,
         kwargs...) where {F}
     if !(ad isa AutoSparseForwardDiff)
-        @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
+        if ad isa AutoSparsePolyesterForwardDiff
+            @warn "$(ad) is only supported if `PolyesterForwardDiff` is explicitly loaded. Using ForwardDiff instead." maxlog=1
+        else
+            @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
+        end
     end
     @unpack ntrials, rng = alg
     fx = fx === nothing ? f(x) : fx
@@ -56,7 +60,11 @@ end
 function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, fx, x;
         kwargs...) where {F}
     if !(ad isa AutoSparseForwardDiff)
-        @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
+        if ad isa AutoSparsePolyesterForwardDiff
+            @warn "$(ad) is only supported if `PolyesterForwardDiff` is explicitly loaded. Using ForwardDiff instead." maxlog=1
+        else
+            @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
+        end
     end
     @unpack ntrials, rng = alg
     cfg = ForwardDiff.JacobianConfig(f, fx, x)