Update with main

ci/dace.yml

@@ -22,12 +22,11 @@ test_model_stencils_aarch64:
   stage: test
   script:
     - pip install dace==$DACE_VERSION
-    - tox -r -e run_model_tests -c model/ -- --backend=$BACKEND $DACE_ORCHESTRATION $COMPONENT  --verbose
+    - tox -r -e run_model_tests -c model/ -- --backend=$BACKEND $COMPONENT  --verbose
   parallel:
     matrix:
     - COMPONENT: [atmosphere/diffusion/tests/diffusion_tests]
       BACKEND: [dace_cpu_noopt]
-      DACE_ORCHESTRATION: ['--dace-orchestration=True', '']
 test_model_datatests_x86_64:
   extends: [.test_model_datatests, .test_template_x86_64]
 test_model_datatests_aarch64:

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py

@@ -25,16 +25,23 @@ def _sum_neighbor_contributions(
     js: fa.CellKField[ta.wpfloat],
     p_cc: fa.CellKField[ta.wpfloat],
 ) -> fa.CellKField[ta.wpfloat]:
-    p_cc_p0 = where(mask1 & (js == 0.0), p_cc, 0.0)
-    p_cc_p1 = where(mask1 & (js == 1.0), p_cc(Koff[1]), 0.0)
-    p_cc_p2 = where(mask1 & (js == 2.0), p_cc(Koff[2]), 0.0)
-    p_cc_p3 = where(mask1 & (js == 3.0), p_cc(Koff[3]), 0.0)
-    p_cc_p4 = where(mask1 & (js == 4.0), p_cc(Koff[4]), 0.0)
-    p_cc_m0 = where(mask2 & (js == 0.0), p_cc(Koff[-1]), 0.0)
-    p_cc_m1 = where(mask2 & (js == 1.0), p_cc(Koff[-2]), 0.0)
-    p_cc_m2 = where(mask2 & (js == 2.0), p_cc(Koff[-3]), 0.0)
-    p_cc_m3 = where(mask2 & (js == 3.0), p_cc(Koff[-4]), 0.0)
-    p_cc_m4 = where(mask2 & (js == 4.0), p_cc(Koff[-5]), 0.0)
+    js_eq0 = js == 0.0
+    js_eq1 = js == 1.0
+    js_eq2 = js == 2.0
+    js_eq3 = js == 3.0
+    js_eq4 = js == 4.0
+
+    p_cc_p0 = where(mask1 & js_eq0, p_cc, 0.0)
+    p_cc_p1 = where(mask1 & js_eq1, p_cc(Koff[1]), 0.0)
+    p_cc_p2 = where(mask1 & js_eq2, p_cc(Koff[2]), 0.0)
+    p_cc_p3 = where(mask1 & js_eq3, p_cc(Koff[3]), 0.0)
+    p_cc_p4 = where(mask1 & js_eq4, p_cc(Koff[4]), 0.0)
+    p_cc_m0 = where(mask2 & js_eq0, p_cc(Koff[-1]), 0.0)
+    p_cc_m1 = where(mask2 & js_eq1, p_cc(Koff[-2]), 0.0)
+    p_cc_m2 = where(mask2 & js_eq2, p_cc(Koff[-3]), 0.0)
+    p_cc_m3 = where(mask2 & js_eq3, p_cc(Koff[-4]), 0.0)
+    p_cc_m4 = where(mask2 & js_eq4, p_cc(Koff[-5]), 0.0)
+
     p_cc_jks = (
         p_cc_p0
         + p_cc_p1
@@ -63,17 +70,21 @@ def _compute_ppm4gpu_fractional_flux(
 ) -> fa.CellKField[ta.wpfloat]:
     js = floor(abs(z_cfl))
     z_cflfrac = abs(z_cfl) - js
+    z_cflfrac_nonzero = z_cflfrac != 0.0
 
     z_cfl_pos = z_cfl > 0.0
-    z_cfl_neg = not z_cfl_pos
+    z_cfl_neg = z_cfl < 0.0
     wsign = where(z_cfl_pos, 1.0, -1.0)
 
+    mask1 = z_cfl_pos & z_cflfrac_nonzero
+    mask2 = z_cfl_neg & z_cflfrac_nonzero
+
     in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat)
 
