Allow allocations for the termination condition version

Author: avik-pal

.buildkite/pipeline.yml

@@ -0,0 +1,15 @@
+steps:
+  - label: "Julia 1"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 30
+    # Don't run Buildkite if the commit message includes the text [skip tests]
+    if: build.message !~ /\[skip tests\]/
+
+env:
+  GROUP: CUDA
+  JULIA_PKG_SERVER: ""

src/nlsolve/lbroyden.jl

@@ -21,7 +21,22 @@ function SimpleLimitedMemoryBroyden(; threshold::Union{Val, Int} = Val(27))
     return SimpleLimitedMemoryBroyden{SciMLBase._unwrap_val(threshold)}()
 end
 
-@views function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleLimitedMemoryBroyden,
+function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleLimitedMemoryBroyden,
+        args...; termination_condition = nothing, kwargs...)
+    if prob.u0 isa SArray
+        if termination_condition === nothing ||
+           termination_condition isa AbsNormTerminationMode
+            return __static_solve(prob, alg, args...; termination_condition, kwargs...)
+        end
+        @warn "Specifying `termination_condition = $(termination_condition)` for \
+               `SimpleLimitedMemoryBroyden` with `SArray` is not non-allocating. Use \
+               either `termination_condition = AbsNormTerminationMode()` or \
+               `termination_condition = nothing`." maxlog=1
+    end
+    return __generic_solve(prob, alg, args...; termination_condition, kwargs...)
+end
+
+@views function __generic_solve(prob::NonlinearProblem, alg::SimpleLimitedMemoryBroyden,
         args...; abstol = nothing, reltol = nothing, maxiters = 1000, alias_u0 = false,
         termination_condition = nothing, kwargs...)
     x = __maybe_unaliased(prob.u0, alias_u0)
@@ -85,17 +100,9 @@ end
 
 # Non-allocating StaticArrays version of SimpleLimitedMemoryBroyden is actually quite
 # finicky, so we'll implement it separately from the generic version
-# We make an exception here and don't support termination conditions
-@views function SciMLBase.__solve(prob::NonlinearProblem{<:SArray},
-        alg::SimpleLimitedMemoryBroyden, args...; abstol = nothing,
-        termination_condition = nothing,
-        maxiters = 1000, kwargs...)
-    if termination_condition !== nothing &&
-       !(termination_condition isa AbsNormTerminationMode)
-        error("SimpleLimitedMemoryBroyden with StaticArrays does not support termination \
-               conditions!")
-    end
-
+# Ignore termination_condition. Don't pass things into internal functions
+function __static_solve(prob::NonlinearProblem{<:SArray}, alg::SimpleLimitedMemoryBroyden,
+        args...; abstol = nothing, maxiters = 1000, kwargs...)
     x = prob.u0
     fx = _get_fx(prob, x)
     threshold = __get_threshold(alg)

test/cuda.jl

@@ -0,0 +1,35 @@
+using SimpleNonlinearSolve, StaticArrays, CUDA, Test
+
+CUDA.allowscalar(false)
+
+f(u, p) = u .* u .- 2
+f!(du, u, p) = du .= u .* u .- 2
+
+@testset "Solving on GPUs" begin
+    for alg in (SimpleNewtonRaphson(), SimpleDFSane(), SimpleTrustRegion(), SimpleBroyden(),
+        SimpleLimitedMemoryBroyden(), SimpleKlement(), SimpleHalley())
+        @info "Testing $alg on CUDA"
+
+        # Static Arrays
+        u0 = @SVector[1.0f0, 1.0f0]
+        probN = NonlinearProblem{false}(f, u0)
+        sol = solve(probN, alg; abstol = 1.0f-6)
+        @test SciMLBase.successful_retcode(sol)
+        @test maximum(abs, sol.resid) ≤ 1.0f-6
+
+        # Regular Arrays
+        u0 = [1.0, 1.0]
+        probN = NonlinearProblem{false}(f, u0)
+        sol = solve(probN, alg; abstol = 1.0f-6)
+        @test SciMLBase.successful_retcode(sol)
+        @test maximum(abs, sol.resid) ≤ 1.0f-6
+
+        # Regular Arrays Inplace
+        alg isa SimpleHalley && continue
+        u0 = [1.0, 1.0]
+        probN = NonlinearProblem{true}(f!, u0)
+        sol = solve(probN, alg; abstol = 1.0f-6)
+        @test SciMLBase.successful_retcode(sol)
+        @test maximum(abs, sol.resid) ≤ 1.0f-6
+    end
+end

test/cuda/Project.toml

@@ -0,0 +1,5 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/runtests.jl

@@ -2,6 +2,12 @@ using SafeTestsets, Test
 
 const GROUP = get(ENV, "GROUP", "All")
 
+function activate_env(env)
+    Pkg.activate(env)
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+end
+
 @time @testset "SimpleNonlinearSolve.jl" begin
     if GROUP == "All" || GROUP == "Core"
         @time @safetestset "Basic Tests" include("basictests.jl")
@@ -10,4 +16,9 @@ const GROUP = get(ENV, "GROUP", "All")
         @time @safetestset "Least Squares Tests" include("least_squares.jl")
         @time @safetestset "23 Test Problems" include("23_test_problems.jl")
     end
+
+    if GROUP == "CUDA"
+        activate_env("cuda")
+        @time @safetestset "CUDA Tests" include("cuda.jl")
+    end
 end