resolved conflicts

base-requirements-dev.txt

@@ -17,7 +17,7 @@ pytest-mpi>=0.6
 ruff>=0.2.2
 setuptools>=40.8.0
 tox >= 3.25
-typing-extensions==4.5.0
+typing-extensions>=4.6.0
 types-cffi>=1.15
 wheel>=0.37.1
 wget>=3.2

ci/default.yml

@@ -30,7 +30,8 @@ test_model_datatests:
     extends: .test_template
     stage: test
     script:
-        - tox -r -e run_model_tests -c model/ --verbose -- $COMPONENT
+        - tox -r -e run_model_tests -c model/ --verbose -- --backend=$BACKEND $COMPONENT
     parallel:
         matrix:
             - COMPONENT: [atmosphere/diffusion/tests/diffusion_tests, atmosphere/dycore/tests/dycore_tests, common/tests, driver/tests]
+              BACKEND: [gtfn_cpu]

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

@@ -22,11 +22,6 @@
 from gt4py.next import as_field
 from gt4py.next.common import Dimension
 from gt4py.next.ffront.fbuiltins import Field, int32
-from gt4py.next.program_processors.runners.gtfn import (
-    run_gtfn,
-    run_gtfn_cached,
-    run_gtfn_imperative,
-)
 
 from icon4py.model.atmosphere.diffusion.diffusion_states import (
     DiffusionDiagnosticState,
@@ -78,6 +73,7 @@
     mo_intp_rbf_rbf_vec_interpol_vertex,
 )
 from icon4py.model.common.states.prognostic_state import PrognosticState
+from icon4py.model.common.model_backend import backend
 
 
 """
@@ -89,11 +85,6 @@
 # flake8: noqa
 log = logging.getLogger(__name__)
 
-cached_backend = run_gtfn_cached
-compiled_backend = run_gtfn
-imperative_backend = run_gtfn_imperative
-backend = run_gtfn_cached  #
-
 
 class DiffusionType(int, Enum):
     """
