Add a sparse Jacobian with Enzyme.jl

Author: amontoison

docs/src/sparse.md

@@ -31,15 +31,15 @@ x = rand(T, 2)
 H = hess(nlp, x)
 ```
 
-The backends available for sparse derivatives (`SparseADJacobian`, `SparseADHessian`, and `SparseReverseADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
+The backends available for sparse derivatives (`SparseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, and `SparseReverseADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
 These arguments specify the sparsity pattern detector and the coloring algorithm, respectively.
 
 - A **`detector`** must be of type `ADTypes.AbstractSparsityDetector`.
   The default detector is `TracerSparsityDetector()` from the package `SparseConnectivityTracer.jl`.
   Prior to version 0.8.0, the default was `SymbolicSparsityDetector()` from `Symbolics.jl`.
 
 - A **`coloring_algorithm`** must be of type `SparseMatrixColorings.GreedyColoringAlgorithm`.
-  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian`.
+  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian`.
   These algorithms are provided by the package `SparseMatrixColorings.jl`.
 
 The `GreedyColoringAlgorithm{:direct}()` performs column coloring for Jacobians and star coloring for Hessians.

src/enzyme.jl

@@ -81,7 +81,7 @@ function EnzymeReverseADJprod(
 end
 
 function Jprod!(b::EnzymeReverseADJprod, Jv, c!, x, v, ::Val)
-  Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x,Jv), Duplicated(x, v))
+  Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.x, Jv), Duplicated(x, v))
   return Jv
 end
 

src/sparse_jacobian.jl

@@ -64,25 +64,6 @@ function SparseADJacobian(
   )
 end
 
-function get_nln_nnzj(b::SparseADJacobian, nvar, ncon)
-  length(b.rowval)
-end
-
-function NLPModels.jac_structure!(
-  b::SparseADJacobian,
-  nlp::ADModel,
-  rows::AbstractVector{<:Integer},
-  cols::AbstractVector{<:Integer},
-)
-  rows .= b.rowval
-  for i = 1:(nlp.meta.nvar)
-    for j = b.colptr[i]:(b.colptr[i + 1] - 1)
-      cols[j] = i
-    end
-  end
-  return rows, cols
-end
-
 function sparse_jac_coord!(
   ℓ!::Function,
   b::SparseADJacobian{Tag},
@@ -111,8 +92,116 @@ function sparse_jac_coord!(
   return vals
 end
 
+struct SparseEnzymeADJacobian{R, C, S} <: ADBackend
+  nvar::Int
+  ncon::Int
+  rowval::Vector{Int}
+  colptr::Vector{Int}
+  nzval::Vector{R}
+  result_coloring::C
+  compressed_jacobian::S
+  v::Vector{R}
+  buffer::Vector{R}
+end
+
+function SparseEnzymeADJacobian(
+  nvar,
+  f,
+  ncon,
+  c!;
+  x0::AbstractVector = rand(nvar),
+  coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
+  detector::AbstractSparsityDetector = TracerSparsityDetector(),
+  kwargs...,
+)
+  output = similar(x0, ncon)
+  J = compute_jacobian_sparsity(c!, output, x0, detector = detector)
+  SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...)
+end
+
+function SparseEnzymeADJacobian(
+  nvar,
+  f,
+  ncon,
+  c!,
+  J::SparseMatrixCSC{Bool, Int};
+  x0::AbstractVector{T} = rand(nvar),
+  coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
+  kwargs...,
+) where {T}
+  # We should support :row and :bidirectional in the future
+  problem = ColoringProblem{:nonsymmetric, :column}()
+  result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
+
+  rowval = J.rowval
+  colptr = J.colptr
+  nzval = T.(J.nzval)
+  compressed_jacobian = similar(x0, ncon)
+  v = similar(x0)
+  buffer = zeros(T, ncon)
+
+  SparseEnzymeADJacobian(
+    nvar,
+    ncon,
+    rowval,
+    colptr,
+    nzval,
+    result_coloring,
+    compressed_jacobian,
+    v,
+    buffer,
+  )
+end
+
+function sparse_jac_coord!(
+  c!::Function,
+  b::SparseEnzymeADJacobian,
+  x::AbstractVector,
+  vals::AbstractVector,
+)
+  # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
+  A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval)
+
+  groups = column_groups(b.result_coloring)
+  for (icol, cols) in enumerate(groups)
+    # Update the seed
+    b.v .= 0
+    for col in cols
+      b.v[col] = 1
+    end
+
+    # b.compressed_jacobian is just a vector Jv here
+    # We don't use the vector mode
+    Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), Duplicated(x, b.v))
+
+    # Update the columns of the Jacobian that have the color `icol`
+    decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring)
+  end
+  vals .= b.nzval
+  return vals
+end
+
+function get_nln_nnzj(b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nvar, ncon)
+  length(b.rowval)
+end
+
+function NLPModels.jac_structure!(
+  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
+  nlp::ADModel,
+  rows::AbstractVector{<:Integer},
+  cols::AbstractVector{<:Integer},
+)
+  rows .= b.rowval
+  for i = 1:(nlp.meta.nvar)
+    for j = b.colptr[i]:(b.colptr[i + 1] - 1)
+      cols[j] = i
+    end
+  end
+  return rows, cols
+end
+
 function NLPModels.jac_coord!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
   nlp::ADModel,
   x::AbstractVector,
   vals::AbstractVector,
