feat[next-dace]: support setting of external gpu stream

pyproject.toml

@@ -434,6 +434,7 @@ url = 'https://gridtools.github.io/pypi/'
 # dace = {index = "gridtools"}
 [tool.uv.sources]
 atlas4py = { index = "test.pypi" }
+dace = { git = "https://github.com/spcl/dace", branch = "new-gpu-codegen-dev" }
 
 # -- versioningit --
 [tool.versioningit]

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

@@ -24,15 +24,18 @@
 
 @dataclasses.dataclass(frozen=True)
 class DaCeBackend(backend.Backend[Any]):
-    """DaCe backend with support for injecting an external workspace at load time."""
+    """DaCe backend with support for injecting an external workspace and GPU stream at load time."""
 
     external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None
+    external_gpu_stream: gtx_wfdcommon.GPUStreamHandle | None = None
 
     def load_artifact(self, artifact: artifacts.CompilationArtifact) -> artifacts.ExecutableProgram:
         program = super().load_artifact(artifact)
         assert isinstance(program, gtx_wfddecoration.DaCeDecoratedProgram)
         # Inject the backend-level workspace so it is used when arguments are constructed.
         program.set_external_workspace(self.external_workspace or {})
+        # Inject the backend-level GPU stream, it is passed as argument to the SDFG.
+        program.set_external_gpu_stream(self.external_gpu_stream)
         return program
 
 
