Fix chunk size picking for StaticArrays

SciML · Nov 27, 2023 · 8324047 · 8324047
1 parent 98e6af3
commit 8324047
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/src/utils.jl b/src/utils.jl
@@ -39,13 +39,23 @@ __standard_tag(::Nothing, x) = ForwardDiff.Tag(SimpleNonlinearSolveTag(), eltype
 __standard_tag(tag::ForwardDiff.Tag, _) = tag
 __standard_tag(tag, x) = ForwardDiff.Tag(tag, eltype(x))
 
+__pick_forwarddiff_chunk(x) = ForwardDiff.Chunk(length(x))
+function __pick_forwarddiff_chunk(x::StaticArray)
+    L = prod(Size(x))
+    if L ≤ ForwardDiff.DEFAULT_CHUNK_THRESHOLD
+        return ForwardDiff.Chunk{L}()
+    else
+        return ForwardDiff.Chunk{ForwardDiff.DEFAULT_CHUNK_THRESHOLD}()
+    end
+end
+
 function __get_jacobian_config(ad::AutoForwardDiff{CS}, f, x) where {CS}
-    ck = (CS === nothing || CS ≤ 0) ? ForwardDiff.Chunk(length(x)) : ForwardDiff.Chunk{CS}()
+    ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
     tag = __standard_tag(ad.tag, x)
     return ForwardDiff.JacobianConfig(f, x, ck, tag)
 end
 function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!, y, x) where {CS}
-    ck = (CS === nothing || CS ≤ 0) ? ForwardDiff.Chunk(length(x)) : ForwardDiff.Chunk{CS}()
+    ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
     tag = __standard_tag(ad.tag, x)
     return ForwardDiff.JacobianConfig(f!, y, x, ck, tag)
 end

diff --git a/test/basictests.jl b/test/basictests.jl
@@ -258,11 +258,9 @@ end
 # --- Allocation Checks ---
 
 ## SimpleDFSane needs to allocate a history vector
-@testset "Allocation Checks: $(_nameof(alg))" for alg in (
-        SimpleNewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2)),
-        SimpleHalley(; autodiff = AutoForwardDiff(; chunksize = 2)),
-        SimpleBroyden(), SimpleKlement(), SimpleLimitedMemoryBroyden(),
-        SimpleTrustRegion(; autodiff = AutoForwardDiff(; chunksize = 2)))
+@testset "Allocation Checks: $(_nameof(alg))" for alg in ( SimpleNewtonRaphson(),
+        SimpleHalley(), SimpleBroyden(), SimpleKlement(), SimpleLimitedMemoryBroyden(),
+        SimpleTrustRegion())
     @check_allocs nlsolve(prob, alg) = DiffEqBase.__solve(prob, alg; abstol = 1e-9)
 
     nlprob_scalar = NonlinearProblem{false}(quadratic_f, 1.0, 2.0)