-    p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc)
-    p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now)
-    z_delta_q_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_delta_q)
-    z_a1_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_a1)
+    p_cc_jks = _sum_neighbor_contributions(mask1, mask2, js, p_cc)
+    p_cellmass_now_jks = _sum_neighbor_contributions(mask1, mask2, js, p_cellmass_now)
+    z_delta_q_jks = _sum_neighbor_contributions(mask1, mask2, js, z_delta_q)
+    z_a1_jks = _sum_neighbor_contributions(mask1, mask2, js, z_a1)
 
     z_q_int = (
         p_cc_jks

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py

@@ -9,17 +9,55 @@
 import gt4py.next as gtx
 from gt4py.next.ffront.fbuiltins import abs, astype, floor, where
 
-from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import (
-    _sum_neighbor_contributions,
-)
 from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common.dimension import Koff
 from icon4py.model.common.type_alias import wpfloat
 
 
 # TODO (dastrm): this stencil has no test
 # TODO (dastrm): this stencil does not strictly match the fortran code
 
 
+@gtx.field_operator
+def _sum_neighbor_contributions_all(
+    mask1: fa.CellKField[bool],
+    mask2: fa.CellKField[bool],
+    js: fa.CellKField[ta.wpfloat],
+    p_cc: fa.CellKField[ta.wpfloat],
+    p_cellmass_now: fa.CellKField[ta.wpfloat],
+) -> fa.CellKField[ta.wpfloat]:
+    js_gt0 = js >= 0.0
+    js_gt1 = js >= 1.0
+    js_gt2 = js >= 2.0
+    js_gt3 = js >= 3.0
+    js_gt4 = js >= 4.0
+
+    prod_p0 = where(mask1 & js_gt0, p_cc * p_cellmass_now, 0.0)
+    prod_p1 = where(mask1 & js_gt1, p_cc(Koff[1]) * p_cellmass_now(Koff[1]), 0.0)
+    prod_p2 = where(mask1 & js_gt2, p_cc(Koff[2]) * p_cellmass_now(Koff[2]), 0.0)
+    prod_p3 = where(mask1 & js_gt3, p_cc(Koff[3]) * p_cellmass_now(Koff[3]), 0.0)
+    prod_p4 = where(mask1 & js_gt4, p_cc(Koff[4]) * p_cellmass_now(Koff[4]), 0.0)
+    prod_m0 = where(mask2 & js_gt0, p_cc(Koff[-1]) * p_cellmass_now(Koff[-1]), 0.0)
+    prod_m1 = where(mask2 & js_gt1, p_cc(Koff[-2]) * p_cellmass_now(Koff[-2]), 0.0)
+    prod_m2 = where(mask2 & js_gt2, p_cc(Koff[-3]) * p_cellmass_now(Koff[-3]), 0.0)
+    prod_m3 = where(mask2 & js_gt3, p_cc(Koff[-4]) * p_cellmass_now(Koff[-4]), 0.0)
+    prod_m4 = where(mask2 & js_gt4, p_cc(Koff[-5]) * p_cellmass_now(Koff[-5]), 0.0)
+
+    prod_jks = (
+        prod_p0
+        + prod_p1
+        + prod_p2
+        + prod_p3
+        + prod_p4
+        + prod_m0
+        + prod_m1
+        + prod_m2
+        + prod_m3
+        + prod_m4
+    )
+    return prod_jks
+
+
 @gtx.field_operator
 def _compute_ppm4gpu_integer_flux(
     p_cc: fa.CellKField[ta.wpfloat],
@@ -33,15 +71,16 @@ def _compute_ppm4gpu_integer_flux(
     js = floor(abs(z_cfl)) - 1.0
 
     z_cfl_pos = z_cfl > 0.0
-    z_cfl_neg = not z_cfl_pos
+    z_cfl_neg = z_cfl < 0.0
     wsign = where(z_cfl_pos, 1.0, -1.0)
 
     in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat)
 
-    p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc)
-    p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now)
+    p_cc_cellmass_now_jks = _sum_neighbor_contributions_all(
+        z_cfl_pos, z_cfl_neg, js, p_cc, p_cellmass_now
+    )
 
-    z_iflx = wsign * p_cc_jks * p_cellmass_now_jks
+    z_iflx = wsign * p_cc_cellmass_now_jks
 
     p_upflux = p_upflux + where(in_slev_bounds, z_iflx / p_dtime, 0.0)
 

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py

@@ -13,7 +13,6 @@
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import KDim, VertexDim
 from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift
-from icon4py.model.common.settings import backend
 
 
 @gtx.field_operator
@@ -46,7 +45,7 @@ def _init_zero_v_k() -> gtx.Field[[dims.VertexDim, dims.KDim], float]:
     return broadcast(0.0, (VertexDim, KDim))
 
 
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED, backend=backend)
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def init_zero_v_k(field: gtx.Field[[dims.VertexDim, dims.KDim], float]):
     _init_zero_v_k(out=field)
 

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py

