Skip to content

Commit

Permalink
Add new backends with DifferentiationInterface.jl
Browse files Browse the repository at this point in the history
  • Loading branch information
amontoison committed Sep 26, 2024
1 parent ffbaba9 commit 8740d3a
Show file tree
Hide file tree
Showing 28 changed files with 856 additions and 327 deletions.
222 changes: 221 additions & 1 deletion .buildkite/pipeline.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,5 +8,225 @@ steps:
cuda: "*"
command: |
julia --color=yes --project -e 'using Pkg; Pkg.add("CUDA"); Pkg.add("NLPModels"); Pkg.add("NLPModelsTest"); Pkg.instantiate()'
julia --color=yes --project -e 'include("test/gpu.jl")'
julia --color=yes --project -e 'include("test/gpu/nvidia.jl")'
timeout_in_minutes: 30

- label: "CPUs -- ForwardDiff.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("ForwardDiff")
Pkg.instantiate()
include("test/backend/ForwardDiff.jl")'
timeout_in_minutes: 30

- label: "CPUs -- ReverseDiff.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("ReverseDiff")
Pkg.instantiate()
include("test/backend/ReverseDiff.jl")'
timeout_in_minutes: 30

- label: "CPUs -- Enzyme.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("Enzyme")
Pkg.instantiate()
include("test/backend/Enzyme.jl")'
timeout_in_minutes: 30

- label: "CPUs -- Zygote.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("Zygote")
Pkg.instantiate()
include("test/backend/Zygote.jl")'
timeout_in_minutes: 30

- label: "CPUs -- Mooncake.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("Mooncake")
Pkg.instantiate()
include("test/backend/Mooncake.jl")'
timeout_in_minutes: 30

- label: "CPUs -- Diffractor.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("Diffractor")
Pkg.instantiate()
include("test/backend/Diffractor.jl")'
timeout_in_minutes: 30

- label: "CPUs -- Tracker.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("Tracker")
Pkg.instantiate()
include("test/backend/Tracker.jl")'
timeout_in_minutes: 30

- label: "CPUs -- Symbolics.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("Symbolics")
Pkg.instantiate()
include("test/backend/Symbolics.jl")'
timeout_in_minutes: 30

- label: "CPUs -- ChainRules.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("ChainRules")
Pkg.instantiate()
include("test/backend/ChainRules.jl")'
timeout_in_minutes: 30

- label: "CPUs -- FastDifferentiation.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("FastDifferentiation")
Pkg.instantiate()
include("test/backend/FastDifferentiation.jl")'
timeout_in_minutes: 30

- label: "CPUs -- FiniteDiff.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("FiniteDiff")
Pkg.instantiate()
include("test/backend/FiniteDiff.jl")'
timeout_in_minutes: 30

- label: "CPUs -- FiniteDifferences.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("FiniteDifferences")
Pkg.instantiate()
include("test/backend/FiniteDifferences.jl")'
timeout_in_minutes: 30

- label: "CPUs -- PolyesterForwardDiff.jl"
plugins:
- JuliaCI/julia#v1:
version: "1.10"
agents:
queue: "juliaecosystem"
os: "linux"
arch: "x86_64"
command: |
julia --color=yes --project -e '
using Pkg
Pkg.add("OptimizationProblems")
Pkg.add("PolyesterForwardDiff")
Pkg.instantiate()
include("test/backend/PolyesterForwardDiff.jl")'
3 changes: 1 addition & 2 deletions .github/workflows/CI.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ jobs:
fail-fast: false
matrix:
version:
- 'lts'
- '1'
- 'pre'
os:
Expand All @@ -28,7 +27,7 @@ jobs:
with:
version: ${{ matrix.version }}
arch: ${{ matrix.arch }}
- uses: actions/cache@v1
- uses: actions/cache@v4
env:
cache-name: cache-artifacts
with:
Expand Down
14 changes: 7 additions & 7 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,21 @@ version = "0.8.7"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6"
Requires = "ae029012-a4dd-5104-9daa-d747884805df"
ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"

