Merge #121

121: Re-design run for signal handlers r=charleskawczynski a=charleskawczynski

This PR re-designs the signal handlers to improve inference.

Co-authored-by: Charles Kawczynski <[email protected]>

Author: bors[bot]

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.5.0"
+version = "0.6.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

src/ClimaTimeSteppers.jl

@@ -56,8 +56,6 @@ array_device(::CuArray) = CUDADevice()
 realview(x::Union{Array, SArray, MArray}) = x
 realview(x::CuArray) = x
 
-export allocate_cache, run!
-
 import DiffEqBase, SciMLBase, LinearAlgebra, DiffEqCallbacks, Krylov
 
 include("sparse_containers.jl")

src/solvers/newtons_method.jl

@@ -564,7 +564,9 @@ function solve_newton!(alg::NewtonsMethod, cache, x, f!, j! = nothing)
     (; max_iters, update_j, krylov_method, convergence_checker, verbose) = alg
     (; update_j_cache, krylov_method_cache, convergence_checker_cache) = cache
     (; Δx, f, j) = cache
-    isnothing(j) || run!(update_j, update_j_cache, NewNewtonSolve(), j!, j, x)
+    if (!isnothing(j)) && needs_update!(update_j, update_j_cache, NewNewtonSolve())
+        j!(j, x)
+    end
     for n in 0:max_iters
         # Update x[n] with Δx[n - 1], and exit the loop if Δx[n] is not needed.
         n > 0 && (x .-= Δx)
@@ -574,7 +576,9 @@ function solve_newton!(alg::NewtonsMethod, cache, x, f!, j! = nothing)
         end
 
         # Compute Δx[n].
-        isnothing(j) || run!(update_j, update_j_cache, NewNewtonIteration(), j!, j, x)
+        if (!isnothing(j)) && needs_update!(update_j, update_j_cache, NewNewtonIteration())
+            j!(j, x)
+        end
         f!(f, x)
         if isnothing(krylov_method)
             if j isa DenseMatrix
@@ -598,5 +602,7 @@ end
 function update!(alg::NewtonsMethod, cache, signal::UpdateSignal, j!)
     (; update_j) = alg
     (; update_j_cache, j) = cache
-    isnothing(j) || run!(update_j, update_j_cache, signal, j!, j)
+    if (!isnothing(j)) && needs_update!(update_j, update_j_cache, signal)
+        j!(j)
+    end
 end

src/solvers/update_signal_handler.jl

@@ -12,15 +12,16 @@ abstract type UpdateSignal end
 """
     UpdateSignalHandler
 
-Updates a value upon receiving an appropriate `UpdateSignal`. This is done by
-calling `run!(::UpdateSignalHandler, cache, ::UpdateSignal, f!, args...)`, where
-`f!` is function such that `f!(args...)` modifies the desired value in-place.
+A boolean indicating if updates a value upon receiving an appropriate
+`UpdateSignal`. This is done by calling
+`needs_update!(::UpdateSignalHandler, cache, ::UpdateSignal)`.
+
 The `cache` can be obtained with `allocate_cache(::UpdateSignalHandler, FT)`,
 where `FT` is the floating-point type of the integrator.
 """
 abstract type UpdateSignalHandler end
 
-run!(::UpdateSignalHandler, cache, ::UpdateSignal, f!, args...) = nothing
+needs_update!(::UpdateSignalHandler, cache, ::UpdateSignal) = false
 
 """
     NewTimeStep(t)
@@ -34,7 +35,7 @@ end
 """
     NewNewtonSolve()
 
-The signal for a new `run!` of Newton's method, which occurs on every implicit
+The signal for a new `needs_update!` of Newton's method, which occurs on every implicit
 Runge-Kutta stage of the integrator.
 """
 struct NewNewtonSolve <: UpdateSignal end
@@ -49,24 +50,24 @@ struct NewNewtonIteration <: UpdateSignal end
 """
     UpdateEvery(update_signal_type)
 
