From 49b5ab9f019fa99193b736028a09c876cae7afdb Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 15 Feb 2024 13:56:48 -0500 Subject: [PATCH 1/3] Add Polyester approximate sparsity detection --- Project.toml | 6 ++- ext/SparseDiffToolsPolyesterExt.jl | 3 ++ ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 48 +++++++++++++++---- src/highlevel/coloring.jl | 12 ++++- 4 files changed, 57 insertions(+), 12 deletions(-) create mode 100644 ext/SparseDiffToolsPolyesterExt.jl diff --git a/Project.toml b/Project.toml index ac09c34a..7c280fa2 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.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" diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl new file mode 100644 index 00000000..f9572fa3 --- /dev/null +++ b/ext/SparseDiffToolsPolyesterExt.jl @@ -0,0 +1,3 @@ +module SparseDiffToolsPolyesterExt + +end \ No newline at end of file diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index cc9cbce6..341192ba 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -1,6 +1,6 @@ module SparseDiffToolsPolyesterForwardDiffExt -using ADTypes, SparseDiffTools, PolyesterForwardDiff +using ADTypes, SparseDiffTools, PolyesterForwardDiff, UnPack, Random, SparseArrays import ForwardDiff import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection, ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache, @@ -17,10 +17,8 @@ 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 @@ -39,10 +37,8 @@ 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) @@ -79,4 +75,38 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacob 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 diff --git a/src/highlevel/coloring.jl b/src/highlevel/coloring.jl index 313e390b..5e4b7828 100644 --- a/src/highlevel/coloring.jl +++ b/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) From 4d0c091fe5f5f8d7a731697929fd3a5ecc4cc6dc Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 15 Feb 2024 15:24:14 -0500 Subject: [PATCH 2/3] Implement Polyester Colored AD for oop functions --- ext/SparseDiffToolsPolyesterExt.jl | 68 ++++++++++++++++++- ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 13 ++-- src/differentiation/compute_jacobian_ad.jl | 16 +++-- 3 files changed, 83 insertions(+), 14 deletions(-) diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl index f9572fa3..9e2a6db4 100644 --- a/ext/SparseDiffToolsPolyesterExt.jl +++ b/ext/SparseDiffToolsPolyesterExt.jl @@ -1,3 +1,69 @@ 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 + for i in start:stop + partial_i = p[i] + color_i = 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 + end + end + end + return J +end + end \ No newline at end of file diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index 341192ba..e30c9950 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -5,7 +5,9 @@ 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 @@ -25,8 +27,6 @@ function sparse_jacobian_cache( 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, @@ -45,7 +45,8 @@ function sparse_jacobian_cache( 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, @@ -58,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 @@ -68,7 +69,7 @@ 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 diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index e5de06ee..f0cae288 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/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 From 9fa344489dcafac4583ce999439ef61f851ebd9c Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 15 Feb 2024 16:25:54 -0500 Subject: [PATCH 3/3] Fix color index --- ext/SparseDiffToolsPolyesterExt.jl | 5 +++-- ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl index 9e2a6db4..ea520810 100644 --- a/ext/SparseDiffToolsPolyesterExt.jl +++ b/ext/SparseDiffToolsPolyesterExt.jl @@ -45,9 +45,9 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F, 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] - color_i = i t_ = reshape(eltype(t).(vecx, ForwardDiff.Partials.(partial_i)), size(t)) fx = f(t_) for j in 1:chunksize @@ -60,10 +60,11 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F, @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 \ No newline at end of file +end diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index e30c9950..3ea13500 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -94,7 +94,8 @@ function (alg::ApproximateJacobianSparsity)( fx, kwargs...) end -function (alg::ApproximateJacobianSparsity)(ad::AutoSparsePolyesterForwardDiff, f::F, fx, x; +function (alg::ApproximateJacobianSparsity)( + ad::AutoSparsePolyesterForwardDiff, f::F, fx, x; kwargs...) where {F} @unpack ntrials, rng = alg ck = __chunksize(ad, x)