diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml index 73a1826ca..b926cb4fd 100644 --- a/.github/workflows/Documentation.yml +++ b/.github/workflows/Documentation.yml @@ -23,7 +23,7 @@ jobs: JULIA_DEBUG: "Documenter" GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key - run: julia --project=docs/ --code-coverage=user docs/make.jl + run: julia --project=docs/ --code-coverage=user --color=yes docs/make.jl - uses: julia-actions/julia-processcoverage@v1 - uses: codecov/codecov-action@v3 with: 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..73e096a4b --- /dev/null +++ b/docs/src/basics/autodiff.md @@ -0,0 +1,62 @@ +# 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. + +!!! note + + The `Sparse` versions of the methods refers to automated sparsity detection. These + methods automatically discover the sparse Jacobian form from the function `f`. Note that + all methods specialize the differentiation on a sparse Jacobian if the sparse Jacobian + is given as `prob.f.jac_prototype` in the `NonlinearFunction` definition, and the + `AutoSparse` here simply refers to whether this `jac_prototype` should be generated + automatically. For more details, see + [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl) and + [Sparsity Detection Manual Entry](@ref sparsity-detection). + +## 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..45507ecdf --- /dev/null +++ b/src/adtypes.jl @@ -0,0 +1,148 @@ +# 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 + +""" + 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 + +""" + 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