[compat]
ADTypes = "1.2.1"
ForwardDiff = "0.9.0, 0.10.0"
NLPModels = "0.18, 0.19, 0.20, 0.21"
Requires = "1"
ReverseDiff = "1"
ADTypes = "1.9.0"
DifferentiationInterface = "0.6.1"
ForwardDiff = "0.10.36"
NLPModels = "0.21.3"
ReverseDiff = "1.15.3"
SparseConnectivityTracer = "0.6.1"
SparseMatrixColorings = "0.4.0"
julia = "^1.6"
julia = "1.6"
8 changes: 4 additions & 4 deletions docs/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
DataFrames = "1"
Documenter = "1.0"
ManualNLPModels = "0.1"
NLPModels = "0.21"
NLPModelsJuMP = "0.13"
OptimizationProblems = "0.8"
NLPModels = "0.21.3"
NLPModelsJuMP = "0.13.2"
OptimizationProblems = "0.9"
Percival = "0.7"
Plots = "1"
SolverBenchmark = "0.6"
Zygote = "0.6.62"
Zygote = "0.6.70"
26 changes: 14 additions & 12 deletions docs/src/backend.md
Original file line number Diff line number Diff line change
@@ -1,23 +1,25 @@
# How to switch backend in ADNLPModels

`ADNLPModels` allows the use of different backends to compute the derivatives required within NLPModel API.
It uses `ForwardDiff.jl`, `ReverseDiff.jl`, and more via optional depencies.
It uses `ForwardDiff.jl`, `ReverseDiff.jl`, and more via optional dependencies.

The backend information is in a structure [`ADNLPModels.ADModelBackend`](@ref) in the attribute `adbackend` of a `ADNLPModel`, it can also be accessed with [`get_adbackend`](@ref).

The functions used internally to define the NLPModel API and the possible backends are defined in the following table:

| Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend |
| --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- |
| `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeADGradient` | -- |
| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | -- | `SparseADJacobian` |
| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | -- | `SparseADHessian`/`SparseReverseADHessian` |
| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | -- | -- |
| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | -- | -- |
| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | -- | -- |
| `directional_second_derivative` | `ForwardDiffADGHjvprod` | -- | -- | -- | -- |

