Resolve tolerances differently for kernel launches

Author: avik-pal

src/SimpleNonlinearSolve.jl

@@ -9,7 +9,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
     import DiffEqBase: AbstractNonlinearTerminationMode,
         AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode,
         NonlinearSafeTerminationReturnCode, get_termination_mode,
-        NONLINEARSOLVE_DEFAULT_NORM, _get_tolerance
+        NONLINEARSOLVE_DEFAULT_NORM
     import ForwardDiff: Dual
     import MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex
     import SciMLBase: AbstractNonlinearAlgorithm, build_solution, isinplace, _unwrap_val
@@ -65,6 +65,15 @@ function SciMLBase.solve(prob::NonlinearProblem, alg::AbstractSimpleNonlinearSol
     return SciMLBase.__solve(prob, alg, args...; kwargs...)
 end
 
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{<:Number, <:SArray}},
+        alg::AbstractSimpleNonlinearSolveAlgorithm, args...; abstol = nothing,
+        reltol = nothing, kwargs...)
+    _abstol = __get_tolerance(prob.u0, abstol, eltype(prob.u0))
+    _reltol = __get_tolerance(prob.u0, reltol, eltype(prob.u0))
+    return SciMLBase.__solve(prob, alg, args...; abstol = _abstol, reltol = _reltol,
+        kwargs...)
+end
+
 @setup_workload begin
     for T in (Float32, Float64)
         prob_no_brack_scalar = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))

src/ad.jl

@@ -8,6 +8,20 @@ function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}
 end
 
 # Handle Ambiguities
+for algType in (SimpleNewtonRaphson, SimpleDFSane, SimpleTrustRegion, SimpleBroyden,
+    SimpleLimitedMemoryBroyden, SimpleKlement, SimpleHalley)
+    @eval begin
+        function SciMLBase.solve(prob::NonlinearProblem{uType, iip,
+                    <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}},
+                alg::$(algType), args...; kwargs...) where {uType, T, V, P, iip}
+            sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...)
+            dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p)
+            return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode,
+                sol.stats, sol.original)
+        end
+    end
+end
+
 for algType in (Bisection, Brent, Alefeld, Falsi, ITP, Ridder)
     @eval begin
         function SciMLBase.solve(prob::IntervalNonlinearProblem{uType, iip,

src/bracketing/bisection.jl

@@ -26,7 +26,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Bisection, args...
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)

src/bracketing/brent.jl

@@ -13,7 +13,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Brent, args...;
     fl, fr = f(left), f(right)
     ϵ = eps(convert(typeof(fl), 1))
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)

src/bracketing/falsi.jl

@@ -12,7 +12,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...;
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)

src/bracketing/itp.jl

@@ -58,7 +58,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::ITP, args...;
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)

src/bracketing/ridder.jl

@@ -12,7 +12,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...;
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)

src/nlsolve/lbroyden.jl

@@ -28,6 +28,12 @@ function SimpleLimitedMemoryBroyden(; threshold::Union{Val, Int} = Val(27),
     return SimpleLimitedMemoryBroyden{_unwrap_val(threshold), _unwrap_val(linesearch)}()
 end
 
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{<:Number, <:SArray}},
+        alg::SimpleLimitedMemoryBroyden, args...; kwargs...)
+    # Don't resolve the `abstol` and `reltol` here
+    return SciMLBase.__solve(prob, alg, args...; kwargs...)
+end
+
 function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleLimitedMemoryBroyden,
         args...; termination_condition = nothing, kwargs...)
     if prob.u0 isa SArray
@@ -120,7 +126,7 @@ function __static_solve(prob::NonlinearProblem{<:SArray}, alg::SimpleLimitedMemo
 
     U, Vᵀ = __init_low_rank_jacobian(vec(x), vec(fx), threshold)
 
-    abstol = DiffEqBase._get_tolerance(abstol, eltype(x))
+    abstol = __get_tolerance(x, abstol, eltype(x))
 
     xo, δx, fo, δf = x, -fx, fx, fx
 

src/utils.jl

@@ -370,3 +370,10 @@ end
 
 @inline __reshape(x::Number, args...) = x
 @inline __reshape(x::AbstractArray, args...) = reshape(x, args...)
+
+# Override cases which might be used in a kernel launch
+__get_tolerance(x, η, ::Type{T}) where {T} = DiffEqBase._get_tolerance(η, T)
+function __get_tolerance(x::Union{SArray, Number}, ::Nothing, ::Type{T}) where {T}
+    η = real(oneunit(T)) * (eps(real(one(T))))^(real(T)(0.8))
+    return T(η)
+end

test/cuda.jl

@@ -37,7 +37,7 @@ f!(du, u, p) = du .= u .* u .- 2
 end
 
 function kernel_function(prob, alg)
-    solve(prob, alg; abstol = 1.0f-6, reltol = 1.0f-6)
+    solve(prob, alg)
     return nothing
 end