Wavefront ray scheduling (stage queues)

Deep dive

How it works

Parts:

• Path pool — structure-of-arrays state per item (ray, throughput, RNG, depth), in memory.
• Stage kernels — extend/intersect, shade (one per expensive material), shadow/connect, generate; each handles one kind of work.
• Queues — preallocated arrays of item indices plus an atomic counter, one per stage input.
• Binner — counting sort or atomic append that routes items to the queue of their material or direction bucket.
• Compactor — writes survivors densely into the next round's queue, so dead items cost nothing.

1. Generate: create rays for fresh items, push them to the extend queue.
2. Extend: intersect every ray in the queue, record the hit, push (item, material) to that material's queue (append with atomic counter, or counting sort).
3. Shade: per material, run its kernel over its queue; sample the next direction; push a shadow ray; Russian roulette decides survival.
4. Shadow/connect: test shadow rays, add unoccluded light to the item's result.
5. Compact survivors into the next extend queue; refill from step 1 to keep the pool full.
6. Loop until no item is live and no work remains.

 generate ─► [extend Q] ─► intersect ─┬─► [mat A Q] ─► shade A ─┐
     ▲                                ├─► [mat B Q] ─► shade B ─┼─► [shadow Q] ─► connect
     │                                └─► [mat C Q] ─► shade C ─┘        │
     └───────── compact survivors / refill ◄──────────────────────────────┘

Appending needs a slot per item. Per-lane atomics serialise; Laine et al. combine them inside each 32-thread warp (ballot), keeping queue entries contiguous (40 % total speed in their car-paint scene). Eisenacher et al. sort large ray batches for traversal coherence, then defer shading until hits are sorted.

Code

CPU model: asserts both schedules give identical per-path results and computes lane utilisation for 8- and 32-lane groups under a stated cost model (a model, not hardware).

use std::time::Instant;

const K: usize = 4; // materials
const COST: [u32; K] = [1, 4, 16, 64]; // work units per shade call
const TERM: [u32; K] = [1 << 28, 1 << 30, 1 << 31, 3 << 30]; // stop prob = TERM / 2^32
const MAX_BOUNCES: u32 = 16;

#[derive(Clone, PartialEq, Debug)]
struct Path { rng: u64, acc: u64, bounces: u32 }

fn next(r: &mut u64) -> u64 { // xorshift64
    *r ^= *r << 13; *r ^= *r >> 7; *r ^= *r << 17; *r
}
fn extend(p: &mut Path) -> usize { (next(&mut p.rng) % K as u64) as usize } // "intersect": which material
fn shade(p: &mut Path, m: usize) -> bool { // returns "still alive"
    for _ in 0..COST[m] {
        p.acc = p.acc.rotate_left(5) ^ next(&mut p.rng).wrapping_mul(0x9E37_79B9_7F4A_7C15);
    }
    p.bounces += 1;
    ((next(&mut p.rng) >> 32) as u32) >= TERM[m] && p.bounces < MAX_BOUNCES
}

fn megakernel(ps: &mut [Path]) { // whole loop per path
    for p in ps { loop { let m = extend(p); if !shade(p, m) { break; } } }
}

/// Lockstep model of `w`-lane groups running the megakernel: a step costs the
/// sum of COST over the distinct materials present among live lanes.
fn mega_model(ps: &mut [Path], w: usize) -> (u64, u64) {
    let (mut useful, mut slots) = (0, 0);
    for g in ps.chunks_mut(w) {
        let mut alive = vec![true; g.len()];
        while alive.iter().any(|&a| a) {
            let mut seen = [false; K];
            for l in 0..g.len() {
                if !alive[l] { continue; }
                let p = &mut g[l];
                let m = extend(p);
                seen[m] = true; useful += COST[m] as u64;
                alive[l] = shade(p, m);
            }
            slots += (w as u32 * (0..K).filter(|&m| seen[m]).map(|m| COST[m]).sum::<u32>()) as u64;
        }
    }
    (useful, slots)
}

/// Wavefront: per bounce, extend all live paths, counting-sort them by material,
/// shade each bin, compact survivors into the next queue. Model: each bin is
/// run in groups of `w` lanes of one material, so only the last group per bin idles.
fn wavefront(ps: &mut [Path], w: usize) -> (u64, u64) {
    let (mut useful, mut slots) = (0, 0);
    let mut live: Vec<u32> = (0..ps.len() as u32).collect();
    let mut binned = vec![0u32; ps.len()];
    while !live.is_empty() {
        let mats: Vec<usize> = live.iter().map(|&i| extend(&mut ps[i as usize])).collect();
        let mut start = [0usize; K + 1]; // counting sort by material
        for &m in &mats { start[m + 1] += 1; }
        for m in 0..K { start[m + 1] += start[m]; }
        let mut fill = start;
        for (&i, &m) in live.iter().zip(&mats) { binned[fill[m]] = i; fill[m] += 1; }
        live.clear(); // survivors are appended in bin order (stream compaction)
        for m in 0..K {
            let bin = &binned[start[m]..start[m + 1]];
            useful += bin.len() as u64 * COST[m] as u64;
            slots += (bin.len().div_ceil(w) * w) as u64 * COST[m] as u64;
            for &i in bin { if shade(&mut ps[i as usize], m) { live.push(i); } }
        }
    }
    (useful, slots)
}