-An `UpdateSignalHandler` that performs the update whenever it is `run!` with an
+An `UpdateSignalHandler` that performs the update whenever it is `needs_update!` with an
 `UpdateSignal` of type `update_signal_type`.
 """
 struct UpdateEvery{U <: UpdateSignal} <: UpdateSignalHandler end
 UpdateEvery(::Type{U}) where {U} = UpdateEvery{U}()
 
 allocate_cache(::UpdateEvery, _) = nothing
 
-run!(alg::UpdateEvery{U}, cache, ::U, f!, args...) where {U <: UpdateSignal} = f!(args...)
+needs_update!(alg::UpdateEvery{U}, cache, ::U) where {U <: UpdateSignal} = true
 
 """
     UpdateEveryN(n, update_signal_type, reset_signal_type = Nothing)
 
-An `UpdateSignalHandler` that performs the update every `n`-th time it is `run!`
+An `UpdateSignalHandler` that performs the update every `n`-th time it is `needs_update!`
 with an `UpdateSignal` of type `update_signal_type`. If `reset_signal_type` is
 specified, then the counter (which gets incremented from 0 to `n` and then gets
 reset to 0 when it is time to perform another update) is reset to 0 whenever the
-signal handler is `run!` with an `UpdateSignal` of type `reset_signal_type`.
+signal handler is `needs_update!` with an `UpdateSignal` of type `reset_signal_type`.
 """
 struct UpdateEveryN{U <: UpdateSignal, R <: Union{Nothing, UpdateSignal}} <: UpdateSignalHandler
     n::Int
@@ -75,26 +76,30 @@ UpdateEveryN(n, ::Type{U}, ::Type{R} = Nothing) where {U, R} = UpdateEveryN{U, R
 
 allocate_cache(::UpdateEveryN, _) = (; counter = Ref(0))
 
-function run!(alg::UpdateEveryN{U}, cache, ::U, f!, args...) where {U}
+function needs_update!(alg::UpdateEveryN{U}, cache, ::U) where {U <: UpdateSignal}
     (; n) = alg
     (; counter) = cache
-    if counter[] == 0
-        f!(args...)
-    end
+    result = counter[] == 0
     counter[] += 1
     if counter[] == n
         counter[] = 0
     end
+    return result
 end
-function run!(alg::UpdateEveryN{<:Any, R}, cache, ::R, f!, args...) where {R}
+function needs_update!(alg::UpdateEveryN{U, R}, cache, ::R) where {U, R <: UpdateSignal}
     (; counter) = cache
     counter[] = 0
+    return false
 end
 
+# Account for method ambiguitiy:
+needs_update!(::UpdateEveryN{U, U}, cache, ::U) where {U <: UpdateSignal} =
+    error("Reset and update signal types cannot be the same.")
+
 """
     UpdateEveryDt(dt)
 
-An `UpdateSignalHandler` that performs the update whenever it is `run!` with an
+An `UpdateSignalHandler` that performs the update whenever it is `needs_update!` with an
 `UpdateSignal` of type `NewTimeStep` and the difference between the current time
 and the previous update time is no less than `dt`.
 """
@@ -105,13 +110,15 @@ end
 # TODO: This assumes that typeof(t) == FT, which might not always be correct.
 allocate_cache(alg::UpdateEveryDt, ::Type{FT}) where {FT} = (; is_first_t = Ref(true), prev_update_t = Ref{FT}())
 
-function run!(alg::UpdateEveryDt, cache, signal::NewTimeStep, f!, args...)
+function needs_update!(alg::UpdateEveryDt, cache, signal::NewTimeStep)
     (; dt) = alg
     (; is_first_t, prev_update_t) = cache
     (; t) = signal
+    result = false
     if is_first_t[] || abs(t - prev_update_t[]) >= dt
-        f!(args...)
+        result = true
         is_first_t[] = false
         prev_update_t[] = t
     end
+    return result
 end