feat[next-dace]: Add synchronization for multi-stream scheduling

[edopao]

Summary

This PR makes DaCe's multi-stream scheduling opt-in and safe for GT4Py's external workspace mode. When max_concurrent_gpu_streams >= 1, the DaCe backend injects CUDA/HIP event-based synchronization around SDFG execution, ordering the internal DaCe streams with respect to a caller-provided external stream (or the default stream).

Motivation

External workspace mode (transient_memory_mode=EXTERNAL) avoids per-call GPU allocations by reusing a caller-provided workspace. This is only safe when no other GPU work touches the workspace between two consecutive calls. A concrete use case that breaks this assumption is overlapping GT4Py program execution with MPI halo exchange: halo exchange runs on one CUDA/HIP stream while the GT4Py program runs on DaCe's internal streams. Without explicit ordering, the external workspace can be reused before previous kernels or the halo exchange finish, causing data races.

Changes

• Replaced the global GT4PY_MAX_CONCURRENT_GPU_STREAMS config with a backend parameter:

  • Added max_concurrent_gpu_streams to make_dace_backend, DaCeBackendFactory, DaCeWorkflowFactory, DaCeTranslationStepFactory, DaCeCompiler, and DaCeTranslator.
  • Negative values are rejected in DaCeTranslator.__post_init__.
  • 0 preserves the historical default-stream-only behavior (sets DaCe's compiler.cuda.max_concurrent_streams to -1).
  • >= 1 enables DaCe's internal stream pool.
  • Ignored on CPU targets.

• Added optional external_sync_stream support:

  • Added external_sync_stream: cupy.cuda.Stream | None = None to make_dace_backend.
  • The stream is stored on DaCeBackend (not the picklable workflow) and injected at artifact load time.
  • When no stream is provided and multi-stream scheduling is enabled, the default stream is used implicitly.
  • When max_concurrent_gpu_streams == 0, external_sync_stream is ignored and no validation occurs.

• Implemented event-based SDFG synchronization:

  • Added add_synchronization in the translation step.
  • When enabled on GPU SDFGs, injects sync_entry and sync_exit tasklets that record per-stream done events on internal DaCe streams and make the external sync stream wait on them.
  • Added the SDFG scalar symbol __external_sync_stream, passed as a uint64 scalar.
  • When max_concurrent_gpu_streams == 0 and a blocking call is requested, emits tasklets that synchronize on the default stream.
  • When max_concurrent_gpu_streams == 0 and an asynchronous call is requested, emits DaCe init code that forces all internal streams to the default stream and skips the tasklets.

• Added async_sdfg_call parameter:

  • Controls whether the SDFG call blocks until GPU execution has completed.
  • Default is True (asynchronous).

• Added runtime stream validation:

  • Rejected on CPU targets.
  • Must be a cupy.cuda.Stream.
  • Must be on the current GPU device.
  • Must pass cudaStreamQuery (accepting both cudaSuccess and cudaErrorNotReady).

• Added safeguards:

  • transient_memory_mode=POOL combined with max_concurrent_gpu_streams >= 1 is rejected, because the memory-pool pass assumes default-stream ordering.

• Added / updated unit tests:

  • Stream synchronization behavior in translation.
  • Stream pointer forwarding and default-stream fallback in compilation.
  • Backend factory parameter propagation and validation.

Default behavior

No behavioral change unless explicitly requested: max_concurrent_gpu_streams=0 (default) produces the same SDFG and execution as before.

Testing

uv run pytest tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py
uv run pytest tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py
uv run pytest tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_streams.py
uv run pytest tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py

Out of scope

• Per-call stream selection.
• Non-cupy stream support (e.g. torch, raw CUDA, HIP handles).
• Device-wide synchronization.
• Multi-stream safety for SCOPED/POOL modes.
• GT4Py-owned cross-call events; the caller is responsible for ordering work on the external stream across calls.

Related documents

• docs/development/ADRs/0028-External_Stream_Synchronization.md

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default