From d2b4e1ece775dab24cea9dbbaba0fedb3dc11aa5 Mon Sep 17 00:00:00 2001 From: ExpandingMan Date: Wed, 14 Aug 2024 18:48:53 -0400 Subject: [PATCH 1/4] make Chunk(::StaticArray) type stable by default --- ext/ForwardDiffStaticArraysExt.jl | 4 ++++ test/GradientTest.jl | 3 +++ test/JacobianTest.jl | 3 +++ 3 files changed, 10 insertions(+) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index f2b1540b..5f82d189 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -10,6 +10,10 @@ using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig vector_mode_jacobian, vector_mode_jacobian!, valtype, value using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult +_chunk(::Length{l}, s::StaticArray) where {l} = Chunk{l}() + +ForwardDiff.Chunk(s::StaticArray) = _chunk(Length(s), s) + @generated function dualize(::Type{T}, x::StaticArray) where T N = length(x) dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index a285b51b..5adfc8c7 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -136,6 +136,9 @@ end @test DiffResults.gradient(sresult1) == DiffResults.gradient(result) @test DiffResults.gradient(sresult2) == DiffResults.gradient(result) @test DiffResults.gradient(sresult3) == DiffResults.gradient(result) + + # make sure this is not a source of type instability + @inferred ForwardDiff.GradientConfig(f, sx) end @testset "exponential function at base zero" begin diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index 8a767a60..7f4e244f 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -222,6 +222,9 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) @test DiffResults.jacobian(sresult1) == DiffResults.jacobian(result) @test DiffResults.jacobian(sresult2) == DiffResults.jacobian(result) @test DiffResults.jacobian(sresult3) == DiffResults.jacobian(result) + + # make sure this is not a source of type instability + @inferred ForwardDiff.JacobianConfig(f, sx) end ######### From 37bad54ebcefb67322c09fdb91f83e028db3cff2 Mon Sep 17 00:00:00 2001 From: ExpandingMan Date: Wed, 14 Aug 2024 21:33:45 -0400 Subject: [PATCH 2/4] try to be less breaking --- ext/ForwardDiffStaticArraysExt.jl | 10 +++++++--- src/prelude.jl | 18 +++++++++++------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 5f82d189..905099e2 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -10,16 +10,20 @@ using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig vector_mode_jacobian, vector_mode_jacobian!, valtype, value using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult -_chunk(::Length{l}, s::StaticArray) where {l} = Chunk{l}() +_get_length_val(::Length{l}) where {l} = Val(l) -ForwardDiff.Chunk(s::StaticArray) = _chunk(Length(s), s) +function ForwardDiff.Chunk(x::StaticArray, th::Val{threshold}=ForwardDiff.DEFAULT_CHUNK_THRESHOLD) where {threshold} + return Chunk(_get_length_val(Length(x)), th) +end @generated function dualize(::Type{T}, x::StaticArray) where T N = length(x) dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...) V = StaticArrays.similar_type(x, Dual{T,eltype(x),N}) return quote - chunk = Chunk{$N}() + # if stars all align this should be type stable, + # we do this instead of Chunk{N}() because this should respect global limit + chunk = Chunk(x) $(Expr(:meta, :inline)) return $V($(dx)) end diff --git a/src/prelude.jl b/src/prelude.jl index 04f62765..1e9c9b2b 100644 --- a/src/prelude.jl +++ b/src/prelude.jl @@ -1,5 +1,5 @@ const NANSAFE_MODE_ENABLED = @load_preference("nansafe_mode", false) -const DEFAULT_CHUNK_THRESHOLD = @load_preference("default_chunk_threshold", 12) +const DEFAULT_CHUNK_THRESHOLD = Val(@load_preference("default_chunk_threshold", 12)) const AMBIGUOUS_TYPES = (AbstractFloat, Irrational, Integer, Rational, Real, RoundingMode) @@ -7,22 +7,26 @@ const UNARY_PREDICATES = Symbol[:isinf, :isnan, :isfinite, :iseven, :isodd, :isr struct Chunk{N} end -const CHUNKS = [Chunk{i}() for i in 1:DEFAULT_CHUNK_THRESHOLD] +const CHUNKS = ntuple(j -> Chunk{j}(), DEFAULT_CHUNK_THRESHOLD) -function Chunk(input_length::Integer, threshold::Integer = DEFAULT_CHUNK_THRESHOLD) - N = pickchunksize(input_length, threshold) - 0 < N <= DEFAULT_CHUNK_THRESHOLD && return CHUNKS[N] +_getval(::Val{n}) where {n} = n + +function Chunk(::Val{input_length}, ::Val{threshold} = DEFAULT_CHUNK_THRESHOLD) where {input_length, threshold} + N = pickchunksize(Val(input_length), Val(threshold)) + 0 < N <= _getval(DEFAULT_CHUNK_THRESHOLD) && return CHUNKS[N] return Chunk{N}() end -function Chunk(x::AbstractArray, threshold::Integer = DEFAULT_CHUNK_THRESHOLD) +function Chunk(x::AbstractArray, ::Val{threshold} = DEFAULT_CHUNK_THRESHOLD) where {threshold} return Chunk(length(x), threshold) end +Chunk(x::AbstractArray, threshold::Integer) = Chunk(x, Val(threshold)) + # Constrained to `N <= threshold`, minimize (in order of priority): # 1. the number of chunks that need to be computed # 2. the number of "left over" perturbations in the final chunk -function pickchunksize(input_length, threshold = DEFAULT_CHUNK_THRESHOLD) +function pickchunksize(::Val{input_length}, ::Val{threshold} = DEFAULT_CHUNK_THRESHOLD) where {input_length, threshold} if input_length <= threshold return input_length else From e8a7b0608b89a7ead9e9e51028189d463da63013 Mon Sep 17 00:00:00 2001 From: ExpandingMan Date: Wed, 14 Aug 2024 21:42:57 -0400 Subject: [PATCH 3/4] try to fix unit tests --- src/prelude.jl | 4 ++++ test/utils.jl | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/prelude.jl b/src/prelude.jl index 1e9c9b2b..e62294f3 100644 --- a/src/prelude.jl +++ b/src/prelude.jl @@ -17,6 +17,10 @@ function Chunk(::Val{input_length}, ::Val{threshold} = DEFAULT_CHUNK_THRESHOLD) return Chunk{N}() end +function Chunk(input_length::Integer, threshold::Integer = _getval(DEFAULT_CHUNK_THRESHOLD)) + return Chunk(Val(input_length), Val(threshold)) +end + function Chunk(x::AbstractArray, ::Val{threshold} = DEFAULT_CHUNK_THRESHOLD) where {threshold} return Chunk(length(x), threshold) end diff --git a/test/utils.jl b/test/utils.jl index 5f4e9bb2..b0f64645 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -1,5 +1,4 @@ import ForwardDiff -using ForwardDiff: DEFAULT_CHUNK_THRESHOLD using Test using Random @@ -7,6 +6,7 @@ using Random # so we don't have to retune EPS for arbitrary inputs Random.seed!(1) +const DEFAULT_CHUNK_THRESHOLD = ForwardDiff._getval(ForwardDiff.DEFAULT_CHUNK_THRESHOLD) const XLEN = DEFAULT_CHUNK_THRESHOLD + 1 const YLEN = div(DEFAULT_CHUNK_THRESHOLD, 2) + 1 const X, Y = rand(XLEN), rand(YLEN) From a4b3e4e75e3e8e1a7d19d42af5500906aa5e48a5 Mon Sep 17 00:00:00 2001 From: ExpandingMan Date: Wed, 14 Aug 2024 21:58:11 -0400 Subject: [PATCH 4/4] add tests needed for codecov --- test/GradientTest.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index 5adfc8c7..df120f82 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -160,6 +160,13 @@ end @test isempty(g_grad_const(zeros(Float64, 0))) end +# chunk now takes Val, but this was added to avoid API breakage +# the constructor does not otherwise get tested +@testset "chunk constructor methods" begin + @test ForwardDiff.Chunk([1.0], 10) == ForwardDiff.Chunk{1}() + @test ForwardDiff.Chunk(ones(10), 4) == ForwardDiff.Chunk{4}() +end + @testset "dimension errors for gradient" begin @test_throws DimensionMismatch ForwardDiff.gradient(identity, 2pi) # input @test_throws DimensionMismatch ForwardDiff.gradient(identity, fill(2pi, 2)) # vector_mode_gradient