The functions `hess_structure!`, `hess_coord!`, `jac_structure!` and `jac_coord!` defined in `ad.jl` are generic to all the backends for now.
| package | ForwardDiff.jl | ReverseDiff.jl | Enzyme.jl | Zygote.jl | Mooncake.jl | Diffractor.jl | Tracker.jl | Symbolics.jl | ChainRules.jl | FastDifferentiation.jl | FiniteDiff.jl | FiniteDifferences.jl | PolyesterForwardDiff.jl |
|--------|----------------|----------------|-----------|-----------|-------------|---------------|------------|--------------|----------------------|------------------------|---------------|----------------------|-------------------------|
| $\nabla f(x)$ | `ForwardDiffADGradient` | `ReverseDiffADGradient` | `EnzymeADGradient` | `ZygoteADGradient` | `MooncakeADGradient` | `DiffractorADGradient` | `TrackerADGradient` | `SymbolicsADGradient` | `ChainRulesADGradient` | `FastDifferentiationADGradient` | `FiniteDiffADGradient` | `FiniteDifferencesADGradient` | `PolyesterForwardDiffADGradient` |
| $J_c(x)*v$ | `ForwardDiffADJprod` | `ReverseDiffADJprod` | `EnzymeADJprod` | `ZygoteADJprod` | `MooncakeADJprod` | `DiffractorADJprod` | `TrackerADJprod` | `SymbolicsADJprod` | `ChainRulesADJprod` | `FastDifferentiationADJprod` | `FiniteDiffADJprod` | `FiniteDifferencesADJprod` | `PolyesterForwardDiffADJprod` |
| $J^T_c(x)*v$ | `ForwardDiffADJtprod` | `ReverseDiffADJtprod` | `EnzymeADJtprod` | `ZygoteADJtprod` | `MooncakeADJtprod` | `DiffractorADJtprod` | `TrackerADJtprod` | `SymbolicsADJtprod` | `ChainRulesADJtprod` | `FastDifferentiationADJtprod` | `FiniteDiffADJtprod` | `FiniteDifferencesADJtprod` | `PolyesterForwardDiffADJtprod` |
| $J_c(x)$ | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `EnzymeADJacobian` | `ZygoteADJacobian` | `MooncakeADJacobian` | `DiffractorADJacobian` | `TrackerADJacobian` | `SymbolicsADJacobian` | `ChainRulesADJacobian` | `FastDifferentiationADJacobian` | `FiniteDiffADJacobian` | `FiniteDifferencesADJacobian` | `PolyesterForwardDiffADJacobian` |
| $\nabla^2 \mathcal{L}(x)*v$ | `ForwardDiffADHvprod` | `ReverseDiffADHvprod` | `EnzymeADHvprod` | `ZygoteADHvprod` | `MooncakeADHvprod` | `DiffractorADHvprod` | `TrackerADHvprod` | `SymbolicsADHvprod` | `ChainRulesADHvprod` | `FastDifferentiationADHvprod` | `FiniteDiffADHvprod` | `FiniteDifferencesADHvprod` | `PolyesterForwardDiffADHvprod` |
| $\nabla^2 \mathcal{L}(x)$ | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `EnzymeADHessian` | `ZygoteADHessian` | `MooncakeADHessian` | `DiffractorADHessian` | `TrackerADHessian` | `SymbolicsADHessian` | `ChainRulesADHessian` | `FastDifferentiationADHessian` | `FiniteDiffADHessian` | `FiniteDifferencesADHessian` | `PolyesterForwardDiffADHessian` |

$\mathcal{L}(x)$ denotes the Lagrangian $f(x) + \lambda^T c(x)$.
Except for the backends based on `ForwardDiff.jl` and `ReverseDiff.jl`, all other backends require the associated AD package to be manually installed by the user to work.
Note that the Jacobians and Hessians computed by the backends above are dense.
The backends `SparseADJacobian`, `SparseADHessian`, and `SparseReverseADHessian` should be used instead if sparse Jacobians and Hessians are required.

```@example ex1
using ADNLPModels
Expand Down
9 changes: 5 additions & 4 deletions src/ADNLPModels.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,16 @@ module ADNLPModels
using LinearAlgebra, SparseArrays

# external
using ADTypes: ADTypes, AbstractColoringAlgorithm, AbstractSparsityDetector
import DifferentiationInterface
using ADTypes: ADTypes, AbstractADType, AbstractColoringAlgorithm, AbstractSparsityDetector, AutoEnzyme, AutoZygote
using ADTypes: AutoForwardDiff, AutoReverseDiff, AutoMooncake, AutoDiffractor, AutoTracker, AutoSymbolics
using ADTypes: AutoChainRules, AutoFastDifferentiation, AutoFiniteDiff, AutoFiniteDifferences, AutoPolyesterForwardDiff
using SparseConnectivityTracer: TracerSparsityDetector
using SparseMatrixColorings
using ForwardDiff, ReverseDiff

# JSO
using NLPModels
using Requires

abstract type AbstractADNLPModel{T, S} <: AbstractNLPModel{T, S} end
abstract type AbstractADNLSModel{T, S} <: AbstractNLSModel{T, S} end
Expand All @@ -27,8 +29,7 @@ include("sparse_hessian.jl")

include("forward.jl")
include("reverse.jl")
include("enzyme.jl")
include("zygote.jl")
include("di.jl")
include("predefined_backend.jl")
include("nlp.jl")

Expand Down
Loading

0 comments on commit 8740d3a

Please sign in to comment.