@@ -155,7 +146,6 @@ def __init__(
         zdiffu_t: bool = True,
         thslp_zdiffu: float = 0.025,
         thhgtd_zdiffu: float = 200.0,
-        hdiff_rcf: bool = True,
         velocity_boundary_diffusion_denom: float = 200.0,
         temperature_boundary_diffusion_denom: float = 135.0,
         max_nudging_coeff: float = 0.02,
@@ -224,10 +214,6 @@ def __init__(
         #       DiffusionConfig and init()
         self.ndyn_substeps: int = n_substeps
 
-        #: If True, compute horizontal diffusion only at the large time step
-        #: Called 'lhdiff_rcf' in mo_nonhydrostatic_nml.f90
-        self.lhdiff_rcf: bool = hdiff_rcf
-
         # namelist mo_gridref_nml.f90
 
         #: Denominator for temperature boundary diffusion
@@ -451,14 +437,12 @@ def _get_start_index_for_w_diffusion() -> int32:
         self.bdy_diff: float = 0.015 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon)
         self.fac_bdydiff_v: float = (
             math.sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator
-            if config.lhdiff_rcf
-            else 1.0 / config.velocity_boundary_diffusion_denominator
         )
 
         self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float
         self.diff_multfac_w: float = min(1.0 / 48.0, params.K4W * config.substep_as_float)
 
-        init_diffusion_local_fields_for_regular_timestep.with_backend(backend)(
+        init_diffusion_local_fields_for_regular_timestep(
             params.K4,
             config.substep_as_float,
             *params.smagorinski_factor,
@@ -503,7 +487,6 @@ def _index_field(dim: Dimension, size=None):
         self.z_nabla2_e = _allocate(EdgeDim, KDim)
         self.z_temp = _allocate(CellDim, KDim)
         self.diff_multfac_smag = _allocate(KDim)
-        self.z_nabla4_e2 = _allocate(EdgeDim, KDim)
         # TODO(Magdalena): this is KHalfDim
         self.vertical_index = _index_field(KDim, self.grid.num_levels + 1)
         self.horizontal_cell_index = _index_field(CellDim)
@@ -532,7 +515,7 @@ def initial_run(
         diff_multfac_vn = zero_field(self.grid, KDim)
         smag_limit = zero_field(self.grid, KDim)
 
-        setup_fields_for_initial_step.with_backend(backend)(
+        setup_fields_for_initial_step(
             self.params.K4,
             self.config.hdiff_efdt_ratio,
             diff_multfac_vn,
@@ -569,15 +552,13 @@ def run(
             smag_limit=self.smag_limit,
             smag_offset=self.smag_offset,
         )
-        if not self.config.lhdiff_rcf:
-            self._sync_cell_fields(prognostic_state)
 
     def _sync_cell_fields(self, prognostic_state):
         """
         Communicate theta_v, exner and w.
 
         communication only done in original code if the following condition applies:
-        IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN
+        IF ( linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN
         """
         log.debug("communication of prognostic cell fields: theta, w, exner - start")
         self._exchange.exchange_and_wait(
@@ -642,12 +623,10 @@ def _do_diffusion_step(
         )
 
         # dtime dependent: enh_smag_factor,
-        scale_k.with_backend(backend)(
-            self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}
-        )
+        scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={})
 
         log.debug("rbf interpolation 1: start")
-        mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)(
+        mo_intp_rbf_rbf_vec_interpol_vertex(
             p_e_in=prognostic_state.vn,
             ptr_coeff_1=self.interpolation_state.rbf_coeff_1,
             ptr_coeff_2=self.interpolation_state.rbf_coeff_2,
@@ -657,7 +636,7 @@ def _do_diffusion_step(
             horizontal_end=vertex_end_local,
             vertical_start=0,
             vertical_end=klevels,
-            offset_provider={"V2E": self.grid.get_offset_provider("V2E")},
+            offset_provider=self.grid.offset_providers,
         )
         log.debug("rbf interpolation 1: end")
 
@@ -667,7 +646,7 @@ def _do_diffusion_step(
         log.debug("communication rbf extrapolation of vn - end")
 
         log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start")
-        calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)(
+        calculate_nabla2_and_smag_coefficients_for_vn(
             diff_multfac_smag=self.diff_multfac_smag,
             tangent_orientation=self.edge_params.tangent_orientation,
             inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths,
@@ -688,10 +667,7 @@ def _do_diffusion_step(
             horizontal_end=edge_end_local_minus2,
             vertical_start=0,
             vertical_end=klevels,
-            offset_provider={
-                "E2C2V": self.grid.get_offset_provider("E2C2V"),
-                "E2ECV": self.grid.get_offset_provider("E2ECV"),
-            },
+            offset_provider=self.grid.offset_providers,
         )
         log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end")
         if (
@@ -701,7 +677,7 @@ def _do_diffusion_step(
             log.debug(
                 "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start"
             )
-            calculate_diagnostic_quantities_for_turbulence.with_backend(backend)(
+            calculate_diagnostic_quantities_for_turbulence(
                 kh_smag_ec=self.kh_smag_ec,
                 vn=prognostic_state.vn,
                 e_bln_c_s=self.interpolation_state.e_bln_c_s,
@@ -714,11 +690,7 @@ def _do_diffusion_step(
                 horizontal_end=cell_end_local,
                 vertical_start=1,
                 vertical_end=klevels,
-                offset_provider={
-                    "C2E": self.grid.get_offset_provider("C2E"),
-                    "C2CE": self.grid.get_offset_provider("C2CE"),
-                    "Koff": KDim,
-                },
+                offset_provider=self.grid.offset_providers,
             )
             log.debug(
                 "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end"
@@ -732,7 +704,7 @@ def _do_diffusion_step(
             log.debug("communication rbf extrapolation of z_nable2_e - end")
 
         log.debug("2nd rbf interpolation: start")
-        mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)(
+        mo_intp_rbf_rbf_vec_interpol_vertex(
             p_e_in=self.z_nabla2_e,
             ptr_coeff_1=self.interpolation_state.rbf_coeff_1,
             ptr_coeff_2=self.interpolation_state.rbf_coeff_2,
@@ -742,7 +714,7 @@ def _do_diffusion_step(
             horizontal_end=vertex_end_local,
             vertical_start=0,
             vertical_end=klevels,
-            offset_provider={"V2E": self.grid.get_offset_provider("V2E")},
+            offset_provider=self.grid.offset_providers,
         )
         log.debug("2nd rbf interpolation: end")
 
@@ -752,7 +724,7 @@ def _do_diffusion_step(
         log.debug("communication rbf extrapolation of z_nable2_e - end")
 
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start")
-        apply_diffusion_to_vn.with_backend(backend)(
+        apply_diffusion_to_vn(
             u_vert=self.u_vert,
             v_vert=self.v_vert,
             primal_normal_vert_v1=self.edge_params.primal_normal_vert[0],
@@ -774,10 +746,7 @@ def _do_diffusion_step(
             horizontal_end=edge_end_local,
             vertical_start=0,
             vertical_end=klevels,
-            offset_provider={
-                "E2C2V": self.grid.get_offset_provider("E2C2V"),
-                "E2ECV": self.grid.get_offset_provider("E2ECV"),
-            },
+            offset_provider=self.grid.offset_providers,
         )
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end")
         log.debug("communication of prognistic.vn : start")
@@ -787,8 +756,8 @@ def _do_diffusion_step(
             "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"
         )
         # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify?
-        copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={})
-        apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_backend(backend)(
+        copy_field(prognostic_state.w, self.w_tmp, offset_provider={})
+        apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
             area=self.cell_params.area,
             geofac_n2s=self.interpolation_state.geofac_n2s,
             geofac_grg_x=self.interpolation_state.geofac_grg_x,
@@ -811,9 +780,7 @@ def _do_diffusion_step(
             horizontal_end=cell_end_halo,
             vertical_start=0,
             vertical_end=klevels,
-            offset_provider={
-                "C2E2CO": self.grid.get_offset_provider("C2E2CO"),
-            },
+            offset_provider=self.grid.offset_providers,
         )
         log.debug(
             "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): end"
@@ -822,7 +789,7 @@ def _do_diffusion_step(
         log.debug(
             "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start"
         )
-        calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend(backend)(
+        calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools(
             theta_v=prognostic_state.theta_v,
             theta_ref_mc=self.metric_state.theta_ref_mc,
             thresh_tdiff=self.thresh_tdiff,
@@ -832,16 +799,13 @@ def _do_diffusion_step(
             horizontal_end=edge_end_halo,
             vertical_start=(klevels - 2),
             vertical_end=klevels,
-            offset_provider={
-                "E2C": self.grid.get_offset_provider("E2C"),
-                "C2E2C": self.grid.get_offset_provider("C2E2C"),
-            },
+            offset_provider=self.grid.offset_providers,
         )
         log.debug(
             "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end"
         )
         log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start")
-        calculate_nabla2_for_theta.with_backend(backend)(
+        calculate_nabla2_for_theta(
             kh_smag_e=self.kh_smag_e,
             inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths,
             theta_v=prognostic_state.theta_v,
@@ -851,18 +815,14 @@ def _do_diffusion_step(
             horizontal_end=cell_end_local,
             vertical_start=0,
             vertical_end=klevels,
-            offset_provider={
-                "C2E": self.grid.get_offset_provider("C2E"),
-                "E2C": self.grid.get_offset_provider("E2C"),
-                "C2CE": self.grid.get_offset_provider("C2CE"),
-            },
+            offset_provider=self.grid.offset_providers,
         )
         log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end")
         log.debug(
             "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start"
         )
         if self.config.apply_zdiffusion_t:
-            truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)(
+            truly_horizontal_diffusion_nabla_of_theta_over_steep_points(
                 mask=self.metric_state.mask_hdiff,
                 zd_vertoffset=self.metric_state.zd_vertoffset,
                 zd_diffcoef=self.metric_state.zd_diffcoef,
@@ -875,18 +835,14 @@ def _do_diffusion_step(
                 horizontal_end=cell_end_local,
                 vertical_start=0,
                 vertical_end=klevels,
-                offset_provider={
-                    "C2CEC": self.grid.get_offset_provider("C2CEC"),
-                    "C2E2C": self.grid.get_offset_provider("C2E2C"),
-                    "Koff": KDim,
-                },
+                offset_provider=self.grid.offset_providers,
             )
 
             log.debug(
                 "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end"
             )
         log.debug("running stencil 16 (update_theta_and_exner): start")
-        update_theta_and_exner.with_backend(backend)(
+        update_theta_and_exner(
             z_temp=self.z_temp,
             area=self.cell_params.area,
             theta_v=prognostic_state.theta_v,

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

@@ -14,12 +14,13 @@
 
 import numpy as np
 from gt4py.next import as_field
-from gt4py.next.common import Dimension, Field
+from gt4py.next.common import Dimension, Field, GridType
 from gt4py.next.ffront.decorator import field_operator, program
 from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum
 
 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
 from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift
+from icon4py.model.common.model_backend import backend
 
 
 # TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py
@@ -33,7 +34,7 @@ def _identity_c_k(field: Field[[CellDim, KDim], float]) -> Field[[CellDim, KDim]
     return field
 
 
-@program
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
 def copy_field(old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float]):
     _identity_c_k(old_f, out=new_f)
 
@@ -48,7 +49,7 @@ def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]
     return field * factor
 
 
-@program
+@program(backend=backend)
 def scale_k(field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float]):
     _scale_k(field, factor, out=scaled_field)
 
@@ -58,7 +59,7 @@ def _set_zero_v_k() -> Field[[VertexDim, KDim], float]:
     return broadcast(0.0, (VertexDim, KDim))
 
 
-@program
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
 def set_zero_v_k(field: Field[[VertexDim, KDim], float]):
     _set_zero_v_k(out=field)
 
@@ -89,7 +90,7 @@ def _setup_fields_for_initial_step(
     return diff_multfac_vn, smag_limit
 
 
-@program
+@program(backend=backend)
 def setup_fields_for_initial_step(
     k4: float,
     hdiff_efdt_ratio: float,
@@ -133,7 +134,7 @@ def _init_diffusion_local_fields_for_regular_timestemp(
     )
 
 
-@program
+@program(backend=backend)
 def init_diffusion_local_fields_for_regular_timestep(
     k4: float,
     dyn_substeps: float,