Fused velocity advection hud (#286)

This PR introduces more fused stencils in the velocity advection. The time limit is increased for the cscs-ci. Unifor naming convection is introduced for k, cell, edge.
C2SM · Dec 6, 2023 · 0b39c16 · 0b39c16
1 parent 4a4fb5a
commit 0b39c16
Show file tree

Hide file tree

Showing 155 changed files with 2,854 additions and 460 deletions.
diff --git a/ci/cscs.yml b/ci/cscs.yml
@@ -29,7 +29,7 @@ variables:
   variables:
     SLURM_JOB_NUM_NODES: 2
     SLURM_NTASKS: 2
-    SLURM_TIMELIMIT: '01:00:00'
+    SLURM_TIMELIMIT: '02:00:00'
 
 build_job:
   extends: .build_template

diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py
@@ -23,7 +23,6 @@ def _btraj_dreg_stencil_01(
     p_vn: Field[[EdgeDim, KDim], float],
     tangent_orientation: Field[[EdgeDim], float],
 ) -> Field[[EdgeDim, KDim], bool]:
-
     tangent_orientation = broadcast(tangent_orientation, (EdgeDim, KDim))
     lvn_sys_pos_true = where(p_vn * tangent_orientation >= 0.0, True, False)
     mask_lcounterclock = broadcast(lcounterclock, (EdgeDim, KDim))

diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py
@@ -24,7 +24,6 @@ def _btraj_dreg_stencil_02(
     edge_cell_length: Field[[ECDim], float],
     p_dt: float,
 ) -> Field[[EdgeDim, KDim], int32]:
-
     lvn_pos = where(p_vn >= 0.0, True, False)
     traj_length = sqrt(p_vn * p_vn + p_vt * p_vt) * p_dt
     e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1]))

diff --git a/...here/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/...here/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py
@@ -33,7 +33,6 @@ def ccw(
     p2_lon: Field[[EdgeDim, KDim], float],
     p2_lat: Field[[EdgeDim, KDim], float],
 ) -> Field[[EdgeDim, KDim], int32]:
-
     dx1 = p1_lon - p0_lon
     dy1 = p1_lat - p0_lat
 
@@ -60,7 +59,6 @@ def lintersect(
     line2_p2_lon: Field[[EdgeDim, KDim], float],
     line2_p2_lat: Field[[EdgeDim, KDim], float],
 ) -> Field[[EdgeDim, KDim], bool]:
-
     intersect1 = ccw(
         line1_p1_lon,
         line1_p1_lat,
@@ -108,7 +106,6 @@ def line_intersect(
     line2_p2_lon: Field[[EdgeDim, KDim], float],
     line2_p2_lat: Field[[EdgeDim, KDim], float],
 ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]:
-
     m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon)
     m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon)
 
@@ -159,7 +156,6 @@ def _divide_flux_area_list_stencil_01(
     Field[[EdgeDim, KDim], float],
     Field[[EdgeDim, KDim], float],
 ]:
-
     arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl
     arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl
     arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl

diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py
@@ -22,7 +22,6 @@ def _face_val_ppm_stencil_01a(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1]))
     zfac = (p_cc(Koff[+1]) - p_cc) / (p_cellhgt_mc_now(Koff[+1]) + p_cellhgt_mc_now)
     z_slope = (
@@ -41,7 +40,6 @@ def _face_val_ppm_stencil_01b(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1]))
     zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now)
     z_slope = (
@@ -58,14 +56,13 @@ def _face_val_ppm_stencil_01b(
 def _face_val_ppm_stencil_01(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
-    vert_idx: Field[[KDim], int32],
+    k: Field[[KDim], int32],
     elev: int32,
 ) -> Field[[CellDim, KDim], float]:
-
-    vert_idx = broadcast(vert_idx, (CellDim, KDim))
+    k = broadcast(k, (CellDim, KDim))
 
     z_slope = where(
-        vert_idx == elev,
+        k == elev,
         _face_val_ppm_stencil_01b(p_cc, p_cellhgt_mc_now),
         _face_val_ppm_stencil_01a(p_cc, p_cellhgt_mc_now),
     )
@@ -77,14 +74,14 @@ def _face_val_ppm_stencil_01(
 def face_val_ppm_stencil_01(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
-    vert_idx: Field[[KDim], int32],
+    k: Field[[KDim], int32],
     elev: int32,
     z_slope: Field[[CellDim, KDim], float],
 ):
     _face_val_ppm_stencil_01(
         p_cc,
         p_cellhgt_mc_now,
-        vert_idx,
+        k,
         elev,
         out=z_slope,
     )
diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py
@@ -22,7 +22,6 @@ def _face_val_ppm_stencil_02a(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (
         p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)
     ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1]))
@@ -34,7 +33,6 @@ def _face_val_ppm_stencil_02a(
 def _face_val_ppm_stencil_02b(
     p_cc: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     p_face = p_cc
     return p_face
 
@@ -43,7 +41,6 @@ def _face_val_ppm_stencil_02b(
 def _face_val_ppm_stencil_02c(
     p_cc: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     p_face = p_cc(Koff[-1])
     return p_face
 
@@ -53,24 +50,23 @@ def _face_val_ppm_stencil_02(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
     p_face_in: Field[[CellDim, KDim], float],
-    vert_idx: Field[[KDim], int32],
+    k: Field[[KDim], int32],
     slev: int32,
     elev: int32,
     slevp1: int32,
     elevp1: int32,
 ) -> Field[[CellDim, KDim], float]:
-
-    vert_idx = broadcast(vert_idx, (CellDim, KDim))
+    k = broadcast(k, (CellDim, KDim))
 
     p_face = where(
-        (vert_idx == slevp1) | (vert_idx == elev),
+        (k == slevp1) | (k == elev),
         _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now),
         p_face_in,
     )
 
-    p_face = where((vert_idx == slev), _face_val_ppm_stencil_02b(p_cc), p_face)
+    p_face = where((k == slev), _face_val_ppm_stencil_02b(p_cc), p_face)
 
-    p_face = where((vert_idx == elevp1), _face_val_ppm_stencil_02c(p_cc), p_face)
+    p_face = where((k == elevp1), _face_val_ppm_stencil_02c(p_cc), p_face)
 
     return p_face
 
@@ -80,7 +76,7 @@ def face_val_ppm_stencil_02(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
     p_face_in: Field[[CellDim, KDim], float],
-    vert_idx: Field[[KDim], int32],
+    k: Field[[KDim], int32],
     slev: int32,
     elev: int32,
     slevp1: int32,
@@ -91,7 +87,7 @@ def face_val_ppm_stencil_02(
         p_cc,
         p_cellhgt_mc_now,
         p_face_in,
-        vert_idx,
+        k,
         slev,
         elev,
         slevp1,

diff --git a/...l/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py b/...l/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py
@@ -22,7 +22,6 @@ def _face_val_ppm_stencil_02a(
     p_cc: Field[[CellDim, KDim], float],
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (
         p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)
     ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1]))

diff --git a/...l/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py b/...l/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py
@@ -21,7 +21,6 @@
 def _face_val_ppm_stencil_02b(
     p_cc: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     p_face = p_cc
     return p_face
 

diff --git a/...l/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py b/...l/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py
@@ -21,7 +21,6 @@
 def _face_val_ppm_stencil_02c(
     p_cc: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     p_face = p_cc(Koff[-1])
     return p_face
 

diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py
@@ -23,7 +23,6 @@ def _face_val_ppm_stencil_05(
     p_cellhgt_mc_now: Field[[CellDim, KDim], float],
     z_slope: Field[[CellDim, KDim], float],
 ) -> Field[[CellDim, KDim], float]:
-
     zgeo1 = p_cellhgt_mc_now(Koff[-1]) / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)
     zgeo2 = 1.0 / (
         p_cellhgt_mc_now(Koff[-2])

diff --git a/...osphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py b/...osphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py
@@ -42,7 +42,6 @@ def _hflux_ffsl_hybrid_stencil_01a(
     z_quad_vector_sum0_10: Field[[EdgeDim, KDim], float],
     patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32],
 ) -> Field[[EdgeDim, KDim], float]:
-
     p_out_e_hybrid_1a = (
         where(
             patch0_cell_rel_idx_dsl == int32(1),

diff --git a/...mosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py b/...mosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py
@@ -23,7 +23,6 @@ def _hflux_ffsl_hybrid_stencil_02(
     p_mass_flx_e: Field[[EdgeDim, KDim], float],
     z_dreg_area: Field[[EdgeDim, KDim], float],
 ) -> Field[[EdgeDim, KDim], float]:
-
     p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area
 
     return p_out_e_hybrid_2

diff --git a/...atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py b/...atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py
@@ -53,7 +53,6 @@ def _hflx_limiter_mo_stencil_03a(
     z_min: Field[[CellDim, KDim], float],
     dbl_eps: float,
 ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
-
     r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps)
     r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps)
 
@@ -71,7 +70,6 @@ def _hflx_limiter_mo_stencil_03(
     z_tracer_new_low: Field[[CellDim, KDim], float],
     dbl_eps: float,
 ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
-
     z_max, z_min = _hflx_limiter_mo_stencil_03_min_max(
         z_tracer_max, z_tracer_min, beta_fct, r_beta_fct
     )
@@ -100,7 +98,6 @@ def hflx_limiter_mo_stencil_03(
     r_p: Field[[CellDim, KDim], float],
     r_m: Field[[CellDim, KDim], float],
 ):
-
     _hflx_limiter_mo_stencil_03(
         z_tracer_max,
         z_tracer_min,

diff --git a/...atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py b/...atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py
@@ -25,7 +25,6 @@ def _hflx_limiter_mo_stencil_05(
     r_m: Field[[CellDim, KDim], float],
     r_p: Field[[CellDim, KDim], float],
 ) -> Field[[EdgeDim, KDim], float]:
-
     z_signum = where((z_anti > 0.0), 1.0, -1.0)
 
     r_frac = 0.5 * (

diff --git a/.../advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py b/.../advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py
@@ -72,7 +72,6 @@ def _prep_gauss_quadrature_c_list_stencil(
     Field[[EdgeDim, KDim], float],
     Field[[EdgeDim, KDim], float],
 ]:
-
     z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1
     z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2
     z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1

diff --git a/...mosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/...mosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py
@@ -50,7 +50,6 @@ def _recon_lsq_cell_c_svd_stencil(
     Field[[CellDim, KDim], float],
     Field[[CellDim, KDim], float],
 ]:
-
     p_coeff_10_dsl = (
         lsq_pseudoinv_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc)
         + lsq_pseudoinv_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc)

diff --git a/...sphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py b/...sphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py
@@ -44,7 +44,6 @@ def _upwind_hflux_miura3_stencil_01(
     p_mass_flx_e: Field[[EdgeDim, KDim], float],
     cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32],
 ) -> Field[[EdgeDim, KDim], float]:
-
     p_out_e_miura3 = (
         (
             where(

diff --git a/...re/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py b/...re/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py
@@ -28,7 +28,6 @@ def _upwind_hflux_miura_cycl_stencil_01(
     p_mass_flx_e: Field[[EdgeDim, KDim], float],
     cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32],
 ) -> Field[[EdgeDim, KDim], float]:
-
     z_tracer_mflx_dsl = (
         where(
             cell_rel_idx_dsl == int32(1),

diff --git a/...re/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/...re/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py
@@ -33,7 +33,6 @@ def _upwind_hflux_miura_cycl_stencil_02(
     Field[[CellDim, KDim], float],
     Field[[CellDim, KDim], float],
 ]:
-
     z_rhofluxdiv_c_out = (
         neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim)
         if nsub == int32(1)

diff --git a/...osphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/...osphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py
@@ -28,7 +28,6 @@ def _upwind_hflux_miura_stencil_01(
     p_mass_flx_e: Field[[EdgeDim, KDim], float],
     cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32],
 ) -> Field[[EdgeDim, KDim], float]:
-
     p_out_e = (
         where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0]))
         + distv_bary_1

diff --git a/...atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py b/...atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py
@@ -26,7 +26,6 @@ def _v_limit_prbl_sm_stencil_02(
     p_face: Field[[CellDim, KDim], float],
     p_cc: Field[[CellDim, KDim], float],
 ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
-
     q_face_up, q_face_low = where(
         l_limit != int32(0),
         where(

diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py
@@ -39,7 +39,6 @@ def ccw(
     p2_lon: np.array,
     p2_lat: np.array,
 ) -> np.array:
-
     dx1 = p1_lon - p0_lon
     dy1 = p1_lat - p0_lat
 
@@ -65,7 +64,6 @@ def lintersect(
     line2_p2_lon: np.array,
     line2_p2_lat: np.array,
 ) -> np.array:
-
     intersect1 = ccw(
         line1_p1_lon,
         line1_p1_lat,
@@ -112,7 +110,6 @@ def line_intersect(
     line2_p2_lon: np.array,
     line2_p2_lat: np.array,
 ) -> tuple[np.array]:
-
     m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon)
     m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon)
 

diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py
@@ -24,10 +24,10 @@
 def face_val_ppm_stencil_01_numpy(
     p_cc: np.array,
     p_cellhgt_mc_now: np.array,
-    vert_idx: np.array,
+    k: np.array,
     elev: int32,
 ):
-    # this is a comment: vert_idx = np.broadcast_to(vert_idx, p_cc.shape)
+    # this is a comment: k = np.broadcast_to(k, p_cc.shape)
 
     # 01a
     zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / (
@@ -55,7 +55,7 @@ def face_val_ppm_stencil_01_numpy(
         + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1
     )
 
-    z_slope = np.where(vert_idx[1:-1] < elev, z_slope_a, z_slope_b)
+    z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b)
 
     return z_slope
 
@@ -65,22 +65,22 @@ def test_face_val_ppm_stencil_01(backend):
     p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1})
     p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1})
 
-    vert_idx = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32))
-    elev = vert_idx[-2]
+    k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32))
+    elev = k[-2]
 
     z_slope = random_field(grid, CellDim, KDim)
 
     ref = face_val_ppm_stencil_01_numpy(
         p_cc.asnumpy(),
         p_cellhgt_mc_now.asnumpy(),
-        vert_idx.asnumpy(),
+        k.asnumpy(),
         elev,
     )
 
     face_val_ppm_stencil_01.with_backend(backend)(
         p_cc,
         p_cellhgt_mc_now,
-        vert_idx,
+        k,
         elev,
         z_slope,
         offset_provider={"Koff": KDim},