fn main() {
    let n = 1 << 20;
    let init: Vec<Path> = (0..n as u64)
        .map(|i| Path { rng: (i + 1).wrapping_mul(0x2545_F491_4F6C_DD1D) | 1, acc: 0, bounces: 0 })
        .collect();
    let (mut a, mut b) = (init.clone(), init.clone());
    let t = Instant::now(); megakernel(&mut a); let ta = t.elapsed();
    let t = Instant::now(); wavefront(&mut b, 32); let tb = t.elapsed();
    assert_eq!(a, b); // identical per-path results
    println!("avg bounces {:.2}", a.iter().map(|p| p.bounces as f64).sum::<f64>() / n as f64);
    println!("time (measured, scalar, 1 thread): megakernel {ta:?}, wavefront {tb:?}");
    for w in [8, 32] {
        let (mut c, mut d) = (init.clone(), init.clone());
        let (u1, s1) = mega_model(&mut c, w);
        let (u2, s2) = wavefront(&mut d, w);
        assert_eq!(c, a); assert_eq!(d, a); assert_eq!(u1, u2);
        println!("lane model w={w}: megakernel {:.1}% , wavefront {:.1}% lane utilisation", 100.0 * u1 as f64 / s1 as f64, 100.0 * u2 as f64 / s2 as f64);
    }
}

Measured on this Mac (Apple M4 Pro, release build, one thread, 1M paths, 3 runs): megakernel 86 to 87 ms, wavefront 116 to 121 ms. Lane model, same run: w=8 megakernel 19.8 % vs wavefront about 100 %; w=32 10.9 % vs about 100 %. Scalar CPU code has no divergence to remove, so the wavefront only pays its bookkeeping.

Complexity and constants

• Work: sum of the path lengths, as in the megakernel. Per round: O(L) for L live items, plus O(K) for K queues (counting sort).
• Cost drivers: bytes of path state read and written per stage (Laine et al.: 212 B per path, a pool of 2^20 paths = 212 MB, each queue 4 MB, on a 2013 Tesla K20); the number of kernel launches per round (stages x materials); how full each bin is.
• In the model, a bin run in w-lane groups wastes at most w - 1 lanes, so only each bin's last group idles.

Efficiency

Cost: per item per round, a read and write of its state in each stage it visits, one queue append, and the stage's compute (212 B state per path, K20 design). The wavefront trades divergence cost for memory traffic.

Levers:

• Structure-of-arrays path state — 80 % total speedup over array-of-structures (Laine et al. 2013, their car-paint scene) — whenever stages touch a few fields of many items.
• Warp-aggregated queue appends — 40 % total speed (same paper, same scene) — GPU queues written by many lanes.
• Separate kernels only for expensive materials, one shared kernel for cheap ones — avoids launch and queue overhead (Laine et al.) — many materials of very different cost.
• Sort queues by ray direction or origin cell — better traversal coherence (Eisenacher et al. 2013) — large, out-of-core ray batches.
• Hardware thread reordering (NVIDIA SER) — 20 to 50 % for a path tracer after a code change (NVIDIA blog, 2022) — GPUs and APIs that expose it.

Naive → standard → state of the art:

• Megakernel per path · divergence and the largest shader's registers set speed (23 % thread use in one scene).
• Wavefront with stage queues and material bins · 53 % thread use in the same scene, 36 to 221 % faster (Laine et al.).
• Wavefront plus ray sorting, or hardware thread reordering · more coherence (Eisenacher et al.; NVIDIA SER).

Implementation notes

• Size every queue for the whole pool so appends cannot overflow; clear by zeroing the counter (Laine et al.).
• Double-buffer the ray queue between rounds (pbrt-v4 uses two ray queues indexed by depth parity).
• Queues hold item indices; stages read the pool, so state is not copied between stages.
• Per-item RNG state lives in the pool, so the result does not depend on the order of the queue (the check in the code).

Pitfalls

• Too many queues: one kernel per material makes launch cost and under-full bins dominate when most materials are rare.
• Cheap materials: launch and queue overhead outweigh the gain, so Laine et al. merge them into one kernel.
• Floating-point sums that depend on queue order change results between runs; fix an order or accept noise.

Variants and alternatives

In production

• Rendering: pbrt-v4 (queues for rays, material evaluation, shadow rays; same code on CPU threads and GPU), Blender Cycles (each path stores its next kernel; active-index and sorted-index kernels build queues, sorted by a shader key).
• Particle transport: OpenMC event_based mode with max_particles_in_flight.
• Databases: MonetDB/X100 passes selection vectors between vectorized primitives instead of copying survivors (Boncz et al. 2005).
• Graphs: GPU breadth-first search manages its fine-grained tasks with prefix sums (Merrill et al. 2012).

Implementations

Sources

• Laine, Karras, Aila, Megakernels Considered Harmful: Wavefront Path Tracing on GPUs, HPG 2013. paper · doi
• Eisenacher, Nichols, Selle, Burley, Sorted Deferred Shading for Production Path Tracing, Computer Graphics Forum 32(4), 2013. doi
• NVIDIA, Improve Shader Performance and In-Game Frame Rates with Shader Execution Reordering, 2022. blog
• Hamilton, Slattery, Evans, Multigroup Monte Carlo on GPUs: Comparison of history- and event-based algorithms, Annals of Nuclear Energy 113, 2018. doi
• Boncz, Zukowski, Nes, MonetDB/X100: Hyper-Pipelining Query Execution, CIDR 2005. pdf
• Merrill, Garland, Grimshaw, Scalable GPU Graph Traversal, PPoPP 2012. paper
• OpenMC documentation, Settings (event_based, max_particles_in_flight). docs
• pbrt-v4 and Blender Cycles source: workqueue.h · kernel.h