Module staging 

Reusable pinned host staging buffer for hot-path H2D / D2H transfers.

StagingBuffer owns one page-locked (PinnedBuffer) host allocation that is allocated on first use, grown only when a larger transfer arrives, and reused for every subsequent transfer. It exists for the workload where per-call transfer overhead dominates: the same few tensor shapes moved host↔device hundreds or thousands of times (a video inference pipeline running the same ONNX models once per frame), where a fresh cuMemAllocHost_v2 per call would cost more than the transfer it enables.

Which API to use

The crate offers two shapes of staged transfer, and the difference between them is worth understanding because it is the difference between a 2.7x speedup and a 2x slowdown:

• upload_with / download_into — the producer writes its data directly into page-locked memory (or reads results directly out of it). No intermediate host copy exists at all, so the DMA engine transfers straight from/to the caller’s working memory. This is the fast path, and it wins at every size (measured 1.55x–2.67x, see below).

• upload / download — convenience wrappers taking an ordinary &[T] / &mut [T]. These must memcpy between the caller’s pageable slice and the pinned buffer, and that extra host copy is not free: a single-threaded host memcpy runs at roughly 10 GB/s, while the CUDA driver’s own pageable path pipelines its chunked staging copy against the DMA and so sustains more. Past about 1 MiB the wrapper therefore loses badly to just letting the driver do it. Rather than expose that as a footgun, these two methods auto-select: they stage through pinned memory only while the transfer is at or below StagingBuffer::auto_stage_max_bytes, and hand larger transfers to the driver’s pageable path. They are never slower than not using a StagingBuffer at all.

Measured (NVIDIA RTX A4000, driver 550.144.03, CUDA 12.4, sm_86)

Median of 100–400 repetitions, release build. “copy-in” is the convenience wrapper (pageable slice → pinned → DMA); “resident” is upload_with (producer fills pinned memory directly).

transfer	size	pageable	copy-in	resident
H2D 112×112×3 f32	150 KiB	18.2 µs	15.4 µs 1.18x	10.7 µs 1.71x
H2D 128×128×3 f32	196 KiB	21.9 µs	18.7 µs 1.17x	12.5 µs 1.75x
H2D 640×640×3 f32	4.7 MiB	311 µs	681 µs 0.46x	201 µs 1.55x
D2H 112×112×3 f32	150 KiB	17.8 µs	14.2 µs 1.25x	10.1 µs 1.77x
D2H 128×128×3 f32	196 KiB	28.1 µs	18.2 µs 1.54x	11.9 µs 2.37x
D2H 640×640×3 f32	4.7 MiB	509 µs	676 µs 0.75x	191 µs 2.67x

The copy-in column is exactly why the auto-select threshold exists.

Stream semantics

Every method here is synchronous from the caller’s thread: it returns only once the transfer has fully landed. That deliberately matches DeviceBuffer::copy_from_host / DeviceBuffer::copy_to_host, so a StagingBuffer is a drop-in replacement, and it is required for a reused staging buffer — the next call overwrites the same pinned bytes, so the DMA reading them must have finished.

All transfers are ordered against stream: an upload is enqueued after work already queued on stream, and a download observes the results of work already queued on stream. This holds on both sides of the auto-select threshold (the pageable fallback synchronises stream explicitly), so changing tensor size can never silently change ordering.

Example

let mut staging = StagingBuffer::new();
let mut d_input = DeviceBuffer::<f32>::alloc(3 * 640 * 640)?;

// Per frame: write preprocessed pixels straight into pinned memory.
staging.upload_with(&mut d_input, 3 * 640 * 640, &stream, |dst: &mut [f32]| {
    for (i, v) in dst.iter_mut().enumerate() {
        *v = i as f32; // real code: normalised pixel data
    }
})?;

Structs

StagingBuffer

A grow-on-demand, reused page-locked host buffer used to stage transfers.

StagingStats

Counters describing how a StagingBuffer has been used.

Constants

DEFAULT_AUTO_STAGE_MAX_BYTES

Default value of StagingBuffer::auto_stage_max_bytes: 512 KiB.

Traits

StagingPod

Types for which every bit pattern is a valid value.