wgpu-primitives

Safe, composable GPU prefix scan and unsigned integer radix sort for Rust applications using wgpu.

wgpu-primitives continues the package previously published as wgpu-algorithms beginning with version 0.2.

Benchmark highlights

With data already resident on an NVIDIA RTX 4070 Ti SUPER, the GPU-buffer APIs delivered the following results at 100 million items (u32 values or KeyValue pairs):

These figures measure the composable resident-buffer path: command encoding, submission, primitive execution, and reusable workspace management are included, while host upload and readback are excluded. See Performance for smaller inputs and round-trip results.

The stable key-value row uses the NVIDIA Vulkan fast path added in version 0.4 and the median of three benchmark-process medians. It has 42.2% lower latency (1.73x speedup) than wgpu_sort at 100 million pairs for this bounded-key workload on the tested NVIDIA Vulkan system. With random full-width u32 keys, the same path measured 15.457 ms, a modest 2.8% latency reduction (1.03x speedup). The remaining headline rows are from version 0.3.

Features

• Inclusive and exclusive u32 prefix scan.
• Stable 2-bit LSD radix sort for u32 values.
• Stable LSD radix sort for (u32 key, u32 value) pairs, with a profiled NVIDIA Vulkan fast path.
• Convenience slice APIs that upload, execute, and read back.
• GPU-buffer APIs that record into an existing command encoder.
• Reusable internal scratch storage.
• No unsafe blocks in library code.

Usage

The convenience context is useful for standalone compute programs:

use wgpu_primitives::{Context, Scanner, Sorter};

#[tokio::main]
async fn main() -> Result<(), wgpu_primitives::Error> {
    let context = Context::init().await?;
    let mut scanner = Scanner::from_context(&context);
    let mut sorter = Sorter::from_context(&context);

    let prefixes = scanner.scan_exclusive(&[3, 1, 4, 1]).await?;
    let sorted = sorter.sort(&[10, 4, 7, 1]).await?;

    assert_eq!(prefixes, [0, 3, 4, 8]);
    assert_eq!(sorted, [1, 4, 7, 10]);
    Ok(())
}

Stable key-value sorting keeps payloads in their original order when keys compare equal:

use wgpu_primitives::{Context, KeyValue, KeyValueSorter};

#[tokio::main]
async fn main() -> Result<(), wgpu_primitives::Error> {
    let context = Context::init().await?;
    let mut sorter = KeyValueSorter::from_context(&context);
    let sorted = sorter
        .sort(&[
            KeyValue::new(2, 10),
            KeyValue::new(1, 20),
            KeyValue::new(2, 30),
        ])
        .await?;

    assert_eq!(
        sorted,
        [
            KeyValue::new(1, 20),
            KeyValue::new(2, 10),
            KeyValue::new(2, 30),
        ]
    );
    Ok(())
}

Applications that already own a wgpu device should reuse it:

let mut scanner = Scanner::new(&device, &queue);
let mut sorter = Sorter::new(&device, &queue);

scanner.record_scan(&mut encoder, &scan_input, &scan_output, item_count)?;
sorter.record_sort(&mut encoder, &sort_input, &sort_output, item_count)?;
queue.submit(Some(encoder.finish()));

KeyValueSorter::new_for_adapter enables measured adapter-specific kernels when adapter metadata is available. KeyValueSorter::new retains the portable path.

record_scan requires COPY_SRC on the input and COPY_DST | STORAGE on the output. record_sort requires STORAGE on both buffers.

Installation

Version 0.4 contains inclusive and exclusive scan, key-only radix sort, stable key-value radix sort, GPU timestamp profiling, and the adapter-selected NVIDIA Vulkan fast path:

[dependencies]

wgpu-primitives = "0.4"

Algorithms

The scan recursively computes per-workgroup inclusive prefixes, scans the workgroup totals, and propagates those totals back through the hierarchy.

The portable radix sort processes two bits per pass. Each of its 16 passes builds four per-workgroup histograms, scans them into global offsets, and stably scatters values between ping-pong buffers.

KeyValueSorter moves values with their keys during every scatter, so equal keys retain their original value order. On discrete NVIDIA Vulkan adapters with 32-wide subgroups, adapter-aware construction selects a dedicated 8-bit path. It builds all four byte histograms in one read, computes their prefixes together, and uses subgroup-assisted stable scatter with partition lookback. When both upper bytes are constant, indirect dispatch skips their identity scatters. Other hardware and backends retain the 4-bit NVIDIA Vulkan or portable 2-bit path.

Performance

Criterion measurements from an RTX 4070 Ti SUPER show why the GPU-buffer API is the primary interface. Resident execution keeps data on the GPU; round-trip execution includes upload, allocation, execution, and readback. GPU acceleration becomes valuable as the workload grows enough to amortize dispatch and transfer overhead.

At 100M items, resident throughput reached 17.96 billion elements/s for inclusive scan, 16.03 billion elements/s for exclusive scan, and 2.287 billion elements/s for key-only sort. The version 0.4 key-value path compares as follows:

See the base benchmark methodology, timestamp baseline, and direct wgpu_sort comparison. The comparison has a committed reproduction harness and machine-readable aggregate snapshot. The version 0.4 full-width profile records a clean-revision confirmation and the measured scatter optimization budget. The Jetson Orin Nano validation adds physical portable-Vulkan correctness and 4-TPC/8-TPC performance results.

GPU profiling

Version 0.4 adds capability-gated hardware timestamp queries to Scanner, Sorter, and KeyValueSorter. Normal execution does not allocate or resolve queries. Profiled calls return labeled dispatch spans, total dispatch time, and elapsed GPU time; the same compute passes carry stable labels for external tools such as NVIDIA Nsight Graphics.

Run the steady-state profile from a source checkout:

$env:WGPU_BACKEND = 'vulkan' # or 'dx12'
cargo run --release --example profile_primitives

At 100M items, the baseline profile attributed 74.8% of key-value dispatch time to stable scatter. The latest bounded-key profile spends 83.6% in two scatter passes, 16.3% in the all-byte histogram, and 0.1% in prefix setup. The finalized direct comparison harness measures the specialized path at 8.605 ms.

Roadmap

Version 0.4 adds per-dispatch GPU timestamp profiling, a measured NVIDIA Vulkan subgroup fast path, a reproducible pinned wgpu_sort comparison, and physical portable-Vulkan validation on Jetson Orin Nano. The next work is ordered by measured impact:

1. Validate more hardware: measure the specialized path on additional discrete NVIDIA Vulkan devices and driver versions.
2. Improve portability: implement and measure adaptive identity-pass elimination on non-subgroup paths without regressing other backends.
3. Build derived primitives: implement stream compaction and selection on top of scan.

New primitives should land with a GPU-buffer API, deterministic boundary tests, CPU-reference validation, and benchmark coverage.

Development

cargo fmt --all --check

cargo clippy --all-targets --all-features -- -D warnings

cargo test --lib --tests

cargo check --examples --benches

cargo package

cargo bench --bench scan -- --noplot

cargo bench --bench sort -- --noplot

cargo bench --bench key_value_sort -- --noplot

Run the independent, pinned wgpu_sort comparison on Windows with:

& .\benchmarks\wgpu-sort-comparison\run.ps1

GPU integration tests skip when no compatible adapter is available. CI installs Mesa's Vulkan software adapter so the shader paths execute on Linux.

License

MIT