AutoSpecialize & NoSpecialize mode JIT recompiling solvers

[martinjrobins]

Describe the bug 🐞

I'm trying to construct a benchmark that only uses precompiled solvers and avoids any JIT compilation of the solvers using the specialization levels, however I can't seem to do this. I'm not sure if this is bug or something I'm not doing correctly, can you please advise.

Minimal Reproducible Example 👇

import DifferentialEquations as DE
using SciMLBase

ngroups = 1
function rober!(du, u, p, t)
        k₁, k₂, k₃ = p
        y₁ = @view u[1:ngroups]
        y₂ = @view u[ngroups+1:2*ngroups]
        y₃ = @view u[2*ngroups+1:3*ngroups]
        dy₁ = @view du[1:ngroups]
        dy₂ = @view du[ngroups+1:2*ngroups]
        dy₃ = @view du[2*ngroups+1:3*ngroups]
        dy₁ .= -k₁ .* y₁ .+ k₃ .* y₂ .* y₃
        dy₂ .= k₁ .* y₁ .- k₂ .* y₂ .^2 .- k₃ .* y₂ .* y₃
        dy₃ .= k₂ .* y₂ .^2
        nothing
end

function setup()
        u0 = vcat(ones(ngroups), zeros(2*ngroups))
        p = [0.04, 3e7, 1e4]
        tspan = (0.0, 1e10)
        prob = DE.ODEProblem{true, SciMLBase.AutoSpecialize}(rober!, u0, tspan, p)
        return prob
end

setup1 = @elapsed model = setup()
println("Setup time 1: ", setup1)
setup1 = @elapsed model = setup()
println("Setup time 1 (2nd run): ", setup1)

alg1 = DE.FBDF()
solve11 = @elapsed sol1 = DE.solve(model, alg=alg1)
println("solve time 1 (alg1): ", solve11)
solve12 = @elapsed sol1 = DE.solve(model, alg=alg1)
println("solve time 2 (alg1): ", solve12)
alg2 = DE.TRBDF2()
solve13 = @elapsed sol1 = DE.solve(model, alg=alg2)
println("solve time 3 (alg2): ", solve13)
solve14 = @elapsed sol1 = DE.solve(model, alg=alg2)
println("solve time 4 (alg2): ", solve14)

Setup time 1: 0.063077009
Setup time 1 (2nd run): 2.1841e-5
solve time 1 (alg1): 1.834793159
solve time 2 (alg1): 0.011158691
solve time 3 (alg2): 2.418359262
solve time 4 (alg2): 0.062127387

Expected behavior

Based on the documentation here and @ChrisRackauckas comment here I would expect that autospecialize would use the precompiled solvers and avoid JIT recompilation of the solvers when they are first used in "solve time 1" and "solve time 3", resulting in solve time 1 & 3 to be similar to solve time 2 & 4. Due to the uncertainty of the "auto" nature of this specialization, I have also tried SciMLBase.NoSpecialize, which only gives increased JIT compilation time (again contrary to the documentation which indicates that NoSpecialize should give the best compile time performance):

Setup time 1: 0.058747097
Setup time 1 (2nd run): 2.2562e-5
solve time 1 (alg1): 6.135955937
solve time 2 (alg1): 0.026723698
solve time 3 (alg2): 3.147232052
solve time 4 (alg2): 0.054692904

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

  [6e4b80f9] BenchmarkTools v1.6.3
  [336ed68f] CSV v0.10.15
  [a93c6f00] DataFrames v1.8.1
  [0c46a032] DifferentialEquations v7.17.0
  [961ee093] ModelingToolkit v11.9.0
  [0bca4576] SciMLBase v2.136.0

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

