Struct StreamGuard 

pub struct StreamGuard { /* private fields */ }

RAII guard that sets a stream as the current CUDA stream and restores the previous stream on drop.

Nestable: inner guards restore the outer guard’s stream, not the default stream. This is critical for DDP where sync_now_nccl temporarily switches to comm_stream inside a compute_stream guard.

let compute = GpuStream::new(Device::CUDA(0), false)?;
let comm = GpuStream::new(Device::CUDA(0), false)?;
{
    let _outer = StreamGuard::new(&compute);
    // All CUDA ops on compute_stream.
    {
        let _inner = StreamGuard::new(&comm);
        // CUDA ops on comm_stream.
    }
    // Restored to compute_stream (not default).
}
// Restored to whatever was current before _outer.

Implementations

impl StreamGuard

pub fn new(stream: &GpuStream) -> Self

Set stream as the current CUDA stream. The previous stream is saved and restored when this guard is dropped.

Trait Implementations

impl Drop for StreamGuard

fn drop(&mut self)

Executes the destructor for this type.

fn pin_drop(self: Pin<&mut Self>)

🔬This is a nightly-only experimental API. (pin_ergonomics)

Execute the destructor for this type, but different to Drop::drop, it requires self to be pinned.