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, dace_gpu]
+      # 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(
@@ -926,6 +926,7 @@ def orchestration_uid(self) -> str:
             "_backend",
             "_exchange",
             "_grid",
+            "compile_time_connectivities",
             *[
                 name
                 for name in self.__dict__.keys()

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

@@ -420,8 +420,12 @@ def test_run_diffusion_single_step(
     backend,
     orchestration,
 ):
-    if not orchestration or not helpers.is_dace(backend):
-        pytest.skip("This test is only executed for orchestration on dace backends")
+    if orchestration and not helpers.is_dace(backend):
+        pytest.skip("Orchestration test requires a dace backend.")
+    if orchestration and data_alloc.is_cupy_device(backend):
+        pytest.xfail("DaCe GPU compilation fails.")
+    if helpers.is_embedded(backend):
+        pytest.xfail("Embedded backend currently fails in remap function.")
     grid = get_grid_for_experiment(experiment, backend)
     cell_geometry = get_cell_geometry_for_experiment(experiment, backend)
     edge_geometry = get_edge_geometry_for_experiment(experiment, backend)
@@ -525,6 +529,8 @@ def test_run_diffusion_multiple_steps(
 ):
     if not helpers.is_dace(backend):
         raise pytest.skip("This test is only executed for dace backends")
+    if data_alloc.is_cupy_device(backend):
+        pytest.xfail("DaCe GPU compilation fails.")
     ######################################################################
     # Diffusion initialization
     ######################################################################
@@ -665,8 +671,12 @@ def test_run_diffusion_initial_step(
     backend,
     orchestration,
 ):
-    if not orchestration or not helpers.is_dace(backend):
-        pytest.skip("This test is only executed for orchestration only on dace backends")
+    if orchestration and not helpers.is_dace(backend):
+        pytest.skip("Orchestration test requires a dace backend.")
+    if orchestration and data_alloc.is_cupy_device(backend):
+        pytest.xfail("DaCe GPU compilation fails.")
+    if helpers.is_embedded(backend):
+        pytest.xfail("Embedded backend currently fails in remap function.")
     grid = get_grid_for_experiment(experiment, backend)
     cell_geometry = get_cell_geometry_for_experiment(experiment, backend)
     edge_geometry = get_edge_geometry_for_experiment(experiment, backend)

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

@@ -34,7 +34,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 = data_alloc.as_numpy(icon_grid.get_offset_provider("C2CEC").table)
+    cec_table = data_alloc.as_numpy(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_fused_velocity_advection_stencil_15_to_18.py

@@ -118,7 +118,10 @@ class TestFusedVelocityAdvectionStencil15To18(StencilTest):
         "z_w_con_c_full",
         "ddt_w_adv",
     )
-    MARKERS = (pytest.mark.embedded_remap_error,)
+    MARKERS = (
+        pytest.mark.requires_concat_where,
+        pytest.mark.embedded_remap_error,
+    )
 
     @staticmethod
     def reference(

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

@@ -42,7 +42,10 @@ class TestFusedVelocityAdvectionStencil1To7(StencilTest):
         "z_w_concorr_me",
         "z_v_grad_w",
     )
-    MARKERS = (pytest.mark.embedded_remap_error,)
+    MARKERS = (
+        pytest.mark.requires_concat_where,
+        pytest.mark.embedded_remap_error,
+    )
 
     @staticmethod
     def _fused_velocity_advection_stencil_1_to_6_numpy(

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

@@ -40,7 +40,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-next]"]
 distributed = ["ghex>=0.3.0", "mpi4py>=3.1.5"]
 io = [
   # external dependencies

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

@@ -135,19 +135,19 @@ 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()
         xp = data_alloc.array_ns(self.config.on_gpu)
         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

@@ -9,10 +9,11 @@
 
 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,
@@ -25,10 +26,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

@@ -55,9 +55,7 @@
 if dace:
     from dace import hooks
     from dace.transformation.passes.simplify import SimplifyPass
-    from gt4py.next.program_processors.runners.dace_common.utility import (
-        connectivity_identifier,
-    )
+    from gt4py.next.program_processors.runners.dace import utils as gtx_dace_utils
 
 
 P = ParamSpec("P")
@@ -169,7 +167,7 @@ def wrapper(*args, **kwargs):
                         {
                             k: v
                             for k, v in grid.offset_providers.items()
-                            if connectivity_identifier(k) in sdfg.arrays
+                            if gtx_dace_utils.connectivity_identifier(k) in sdfg.arrays
                         },
                     ),
                 }
@@ -181,23 +179,13 @@ def wrapper(*args, **kwargs):
                 }
 
                 sdfg_args = dace_program._create_sdfg_args(sdfg, updated_args, updated_kwargs)
-                if func_is_method:
-                    del sdfg_args[self_name]
 
                 with dace.config.temporary_config():
                     configure_dace_temp_env(default_build_folder, self._backend)
                     return compiled_sdfg(**sdfg_args)
             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 +308,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 +520,9 @@ def dace_specific_kwargs(
         return {
             # connectivity tables at runtime
             **{
-                connectivity_identifier(k): v.table
+                gtx_dace_utils.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 +611,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

@@ -115,7 +115,7 @@ def add_halo_tasklet(
         )
 
     # Dummy return: preserve same interface with non-DaCe version
-    sdfg.add_scalar(name="__return", dtype=dace.int32)
+    sdfg.add_array(name="__return", shape=[1], dtype=dace.int32)
     ret = state.add_write("__return")
     state.add_edge(tasklet, "__out", ret, None, dace.memlet.Memlet(data="__return", subset="0"))
     out_connectors["__out"] = dace.int32
@@ -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",

model/common/src/icon4py/model/common/states/factory.py

@@ -38,6 +38,7 @@
 TODO: @halungge: allow to read configuration data
 
 """
+
 import collections
 import enum
 import functools
@@ -343,8 +344,8 @@ def _get_offset_providers(self, grid: icon_grid.IconGrid) -> dict[str, gtx.Field
                 horizontal_offsets = {
                     k: v
                     for k, v in grid.offset_providers.items()
-                    if isinstance(v, gtx.NeighborTableOffsetProvider)
-                    and v.origin_axis.kind == gtx.DimensionKind.HORIZONTAL
+                    if isinstance(v, gtx.Connectivity)
+                    and v.domain.dims[0].kind == gtx.DimensionKind.HORIZONTAL
                 }
                 offset_providers.update(horizontal_offsets)
             if dim.kind == gtx.DimensionKind.VERTICAL:
@@ -444,8 +445,8 @@ def _get_offset_providers(self, grid: icon_grid.IconGrid) -> dict[str, gtx.Field
                 horizontal_offsets = {
                     k: v
                     for k, v in grid.offset_providers.items()
-                    if isinstance(v, gtx.NeighborTableOffsetProvider)
-                    and v.origin_axis.kind == gtx.DimensionKind.HORIZONTAL
+                    if isinstance(v, gtx.Connectivity)
+                    and v.domain.dims[0].kind == gtx.DimensionKind.HORIZONTAL
                 }
                 offset_providers.update(horizontal_offsets)
             if dim.kind == gtx.DimensionKind.VERTICAL: