diff --git a/Project.toml b/Project.toml index 6676a0ef..d164d8a7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" authors = ["Pankaj Mishra ", "Chris Rackauckas "] -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" diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 0d971507..238c9323 100644 --- a/src/SparseDiffTools.jl +++ b/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! diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 8c8e3e51..11ac4b8c 100644 --- a/src/highlevel/common.jl +++ b/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,13 +199,21 @@ 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...) @@ -203,6 +221,10 @@ function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianC 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 diff --git a/src/highlevel/finite_diff.jl b/src/highlevel/finite_diff.jl index 31114ae0..3a44258a 100644 --- a/src/highlevel/finite_diff.jl +++ b/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 diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl index 2ae378d6..08aec0d0 100644 --- a/src/highlevel/forward_mode.jl +++ b/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 diff --git a/test/allocs/Project.toml b/test/allocs/Project.toml new file mode 100644 index 00000000..2ce39fac --- /dev/null +++ b/test/allocs/Project.toml @@ -0,0 +1,2 @@ +[deps] +AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" diff --git a/test/runtests.jl b/test/runtests.jl index 0708f4bc..52b53f92 100644 --- a/test/runtests.jl +++ b/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 diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index 5b0a0f4a..d6d4273a 100644 --- a/test/test_sparse_jacobian.jl +++ b/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