@@ -122,7 +211,7 @@ function NLPModels.jac_coord!(
 end
 
 function NLPModels.jac_structure_residual!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
   nls::AbstractADNLSModel,
   rows::AbstractVector{<:Integer},
   cols::AbstractVector{<:Integer},
@@ -137,7 +226,7 @@ function NLPModels.jac_structure_residual!(
 end
 
 function NLPModels.jac_coord_residual!(
-  b::SparseADJacobian,
+  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
   nls::AbstractADNLSModel,
   x::AbstractVector,
   vals::AbstractVector,

src/sparsity_pattern.jl

@@ -79,7 +79,11 @@ end
 
 function get_sparsity_pattern(model::ADModel, ::Val{:jacobian})
   backend = model.adbackend.jacobian_backend
-  validate_sparse_backend(backend, SparseADJacobian, "Jacobian")
+  validate_sparse_backend(
+    backend,
+    Union{SparseADJacobian, SparseEnzymeADJacobian},
+    "Jacobian",
+  )
   m = model.meta.ncon
   n = model.meta.nvar
   colptr = backend.colptr
@@ -102,7 +106,11 @@ end
 
 function get_sparsity_pattern(model::AbstractADNLSModel, ::Val{:jacobian_residual})
   backend = model.adbackend.jacobian_residual_backend
-  validate_sparse_backend(backend, SparseADJacobian, "Jacobian of the residual")
+  validate_sparse_backend(
+    backend,
+    Union{SparseADJacobian, SparseEnzymeADJacobian},
+    "Jacobian of the residual",
+  )
   m = model.nls_meta.nequ
   n = model.meta.nvar
   colptr = backend.colptr

test/sparse_jacobian.jl

@@ -1,5 +1,7 @@
 list_sparse_jac_backend =
-  ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()))
+  ((ADNLPModels.SparseADJacobian, Dict()),
+   (ADNLPModels.SparseEnzymeADJacobian, Dict()),
+   (ADNLPModels.ForwardDiffADJacobian, Dict()))
 
 dt = (Float32, Float64)
 
@@ -51,7 +53,7 @@ dt = (Float32, Float64)
     0 1
   ]
 
-  if backend == ADNLPModels.SparseADJacobian
+  if backend != ADNLPModels.ForwardDiffADJacobian
     J_sp = get_sparsity_pattern(nlp, :jacobian)
     @test J_sp == SparseMatrixCSC{Bool, Int}([
       1 0

test/sparse_jacobian_nls.jl

@@ -1,5 +1,7 @@
 list_sparse_jac_backend =
-  ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()))
+  ((ADNLPModels.SparseADJacobian, Dict()),
+   (ADNLPModels.SparseEnzymeADJacobian, Dict()),
+   (ADNLPModels.ForwardDiffADJacobian, Dict()))
 
 dt = (Float32, Float64)
 
@@ -43,7 +45,7 @@ dt = (Float32, Float64)
     0 1
   ]
 
-  if backend == ADNLPModels.SparseADJacobian
+  if backend != ADNLPModels.ForwardDiffADJacobian
     J_sp = get_sparsity_pattern(nls, :jacobian_residual)
     @test J_sp == SparseMatrixCSC{Bool, Int}([
       1 0