diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index dd92f3bed..bc9b99ba2 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.35.0" +version = "2.35.1" authors = ["Avik Pal and contributors"] [deps] diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl index 636cddc37..2f21ba409 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl +++ b/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...) diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl index bea98b707..1205c0719 100644 --- a/lib/NonlinearSolveBase/src/linear_solve.jl +++ b/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) diff --git a/lib/NonlinearSolveFirstOrder/src/solve.jl b/lib/NonlinearSolveFirstOrder/src/solve.jl index 2e35c4400..e5aa73643 100644 --- a/lib/NonlinearSolveFirstOrder/src/solve.jl +++ b/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... ) diff --git a/lib/NonlinearSolveFirstOrder/src/trust_region.jl b/lib/NonlinearSolveFirstOrder/src/trust_region.jl index 7a90940f0..7421e860c 100644 --- a/lib/NonlinearSolveFirstOrder/src/trust_region.jl +++ b/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, diff --git a/lib/NonlinearSolveFirstOrder/test/operator_jacobian.jl b/lib/NonlinearSolveFirstOrder/test/operator_jacobian.jl index 882244f6c..d479a89b6 100644 --- a/lib/NonlinearSolveFirstOrder/test/operator_jacobian.jl +++ b/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.