DO NOT MERGE: New DaCe GPU Codegen

ci/cscs-ci.yml

@@ -48,6 +48,7 @@ stages:
     DOCKER_BUILD_ARGS: '["BASE_IMAGE", "CACHE_DIR", "EXTRA_APTGET", "EXTRA_UV_ENV_VARS", "EXTRA_UV_PIP_ARGS", "EXTRA_UV_SYNC_ARGS", "PY_VERSION", "UV_VERSION", "WORKDIR_PATH" ]'
     PERSIST_IMAGE_NAME: ${CSCS_REGISTRY_PATH}/public/${ARCH}/base/gt4py-ci-${PY_VERSION}  # The $DOCKER_TAG tag is added in the before_script of .dynamic-image-name
     WATCH_FILECHANGES: 'ci/Dockerfile ci/cscs-ci.yml uv.lock'
+    DACE_compiler_cuda_implementation: experimental
   parallel:
     matrix:
       - PY_VERSION: *test_python_versions
@@ -57,6 +58,7 @@ stages:
     # jfrog.svc.cscs.ch/dockerhub/nvidia is the cached version of docker.io/nvidia
     BASE_IMAGE: jfrog.svc.cscs.ch/dockerhub/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}
     EXTRA_UV_SYNC_ARGS: "--extra cuda12"
+    DACE_compiler_cuda_implementation: experimental
 
 .build_extra_rocm:
   variables:
@@ -66,6 +68,7 @@ stages:
     EXTRA_UV_ENV_VARS: "CUPY_INSTALL_USE_HIP=1 HCC_AMDGPU_TARGET=gfx942 ROCM_HOME=/opt/rocm"
     KUBERNETES_MEMORY_REQUEST: "64Gi"
     KUBERNETES_MEMORY_LIMIT: "64Gi"
+    DACE_compiler_cuda_implementation: experimental
 
 build_cscs_gh200:
   extends:
@@ -89,13 +92,15 @@ build_cscs_amd_rocm:
     TEST_VARIANTS: 'cpu'  # Extended jobs should redefine which variants (cpu, cuda12, rocm6) to test 
     USE_MPI: 0 # TODO(havogt): to workaround the libfabric hook injecting incompatible libraries
     SLURM_JOB_NUM_NODES: 1
-    SLURM_TIMELIMIT: 5
+    SLURM_TIMELIMIT: 10
+    DACE_compiler_cuda_implementation: experimental
   parallel:
     matrix:
-      - SUBPACKAGE: [cartesian]
-        VARIANT: ['internal', 'dace']
-        SUBVARIANT: ['cuda12', 'rocm7', 'cpu']
-        PY_VERSION: *test_python_versions
+      # TODO(phimuell): `cartesian` does not work with the new code generator, no idea why.
+      #- SUBPACKAGE: [cartesian]
+      #  VARIANT: ['internal', 'dace']
+      #  SUBVARIANT: ['cuda12', 'rocm7', 'cpu']
+      #  PY_VERSION: *test_python_versions
       - SUBPACKAGE: eve
         PY_VERSION: *test_python_versions
       - SUBPACKAGE: next
@@ -139,6 +144,7 @@ test_cscs_gh200:
     GT4PY_BUILD_JOBS: 8
     # Limit test parallelism to avoid "OSError: too many open files" in the gt4py build stage.
     PYTEST_XDIST_AUTO_NUM_WORKERS: 32
+    DACE_compiler_cuda_implementation: experimental
   rules:
     - *exclude_variants_rules
     - if: $SUBPACKAGE == 'next' && $VARIANT == 'dace' && $DETAIL == 'nomesh'
@@ -167,6 +173,7 @@ test_cscs_amd_rocm:
     CMAKE_PREFIX_PATH: /opt/rocm # for next
     CUDA_HOME: /opt/rocm # for cartesian
     SLURM_TIMELIMIT: 20 # relaxed relative to gh200 as there is no pressure on the queue
+    DACE_compiler_cuda_implementation: experimental
   rules:
     - *exclude_variants_rules
     - if: $SUBPACKAGE == 'cartesian' && $VARIANT == 'internal' && $SUBVARIANT == 'cpu'

pyproject.toml

