From bfffd5d9c61b04f8b939f73eb02696d2568e7921 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 29 Aug 2026 15:37:04 -0400 Subject: [PATCH 1/2] Trace the nonlinear solver loops under Reactant Run every algorithm's ordinary `solve` dispatch under `Reactant.@compile`. The shared solver loop is a `ReactantCore.@trace while` over `CommonSolve.step!`, return codes are traced enums (EnzymeAD/Reactant.jl#3232, SciMLBase#1563) so the solution is built by `SciMLBase.NonlinearSolution` on every path, and the loop-carried cache state is refreshed by one reflective `dealias_traced!` instead of per-cache field lists. Trust-region, Levenberg-Marquardt, geodesic and dogleg updates are written once with `ifelse`; polyalgorithms are traced as a chain of their members. Forward-mode `AutoEnzyme` is preferred during compilation. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Agent-Harness: Claude Code 2.1.251 Agent-Model: claude-fable-5 Agent-Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE --- AGENTS.md | 56 +++++ docs/src/devdocs/internal_interfaces.md | 11 + docs/src/tutorials/nonlinear_solve_gpus.md | 47 ++++- lib/NonlinearSolveBase/Project.toml | 6 +- .../src/NonlinearSolveBase.jl | 4 + lib/NonlinearSolveBase/src/abstract_types.jl | 2 +- lib/NonlinearSolveBase/src/autodiff.jl | 14 +- lib/NonlinearSolveBase/src/descent/common.jl | 4 +- .../src/descent/damped_newton.jl | 11 +- lib/NonlinearSolveBase/src/descent/dogleg.jl | 64 +++--- .../src/descent/geodesic_acceleration.jl | 20 +- lib/NonlinearSolveBase/src/descent/newton.jl | 1 + lib/NonlinearSolveBase/src/linear_solve.jl | 40 ++++ lib/NonlinearSolveBase/src/polyalg.jl | 23 +- lib/NonlinearSolveBase/src/reactant.jl | 79 +++++++ lib/NonlinearSolveBase/src/solve.jl | 96 +++++++-- .../src/termination_conditions.jl | 26 ++- lib/NonlinearSolveBase/src/tracing.jl | 4 + lib/NonlinearSolveBase/test/qa/qa.jl | 4 + lib/NonlinearSolveFirstOrder/Project.toml | 6 +- .../src/NonlinearSolveFirstOrder.jl | 3 +- .../src/levenberg_marquardt.jl | 38 ++-- lib/NonlinearSolveFirstOrder/src/solve.jl | 197 ++++++++++-------- .../src/trust_region.jl | 167 ++++++++------- lib/NonlinearSolveQuasiNewton/Project.toml | 4 +- lib/NonlinearSolveQuasiNewton/src/solve.jl | 19 +- .../Project.toml | 4 +- .../src/solve.jl | 17 +- lib/SCCNonlinearSolve/Project.toml | 2 +- .../src/SCCNonlinearSolve.jl | 16 +- lib/SciMLJacobianOperators/Project.toml | 2 +- .../src/SciMLJacobianOperators.jl | 2 +- test/Reactant/Project.toml | 37 ++++ test/Reactant/reactant_tests.jl | 166 +++++++++++++++ test/runtests.jl | 31 +++ test/test_groups.toml | 4 + 36 files changed, 932 insertions(+), 295 deletions(-) create mode 100644 AGENTS.md create mode 100644 lib/NonlinearSolveBase/src/reactant.jl create mode 100644 test/Reactant/Project.toml create mode 100644 test/Reactant/reactant_tests.jl diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..13d667f9c3 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,56 @@ +# Agent notes for NonlinearSolve.jl + +Repository-specific conventions and pitfalls for automated contributors. Global operating +rules live outside this repository; keep this file to facts about this codebase. + +## Validation before pushing + +- Every sublibrary under `lib/` has its own project and test groups. Run + `GROUP=Core julia --project=lib/ -e 'using Pkg; Pkg.test()'` and `GROUP=QA` for + each sublibrary whose sources you changed; the root `Pkg.test()` resolves the *released* + sublibraries and does not test local `lib/` changes. +- Sublibrary `[sources]` entries are relative paths (`../NonlinearSolveBase`). Pkg resolves + them against the *active* project, so scratch environments used to test a sublibrary + must live under `lib/` (a sibling directory), never under `/tmp` or `~/tmp`. +- Never `Pkg.develop` into a sublibrary's own project: Pkg rewrites its `Project.toml` + (`[sources]` with absolute paths, extra dependencies, reordered sections). Use a scratch + project under `lib/` and develop the sublibraries into that. +- Format with Runic.jl (`julia -e 'using Runic; Runic.main(["--inplace", files...])'`), not + the unrelated `runic` binary that may be on `PATH`, and run `typos`. +- When building commits from a working tree that predates recent `master` commits, merge + `master` first. A diff against `origin/master` taken from an older base silently reverts + the newer commits in every file it touches. + +## Reactant support (`ReactantCore.@trace` in the solver loops) + +The solver loops trace under `Reactant.@compile` by running the *ordinary* code path; the +helpers are in `lib/NonlinearSolveBase/src/reactant.jl`. + +- `ReactantCore.within_compile()` is a compile-time `false` outside Reactant, so a helper + that checks it costs nothing on the host. Call it only from ordinary functions: inside a + `ReactantCore.@trace while`/`if` body it returns `false` even while tracing, because the + macro captures every symbol of the body, the module included, as a loop variable. Loop + bodies call self-gating helpers (`dealias_traced!`, `Utils.fresh`, ...) unconditionally. +- Solver caches are `@concrete`, so scalar loop state (`nsteps`, `force_stop`, `retcode`, + trust-region counters) must be traced at construction with `maybe_traced`; promoting a + field after the cache exists fails with a type error. +- Reactant records one path per traced object among loop-carried values and requires the + same set after each iteration. Aliases between cache fields break this; `dealias_traced!` + refreshes every traced leaf at the loop boundary. `@bb copyto!(dst, src)` and `@bb copy` + rebind (`dst = src`) for traced arrays, so they create aliases. +- Write decisions on traced values with `ifelse`/`select`, not `if`. A `@trace if` that + mutates the cache is avoided in the loop path; where a `@trace if` is used, its branches + must assign only loop-state leaves (a whole solution object as a branch output has to be + materialized for the untaken side), and no variable in scope may be named `args`, which + the macro uses for its captured-variable bundle. +- Under compilation the returned `NonlinearSolution` has `stats === nothing` (`NLStats` + holds `Int`s) and `prob === nothing` (a problem's `Base.Pairs` keyword arguments cannot + be rebuilt by Reactant's result codegen); the default termination mode is + `AbsNormTerminationMode`, the Jacobian-reuse policy is switched off, and initialization + failure cannot be reported. +- Not traceable at present, with the reason recorded next to the Reactant test matrix in + `test/Reactant/reactant_tests.jl`: line-search globalization (`norm(x, Inf)` scalar-indexes + under Reactant), `DFSane`, and trust-region schemes needing a reverse-mode + vector-Jacobian product (`RobustMultiNewton`). +- The Reactant test group pins unreleased branches of Reactant, SciMLBase and + DifferentiationInterface in `test/runtests.jl`; update the pins when those release. diff --git a/docs/src/devdocs/internal_interfaces.md b/docs/src/devdocs/internal_interfaces.md index 9c2cb39fee..df83f7541b 100644 --- a/docs/src/devdocs/internal_interfaces.md +++ b/docs/src/devdocs/internal_interfaces.md @@ -152,3 +152,14 @@ NonlinearSolveBase.get_reltol NonlinearSolveBase.AbstractNonlinearTerminationMode NonlinearSolveBase.AbstractSafeNonlinearTerminationMode ``` + +## Reactant Support + +Helpers used by the solver packages so that their loops trace under `Reactant.@compile`. + +```@docs +NonlinearSolveBase.maybe_traced +NonlinearSolveBase.dealias_traced! +NonlinearSolveBase.select +NonlinearSolveBase.build_nonlinear_solution +``` diff --git a/docs/src/tutorials/nonlinear_solve_gpus.md b/docs/src/tutorials/nonlinear_solve_gpus.md index 3afb64893c..5c4ab3bb19 100644 --- a/docs/src/tutorials/nonlinear_solve_gpus.md +++ b/docs/src/tutorials/nonlinear_solve_gpus.md @@ -9,7 +9,7 @@ NonlinearSolve.jl supports GPU acceleration on a wide array of devices, such as: | Intel | OneAPI | [OneAPI.jl](https://github.com/JuliaGPU/oneAPI.jl) | `oneAPI.oneAPIBackend()` | | Apple (M-Series) | Metal | [Metal.jl](https://github.com/JuliaGPU/Metal.jl) | `Metal.MetalBackend()` | -To use NonlinearSolve.jl on GPUs, there are two distinctly different approaches: +To use NonlinearSolve.jl on GPUs, there are three distinctly different approaches: 1. You can build a `NonlinearProblem` / `NonlinearLeastSquaresProblem` where the elements of the problem, i.e. `u0` and `p`, are defined on GPUs. This will make the evaluations @@ -20,6 +20,9 @@ To use NonlinearSolve.jl on GPUs, there are two distinctly different approaches: system over a large number of inputs. This is useful for cases where you have a small `NonlinearProblem` / `NonlinearLeastSquaresProblem` which you want to solve over a large number of initial guesses or parameters. + 3. You can compile a complete `solve` call with + [Reactant.jl](https://enzymead.github.io/Reactant.jl/stable/). This keeps the nonlinear + iterations inside one compiled program and lets their count depend on runtime inputs. For a deeper dive into the computational difference between these techniques and why it leads to different pros/cons, see the @@ -28,7 +31,7 @@ In particular, the second form is unique to NonlinearSolve.jl and offers orders performance improvements over libraries in Jax and PyTorch which are restricted to only using the first form. -In this tutorial we will highlight both use cases in separate parts. +In this tutorial we will highlight these use cases in separate parts. !!! note @@ -67,6 +70,46 @@ notice that `cu` arrays automatically default to `Float32` precision. Since Nonl respects the user's chosen types, this changes NonlinearSolve.jl to use `Float32` precision, and thus the tolerances are adjusted accordingly. +## Whole-solve compilation with Reactant.jl + +Reactant arrays can be passed through the standard `NonlinearProblem` and `solve` APIs: + +```julia +import NonlinearSolve as NLS +import Reactant + +f(u, p) = u .* u .- p + +function reactant_solve(u0, p) + prob = NLS.NonlinearProblem(f, u0, p) + return NLS.solve(prob, NLS.SimpleBroyden()) +end + +u0 = Reactant.to_rarray(Float32[1, 1]) +p = Reactant.to_rarray(Float32[2]) +sol = Reactant.@jit reactant_solve(u0, p) +``` + +The nonlinear iteration uses a traced `while` operation. The array shapes are fixed for a +compiled executable, but convergence and the number of quasi-Newton steps are determined at +runtime. Use `Reactant.@compile` instead of `Reactant.@jit` when the executable will be +called repeatedly with new initial values or parameters of the same shape. + +Algorithms follow their normal `solve` dispatch during Reactant compilation. There is no +separate allowlist or fallback algorithm: unsupported operations report their errors from +Reactant or the package that implements them. `SimpleBroyden` and `SimpleKlement` are tested +with square, out-of-place `NonlinearProblem`s. Jacobian-based algorithms additionally depend +on Reactant support in their configured differentiation backend; NonlinearSolve.jl does not +provide Reactant-specific differentiation overloads. + +The returned solution is an ordinary `NonlinearSolution`. `u` and `resid` are device +arrays, and `retcode` is a device scalar (`ConcreteRNumber{ReturnCode.T}`) that can be +compared against `ReturnCode` values or converted with `ReturnCode.T(sol.retcode)`; +`SciMLBase.successful_retcode(sol)` works as usual. `stats` is `nothing`, since `NLStats` +counts host evaluations, which inside a compiled program only happen once while tracing, +and `prob` is `nothing`, since a problem's keyword arguments cannot be returned from a +compiled program. + ## GPU Acceleration over Large Parameter Searches using KernelAbstractions.jl If one has a "small" (200 equations or less) system of equations which they wish to solve diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index 3e2443b4d9..ce051f144b 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.49.3" +version = "2.50.0" authors = ["Avik Pal and contributors"] [deps] @@ -10,6 +10,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" @@ -25,6 +26,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433" RespecializeParams = "9fe22ead-9e00-4db2-8b46-706a60d40f5e" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" @@ -74,6 +76,7 @@ ChainRulesCore = "1" CommonSolve = "0.2.4" Compat = "4.15" ConcreteStructs = "0.2.3" +ConstructionBase = "1.5" DifferentiationInterface = "0.7.3" EnumX = "1" Enzyme = "0.13.90" @@ -95,6 +98,7 @@ PrecompileTools = "1.2" Preferences = "1.4" Printf = "1.10" RecursiveArrayTools = "4" +ReactantCore = "0.1.21" ReverseDiff = "1.15" RespecializeParams = "1" SciMLBase = "3.46" diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index 761dab5e9a..22625ac5fd 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -37,9 +37,11 @@ import RespecializeParams using StaticArraysCore: StaticArray, SMatrix, SArray, MArray using CommonSolve: CommonSolve, init +using ConstructionBase: ConstructionBase using EnzymeCore: EnzymeCore using MaybeInplace: @bb using RecursiveArrayTools: RecursiveArrayTools, AbstractVectorOfArray, ArrayPartition +using ReactantCore: ReactantCore using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinearProblem, AbstractNonlinearAlgorithm, _concrete_solve_adjoint, _concrete_solve_forward, NonlinearProblem, NonlinearLeastSquaresProblem, @@ -76,6 +78,7 @@ include("utils.jl") include("verbosity.jl") include("abstract_types.jl") +include("reactant.jl") include("common_defaults.jl") include("termination_conditions.jl") @@ -113,6 +116,7 @@ include("forward_diff.jl") ) @compat(public, (get_abstol, get_reltol)) +@compat(public, (maybe_traced, dealias_traced!, select, build_nonlinear_solution)) @compat(public, (AbstractNonlinearTerminationMode, AbstractSafeNonlinearTerminationMode)) @compat(public, (nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution)) @compat( diff --git a/lib/NonlinearSolveBase/src/abstract_types.jl b/lib/NonlinearSolveBase/src/abstract_types.jl index 49ff720156..28350339a9 100644 --- a/lib/NonlinearSolveBase/src/abstract_types.jl +++ b/lib/NonlinearSolveBase/src/abstract_types.jl @@ -720,7 +720,7 @@ function has_time_limit(cache::AbstractNonlinearSolveCache) end function not_terminated(cache::AbstractNonlinearSolveCache) - return !cache.force_stop && cache.nsteps < cache.maxiters + return (!cache.force_stop) & (cache.nsteps < cache.maxiters) end _prepare_reinit_parameters(p, ::Any) = SciMLBase.unwrap_parameters(p) diff --git a/lib/NonlinearSolveBase/src/autodiff.jl b/lib/NonlinearSolveBase/src/autodiff.jl index 7c9e649dd0..f6ca6de9b8 100644 --- a/lib/NonlinearSolveBase/src/autodiff.jl +++ b/lib/NonlinearSolveBase/src/autodiff.jl @@ -25,7 +25,8 @@ Choose a forward-mode-compatible automatic differentiation backend for `prob`. If `ad` is an `AbstractADType`, the backend is returned when it is available and compatible with the problem. If `ad === nothing`, NonlinearSolveBase selects the first available -compatible backend from its preferred forward-mode list. +compatible backend from its preferred forward-mode list. During Reactant compilation, +forward-mode `AutoEnzyme` is preferred when available. ### Arguments @@ -76,6 +77,10 @@ function select_forward_mode_autodiff( prob::AbstractNonlinearProblem, ::Nothing; warn_check_mode::Bool = true ) + if ReactantCore.within_compile() + ad = ADTypes.AutoEnzyme(; mode = EnzymeCore.Forward) + !incompatible_backend_and_problem(prob, ad) && return ad + end idx = findfirst(!Base.Fix1(incompatible_backend_and_problem, prob), ForwardADs) idx !== nothing && return ForwardADs[idx] throw(ArgumentError("No forward mode AD backend is compatible with the chosen problem. \ @@ -142,7 +147,8 @@ Choose an automatic differentiation backend for constructing Jacobians for `prob If `ad === nothing`, NonlinearSolveBase prefers a compatible forward-mode backend that is not finite differencing, then falls back to compatible reverse-mode or finite-difference -backends. +backends. During Reactant compilation, forward-mode `AutoEnzyme` is preferred when +available. ### Arguments @@ -166,6 +172,10 @@ function select_jacobian_autodiff(prob::AbstractNonlinearProblem, ad::AbstractAD end function select_jacobian_autodiff(prob::AbstractNonlinearProblem, ::Nothing) + if ReactantCore.within_compile() + ad = ADTypes.AutoEnzyme(; mode = EnzymeCore.Forward) + !incompatible_backend_and_problem(prob, ad) && return ad + end idx = findfirst(!Base.Fix1(incompatible_backend_and_problem, prob), ForwardADs) idx !== nothing && !is_finite_differences_backend(ForwardADs[idx]) && return ForwardADs[idx] diff --git a/lib/NonlinearSolveBase/src/descent/common.jl b/lib/NonlinearSolveBase/src/descent/common.jl index ebf3866f1a..596dc9da8c 100644 --- a/lib/NonlinearSolveBase/src/descent/common.jl +++ b/lib/NonlinearSolveBase/src/descent/common.jl @@ -20,13 +20,13 @@ Construct a `DescentResult` object. @concrete struct DescentResult δu u - success::Bool + success linsolve_success::Bool extras end function DescentResult(; - δu = missing, u = missing, success::Bool = true, linsolve_success::Bool = true, + δu = missing, u = missing, success = true, linsolve_success::Bool = true, extras = (;) ) @assert δu !== missing || u !== missing diff --git a/lib/NonlinearSolveBase/src/descent/damped_newton.jl b/lib/NonlinearSolveBase/src/descent/damped_newton.jl index c998704540..a5650700b1 100644 --- a/lib/NonlinearSolveBase/src/descent/damped_newton.jl +++ b/lib/NonlinearSolveBase/src/descent/damped_newton.jl @@ -75,9 +75,8 @@ function InternalAPI.init( @assert pre_inverted isa Val{false} "Precomputed Inverse for Non-Square Jacobian doesn't make sense." @bb δu = zero(u) - δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i - @bb δu_ = zero(u) - end + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : + collect(ntuple(_ -> zero(u), Utils.unwrap_val(shared) - 1)) normal_form_damping = returns_norm_form_damping(alg.damping_fn) normal_form_linsolve = needs_square_A(alg.linsolve, u) @@ -276,7 +275,7 @@ function InternalAPI.solve!( copyto!(@view(cache.J[1:size(J, 1), :]), J) cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) else - cache.J = Utils.faster_vcat(J, sqrt.(D)) + cache.J = Utils.faster_vcat(J, sqrt_damping(D)) end end A = cache.J @@ -340,6 +339,10 @@ function InternalAPI.solve!( return DescentResult(; δu) end +sqrt_damping(D::Number) = sqrt(D) +sqrt_damping(D::Diagonal) = Diagonal(sqrt.(D.diag)) +sqrt_damping(D) = sqrt.(D) + dampen_jacobian!!(::Any, J::Union{AbstractSciMLOperator, Number}, D) = J + D # Scalar damping (identity-style `(1/α) I` damping) only touches the diagonal of `J`. diff --git a/lib/NonlinearSolveBase/src/descent/dogleg.jl b/lib/NonlinearSolveBase/src/descent/dogleg.jl index fba06674be..af98821418 100644 --- a/lib/NonlinearSolveBase/src/descent/dogleg.jl +++ b/lib/NonlinearSolveBase/src/descent/dogleg.jl @@ -97,8 +97,11 @@ function InternalAPI.solve!( cache.newton_cache, J, fu, u, idx; skip_solve, kwargs... ).δu - # Newton's Step within the trust region - if cache.internalnorm(δu_newton) ≤ trust_region + # Under Reactant no branch can be taken on the traced norms, so every candidate step is + # formed and the result selected at the end; on the ordinary path the early returns + # avoid the extra work. + newton_step = cache.internalnorm(δu_newton) ≤ trust_region + if !ReactantCore.within_compile() && newton_step @bb copyto!(δu, δu_newton) set_du!(cache, δu, idx) return DescentResult(; δu, extras = (; δuJᵀJδu = T(NaN))) @@ -106,27 +109,13 @@ function InternalAPI.solve!( # Take intersection of steepest descent direction and trust region if Cauchy point # lies outside of trust region - if normal_form(cache) - δu_cauchy = cache.newton_cache.Jᵀfu_cache - JᵀJ = cache.newton_cache.JᵀJ_cache - @bb @. δu_cauchy *= -1 - - l_grad = cache.internalnorm(δu_cauchy) - @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - δuJᵀJδu = Utils.safe_dot(δu_cauchy, cache.δu_cache_mul) - else - δu_cauchy = InternalAPI.solve!( - cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs... - ).δu - J_ = preinverted_jacobian(cache) ? inv(J) : J - l_grad = cache.internalnorm(δu_cauchy) - @bb cache.Jᵀδu_cache = J_ × vec(δu_cauchy) - δuJᵀJδu = Utils.safe_dot(cache.Jᵀδu_cache, cache.Jᵀδu_cache) - end + δu_cauchy, l_grad, δuJᵀJδu = dogleg_cauchy_step!( + cache, J, fu, u, idx; skip_solve, kwargs... + ) d_cauchy = (l_grad^3) / δuJᵀJδu - - if d_cauchy ≥ trust_region - λ = trust_region / l_grad + cauchy_step = d_cauchy ≥ trust_region + λ = trust_region / l_grad + if !ReactantCore.within_compile() && cauchy_step @bb @. δu = λ * δu_cauchy set_du!(cache, δu, idx) return DescentResult(; δu, extras = (; δuJᵀJδu = λ^2 * δuJᵀJδu)) @@ -142,10 +131,33 @@ function InternalAPI.solve!( a = Utils.safe_dot(cache.δu_cache_2, cache.δu_cache_2) b = 2 * Utils.safe_dot(cache.δu_cache_1, cache.δu_cache_2) c = d_cauchy^2 - trust_region^2 - aux = max(0, b^2 - 4 * a * c) + aux = max(zero(a), b^2 - 4 * a * c) τ = (-b + sqrt(aux)) / (2 * a) - - @bb @. δu = cache.δu_cache_1 + τ * cache.δu_cache_2 + @bb @. cache.δu_cache_mul = cache.δu_cache_1 + τ * cache.δu_cache_2 + @bb @. δu = ifelse( + newton_step, δu_newton, ifelse(cauchy_step, λ * δu_cauchy, cache.δu_cache_mul) + ) set_du!(cache, δu, idx) - return DescentResult(; δu, extras = (; δuJᵀJδu = T(NaN))) + δuJᵀJδu_result = ifelse((!newton_step) & cauchy_step, λ^2 * δuJᵀJδu, T(NaN)) + return DescentResult(; δu, extras = (; δuJᵀJδu = δuJᵀJδu_result)) +end + +function dogleg_cauchy_step!(cache::DoglegCache, J, fu, u, idx; skip_solve, kwargs...) + if normal_form(cache) + δu_cauchy = cache.newton_cache.Jᵀfu_cache + JᵀJ = cache.newton_cache.JᵀJ_cache + @bb @. δu_cauchy *= -1 + l_grad = cache.internalnorm(δu_cauchy) + @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) + δuJᵀJδu = Utils.safe_dot(δu_cauchy, cache.δu_cache_mul) + else + δu_cauchy = InternalAPI.solve!( + cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs... + ).δu + J_ = preinverted_jacobian(cache) ? inv(J) : J + l_grad = cache.internalnorm(δu_cauchy) + @bb cache.Jᵀδu_cache = J_ × vec(δu_cauchy) + δuJᵀJδu = Utils.safe_dot(cache.Jᵀδu_cache, cache.Jᵀδu_cache) + end + return δu_cauchy, l_grad, δuJᵀJδu end diff --git a/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl index f3211728de..8f69151a53 100644 --- a/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl +++ b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl @@ -46,7 +46,7 @@ get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) Jv fu_cache u_cache - last_step_accepted::Bool + last_step_accepted end function InternalAPI.reinit_self!(cache::GeodesicAccelerationCache; p = cache.p, kwargs...) @@ -77,9 +77,8 @@ function InternalAPI.init( ) where {F} T = promote_type(eltype(u), eltype(fu)) @bb δu = zero(u) - δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i - @bb δu_ = zero(u) - end + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : + collect(ntuple(_ -> zero(u), Utils.unwrap_val(shared) - 1)) descent_cache = InternalAPI.init( prob, alg.descent, J, fu, u; shared = Val(2 * Utils.unwrap_val(shared)), pre_inverted, linsolve_kwargs, @@ -89,9 +88,10 @@ function InternalAPI.init( @bb Jv = similar(fu) @bb fu_cache = copy(fu) @bb u_cache = similar(u) + last_step_accepted = maybe_traced(false) return GeodesicAccelerationCache( δu, δus, descent_cache, prob.f, prob.p, T(alg.α), internalnorm, - T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false + T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, last_step_accepted ) end @@ -124,13 +124,9 @@ function InternalAPI.solve!( norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) - if 2 * norm_a ≤ norm_v * cache.α - @bb @. δu = v + a / 2 - set_du!(cache, δu, idx) - cache.last_step_accepted = true - else - cache.last_step_accepted = false - end + cache.last_step_accepted = 2 * norm_a ≤ norm_v * cache.α + @bb @. δu = ifelse(cache.last_step_accepted, v + a / 2, δu) + set_du!(cache, δu, idx) return DescentResult(; δu, success = cache.last_step_accepted, extras = (; a, v)) end diff --git a/lib/NonlinearSolveBase/src/descent/newton.jl b/lib/NonlinearSolveBase/src/descent/newton.jl index aed3700914..276c912c7a 100644 --- a/lib/NonlinearSolveBase/src/descent/newton.jl +++ b/lib/NonlinearSolveBase/src/descent/newton.jl @@ -25,6 +25,7 @@ end @internal_caches NewtonDescentCache :lincache + function InternalAPI.init( prob::AbstractNonlinearProblem, alg::NewtonDescent, J, fu, u; stats, shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl index becc654fd0..ab49306362 100644 --- a/lib/NonlinearSolveBase/src/linear_solve.jl +++ b/lib/NonlinearSolveBase/src/linear_solve.jl @@ -20,8 +20,23 @@ end stats::NLStats end +@concrete mutable struct ReactantLinearSolveCache <: AbstractLinearSolverCache + A + b + u + p + linsolve + kwargs + stats::NLStats +end + SciMLBase.reinit!(::NativeJLLinearSolveCache; kwargs...) = nothing +function _without_linear_tolerances(kwargs::NamedTuple) + names = filter(name -> name !== :abstol && name !== :reltol, keys(kwargs)) + return NamedTuple{names}(map(name -> kwargs[name], names)) +end + """ construct_linear_solver(alg, linsolve, A, b, u, p; stats, kwargs...) @@ -76,6 +91,9 @@ function construct_linear_solver( ) if (A isa Number && b isa Number) || (A isa Diagonal) return NativeJLLinearSolveCache(A, b, stats) + elseif ReactantCore.within_compile() + reactant_kwargs = _without_linear_tolerances((; kwargs...)) + return ReactantLinearSolveCache(A, b, u, p, linsolve, reactant_kwargs, stats) elseif linsolve isa typeof(\) return NativeJLLinearSolveCache(A, b, stats) elseif linsolve === nothing @@ -117,6 +135,28 @@ function construct_linear_solver( return LinearSolveJLCache(lincache, linsolve, stats) end +function (cache::ReactantLinearSolveCache)(; + A = nothing, b = nothing, linu = nothing, reuse_A_if_factorization = false, + kwargs... + ) + cache.stats.nsolve += 1 + A === nothing || (cache.A = A) + b === nothing || (cache.b = b) + linu === nothing || (cache.u = linu) + + linprob = LinearProblem( + cache.A, cache.b, LinearSolveParameters(cache.u, cache.p); u0 = cache.u + ) + solve_kwargs = merge(cache.kwargs, (; kwargs...)) + linres = cache.linsolve === nothing ? + SciMLBase.solve(linprob; solve_kwargs...) : + SciMLBase.solve(linprob, cache.linsolve; solve_kwargs...) + cache.u = linres.u + return LinearSolveResult( + ; u = linres.u, success = linres.retcode !== ReturnCode.Failure + ) +end + """ alias_A_for_refactorization(linsolve, A)::Bool diff --git a/lib/NonlinearSolveBase/src/polyalg.jl b/lib/NonlinearSolveBase/src/polyalg.jl index cba71388dd..ec95f525b8 100644 --- a/lib/NonlinearSolveBase/src/polyalg.jl +++ b/lib/NonlinearSolveBase/src/polyalg.jl @@ -405,19 +405,28 @@ function build_solution_less_specialize( store_original::Val = Val(false), kwargs... ) if store_original isa Val{true} + return less_specialized_solution( + Any, u, resid, prob, alg, retcode, original, left, right, stats, trace + ) + end + return less_specialized_solution( + Nothing, u, resid, prob, alg, retcode, nothing, left, right, stats, trace + ) +end + +# SciMLBase ≥ 3.51 appends the return code type to `NonlinearSolution`'s parameters. +let retcode_param = fieldtype(SciMLBase.NonlinearSolution, :retcode) === ReturnCode.T ? + () : (:(typeof(retcode)),) + @eval function less_specialized_solution( + ::Type{O}, u, resid, prob, alg, retcode, original, left, right, stats, trace + ) where {O} return SciMLBase.NonlinearSolution{ eltype(eltype(u)), ndims(u), typeof(u), typeof(resid), typeof(prob), - typeof(alg), Any, typeof(left), typeof(stats), typeof(trace), + typeof(alg), O, typeof(left), typeof(stats), typeof(trace), $(retcode_param...), }( u, resid, prob, alg, retcode, original, left, right, stats, trace ) end - return SciMLBase.NonlinearSolution{ - eltype(eltype(u)), ndims(u), typeof(u), typeof(resid), typeof(prob), - typeof(alg), Nothing, typeof(left), typeof(stats), typeof(trace), - }( - u, resid, prob, alg, retcode, nothing, left, right, stats, trace - ) end function findmin_caches(prob::AbstractNonlinearProblem, caches) diff --git a/lib/NonlinearSolveBase/src/reactant.jl b/lib/NonlinearSolveBase/src/reactant.jl new file mode 100644 index 0000000000..1220a87516 --- /dev/null +++ b/lib/NonlinearSolveBase/src/reactant.jl @@ -0,0 +1,79 @@ +# Support for running the solver loops under `Reactant.@compile`/`@jit`. Outside a +# Reactant compilation `ReactantCore.within_compile()` is a compile-time `false`, so every +# helper here folds away on the ordinary Julia path. +# +# `within_compile()` must only be called from ordinary functions, never directly inside a +# `ReactantCore.@trace` body: the macro captures every symbol of the body, the module +# included, as a loop-carried variable, and the call then no longer goes through Reactant's +# overlay and returns `false` while tracing. The helpers below gate themselves, so loop +# bodies call them unconditionally. + +""" + maybe_traced(x) + +Return `x` promoted to a traced scalar when called during a Reactant compilation and `x` +itself otherwise. Loop-carried scalars (`nsteps`, `force_stop`, `retcode`, ...) have to be +traced before a `ReactantCore.@trace while` loop for their final values to be visible after +it, and since solver caches are `@concrete`, at the point the cache is constructed. +""" +maybe_traced(x) = ReactantCore.within_compile() ? ReactantCore.promote_to_traced(x) : x + +""" + dealias_traced!(x) + +Replace every traced array or scalar reachable from `x` with a fresh copy so that no two +places share a traced value. Reactant records only one path per traced object among the +values carried by a `@trace while` loop and requires the set of paths to be the same before +and after each iteration; solver caches alias freely (`u_cache`/`u`, `p` in several caches, +the problem inside the trace) and steps rebind fields, so every value is made distinct at +the loop boundary instead. Mutable structs are updated in place, immutable ones are rebuilt. +""" +function dealias_traced!(x) + ReactantCore.within_compile() || return x + ReactantCore.is_traced(x) || return x + # Traced arrays are dense; structured wrappers (`Diagonal`, ...) are walked as structs so + # that they keep their type. + x isa DenseArray && return x .+ zero(eltype(x)) + x isa Number && return x + zero(x) + x isa Union{Tuple, NamedTuple} && return map(dealias_traced!, x) + T = typeof(x) + if ismutable(x) + for name in fieldnames(T) + isdefined(x, name) || continue + setfield!(x, name, dealias_traced!(getfield(x, name))) + end + return x + end + names = fieldnames(T) + isempty(names) && return x + values = map(name -> dealias_traced!(getfield(x, name)), names) + return ConstructionBase.setproperties(x, NamedTuple{names}(values)) +end + +""" + build_nonlinear_solution(prob, alg, u, resid; retcode, stats = nothing, trace = nothing) + +`SciMLBase.build_solution` for a nonlinear problem, except that during a Reactant +compilation the solution stores `nothing` as its problem: a problem carries its keyword +arguments as a `Base.Pairs`, which Reactant cannot rebuild when it returns the solution from +the compiled program. +""" +function build_nonlinear_solution( + prob, alg, u, resid; retcode, stats = nothing, trace = nothing + ) + stored_prob = ReactantCore.within_compile() ? nothing : prob + return SciMLBase.NonlinearSolution( + u, resid, stored_prob, alg, retcode, nothing, nothing, nothing, stats, trace + ) +end + +""" + select(cond, a, b) + +`ifelse(cond, a, b)` that also works when `cond` is a traced scalar and `a`, `b` are +arrays, in which case the selection is elementwise. +""" +select(cond::Bool, a, b) = ifelse(cond, a, b) +# Two methods only: a condition of unknown type must not turn this into a dynamic dispatch +# that loses the result type, which is that of `a`/`b` either way. +select(cond, a, b) = a isa AbstractArray ? ifelse.(cond, a, b) : ifelse(cond, a, b) diff --git a/lib/NonlinearSolveBase/src/solve.jl b/lib/NonlinearSolveBase/src/solve.jl index a0d5545342..4e57b39ac3 100644 --- a/lib/NonlinearSolveBase/src/solve.jl +++ b/lib/NonlinearSolveBase/src/solve.jl @@ -360,18 +360,21 @@ end function _run_cache_to_completion!( cache::AbstractNonlinearSolveCache, step_observer = nothing ) - cache.retcode == ReturnCode.InitialFailure && return cache - while not_terminated(cache) + # A traced return code cannot be branched on; initialization failures are host events. + !ReactantCore.within_compile() && cache.retcode == ReturnCode.InitialFailure && return cache + dealias_traced!(cache) + ReactantCore.@trace track_numbers = false while not_terminated(cache) CommonSolve.step!(cache) + dealias_traced!(cache) _observe_nonlinear_step!(step_observer, cache) end # The solver might have set a different `retcode` - if cache.retcode == ReturnCode.Default - cache.retcode = ifelse( - cache.nsteps ≥ cache.maxiters, ReturnCode.MaxIters, ReturnCode.Success - ) - end + cache.retcode = ifelse( + cache.retcode == ReturnCode.Default, + ifelse(cache.nsteps ≥ cache.maxiters, ReturnCode.MaxIters, ReturnCode.Success), + cache.retcode + ) # A driver may have stepped with `evaluate_residual = false`; the residual has to be # brought forward before it is reported, since nothing downstream re-evaluates it. @@ -402,9 +405,12 @@ end end function _solution_from_cache(cache::AbstractNonlinearSolveCache; transform_bounds::Bool) - sol = SciMLBase.build_solution( + # `NLStats` holds plain integers, so under Reactant it would only count trace-time + # evaluations. + stats = ReactantCore.within_compile() ? nothing : cache.stats + sol = build_nonlinear_solution( cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, cache.stats, cache.trace + cache.retcode, stats, cache.trace ) # Inverse bounds transform: if the problem function was wrapped with a @@ -424,7 +430,10 @@ end function _solve_without_solution!(cache::AbstractNonlinearSolveCache) if applicable(InternalAPI.step!, cache) && hasfield(typeof(cache), :termination_cache) && hasfield(typeof(cache), :trace) - cache.retcode == ReturnCode.InitialFailure && return cache + # A traced return code cannot be branched on; initialization failures are host + # events. + !ReactantCore.within_compile() && cache.retcode == ReturnCode.InitialFailure && + return cache _run_cache_to_completion!(cache) return _has_bounded_wrapper(cache) ? _solution_from_cache(cache; transform_bounds = true) : cache @@ -433,7 +442,7 @@ function _solve_without_solution!(cache::AbstractNonlinearSolveCache) end function CommonSolve.solve!(cache::AbstractNonlinearSolveCache) - if cache.retcode == ReturnCode.InitialFailure + if !ReactantCore.within_compile() && cache.retcode == ReturnCode.InitialFailure return _solution_from_cache(cache; transform_bounds = false) end @@ -454,7 +463,9 @@ end @inline _solve_result_stats(cache::AbstractNonlinearSolveCache) = cache.stats @inline function _solve_result_original(cache::AbstractNonlinearSolveCache) return _solution_from_cache( - cache; transform_bounds = cache.retcode != ReturnCode.InitialFailure + cache; + transform_bounds = ReactantCore.within_compile() || + cache.retcode != ReturnCode.InitialFailure ) end @@ -481,7 +492,7 @@ end push!( calls, quote - if cache.retcode == ReturnCode.InitialFailure + if !ReactantCore.within_compile() && cache.retcode == ReturnCode.InitialFailure u = $(SII.state_values)(cache)::_uType return build_solution_less_specialize( cache.prob, cache.alg, u, @@ -622,9 +633,63 @@ function SciMLBase.__solve( prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm, args...; kwargs... ) + ReactantCore.within_compile() && return _traced_polysolve(prob, alg, args...; kwargs...) return __generated_polysolve(prob, alg, args...; kwargs...) end +# Under Reactant every algorithm of the polyalgorithm is compiled; at run time a later one +# only executes when none of the earlier ones succeeded, and the smallest residual seen is +# kept as the fallback, mirroring `__generated_polysolve`. The splat is not called `args`: +# `ReactantCore.@trace if` binds its captured variables through a parameter of that name. +function _traced_polysolve( + prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm{Val{N}}, solve_args...; + initializealg = NonlinearSolveDefaultInit(), kwargs... + ) where {N} + prob, success = run_initialization!(prob, initializealg, prob) + if !success + u = SII.state_values(prob) + return SciMLBase.build_solution( + prob, alg, u, Utils.evaluate_f(prob, u); retcode = ReturnCode.InitialFailure + ) + end + sol = SciMLBase.__solve(prob, alg.algs[alg.start_index], solve_args...; kwargs...) + u, resid, retcode = sol.u, sol.resid, sol.retcode + best_norm = _poly_resid_norm(prob, resid) + done = _poly_success(retcode) + for i in (alg.start_index + 1):N + alg_i = alg.algs[i] + # The branch assigns only the loop state: a solution object as branch output would + # have to be materialized for the untaken side as well. + ReactantCore.@trace track_numbers = false if !done + u, resid, retcode, best_norm, done = _traced_polysolve_member( + prob, alg_i, solve_args, kwargs, u, resid, retcode, best_norm + ) + end + end + return build_nonlinear_solution(prob, alg, u, resid; retcode) +end + +function _traced_polysolve_member(prob, alg, solve_args, kwargs, u, resid, retcode, best_norm) + sol = SciMLBase.__solve(prob, alg, solve_args...; kwargs...) + success = _poly_success(sol.retcode) + resid_norm = _poly_resid_norm(prob, sol.resid) + better = success | (resid_norm < best_norm) + return select(better, sol.u, u), select(better, sol.resid, resid), + ifelse(better, sol.retcode, retcode), ifelse(better, resid_norm, best_norm), success +end + +function _poly_success(retcode) + return (retcode == ReturnCode.Success) | (retcode == ReturnCode.Terminated) | + (retcode == ReturnCode.FloatingPointLimit) +end + +function _poly_resid_norm(prob::AbstractNonlinearProblem, resid) + # Reductions rather than `norm`, which iterates and cannot be traced. + fx = prob isa NonlinearLeastSquaresProblem ? sqrt(sum(abs2, resid)) : + maximum(abs, resid) + return ifelse(isnan(fx), oftype(fx, Inf), fx) +end + function SciMLBase.__solve( prob::AbstractNonlinearProblem, args...; default_set = false, second_time = false, kwargs... @@ -834,7 +899,7 @@ NonlinearSolve.step!(cache) ``` """ function CommonSolve.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) - not_terminated(cache) || return + ReactantCore.within_compile() || not_terminated(cache) || return has_time_limit(cache) && (time_start = time()) @@ -960,7 +1025,7 @@ function SciMLBase.__init( end function CommonSolve.solve!(cache::NonlinearSolveNoInitCache) - if cache.retcode == ReturnCode.InitialFailure + if !ReactantCore.within_compile() && cache.retcode == ReturnCode.InitialFailure u = SII.state_values(cache) return SciMLBase.build_solution( cache.prob, cache.alg, u, Utils.evaluate_f(cache.prob, u); cache.retcode @@ -1035,6 +1100,7 @@ function _solve_forward( end function maybe_wrap_f(prob::AbstractNonlinearProblem) + ReactantCore.within_compile() && return prob # AutoDePSpecialize opaque-`p` path (packs `p` + wraps `f` together). opaque = maybe_opaque_wrap(prob) opaque === nothing || return opaque diff --git a/lib/NonlinearSolveBase/src/termination_conditions.jl b/lib/NonlinearSolveBase/src/termination_conditions.jl index dee966e59f..68132b9715 100644 --- a/lib/NonlinearSolveBase/src/termination_conditions.jl +++ b/lib/NonlinearSolveBase/src/termination_conditions.jl @@ -47,7 +47,7 @@ residual_only_termination_mode(::AbsNormTerminationMode) = true # Core Implementation @concrete mutable struct NonlinearTerminationModeCache{uType, T} u::uType - retcode::ReturnCode.T + retcode abstol::T reltol::T best_objective_value::T @@ -172,7 +172,7 @@ function CommonSolve.init( leastsq = typeof(prob) <: NonlinearLeastSquaresProblem return NonlinearTerminationModeCache( - u_unaliased, ReturnCode.Default, abstol, reltol, best_value, mode, + u_unaliased, maybe_traced(ReturnCode.Default), abstol, reltol, best_value, mode, initial_objective, objectives_trace, 0, saved_value_prototype, u0_norm, step_norm_trace, max_stalled_steps, u_diff_cache, leastsq ) @@ -193,7 +193,7 @@ function SciMLBase.reinit!( cache.u .= u end end - cache.retcode = ReturnCode.Default + cache.retcode = maybe_traced(ReturnCode.Default) cache.abstol = get_tolerance(u, abstol, T) cache.reltol = get_tolerance(u, reltol, T) @@ -233,11 +233,9 @@ end function (cache::NonlinearTerminationModeCache)( mode::AbstractNonlinearTerminationMode, du, u, uprev, abstol, reltol, args... ) - if check_convergence(mode, du, u, uprev, abstol, reltol) - cache.retcode = ReturnCode.Success - return true - end - return false + converged = check_convergence(mode, du, u, uprev, abstol, reltol) + cache.retcode = ifelse(converged, ReturnCode.Success, cache.retcode) + return converged end function (cache::NonlinearTerminationModeCache)( @@ -385,10 +383,13 @@ end function default_termination_mode( ::Union{ImmutableNonlinearProblem, NonlinearProblem}, ::Val{:regular} ) + ReactantCore.within_compile() && + return AbsNormTerminationMode(Base.Fix1(maximum, abs)) return AbsNormSafeBestTerminationMode(Base.Fix1(maximum, abs); max_stalled_steps = 32) end function default_termination_mode(::NonlinearLeastSquaresProblem, ::Val{:regular}) + ReactantCore.within_compile() && return AbsNormTerminationMode(Base.Fix2(norm, 2)) return AbsNormSafeBestTerminationMode(Base.Fix2(norm, 2); max_stalled_steps = 32) end @@ -418,11 +419,14 @@ function check_and_update!(cache, fu, u, uprev) end function check_and_update!(tc_cache, cache, fu, u, uprev, mode) - return if tc_cache(fu, u, uprev) - cache.retcode = tc_cache.retcode + converged = tc_cache(fu, u, uprev) + cache.retcode = ifelse(converged, tc_cache.retcode, cache.retcode) + cache.force_stop = converged | cache.force_stop + # Only the best-iterate modes have anything to copy back; they are not traced. + if mode isa AbstractSafeBestNonlinearTerminationMode && converged update_from_termination_cache!(tc_cache, cache, mode, u) - cache.force_stop = true end + return nothing end function update_from_termination_cache!(tc_cache, cache, u = get_u(cache)) diff --git a/lib/NonlinearSolveBase/src/tracing.jl b/lib/NonlinearSolveBase/src/tracing.jl index b0f7d0155b..167fbebbd6 100644 --- a/lib/NonlinearSolveBase/src/tracing.jl +++ b/lib/NonlinearSolveBase/src/tracing.jl @@ -236,6 +236,10 @@ function init_nonlinearsolve_trace( J = uses_jac_inverse isa Val{true} ? (trace_level.trace_mode isa Val{:minimal} ? J : LinearAlgebra.pinv(J)) : J history = init_trace_history(prob, show_trace, trace_level, store_trace, u, fu, J, δu) + # Under Reactant the problem's arrays are traced and must not be carried by the trace: + # Reactant does not register them as loop-carried values, so `prob` would only be used + # for its type, and `nothing` selects the same code paths. + prob = ReactantCore.within_compile() ? nothing : prob return NonlinearSolveTrace(show_trace, store_trace, history, trace_level, prob) end diff --git a/lib/NonlinearSolveBase/test/qa/qa.jl b/lib/NonlinearSolveBase/test/qa/qa.jl index 082307f6b7..ee90870c12 100644 --- a/lib/NonlinearSolveBase/test/qa/qa.jl +++ b/lib/NonlinearSolveBase/test/qa/qa.jl @@ -43,8 +43,12 @@ run_qa( # nlls_generate_vjp_function, nodual_value, nonlinearsolve_∂f_∂p, # nonlinearsolve_∂f_∂u, reinit!, restructure, safe_reshape, safe_similar, # sparse_or_structured_prototype, structural_sparse + # ReactantCore: `is_traced` (the predicate behind `@trace`; not declared public + # upstream, and it cannot be reproduced here because Reactant defines its + # methods for the traced types) all_qualified_accesses_are_public = (; ignore = ( + :is_traced, :ChainRulesOriginator, :DAEInitializationAlgorithm, :EnzymeOriginator, :NonNumberEltypeError, :OverrideInitData, :Void, :get_root_indp, :has_colorvec, :isdualtype, diff --git a/lib/NonlinearSolveFirstOrder/Project.toml b/lib/NonlinearSolveFirstOrder/Project.toml index 463ee6299c..3fe87ed64f 100644 --- a/lib/NonlinearSolveFirstOrder/Project.toml +++ b/lib/NonlinearSolveFirstOrder/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolveFirstOrder" uuid = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d" authors = ["Avik Pal and contributors"] -version = "2.5.0" +version = "2.5.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -16,6 +16,7 @@ LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" @@ -45,8 +46,9 @@ LinearAlgebra = "1.10" LinearSolve = "5.4" MaybeInplace = "0.1.4" NonlinearProblemLibrary = "0.1.2" -NonlinearSolveBase = "2.45" +NonlinearSolveBase = "2.50" PrecompileTools = "1.2" +ReactantCore = "0.1.21" Random = "1.10" Reexport = "1.2.2" SciMLBase = "3.37" diff --git a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl index a9af5eef6a..4ef1c27460 100644 --- a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl +++ b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl @@ -21,12 +21,13 @@ module NonlinearSolveFirstOrder using ConcreteStructs: @concrete using PrecompileTools: @compile_workload, @setup_workload +using ReactantCore: ReactantCore using Reexport: @reexport using Setfield: @set! using ADTypes: ADTypes using ArrayInterface: ArrayInterface -using LinearAlgebra: LinearAlgebra, Diagonal, dot, diagind +using LinearAlgebra: LinearAlgebra, Diagonal, diag, dot, diagind using LineSearch: BackTracking using StaticArraysCore: SArray diff --git a/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl index 2ec8e10dc9..088c21ee97 100644 --- a/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl +++ b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl @@ -84,7 +84,7 @@ function InternalAPI.init( else @bb J_diag_cache = similar(u) end - J_damped = T(initial_damping) .* DᵀD + J_damped = scale_levenberg_marquardt_diagonal(T(initial_damping), DᵀD) return LevenbergMarquardtDampingCache( T(f.increase_factor), T(f.decrease_factor), T(f.min_damping), T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache, J_damped, f, @@ -115,7 +115,7 @@ function InternalAPI.reinit!(cache::LevenbergMarquardtDampingCache, args...; kwa cache.DᵀD = Diagonal(ones(typeof(cache.DᵀD.diag)) * cache.min_damping) end end - cache.J_damped = cache.λ .* cache.DᵀD + cache.J_damped = scale_levenberg_marquardt_diagonal(cache.λ, cache.DᵀD) return end @@ -156,7 +156,7 @@ function InternalAPI.solve!( cache.DᵀD = update_levenberg_marquardt_diagonal!!( cache.DᵀD, Utils.safe_vec(cache.J_diag_cache) ) - @bb @. cache.J_damped = cache.λ * cache.DᵀD + cache.J_damped = scale_levenberg_marquardt_diagonal(cache.λ, cache.DᵀD) return cache.J_damped end @@ -164,17 +164,16 @@ function InternalAPI.solve!( cache::LevenbergMarquardtDampingCache, JᵀJ, fu, ::Val{true}; kwargs... ) cache.DᵀD = update_levenberg_marquardt_diagonal!!(cache.DᵀD, JᵀJ) - @bb @. cache.J_damped = cache.λ * cache.DᵀD + cache.J_damped = scale_levenberg_marquardt_diagonal(cache.λ, cache.DᵀD) return cache.J_damped end function NonlinearSolveBase.callback_into_cache!( topcache, cache::LevenbergMarquardtDampingCache, args... ) - if NonlinearSolveBase.last_step_accepted(topcache.trustregion_cache) && - NonlinearSolveBase.last_step_accepted(topcache.descent_cache) - cache.λ_factor = 1 / cache.decrease_factor - end + accepted = NonlinearSolveBase.last_step_accepted(topcache.trustregion_cache) & + NonlinearSolveBase.last_step_accepted(topcache.descent_cache) + cache.λ_factor = ifelse(accepted, 1 / cache.decrease_factor, cache.λ_factor) cache.λ *= cache.λ_factor return cache.λ_factor = cache.increase_factor end @@ -217,8 +216,9 @@ function InternalAPI.init( @bb v = copy(u) @bb u_cache = similar(u) @bb fu_cache = similar(fu) + last_step_accepted = NonlinearSolveBase.maybe_traced(false) return LevenbergMarquardtTrustRegionCache( - f, p, T(Inf), v, T(Inf), internalnorm, T(alg.β_uphill), false, + f, p, T(Inf), v, T(Inf), internalnorm, T(alg.β_uphill), last_step_accepted, u_cache, fu_cache, stats ) end @@ -232,7 +232,7 @@ end norm_v_old internalnorm β_uphill - last_step_accepted::Bool + last_step_accepted u_cache fu_cache stats::NLStats @@ -262,13 +262,10 @@ function InternalAPI.solve!( loss = cache.internalnorm(cache.fu_cache) - if (1 - β)^cache.β_uphill * loss ≤ cache.loss_old # Accept Step - cache.last_step_accepted = true - cache.norm_v_old = norm_v - @bb copyto!(cache.v_cache, v) - else - cache.last_step_accepted = false - end + accepted = (1 - β)^cache.β_uphill * loss ≤ cache.loss_old + cache.last_step_accepted = accepted + cache.norm_v_old = ifelse(accepted, norm_v, cache.norm_v_old) + @bb @. cache.v_cache = ifelse(accepted, v, cache.v_cache) return cache.last_step_accepted, cache.u_cache, cache.fu_cache end @@ -296,7 +293,7 @@ function update_levenberg_marquardt_diagonal!!(y::Diagonal, x::AbstractVecOrMat) return y end ndims(x) == 1 && return Diagonal(max.(y.diag, x)) - return Diagonal(max.(y.diag, @view(x[diagind(x)]))) + return Diagonal(max.(y.diag, diag(x))) end init_levenberg_marquardt_diagonal(u::Number, v) = oftype(u, v) @@ -306,3 +303,8 @@ function init_levenberg_marquardt_diagonal(u, v) d .= v return Diagonal(d) end + +scale_levenberg_marquardt_diagonal(λ, DᵀD::Number) = λ * DᵀD +function scale_levenberg_marquardt_diagonal(λ, DᵀD::Diagonal) + return Diagonal(λ .* DᵀD.diag) +end diff --git a/lib/NonlinearSolveFirstOrder/src/solve.jl b/lib/NonlinearSolveFirstOrder/src/solve.jl index 988e6c0204..a3193f7ff6 100644 --- a/lib/NonlinearSolveFirstOrder/src/solve.jl +++ b/lib/NonlinearSolveFirstOrder/src/solve.jl @@ -76,7 +76,7 @@ end # Counters stats::NLStats - nsteps::Int + nsteps maxiters::Int maxtime max_shrink_times::Int @@ -94,8 +94,8 @@ end # Termination & Tracking termination_cache trace - retcode::ReturnCode.T - force_stop::Bool + retcode + force_stop kwargs initializealg @@ -121,12 +121,12 @@ function InternalAPI.reinit_self!( Utils.reinit_common!(cache, u0, p, alias_u0) InternalAPI.reinit!(cache.stats) - cache.nsteps = 0 + cache.nsteps = NonlinearSolveBase.maybe_traced(0) cache.maxiters = maxiters cache.maxtime = maxtime cache.total_time = 0.0 - cache.force_stop = false - cache.retcode = ReturnCode.Default + cache.force_stop = NonlinearSolveBase.maybe_traced(false) + cache.retcode = NonlinearSolveBase.maybe_traced(ReturnCode.Default) cache.make_new_jacobian = true cache.fu_deferred = false reset_jacobian_reuse!(cache.jacobian_reuse_cache, cache.fu) @@ -291,9 +291,12 @@ function SciMLBase.__init( ) end - jacobian_reuse_cache = init_jacobian_reuse_cache( - resolve_jacobian_reuse(alg.jacobian_reuse, u), J, fu, internalnorm - ) + # A reuse decision compares residual norms, which are traced under Reactant; there + # the Jacobian is recomputed on every step instead. + jacobian_reuse = resolve_jacobian_reuse(alg.jacobian_reuse, u) + ReactantCore.within_compile() && + (jacobian_reuse = without_jacobian_reuse(jacobian_reuse)) + jacobian_reuse_cache = init_jacobian_reuse_cache(jacobian_reuse, J, fu, internalnorm) trace = NonlinearSolveBase.init_nonlinearsolve_trace( prob, alg, u, fu, J, du; kwargs... @@ -303,8 +306,10 @@ function SciMLBase.__init( fu, u, u_cache, prob.p, alg, prob, globalization, jac_cache, descent_cache, forcing_cache, jacobian_reuse_cache, linesearch_cache, trustregion_cache, - stats, 0, maxiters, maxtime, alg.max_shrink_times, timer, - 0.0, true, false, termination_cache, trace, ReturnCode.Default, false, kwargs, + stats, NonlinearSolveBase.maybe_traced(0), maxiters, maxtime, + alg.max_shrink_times, timer, 0.0, true, false, termination_cache, trace, + NonlinearSolveBase.maybe_traced(ReturnCode.Default), + NonlinearSolveBase.maybe_traced(false), kwargs, initializealg, verbose ) NonlinearSolveBase.run_initialization!(cache) @@ -406,96 +411,110 @@ function InternalAPI.step!( δu, descent_intermediates = descent_result.δu, descent_result.extras - if descent_result.success - if has_forcing - post_step_forcing!(cache.forcing_cache, J, cache.u, cache.fu, δu, cache.nsteps) - end + α = NonlinearSolveBase.maybe_traced(false) + ReactantCore.@trace track_numbers = false if descent_result.success + α = _perform_first_order_step!( + cache, J, δu, descent_intermediates, has_forcing, defer_residual, policy_driven + ) + else + # Under Reactant every step is traced, so the Jacobian is always recomputed. + cache.make_new_jacobian = ReactantCore.within_compile() || + jacobian_is_stale(cache.jacobian_reuse_cache) + end + # The line search asked for a retry on a fresh Jacobian; that step already finished. + α === nothing && return - accepted_step = false - if cache.globalization isa Val{:LineSearch} - @static_timeit cache.timer "linesearch" begin - linesearch_sol = CommonSolve.solve!(cache.linesearch_cache, cache.u, δu) - linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) - α = linesearch_sol.step_size - end - if linesearch_failed && policy_driven && - jacobian_is_stale(cache.jacobian_reuse_cache) - @SciMLMessage("Line Search Failed with stale Jacobian information. Retrying with updated Jacobian.", cache.verbose, :linsolve_failed_noncurrent) - cache.make_new_jacobian = true - InternalAPI.step!(cache; recompute_jacobian = true) - return - elseif linesearch_failed - cache.retcode = ReturnCode.InternalLineSearchFailed - cache.force_stop = true - end - @static_timeit cache.timer "step" begin - @bb axpy!(α, δu, cache.u) + update_trace!(cache, α) + @bb copyto!(cache.u_cache, cache.u) + + NonlinearSolveBase.callback_into_cache!(cache) + + return nothing +end + +function _perform_first_order_step!( + cache, J, δu, descent_intermediates, has_forcing, defer_residual, policy_driven + ) + if has_forcing + post_step_forcing!(cache.forcing_cache, J, cache.u, cache.fu, δu, cache.nsteps) + end + + # Under Reactant acceptance is traced and cannot be branched on, so every step counts + # as accepted for the Jacobian policy (which is switched off there anyway). + accepted_step = false + if cache.globalization isa Val{:LineSearch} + @static_timeit cache.timer "linesearch" begin + linesearch_sol = CommonSolve.solve!(cache.linesearch_cache, cache.u, δu) + linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) + α = linesearch_sol.step_size + end + if linesearch_failed && policy_driven && + jacobian_is_stale(cache.jacobian_reuse_cache) + @SciMLMessage("Line Search Failed with stale Jacobian information. Retrying with updated Jacobian.", cache.verbose, :linsolve_failed_stale_jac) + cache.make_new_jacobian = true + InternalAPI.step!(cache; recompute_jacobian = true) + return nothing + elseif linesearch_failed + cache.retcode = ReturnCode.InternalLineSearchFailed + cache.force_stop = true + end + @static_timeit cache.timer "step" begin + @bb axpy!(α, δu, cache.u) + cache.u = NonlinearSolveBase.apply_postcondition!!( + cache.u, cache.u_cache, cache + ) + Utils.evaluate_f!(cache, cache.u, cache.p) + end + accepted_step = !linesearch_failed + elseif cache.globalization isa Val{:TrustRegion} + @static_timeit cache.timer "trustregion" begin + tr_accepted, u_new, + fu_new = InternalAPI.solve!( + cache.trustregion_cache, J, cache.fu, cache.u, δu, descent_intermediates + ) + @bb @. cache.u = ifelse(tr_accepted, u_new, cache.u) + if NonlinearSolveBase.get_postcondition(cache) === nothing + @bb @. cache.fu = ifelse(tr_accepted, fu_new, cache.fu) + elseif ReactantCore.within_compile() || tr_accepted cache.u = NonlinearSolveBase.apply_postcondition!!( cache.u, cache.u_cache, cache ) Utils.evaluate_f!(cache, cache.u, cache.p) end - accepted_step = !linesearch_failed - elseif cache.globalization isa Val{:TrustRegion} - @static_timeit cache.timer "trustregion" begin - tr_accepted, u_new, - fu_new = InternalAPI.solve!( - cache.trustregion_cache, J, cache.fu, cache.u, δu, descent_intermediates + α = tr_accepted + accepted_step = ReactantCore.within_compile() || tr_accepted + accepted_step || + (cache.make_new_jacobian = jacobian_is_stale(cache.jacobian_reuse_cache)) + if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && + cache.max_shrink_times < typemax(Int) + exceeded = cache.trustregion_cache.shrink_counter > cache.max_shrink_times + cache.retcode = ifelse( + exceeded, ReturnCode.ShrinkThresholdExceeded, cache.retcode ) - if tr_accepted - @bb copyto!(cache.u, u_new) - if NonlinearSolveBase.get_postcondition(cache) === nothing - @bb copyto!(cache.fu, fu_new) - else - cache.u = NonlinearSolveBase.apply_postcondition!!( - cache.u, cache.u_cache, cache - ) - Utils.evaluate_f!(cache, cache.u, cache.p) - end - α = true - accepted_step = true - else - α = false - cache.make_new_jacobian = - jacobian_is_stale(cache.jacobian_reuse_cache) - end - if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && - cache.trustregion_cache.shrink_counter > cache.max_shrink_times - cache.retcode = ReturnCode.ShrinkThresholdExceeded - cache.force_stop = true - end + cache.force_stop = exceeded | cache.force_stop end - elseif cache.globalization isa Val{:None} - @static_timeit cache.timer "step" begin - @bb axpy!(1, δu, cache.u) - cache.u = NonlinearSolveBase.apply_postcondition!!( - cache.u, cache.u_cache, cache - ) - defer_residual || Utils.evaluate_f!(cache, cache.u, cache.p) - end - α = true - accepted_step = true - else - error("Unknown Globalization Strategy: $(cache.globalization). Allowed values \ - are (:LineSearch, :TrustRegion, :None)") end - if defer_residual - cache.fu_deferred = true - else - accepted_step && schedule_next_jacobian!(cache) - NonlinearSolveBase.check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + elseif cache.globalization isa Val{:None} + @static_timeit cache.timer "step" begin + @bb axpy!(1, δu, cache.u) + cache.u = NonlinearSolveBase.apply_postcondition!!( + cache.u, cache.u_cache, cache + ) + defer_residual || Utils.evaluate_f!(cache, cache.u, cache.p) end + α = true + accepted_step = true else - α = false - cache.make_new_jacobian = jacobian_is_stale(cache.jacobian_reuse_cache) + error("Unknown Globalization Strategy: $(cache.globalization). Allowed values \ + are (:LineSearch, :TrustRegion, :None)") end - - update_trace!(cache, α) - @bb copyto!(cache.u_cache, cache.u) - - NonlinearSolveBase.callback_into_cache!(cache) - - return nothing + if defer_residual + cache.fu_deferred = true + else + accepted_step && schedule_next_jacobian!(cache) + NonlinearSolveBase.check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + end + return α end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) diff --git a/lib/NonlinearSolveFirstOrder/src/trust_region.jl b/lib/NonlinearSolveFirstOrder/src/trust_region.jl index 510b85df2b..8da1eca2bd 100644 --- a/lib/NonlinearSolveFirstOrder/src/trust_region.jl +++ b/lib/NonlinearSolveFirstOrder/src/trust_region.jl @@ -255,11 +255,14 @@ function InternalAPI.init( @bb u_cache = similar(u) @bb fu_cache = similar(fu) @bb Jδu_cache = similar(fu) + last_step_accepted = NonlinearSolveBase.maybe_traced(false) + shrink_counter = NonlinearSolveBase.maybe_traced(0) return GenericTrustRegionSchemeCache( alg.method, f, p, mtr, itr, itr, stt, sht, et, shf, ef, p1, p2, p3, p4, ϵ, T(0), vjp_operator, jvp_operator, Jᵀfu_cache, Jδu_cache, - δu_cache, internalnorm, u_cache, fu_cache, false, 0, stats, alg + δu_cache, internalnorm, u_cache, fu_cache, last_step_accepted, shrink_counter, + stats, alg ) end @@ -289,8 +292,8 @@ end internalnorm u_cache fu_cache - last_step_accepted::Bool - shrink_counter::Int + last_step_accepted + shrink_counter stats::NLStats alg end @@ -319,9 +322,9 @@ function InternalAPI.reinit!( cache.initial_trust_radius = T(cache.p1 * cache.internalnorm(cache.Jᵀfu_cache)) end end - cache.last_step_accepted = false + cache.last_step_accepted = zero(cache.last_step_accepted) cache.trust_region = cache.initial_trust_radius - return cache.shrink_counter = 0 + return cache.shrink_counter = zero(cache.shrink_counter) end # Defaults @@ -409,8 +412,18 @@ function InternalAPI.solve!( cache.fu_cache = Utils.evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) cache.stats.nf += 1 - if hasfield(typeof(descent_stats), :δuJᵀJδu) && !isnan(descent_stats.δuJᵀJδu) - δuJᵀJδu = descent_stats.δuJᵀJδu + if hasfield(typeof(descent_stats), :δuJᵀJδu) + # `isnan` of a traced value cannot pick a branch, so under Reactant the product is + # always formed and selected from. + if ReactantCore.within_compile() || isnan(descent_stats.δuJᵀJδu) + @bb cache.Jδu_cache = J × vec(δu) + computed_δuJᵀJδu = Utils.safe_dot(cache.Jδu_cache, cache.Jδu_cache) + δuJᵀJδu = ifelse( + isnan(descent_stats.δuJᵀJδu), computed_δuJᵀJδu, descent_stats.δuJᵀJδu + ) + else + δuJᵀJδu = descent_stats.δuJᵀJδu + end else @bb cache.Jδu_cache = J × vec(δu) δuJᵀJδu = Utils.safe_dot(cache.Jδu_cache, cache.Jδu_cache) @@ -420,83 +433,88 @@ function InternalAPI.solve!( denom = Utils.safe_dot(δu, cache.Jᵀfu_cache) + δuJᵀJδu / 2 cache.ρ = num / denom - if cache.ρ > cache.step_threshold - cache.last_step_accepted = true - else - cache.last_step_accepted = false - end + cache.last_step_accepted = cache.ρ > cache.step_threshold if cache.method isa RUS.__Simple - if cache.ρ < cache.shrink_threshold - cache.trust_region *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ > cache.expand_threshold && cache.ρ > cache.step_threshold - cache.trust_region = cache.expand_factor * cache.trust_region - end - end + shrink = cache.ρ < cache.shrink_threshold + expand = (cache.ρ > cache.expand_threshold) & (cache.ρ > cache.step_threshold) + cache.trust_region = ifelse( + shrink, cache.shrink_factor * cache.trust_region, + ifelse(expand, cache.expand_factor * cache.trust_region, cache.trust_region) + ) + cache.shrink_counter = ifelse( + shrink, cache.shrink_counter + one(cache.shrink_counter), + zero(cache.shrink_counter) + ) elseif cache.method isa RUS.__NLsolve - if cache.ρ < cache.shrink_threshold - cache.trust_region *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ ≥ cache.expand_threshold - cache.trust_region = cache.expand_factor * cache.internalnorm(δu) - elseif cache.ρ ≥ cache.p1 - cache.trust_region = max( - cache.trust_region, cache.expand_factor * cache.internalnorm(δu) + shrink = cache.ρ < cache.shrink_threshold + δu_norm = cache.internalnorm(δu) + cache.trust_region = ifelse( + shrink, cache.shrink_factor * cache.trust_region, + ifelse( + cache.ρ ≥ cache.expand_threshold, cache.expand_factor * δu_norm, + ifelse( + cache.ρ ≥ cache.p1, max(cache.trust_region, cache.expand_factor * δu_norm), + cache.trust_region ) - end - end + ) + ) + cache.shrink_counter = ifelse( + shrink, cache.shrink_counter + one(cache.shrink_counter), + zero(cache.shrink_counter) + ) elseif cache.method isa RUS.__NocedalWright - if cache.ρ < cache.shrink_threshold - cache.trust_region = cache.shrink_factor * cache.internalnorm(δu) - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ > cache.expand_threshold && - abs(cache.internalnorm(δu) - cache.trust_region) < 1.0e-6 * cache.trust_region - cache.trust_region = cache.expand_factor * cache.trust_region - end - end + shrink = cache.ρ < cache.shrink_threshold + δu_norm = cache.internalnorm(δu) + expand = (cache.ρ > cache.expand_threshold) & + (abs(δu_norm - cache.trust_region) < 1.0e-6 * cache.trust_region) + cache.trust_region = ifelse( + shrink, cache.shrink_factor * δu_norm, + ifelse(expand, cache.expand_factor * cache.trust_region, cache.trust_region) + ) + cache.shrink_counter = ifelse( + shrink, cache.shrink_counter + one(cache.shrink_counter), + zero(cache.shrink_counter) + ) elseif cache.method isa RUS.__Hei tr_new = rfunc_adaptive_trust_region( cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, cache.p2 ) * cache.internalnorm(δu) - if tr_new < cache.trust_region - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - end + cache.shrink_counter = ifelse( + tr_new < cache.trust_region, cache.shrink_counter + one(cache.shrink_counter), + zero(cache.shrink_counter) + ) cache.trust_region = tr_new elseif cache.method isa RUS.__Yuan - if cache.ρ < cache.shrink_threshold - cache.p1 = cache.p2 * cache.p1 - cache.shrink_counter += 1 - else - if cache.ρ ≥ cache.expand_threshold && - 2 * cache.internalnorm(δu) > cache.trust_region - cache.p1 = cache.p3 * cache.p1 - end - cache.shrink_counter = 0 - end + shrink = cache.ρ < cache.shrink_threshold + expand = (cache.ρ ≥ cache.expand_threshold) & + (2 * cache.internalnorm(δu) > cache.trust_region) + cache.p1 = ifelse( + shrink, cache.p2 * cache.p1, ifelse(expand, cache.p3 * cache.p1, cache.p1) + ) + cache.shrink_counter = ifelse( + shrink, cache.shrink_counter + one(cache.shrink_counter), + zero(cache.shrink_counter) + ) operator = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) @bb cache.Jᵀfu_cache = operator × vec(cache.fu_cache) cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) elseif cache.method isa RUS.__Fan - if cache.ρ < cache.shrink_threshold - cache.p1 *= cache.p2 - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - cache.ρ > cache.expand_threshold && - (cache.p1 = min(cache.p1 * cache.p3, cache.p4)) - end + shrink = cache.ρ < cache.shrink_threshold + cache.p1 = ifelse( + shrink, cache.p1 * cache.p2, + ifelse(cache.ρ > cache.expand_threshold, min(cache.p1 * cache.p3, cache.p4), cache.p1) + ) + cache.shrink_counter = ifelse( + shrink, cache.shrink_counter + one(cache.shrink_counter), + zero(cache.shrink_counter) + ) cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu_cache)^T(0.99)) elseif cache.method isa RUS.__Bastin - if cache.ρ > cache.step_threshold + accepted = cache.ρ > cache.step_threshold + # The operator products are only needed for an accepted step; under Reactant they + # are always formed and the result selected. + if ReactantCore.within_compile() || accepted jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, cache.p) vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) @bb cache.Jδu_cache = jvp_op × vec(cache.δu_cache) @@ -506,14 +524,15 @@ function InternalAPI.solve!( denom_2 = dot(Utils.safe_vec(cache.Jᵀfu_cache), cache.Jᵀfu_cache) denom = denom_1 + denom_2 / 2 ρ = num / denom - if ρ ≥ cache.expand_threshold - cache.trust_region = cache.p1 * cache.internalnorm(cache.δu_cache) - end - cache.shrink_counter = 0 - else - cache.trust_region *= cache.p2 - cache.shrink_counter += 1 + expand = accepted & (ρ ≥ cache.expand_threshold) + cache.trust_region = ifelse( + expand, cache.p1 * cache.internalnorm(cache.δu_cache), cache.trust_region + ) end + cache.trust_region = ifelse(accepted, cache.trust_region, cache.trust_region * cache.p2) + cache.shrink_counter = ifelse( + accepted, zero(cache.shrink_counter), cache.shrink_counter + one(cache.shrink_counter) + ) end cache.trust_region = min(cache.trust_region, cache.max_trust_radius) diff --git a/lib/NonlinearSolveQuasiNewton/Project.toml b/lib/NonlinearSolveQuasiNewton/Project.toml index 9aca763782..8d000b5ad4 100644 --- a/lib/NonlinearSolveQuasiNewton/Project.toml +++ b/lib/NonlinearSolveQuasiNewton/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolveQuasiNewton" uuid = "9a2c21bd-3a47-402d-9113-8faf9a0ee114" authors = ["Avik Pal and contributors"] -version = "1.15.3" +version = "1.15.4" [deps] ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" @@ -43,7 +43,7 @@ LinearAlgebra = "1.10" LinearSolve = "5.4" MaybeInplace = "0.1.4" NonlinearProblemLibrary = "0.1.2" -NonlinearSolveBase = "2.43" +NonlinearSolveBase = "2.50" PrecompileTools = "1.2" Reexport = "1.2.2" SciMLBase = "3.37" diff --git a/lib/NonlinearSolveQuasiNewton/src/solve.jl b/lib/NonlinearSolveQuasiNewton/src/solve.jl index 4cc90989d2..ff54c249a9 100644 --- a/lib/NonlinearSolveQuasiNewton/src/solve.jl +++ b/lib/NonlinearSolveQuasiNewton/src/solve.jl @@ -75,7 +75,7 @@ end # Counters stats::NLStats - nsteps::Int + nsteps nresets::Int max_resets::Int maxiters::Int @@ -90,8 +90,8 @@ end # Termination & Tracking termination_cache trace - retcode::ReturnCode.T - force_stop::Bool + retcode + force_stop force_reinit::Bool kwargs @@ -128,15 +128,15 @@ function InternalAPI.reinit_self!( ) InternalAPI.reinit!(cache.stats) - cache.nsteps = 0 + cache.nsteps = NonlinearSolveBase.maybe_traced(0) cache.nresets = 0 cache.steps_since_last_reset = 0 cache.maxiters = maxiters cache.maxtime = maxtime cache.total_time = 0.0 - cache.force_stop = false + cache.force_stop = NonlinearSolveBase.maybe_traced(false) cache.force_reinit = false - cache.retcode = ReturnCode.Default + cache.retcode = NonlinearSolveBase.maybe_traced(ReturnCode.Default) NonlinearSolveBase.reset!(cache.trace) SciMLBase.reinit!( @@ -283,9 +283,10 @@ function SciMLBase.__init( fu, u, u_cache, prob.p, J, alg, prob, globalization, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache, - linsolve_workspace, stats, 0, 0, alg.max_resets, maxiters, maxtime, - alg.max_shrink_times, 0, timer, 0.0, termination_cache, trace, - ReturnCode.Default, false, false, kwargs, initializealg, verbose + linsolve_workspace, stats, NonlinearSolveBase.maybe_traced(0), 0, + alg.max_resets, maxiters, maxtime, alg.max_shrink_times, 0, timer, 0.0, + termination_cache, trace, NonlinearSolveBase.maybe_traced(ReturnCode.Default), + NonlinearSolveBase.maybe_traced(false), false, kwargs, initializealg, verbose ) NonlinearSolveBase.run_initialization!(cache) end diff --git a/lib/NonlinearSolveSpectralMethods/Project.toml b/lib/NonlinearSolveSpectralMethods/Project.toml index c90316c3a6..1700984a45 100644 --- a/lib/NonlinearSolveSpectralMethods/Project.toml +++ b/lib/NonlinearSolveSpectralMethods/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolveSpectralMethods" uuid = "26075421-4e9a-44e1-8bd1-420ed7ad02b2" authors = ["Avik Pal and contributors"] -version = "1.8.1" +version = "1.8.2" [deps] CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" @@ -32,7 +32,7 @@ InteractiveUtils = "<0.0.1, 1" LineSearch = "0.1.4" MaybeInplace = "0.1.4" NonlinearProblemLibrary = "0.1.2" -NonlinearSolveBase = "2.41" +NonlinearSolveBase = "2.50" PrecompileTools = "1.2" Reexport = "1.2.2" SciMLBase = "3.37" diff --git a/lib/NonlinearSolveSpectralMethods/src/solve.jl b/lib/NonlinearSolveSpectralMethods/src/solve.jl index b2b8a3eab9..28b2c2084f 100644 --- a/lib/NonlinearSolveSpectralMethods/src/solve.jl +++ b/lib/NonlinearSolveSpectralMethods/src/solve.jl @@ -56,7 +56,7 @@ NonlinearSolveBase.supports_postcondition(::GeneralizedDFSane) = true # Counters stats::NLStats - nsteps::Int + nsteps maxiters::Int maxtime @@ -67,8 +67,8 @@ NonlinearSolveBase.supports_postcondition(::GeneralizedDFSane) = true # Termination & Tracking termination_cache trace - retcode::ReturnCode.T - force_stop::Bool + retcode + force_stop kwargs initializealg @@ -114,11 +114,11 @@ function InternalAPI.reinit_self!( ) InternalAPI.reinit!(cache.stats) - cache.nsteps = 0 + cache.nsteps = NonlinearSolveBase.maybe_traced(0) cache.maxiters = maxiters cache.maxtime = maxtime - cache.force_stop = false - cache.retcode = ReturnCode.Default + cache.force_stop = NonlinearSolveBase.maybe_traced(false) + cache.retcode = NonlinearSolveBase.maybe_traced(ReturnCode.Default) return end @@ -189,8 +189,9 @@ function SciMLBase.__init( cache = GeneralizedDFSaneCache( fu, fu_cache, u, u_cache, prob.p, du, alg, prob, σ_n, T(alg.σ_min), T(alg.σ_max), - linesearch_cache, stats, 0, maxiters, maxtime, timer, 0.0, - tc_cache, trace, ReturnCode.Default, false, kwargs, initializealg, verbose + linesearch_cache, stats, NonlinearSolveBase.maybe_traced(0), maxiters, maxtime, + timer, 0.0, tc_cache, trace, NonlinearSolveBase.maybe_traced(ReturnCode.Default), + NonlinearSolveBase.maybe_traced(false), kwargs, initializealg, verbose ) NonlinearSolveBase.run_initialization!(cache) end diff --git a/lib/SCCNonlinearSolve/Project.toml b/lib/SCCNonlinearSolve/Project.toml index e75bcb17ed..9aabf0bce8 100644 --- a/lib/SCCNonlinearSolve/Project.toml +++ b/lib/SCCNonlinearSolve/Project.toml @@ -1,6 +1,6 @@ name = "SCCNonlinearSolve" uuid = "9dfe8606-65a1-4bb3-9748-cb89d1561431" -version = "1.15.2" +version = "1.15.3" authors = ["Avik Pal and contributors"] [deps] diff --git a/lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl b/lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl index 8a3a9a14cc..b1d1a15f3b 100644 --- a/lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl +++ b/lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl @@ -83,6 +83,17 @@ probvec(prob::LinearProblem) = prob.b iteratively_build_sols(alg, sols; kwargs...) = sols +# SciMLBase ≥ 3.51 appends the return code type to `NonlinearSolution`'s parameters. +let retcode_param = fieldtype(SciMLBase.NonlinearSolution, :retcode) === SciMLBase.ReturnCode.T ? + () : (SciMLBase.ReturnCode.T,) + @eval function nonlinear_solution_type(::Type{T}, ::Type{uType}, ::Type{rType}) where {T, uType, rType} + return SciMLBase.NonlinearSolution{ + T, 1, uType, rType, NamedTuple{(:p,), Tuple{Nothing}}, + Nothing, Nothing, Nothing, Nothing, Nothing, $(retcode_param...), + } + end +end + function solve_single_scc(alg, prob, explicitfun, sols; kwargs...) SciMLBase.invoke_with_despecialized_parameters( explicitfun, (prob.p, sols) @@ -131,10 +142,7 @@ function iteratively_build_sols(alg, probs::AbstractVector, explicitfuns::Abstra uType = typeof(probvec(prob1)) T = eltype(uType) rType = uType # resid has same type as u for nonlinear problems - ST = SciMLBase.NonlinearSolution{ - T, 1, uType, rType, - NamedTuple{(:p,), Tuple{Nothing}}, Nothing, Nothing, Nothing, Nothing, Nothing, - } + ST = nonlinear_solution_type(T, uType, rType) sols = Vector{ST}(undef, length(probs)) for i in eachindex(probs) sols[i] = solve_single_scc(alg, probs[i], explicitfuns[i], view(sols, 1:(i - 1)); kwargs...)::ST diff --git a/lib/SciMLJacobianOperators/Project.toml b/lib/SciMLJacobianOperators/Project.toml index 0d27433615..2f2feb2253 100644 --- a/lib/SciMLJacobianOperators/Project.toml +++ b/lib/SciMLJacobianOperators/Project.toml @@ -1,6 +1,6 @@ name = "SciMLJacobianOperators" uuid = "19f34311-ddf3-4b8b-af20-060888a46c0e" -version = "0.1.18" +version = "0.1.19" authors = ["Avik Pal and contributors"] [deps] diff --git a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl index 0ee3dad235..2f2e6d8178 100644 --- a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl +++ b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl @@ -83,7 +83,7 @@ using DifferentiationInterface.jl and multiply by `v`. See also [`VecJacOperator`](@ref) and [`JacVecOperator`](@ref). """ -@concrete struct JacobianOperator{iip, T <: Real} <: AbstractJacobianOperator{T} +@concrete struct JacobianOperator{iip, T <: Number} <: AbstractJacobianOperator{T} mode <: AbstractMode jvp_op diff --git a/test/Reactant/Project.toml b/test/Reactant/Project.toml new file mode 100644 index 0000000000..fbf3fedf76 --- /dev/null +++ b/test/Reactant/Project.toml @@ -0,0 +1,37 @@ +[deps] +BracketingNonlinearSolve = "70df07ce-3d50-431d-a3e7-ca6ddb60ac1e" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" +NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" +NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d" +NonlinearSolveQuasiNewton = "9a2c21bd-3a47-402d-9113-8faf9a0ee114" +NonlinearSolveSpectralMethods = "26075421-4e9a-44e1-8bd1-420ed7ad02b2" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +BracketingNonlinearSolve = {path = "../../lib/BracketingNonlinearSolve"} +NonlinearSolveBase = {path = "../../lib/NonlinearSolveBase"} +NonlinearSolveFirstOrder = {path = "../../lib/NonlinearSolveFirstOrder"} +NonlinearSolveQuasiNewton = {path = "../../lib/NonlinearSolveQuasiNewton"} +NonlinearSolveSpectralMethods = {path = "../../lib/NonlinearSolveSpectralMethods"} +SimpleNonlinearSolve = {path = "../../lib/SimpleNonlinearSolve"} + +[compat] +BracketingNonlinearSolve = "1.6" +DifferentiationInterface = "0.7.3" +Enzyme = "0.13.90" +NonlinearSolve = "4.28" +NonlinearSolveBase = "2.47" +NonlinearSolveFirstOrder = "2" +NonlinearSolveQuasiNewton = "1.15.1" +NonlinearSolveSpectralMethods = "1.6" +Reactant = "0.2.284" +SciMLBase = "3.37" +SimpleNonlinearSolve = "2.14" +Test = "1.10" +julia = "1.10" diff --git a/test/Reactant/reactant_tests.jl b/test/Reactant/reactant_tests.jl new file mode 100644 index 0000000000..e7f995086d --- /dev/null +++ b/test/Reactant/reactant_tests.jl @@ -0,0 +1,166 @@ +using NonlinearSolve +using Enzyme +using Reactant +using SciMLBase +using Test + +f(u, p) = u .* u .- p +jac(u, p) = reshape(2 .* u, :, 1) .* Float32[1 0; 0 1] +nonlinear_function = NonlinearFunction(f; jac) +autodiff_nonlinear_function = NonlinearFunction(f) + +function solve_newton(u, p) + return solve(NonlinearProblem(nonlinear_function, u, p), NewtonRaphson()) +end + +function solve_trust_region(u, p) + return solve(NonlinearProblem(nonlinear_function, u, p), TrustRegion()) +end + +function solve_default(u, p) + return solve(NonlinearProblem(nonlinear_function, u, p)) +end + +function solve_gauss_newton(u, p) + return solve( + NonlinearLeastSquaresProblem(nonlinear_function, u, p), GaussNewton() + ) +end + +function solve_autodiff_newton(u, p) + return solve( + NonlinearProblem(autodiff_nonlinear_function, u, p), NewtonRaphson() + ) +end + +function solve_autodiff_trust_region(u, p) + return solve( + NonlinearProblem(autodiff_nonlinear_function, u, p), TrustRegion() + ) +end + +function solve_autodiff_default(u, p) + return solve(NonlinearProblem(autodiff_nonlinear_function, u, p)) +end + +function solve_autodiff_gauss_newton(u, p) + return solve( + NonlinearLeastSquaresProblem(autodiff_nonlinear_function, u, p), + GaussNewton() + ) +end + +u0 = Reactant.to_rarray(Float32[1, 1]) +p0 = Reactant.to_rarray(Float32[2]) +compiled_newton = Reactant.@compile solve_newton(u0, p0) +compiled_trust_region = Reactant.@compile solve_trust_region(u0, p0) +compiled_default = Reactant.@compile solve_default(u0, p0) +compiled_gauss_newton = Reactant.@compile solve_gauss_newton(u0, p0) +compiled_autodiff_newton = Reactant.@compile solve_autodiff_newton(u0, p0) +compiled_autodiff_trust_region = Reactant.@compile solve_autodiff_trust_region(u0, p0) +compiled_autodiff_default = Reactant.@compile solve_autodiff_default(u0, p0) +compiled_autodiff_gauss_newton = Reactant.@compile solve_autodiff_gauss_newton(u0, p0) + + +# A polyalgorithm's members keep `autodiff = nothing`; the backend is chosen when each +# member is solved, so the choice is only visible on a directly solved algorithm. +for (compiled, name, uses_enzyme) in ( + (compiled_newton, :NewtonRaphson, false), + (compiled_trust_region, :TrustRegion, false), + (compiled_default, nothing, false), + (compiled_gauss_newton, :GaussNewton, false), + (compiled_autodiff_newton, :NewtonRaphson, true), + (compiled_autodiff_trust_region, :TrustRegion, true), + (compiled_autodiff_default, nothing, false), + (compiled_autodiff_gauss_newton, :GaussNewton, true), + ) + sol = compiled( + Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[2]) + ) + @test sol.u isa Reactant.ConcreteRArray + @test Array(sol.u) ≈ fill(sqrt(2.0f0), 2) + @test maximum(abs, Array(sol.resid)) ≤ 1.0f-5 + @test sol.retcode == ReturnCode.Success + @test SciMLBase.successful_retcode(sol) + if name === nothing + @test sol.alg isa NonlinearSolvePolyAlgorithm + else + @test sol.alg.name === name + end + uses_enzyme && @test sol.alg.autodiff isa AutoEnzyme + @test sol.prob === nothing + @test sol.stats === nothing +end + + +function solve_newton_one_step(u, p) + return solve( + NonlinearProblem(nonlinear_function, u, p), NewtonRaphson(); maxiters = 1 + ) +end + + +sol_newton_maxiters = Reactant.@jit solve_newton_one_step( + Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[2]) +) +@test sol_newton_maxiters.retcode == ReturnCode.MaxIters +@test !SciMLBase.successful_retcode(sol_newton_maxiters) + +struct CompiledProblemSolve{P, F, A} + problem_type::P + f::F + alg::A +end + +function (s::CompiledProblemSolve)(u, p) + prob = s.problem_type(s.f, u, p) + return s.alg === nothing ? solve(prob; abstol = 1.0f-5) : + solve(prob, s.alg; abstol = 1.0f-5) +end + +# Not compiled here: the least-squares polyalgorithms (and the default least-squares solve) +# contain a member with a line search, whose initialization calls `norm(x, Inf)`, which +# Reactant's overload scalar-indexes; `RobustMultiNewton` contains trust-region schemes whose +# vector-Jacobian products need a reverse-mode pullback, which DifferentiationInterface's +# Enzyme backend does not route through Reactant. +reactant_solver_cases = ( + (:NewtonRaphson, NonlinearProblem, NewtonRaphson()), + (:TrustRegion, NonlinearProblem, TrustRegion()), + (:LevenbergMarquardt, NonlinearProblem, LevenbergMarquardt()), + ( + :LevenbergMarquardtWithoutGeodesic, + NonlinearProblem, + LevenbergMarquardt(; disable_geodesic = Val(true)), + ), + (:PseudoTransient, NonlinearProblem, PseudoTransient()), + ( + :FastShortcutNonlinearPolyalg, + NonlinearProblem, + FastShortcutNonlinearPolyalg(Float32; u0_len = 2), + ), + ( + :NonlinearSolvePolyAlgorithm, + NonlinearProblem, + NonlinearSolvePolyAlgorithm((NewtonRaphson(), TrustRegion())), + ), + (:DefaultNonlinearSolve, NonlinearProblem, nothing), + (:GaussNewton, NonlinearLeastSquaresProblem, GaussNewton()), + (:LeastSquaresTrustRegion, NonlinearLeastSquaresProblem, TrustRegion()), + ( + :LeastSquaresLevenbergMarquardt, + NonlinearLeastSquaresProblem, + LevenbergMarquardt(), + ), +) + +@testset "Analytical Jacobian: $name" for (name, problem_type, alg) in reactant_solver_cases + compiled = Reactant.compile( + CompiledProblemSolve(problem_type, nonlinear_function, alg), (u0, p0) + ) + sol = compiled( + Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[2]) + ) + @test sol.retcode == ReturnCode.Success + @test Array(sol.u) ≈ fill(sqrt(2.0f0), 2) + @test maximum(abs, Array(sol.resid)) ≤ 1.0f-5 +end diff --git a/test/runtests.jl b/test/runtests.jl index 1890ef890e..2b86cf456f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -174,6 +174,37 @@ else return @time @safetestset "Termination Conditions: Allocations" include("gpu/cuda_tests__item2.jl") end, ), + "Reactant" => (; + env = joinpath(@__DIR__, "Reactant"), + body = function () + # Unreleased fixes this group depends on: traced enums + # (EnzymeAD/Reactant.jl#3232), the `retcode` type parameter + # (SciML/SciMLBase.jl#1563) and Enzyme Jacobians on Reactant arrays in + # DifferentiationInterface. + Pkg.add( + [ + Pkg.PackageSpec(; + url = "https://github.com/ChrisRackauckas-Claude/Reactant.jl.git", + rev = "traced-enums", subdir = "lib/ReactantCore", + ), + Pkg.PackageSpec(; + url = "https://github.com/ChrisRackauckas-Claude/Reactant.jl.git", + rev = "traced-enums", + ), + Pkg.PackageSpec(; + url = "https://github.com/ChrisRackauckas-Claude/SciMLBase.jl.git", + rev = "traced-retcode", + ), + Pkg.PackageSpec(; + url = "https://github.com/ChrisRackauckas-Claude/DifferentiationInterface.jl.git", + rev = "571fc1780c49d00addbeb31b491ab6e2b3273ee0", + subdir = "DifferentiationInterface", + ), + ] + ) + return @time @safetestset "Reactant Integration" include("Reactant/reactant_tests.jl") + end, + ), ), # QA (Aqua/ExplicitImports via SciMLTesting.run_qa) lives in an isolated sub-env # under test/qa so its compat bounds don't constrain the base resolve. Excluded diff --git a/test/test_groups.toml b/test/test_groups.toml index 0c41f0d987..ecbe5ec043 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -50,3 +50,7 @@ os = ["ubuntu-latest", "macos-latest"] [CUDA] versions = ["1"] runner = ["self-hosted", "Linux", "X64", "gpu"] + +[Reactant] +versions = ["1"] +os = ["ubuntu-latest"] From 67e5133a21a4cf4abc49b5b38d99d97998fb8010 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 29 Aug 2026 15:46:40 -0400 Subject: [PATCH 2/2] SimpleNonlinearSolve: single branchless loops that trace under Reactant Write each simple solver iteration once as a `ReactantCore.@trace while` loop without early returns, so the same code runs on the host and under `Reactant.@compile`. Loop state is prepared by `Utils.init_loop_state`/ `Utils.fresh` and the solution by `Utils.simple_solution`; the Reactant-only shadow solution, copies and duplicated loops are removed. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Agent-Harness: Claude Code 2.1.251 Agent-Model: claude-fable-5 Agent-Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE --- lib/SimpleNonlinearSolve/Project.toml | 6 +- .../src/SimpleNonlinearSolve.jl | 2 + lib/SimpleNonlinearSolve/src/broyden.jl | 17 ++- lib/SimpleNonlinearSolve/src/halley.jl | 65 ++++++----- lib/SimpleNonlinearSolve/src/klement.jl | 19 ++- lib/SimpleNonlinearSolve/src/raphson.jl | 14 ++- lib/SimpleNonlinearSolve/src/trust_region.jl | 110 +++++++++--------- lib/SimpleNonlinearSolve/src/utils.jl | 50 +++++++- test/Reactant/reactant_tests.jl | 56 +++++++++ 9 files changed, 237 insertions(+), 102 deletions(-) diff --git a/lib/SimpleNonlinearSolve/Project.toml b/lib/SimpleNonlinearSolve/Project.toml index 244d2eefb4..a023079dd6 100644 --- a/lib/SimpleNonlinearSolve/Project.toml +++ b/lib/SimpleNonlinearSolve/Project.toml @@ -1,7 +1,7 @@ name = "SimpleNonlinearSolve" uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7" authors = ["SciML"] -version = "2.14.1" +version = "2.14.2" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1" @@ -57,10 +58,11 @@ LineSearch = "0.1.9" LinearAlgebra = "1.10" MaybeInplace = "0.1.4" NonlinearProblemLibrary = "0.1.2" -NonlinearSolveBase = "2.41" +NonlinearSolveBase = "2.50" PolyesterForwardDiff = "0.1.3" PrecompileTools = "1.2" Random = "1.10" +ReactantCore = "0.1.21" Reexport = "1.2.2" ReverseDiff = "1.15" SciMLBase = "3.33" diff --git a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl index 86d78d64cf..66696b5f9c 100644 --- a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl +++ b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl @@ -21,6 +21,7 @@ module SimpleNonlinearSolve using ConcreteStructs: @concrete using PrecompileTools: @compile_workload, @setup_workload +using ReactantCore: ReactantCore using Reexport: @reexport using Setfield: @set! @@ -63,6 +64,7 @@ const NLBUtils = NonlinearSolveBase.Utils is_extension_loaded(::Val) = false + include("utils.jl") include("broyden.jl") diff --git a/lib/SimpleNonlinearSolve/src/broyden.jl b/lib/SimpleNonlinearSolve/src/broyden.jl index a18e1d9246..c5456b2bb5 100644 --- a/lib/SimpleNonlinearSolve/src/broyden.jl +++ b/lib/SimpleNonlinearSolve/src/broyden.jl @@ -41,8 +41,7 @@ function SciMLBase.__solve( fx = NLBUtils.evaluate_f(prob, x) T = promote_type(eltype(fx), eltype(x)) - iszero(fx) && - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) + solved = iszero(fx) @bb xo = copy(x) @bb δx = similar(x) @@ -70,7 +69,11 @@ function SciMLBase.__solve( ls_cache = _linesearch_cache(prob, alg.linesearch, fx, x; kwargs...) - for _ in 1:maxiters + retcode, iterations, x, xo, δx, δf, fprev, J⁻¹, J⁻¹δf, xᵀJ⁻¹, δJ⁻¹n, + δJ⁻¹ = Utils.init_loop_state( + ReturnCode.Default, x, xo, δx, δf, fprev, J⁻¹, J⁻¹δf, xᵀJ⁻¹, δJ⁻¹n, δJ⁻¹ + ) + ReactantCore.@trace track_numbers = false while (!solved) & (iterations < maxiters) @bb δx = J⁻¹ × vec(fprev) @bb δx .*= -1 @@ -86,8 +89,8 @@ function SciMLBase.__solve( @bb @. δf = fx - fprev # Termination Checks - solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) - solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) + solved, retcode, fx, x = Utils.check_termination(tc_cache, fx, x, xo, prob) + fx, x = Utils.fresh(fx), Utils.fresh(x) @bb J⁻¹δf = J⁻¹ × vec(δf) d = dot(δx, J⁻¹δf) @@ -102,7 +105,9 @@ function SciMLBase.__solve( @bb copyto!(xo, x) @bb copyto!(fprev, fx) + xo, fprev = Utils.fresh(xo), Utils.fresh(fprev) + iterations += 1 end - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) + return Utils.simple_solution(prob, alg, x, fx, x, fx, retcode, solved) end diff --git a/lib/SimpleNonlinearSolve/src/halley.jl b/lib/SimpleNonlinearSolve/src/halley.jl index c7e8f1d7e8..83196a3be6 100644 --- a/lib/SimpleNonlinearSolve/src/halley.jl +++ b/lib/SimpleNonlinearSolve/src/halley.jl @@ -21,7 +21,11 @@ A low-overhead implementation of Halley's Method. end function configure_autodiff(prob, alg::SimpleHalley) - autodiff = something(alg.autodiff, AutoForwardDiff()) + autodiff = if alg.autodiff === nothing && ReactantCore.within_compile() + NonlinearSolveBase.select_jacobian_autodiff(prob, nothing) + else + something(alg.autodiff, AutoForwardDiff()) + end autodiff = SciMLBase.has_jac(prob.f) ? autodiff : NonlinearSolveBase.select_jacobian_autodiff(prob, autodiff) @set! alg.autodiff = autodiff @@ -42,8 +46,7 @@ function SciMLBase.__solve( fx = NLBUtils.evaluate_f(prob, x) T = promote_type(eltype(fx), eltype(x)) - iszero(fx) && - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) + solved = iszero(fx) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( @@ -64,39 +67,49 @@ function SciMLBase.__solve( end J = Utils.compute_jacobian!!(nothing, prob, autodiff, fx_cache, x, jac_cache) - for _ in 1:maxiters + retcode, iterations, x, fx, xo, J, cᵢ = Utils.init_loop_state( + ReturnCode.Default, x, fx, xo, J, cᵢ + ) + fx_sol, x_sol = Utils.fresh(fx), Utils.fresh(x) + unstable = false + + ReactantCore.@trace track_numbers = false while (!solved) & (!unstable) & + (iterations < maxiters) NLBUtils.can_setindex(x) || (A = J) # Factorize Once and Reuse - J_fact = if J isa Number - J + if J isa Number + J_fact = J else - fact = LinearAlgebra.lu(J; check = false) - !LinearAlgebra.issuccess(fact) && return SciMLBase.build_solution( - prob, alg, x, fx; retcode = ReturnCode.Unstable - ) - fact + J_fact = LinearAlgebra.lu(J; check = false) + unstable = !Utils.factorization_succeeded(J_fact) end - aᵢ = J_fact \ NLBUtils.safe_vec(fx) - hvvp = Utils.compute_hvvp( - prob, autodiff, fx_cache, x, NLBUtils.restructure(x, aᵢ) - ) - bᵢ = J_fact \ NLBUtils.safe_vec(hvvp) + if !unstable + aᵢ = J_fact \ NLBUtils.safe_vec(fx) + hvvp = Utils.compute_hvvp( + prob, autodiff, fx_cache, x, NLBUtils.restructure(x, aᵢ) + ) + bᵢ = J_fact \ NLBUtils.safe_vec(hvvp) - cᵢ_ = NLBUtils.safe_vec(cᵢ) - @bb @. cᵢ_ = (aᵢ * aᵢ) / (-aᵢ + (T(0.5) * bᵢ)) - cᵢ = NLBUtils.restructure(cᵢ, cᵢ_) + cᵢ_ = NLBUtils.safe_vec(cᵢ) + @bb @. cᵢ_ = (aᵢ * aᵢ) / (-aᵢ + (T(0.5) * bᵢ)) + cᵢ = NLBUtils.restructure(cᵢ, cᵢ_) - solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) - solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) + solved, retcode, fx_sol, + x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) + fx_sol, x_sol = Utils.fresh(fx_sol), Utils.fresh(x_sol) - @bb @. x += cᵢ - @bb copyto!(xo, x) + @bb @. x += ifelse(solved, zero(cᵢ), cᵢ) + @bb copyto!(xo, x) + xo = Utils.fresh(xo) - fx = NLBUtils.evaluate_f!!(prob, fx, x) - J = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) + fx = NLBUtils.evaluate_f!!(prob, fx, x) + J = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) + end + iterations += 1 end - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) + unstable && return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Unstable) + return Utils.simple_solution(prob, alg, x, fx, x_sol, fx_sol, retcode, solved) end diff --git a/lib/SimpleNonlinearSolve/src/klement.jl b/lib/SimpleNonlinearSolve/src/klement.jl index 3177ed807e..41a0222daa 100644 --- a/lib/SimpleNonlinearSolve/src/klement.jl +++ b/lib/SimpleNonlinearSolve/src/klement.jl @@ -18,6 +18,7 @@ function SciMLBase.__solve( x = NLBUtils.maybe_unaliased(prob.u0, _alias_u0) T = eltype(x) fx = NLBUtils.evaluate_f(prob, x) + solved = iszero(fx) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( @@ -32,8 +33,14 @@ function SciMLBase.__solve( J = one.(x) @bb δx² = similar(x) - for _ in 1:maxiters - any(iszero, J) && (J = Utils.identity_jacobian!!(J)) + retcode, iterations, x, δx, fprev, xo, J, δx² = Utils.init_loop_state( + ReturnCode.Default, x, δx, fprev, xo, J, δx² + ) + + ReactantCore.@trace track_numbers = false while (!solved) & (iterations < maxiters) + # `J` is the diagonal Jacobian approximation, so resetting it is a broadcast. + reset_jacobian = any(iszero, J) + @bb @. J = ifelse(reset_jacobian, one(J), J) @bb @. δx = fprev / J @@ -41,8 +48,8 @@ function SciMLBase.__solve( fx = NLBUtils.evaluate_f!!(prob, fx, x) # Termination Checks - solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) - solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) + solved, retcode, fx, x = Utils.check_termination(tc_cache, fx, x, xo, prob) + fx, x = Utils.fresh(fx), Utils.fresh(x) @bb δx .*= -1 @bb @. δx² = δx^2 * J^2 @@ -50,7 +57,9 @@ function SciMLBase.__solve( @bb copyto!(fprev, fx) @bb copyto!(xo, x) + fprev, xo = Utils.fresh(fprev), Utils.fresh(xo) + iterations += 1 end - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) + return Utils.simple_solution(prob, alg, x, fx, x, fx, retcode, solved) end diff --git a/lib/SimpleNonlinearSolve/src/raphson.jl b/lib/SimpleNonlinearSolve/src/raphson.jl index f29158ddf2..50e2cd7cb2 100644 --- a/lib/SimpleNonlinearSolve/src/raphson.jl +++ b/lib/SimpleNonlinearSolve/src/raphson.jl @@ -52,8 +52,7 @@ function SciMLBase.__solve( x = NLBUtils.maybe_unaliased(prob.u0, _alias_u0) fx = NLBUtils.evaluate_f(prob, x) - iszero(fx) && - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) + solved = iszero(fx) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( @@ -66,17 +65,22 @@ function SciMLBase.__solve( jac_cache = Utils.prepare_jacobian(prob, autodiff, fx_cache, x) J = Utils.compute_jacobian!!(nothing, prob, autodiff, fx_cache, x, jac_cache) - for _ in 1:maxiters + retcode, iterations, x, fx, xo, J = Utils.init_loop_state(ReturnCode.Default, x, fx, xo, J) + fx_sol, x_sol = Utils.fresh(fx), Utils.fresh(x) + + ReactantCore.@trace track_numbers = false while (!solved) & (iterations < maxiters) @bb copyto!(xo, x) + xo = Utils.fresh(xo) δx = NLBUtils.restructure(x, J \ NLBUtils.safe_vec(fx)) @bb x .-= δx solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) - solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) + fx_sol, x_sol = Utils.fresh(fx_sol), Utils.fresh(x_sol) fx = NLBUtils.evaluate_f!!(prob, fx, x) J = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) + iterations += 1 end - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) + return Utils.simple_solution(prob, alg, x, fx, x_sol, fx_sol, retcode, solved) end diff --git a/lib/SimpleNonlinearSolve/src/trust_region.jl b/lib/SimpleNonlinearSolve/src/trust_region.jl index 9b54fff202..ac900b9b13 100644 --- a/lib/SimpleNonlinearSolve/src/trust_region.jl +++ b/lib/SimpleNonlinearSolve/src/trust_region.jl @@ -112,8 +112,8 @@ function SciMLBase.__solve( ) # Set default trust region radius if not specified by user. - iszero(Δₘₐₓ) && (Δₘₐₓ = max(L2_NORM(fx), maximum(x) - minimum(x))) - if iszero(Δ) + Δₘₐₓ = ifelse(iszero(Δₘₐₓ), max(L2_NORM(fx), maximum(x) - minimum(x)), Δₘₐₓ) + if iszero(alg.initial_trust_radius) if NLBUtils.unwrap_val(alg.nlsolve_update_rule) norm_x = L2_NORM(x) Δ = T(ifelse(norm_x > 0, norm_x, 1)) @@ -136,66 +136,64 @@ function SciMLBase.__solve( solved, retcode, fx_sol, x_sol = Utils.check_termination( tc_cache, fx, x, xo, prob ) - solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) + retcode, iterations, x, xo, fx, J, H, g, Hδ, δsd, δN_δsd, δN = Utils.init_loop_state( + retcode, x, xo, fx, J, H, g, Hδ, δsd, δN_δsd, δN + ) + dogleg_cache = (; δsd, δN_δsd, δN) + fx_sol, x_sol = Utils.fresh(fx_sol), Utils.fresh(x_sol) + shrink_counter = NonlinearSolveBase.maybe_traced(shrink_counter) - for _ in 1:maxiters + ReactantCore.@trace track_numbers = false while (!solved) & + (iterations < maxiters) & (shrink_counter ≤ max_shrink_times) # Solve the trust region subproblem. δ = dogleg_method!!(dogleg_cache, J, fx, g, Δ) @bb @. x = xo + δ - fx = NLBUtils.evaluate_f!!(prob, fx, x) - fₖ₊₁ = L2_NORM(fx)^2 / T(2) # Compute the ratio of the actual to predicted reduction. @bb Hδ = H × vec(δ) r = (fₖ₊₁ - fₖ) / (dot(δ, g) + (dot(δ, Hδ) / T(2))) - # Update the trust region radius. - if r ≥ η₂ - shrink_counter = 0 - else - Δ = t₁ * Δ - shrink_counter += 1 - shrink_counter > max_shrink_times && return SciMLBase.build_solution( - prob, alg, x, fx; retcode = ReturnCode.ShrinkThresholdExceeded - ) - end - - if r ≥ η₁ - # Termination Checks - solved, retcode, fx_sol, - x_sol = Utils.check_termination( - tc_cache, fx, x, xo, prob - ) - solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) - - # Take the step. - @bb copyto!(xo, x) - - J = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) - fx = NLBUtils.evaluate_f!!(prob, fx, x) - - # Update the trust region radius. - if !NLBUtils.unwrap_val(alg.nlsolve_update_rule) && r > η₃ - Δ = min(t₂ * Δ, Δₘₐₓ) - end - fₖ = fₖ₊₁ - - @bb H = transpose(J) × J - @bb g = transpose(J) × vec(fx) + shrink = r < η₂ + shrink_counter = ifelse(shrink, shrink_counter + 1, 0) + Δ = ifelse(shrink, t₁ * Δ, Δ) + + # The step is only accepted, and termination only checked, for a sufficient + # reduction; the candidate quantities are always formed and then selected. + accepted = r ≥ η₁ + candidate_solved, retcode, fx_sol, + x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) + fx_sol, x_sol = Utils.fresh(fx_sol), Utils.fresh(x_sol) + solved = accepted & candidate_solved + @bb @. xo = ifelse(accepted, x, xo) + J_candidate = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) + fx = NLBUtils.evaluate_f!!(prob, fx, x) + @bb @. J = ifelse(accepted, J_candidate, J) + if !NLBUtils.unwrap_val(alg.nlsolve_update_rule) + expand = accepted & (r > η₃) + Δ = ifelse(expand, min(t₂ * Δ, Δₘₐₓ), Δ) end + fₖ = ifelse(accepted, fₖ₊₁, fₖ) + H_candidate = transpose(J) * J + g_candidate = NLBUtils.restructure(x, transpose(J) * NLBUtils.safe_vec(fx)) + @bb @. H = ifelse(accepted, H_candidate, H) + @bb @. g = ifelse(accepted, g_candidate, g) if NLBUtils.unwrap_val(alg.nlsolve_update_rule) - if r > η₃ - Δ = t₂ * L2_NORM(δ) - elseif r > 0.5 - Δ = max(Δ, t₂ * L2_NORM(δ)) - end + expanded = t₂ * L2_NORM(δ) + Δ = ifelse(r > η₃, expanded, ifelse(r > 0.5, max(Δ, expanded), Δ)) end + iterations += 1 end - return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) + failure_retcode = ifelse( + shrink_counter > max_shrink_times, ReturnCode.ShrinkThresholdExceeded, + ReturnCode.MaxIters + ) + return Utils.simple_solution( + prob, alg, x, fx, x_sol, fx_sol, retcode, solved; failure_retcode + ) end function dogleg_method!!(cache, J, f::F, g, Δ) where {F} @@ -204,15 +202,18 @@ function dogleg_method!!(cache, J, f::F, g, Δ) where {F} # Compute the Newton step @bb δN .= NLBUtils.restructure(δN, J \ NLBUtils.safe_vec(f)) @bb δN .*= -1 - # Test if the full step is within the trust region - (L2_NORM(δN) ≤ Δ) && return δN + + # Under Reactant no branch can be taken on the traced norms, so every candidate step is + # formed and the result selected at the end; on the ordinary path the early returns + # avoid the extra work. + newton_step = L2_NORM(δN) ≤ Δ + !ReactantCore.within_compile() && newton_step && return δN # Calculate Cauchy point, optimum along the steepest descent direction - @bb δsd .= g - @bb @. δsd *= -1 + @bb @. δsd = -g norm_δsd = L2_NORM(δsd) - - if (norm_δsd ≥ Δ) + cauchy_step = norm_δsd ≥ Δ + if !ReactantCore.within_compile() && cauchy_step @bb @. δsd *= Δ / norm_δsd return δsd end @@ -223,7 +224,8 @@ function dogleg_method!!(cache, J, f::F, g, Δ) where {F} dot_δsd_δN_δsd = dot(δsd, δN_δsd) dot_δsd = dot(δsd, δsd) fact = dot_δsd_δN_δsd^2 - dot_δN_δsd * (dot_δsd - Δ^2) - tau = (-dot_δsd_δN_δsd + sqrt(fact)) / dot_δN_δsd - @bb @. δsd += tau * δN_δsd - return δsd + tau = (-dot_δsd_δN_δsd + sqrt(max(zero(fact), fact))) / dot_δN_δsd + @bb @. δN_δsd = δsd + tau * δN_δsd + @bb @. δN = ifelse(newton_step, δN, ifelse(cauchy_step, δsd * Δ / norm_δsd, δN_δsd)) + return δN end diff --git a/lib/SimpleNonlinearSolve/src/utils.jl b/lib/SimpleNonlinearSolve/src/utils.jl index bd2d81b06f..57626f72dc 100644 --- a/lib/SimpleNonlinearSolve/src/utils.jl +++ b/lib/SimpleNonlinearSolve/src/utils.jl @@ -4,6 +4,7 @@ using ArrayInterface: ArrayInterface using DifferentiationInterface: DifferentiationInterface, Constant using FastClosures: @closure using LinearAlgebra: LinearAlgebra, I, diagind +using ReactantCore: ReactantCore using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearTerminationMode, AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode @@ -198,13 +199,54 @@ function compute_hvvp(prob, autodiff, fx, x, dir) return only(DI.pushforward(jvp_fn, autodiff, x, (dir,), Constant(prob.p))) end -function nonlinear_solution_new_alg( - sol::SciMLBase.NonlinearSolution{T, N, uType, R, P, A, O, uType2, S, Tr}, alg - ) where {T, N, uType, R, P, A, O, uType2, S, Tr} - return SciMLBase.NonlinearSolution{T, N, uType, R, P, typeof(alg), O, uType2, S, Tr}( +function nonlinear_solution_new_alg(sol::SciMLBase.NonlinearSolution, alg) + return SciMLBase.NonlinearSolution( sol.u, sol.resid, sol.prob, alg, sol.retcode, sol.original, sol.left, sol.right, sol.stats, sol.trace ) end +""" + init_loop_state(retcode, xs...) + +Prepare the state carried by a solver loop: `retcode` and an iteration counter become +traced under Reactant, and each array in `xs` is made distinct from every other one. Returns +`(retcode, iterations, xs...)`. +""" +function init_loop_state(retcode, xs...) + return NonlinearSolveBase.maybe_traced(retcode), NonlinearSolveBase.maybe_traced(0), + map(fresh, xs)... +end + +# Also needed after `@bb copyto!(dst, src)`, which rebinds `dst = src` for arrays that +# cannot be indexed into (traced arrays included) and so would alias the two. +fresh(x) = NonlinearSolveBase.dealias_traced!(x) + +# A traced factorization cannot report failure; the residual shows it instead. +factorization_succeeded(fact) = ReactantCore.within_compile() || LinearAlgebra.issuccess(fact) + +""" + simple_solution(prob, alg, x, fx, x_sol, fx_sol, retcode, solved) + +Build the solution of a simple solver loop. `x_sol`/`fx_sol` are the iterate reported by +the termination check and are returned when `solved`; otherwise the last iterate `x`/`fx` +is returned with `failure_retcode`. A `retcode` still at `Default` on success means the +initial guess was already a root. +""" +function simple_solution( + prob, alg, x, fx, x_sol, fx_sol, retcode, solved; + failure_retcode = ReturnCode.MaxIters + ) + # `solved` may not infer when the Jacobian type is only known at run time; the results + # have the types of their inputs on both the host and the traced path. + retcode = ifelse( + solved, + ifelse(retcode == ReturnCode.Default, ReturnCode.Success, retcode), + failure_retcode + )::typeof(retcode) + u = NonlinearSolveBase.select(solved, x_sol, x)::typeof(x) + resid = NonlinearSolveBase.select(solved, fx_sol, fx)::typeof(fx) + return NonlinearSolveBase.build_nonlinear_solution(prob, alg, u, resid; retcode) +end + end diff --git a/test/Reactant/reactant_tests.jl b/test/Reactant/reactant_tests.jl index e7f995086d..bdf8466c0e 100644 --- a/test/Reactant/reactant_tests.jl +++ b/test/Reactant/reactant_tests.jl @@ -9,6 +9,14 @@ jac(u, p) = reshape(2 .* u, :, 1) .* Float32[1 0; 0 1] nonlinear_function = NonlinearFunction(f; jac) autodiff_nonlinear_function = NonlinearFunction(f) +function solve_simple_broyden(u, p) + return solve(NonlinearProblem(nonlinear_function, u, p), SimpleBroyden()) +end + +function solve_simple_klement(u, p) + return solve(NonlinearProblem(nonlinear_function, u, p), SimpleKlement()) +end + function solve_newton(u, p) return solve(NonlinearProblem(nonlinear_function, u, p), NewtonRaphson()) end @@ -52,6 +60,8 @@ end u0 = Reactant.to_rarray(Float32[1, 1]) p0 = Reactant.to_rarray(Float32[2]) +compiled_broyden = Reactant.@compile solve_simple_broyden(u0, p0) +compiled_klement = Reactant.@compile solve_simple_klement(u0, p0) compiled_newton = Reactant.@compile solve_newton(u0, p0) compiled_trust_region = Reactant.@compile solve_trust_region(u0, p0) compiled_default = Reactant.@compile solve_default(u0, p0) @@ -61,6 +71,29 @@ compiled_autodiff_trust_region = Reactant.@compile solve_autodiff_trust_region(u compiled_autodiff_default = Reactant.@compile solve_autodiff_default(u0, p0) compiled_autodiff_gauss_newton = Reactant.@compile solve_autodiff_gauss_newton(u0, p0) +for (compiled, Alg) in ( + (compiled_broyden, SimpleBroyden), (compiled_klement, SimpleKlement), + ) + sol = compiled( + Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[2]) + ) + sol_far = compiled( + Reactant.to_rarray(Float32[10, 10]), Reactant.to_rarray(Float32[2]) + ) + + @test sol.u isa Reactant.ConcreteRArray + @test Array(sol.u) ≈ fill(sqrt(2.0f0), 2) + @test maximum(abs, Array(sol.resid)) ≤ 1.0f-5 + @test sol.retcode == ReturnCode.Success + @test SciMLBase.successful_retcode(sol) + @test sol.alg isa Alg + @test sol.prob === nothing + @test sol.stats === nothing + + @test Array(sol_far.u) ≈ fill(sqrt(2.0f0), 2) + @test sol_far.retcode == ReturnCode.Success +end + # A polyalgorithm's members keep `autodiff = nothing`; the backend is chosen when each # member is solved, so the choice is only visible on a directly solved algorithm. @@ -92,6 +125,17 @@ for (compiled, name, uses_enzyme) in ( @test sol.stats === nothing end +sol_converged = compiled_broyden( + Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[1]) +) +@test sol_converged.retcode == ReturnCode.Success +@test Array(sol_converged.u) == Float32[1, 1] + +function solve_one_step(u, p) + return solve( + NonlinearProblem(nonlinear_function, u, p), SimpleKlement(); maxiters = 1 + ) +end function solve_newton_one_step(u, p) return solve( @@ -99,6 +143,11 @@ function solve_newton_one_step(u, p) ) end +sol_maxiters = Reactant.@jit solve_one_step( + Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[2]) +) +@test sol_maxiters.retcode == ReturnCode.MaxIters +@test !SciMLBase.successful_retcode(sol_maxiters) sol_newton_maxiters = Reactant.@jit solve_newton_one_step( Reactant.to_rarray(Float32[1, 1]), Reactant.to_rarray(Float32[2]) @@ -133,6 +182,12 @@ reactant_solver_cases = ( LevenbergMarquardt(; disable_geodesic = Val(true)), ), (:PseudoTransient, NonlinearProblem, PseudoTransient()), + (:SimpleNewtonRaphson, NonlinearProblem, SimpleNewtonRaphson()), + (:SimpleTrustRegion, NonlinearProblem, SimpleTrustRegion()), + (:SimpleBroyden, NonlinearProblem, SimpleBroyden()), + (:SimpleKlement, NonlinearProblem, SimpleKlement()), + (:SimpleLimitedMemoryBroyden, NonlinearProblem, SimpleLimitedMemoryBroyden()), + (:SimpleHalley, NonlinearProblem, SimpleHalley()), ( :FastShortcutNonlinearPolyalg, NonlinearProblem, @@ -151,6 +206,7 @@ reactant_solver_cases = ( NonlinearLeastSquaresProblem, LevenbergMarquardt(), ), + (:SimpleGaussNewton, NonlinearLeastSquaresProblem, SimpleGaussNewton()), ) @testset "Analytical Jacobian: $name" for (name, problem_type, alg) in reactant_solver_cases