Skip to content

feat: use DI for structured Jacobians #470

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Parent: feat: use DifferentiationInterface for sparse AD
merged 3 commits into from
Oct 2, 2024
Merged

feat: use DI for structured Jacobians #470

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
35d277b
feat: using DI for structured Jacobians
avik-pal Oct 2, 2024
44143ec
docs: add a table to guarantee selections
avik-pal Oct 2, 2024
8fc028b
chore: apply suggestions from code review
avik-pal Oct 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
Expand Down Expand Up @@ -146,6 +145,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Expand All @@ -155,4 +155,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

[targets]
test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SparseDiffTools", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
3 changes: 0 additions & 3 deletions src/NonlinearSolve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,6 @@ using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJac
## Sparse AD Support
using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
using SparseConnectivityTracer: TracerSparsityDetector # This can be dropped in the next release
using SparseDiffTools: SparseDiffTools, JacPrototypeSparsityDetection,
PrecomputedJacobianColorvec, init_jacobian, sparse_jacobian,
sparse_jacobian!, sparse_jacobian_cache
using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm,
LargestFirst

Expand Down
118 changes: 20 additions & 98 deletions src/internal/jacobian.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`.
stats::NLStats
autodiff
di_extras
sdifft_extras
end

function reinit_cache!(cache::JacobianCache{iip}, args...; p = cache.p,
Expand All @@ -63,31 +62,13 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,

if !has_analytic_jac && needs_jac
autodiff = construct_concrete_adtype(f, autodiff)
using_sparsedifftools = autodiff isa StructuredMatrixAutodiff
# SparseMatrixColorings can't handle structured matrices
if using_sparsedifftools
di_extras = nothing
uf = JacobianWrapper{iip}(f, p)
sdifft_extras = if iip
sparse_jacobian_cache(
autodiff.autodiff, autodiff.sparsity_detection, uf, fu, u)
else
sparse_jacobian_cache(autodiff.autodiff, autodiff.sparsity_detection,
uf, __maybe_mutable(u, autodiff); fx = fu)
end
autodiff = autodiff.autodiff # For saving we unwrap
di_extras = if iip
DI.prepare_jacobian(f, fu, autodiff, u, Constant(prob.p))
else
sdifft_extras = nothing
di_extras = if iip
DI.prepare_jacobian(f, fu, autodiff, u, Constant(prob.p))
else
DI.prepare_jacobian(f, autodiff, u, Constant(prob.p))
end
DI.prepare_jacobian(f, autodiff, u, Constant(prob.p))
end
else
using_sparsedifftools = false
di_extras = nothing
sdifft_extras = nothing
end

J = if !needs_jac
Expand All @@ -98,22 +79,18 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff)
else
if f.jac_prototype === nothing
if !using_sparsedifftools
# While this is technically wasteful, it gives out the type of the Jacobian
# which is needed to create the linear solver cache
stats.njacs += 1
if has_analytic_jac
__similar(
fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
# While this is technically wasteful, it gives out the type of the Jacobian
# which is needed to create the linear solver cache
stats.njacs += 1
if has_analytic_jac
__similar(
fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
else
if iip
DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
else
if iip
DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
else
DI.jacobian(f, autodiff, u, Constant(p))
end
DI.jacobian(f, autodiff, u, Constant(p))
end
else
zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true)))
end
else
jac_proto = if eltype(f.jac_prototype) <: Bool
Expand All @@ -126,20 +103,19 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
end
end

return JacobianCache{iip}(
J, f, fu, u, p, stats, autodiff, di_extras, sdifft_extras)
return JacobianCache{iip}(J, f, fu, u, p, stats, autodiff, di_extras)
end

function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
autodiff = nothing, kwargs...) where {F}
fu = f(u, p)
if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing, nothing)
return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same question about cases where DI preparation is skipped

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only where the user provides an analytic jacobian. This one is a special case for scalars when DI is skipped if jacobian/jvp/vjp is provided.

end
autodiff = get_dense_ad(get_concrete_forward_ad(
autodiff, prob; check_forward_mode = false))
di_extras = DI.prepare_derivative(f, get_dense_ad(autodiff), u, Constant(prob.p))
return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras, nothing)
return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras)
end

(cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p)
Expand Down Expand Up @@ -172,27 +148,16 @@ function (cache::JacobianCache{iip})(
if iip
if SciMLBase.has_jac(cache.f)
cache.f.jac(J, u, p)
elseif cache.di_extras !== nothing
else
DI.jacobian!(
cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p))
else
uf = JacobianWrapper{iip}(cache.f, p)
sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, cache.fu, u)
end
return J
else
if SciMLBase.has_jac(cache.f)
return cache.f.jac(u, p)
elseif cache.di_extras !== nothing
return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p))
else
uf = JacobianWrapper{iip}(cache.f, p)
if __can_setindex(typeof(J))
sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, u)
return J
else
return sparse_jacobian(cache.autodiff, cache.sdifft_extras, uf, u)
end
return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p))
end
end
end
Expand All @@ -207,10 +172,6 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType)
end
return ad # No sparse AD
else
if ArrayInterface.isstructured(f.jac_prototype)
return select_fastest_structured_matrix_autodiff(f.jac_prototype, f, ad)
end

return AutoSparse(
ad;
sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype),
Expand All @@ -227,10 +188,6 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType)
Base.depwarn("`sparsity::typeof($(typeof(f.sparsity)))` is deprecated. \
Pass it as `jac_prototype` instead.",
:NonlinearSolve)
if ArrayInterface.isstructured(f.sparsity)
return select_fastest_structured_matrix_autodiff(f.sparsity, f, ad)
end

return AutoSparse(
ad;
sparsity_detector = KnownJacobianSparsityDetector(f.sparsity),
Expand All @@ -252,11 +209,8 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType)
coloring_algorithm = GreedyColoringAlgorithm(LargestFirst())
)
else
if ArrayInterface.isstructured(f.jac_prototype)
return select_fastest_structured_matrix_autodiff(f.jac_prototype, f, ad)
end

if f.jac_prototype isa AbstractSparseMatrix
if f.jac_prototype isa AbstractSparseMatrix ||
ArrayInterface.isstructured(f.jac_prototype)
if !(sparsity_detector isa NoSparsityDetector)
@warn "`jac_prototype` is a sparse matrix but sparsity = $(f.sparsity) \
has also been specified. Ignoring sparsity field and using \
Expand All @@ -275,38 +229,6 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType)
end
end

@concrete struct StructuredMatrixAutodiff <: AbstractADType
autodiff <: AbstractADType
sparsity_detection
end

function select_fastest_structured_matrix_autodiff(
prototype::AbstractMatrix, f::NonlinearFunction, ad::AbstractADType)
sparsity_detection = if SciMLBase.has_colorvec(f)
PrecomputedJacobianColorvec(;
jac_prototype = prototype,
f.colorvec,
partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode
)
else
if ArrayInterface.fast_matrix_colors(prototype)
colorvec = if ADTypes.mode(ad) isa ADTypes.ForwardMode
ArrayInterface.matrix_colors(prototype)
else
ArrayInterface.matrix_colors(prototype')
end
PrecomputedJacobianColorvec(;
jac_prototype = prototype,
colorvec,
partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode
)
else
JacPrototypeSparsityDetection(; jac_prototype = prototype)
end
end
return StructuredMatrixAutodiff(AutoSparse(ad), sparsity_detection)
end

function select_fastest_coloring_algorithm(
prototype, f::NonlinearFunction, ad::AbstractADType)
if SciMLBase.has_colorvec(f)
Expand Down
Loading