Test GTIR-DaCe backend

.gitignore

@@ -17,7 +17,7 @@ simple_mesh*.nc
 **/docs/_source/*.rst
 
 ### GT4Py ####
-.gt_cache/
+.gt4py_cache/
 
 # DaCe
 .dacecache

ci/dace.yml

@@ -4,11 +4,13 @@ include:
 .test_model_stencils:
   stage: test
   script:
-    - nox -s test_model_stencils-3.10 -- --backend=$BACKEND --grid=$GRID
+    - nox -s "test_model_stencils-3.10(subpackage='$COMPONENT')" -- --backend=$BACKEND --grid=$GRID
   parallel:
     matrix:
     - BACKEND: [dace_cpu, dace_gpu]
       GRID: [simple_grid, icon_grid]
+      # TODO(edopao): Add more components once they work fine with DaCe 
+      COMPONENT: [atmosphere/diffusion, atmosphere/dycore]
 test_model_stencils_x86_64:
   extends: [.test_model_stencils, .test_template_x86_64]
 test_model_stencils_aarch64:
@@ -20,9 +22,9 @@ test_model_stencils_aarch64:
     - nox -s "test_model_datatest-3.10(subpackage='$COMPONENT')" -- --backend=$BACKEND
   parallel:
     matrix:
-    # TODO(edopao): Add more components once they work fine with DaCe 
-    - COMPONENT: [atmosphere/diffusion, atmosphere/dycore]
-      BACKEND: [dace_cpu]
+    - BACKEND: [dace_cpu]
+      # TODO(edopao): Add more components once they work fine with DaCe 
+      COMPONENT: [atmosphere/diffusion, atmosphere/dycore]
 test_model_datatests_x86_64:
   extends: [.test_model_datatests, .test_template_x86_64]
 test_model_datatests_aarch64:
@@ -31,12 +33,13 @@ test_model_datatests_aarch64:
 .benchmark_model_stencils:
   stage: benchmark
   script:
-    # force execution of tests where validation is expected to fail, because the reason for failure is wrong numpy reference
-    - nox -s benchmark_model-3.10  -- --backend=$BACKEND --grid=$GRID --runxfail
+    - nox -s "benchmark_model-3.10(subpackage='$COMPONENT')"  -- --backend=$BACKEND --grid=$GRID
   parallel:
     matrix:
     - BACKEND: [dace_cpu, dace_gpu]
       GRID: [icon_grid, icon_grid_global]
+      # TODO(edopao): Add more components once they work fine with DaCe 
+      COMPONENT: [atmosphere/diffusion, atmosphere/dycore]
 benchmark_model_stencils_x86_64:
   extends: [.benchmark_model_stencils, .test_template_x86_64]
 benchmark_model_stencils_aarch64:

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

@@ -463,9 +463,9 @@ def __init__(
 
         self._determine_horizontal_domains()
 
-        self.compile_time_connectivities = dace_orchestration.build_compile_time_connectivities(
-            self._grid.offset_providers
-        )
+        # TODO(edopao): we should call gtx.common.offset_provider_to_type()
+        # but this requires some changes in type inference.
+        self.compile_time_connectivities = self._grid.offset_providers
 
     def _allocate_temporary_fields(self):
         self.diff_multfac_vn = data_alloc.allocate_zero_field(
@@ -925,6 +925,7 @@ def orchestration_uid(self) -> str:
             "_backend",
             "_exchange",
             "_grid",
+            "compile_time_connectivities",
             *[
                 name
                 for name in self.__dict__.keys()

model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py

@@ -33,7 +33,7 @@ def reference(
         **kwargs,
     ) -> dict:
         c2e = grid.connectivities[dims.C2EDim]
-        c2ce = grid.get_offset_provider("C2CE").table
+        c2ce = grid.get_offset_provider("C2CE").ndarray
 
         geofac_div = np.expand_dims(geofac_div, axis=-1)
         e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1)

model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py

@@ -13,7 +13,7 @@
 from icon4py.model.common.decomposition import definitions
 from icon4py.model.common.grid import vertical as v_grid
 from icon4py.model.common.utils import data_allocation as data_alloc
-from icon4py.model.testing import datatest_utils, parallel_helpers
+from icon4py.model.testing import datatest_utils, helpers, parallel_helpers
 
 from .. import utils
 
@@ -22,7 +22,7 @@
 @pytest.mark.parametrize("experiment", [datatest_utils.REGIONAL_EXPERIMENT])
 @pytest.mark.parametrize("ndyn_substeps", [2])
 @pytest.mark.parametrize("linit", [True, False])
-@pytest.mark.parametrize("orchestration", [True, False])
+@pytest.mark.parametrize("orchestration", [False, True])
 def test_parallel_diffusion(
     experiment,
     step_date_init,
@@ -44,7 +44,7 @@ def test_parallel_diffusion(
     backend,
     orchestration,
 ):
-    if orchestration and ("dace" not in backend.name.lower()):
+    if orchestration and not helpers.is_dace(backend):
         raise pytest.skip("This test is only executed for `dace backends.")
     caplog.set_level("INFO")
     parallel_helpers.check_comm_size(processor_props)
@@ -174,7 +174,7 @@ def test_parallel_diffusion_multiple_steps(
     caplog,
     backend,
 ):
-    if "dace" not in backend.name.lower():
+    if not helpers.is_dace(backend):
         raise pytest.skip("This test is only executed for `dace backends.")
     ######################################################################
     # Diffusion initialization

model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py

@@ -388,7 +388,8 @@ def test_verify_diffusion_init_against_savepoint(
         (dt_utils.GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"),
     ],
 )
-@pytest.mark.parametrize("ndyn_substeps, orchestration", [(2, [True, False])])
+@pytest.mark.parametrize("ndyn_substeps", [2])
+@pytest.mark.parametrize("orchestration", [False, True])
 def test_run_diffusion_single_step(
     savepoint_diffusion_init,
     savepoint_diffusion_exit,
@@ -403,7 +404,7 @@ def test_run_diffusion_single_step(
     backend,
     orchestration,
 ):
-    if orchestration and ("dace" not in backend.name.lower()):
+    if orchestration and not helpers.is_dace(backend):
         pytest.skip(f"running backend = '{backend.name}': orchestration only on dace backends")
     grid = get_grid_for_experiment(experiment, backend)
     cell_geometry = get_cell_geometry_for_experiment(experiment, backend)
@@ -502,8 +503,8 @@ def test_run_diffusion_multiple_steps(
     backend,
     icon_grid,
 ):
-    if "dace" not in backend.name.lower():
-        raise pytest.skip("This test is only executed for DaCe backends.")
+    if not helpers.is_dace(backend):
+        pytest.skip(f"running backend = '{backend.name}': orchestration only on dace backends")
     ######################################################################
     # Diffusion initialization
     ######################################################################
@@ -624,7 +625,8 @@ def test_run_diffusion_multiple_steps(
 
 @pytest.mark.datatest
 @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT])
-@pytest.mark.parametrize("linit, orchestration", [(True, [True, False])])
+@pytest.mark.parametrize("linit", [True, True])
+@pytest.mark.parametrize("orchestration", [False, True])
 def test_run_diffusion_initial_step(
     experiment,
     linit,
@@ -639,7 +641,7 @@ def test_run_diffusion_initial_step(
     backend,
     orchestration,
 ):
-    if orchestration and ("dace" not in backend.name.lower()):
+    if orchestration and not helpers.is_dace(backend):
         pytest.skip(f"running backend = '{backend.name}': orchestration only on dace backends")
     grid = get_grid_for_experiment(experiment, backend)
     cell_geometry = get_cell_geometry_for_experiment(experiment, backend)

model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_states.py

@@ -30,7 +30,7 @@ def test_verify_geofac_n2s_field_manipulation(interpolation_savepoint, icon_grid
     geofac_c = interpolation_state.geofac_n2s_c.asnumpy()
     geofac_nbh = interpolation_state.geofac_n2s_nbh.asnumpy()
     assert np.count_nonzero(geofac_nbh) > 0
-    cec_table = icon_grid.get_offset_provider("C2CEC").table
+    cec_table = icon_grid.get_offset_provider("C2CEC").ndarray
     assert np.allclose(geofac_c, geofac_n2s[:, 0])
     assert geofac_nbh[cec_table].shape == geofac_n2s[:, 1:].shape
     assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:])

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_solve_nonhydro_stencil_39_40.py

@@ -6,7 +6,6 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 import gt4py.next as gtx
-from gt4py.next.ffront.decorator import GridType, field_operator, program
 from gt4py.next.ffront.fbuiltins import where
 
 from icon4py.model.atmosphere.dycore.stencils.compute_contravariant_correction_of_w import (
@@ -19,7 +18,7 @@
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
-@field_operator
+@gtx.field_operator
 def _fused_solve_nonhydro_stencil_39_40(
     e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat],
     z_w_concorr_me: fa.EdgeKField[vpfloat],
@@ -39,7 +38,7 @@ def _fused_solve_nonhydro_stencil_39_40(
     return w_concorr_c
 
 
-@program(grid_type=GridType.UNSTRUCTURED)
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def fused_solve_nonhydro_stencil_39_40(
     e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat],
     z_w_concorr_me: fa.EdgeKField[vpfloat],

model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_interpolated_horizontal_advection_of_w.py

@@ -22,7 +22,7 @@ def add_interpolated_horizontal_advection_of_w_numpy(
     grid, e_bln_c_s: np.array, z_v_grad_w: np.array, ddt_w_adv: np.array, **kwargs
 ) -> np.array:
     e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1)
-    c2ce = grid.get_offset_provider("C2CE").table
+    c2ce = grid.get_offset_provider("C2CE").ndarray
 
     ddt_w_adv = ddt_w_adv + np.sum(
         z_v_grad_w[grid.connectivities[dims.C2EDim]] * e_bln_c_s[c2ce],

model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_divergence_of_fluxes_of_rho_and_theta.py

@@ -32,7 +32,7 @@ def reference(
     ) -> tuple[np.array]:
         c2e = grid.connectivities[dims.C2EDim]
         geofac_div = np.expand_dims(geofac_div, axis=-1)
-        c2ce = grid.get_offset_provider("C2CE").table
+        c2ce = grid.get_offset_provider("C2CE").ndarray
 
         z_flxdiv_mass = np.sum(
             geofac_div[c2ce] * mass_fl_e[c2e],

model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_15_to_18.py

@@ -41,6 +41,7 @@ class TestFusedVelocityAdvectionStencil15To18(StencilTest):
         "z_w_con_c_full",
         "ddt_w_adv",
     )
+    MARKER = (pytest.mark.requires_concat_where,)
 
     @staticmethod
     def _fused_velocity_advection_stencil_16_to_18(

model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py

@@ -22,7 +22,7 @@ def interpolate_to_cell_center_numpy(
     grid, interpolant: np.array, e_bln_c_s: np.array, **kwargs
 ) -> np.array:
     e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1)
-    c2ce = grid.get_offset_provider("C2CE").table
+    c2ce = grid.get_offset_provider("C2CE").ndarray
 
     interpolation = np.sum(
         interpolant[grid.connectivities[dims.C2EDim]] * e_bln_c_s[c2ce],

model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py

@@ -7,7 +7,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 import gt4py.next as gtx
 import numpy as np
-from gt4py.next.program_processors.runners.gtfn import run_gtfn
+from gt4py.next.program_processors.runners.dace import run_dace_cpu
 
 from icon4py.model.atmosphere.dycore.stencils.solve_tridiagonal_matrix_for_w_forward_sweep import (
     solve_tridiagonal_matrix_for_w_forward_sweep,
@@ -87,7 +87,7 @@ def test_solve_tridiagonal_matrix_for_w_forward_sweep():
     v_start = 1
     v_end = gtx.int32(grid.num_levels)
     # TODO we run this test with the C++ backend as the `embedded` backend doesn't handle this pattern
-    solve_tridiagonal_matrix_for_w_forward_sweep.with_backend(run_gtfn)(
+    solve_tridiagonal_matrix_for_w_forward_sweep.with_backend(run_dace_cpu)(
         vwind_impl_wgt=vwind_impl_wgt,
         theta_v_ic=theta_v_ic,
         ddqz_z_half=ddqz_z_half,

model/common/pyproject.toml

@@ -38,7 +38,7 @@ version = "0.0.6"
 all = ["icon4py-common[dace,distributed,io]"]
 cuda11 = ['cupy-cuda11x>=13.0', 'gt4py[cuda11]']
 cuda12 = ['cupy-cuda12x>=13.0', 'gt4py[cuda12]']
-dace = ["dace<1.0", "gt4py[dace]"] # TODO(egparedes): DaCe max constraint should be transformed to min constraint after updating gt4py
+dace = ["dace>=1.0", "gt4py[dace]"]
 distributed = ["ghex>=0.3.0", "mpi4py>=3.1.5"]
 io = [
   # external dependencies

model/common/src/icon4py/model/common/grid/base.py

@@ -136,18 +136,18 @@ def _get_offset_provider(self, dim, from_dim, to_dim):
         ), 'Neighbor table\'s "{}" data type must be gtx.int32. Instead it\'s "{}"'.format(
             dim, self.connectivities[dim].dtype
         )
-        return gtx.NeighborTableOffsetProvider(
-            self.connectivities[dim],
-            from_dim,
+        return gtx.as_connectivity(
+            [from_dim, dim],
             to_dim,
-            self.size[dim],
-            has_skip_values=self._has_skip_values(dim),
+            self.connectivities[dim],
+            skip_value=-1 if self._has_skip_values(dim) else None,
         )
 
     def _get_offset_provider_for_sparse_fields(self, dim, from_dim, to_dim):
         if dim not in self.connectivities:
             raise MissingConnectivity()
         return grid_utils.neighbortable_offset_provider_for_1d_sparse_fields(
+            dim,
             self.connectivities[dim].shape,
             from_dim,
             to_dim,

model/common/src/icon4py/model/common/grid/utils.py

@@ -7,10 +7,11 @@
 # SPDX-License-Identifier: BSD-3-Clause
 import gt4py.next as gtx
 import numpy as np
-from gt4py.next import Dimension, NeighborTableOffsetProvider
+from gt4py.next import Dimension
 
 
 def neighbortable_offset_provider_for_1d_sparse_fields(
+    dim: Dimension,
     old_shape: tuple[int, int],
     origin_axis: Dimension,
     neighbor_axis: Dimension,
@@ -22,10 +23,9 @@ def neighbortable_offset_provider_for_1d_sparse_fields(
     ), 'Neighbor table\'s ("{}" to "{}") data type for 1d sparse fields must be gtx.int32. Instead it\'s "{}"'.format(
         origin_axis, neighbor_axis, table.dtype
     )
-    return NeighborTableOffsetProvider(
-        table,
-        origin_axis,
+    return gtx.as_connectivity(
+        [origin_axis, dim],
         neighbor_axis,
-        table.shape[1],
-        has_skip_values=has_skip_values,
+        table,
+        skip_value=-1 if has_skip_values else None,
     )

model/common/src/icon4py/model/common/orchestration/decorator.py

@@ -190,14 +190,6 @@ def wrapper(*args, **kwargs):
             else:
                 return fuse_func(*args, **kwargs)
 
-        # Pytest does not clear the cache between runs in a proper way -pytest.mark.parametrize(...)-.
-        # This leads to corrupted cache and subsequent errors.
-        # To avoid this, we provide a way to clear the cache.
-        def clear_cache():
-            orchestrator_cache.clear()
-
-        wrapper.clear_cache = clear_cache
-
         return wrapper
 
     return _decorator(func) if func else _decorator
@@ -320,21 +312,6 @@ def wait(comm_handle: Union[int, decomposition.ExchangeResult]):
         comm_handle.wait()
 
 
-def build_compile_time_connectivities(
-    offset_providers: dict[str, gtx.common.Connectivity],
-) -> dict[str, gtx.common.Connectivity]:
-    connectivities = {}
-    for k, v in offset_providers.items():
-        if hasattr(v, "table"):
-            connectivities[k] = gtx.otf.arguments.CompileTimeConnectivity(
-                v.max_neighbors, v.has_skip_values, v.origin_axis, v.neighbor_axis, v.table.dtype
-            )
-        else:
-            connectivities[k] = v
-
-    return connectivities
-
-
 if dace:
 
     def to_dace_annotations(fuse_func: Callable) -> dict[str, Any]:
@@ -547,9 +524,9 @@ def dace_specific_kwargs(
         return {
             # connectivity tables at runtime
             **{
-                connectivity_identifier(k): v.table
+                connectivity_identifier(k): v.ndarray
                 for k, v in offset_providers.items()
-                if hasattr(v, "table")
+                if hasattr(v, "ndarray")
             },
             # GHEX C++ ptrs
             "__context_ptr": expose_cpp_ptr(exchange_obj._context)
@@ -638,8 +615,8 @@ def _concretize_symbols_for_dace_structure(dace_cls, orig_cls):
 
         return {
             **{
-                "CellDim_sym": grid.offset_providers["C2E"].table.shape[0],
-                "EdgeDim_sym": grid.offset_providers["E2C"].table.shape[0],
+                "CellDim_sym": grid.offset_providers["C2E"].ndarray.shape[0],
+                "EdgeDim_sym": grid.offset_providers["E2C"].ndarray.shape[0],
                 "KDim_sym": grid.num_levels,
             },
             **concretize_symbols_for_dace_structure,

model/common/src/icon4py/model/common/orchestration/halo_exchange.py

@@ -282,7 +282,7 @@ def __sdfg__(self, *args, **kwargs) -> dace.SDFG:
         sdfg = dace.SDFG("DummyNestedSDFG")
         state = sdfg.add_state()
 
-        sdfg.add_scalar(name="__return", dtype=dace.int32)
+        sdfg.add_array(name="__return", shape=[1], dtype=dace.int32)
 
         tasklet = dace.sdfg.nodes.Tasklet(
             "DummyNestedSDFG",