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

Dycore in mixed-precision #301

Merged
merged 20 commits into from
Nov 27, 2023
Merged

Dycore in mixed-precision #301

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
a4aaff1
Add '--enable-mixed-precision' option in CLI to compile stencils in m…
Oct 16, 2023
6696bc7
Add Type Alias (wpfloat, vpfloat) definition to be imported in stencils
Oct 16, 2023
a7d925e
Adapt types and casting in dycore stencils for mixed-precision
Oct 16, 2023
20613d9
Add pytest '--enable-mixed-precision' to test in mixed-mode
Oct 20, 2023
8ab5341
Add dtype argument to random_field
Oct 20, 2023
2678197
Adapt unit tests to test in mixed-precision mode
Oct 20, 2023
331b7d7
Run pre-commit
Oct 20, 2023
40dfd7b
Test --enable-mixed-precision option in icon4pygen
Oct 30, 2023
0cec8af
Merge remote-tracking branch 'origin/main' into mixed_precision_dycore
Oct 30, 2023
6765232
Merge remote-tracking branch 'origin/main' into mixed_precision_dycore
Nov 3, 2023
7ac0f83
Address requested changes
Nov 3, 2023
404115c
Merge remote-tracking branch 'origin/main' into mixed_precision_dycore
Nov 13, 2023
c4e0de9
Merge remote-tracking branch 'origin/main' into mixed_precision_dycore
Nov 13, 2023
270a445
pre-commit and fixes
Nov 13, 2023
5cfcbcb
Merge branch 'mixed_precision_dycore' of github.com:C2SM/icon4py into…
Nov 13, 2023
1ca0f29
Merge branch 'main' of github.com:C2SM/icon4py into mixed_precision_d…
Nov 17, 2023
144d1ae
Merge branch 'main' of github.com:C2SM/icon4py into mixed_precision_d…
Nov 17, 2023
3196ed4
Fix order of tests in test_codegen
Nov 20, 2023
2809583
Fix test_mixed_precision with reload fct
Nov 20, 2023
364c2a6
Merge branch 'main' of github.com:C2SM/icon4py into mixed_precision_d…
Nov 24, 2023
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
49 changes: 25 additions & 24 deletions ...dvection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,29 +13,30 @@

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, int32, where
from gt4py.next.ffront.fbuiltins import Field, astype, int32, where

from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim
from icon4py.model.common.type_alias import vpfloat, wpfloat


@field_operator
def _mo_advection_traj_btraj_compute_o1_dsl(
p_vn: Field[[EdgeDim, KDim], float],
p_vt: Field[[EdgeDim, KDim], float],
p_vn: Field[[EdgeDim, KDim], wpfloat],
p_vt: Field[[EdgeDim, KDim], wpfloat],
cell_idx: Field[[ECDim], int32],
cell_blk: Field[[ECDim], int32],
pos_on_tplane_e_1: Field[[ECDim], float],
pos_on_tplane_e_2: Field[[ECDim], float],
primal_normal_cell_1: Field[[ECDim], float],
dual_normal_cell_1: Field[[ECDim], float],
primal_normal_cell_2: Field[[ECDim], float],
dual_normal_cell_2: Field[[ECDim], float],
p_dthalf: float,
pos_on_tplane_e_1: Field[[ECDim], wpfloat],
pos_on_tplane_e_2: Field[[ECDim], wpfloat],
primal_normal_cell_1: Field[[ECDim], wpfloat],
dual_normal_cell_1: Field[[ECDim], wpfloat],
primal_normal_cell_2: Field[[ECDim], wpfloat],
dual_normal_cell_2: Field[[ECDim], wpfloat],
p_dthalf: wpfloat,
) -> tuple[
Field[[EdgeDim, KDim], int32],
Field[[EdgeDim, KDim], int32],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], vpfloat],
Field[[EdgeDim, KDim], vpfloat],
]:
lvn_pos = where(p_vn > 0.0, True, False)

Expand Down Expand Up @@ -66,26 +67,26 @@ def _mo_advection_traj_btraj_compute_o1_dsl(
+ z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]),
)

return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2
return p_cell_idx, p_cell_blk, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat)


