Add sparse Enzyme Mode

Author: avik-pal

Project.toml

@@ -1,9 +1,6 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
-authors = [
-    "Pankaj Mishra <pankajmishra1511@gmail.com>",
-    "Chris Rackauckas <contact@chrisrackauckas.com>",
-]
+authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
 version = "2.5.0"
 
 [deps]
@@ -28,10 +25,12 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [weakdeps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
+SparseDiffToolsEnzymeExt = "Enzyme"
 SparseDiffToolsSymbolicsExt = "Symbolics"
 SparseDiffToolsZygoteExt = "Zygote"
 
@@ -41,6 +40,7 @@ Adapt = "3.0"
 ArrayInterface = "7.4.2"
 Compat = "4"
 DataStructures = "0.18"
+Enzyme = "0.11"
 FiniteDiff = "2.8.1"
 ForwardDiff = "0.10"
 Graphs = "1"
@@ -61,6 +61,7 @@ ArrayInterfaceBandedMatrices = "2e50d22c-5be1-4042-81b1-c572ed69783d"
 ArrayInterfaceBlockBandedMatrices = "5331f1e9-51c7-46b0-a9b0-df4434785e0a"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -71,17 +72,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = [
-    "Test",
-    "ArrayInterfaceBandedMatrices",
-    "ArrayInterfaceBlockBandedMatrices",
-    "BandedMatrices",
-    "BlockBandedMatrices",
-    "IterativeSolvers",
-    "Pkg",
-    "Random",
-    "SafeTestsets",
-    "Symbolics",
-    "Zygote",
-    "StaticArrays",
-]
+test = ["Test", "ArrayInterfaceBandedMatrices", "ArrayInterfaceBlockBandedMatrices", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]

ext/SparseDiffToolsEnzymeExt.jl

@@ -0,0 +1,58 @@
+module SparseDiffToolsEnzymeExt
+
+import ArrayInterface: fast_scalar_indexing
+import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient!, AutoSparseEnzyme
+# FIXME: For Enzyme we currently assume reverse mode
+import ADTypes: AutoEnzyme
+using Enzyme
+
+using ForwardDiff
+
+## Satisfying High-Level Interface for Sparse Jacobians
+function __gradient(::Union{AutoSparseEnzyme, AutoEnzyme}, f, x, cols)
+    dx = zero(x)
+    autodiff(Reverse, __f̂, Const(f), Duplicated(x, dx), Const(cols))
+    return vec(dx)
+end
+
+function __gradient!(::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x, cols)
+    dx = zero(x)
+    dfx = zero(fx)
+    autodiff(Reverse, __f̂, Active, Const(f!), Duplicated(fx, dfx), Duplicated(x, dx),
+        Const(cols))
+    return dx
+end
+
+function __jacobian!(J::AbstractMatrix, ::Union{AutoSparseEnzyme, AutoEnzyme}, f, x)
+    J .= jacobian(Reverse, f, x, Val(size(J, 1)))
+    return J
+end
+
+@views function __jacobian!(J, ad::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x)
+    # This version is slowish not sure how to do jacobians for inplace functions
+    @warn "Current code for computing jacobian for inplace functions in Enzyme is slow." maxlog=1
+    dfx = zero(fx)
+    dx = zero(x)
+
+    function __f_row_idx(f!, fx, x, row_idx)
+        f!(fx, x)
+        if fast_scalar_indexing(fx)
+            return fx[row_idx]
+        else
+            # Avoid the GPU Arrays scalar indexing error
+            return sum(selectdim(fx, 1, row_idx:row_idx))
+        end
+    end
+
+    for row_idx in 1:size(J, 1)
+        autodiff(Reverse, __f_row_idx, Const(f!), DuplicatedNoNeed(fx, dfx),
+            Duplicated(x, dx), Const(row_idx))
+        J[row_idx, :] .= dx
+        fill!(dfx, 0)
+        fill!(dx, 0)
+    end
+
+    return J
+end
+
+end

src/SparseDiffTools.jl

@@ -87,6 +87,8 @@ export update_coefficients, update_coefficients!, value!
 
 # High Level Interface: sparse_jacobian
 export AutoSparseZygote  # FIXME: Remove once https://github.com/SciML/ADTypes.jl/pull/16 is merged
+export AutoSparseEnzyme
+
 export NoSparsityDetection,
     SymbolicsSparsityDetection, JacPrototypeSparsityDetection, AutoSparsityDetection
 export sparse_jacobian, sparse_jacobian_cache, sparse_jacobian!

src/highlevel/common.jl

@@ -1,6 +1,8 @@
 const AbstractSparseADType = Union{AbstractSparseForwardMode, AbstractSparseReverseMode,
     AbstractSparseFiniteDifferences}
 
+struct AutoSparseEnzyme <: AbstractSparseReverseMode end
+
 # Sparsity Detection
 abstract type AbstractMaybeSparsityDetection end
 abstract type AbstractSparsityDetection <: AbstractMaybeSparsityDetection end

src/highlevel/reverse_mode.jl

@@ -7,17 +7,17 @@ struct ReverseModeJacobianCache{CO, CA, J, FX, X, I} <: AbstractMaybeSparseJacob
     idx_vec::I
 end
 
-function sparse_jacobian_cache(ad::AbstractReverseMode, sd::AbstractMaybeSparsityDetection,
-    f, x; fx = nothing)
+function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode},
+    sd::AbstractMaybeSparsityDetection, f, x; fx = nothing)
     fx = fx === nothing ? similar(f(x)) : fx
     coloring_result = sd(ad, f, x)
     jac_prototype = __getfield(coloring_result, Val(:jacobian_sparsity))
     return ReverseModeJacobianCache(coloring_result, nothing, jac_prototype, fx, x,
         collect(1:length(fx)))
 end
 