@@ -70,14 +73,16 @@ def make_dace_toolchain(
         allocator=cfg.make_allocator(),
         transforms=backend.DEFAULT_TRANSFORMS,
         external_workspace=cfg.external_workspace,
+        external_gpu_stream=cfg.external_gpu_stream,
     )
 
 
 def make_dace_backend(
     gpu: bool,
     auto_optimize: bool = True,
-    async_sdfg_call: bool = True,
+    sync_sdfg_call: bool = False,
     optimization_args: dict[str, Any] | None = None,
+    external_gpu_stream: gtx_wfdcommon.GPUStreamHandle | None = None,
     external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None,
     unstructured_horizontal_has_unit_stride: bool = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE,
     use_metrics: bool = True,
@@ -94,10 +99,17 @@ def make_dace_backend(
     Args:
         gpu: Enable GPU transformations and code generation.
         auto_optimize: Enable the SDFG auto-optimize pipeline.
-        async_sdfg_call: Make an asynchronous SDFG call on GPU to allow overlapping
-            of GPU kernel execution with the Python driver code.
+        sync_sdfg_call: Force a synchronous SDFG call on GPU, instead of allowing
+            the GPU kernel execution to overlap with the Python driver code.
         optimization_args: A `dict` containing configuration parameters for
             the SDFG auto-optimize pipeline, see `gt_auto_optimize()`.
+        external_gpu_stream: An external GPU stream to be used for the SDFG call,
+            which allows to synchronize the execution of GPU kernels with external
+            workload and to share the stream memory pool. It can be the pointer
+            value of the stream, as an integer, or a stream object such as
+            `cupy.cuda.Stream`, see `GPUStreamHandle`. The stream handle is passed
+            to the SDFG at runtime, on the first call, thus it does not affect the
+            generated code. If not provided, the default stream is used.
         external_workspace: Workspace memory externally allocated, which is used
             for SDFG's transient arrays when `transient_memory_mode` is `EXTERNAL`.
         unstructured_horizontal_has_unit_stride: When the memory layout has unit stride
@@ -131,12 +143,13 @@ def make_dace_backend(
             gpu=gpu,
             auto_optimize=auto_optimize,
             external_workspace=external_workspace,
+            external_gpu_stream=external_gpu_stream,
             unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride,
         ),
         translation=functools.partial(
             gtx_wfdfactory.make_dace_translator,
             auto_optimize_args=optimization_args,
-            async_sdfg_call=async_sdfg_call,
+            sync_sdfg_call=sync_sdfg_call,
             use_metrics=use_metrics,
             disable_field_origin_on_program_arguments=use_zero_origin,
             use_max_domain_range_on_unstructured_shift=use_max_domain_range_on_unstructured_shift,

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

@@ -11,6 +11,7 @@
 from typing import Any, Final, Generator, Optional, TypeAlias
 
 import dace
+from typing_extensions import Protocol, runtime_checkable
 
 from gt4py._core import definitions as core_defs
 from gt4py.eve import xtyping
@@ -34,6 +35,14 @@
 """DaCe datatype of `SDFG_ARG_METRIC_COMPUTE_TIME` argument."""
 
 
+SDFG_ARG_EXTERNAL_GPU_STREAM: Final[str] = "gt_external_gpu_stream"
+"""Name of SDFG argument to input an external GPU stream to GT4Py."""
+
+
+SDFG_ARG_EXTERNAL_GPU_STREAM_DTYPE: Final[dace.dtypes.typeclass] = dace.int64
+"""DaCe datatype of `SDFG_ARG_EXTERNAL_GPU_STREAM` argument, see `GPUStreamHandle`."""
+
+
 ExternalWorkspace: TypeAlias = dict[
     core_defs.DeviceType, xtyping.ArrayInterface | xtyping.CUDAArrayInterface
 ]
@@ -45,6 +54,59 @@
 """
 
 
+@runtime_checkable
+class GPUStream(Protocol):
+    """A GPU stream object that exposes its handle, for example `cupy.cuda.Stream`."""
+
+    @property
+    def ptr(self) -> int:
+        """The pointer value of the stream (`cudaStream_t` or `hipStream_t`)."""
+        ...
+
+    def __hash__(self) -> int:
+        return self.ptr
+
+
+GPUStreamHandle: TypeAlias = int | GPUStream
+""" Handle of a GPU stream (`cudaStream_t` or `hipStream_t`).
+
+    Either the pointer value of the stream as an integer, or a stream object that
+    exposes it as the `ptr` attribute, for example `cupy.cuda.Stream`. It is passed
+    to the SDFG at runtime as an integer, see `SDFG_ARG_EXTERNAL_GPU_STREAM` and
+    `get_gpu_stream_ptr()`.
+"""
+
+
+DEFAULT_GPU_STREAM: Final[int] = 0
+"""The pointer value of the default GPU stream, i.e. `nullptr`."""
+
+
+def get_gpu_stream_ptr(stream: GPUStreamHandle | None) -> int:
+    """Return the pointer value of a GPU stream, as an integer.
+
+    Args:
+        stream: The GPU stream handle. `None` selects the default stream.
+
+    Returns:
+        The pointer value of the stream, `DEFAULT_GPU_STREAM` for the default stream.
+
+    Raises:
+        TypeError: If `stream` is neither an integer nor a stream object that exposes
+            the `ptr` attribute.
+    """
+    if stream is None:
+        return DEFAULT_GPU_STREAM
+    # NOTE: `bool` is a subclass of `int`, but it is not a valid stream handle.
+    if isinstance(stream, int) and not isinstance(stream, bool):
+        return stream
+    if isinstance(stream, GPUStream) and isinstance(stream.ptr, int):
+        return stream.ptr
+    raise TypeError(
+        f"Invalid GPU stream: should be of type 'int' or a stream object with an integer "
+        f"'ptr' attribute, for example 'cupy.cuda.Stream', got '{type(stream).__name__}'."
+    )
+
+
 def set_dace_config(
     device_type: core_defs.DeviceType,
     cmake_build_type: Optional[gtx_config.CMakeBuildType] = None,
@@ -138,14 +200,27 @@ def set_dace_config(
     # 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
-    #  generated wrong code. The current approach is to use the default stream, i.e.
-    #  setting `max_concurrent_streams` to `-1`. However, the draw back is, that
-    #  apparently then all synchronization is disabled, even the one at the very
-    #  end of the SDFG call. To correct for that we are using either
-    #  `make_sdfg_call_sync()` or `make_sdfg_call_async()`, see there or in
-    #  [DaCe issue#2120](https://github.com/spcl/dace/issues/2120) for more.
+    #  generated wrong code. The current approach is to use a single stream, i.e.
+    #  setting `max_concurrent_streams` to `-1`. With the experimental CUDA codegen
+    #  (see below), all GPU work is then scheduled on the streams stored in
+    #  `__state->gpu_context->streams`, which are all set to the same stream: the
+    #  default stream or an external stream, see `add_external_gpu_stream_arg()`.
+    #  The synchronization at the end of the SDFG call is controlled by gt4py, see
+    #  `make_sdfg_call_sync()` and [DaCe issue#2120](https://github.com/spcl/dace/issues/2120).
     dace.Config.set("compiler.cuda.max_concurrent_streams", value=-1)
 
+    # Use the experimental CUDA codegen: unlike the legacy one, it launches kernels
+    #  and copies on the streams stored in `__state->gpu_context->streams`, even
+    #  when `max_concurrent_streams` is `-1`. This allows to replace the default
+    #  stream with an external stream, see `add_external_gpu_stream_arg()`.
+    dace.Config.set("compiler.cuda.implementation", value="experimental")
+
+    # The experimental CUDA codegen synchronizes the GPU streams at the end of
+    #  the SDFG by default, which makes the SDFG call synchronous. We disable it
+    #  because, when a synchronous call is requested, gt4py adds the stream
+    #  synchronization itself, see `make_sdfg_call_sync()`.
+    dace.Config.set("compiler.cuda.synchronize_on_exit", value=False)
+
     # This assumes that a process will only use one type of GPU.
     if device_type == core_defs.DeviceType.ROCM:
         dace.Config.set("compiler.cuda.backend", value="hip")

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

@@ -34,7 +34,9 @@ class DaCeDecoratedProgram:
     program. External workspace memory (when the SDFG uses
     ``TransientMemoryMode.EXTERNAL``) is installed onto the underlying
     ``CompiledDaceProgram`` before the first call via `set_external_workspace`;
-    its lifetime is owned by the caller, not by this wrapper.
+    its lifetime is owned by the caller, not by this wrapper. Similarly, an external
+    GPU stream can be set before the first call via `set_external_gpu_stream`; it
+    is passed as argument to the SDFG and applied on SDFG initialization.
     """
 
     def __init__(
@@ -43,6 +45,7 @@ def __init__(
         device_type: core_defs.DeviceType = core_defs.DeviceType.CPU,
     ) -> None:
         self._fun = fun
+        self._external_gpu_stream: gtx_wfdcommon.GPUStreamHandle | None = None
         # Retrieve metrics level from GT4Py environment variable.
         self._collect_time = metrics.is_level_enabled(metrics.PERFORMANCE)
         self._collect_time_arg = np.array(
@@ -85,6 +88,10 @@ def __call__(
                 gtx_wfdcommon.SDFG_ARG_METRIC_LEVEL: metrics.get_current_level(),
                 gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME: self._collect_time_arg,
             }
+            if gtx_wfdcommon.SDFG_ARG_EXTERNAL_GPU_STREAM in self._fun.sdfg_program.sdfg.symbols:
+                this_call_args[gtx_wfdcommon.SDFG_ARG_EXTERNAL_GPU_STREAM] = (
+                    gtx_wfdcommon.get_gpu_stream_ptr(self._external_gpu_stream)
+                )
             self._fun.construct_arguments(**this_call_args)
 
         # Perform the call to the SDFG.
@@ -98,6 +105,38 @@ def __call__(
     def set_external_workspace(self, external_workspace: gtx_wfdcommon.ExternalWorkspace) -> None:
         """Set the external workspace for the underlying compiled program.
 
-        This method should be called before the first call to the program.
+        This method should be called before the first call to the program, because
+        the workspace is installed when the SDFG arguments are constructed.
+
+        Args:
+            external_workspace: Mapping from device types to the array-like objects
+                used as workspace memory for the SDFG transient arrays, when the
+                transient memory mode is `EXTERNAL`.
         """
         self._fun.external_workspace = external_workspace
+
+    def set_external_gpu_stream(
+        self, external_gpu_stream: gtx_wfdcommon.GPUStreamHandle | None
+    ) -> None:
+        """Set the external GPU stream used by the underlying compiled program.
+
+        This method should be called before the first call to the program, because
+        the stream is applied on SDFG initialization.
+
+        Args:
+            external_gpu_stream: The handle of the external GPU stream, either as an
+                integer or as a stream object, see `GPUStreamHandle`. `None` selects
+                the default stream.
+
+        Raises:
+            RuntimeError: If the program was already called.
+            TypeError: If `external_gpu_stream` is not a valid GPU stream handle.
+        """
+        if self._fun.csdfg_argv is not None:
+            raise RuntimeError(
+                "The external GPU stream must be set before the first call to the program."
+            )
+        # Validate the stream handle, but store it as given: a stream object is thus
+        #  kept alive as long as the program, which uses its pointer value.
+        gtx_wfdcommon.get_gpu_stream_ptr(external_gpu_stream)
+        self._external_gpu_stream = external_gpu_stream

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

@@ -43,18 +43,26 @@ class DaCeConfig(next_backend.ToolchainConfig):
     #: `EXTERNAL` transient memory mode, which stores the transients in it. The
     #: workspace is a dict, so a config holding one is not hashable.
     external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None
+    #: GPU stream allocated outside the SDFG. The toolchain injects it into the
+    #: loaded programs, and the translation step then passes it to the SDFG as an
+    #: argument. The stream is an integer, so a config holding one is hashable.
+    external_gpu_stream: gtx_wfdcommon.GPUStreamHandle | None = None
 
     #: Name of the function that binds the SDFG arguments. The bindings and the
     #: compilation steps must agree on it.
     bind_func_name: ClassVar[str] = "update_sdfg_args"
 
+    def __post_init__(self) -> None:
+        # Validate the stream handle early, it is only used when the program is loaded.
+        gtx_wfdcommon.get_gpu_stream_ptr(self.external_gpu_stream)
+
 
 def make_dace_translator(
     cfg: DaCeConfig,
     /,
     *,
     auto_optimize_args: dict[str, Any] | None = None,
-    async_sdfg_call: bool = True,
+    sync_sdfg_call: bool = True,
     use_metrics: bool = True,
     disable_field_origin_on_program_arguments: bool = False,
     use_max_domain_range_on_unstructured_shift: bool | None = None,
@@ -70,8 +78,8 @@ def make_dace_translator(
             see `gt_auto_optimize()`. The parameters `DaCeTranslator` derives
             itself cannot be set here. With an external workspace in `cfg`, the
             `transient_memory_mode` defaults to `EXTERNAL`.
-        async_sdfg_call: Make an asynchronous SDFG call, to overlap GPU kernel
-            execution with the Python driver code. Only effective on GPU.
+        sync_sdfg_call: Make a synchronous SDFG call, to ensure that the GPU kernel
+            execution is completed before the Python driver code continues. Only effective on GPU.
         use_metrics: Add SDFG instrumentation for stencil compute time.
         disable_field_origin_on_program_arguments: Assume that all fields passed
             as program arguments have zero-based origin, which skips the range
@@ -116,7 +124,7 @@ def make_dace_translator(
         device_type=cfg.device_type,
         auto_optimize=cfg.auto_optimize,
         auto_optimize_args=optimization_args,
-        async_sdfg_call=async_sdfg_call and cfg.gpu,
+        sync_sdfg_call=sync_sdfg_call,
         unstructured_horizontal_has_unit_stride=cfg.unstructured_horizontal_has_unit_stride,
         use_metrics=use_metrics,
         disable_field_origin_on_program_arguments=disable_field_origin_on_program_arguments,