[edopao]

Adds support for running the DaCe-generated SDFG on a GPU stream provided by the caller, instead of always using DaCe's default stream. This lets the SDFG's GPU work be interleaved/synchronized with an external (non-GT4Py) workload and share that workload's stream memory pool.

• make_dace_backend() gains a new external_gpu_stream: int | None parameter. The stream handle is a pointer, which in general changes from run to run of the application (except for the default stream), therefore it is a runtime argument of the SDFG, not a compile-time value:
  • The translation stage adds the SDFG symbol gt_external_gpu_stream (int64) to every GPU SDFG, via add_external_gpu_stream_arg(). The SDFG init code casts it with reinterpret_cast<gpuStream_t> and passes it to DaCe's __dace_gpu_set_all_streams(), which writes it into every entry of __state->gpu_context->streams; the value 0 is the default stream (nullptr). It is a symbol, not a scalar data container, because DaCe only passes the symbols to the SDFG init functions.
  • The stream is stored on DaCeBackend.external_gpu_stream and injected into the loaded program by load_artifact() through DaCeDecoratedProgram.set_external_gpu_stream(), in the same way as the external workspace. On the first call, it is passed to the SDFG argument vector (0 when not set), and applied by the SDFG initialization.
  • The generated code does not depend on the stream value, so the compiled SDFG is not bound to a specific stream and can be reused from the build cache across runs.
• Renamed DaCeTranslator.async_sdfg_call / make_dace_backend(async_sdfg_call=...) to sync_sdfg_call, inverting the flag so its name matches its effect (sync_sdfg_call=True forces a synchronous SDFG call; the default False preserves the previous async-by-default behavior).
• Removed make_sdfg_call_async(): setting all DaCe streams is no longer tied to the async/sync choice — it is now handled unconditionally (for both sync and async calls) by the SDFG init code added in add_external_gpu_stream_arg().
• make_sdfg_call_sync()'s synchronization tasklet now loops over every stream tracked in DaCe's GPU context (__state->gpu_context->streams[0..num_streams)) and calls StreamSynchronize on each one, rather than only the default stream, since an external stream may not be DaCe's stream 0.
• The metrics timers (use_metrics) now synchronize only the SDFG GPU streams, instead of calling DeviceSynchronize(), so that the measured time does not include unrelated work on other streams, e.g. external workload queued by the caller.

Switch to the experimental DaCe GPU codegen

With compiler.cuda.max_concurrent_streams = -1, the legacy DaCe CUDA codegen hard-codes nullptr (the default stream) for kernel launches and copies, and never reads gpu_context->streams. Overriding the stream array therefore had no effect. To make the external stream actually used:

• set_dace_config() selects compiler.cuda.implementation = "experimental". The experimental codegen launches kernels and copies on __state->gpu_context->streams[...] also with max_concurrent_streams = -1, and the SDFG init code runs after the stream array is initialized, so the override takes effect.
• dace is temporarily pinned to the fix-experimental-cuda-gpu-copy-fallback branch of edopao/dace, which is based on the new-gpu-codegen-dev branch of spcl/dace and fixes a bug in the experimental codegen: a GPU-to-GPU copy whose write overlap with another writer could not be decided symbolically (e.g. concat_where with dynamic domains) was emitted as a host-side copy between device buffers, which segfaulted.
• set_dace_config() sets compiler.cuda.synchronize_on_exit = False. Otherwise the experimental codegen synchronizes the GPU stream at the end of every SDFG call, making the call always synchronous. When sync_sdfg_call=True, the synchronization is still added by make_sdfg_call_sync(). Note that DaCe still synchronizes the stream before the host reads data computed on GPU.

Tests

• test_generate_sdfg_external_gpu_stream is parametrized over a program lowered to a GPU kernel and a program lowered to a GPU copy. It checks that the SDFG declares the gt_external_gpu_stream symbol, that it is in the init signature and applied by the init code, and it generates the code with the GT4Py DaCe configuration to check that kernel launches and Memcpy*Async calls take their stream from __state->gpu_context->streams (i.e. no nullptr). This test fails with the legacy codegen.
• test_decorated_program_passes_external_gpu_stream_arg checks that the stream set on the decorated program is passed as SDFG argument on the first call (0 when not set), and that it cannot be changed after the first call. test_decorated_program_skips_external_gpu_stream_arg_on_cpu checks that the argument is not passed to CPU SDFGs.
• Fixed _are_streams_synchronized(): it used re.match, which only matches at the beginning of the generated file and thus always returned False, and it generated code outside of the GT4Py DaCe configuration context. It now uses re.search within dace_context().
• With the fixed helper, test_generate_sdfg_async_call_multi_state now checks that the GPU stream is synchronized after the device-to-host copy of a value that is used on the InterState edge or in a CPU tasklet (this was previously documented as a known wrong behavior, see No Synchronization When Using The "DefaultStream" spcl/dace#2120). The case where the value is only used in the InterState edge condition is marked as xfail(strict=True), because the experimental codegen does not emit the synchronization there.

Open points

• The external stream is applied by the SDFG init code, which runs once on the first call of each compiled program. Thus the stream is fixed for the lifetime of the program in the process: set_external_gpu_stream() raises an error if called after the first call.
• The unit tests in test_dace_translation.py only generate code, thus the requires_gpu marks were removed from the CUDA/ROCM parametrization so that they also run in CPU-only CI.
• The experimental DaCe codegen is a work-in-progress implementation; the GPU test suite (test_next on CSCS) is the real validation.
• The dace pin on the edopao/dace fork is temporary, until the copy fix is merged upstream in new-gpu-codegen-dev.
• Missing stream synchronization in DaCe when a device-to-host copy result is read in an InterState edge condition (see the xfail above) — to be reported upstream.
• The experimental codegen falls back to the default stream (nullptr) for kernels in library nodes expanded to nested SDFGs, which have no GPU stream connected; such kernels would not run on the external stream.

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default