Status `~/git/pydiffsol/bench/Manifest.toml`
  [47edcb42] ADTypes v1.21.0
  [6e696c72] AbstractPlutoDingetjes v1.3.2
  [1520ce14] AbstractTrees v0.4.5
  [7d9f7c33] Accessors v0.1.43
  [79e6a3ab] Adapt v4.4.0
  [66dad0bd] AliasTables v1.1.3
  [a95523ee] AlmostBlockDiagonals v0.1.10
  [ec485272] ArnoldiMethod v0.4.0
  [4fba245c] ArrayInterface v7.22.0
  [4c555306] ArrayLayouts v1.12.2
  [aae01518] BandedMatrices v1.11.0
  [6e4b80f9] BenchmarkTools v1.6.3
  [e2ed5e7c] Bijections v0.2.2
  [caf10ac8] BipartiteGraphs v0.1.6
  [62783981] BitTwiddlingConvenienceFunctions v0.1.6
  [8e7c35d0] BlockArrays v1.9.3
  [764a87c0] BoundaryValueDiffEq v5.18.0
  [7227322d] BoundaryValueDiffEqAscher v1.8.0
  [56b672f2] BoundaryValueDiffEqCore v1.11.1
  [85d9eb09] BoundaryValueDiffEqFIRK v1.9.0
  [1a22d4ce] BoundaryValueDiffEqMIRK v1.9.0
  [9255f1d6] BoundaryValueDiffEqMIRKN v1.8.0
  [ed55bfe0] BoundaryValueDiffEqShooting v1.9.0
  [70df07ce] BracketingNonlinearSolve v1.6.2
  [fa961155] CEnum v0.5.0
  [2a0fbf3d] CPUSummary v0.2.7
  [336ed68f] CSV v0.10.15
  [d360d2e6] ChainRulesCore v1.26.0
  [fb6a15b2] CloseOpenIntervals v0.1.13
  [944b1d66] CodecZlib v0.7.8
⌅ [861a8166] Combinatorics v1.0.2
  [38540f10] CommonSolve v0.2.6
  [bbf7d656] CommonSubexpressions v0.3.1
  [f70d9fcc] CommonWorldInvalidations v1.0.0
  [34da2185] Compat v4.18.1
  [b152e2b5] CompositeTypes v0.1.4
  [a33af91c] CompositionsBase v0.1.2
  [2569d6c7] ConcreteStructs v0.2.3
  [187b0558] ConstructionBase v1.6.0
  [adafc99b] CpuId v0.3.1
  [a8cc5b0e] Crayons v4.1.1
  [9a962f9c] DataAPI v1.16.0
  [a93c6f00] DataFrames v1.8.1
  [864edb3b] DataStructures v0.19.3
  [e2d170a0] DataValueInterfaces v1.0.0
  [bcd4f6db] DelayDiffEq v5.67.1
  [2b5f629d] DiffEqBase v6.199.0
  [459566f4] DiffEqCallbacks v4.12.0
  [77a26b50] DiffEqNoiseProcess v5.26.0
  [163ba53b] DiffResults v1.1.0
  [b552c78f] DiffRules v1.15.1
  [0c46a032] DifferentialEquations v7.17.0
  [a0c0ee7d] DifferentiationInterface v0.7.15
  [31c24e10] Distributions v0.25.123
  [ffbed154] DocStringExtensions v0.9.5
  [5b8099bc] DomainSets v0.7.16
  [7c1d4256] DynamicPolynomials v0.6.4
  [4e289a0a] EnumX v1.0.6
  [f151be2c] EnzymeCore v0.8.18
  [d4d017d3] ExponentialUtilities v1.30.0
  [e2ba6199] ExprTools v0.1.10
  [55351af7] ExproniconLite v0.10.14
  [9d29842c] FastAlmostBandedMatrices v0.1.6
  [7034ab61] FastBroadcast v0.3.5
  [9aa1b823] FastClosures v0.3.2
  [442a2c76] FastGaussQuadrature v1.1.0
  [a4df4552] FastPower v1.3.1
  [48062228] FilePathsBase v0.9.24
  [1a297f60] FillArrays v1.16.0
  [64ca27bc] FindFirstFunctions v1.8.0
  [6a86dc24] FiniteDiff v2.29.0
  [f6369f11] ForwardDiff v1.3.2
  [069b7b12] FunctionWrappers v1.1.3
  [77dc65aa] FunctionWrappersWrappers v0.1.3
  [46192b85] GPUArraysCore v0.2.0
  [c145ed77] GenericSchur v0.5.6
  [86223c79] Graphs v1.13.4
  [34004b35] HypergeometricFunctions v0.3.28
  [615f187c] IfElse v0.1.1
  [3263718b] ImplicitDiscreteSolve v1.6.0
  [d25df0c9] Inflate v0.1.5
  [842dd82b] InlineStrings v1.4.5
  [18e54dd8] IntegerMathUtils v0.1.3
  [8197267c] IntervalSets v0.7.13
  [3587e190] InverseFunctions v0.1.17
  [41ab1584] InvertedIndices v1.3.1
  [92d709cd] IrrationalConstants v0.2.6
  [82899510] IteratorInterfaceExtensions v1.0.0
  [692b3bcd] JLLWrappers v1.7.1
