Use public APIs for AutoDePSpecialize wrappers

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 <avikpal@mit.edu> 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"

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,
 ]

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,26 +288,27 @@ 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
     p = prob.p
     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

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