@@ -22,7 +22,6 @@
     _update_theta_and_exner,
 )
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -58,7 +57,7 @@ def _apply_diffusion_to_theta_and_exner(
     return theta_v, exner
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def apply_diffusion_to_theta_and_exner(
     kh_smag_e: fa.EdgeKField[vpfloat],
     inv_dual_edge_length: fa.EdgeField[wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py

@@ -11,7 +11,6 @@
 from gt4py.next.ffront.fbuiltins import astype
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -32,7 +31,7 @@ def _apply_nabla2_and_nabla4_global_to_vn(
     return vn_wp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def apply_nabla2_and_nabla4_global_to_vn(
     area_edge: fa.EdgeField[wpfloat],
     kh_smag_e: fa.EdgeKField[vpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py

@@ -11,7 +11,6 @@
 from gt4py.next.ffront.fbuiltins import astype, broadcast, maximum
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -38,7 +37,7 @@ def _apply_nabla2_and_nabla4_to_vn(
     return vn_wp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def apply_nabla2_and_nabla4_to_vn(
     area_edge: fa.EdgeField[wpfloat],
     kh_smag_e: fa.EdgeKField[vpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py

@@ -10,7 +10,6 @@
 from gt4py.next.ffront.decorator import field_operator, program
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import wpfloat
 
 
@@ -25,7 +24,7 @@ def _apply_nabla2_to_vn_in_lateral_boundary(
     return vn_wp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def apply_nabla2_to_vn_in_lateral_boundary(
     z_nabla2_e: fa.EdgeKField[wpfloat],
     area_edge: fa.EdgeField[wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import C2E2CO, C2E2CODim
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -32,7 +31,7 @@ def _apply_nabla2_to_w(
     return w_wp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def apply_nabla2_to_w(
     area: fa.CellField[wpfloat],
     z_nabla2_c: fa.CellKField[vpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py

@@ -11,7 +11,6 @@
 from gt4py.next.ffront.fbuiltins import astype, broadcast
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -29,7 +28,7 @@ def _apply_nabla2_to_w_in_upper_damping_layer(
     return w_wp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def apply_nabla2_to_w_in_upper_damping_layer(
     w: fa.CellKField[wpfloat],
     diff_multfac_n2w: fa.KField[wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import field_type_aliases as fa
 from icon4py.model.common.dimension import Koff
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -33,7 +32,7 @@ def _calculate_diagnostics_for_turbulence(
     return astype((div_ic_wp, hdef_ic_wp), vpfloat)
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def calculate_diagnostics_for_turbulence(
     div: fa.CellKField[vpfloat],
     kh_c: fa.CellKField[vpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import C2E2CO, C2E2CODim
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -27,7 +26,7 @@ def _calculate_horizontal_gradients_for_turbulence(
     return astype((dwdx_wp, dwdy_wp), vpfloat)
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def calculate_horizontal_gradients_for_turbulence(
     w: fa.CellKField[wpfloat],
     geofac_grg_x: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import C2E2CO, C2E2CODim
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -24,7 +23,7 @@ def _calculate_nabla2_for_w(
     return astype(z_nabla2_c_wp, vpfloat)
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def calculate_nabla2_for_w(
     w: fa.CellKField[wpfloat],
     geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import E2C
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -28,7 +27,7 @@ def _calculate_nabla2_for_z(
     return z_nabla2_e_wp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def calculate_nabla2_for_z(
     kh_smag_e: fa.EdgeKField[vpfloat],
     inv_dual_edge_length: fa.EdgeField[wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import C2CE, C2E, C2EDim
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -25,7 +24,7 @@ def _calculate_nabla2_of_theta(
     return astype(z_temp_wp, vpfloat)
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def calculate_nabla2_of_theta(
     z_nabla2_e: fa.EdgeKField[wpfloat],
     geofac_div: gtx.Field[gtx.Dims[dims.CEDim], wpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
@@ -56,7 +55,7 @@ def _calculate_nabla4(
     return astype(z_nabla4_e2_wp, vpfloat)
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def calculate_nabla4(
     u_vert: fa.VertexKField[vpfloat],
     v_vert: fa.VertexKField[vpfloat],

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py

@@ -12,7 +12,6 @@
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import E2C, E2CDim
-from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import vpfloat
 
 
@@ -25,7 +24,7 @@ def _enhance_diffusion_coefficient_for_grid_point_cold_pools(
     return kh_smag_e_vp
 
 
-@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+@program(grid_type=GridType.UNSTRUCTURED)
 def enhance_diffusion_coefficient_for_grid_point_cold_pools(
     kh_smag_e: fa.EdgeKField[vpfloat],
     enh_diffu_3d: fa.CellKField[vpfloat],