Preserve inference for operator-backed linear caches

lib/NonlinearSolveBase/Project.toml

@@ -1,6 +1,6 @@
 name = "NonlinearSolveBase"
 uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
-version = "2.35.0"
+version = "2.35.1"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]

lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl

@@ -6,6 +6,7 @@ using CommonSolve: CommonSolve, init, solve!
 using LinearSolve: LinearSolve, QRFactorization, SciMLLinearSolveAlgorithm
 using SciMLBase: SciMLBase, ReturnCode, LinearProblem, LinearAliasSpecifier
 using SciMLLogging: @SciMLMessage
+using SciMLOperators: AbstractSciMLOperator
 
 using LinearAlgebra: ColumnNorm, Symmetric
 
@@ -109,6 +110,23 @@ function set_lincache_A!(lincache, new_A)
     lincache.A = new_A
     return
 end
+function set_lincache_A!(lincache, new_A::AbstractSciMLOperator)
+    if lincache.A isa AbstractSciMLOperator
+        lincache.A = new_A
+        return
+    end
+
+    # A concrete cache owns its materialization, so refresh that buffer rather than
+    # rebinding it to the externally maintained operator.
+    A = convert(AbstractMatrix, new_A)
+    if ArrayInterface.can_setindex(lincache.A)
+        copyto!(lincache.A, A)
+        lincache.A = lincache.A
+    else
+        lincache.A = A
+    end
+    return
+end
 
 function LinearSolve.update_tolerances!(cache::LinearSolveJLCache; kwargs...)
     return LinearSolve.update_tolerances!(cache.lincache; kwargs...)

lib/NonlinearSolveBase/src/linear_solve.jl

