Skip to content

Commit

Permalink
Merge pull request #223 from CliMA/ck/rm_nvtx
Browse files Browse the repository at this point in the history 
Add `post_explicit!` call to end of step, rm nvtx range macros
  • Loading branch information
@charleskawczynski
charleskawczynski authored Oct 4, 2023
2 parents efec3c3 + 57c3828 commit a256536
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 132 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
219 changes: 88 additions & 131 deletions src/solvers/imex_ark.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,177 +59,134 @@ 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

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
1 change: 1 addition & 0 deletions src/solvers/imex_ssprk.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

2 comments on commit a256536

@charleskawczynski
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/92762

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.7.12 -m "<description of version>" a2565367aefdc29a3bf7c0ec8b928f03884b942d
git push origin v0.7.12

Please sign in to comment.