Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Manually fuse some broadcast expressions in diagnostic edmf #3366

Merged
merged 1 commit into from
Oct 8, 2024
Merged

Manually fuse some broadcast expressions in diagnostic edmf #3366

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 58 additions & 54 deletions src/cache/diagnostic_edmf_precomputed_quantities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,51 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
return nothing
end

function compute_u³ʲ_u³ʲ(
u³ʲ_prev_halflevel,
J_prev_halflevel,
J_halflevel,
J_prev_level,
∇Φ³_data_prev_level,
ρʲ_prev_level,
ρ_prev_level,
entrʲ_prev_level,
turb_entrʲ_prev_level,
u³⁰_data_prev_halflevel,
nh_pressure³ʲ_data_prev_halflevel,
)
u³ʲ_u³ʲ =
(1 / (J_halflevel^2)) *
(J_prev_halflevel^2 * u³ʲ_prev_halflevel * u³ʲ_prev_halflevel)

u³ʲ_u³ʲ -=
(1 / (J_halflevel^2)) * (
J_prev_level^2 *
2 *
(
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) /
ρʲ_prev_level
)
)

u³ʲ_u³ʲ +=
(1 / (J_halflevel^2)) * (
J_prev_level^2 *
2 *
(
(entrʲ_prev_level + turb_entrʲ_prev_level) *
u³⁰_data_prev_halflevel -
(entrʲ_prev_level + turb_entrʲ_prev_level) * u³ʲ_prev_halflevel
)
)

u³ʲ_u³ʲ -=
(1 / (J_halflevel^2)) *
(J_prev_level^2 * 2 * nh_pressure³ʲ_data_prev_halflevel)
return u³ʲ_u³ʲ
end


NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
Y,
p,
Expand Down Expand Up @@ -490,60 +535,19 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
end

u³ʲ_datau³ʲ_data = p.scratch.temp_data_level
# Using constant exponents in broadcasts allocate, so we use
# local_geometry_halflevel.J * local_geometry_halflevel.J instead.
# See ClimaCore.jl issue #1126.
@. u³ʲ_datau³ʲ_data =
(
1 /
(local_geometry_halflevel.J * local_geometry_halflevel.J)
) * (
local_geometry_prev_halflevel.J *
local_geometry_prev_halflevel.J *
u³ʲ_data_prev_halflevel *
u³ʲ_data_prev_halflevel
)

@. u³ʲ_datau³ʲ_data -=
(
1 /
(local_geometry_halflevel.J * local_geometry_halflevel.J)
) * (
local_geometry_prev_level.J *
local_geometry_prev_level.J *
2 *
(
∇Φ³_data_prev_level * (ρʲ_prev_level - ρ_prev_level) /
ρʲ_prev_level
)
)

@. u³ʲ_datau³ʲ_data +=
(
1 /
(local_geometry_halflevel.J * local_geometry_halflevel.J)
) * (
local_geometry_prev_level.J *
local_geometry_prev_level.J *
2 *
(
(entrʲ_prev_level + turb_entrʲ_prev_level) *
u³⁰_data_prev_halflevel -
(entrʲ_prev_level + turb_entrʲ_prev_level) *
u³ʲ_data_prev_halflevel
)
)

@. u³ʲ_datau³ʲ_data -=
(
1 /
(local_geometry_halflevel.J * local_geometry_halflevel.J)
) * (
local_geometry_prev_level.J *
local_geometry_prev_level.J *
2 *
nh_pressure³ʲ_data_prev_halflevel
)
@. u³ʲ_datau³ʲ_data = compute_u³ʲ_u³ʲ(
u³ʲ_data_prev_halflevel,
local_geometry_prev_halflevel.J,
local_geometry_halflevel.J,
local_geometry_prev_level.J,
∇Φ³_data_prev_level,
ρʲ_prev_level,
ρ_prev_level,
entrʲ_prev_level,
turb_entrʲ_prev_level,
u³⁰_data_prev_halflevel,
nh_pressure³ʲ_data_prev_halflevel,
)

# get u³ʲ to calculate divergence term for detrainment,
# u³ʲ will be clipped later after we get area fraction
Expand Down
Loading