From 174051bc9ba5420d5a879eed7f000a095c981d49 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Oct 2023 07:44:05 -0700 Subject: [PATCH 1/2] Add post_explicit call at end of step, rm nvtx --- src/solvers/imex_ark.jl | 219 +++++++++++++++----------------------- src/solvers/imex_ssprk.jl | 1 + 2 files changed, 89 insertions(+), 131 deletions(-) diff --git a/src/solvers/imex_ark.jl b/src/solvers/imex_ark.jl index 5250f2cb..8a3e8f3c 100644 --- a/src/solvers/imex_ark.jl +++ b/src/solvers/imex_ark.jl @@ -59,132 +59,102 @@ function step_u!(integrator, cache::IMEXARKCache, f, name) s = length(b_exp) if !isnothing(T_imp!) && !isnothing(newtons_method) - NVTX.@range "update!" color = colorant"yellow" begin - (; update_j) = newtons_method - jacobian = newtons_method_cache.j - if (!isnothing(jacobian)) && needs_update!(update_j, NewTimeStep(t)) - if γ isa Nothing - sdirk_error(name) - else - T_imp!.Wfact(jacobian, u, p, dt * γ, t) - end + (; update_j) = newtons_method + jacobian = newtons_method_cache.j + if (!isnothing(jacobian)) && needs_update!(update_j, NewTimeStep(t)) + if γ isa Nothing + sdirk_error(name) + else + T_imp!.Wfact(jacobian, u, p, dt * γ, t) end end end for i in 1:s - NVTX.@range "stage" payload = i begin - t_exp = t + dt * c_exp[i] - t_imp = t + dt * c_imp[i] + t_exp = t + dt * c_exp[i] + t_imp = t + dt * c_imp[i] - NVTX.@range "U = u" color = colorant"yellow" begin - @. U = u - end + @. U = u - if !isnothing(T_lim!) # Update based on limited tendencies from previous stages - NVTX.@range "U+=dt*a_exp*T_lim" color = colorant"yellow" begin - for j in 1:(i - 1) - iszero(a_exp[i, j]) && continue - @. U += dt * a_exp[i, j] * T_lim[j] - end - end - NVTX.@range "lim!" color = colorant"yellow" begin - lim!(U, p, t_exp, u) - end - end - - if !isnothing(T_exp!) # Update based on explicit tendencies from previous stages - NVTX.@range "U+=dt*a_exp*T_exp" color = colorant"yellow" begin - for j in 1:(i - 1) - iszero(a_exp[i, j]) && continue - @. U += dt * a_exp[i, j] * T_exp[j] - end - end + if !isnothing(T_lim!) # Update based on limited tendencies from previous stages + for j in 1:(i - 1) + iszero(a_exp[i, j]) && continue + @. U += dt * a_exp[i, j] * T_lim[j] end + lim!(U, p, t_exp, u) + end - if !isnothing(T_imp!) # Update based on implicit tendencies from previous stages - NVTX.@range "U+=dt*a_imp*T_imp" color = colorant"yellow" begin - for j in 1:(i - 1) - iszero(a_imp[i, j]) && continue - @. U += dt * a_imp[i, j] * T_imp[j] - end - end + if !isnothing(T_exp!) # Update based on explicit tendencies from previous stages + for j in 1:(i - 1) + iszero(a_exp[i, j]) && continue + @. U += dt * a_exp[i, j] * T_exp[j] end + end - NVTX.@range "dss!" color = colorant"yellow" begin - dss!(U, p, t_exp) + if !isnothing(T_imp!) # Update based on implicit tendencies from previous stages + for j in 1:(i - 1) + iszero(a_imp[i, j]) && continue + @. U += dt * a_imp[i, j] * T_imp[j] end + end - if !(!isnothing(T_imp!) && !iszero(a_imp[i, i])) # Implicit solve - post_explicit!(U, p, t_imp) - else - @assert !isnothing(newtons_method) - NVTX.@range "temp = U" color = colorant"yellow" begin - @. temp = U + dss!(U, p, t_exp) + + if !(!isnothing(T_imp!) && !iszero(a_imp[i, i])) # Implicit solve + post_explicit!(U, p, t_imp) + else + @assert !isnothing(newtons_method) + @. temp = U + post_explicit!(U, p, t_imp) + # TODO: can/should we remove these closures? + implicit_equation_residual! = + (residual, Ui) -> begin + T_imp!(residual, Ui, p, t_imp) + @. residual = temp + dt * a_imp[i, i] * residual - Ui end - post_explicit!(U, p, t_imp) - # TODO: can/should we remove these closures? - implicit_equation_residual! = - (residual, Ui) -> begin - NVTX.@range "T_imp!" color = colorant"yellow" begin - T_imp!(residual, Ui, p, t_imp) - end - NVTX.@range "residual=temp+dt*a_imp*residual-Ui" color = colorant"yellow" begin - @. residual = temp + dt * a_imp[i, i] * residual - Ui - end + implicit_equation_jacobian! = (jacobian, Ui) -> T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) + call_post_implicit! = Ui -> begin + post_implicit!(Ui, p, t_imp) + end + call_post_implicit_last! = + Ui -> begin + if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && !iszero(a_imp[i, i]) + # If T_imp[i] is being treated implicitly, ensure that it + # exactly satisfies the implicit equation. + @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) end - implicit_equation_jacobian! = (jacobian, Ui) -> T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) - call_post_implicit! = Ui -> begin post_implicit!(Ui, p, t_imp) end - call_post_implicit_last! = - Ui -> begin - if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && !iszero(a_imp[i, i]) - # If T_imp[i] is being treated implicitly, ensure that it - # exactly satisfies the implicit equation. - @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) - end - post_implicit!(Ui, p, t_imp) - end - NVTX.@range "solve_newton!" color = colorant"yellow" begin - solve_newton!( - newtons_method, - newtons_method_cache, - U, - implicit_equation_residual!, - implicit_equation_jacobian!, - call_post_implicit!, - call_post_implicit_last!, - ) - end - end + solve_newton!( + newtons_method, + newtons_method_cache, + U, + implicit_equation_residual!, + implicit_equation_jacobian!, + call_post_implicit!, + call_post_implicit_last!, + ) + end - # We do not need to DSS U again because the implicit solve should - # give the same results for redundant columns (as long as the implicit - # tendency only acts in the vertical direction). + # We do not need to DSS U again because the implicit solve should + # give the same results for redundant columns (as long as the implicit + # tendency only acts in the vertical direction). - if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i]) - if iszero(a_imp[i, i]) && !isnothing(T_imp!) - # If its coefficient is 0, T_imp[i] is effectively being - # treated explicitly. - NVTX.@range "T_imp!" color = colorant"yellow" begin - T_imp!(T_imp[i], U, p, t_imp) - end - end + if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i]) + if iszero(a_imp[i, i]) && !isnothing(T_imp!) + # If its coefficient is 0, T_imp[i] is effectively being + # treated explicitly. + T_imp!(T_imp[i], U, p, t_imp) end + end - if !all(iszero, a_exp[:, i]) || !iszero(b_exp[i]) - if !isnothing(T_lim!) - NVTX.@range "T_lim!" color = colorant"yellow" begin - T_lim!(T_lim[i], U, p, t_exp) - end - end - if !isnothing(T_exp!) - NVTX.@range "T_exp!" color = colorant"yellow" begin - T_exp!(T_exp[i], U, p, t_exp) - end - end + if !all(iszero, a_exp[:, i]) || !iszero(b_exp[i]) + if !isnothing(T_lim!) + T_lim!(T_lim[i], U, p, t_exp) + end + if !isnothing(T_exp!) + T_exp!(T_exp[i], U, p, t_exp) end end end @@ -192,44 +162,31 @@ function step_u!(integrator, cache::IMEXARKCache, f, name) t_final = t + dt if !isnothing(T_lim!) # Update based on limited tendencies from previous stages - NVTX.@range "temp=u" color = colorant"yellow" begin - @. temp = u - end - NVTX.@range "temp+=dt*b_exp*T_lim" color = colorant"yellow" begin - for j in 1:s - iszero(b_exp[j]) && continue - @. temp += dt * b_exp[j] * T_lim[j] - end - end - NVTX.@range "lim!" color = colorant"yellow" begin - lim!(temp, p, t_final, u) - end - NVTX.@range "u=temp" color = colorant"yellow" begin - @. u = temp + @. temp = u + for j in 1:s + iszero(b_exp[j]) && continue + @. temp += dt * b_exp[j] * T_lim[j] end + lim!(temp, p, t_final, u) + @. u = temp end if !isnothing(T_exp!) # Update based on explicit tendencies from previous stages - NVTX.@range "u+=dt*b_exp*T_exp" color = colorant"yellow" begin - for j in 1:s - iszero(b_exp[j]) && continue - @. u += dt * b_exp[j] * T_exp[j] - end + for j in 1:s + iszero(b_exp[j]) && continue + @. u += dt * b_exp[j] * T_exp[j] end end if !isnothing(T_imp!) # Update based on implicit tendencies from previous stages - NVTX.@range "u+=dt*b_imp*T_imp" color = colorant"yellow" begin - for j in 1:s - iszero(b_imp[j]) && continue - @. u += dt * b_imp[j] * T_imp[j] - end + for j in 1:s + iszero(b_imp[j]) && continue + @. u += dt * b_imp[j] * T_imp[j] end end - NVTX.@range "dss!" color = colorant"yellow" begin - dss!(u, p, t_final) - end + dss!(u, p, t_final) + post_explicit!(u, p, t_final) return u end diff --git a/src/solvers/imex_ssprk.jl b/src/solvers/imex_ssprk.jl index 427ac1bc..5f5982ba 100644 --- a/src/solvers/imex_ssprk.jl +++ b/src/solvers/imex_ssprk.jl @@ -185,6 +185,7 @@ function step_u!(integrator, cache::IMEXSSPRKCache, f, name) end dss!(u, p, t_final) + post_explicit!(u, p, t_final) return u end From 57c3828df02a0eca0f2cf60133e8ec96ee508232 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Oct 2023 07:44:23 -0700 Subject: [PATCH 2/2] Bump patch version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 07396b8c..9a573241 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ClimaTimeSteppers" uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79" authors = ["Climate Modeling Alliance"] -version = "0.7.11" +version = "0.7.12" [deps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"