diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index 6abd48606..48fa3f23f 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -1,6 +1,6 @@ name = "NonlinearSolveBase" uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" -version = "2.38.1" +version = "2.38.2" authors = ["Avik Pal and contributors"] [deps] @@ -79,7 +79,7 @@ EnzymeCore = "0.8.16" FastClosures = "0.3.2" ForwardDiff = "0.10.36, 1" FunctionWrappers = "1.1.2" -FunctionWrappersWrappers = "1" +FunctionWrappersWrappers = "1.9.3" InteractiveUtils = "<0.0.1, 1" LineSearch = "0.1.4" LinearAlgebra = "1.10" diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl index d5fc68995..249af1dd5 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl @@ -8,7 +8,6 @@ using FastClosures: @closure using ForwardDiff: ForwardDiff, Dual, pickchunksize using FunctionWrappers: FunctionWrappers import FunctionWrappersWrappers -import RespecializeParams using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem, NonlinearProblem, NonlinearLeastSquaresProblem, ImmutableNonlinearProblem, remake using Setfield: @set @@ -22,6 +21,19 @@ import NonlinearSolveBase: wrapfun_iip, standardize_forwarddiff_tag const DI = DifferentiationInterface +_cache_storage_type( + ::FunctionWrappersWrappers.FunctionWrappersWrapper{FW, P, CS} +) where {FW, P, CS} = CS + +# Derive the concrete storage type through the public cache-mode constructor. +const FWW_SINGLE_CACHE_STORAGE_TYPE = _cache_storage_type( + FunctionWrappersWrappers.FunctionWrappersWrapper( + identity, (Tuple{Float64},), (Float64,); + cache = FunctionWrappersWrappers.SingleCache(), + policy = FunctionWrappersWrappers.AllowNonIsBits(), + ) +) + # --- AutoSpecialize / norecompile infrastructure for ForwardDiff --- const dualT = ForwardDiff.Dual{ @@ -79,14 +91,13 @@ function _make_fww_iip( @nospecialize(vff), ::Type{A1}, ::Type{A2}, ::Type{A3}, ::Type{A4} ) where {A1, A2, A3, A4} FW = FunctionWrappers.FunctionWrapper - fwt = ( - FW{Nothing, A1}(vff), FW{Nothing, A2}(vff), - FW{Nothing, A3}(vff), FW{Nothing, A4}(vff), - ) - cs = FunctionWrappersWrappers.SingleCacheStorage() + FWT = Tuple{ + FW{Nothing, A1}, FW{Nothing, A2}, + FW{Nothing, A3}, FW{Nothing, A4}, + } return FunctionWrappersWrappers.FunctionWrappersWrapper{ - typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs), - }(fwt, cs) + FWT, FunctionWrappersWrappers.AllowNonIsBits, FWW_SINGLE_CACHE_STORAGE_TYPE, + }(vff) end # IIP wrapfun: wraps f(du, u, p) with dual-aware type combinations. @@ -113,23 +124,6 @@ end ) end -# Opaque-`p` variant of `wrapfun_iip` for the AutoDePSpecialize path: same -# dual-aware shapes as above but with the parameter slot de-specialized to -# `RespecializeParams.OpaqueParams` (via `wrap_void_opaque`, which substitutes -# the third slot). `p` is never a `Dual` on this path (opaque-ification is -# skipped for dual state/params), so only the plain and Jacobian-`Dual`-`u` -# signatures are needed. -@inline function NonlinearSolveBase.wrapfun_iip_opaque( - ff, ::Type{P}, inputs::Tuple{T1, T2, T3} - ) where {P, T1 <: AbstractArray, T2 <: AbstractArray, T3} - T = eltype(T1) - dT = dualgen(T) - VdT = typeof(similar(inputs[1], dT)) - return RespecializeParams.wrap_void_opaque( - ff, P, (Tuple{T1, T2, T3}, Tuple{VdT, VdT, T3}) - ) -end - const GENERAL_SOLVER_TYPES = [ Nothing, NonlinearSolvePolyAlgorithm, ] diff --git a/lib/NonlinearSolveBase/src/autospecialize.jl b/lib/NonlinearSolveBase/src/autospecialize.jl index c840a4f54..824f47023 100644 --- a/lib/NonlinearSolveBase/src/autospecialize.jl +++ b/lib/NonlinearSolveBase/src/autospecialize.jl @@ -232,8 +232,8 @@ indexing read the concrete parameter and are not supported (reverse-mode support is a planned follow-up). Recover the payload with `RespecializeParams.unpack(sol.prob.p, typeof(p))`. -Problems carrying a symbolic system (`SciMLBase.has_sys(prob.f)`, e.g. -ModelingToolkit) are declined in [`maybe_opaque_wrap`] rather than here, since +Problems carrying a symbolic system (e.g. ModelingToolkit) are declined in +[`maybe_opaque_wrap`] rather than here, since that policy needs the function, not just `p`: their `p` must stay concrete for initialization and symbolic indexing. """ @@ -268,12 +268,8 @@ get_raw_f(f::AutoDePSpecializeCallable) = RespecializeParams.OpaqueVoid(f.ptype, _autodep_paramtype(f::AutoDePSpecializeCallable) = f.ptype -# Base (no ForwardDiff) opaque wrapper: single-signature. The ForwardDiff -# extension overrides this with the dual-aware variants. -function wrapfun_iip_opaque(ff, ::Type{P}, inputs::Tuple) where {P} - sig = Tuple{typeof(inputs[1]), typeof(inputs[2]), typeof(inputs[3])} - return RespecializeParams.wrap_void_opaque(ff, P, (sig,)) -end +@inline _has_symbolic_system(f) = + hasfield(typeof(f), :sys) && getfield(f, :sys) !== nothing """ maybe_opaque_wrap(prob) -> prob or nothing @@ -292,8 +288,7 @@ function maybe_opaque_wrap(prob::AbstractNonlinearProblem) SciMLBase.specialization(prob.f) === SciMLBase.AutoDePSpecialize || return nothing EnzymeCore.within_autodiff() && return nothing is_fw_wrapped(prob.f.f) && return nothing - (prob isa NonlinearProblem || prob isa SciMLBase.ImmutableNonlinearProblem) || - return nothing + SciMLBase.problem_type(prob) isa SciMLBase.StandardNonlinearProblem || return nothing SciMLBase.isinplace(prob) || return nothing u0 = prob.u0 @@ -301,17 +296,19 @@ function maybe_opaque_wrap(prob::AbstractNonlinearProblem) u0 isa AbstractArray || return nothing SciMLBase.isdualtype(eltype(u0)) && return nothing should_opaque_p(p) || return nothing - # Symbolic-system problems (`has_sys`, e.g. ModelingToolkit): decline. Their + # Symbolic-system problems (e.g. ModelingToolkit): decline. Their # `p` (MTKParameters) must stay concrete for the initialization pipeline and # symbolic parameter indexing; an opaque container breaks both. They gain # nothing (structured params are non-isbits and already type-uniform), so # falling back to the normal specialization is the correct no-op. - SciMLBase.has_sys(prob.f) && return nothing + _has_symbolic_system(prob.f) && return nothing P = typeof(p) orig = prob.f.f - fw = wrapfun_iip_opaque(orig, P, (u0, u0, p)) + opaque_p = RespecializeParams.pack_auto(p) + opaque_f = RespecializeParams.OpaqueVoid(P, orig) + fw = wrapfun_iip(opaque_f, (u0, u0, opaque_p)) newprob = @set prob.f.f = AutoDePSpecializeCallable{typeof(fw)}(fw, orig, P) - @set! newprob.p = RespecializeParams.pack_auto(p) + @set! newprob.p = opaque_p return newprob end diff --git a/lib/NonlinearSolveBase/test/autodepspecialize.jl b/lib/NonlinearSolveBase/test/autodepspecialize.jl index 55cfad6a9..9f2485c6e 100644 --- a/lib/NonlinearSolveBase/test/autodepspecialize.jl +++ b/lib/NonlinearSolveBase/test/autodepspecialize.jl @@ -1,5 +1,5 @@ using NonlinearSolveBase, SciMLBase, RespecializeParams, ForwardDiff, Test -using SymbolicIndexingInterface: SymbolCache +using SymbolicIndexingInterface: SymbolCache, symbolic_container # An `isbits` struct parameter. Under `AutoDePSpecialize`, `get_concrete_problem` # should pack `p` into a `RespecializeParams.OpaqueParams` and wrap the residual @@ -38,6 +38,18 @@ u0 = [1.0, 1.0] res = [0.0, 0.0] cp.f.f(res, [2.0, 3.0], cp.p) @test res ≈ [4.0 - 2.0, 9.0 - 3.0] + + DualT = ForwardDiff.Dual{ + ForwardDiff.Tag{NonlinearSolveBase.NonlinearSolveTag, Float64}, Float64, 1, + } + dual_u = DualT[ + DualT(2.0, ForwardDiff.Partials((1.0,))), + DualT(3.0, ForwardDiff.Partials((0.0,))), + ] + dual_res = similar(dual_u) + cp.f.f(dual_res, dual_u, cp.p) + @test ForwardDiff.value.(dual_res) ≈ res + @test first.(ForwardDiff.partials.(dual_res)) ≈ [4.0, 0.0] end @testset "AutoSpecialize / FullSpecialize leave p untouched" begin @@ -78,23 +90,23 @@ u0 = [1.0, 1.0] end @testset "symbolic-system problems are declined (MTK safety)" begin - # A problem whose `f` carries a symbolic system (`has_sys`, as every - # ModelingToolkit problem does) must NOT be opaque-ified: its `p` has to + # A problem whose `f` carries a symbolic system, as every ModelingToolkit + # problem does, must NOT be opaque-ified: its `p` has to # stay concrete for initialization and symbolic parameter indexing. # `SymbolCache` stands in for an MTK `System` without the dependency. fsym!(res, u, p) = (res[1] = u[1]^2 - 2.0; res[2] = u[2]^2 - 3.0; nothing) - ff = NonlinearFunction{true, SciMLBase.AutoDePSpecialize}( - fsym!; sys = SymbolCache([:x, :y], [:a, :b]) - ) - @test SciMLBase.has_sys(ff) + for sys in (SymbolCache([:x, :y], [:a, :b]), SymbolCache()) + ff = NonlinearFunction{true, SciMLBase.AutoDePSpecialize}(fsym!; sys) + @test symbolic_container(ff) === sys - cp = concretize(NonlinearProblem(ff, u0, VecQ([2.0]))) # non-isbits p - @test cp.p isa VecQ - @test !(cp.p isa RespecializeParams.OpaqueRef) + cp = concretize(NonlinearProblem(ff, u0, VecQ([2.0]))) # non-isbits p + @test cp.p isa VecQ + @test !(cp.p isa RespecializeParams.OpaqueRef) - cpi = concretize(NonlinearProblem(ff, u0, DePP(2.0, 3.0))) # isbits p - @test cpi.p isa DePP - @test !(cpi.p isa RespecializeParams.OpaqueParams) + cpi = concretize(NonlinearProblem(ff, u0, DePP(2.0, 3.0))) # isbits p + @test cpi.p isa DePP + @test !(cpi.p isa RespecializeParams.OpaqueParams) + end end @testset "already-packed p is not re-wrapped (idempotent)" begin