Stop precomputing chunk sizes into a nonconcrete array (#708)

This prevents constant propagation of the length of a static array. Alternative to #707 Co-authored-by: David Widmann <[email protected]>
JuliaDiff · Aug 19, 2024 · dd83a19 · dd83a19
1 parent 2ff6808
commit dd83a19
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/src/prelude.jl b/src/prelude.jl
@@ -14,12 +14,9 @@ const BINARY_PREDICATES = Symbol[:isequal, :isless, :<, :>, :(==), :(!=), :(<=),
 
 struct Chunk{N} end
 
-const CHUNKS = [Chunk{i}() for i in 1: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]
-    return Chunk{N}()
+    Base.@nif 12 d->(N == d) d->(Chunk{d}()) d->(Chunk{N}())
 end
 
 function Chunk(x::AbstractArray, threshold::Integer = DEFAULT_CHUNK_THRESHOLD)

diff --git a/test/GradientTest.jl b/test/GradientTest.jl
@@ -139,6 +139,9 @@ for T in (StaticArrays.SArray, StaticArrays.MArray)
     @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
@@ -224,6 +224,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
 
 @testset "dimension errors for jacobian" begin
@@ -235,7 +238,7 @@ end
 @testset "eigen" begin
     @test ForwardDiff.jacobian(x -> eigvals(SymTridiagonal(x, x[1:end-1])), [1.,2.]) ≈ [(1 - 3/sqrt(5))/2 (1 - 1/sqrt(5))/2 ; (1 + 3/sqrt(5))/2 (1 + 1/sqrt(5))/2]
     @test ForwardDiff.jacobian(x -> eigvals(Symmetric(x*x')), [1.,2.]) ≈ [0 0; 2 4]
-    
+
     x0 = [1.0, 2.0];
     ev1(x) = eigen(Symmetric(x*x')).vectors[:,1]
     @test ForwardDiff.jacobian(ev1, x0) ≈ Calculus.finite_difference_jacobian(ev1, x0)