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.

Features

• Inclusive and exclusive u32 prefix scan.
• Stable 2-bit LSD radix sort for u32 values.
• Stable 2-bit LSD radix sort for (u32 key, u32 value) pairs.
• 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()));

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.3 contains inclusive and exclusive scan, key-only radix sort, and stable key-value radix sort:

[dependencies]

wgpu-primitives = "0.3"

Algorithms

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

The 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 runs the same stable passes over interleaved KeyValue items. Values move with their keys during every scatter, so equal keys retain their original value order.

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, 2.287 billion elements/s for key-only sort, and 906.83 million pairs/s for stable key-value sort. See the full methodology, confidence intervals, backend comparison, and memory accounting.

Roadmap

Version 0.2 established the public GPU-buffer APIs, deterministic GPU tests, reusable workspace, cross-backend benchmarks, and the wgpu-primitives package name. Version 0.3 adds exclusive scan and stable key-value radix sort. The next work is ordered by how much it improves the crate as a reusable primitive library:

1. Measure kernels directly: add GPU timestamp-query benchmarks and per-pass profiling so optimization decisions are separated from command submission and synchronization cost.
2. Reduce runtime overhead: remove the radix sort's per-invocation uniform-buffer allocation and tune workgroup/radix configurations from measurements across DX12, Vulkan, and Metal.
3. Build derived primitives: implement stream compaction and selection on top of scan after the lower-level APIs and performance contracts are stable.
4. Broaden hardware evidence: publish reproducible benchmark reports from integrated and discrete GPUs, including crossover sizes, memory use, and resident versus round-trip behavior.

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

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