Merge pull request #276 from avik-pal/ap/improve_immutables

Non allocating versions for StaticArrays

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "2.13.0"
+version = "2.14.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -45,13 +45,13 @@ Enzyme = "0.11"
 FiniteDiff = "2.8.1"
 ForwardDiff = "0.10"
 Graphs = "1"
-LinearAlgebra = "1.6"
+LinearAlgebra = "<0.0.1, 1"
 PackageExtensionCompat = "1"
-Random = "1.6"
+Random = "<0.0.1, 1"
 Reexport = "1"
 SciMLOperators = "0.3.7"
 Setfield = "1"
-SparseArrays = "1.6"
+SparseArrays = "<0.0.1, 1"
 StaticArrayInterface = "1.3"
 StaticArrays = "1"
 Symbolics = "5.5"

src/SparseDiffTools.jl

@@ -16,7 +16,7 @@ import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 using ArrayInterface, SparseArrays
 import ArrayInterface: matrix_colors
 import StaticArrays
-import StaticArrays: StaticArray
+import StaticArrays: StaticArray, SArray, MArray, Size
 # Others
 using SciMLOperators, LinearAlgebra, Random
 import DataStructures: DisjointSets, find_root!, union!

src/highlevel/common.jl

@@ -173,6 +173,13 @@ A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`.
 """
 function sparse_jacobian_cache end
 
+function sparse_jacobian_static_array(ad, cache, f, x::SArray)
+    # Not the most performant fallback
+    J = init_jacobian(cache)
+    sparse_jacobian!(J, ad, cache, f, MArray(x))
+    return J
+end
+
 """
     sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f, x; fx=nothing)
     sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!, fx, x)
@@ -181,6 +188,9 @@ Sequentially calls `sparse_jacobian_cache` and `sparse_jacobian!` to compute the
 `f` at `x`. Use this if the jacobian for `f` is computed exactly once. In all other
 cases, use `sparse_jacobian_cache` once to generate the cache and use `sparse_jacobian!`
 with the same cache to compute the jacobian.
+
+If `x` is a StaticArray, then this function tries to use a non-allocating implementation for
+the jacobian computation. This is possible only for a limited backends currently.
 """
 function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, args...;
         kwargs...)
@@ -189,20 +199,32 @@ function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection,
     sparse_jacobian!(J, ad, cache, args...)
     return J
 end
+function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f,
+        x::SArray; kwargs...)
+    cache = sparse_jacobian_cache(ad, sd, f, x; kwargs...)
+    return sparse_jacobian_static_array(ad, cache, f, x)
+end
 
 """
     sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f, x)
     sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f!, fx, x)
 
 Use the sparsity detection `cache` for computing the sparse Jacobian. This allocates a new
-Jacobian at every function call
+Jacobian at every function call.
+
+If `x` is a StaticArray, then this function tries to use a non-allocating implementation for
+the jacobian computation. This is possible only for a limited backends currently.
 """
 function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache,
         args...)
     J = init_jacobian(cache)
     sparse_jacobian!(J, ad, cache, args...)
     return J
 end
+function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f,
+        x::SArray)
+    return sparse_jacobian_static_array(ad, cache, f, x)
+end
 
 """
     sparse_jacobian!(J::AbstractMatrix, ad::AbstractADType, sd::AbstractSparsityDetection,
@@ -247,14 +269,18 @@ function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) w
     C isa ForwardDiff.Chunk && return C
     return __chunksize(Val(C), x)
 end
-__chunksize(::Val{nothing}, x) = ForwardDiff.Chunk(x)
+__chunksize(::Val{nothing}, x) = __chunksize(x)
 function __chunksize(::Val{C}, x) where {C}
     if C isa Integer && !(C isa Bool)
-        return C ≤ 0 ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{C}()
+        return C ≤ 0 ? __chunksize(x) : ForwardDiff.Chunk{C}()
     else
         error("$(C)::$(typeof(C)) is not a valid chunksize!")
     end
 end
+
+__chunksize(x) = ForwardDiff.Chunk(x)
+__chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}()
+
 function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C}
     C === nothing && return nothing
     C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C)
@@ -273,18 +299,36 @@ end
     return :(nothing)
 end
 
-function init_jacobian(c::AbstractMaybeSparseJacobianCache)
+"""
+    init_jacobian(cache::AbstractMaybeSparseJacobianCache;
+        preserve_immutable::Val = Val(false))
+
+Initialize the Jacobian based on the cache. Uses sparse jacobians if possible.
+
+If `preserve_immutable` is `true`, then the Jacobian returned might be immutable, this is
+relevant if the inputs are immutable like `StaticArrays`.
+"""
+function init_jacobian(c::AbstractMaybeSparseJacobianCache;
+        preserve_immutable::Val = Val(false))
     T = promote_type(eltype(c.fx), eltype(c.x))
-    return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x)
+    return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x;
+        preserve_immutable)
 end
-init_jacobian(::Nothing, ::Type{T}, fx, x) where {T} = similar(fx, T, length(fx), length(x))
-function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray) where {T}
-    # We want to construct a MArray to preserve types
-    J = StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
-    return J
+function init_jacobian(::Nothing, ::Type{T}, fx, x; kwargs...) where {T}
+    return similar(fx, T, length(fx), length(x))
+end
+function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray;
+        preserve_immutable::Val{PI} = Val(true)) where {T, PI}
+    if PI
+        return StaticArrays.SArray{Tuple{length(fx), length(x)}, T}(I)
+    else
+        return StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
+    end
+end
+function init_jacobian(J, ::Type{T}, fx, x; kwargs...) where {T}
+    return similar(J, T, size(J, 1), size(J, 2))
 end
-init_jacobian(J, ::Type{T}, _, _) where {T} = similar(J, T, size(J, 1), size(J, 2))
-init_jacobian(J::SparseMatrixCSC, ::Type{T}, _, _) where {T} = T.(J)
+init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
 
 __maybe_copy_x(_, x) = x
 __maybe_copy_x(_, ::Nothing) = nothing

src/highlevel/finite_diff.jl

@@ -48,3 +48,7 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::FiniteDiffJacobianCache,
     FiniteDiff.finite_difference_jacobian!(J, f!, x, cache.cache)
     return J
 end
+
+function sparse_jacobian_static_array(_, cache::FiniteDiffJacobianCache, f, x::SArray)
+    return FiniteDiff.finite_difference_jacobian(f, x, cache.cache)
+end

src/highlevel/forward_mode.jl

@@ -71,3 +71,11 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache,
     end
     return J
 end
+
+function sparse_jacobian_static_array(_, cache::ForwardDiffJacobianCache, f, x::SArray)
+    if cache.cache isa ForwardColorJacCache
+        return forwarddiff_color_jacobian(f, x, cache.cache)
+    else
+        return ForwardDiff.jacobian(f, x, cache.cache)
+    end
+end

test/allocs/Project.toml

@@ -0,0 +1,2 @@
+[deps]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"

test/runtests.jl

@@ -5,8 +5,8 @@ const GROUP = get(ENV, "GROUP", "All")
 const is_APPVEYOR = (Sys.iswindows() && haskey(ENV, "APPVEYOR"))
 const is_TRAVIS = haskey(ENV, "TRAVIS")
 
-function activate_gpu_env()
-    Pkg.activate("gpu")
+function activate_env(env)
+    Pkg.activate(env)
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
     Pkg.instantiate()
 end
@@ -42,6 +42,7 @@ if GROUP == "Core" || GROUP == "All"
 end
 
 if GROUP == "InterfaceI" || GROUP == "All"
+    VERSION ≥ v"1.9" && activate_env("allocs")
     @time @safetestset "Jac Vecs and Hes Vecs" begin
         include("test_jaches_products.jl")
     end
@@ -54,7 +55,7 @@ if GROUP == "InterfaceI" || GROUP == "All"
 end
 
 if GROUP == "GPU"
-    activate_gpu_env()
+    activate_env("gpu")
     @time @safetestset "GPU AD" begin
         include("test_gpu_ad.jl")
     end

test/test_sparse_jacobian.jl

@@ -1,6 +1,6 @@
 ## Sparse Jacobian tests
-using SparseDiffTools, Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme
-using Test
+using SparseDiffTools,
+    Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays
 
 @views function fdiff(y, x) # in-place
     L = length(x)
@@ -163,3 +163,35 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
         end
     end
 end
+
+@static if VERSION ≥ v"1.9"
+    using AllocCheck
+end
+
+@static if VERSION ≥ v"1.9"
+    # Testing that the non-sparse jacobian's are non-allocating.
+    fvcat(x) = vcat(x, x)
+
+    x_sa = @SVector randn(Float32, 10)
+
+    J_true_sa = ForwardDiff.jacobian(fvcat, x_sa)
+
+    AllocCheck.@check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F}
+        return sparse_jacobian(ad, sd, f, x)
+    end
+
+    @testset "Static Arrays" begin
+        @testset "No Allocations: $(difftype)" for difftype in (AutoSparseForwardDiff(),
+            AutoForwardDiff())
+            J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa)
+            @test J ≈ J_true_sa
+        end
+
+        @testset "Other Backends: $(difftype)" for difftype in (AutoSparseZygote(),
+            AutoZygote(), AutoSparseEnzyme(), AutoEnzyme(), AutoSparseFiniteDiff(),
+            AutoFiniteDiff())
+            J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa)
+            @test J ≈ J_true_sa
+        end
+    end
+end