From efa1c67ed65d1daa0bf947683575e5bd0c5b5539 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 30 Jan 2024 06:51:48 -0500 Subject: [PATCH] Documentation for AD Backends --- docs/Project.toml | 2 + docs/pages.jl | 1 + docs/src/basics/autodiff.md | 51 ++++++++++++ src/NonlinearSolve.jl | 3 +- src/adtypes.jl | 150 ++++++++++++++++++++++++++++++++++++ 5 files changed, 206 insertions(+), 1 deletion(-) create mode 100644 docs/src/basics/autodiff.md create mode 100644 src/adtypes.jl diff --git a/docs/Project.toml b/docs/Project.toml index 1a82e485c..3e6cef977 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,4 +1,5 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" @@ -22,6 +23,7 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" [compat] +ADTypes = "0.2.6" AlgebraicMultigrid = "0.5, 0.6" ArrayInterface = "6, 7" BenchmarkTools = "1" diff --git a/docs/pages.jl b/docs/pages.jl index 9b15d694a..5c77d7f2c 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -13,6 +13,7 @@ pages = [ "basics/nonlinear_functions.md", "basics/solve.md", "basics/nonlinear_solution.md", + "basics/autodiff.md", "basics/termination_condition.md", "basics/diagnostics_api.md", "basics/sparsity_detection.md", diff --git a/docs/src/basics/autodiff.md b/docs/src/basics/autodiff.md new file mode 100644 index 000000000..a49af7578 --- /dev/null +++ b/docs/src/basics/autodiff.md @@ -0,0 +1,51 @@ +# Automatic Differentiation Backends + +## Summary of Finite Differencing Backends + + - [`AutoFiniteDiff`](@ref): Finite differencing, not optimal but always applicable. + - [`AutoSparseFiniteDiff`](@ref): Sparse version of [`AutoFiniteDiff`](@ref). + +## Summary of Forward Mode AD Backends + + - [`AutoForwardDiff`](@ref): The best choice for dense problems. + - [`AutoSparseForwardDiff`](@ref): Sparse version of [`AutoForwardDiff`](@ref). + - [`AutoPolyesterForwardDiff`](@ref): Might be faster than [`AutoForwardDiff`](@ref) for + large problems. Requires `PolyesterForwardDiff.jl` to be installed and loaded. + +## Summary of Reverse Mode AD Backends + + - [`AutoZygote`](@ref): The fastest choice for non-mutating array-based (BLAS) functions. + - [`AutoSparseZygote`](@ref): Sparse version of [`AutoZygote`](@ref). + - [`AutoEnzyme`](@ref): Uses `Enzyme.jl` Reverse Mode and should be considered + experimental. + +!!! note + + If `PolyesterForwardDiff.jl` is installed and loaded, then `SimpleNonlinearSolve.jl` + will automatically use `AutoPolyesterForwardDiff` as the default AD backend. + +## API Reference + +### Finite Differencing Backends + +```@docs +AutoFiniteDiff +AutoSparseFiniteDiff +``` + +### Forward Mode AD Backends + +```@docs +AutoForwardDiff +AutoSparseForwardDiff +AutoPolyesterForwardDiff +``` + +### Reverse Mode AD Backends + +```@docs +AutoZygote +AutoSparseZygote +AutoEnzyme +NonlinearSolve.AutoSparseEnzyme +``` diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 2f3d0cf13..a2321ef83 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -24,7 +24,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem, AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace - import SparseDiffTools: AbstractSparsityDetection + import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix end @@ -40,6 +40,7 @@ const True = Val(true) const False = Val(false) include("abstract_types.jl") +include("adtypes.jl") include("timer_outputs.jl") include("internal/helpers.jl") diff --git a/src/adtypes.jl b/src/adtypes.jl new file mode 100644 index 000000000..98c85265e --- /dev/null +++ b/src/adtypes.jl @@ -0,0 +1,150 @@ +# This just documents the AD types from ADTypes.jl + +""" + AutoFiniteDiff(; fdtype = Val(:forward), fdjtype = fdtype, fdhtype = Val(:hcentral)) + +This uses [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl). While not necessarily +the most efficient, this is the only choice that doesn't require the `f` function to be +automatically differentiable, which means it applies to any choice. However, because it's +using finite differencing, one needs to be careful as this procedure introduces numerical +error into the derivative estimates. + + - Compatible with GPUs + - Can be used for Jacobian-Vector Products (JVPs) + - Can be used for Vector-Jacobian Products (VJPs) + - Supports both inplace and out-of-place functions + +### Keyword Arguments + + - `fdtype`: the method used for defining the gradient + - `fdjtype`: the method used for defining the Jacobian of constraints. + - `fdhtype`: the method used for defining the Hessian +""" +AutoFiniteDiff + +""" + AutoSparseFiniteDiff() + +Sparse Version of [`AutoFiniteDiff`](@ref) that uses +[FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) and the column color vector of +the Jacobian Matrix to efficiently compute the Sparse Jacobian. + + - Supports both inplace and out-of-place functions +""" +AutoSparseFiniteDiff + +""" + AutoForwardDiff(; chunksize = nothing, tag = nothing) + AutoForwardDiff{chunksize, tagType}(tag::tagType) + +This uses the [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) package. It is +the fastest choice for square or wide systems. It is easy to use and compatible with most +Julia functions which have loose type restrictions. + + - Compatible with GPUs + - Can be used for Jacobian-Vector Products (JVPs) + - Supports both inplace and out-of-place functions + +For type-stability of internal operations, a positive `chunksize` must be provided. + +### Keyword Arguments + + - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this + number to a high value will lead to slowdowns. Use + [`NonlinearSolve.pickchunksize`](@ref) to get a proper value. + - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag. +""" +AutoForwardDiff + +""" + AutoSparseForwardDiff(; chunksize = nothing, tag = nothing) + AutoSparseForwardDiff{chunksize, tagType}(tag::tagType) + +Sparse Version of [`AutoForwardDiff`](@ref) that uses +[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and the column color vector of +the Jacobian Matrix to efficiently compute the Sparse Jacobian. + + - Supports both inplace and out-of-place functions + +For type-stability of internal operations, a positive `chunksize` must be provided. + +### Keyword Arguments + + - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this + number to a high value will lead to slowdowns. Use + [`NonlinearSolve.pickchunksize`](@ref) to get a proper value. + - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag. +""" +AutoSparseForwardDiff + +@static if isdefined(ADTypes, :AutoPolyesterForwardDiff) + """ + AutoPolyesterForwardDiff(; chunksize = nothing) + + Uses [`PolyesterForwardDiff.jl`](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) + to compute the jacobian. This is essentially parallelized `ForwardDiff.jl`. + + - Supports both inplace and out-of-place functions + + ### Keyword Arguments + + - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting + this number to a high value will lead to slowdowns. Use + [`NonlinearSolve.pickchunksize`](@ref) to get a proper value. + """ + AutoPolyesterForwardDiff +end + +""" + AutoZygote() + +Uses [`Zygote.jl`](https://github.com/FluxML/Zygote.jl) package. This is the staple +reverse-mode AD that handles a large portion of Julia with good efficiency. + + - Compatible with GPUs + - Can be used for Vector-Jacobian Products (VJPs) + - Supports only out-of-place functions + +For VJPs this is the current best choice. This is the most efficient method for long +jacobians. +""" +AutoZygote + +""" + AutoSparseZygote() + +Sparse version of [`AutoZygote`](@ref) that uses +[`Zygote.jl`](https://github.com/FluxML/Zygote.jl) and the row color vector of +the Jacobian Matrix to efficiently compute the Sparse Jacobian. + + - Supports only out-of-place functions + +This is efficient only for long jacobians or if the maximum value of the row color vector is +significantly lower than the maximum value of the column color vector. +""" +AutoSparseZygote + +""" + AutoEnzyme() + +Uses reverse mode [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). This is currently +experimental, and not extensively tested on our end. We only support Jacobian construction +and VJP support is currently not implemented. + + - Supports both inplace and out-of-place functions +""" +AutoEnzyme + +""" + AutoSparseEnzyme() + +Sparse version of [`AutoEnzyme`](@ref) that uses +[Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) and the row color vector of +the Jacobian Matrix to efficiently compute the Sparse Jacobian. + + - Supports both inplace and out-of-place functions + +This is efficient only for long jacobians or if the maximum value of the row color vector is +significantly lower than the maximum value of the column color vector. +""" +AutoSparseEnzyme