@program(grid_type=GridType.UNSTRUCTURED)
def mo_advection_traj_btraj_compute_o1_dsl(
p_vn: Field[[EdgeDim, KDim], float],
p_vt: Field[[EdgeDim, KDim], float],
p_vn: Field[[EdgeDim, KDim], wpfloat],
p_vt: Field[[EdgeDim, KDim], wpfloat],
cell_idx: Field[[ECDim], int32],
cell_blk: Field[[ECDim], int32],
pos_on_tplane_e_1: Field[[ECDim], float],
pos_on_tplane_e_2: Field[[ECDim], float],
primal_normal_cell_1: Field[[ECDim], float],
dual_normal_cell_1: Field[[ECDim], float],
primal_normal_cell_2: Field[[ECDim], float],
dual_normal_cell_2: Field[[ECDim], float],
pos_on_tplane_e_1: Field[[ECDim], wpfloat],
pos_on_tplane_e_2: Field[[ECDim], wpfloat],
primal_normal_cell_1: Field[[ECDim], wpfloat],
dual_normal_cell_1: Field[[ECDim], wpfloat],
primal_normal_cell_2: Field[[ECDim], wpfloat],
dual_normal_cell_2: Field[[ECDim], wpfloat],
p_cell_idx: Field[[EdgeDim, KDim], int32],
p_cell_blk: Field[[EdgeDim, KDim], int32],
p_distv_bary_1: Field[[EdgeDim, KDim], float],
p_distv_bary_2: Field[[EdgeDim, KDim], float],
p_dthalf: float,
p_distv_bary_1: Field[[EdgeDim, KDim], vpfloat],
p_distv_bary_2: Field[[EdgeDim, KDim], vpfloat],
p_dthalf: wpfloat,
):
_mo_advection_traj_btraj_compute_o1_dsl(
p_vn,
Expand Down
49 changes: 25 additions & 24 deletions ...ion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,24 +27,25 @@
_update_theta_and_exner,
)
from icon4py.model.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim
from icon4py.model.common.type_alias import vpfloat, wpfloat