@@ -91,9 +91,14 @@ function construct_linear_solver(
         # workspace), so no defensive copy of A is made
         linprob = LinearProblem(A, b, LinearSolveParameters(u_fixed, p); u0 = u_cache)
     elseif A isa AbstractSciMLOperator
-        # A SciMLOperator `A` is externally maintained (refreshed in place by
-        # `update_coefficients!`), so alias it: copying would sever those in-place updates
-        # and, for some operators, change the concrete type (breaking the later `A`-rebind).
+        # An owned materialization protects operator storage from factorizations and keeps
+        # large operator closures out of the default solver's cache type. Explicit
+        # matrix-free solvers retain the operator and its coefficient updates.
+        A = if linsolve === nothing || needs_concrete_A(linsolve)
+            copy(convert(AbstractMatrix, A))
+        else
+            A
+        end
         linprob = LinearProblem(A, b, LinearSolveParameters(u_fixed, p); u0 = u_cache)
         alias = LinearAliasSpecifier(alias_A = true, alias_b = false)
     elseif alias_A_for_refactorization(linsolve, A)

lib/NonlinearSolveFirstOrder/src/solve.jl

@@ -131,6 +131,12 @@ NonlinearSolveBase.@internal_caches(
     :jac_cache, :descent_cache, :linesearch_cache, :trustregion_cache, :forcing_cache
 )
 
+function maybe_unwrap_trustregion_prob(prob, alg, vjp_autodiff, jvp_autodiff)
+    return NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(
+        prob, vjp_autodiff, jvp_autodiff
+    )
+end
+
 function SciMLBase.__init(
         prob::AbstractNonlinearProblem, alg::GeneralizedFirstOrderAlgorithm, args...;
         stats = NLStats(0, 0, 0, 0, 0), alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false), maxiters = 1000,
@@ -176,12 +182,6 @@ function SciMLBase.__init(
         verbose = NonlinearVerbosity(verbose)
     end
 
-    # Enzyme cannot differentiate through FunctionWrappers' llvmcall.
-    # Create unwrapped prob for all AD-related constructions when using Enzyme.
-    _ad_prob = NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(
-        prob, alg.autodiff, alg.jvp_autodiff, alg.vjp_autodiff
-    )
-
     timer = get_timer_output()
     @static_timeit timer "cache construction" begin
         u = Utils.maybe_unaliased(prob.u0, alias_u0)
@@ -197,7 +197,7 @@ function SciMLBase.__init(
         linsolve_kwargs = merge((; verbose = verbose.linear_verbosity, abstol, reltol), linsolve_kwargs)
 
         jac_cache = NonlinearSolveBase.construct_jacobian_cache(
-            _ad_prob, alg, _ad_prob.f, fu, u, _ad_prob.p;
+            prob, alg, prob.f, fu, u, prob.p;
             stats, alg.autodiff, linsolve, alg.jvp_autodiff, alg.vjp_autodiff
         )
         J = reused_jacobian(jac_cache, u)
@@ -227,13 +227,16 @@ function SciMLBase.__init(
             # Standardize AD tags so VecJac/JacVec operators use NonlinearSolveTag
             # when the function is wrapped by AutoSpecialize.
             _tr_vjp_ad = NonlinearSolveBase.standardize_forwarddiff_tag(
-                alg.vjp_autodiff, _ad_prob
+                alg.vjp_autodiff, prob
             )
             _tr_jvp_ad = NonlinearSolveBase.standardize_forwarddiff_tag(
-                alg.jvp_autodiff, _ad_prob
+                alg.jvp_autodiff, prob
+            )
+            _tr_prob = maybe_unwrap_trustregion_prob(
+                prob, alg.trustregion, _tr_vjp_ad, _tr_jvp_ad
             )
             trustregion_cache = InternalAPI.init(
-                _ad_prob, alg.trustregion, _ad_prob.f, fu, u, _ad_prob.p;
+                _tr_prob, alg.trustregion, _tr_prob.f, fu, u, _tr_prob.p;
                 vjp_autodiff = _tr_vjp_ad, jvp_autodiff = _tr_jvp_ad,
                 stats, internalnorm, kwargs...
             )
@@ -257,20 +260,25 @@ function SciMLBase.__init(
             else
                 alg.jvp_autodiff
             end
-            _ls_ad = NonlinearSolveBase.standardize_forwarddiff_tag(ls_ad, _ad_prob)
+            _ls_ad = NonlinearSolveBase.standardize_forwarddiff_tag(ls_ad, prob)
+            _ls_prob = NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(prob, _ls_ad)
             linesearch_cache = CommonSolve.init(
-                _ad_prob, alg.linesearch, fu, u; stats, internalnorm,
+                _ls_prob, alg.linesearch, fu, u; stats, internalnorm,
                 autodiff = _ls_ad, kwargs...
             )
             globalization = Val(:LineSearch)
         end
 
         if has_forcing
+            _forcing_ad = ifelse(
+                provided_jvp_autodiff, alg.jvp_autodiff, alg.vjp_autodiff
+            )
+            _forcing_prob = NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(
+                prob, _forcing_ad
+            )
             forcing_cache = InternalAPI.init(
-                _ad_prob, alg.forcing, fu, u, u, _ad_prob.p; stats, internalnorm,
-                autodiff = ifelse(
-                    provided_jvp_autodiff, alg.jvp_autodiff, alg.vjp_autodiff
-                ),
+                _forcing_prob, alg.forcing, fu, u, u, _forcing_prob.p;
+                stats, internalnorm, autodiff = _forcing_ad,
                 verbose,
                 kwargs...
             )

lib/NonlinearSolveFirstOrder/src/trust_region.jl

@@ -201,6 +201,21 @@ the value used in the respective paper.
     initial_trust_radius = nothing
 end
 
+function maybe_unwrap_trustregion_prob(
+        prob, alg::GenericTrustRegionScheme, vjp_autodiff, jvp_autodiff
+    )
+    # Only these schemes construct AD-backed operators. Unwrapping the others can
+    # type-erase AutoSpecialize's function and prevent concrete cache inference.
+    if alg.method isa RUS.__Bastin
+        return NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(
+            prob, vjp_autodiff, jvp_autodiff
+        )
+    elseif alg.method isa RUS.__Yuan
+        return NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(prob, vjp_autodiff)
+    end
+    return prob
+end
+
 function InternalAPI.init(
         prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f, fu, u, p,
         args...; stats, internalnorm::F = L2_NORM, vjp_autodiff = nothing,

lib/NonlinearSolveFirstOrder/test/operator_jacobian.jl

@@ -1,10 +1,11 @@
 using NonlinearSolveFirstOrder, LinearSolve, SciMLOperators, SciMLBase
+using Enzyme # Activate the preferred reverse-mode backend before checking cache inference.
 using LinearAlgebra, SparseArrays, Test
 using SciMLOperators: AbstractSciMLOperator, isconvertible
 
 # A `jac_prototype` that is an `AbstractSciMLOperator` is handed to the solver as the
 # Jacobian directly: an iterative solver applies it matrix-free via `mul!`, while a
-# factorization materializes it lazily via `convert(AbstractMatrix, ·)`. The choice is
+# factorization materializes it via `convert(AbstractMatrix, ·)`. The choice is
 # routed by `needs_concrete_A(linsolve)` (like NLNewton); `isconvertible(op)` guards the
 # factorization path.
 
@@ -19,6 +20,11 @@ const xref = Wmat \ bvec
     @test isconvertible(mop)
     prob = NonlinearProblem(NonlinearFunction(resid!; jac_prototype = mop), zeros(N))
 
+    cache = @inferred init(prob, TrustRegion())
+    sol = solve!(cache)
+    @test SciMLBase.successful_retcode(sol)
+    @test sol.u ≈ xref
+
     for ls in (KrylovJL_GMRES(), LUFactorization(), KLUFactorization())
         cache = init(prob, NewtonRaphson(linsolve = ls))
         # The cache holds the operator itself, never a residual-derived JacobianOperator.