@@ -100,7 +100,7 @@ dependencies = [
   'click>=8.0.0',
   'cmake>=3.22',
   'cytoolz>=1.0.1',
-  'dace>=2.0.0a3',
+  'dace==2.3.24',
   'deepdiff>=8.1.0',
   'devtools>=0.6',
   'factory-boy>=3.3.3',
@@ -478,6 +478,9 @@ url = 'https://gridtools.github.io/pypi/'
 # dace = {index = "gridtools"}
 [tool.uv.sources]
 atlas4py = {index = "test.pypi"}
+dace = [
+  {git = "https://github.com/philip-paul-mueller/dace", branch = "phimuell__new-gpu-codegen-dev"}
+]
 
 # -- versioningit --
 [tool.versioningit]

src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py

@@ -1057,6 +1057,16 @@ def gt_gpu_apply_mempool(sdfg: dace.SDFG) -> None:
     Args:
         sdfg: The SDFG that should be processed.
     """
+
+    # TODO(phimuell): Reverse once the new codegen has caught up.
+    gpu_backend = dace.Config.get("compiler.cuda.backend")
+    if gpu_backend != "cuda":
+        warnings.warn(
+            f"GPU Memory-Pool is only implemented for `CUDA` and not for `{gpu_backend}`.",
+            stacklevel=0,
+        )
+        return
+
     for _, _, desc in sdfg.arrays_recursive():
         if (
             isinstance(desc, dace.data.Array)

src/gt4py/next/program_processors/runners/dace/workflow/common.py

@@ -109,6 +109,15 @@ def set_dace_config(
     #  This setting allows to throw an exception if any implicit Copy-Map slips thorugh.
     dace.Config.set("compiler.cuda.allow_implicit_memlet_to_map", value=False)
 
+    # Use the new GPU code generator
+    # NOTE: In the CI file we export the variable to force the experimental code gen to be used.
+    dace.Config.set("compiler.cuda.implementation", value="experimental")
+
+    # Skip GPU Sync at the end.
+    # NOTE: That this will most likely break the UNIT tests, but should not be a problem
+    #   for the blueline.
+    dace.Config.set("compiler", "cuda", "synchronize_on_exit", value=False)
+
     # In some stencils, for example `apply_diffusion_to_w`, the cuda codegen messes
     #  up with the cuda streams, i.e. it allocates N streams but uses N+1. The first
     #  idea was to use just one stream. However, even in that case the generator

src/gt4py/next/program_processors/runners/dace/workflow/translation.py

@@ -106,6 +106,10 @@ def make_sdfg_call_async(sdfg: dace.SDFG, gpu: bool) -> None:
     Todo: Revisit this function once DaCe changes its behaviour in this regard.
     """
 
+    # TODO(phimuell, edopao): Revisit this function after we understand the new
+    #   code generator better.
+    return
+
     # This is only a problem on GPU.
     # TODO(phimuell): Figuring out what about OpenMP.
     if not gpu:
@@ -282,6 +286,10 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None:
     work that runs on the GPU. Furthermore, all work is scheduled on the default stream.
     """
 
+    # TODO(phimuell, edopao): Revisit this function after we understand the new
+    #   code generator better.
+    return
+
     if not gpu:
         # This is only a problem on GPU. Dace uses OpenMP on CPU and
         # the OpenMP parallel region creates a synchronization point.

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py

@@ -201,6 +201,7 @@ def _check_cpu_sdfg_call(sdfg: dace.SDFG) -> None:
     assert not _are_streams_synchronized(sdfg)
 
 
+@pytest.mark.skip("To revisit after switch to new code gen.")
 @pytest.mark.parametrize(
     "make_async_sdfg_call",
     [False, True],
@@ -242,6 +243,7 @@ def test_generate_sdfg_async_call(make_async_sdfg_call: bool, device_type: core_
         _check_sdfg_without_async_call(sdfg)
 
 
+@pytest.mark.skip("To revisit after switch to new code gen.")
 def test_generate_sdfg_async_call_no_map(device_type: core_defs.DeviceType):
     """Verify that the flag `async_sdfg_call=True` has no effect on an SDFG that does not contain any GPU map."""
 
@@ -367,6 +369,7 @@ def _make_multi_state_sdfg_3(
     return sdfg, first_state, second_state
 
 
+@pytest.mark.skip("To revisit after switch to new code gen.")
 @pytest.mark.parametrize(
     "multi_state_config",
     [

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_gpu_utils.py

@@ -159,7 +159,6 @@ def test_set_gpu_properties(method: int):
     sdfg = dace.SDFG(gtx_transformations_utils.unique_name("gpu_properties_test"))
     state = sdfg.add_state(is_start_block=True)
 
-    map_entries: dict[int, dace_nodes.MapEntry] = {}
     for dim in [1, 2, 3, 4]:
         shape = (10,) * dim
         sdfg.add_array(
@@ -168,15 +167,15 @@ def test_set_gpu_properties(method: int):
         sdfg.add_array(
             f"B_{dim}", shape=shape, dtype=dace.float64, storage=dace.StorageType.GPU_Global
         )
-        _, me, _ = state.add_mapped_tasklet(
+        state.add_mapped_tasklet(
             f"map_{dim}",
             map_ranges={f"__i{i}": f"0:{s}" for i, s in enumerate(shape)},
             inputs={"__in": dace.Memlet(f"A_{dim}[{','.join(f'__i{i}' for i in range(dim))}]")},
             code="__out = math.cos(__in)",
             outputs={"__out": dace.Memlet(f"B_{dim}[{','.join(f'__i{i}' for i in range(dim))}]")},
             external_edges=True,
         )
-        map_entries[dim] = me
+    del state
     sdfg.validate()
 
     if method == 0:
@@ -204,6 +203,11 @@ def test_set_gpu_properties(method: int):
     else:
         raise ValueError(f"Unknown method {method}")
 
+    # Because of the inplace reconstruction all references to graph objects are destroyed.
+    map_entries: dict[int, dace_nodes.MapEntry] = {}
+    for node in sdfg.states()[0].nodes():
+        if isinstance(node, dace_nodes.MapEntry):
+            map_entries[int(node.label[4])] = node
     map1, map2, map3, map4 = (map_entries[d].map for d in [1, 2, 3, 4])
 
     # It takes the normal block size and does not regulate anything.
@@ -259,6 +263,7 @@ def test_set_gpu_properties_1D():
         map_entries[dim] = me
     sdfg.validate()
 
+    # `get_set_gpu_blocksize()` is non destructive, so `map_entries` are still pointing into the SDFG.
     sdfg.apply_gpu_transformations()
     gtx_dace_fieldview_gpu_utils.gt_set_gpu_blocksize(
         sdfg=sdfg,
@@ -323,6 +328,7 @@ def test_set_gpu_properties_2D_3D():
         map_entries[dim] = me
     sdfg.validate()
 
+    # `get_set_gpu_blocksize()` is non destructive, so `map_entries` are still pointing into the SDFG.
     sdfg.apply_gpu_transformations()
     gtx_dace_fieldview_gpu_utils.gt_set_gpu_blocksize(
         sdfg=sdfg,