⌃ [682c06a0] JSON v0.21.4
  [ae98c720] Jieko v0.2.1
  [ccbc3e58] JumpProcesses v9.21.1
  [ba0b0d4f] Krylov v0.10.5
  [b964fa9f] LaTeXStrings v1.4.0
  [10f19ff3] LayoutPointers v0.1.17
  [5078a376] LazyArrays v2.9.5
  [2d8b4e74] LevyArea v1.0.0
  [87fe0de2] LineSearch v0.1.6
  [d3d80556] LineSearches v7.6.0
  [7ed4a6bd] LinearSolve v3.57.0
  [2ab3a3ac] LogExpFunctions v0.3.29
  [e6f89c97] LoggingExtras v1.2.0
  [1914dd2f] MacroTools v0.5.16
  [d125e4d3] ManualMemory v0.1.8
  [a3b82374] MatrixFactorizations v3.1.3
  [bb5d69b7] MaybeInplace v0.1.4
  [e1d29d7a] Missings v1.2.0
  [961ee093] ModelingToolkit v11.9.0
  [7771a370] ModelingToolkitBase v1.8.1
  [6bb917b9] ModelingToolkitTearing v1.2.4
  [2e0e35c7] Moshi v0.3.7
  [46d2c3a1] MuladdMacro v0.2.4
  [102ac46a] MultivariatePolynomials v0.5.13
  [d8a4904e] MutableArithmetics v1.6.7
  [d41bc354] NLSolversBase v8.0.0
  [77ba4419] NaNMath v1.1.3
  [8913a72c] NonlinearSolve v4.14.0
  [be0214bd] NonlinearSolveBase v2.11.1
  [5959db7a] NonlinearSolveFirstOrder v1.11.1
  [9a2c21bd] NonlinearSolveQuasiNewton v1.12.0
  [26075421] NonlinearSolveSpectralMethods v1.6.0
  [6fe1bfb0] OffsetArrays v1.17.0
  [bac558e1] OrderedCollections v1.8.1
  [1dea7af3] OrdinaryDiffEq v6.106.0
  [89bda076] OrdinaryDiffEqAdamsBashforthMoulton v1.9.0
  [6ad6398a] OrdinaryDiffEqBDF v1.14.0
  [bbf590c4] OrdinaryDiffEqCore v3.2.0
  [50262376] OrdinaryDiffEqDefault v1.12.0
  [4302a76b] OrdinaryDiffEqDifferentiation v1.22.0
  [9286f039] OrdinaryDiffEqExplicitRK v1.8.0
  [e0540318] OrdinaryDiffEqExponentialRK v1.12.0
  [becaefa8] OrdinaryDiffEqExtrapolation v1.13.0
  [5960d6e9] OrdinaryDiffEqFIRK v1.20.0
  [101fe9f7] OrdinaryDiffEqFeagin v1.8.0
  [d3585ca7] OrdinaryDiffEqFunctionMap v1.9.0
  [d28bc4f8] OrdinaryDiffEqHighOrderRK v1.9.0
  [9f002381] OrdinaryDiffEqIMEXMultistep v1.11.0
  [521117fe] OrdinaryDiffEqLinear v1.10.0
  [1344f307] OrdinaryDiffEqLowOrderRK v1.10.0
  [b0944070] OrdinaryDiffEqLowStorageRK v1.11.0
  [127b3ac7] OrdinaryDiffEqNonlinearSolve v1.19.0
  [c9986a66] OrdinaryDiffEqNordsieck v1.8.0
  [5dd0a6cf] OrdinaryDiffEqPDIRK v1.10.0
  [5b33eab2] OrdinaryDiffEqPRK v1.8.0
  [04162be5] OrdinaryDiffEqQPRK v1.8.0
  [af6ede74] OrdinaryDiffEqRKN v1.9.0
  [43230ef6] OrdinaryDiffEqRosenbrock v1.22.0
  [2d112036] OrdinaryDiffEqSDIRK v1.11.0
  [669c94d9] OrdinaryDiffEqSSPRK v1.11.0
  [e3e12d00] OrdinaryDiffEqStabilizedIRK v1.10.0
  [358294b1] OrdinaryDiffEqStabilizedRK v1.8.0
  [fa646aed] OrdinaryDiffEqSymplecticRK v1.11.0
  [b1df2697] OrdinaryDiffEqTsit5 v1.9.0
  [79d7bb75] OrdinaryDiffEqVerner v1.10.0
  [90014a1f] PDMats v0.11.37
  [69de0a69] Parsers v2.8.3
  [e409e4f3] PoissonRandom v0.4.7
  [f517fe37] Polyester v0.7.18
  [1d0040c9] PolyesterWeave v0.2.2
  [2dfb63ee] PooledArrays v1.4.3
⌅ [d236fae5] PreallocationTools v0.4.34
  [aea7be01] PrecompileTools v1.3.3
  [21216c6a] Preferences v1.5.1
  [08abe8d2] PrettyTables v3.1.2
  [27ebfcd6] Primes v0.5.7
  [43287f4e] PtrArrays v1.3.0
  [1fd47b50] QuadGK v2.11.2
  [74087812] Random123 v1.7.1
  [e6cf234a] RandomNumbers v1.6.0
  [988b38a3] ReadOnlyArrays v0.2.0
  [795d4caa] ReadOnlyDicts v1.0.1
  [3cdcf5f2] RecipesBase v1.3.4
  [731186ca] RecursiveArrayTools v3.47.0
  [189a3867] Reexport v1.2.2
  [ae029012] Requires v1.3.1
  [ae5879a3] ResettableStacks v1.2.0
  [79098fc4] Rmath v0.9.0
  [7e49a35a] RuntimeGeneratedFunctions v0.5.16
  [9dfe8606] SCCNonlinearSolve v1.9.0
  [94e857df] SIMDTypes v0.1.0
  [0bca4576] SciMLBase v2.136.0
  [19f34311] SciMLJacobianOperators v0.1.12
  [a6db7da4] SciMLLogging v1.8.0
  [c0aeaf25] SciMLOperators v1.14.1
  [431bcebd] SciMLPublic v1.0.1
  [53ae85a6] SciMLStructures v1.10.0
  [91c51154] SentinelArrays v1.4.9
  [efcf1570] Setfield v1.1.2
  [727e6d20] SimpleNonlinearSolve v2.10.0
  [699a6c99] SimpleTraits v0.9.5
  [a2af1166] SortingAlgorithms v1.2.2
  [9f842d2f] SparseConnectivityTracer v1.1.3
  [0a514795] SparseMatrixColorings v0.4.23
  [276daf66] SpecialFunctions v2.6.1
  [64909d44] StateSelection v1.3.0
  [aedffcd0] Static v1.3.1
  [0d7ed370] StaticArrayInterface v1.8.0
  [90137ffa] StaticArrays v1.9.16
  [1e83bf80] StaticArraysCore v1.4.4
  [10745b16] Statistics v1.11.1
  [82ae8749] StatsAPI v1.8.0
  [2913bbd2] StatsBase v0.34.10
  [4c63d2b9] StatsFuns v1.5.2
  [9672c7b4] SteadyStateDiffEq v2.9.0
  [789caeaf] StochasticDiffEq v6.92.0
  [7792a7ef] StrideArraysCore v0.5.8
  [892a3eda] StringManipulation v0.4.2
  [c3572dad] Sundials v5.1.0
  [3384d301] SymbolicCompilerPasses v0.1.1
  [2efcf032] SymbolicIndexingInterface v0.3.46
  [19f23fe9] SymbolicLimits v1.1.0
  [d1185830] SymbolicUtils v4.15.1
  [0c5d862f] Symbolics v7.11.0
  [3783bdb8] TableTraits v1.0.1
  [bd369af6] Tables v1.12.1
  [ed4db957] TaskLocalValues v0.1.3
  [8ea1fca8] TermInterface v2.0.0
  [8290d209] ThreadingUtilities v0.5.5
  [a759f4b9] TimerOutputs v0.5.29
  [3bb67fe8] TranscodingStreams v0.11.3
  [781d530d] TruncatedStacktraces v1.4.0
  [3a884ed6] UnPack v1.0.2
  [d30d5f5c] WeakCacheSets v0.1.0
  [ea10d353] WeakRefStrings v1.4.2
  [76eceee3] WorkerUtilities v1.6.1
  [1d5cc7b8] IntelOpenMP_jll v2025.2.0+0
  [856f044c] MKL_jll v2025.2.0+0
  [656ef2d0] OpenBLAS32_jll v0.3.30+0
  [efe28fd5] OpenSpecFun_jll v0.5.6+0
  [f50d1b31] Rmath_jll v0.5.1+0
  [ca45d3f4] SuiteSparse32_jll v7.12.1+0
  [fb77eaff] Sundials_jll v7.5.0+0
  [1317d2d5] oneTBB_jll v2022.0.0+1
  [0dad84c5] ArgTools v1.1.2
  [56f22d72] Artifacts v1.11.0
  [2a0f44e3] Base64 v1.11.0
  [ade2ca70] Dates v1.11.0
  [8ba89e20] Distributed v1.11.0
  [f43a241f] Downloads v1.6.0
  [7b1f6079] FileWatching v1.11.0
  [9fa8497b] Future v1.11.0
  [b77e0a4c] InteractiveUtils v1.11.0
  [ac6e5ff7] JuliaSyntaxHighlighting v1.12.0
  [4af54fe1] LazyArtifacts v1.11.0
  [b27032c2] LibCURL v0.6.4
  [76f85450] LibGit2 v1.11.0
  [8f399da3] Libdl v1.11.0
  [37e2e46d] LinearAlgebra v1.12.0
  [56ddb016] Logging v1.11.0
  [d6f4376e] Markdown v1.11.0
  [a63ad114] Mmap v1.11.0
  [ca575930] NetworkOptions v1.3.0
  [44cfe95a] Pkg v1.12.0
  [de0858da] Printf v1.11.0
  [9abbd945] Profile v1.11.0
  [3fa0cd96] REPL v1.11.0
  [9a3f8284] Random v1.11.0
  [ea8e919c] SHA v0.7.0
  [9e88b42a] Serialization v1.11.0
  [6462fe0b] Sockets v1.11.0
  [2f01184e] SparseArrays v1.12.0
  [f489334b] StyledStrings v1.11.0
  [4607b0f0] SuiteSparse
  [fa267f1f] TOML v1.0.3
  [a4e569a6] Tar v1.10.0
  [8dfed614] Test v1.11.0
  [cf7118a7] UUIDs v1.11.0
  [4ec0a83e] Unicode v1.11.0
  [e66e0078] CompilerSupportLibraries_jll v1.3.0+1
  [deac9b47] LibCURL_jll v8.11.1+1
  [e37daf67] LibGit2_jll v1.9.0+0
  [29816b5a] LibSSH2_jll v1.11.3+1
  [14a3606d] MozillaCACerts_jll v2025.5.20
  [4536629a] OpenBLAS_jll v0.3.29+0
  [05823500] OpenLibm_jll v0.8.7+0
  [458c3c95] OpenSSL_jll v3.5.1+0
  [bea87d4a] SuiteSparse_jll v7.8.3+2
  [83775a58] Zlib_jll v1.3.1+2
  [8e850b90] libblastrampoline_jll v5.15.0+0
  [8e850ede] nghttp2_jll v1.64.0+1
  [3f19e933] p7zip_jll v17.5.0+2
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated -m`

• Output of versioninfo()

Julia Version 1.12.1
Commit ba1e628ee49 (2025-10-17 13:02 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 32 × AMD EPYC 7343 16-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-18.1.7 (ORCJIT, znver3)
  GC: Built with stock GC
Threads: 1 default, 1 interactive, 1 GC (on 32 virtual cores)

[aryansri05]

I managed to reproduce the JIT lag reported in issue #1125 on my machine. It’s pretty clear NoSpecialize isn't holding the line like it should.
Since this is an in-place ROBER problem, do you think the specialization is leaking because of how FunctionWrappers are being handled in SciMLBase, or is it more likely that the solver's cache in OrdinaryDiffEq is triggering a new JIT session? Just want to make sure I’m looking at the right repo before I start hacking on a fix

[ChrisRackauckas]

NoSpecialize isn't using FunctionWrappers, it's supposed to just Any in the integrator type. So something probably regressed there (or changed in the Julia compiler), so I'd look into how that's being stored in the ODEIntegrator, i.e. integ.f should have a type parameter of Any or Void{Any} as a wrapper.

[aryansri05]

Thanks for the lead, Chris. That makes sense—if integ.f isn't correctly typed as Any or Void{Any}, it would explain why we're seeing specialization regardless of the flag.

I'll start by checking the ODEIntegrator struct definition and the init call in OrdinaryDiffEq.jl to see where the type parameter for f might be getting narrowed

One quick question: Is the integrator type constructed primarily in src/solve.jl, or should I look into the specific solver files (like src/solvers/bdf_utils.jl for FBDF) to see if they are overriding the default f storage?