make Chunk(::StaticArray) type stable by default

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

src/prelude.jl

@@ -1,28 +1,36 @@
 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)
 
 const UNARY_PREDICATES = Symbol[:isinf, :isnan, :isfinite, :iseven, :isodd, :isreal, :isinteger]
 
 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

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

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
 
 #########

test/utils.jl

@@ -1,12 +1,12 @@
 import ForwardDiff
-using ForwardDiff: DEFAULT_CHUNK_THRESHOLD
 using Test
 using Random
 
 # seed RNG, thus making result inaccuracies deterministic
 # 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)