SimpleNonlinearSolve: single branchless loops that trace under Reactant

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/<Sublib> -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.

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
+```

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

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

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(

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)

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]

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

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`.

lib/NonlinearSolveBase/src/descent/dogleg.jl

@@ -97,36 +97,25 @@ 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)))
     end
 
     # 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