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_external_stream_sync in the translation step.
  • When enabled on GPU SDFGs, injects entry and exit tasklets that:
    • Wait on a GT4Py-owned cross-call event before touching the external workspace.
    • Record per-stream done events on internal DaCe streams.
    • Make the external sync stream wait on those events.
    • Record the cross-call event on the external sync stream.
  • Added SDFG scalar symbols __external_ws_event and __external_sync_stream, passed as uint64 scalars.
  • When max_concurrent_gpu_streams == 0, emits DaCe init code that forces all internal streams to the default stream instead.

• 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.

Related documents

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

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default