@field_operator
def _apply_diffusion_to_theta_and_exner(
kh_smag_e: Field[[EdgeDim, KDim], float],
inv_dual_edge_length: Field[[EdgeDim], float],
theta_v_in: Field[[CellDim, KDim], float],
geofac_div: Field[[CEDim], float],
kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
inv_dual_edge_length: Field[[EdgeDim], wpfloat],
theta_v_in: Field[[CellDim, KDim], wpfloat],
geofac_div: Field[[CEDim], wpfloat],
mask: Field[[CellDim, KDim], bool],
zd_vertoffset: Field[[CECDim, KDim], int32],
zd_diffcoef: Field[[CellDim, KDim], float],
geofac_n2s_c: Field[[CellDim], float],
geofac_n2s_nbh: Field[[CECDim], float],
vcoef: Field[[CECDim, KDim], float],
area: Field[[CellDim], float],
exner: Field[[CellDim, KDim], float],
rd_o_cvd: float,
) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
zd_diffcoef: Field[[CellDim, KDim], wpfloat],
geofac_n2s_c: Field[[CellDim], wpfloat],
geofac_n2s_nbh: Field[[CECDim], wpfloat],
vcoef: Field[[CECDim, KDim], wpfloat],
area: Field[[CellDim], wpfloat],
exner: Field[[CellDim, KDim], wpfloat],
rd_o_cvd: vpfloat,
) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v_in)
z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div)
z_temp = _truly_horizontal_diffusion_nabla_of_theta_over_steep_points(
Expand All @@ -63,20 +64,20 @@ def _apply_diffusion_to_theta_and_exner(

@program
def apply_diffusion_to_theta_and_exner(
kh_smag_e: Field[[EdgeDim, KDim], float],
inv_dual_edge_length: Field[[EdgeDim], float],
theta_v_in: Field[[CellDim, KDim], float],
geofac_div: Field[[CEDim], float],
kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
inv_dual_edge_length: Field[[EdgeDim], wpfloat],
theta_v_in: Field[[CellDim, KDim], wpfloat],
geofac_div: Field[[CEDim], wpfloat],
mask: Field[[CellDim, KDim], bool],
zd_vertoffset: Field[[CECDim, KDim], int32],
zd_diffcoef: Field[[CellDim, KDim], float],
geofac_n2s_c: Field[[CellDim], float],
geofac_n2s_nbh: Field[[CECDim], float],
vcoef: Field[[CECDim, KDim], float],
area: Field[[CellDim], float],
theta_v: Field[[CellDim, KDim], float],
exner: Field[[CellDim, KDim], float],
rd_o_cvd: float,
zd_diffcoef: Field[[CellDim, KDim], wpfloat],
geofac_n2s_c: Field[[CellDim], wpfloat],
geofac_n2s_nbh: Field[[CECDim], wpfloat],
vcoef: Field[[CECDim, KDim], wpfloat],
area: Field[[CellDim], wpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
exner: Field[[CellDim, KDim], wpfloat],
rd_o_cvd: vpfloat,
):
_apply_diffusion_to_theta_and_exner(
kh_smag_e,
Expand Down
59 changes: 30 additions & 29 deletions ...sphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,28 +25,29 @@
)
from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4
from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim
from icon4py.model.common.type_alias import vpfloat, wpfloat


@field_operator
def _apply_diffusion_to_vn(
u_vert: Field[[VertexDim, KDim], float],
v_vert: Field[[VertexDim, KDim], float],
primal_normal_vert_v1: Field[[ECVDim], float],
primal_normal_vert_v2: Field[[ECVDim], float],
z_nabla2_e: Field[[EdgeDim, KDim], float],
inv_vert_vert_length: Field[[EdgeDim], float],
inv_primal_edge_length: Field[[EdgeDim], float],
area_edge: Field[[EdgeDim], float],
kh_smag_e: Field[[EdgeDim, KDim], float],
diff_multfac_vn: Field[[KDim], float],
nudgecoeff_e: Field[[EdgeDim], float],
vn: Field[[EdgeDim, KDim], float],
u_vert: Field[[VertexDim, KDim], vpfloat],
v_vert: Field[[VertexDim, KDim], vpfloat],
primal_normal_vert_v1: Field[[ECVDim], wpfloat],
primal_normal_vert_v2: Field[[ECVDim], wpfloat],
z_nabla2_e: Field[[EdgeDim, KDim], wpfloat],
inv_vert_vert_length: Field[[EdgeDim], wpfloat],
inv_primal_edge_length: Field[[EdgeDim], wpfloat],
area_edge: Field[[EdgeDim], wpfloat],
kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
diff_multfac_vn: Field[[KDim], wpfloat],
nudgecoeff_e: Field[[EdgeDim], wpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
horz_idx: Field[[EdgeDim], int32],
nudgezone_diff: float,
fac_bdydiff_v: float,
nudgezone_diff: vpfloat,
fac_bdydiff_v: wpfloat,
start_2nd_nudge_line_idx_e: int32,
limited_area: bool,
) -> Field[[EdgeDim, KDim], float]:
) -> Field[[EdgeDim, KDim], wpfloat]:
z_nabla4_e2 = _calculate_nabla4(
u_vert,
v_vert,
Expand Down Expand Up @@ -93,21 +94,21 @@ def _apply_diffusion_to_vn(

@program(grid_type=GridType.UNSTRUCTURED)
def apply_diffusion_to_vn(
u_vert: Field[[VertexDim, KDim], float],
v_vert: Field[[VertexDim, KDim], float],
primal_normal_vert_v1: Field[[ECVDim], float],
primal_normal_vert_v2: Field[[ECVDim], float],
z_nabla2_e: Field[[EdgeDim, KDim], float],
inv_vert_vert_length: Field[[EdgeDim], float],
inv_primal_edge_length: Field[[EdgeDim], float],
area_edge: Field[[EdgeDim], float],
kh_smag_e: Field[[EdgeDim, KDim], float],
diff_multfac_vn: Field[[KDim], float],
nudgecoeff_e: Field[[EdgeDim], float],
vn: Field[[EdgeDim, KDim], float],
u_vert: Field[[VertexDim, KDim], vpfloat],
v_vert: Field[[VertexDim, KDim], vpfloat],
primal_normal_vert_v1: Field[[ECVDim], wpfloat],
primal_normal_vert_v2: Field[[ECVDim], wpfloat],
z_nabla2_e: Field[[EdgeDim, KDim], wpfloat],
inv_vert_vert_length: Field[[EdgeDim], wpfloat],
inv_primal_edge_length: Field[[EdgeDim], wpfloat],
area_edge: Field[[EdgeDim], wpfloat],
kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
diff_multfac_vn: Field[[KDim], wpfloat],
nudgecoeff_e: Field[[EdgeDim], wpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
horz_idx: Field[[EdgeDim], int32],
nudgezone_diff: float,
fac_bdydiff_v: float,
nudgezone_diff: vpfloat,
fac_bdydiff_v: wpfloat,
start_2nd_nudge_line_idx_e: int32,
limited_area: bool,
horizontal_start: int32,
Expand Down
45 changes: 23 additions & 22 deletions ...iffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,28 +25,29 @@
_calculate_nabla2_for_w,
)
from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim
from icon4py.model.common.type_alias import vpfloat, wpfloat


@field_operator
def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
area: Field[[CellDim], float],
geofac_n2s: Field[[CellDim, C2E2CODim], float],
geofac_grg_x: Field[[CellDim, C2E2CODim], float],
geofac_grg_y: Field[[CellDim, C2E2CODim], float],
w_old: Field[[CellDim, KDim], float],
dwdx: Field[[CellDim, KDim], float],
dwdy: Field[[CellDim, KDim], float],
diff_multfac_w: float,
diff_multfac_n2w: Field[[KDim], float],
area: Field[[CellDim], wpfloat],
geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat],
geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat],
geofac_grg_y: Field[[CellDim, C2E2CODim], wpfloat],
w_old: Field[[CellDim, KDim], wpfloat],
dwdx: Field[[CellDim, KDim], vpfloat],
dwdy: Field[[CellDim, KDim], vpfloat],
diff_multfac_w: wpfloat,
diff_multfac_n2w: Field[[KDim], wpfloat],
vert_idx: Field[[KDim], int32],
horz_idx: Field[[CellDim], int32],
nrdmax: int32,
interior_idx: int32,
halo_idx: int32,
) -> tuple[
Field[[CellDim, KDim], float],
Field[[CellDim, KDim], float],
Field[[CellDim, KDim], float],
Field[[CellDim, KDim], wpfloat],
Field[[CellDim, KDim], vpfloat],
Field[[CellDim, KDim], vpfloat],
]:
vert_idx = broadcast(vert_idx, (CellDim, KDim))

Expand Down Expand Up @@ -78,16 +79,16 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(

@program
def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
area: Field[[CellDim], float],
geofac_n2s: Field[[CellDim, C2E2CODim], float],
geofac_grg_x: Field[[CellDim, C2E2CODim], float],
geofac_grg_y: Field[[CellDim, C2E2CODim], float],
w_old: Field[[CellDim, KDim], float],
w: Field[[CellDim, KDim], float],
dwdx: Field[[CellDim, KDim], float],
dwdy: Field[[CellDim, KDim], float],
diff_multfac_w: float,
diff_multfac_n2w: Field[[KDim], float],
area: Field[[CellDim], wpfloat],
geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat],
geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat],
geofac_grg_y: Field[[CellDim, C2E2CODim], wpfloat],
w_old: Field[[CellDim, KDim], wpfloat],
w: Field[[CellDim, KDim], wpfloat],
dwdx: Field[[CellDim, KDim], vpfloat],
dwdy: Field[[CellDim, KDim], vpfloat],
diff_multfac_w: wpfloat,
diff_multfac_n2w: Field[[KDim], wpfloat],
vert_idx: Field[[KDim], int32],
horz_idx: Field[[CellDim], int32],
nrdmax: int32,
Expand Down
37 changes: 21 additions & 16 deletions ...n/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,32 +13,37 @@

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field
from gt4py.next.ffront.fbuiltins import Field, astype

from icon4py.model.common.dimension import EdgeDim, KDim
from icon4py.model.common.type_alias import vpfloat, wpfloat


@field_operator
def _apply_nabla2_and_nabla4_global_to_vn(
area_edge: Field[[EdgeDim], float],
kh_smag_e: Field[[EdgeDim, KDim], float],
z_nabla2_e: Field[[EdgeDim, KDim], float],
z_nabla4_e2: Field[[EdgeDim, KDim], float],
diff_multfac_vn: Field[[KDim], float],
vn: Field[[EdgeDim, KDim], float],
) -> Field[[EdgeDim, KDim], float]:
vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge)
return vn
area_edge: Field[[EdgeDim], wpfloat],
kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
z_nabla2_e: Field[[EdgeDim, KDim], wpfloat],
z_nabla4_e2: Field[[EdgeDim, KDim], vpfloat],
diff_multfac_vn: Field[[KDim], wpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
) -> Field[[EdgeDim, KDim], wpfloat]:
kh_smag_e_wp, z_nabla4_e2_wp = astype((kh_smag_e, z_nabla4_e2), wpfloat)

vn_wp = vn + area_edge * (
kh_smag_e_wp * z_nabla2_e - diff_multfac_vn * z_nabla4_e2_wp * area_edge
)
return vn_wp


@program(grid_type=GridType.UNSTRUCTURED)
def apply_nabla2_and_nabla4_global_to_vn(
area_edge: Field[[EdgeDim], float],
kh_smag_e: Field[[EdgeDim, KDim], float],
z_nabla2_e: Field[[EdgeDim, KDim], float],
z_nabla4_e2: Field[[EdgeDim, KDim], float],
diff_multfac_vn: Field[[KDim], float],
vn: Field[[EdgeDim, KDim], float],
area_edge: Field[[EdgeDim], wpfloat],
kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
z_nabla2_e: Field[[EdgeDim, KDim], wpfloat],
z_nabla4_e2: Field[[EdgeDim, KDim], vpfloat],
diff_multfac_vn: Field[[KDim], wpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
):
_apply_nabla2_and_nabla4_global_to_vn(
area_edge,
Expand Down
Loading
Loading