Interoperability: CUDA Streams

[ax3l]

CUDA streams the same way as FFT plans and BLAS handles and MPI communicators are a shared resource and we need to be able to expose them (easy) but also be able to be constrained to pre-initialized handles, if we need them to be. We can already do this pretty well now for FFTs, MPI and some BLAS stuff.

Functions like ParallelFor will call Gpu::gpuStream to get a stream for its kernel. But the launch function allows one to pass a stream.

We could expose a function that let's the user the set the stream before ParallelFor.

Currently, AMReX builds 4 streams by default. MFIter defines the strategy (round robin) and sets the values of Gpu::gpuStream.

[passionworkeer]

A workable API shape might mirror the existing handle-style interop points: keep the current default round-robin behavior, but add an opt-in stream guard / setter around Gpu::gpuStream so Python callers can temporarily bind AMReX launches to an external stream. Two details seem important for correctness: first, expose both ‘get current stream’ and ‘run on provided stream’ paths, so integrations can either reuse or override; second, make the override scoped (RAII / context-manager style) to avoid leaking stream state across unrelated MFIter/ParallelFor calls. For validation, I’d test that a kernel launched under an injected stream preserves ordering with upstream work on the same stream, while the default path still round-robins exactly as before when no external stream is provided.