-function sparse_jacobian_cache(ad::AbstractReverseMode, sd::AbstractMaybeSparsityDetection,
-    f!, fx, x)
+function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode},
+    sd::AbstractMaybeSparsityDetection, f!, fx, x)
     coloring_result = sd(ad, f!, fx, x)
     jac_prototype = __getfield(coloring_result, Val(:jacobian_sparsity))
     return ReverseModeJacobianCache(coloring_result, nothing, jac_prototype, fx, x,

test/test_sparse_jacobian.jl

@@ -1,19 +1,20 @@
 ## Sparse Jacobian tests
-using SparseDiffTools, Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote
+using SparseDiffTools, Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme
 using Test
 
 @views function fdiff(y, x) # in-place
-    y[(begin + 1):(end - 1)] .= x[begin:(end - 2)] .- 2 .* x[(begin + 1):(end - 1)] .+
-                                x[(begin + 2):end]
-    y[begin] = -2 * x[begin] + x[begin + 1]
-    y[end] = x[end - 1] - 2 * x[end]
+    L = length(x)
+    y[2:(L - 1)] .= x[1:(L - 2)] .- 2 .* x[2:(L - 1)] .+ x[3:L]
+    y[1] = -2 * x[1] + x[2]
+    y[L] = x[L - 1] - 2 * x[L]
     return nothing
 end
 
 @views function fdiff(x) # out-of-place
-    y₂ = x[begin:(end - 2)] .- 2 .* x[(begin + 1):(end - 1)] .+ x[(begin + 2):end]
+    L = length(x)
+    y₂ = x[1:(L - 2)] .- 2 .* x[2:(L - 1)] .+ x[3:L]
     y₁ = -2x[1] + x[2]
-    y₃ = x[end - 1] - 2x[end]
+    y₃ = x[L - 1] - 2x[L]
     return vcat(y₁, y₂, y₃)
 end
 
@@ -36,7 +37,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars
 
         @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in (AutoSparseZygote(),
             AutoZygote(), AutoSparseForwardDiff(), AutoForwardDiff(),
-            AutoSparseFiniteDiff(), AutoFiniteDiff())
+            AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme())
             @testset "Cache & Reuse" begin
                 cache = sparse_jacobian_cache(difftype, sd, fdiff, x)
                 J = SparseDiffTools.__init_𝒥(cache)
@@ -88,7 +89,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(jac_prototype = J_spars
         @info "Inplace Place Function"
 
         @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (AutoSparseForwardDiff(),
-            AutoForwardDiff(), AutoSparseFiniteDiff(), AutoFiniteDiff())
+            AutoForwardDiff(), AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(),
+            AutoSparseEnzyme())
             y = similar(x)
             cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x)