diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index f2b1540b..905099e2 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -10,12 +10,20 @@ using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig vector_mode_jacobian, vector_mode_jacobian!, valtype, value using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult +_get_length_val(::Length{l}) where {l} = Val(l) + +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..e62294f3 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,30 @@ 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(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 +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 diff --git a/test/GradientTest.jl b/test/GradientTest.jl index a285b51b..df120f82 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 @@ -157,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 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 ######### 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)