Fix allocation regression: replace try-catch with type-based dispatch

The try-catch in AutoSpecializeCallable prevented inlining and added
~32 bytes per call, exceeding the 64-byte @ballocated budget in
NonlinearSolveFirstOrder, QuasiNewton, and SpectralMethods tests.

Replace with explicit dispatch methods for known argument types
(Vector{Float64}, Float64, NullParameters, and ForwardDiff duals),
routing to f.fw for zero-allocation calls. Unsupported types fall
back to f.orig via vararg dispatch. Also fix @test_broken -> @test
for @inferred solve(prob) which now passes.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl

@@ -13,7 +13,7 @@ using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem,
 
 using NonlinearSolveBase: NonlinearSolveBase, ImmutableNonlinearProblem, Utils, InternalAPI,
     NonlinearSolvePolyAlgorithm, NonlinearSolveForwardDiffCache,
-    NonlinearSolveTag, is_fw_wrapped
+    NonlinearSolveTag, AutoSpecializeCallable, is_fw_wrapped
 
 import NonlinearSolveBase: wrapfun_iip, standardize_forwarddiff_tag
 
@@ -28,6 +28,35 @@ dualgen(::Type{T}) where {T} = ForwardDiff.Dual{
     ForwardDiff.Tag{NonlinearSolveTag, T}, T, 1,
 }
 
+# Fast-path dispatch for IIP calls with NonlinearSolveTag duals.
+# These bypass the generic fallback path, calling directly into FunctionWrappersWrapper
+# for zero-allocation dispatch.
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{dualT}, u::Vector{dualT}, p::Vector{Float64},
+    )
+    return f.fw(du, u, p)
+end
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{dualT}, u::Vector{dualT}, p::SciMLBase.NullParameters,
+    )
+    return f.fw(du, u, p)
+end
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{dualT}, u::Vector{dualT}, p::Vector{dualT},
+    )
+    return f.fw(du, u, p)
+end
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{dualT}, u::Vector{Float64}, p::Vector{dualT},
+    )
+    return f.fw(du, u, p)
+end
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{dualT}, u::Vector{dualT}, p::Float64,
+    )
+    return f.fw(du, u, p)
+end
+
 # Helper: build the canonical AutoForwardDiff for wrapped functions (chunksize=1 + tag).
 function _wrapped_forwarddiff_ad()
     tag = ForwardDiff.Tag(NonlinearSolveTag(), Float64)

lib/NonlinearSolveBase/src/autospecialize.jl

@@ -40,22 +40,29 @@ struct AutoSpecializeCallable{FW} <: Function
     orig::Any  # type-erased: all wrapped functions share the same Julia type
 end
 
-# FunctionWrappersWrapper throws NoFunctionWrapperFoundError (or ErrorException
-# in older versions) when no signature matches. Fall back to the original function
-# for unsupported argument types (e.g., external packages like LeastSquaresOptim
-# doing their own ForwardDiff with different tags/chunksizes, or JVP paths that
-# bypass tag standardization).
-function (f::AutoSpecializeCallable)(args...)
-    try
-        return f.fw(args...)
-    catch e
-        if e isa FunctionWrappersWrappers.NoFunctionWrapperFoundError ||
-                (e isa ErrorException && contains(e.msg, "No matching function wrapper"))
-            return f.orig(args...)
-        end
-        rethrow()
-    end
+# Fast-path dispatch for IIP calls with Vector{Float64} arguments.
+# These call directly through the FunctionWrappersWrapper for zero-allocation dispatch.
+# The ForwardDiff extension adds analogous methods for dual-number argument types.
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64},
+    )
+    return f.fw(du, u, p)
+end
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{Float64}, u::Vector{Float64}, p::Float64,
+    )
+    return f.fw(du, u, p)
 end
+@inline function (f::AutoSpecializeCallable)(
+        du::Vector{Float64}, u::Vector{Float64}, p::SciMLBase.NullParameters,
+    )
+    return f.fw(du, u, p)
+end
+
+# Fallback: call original function for unsupported argument types (e.g., external
+# packages like LeastSquaresOptim doing their own ForwardDiff with different
+# tags/chunksizes, or JVP paths that bypass tag standardization).
+@inline (f::AutoSpecializeCallable)(args...) = f.orig(args...)
 
 """
     is_fw_wrapped(f) -> Bool

test/core_tests.jl

@@ -317,19 +317,15 @@ end
 
     @test all(!isnan, sol.u)
     @test !SciMLBase.successful_retcode(sol.retcode)
-    # IIP Vector{Float64} problems go through AutoSpecialize FunctionWrapper wrapping,
-    # which changes the function type in the solution, breaking type inference.
-    @test_broken (@inferred solve(prob)) isa Any
+    @test (@inferred solve(prob)) isa Any
 
     u0 = [0.0, 0.0, 0.0]
     prob = NonlinearProblem(f1_infeasible, u0)
     sol = solve(prob)
 
     @test all(!isnan, sol.u)
     @test !SciMLBase.successful_retcode(sol.retcode)
-    # OOP Vector{Float64} problems go through AutoSpecialize FunctionWrapper wrapping,
-    # which uses try-catch fallback for mismatched dual tags, breaking type inference.
-    @test_broken (@inferred solve(prob)) isa Any
+    @test (@inferred solve(prob)) isa Any
 
     u0 = @SVector [0.0, 0.0, 0.0]
     prob = NonlinearProblem(f1_infeasible, u0)