indicatrix 0.8.0

A physically-based spectral gemstone renderer that turns GemCAD-style cutting instructions into rendered output.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
// The wavefront transport pipeline -- an ALTERNATIVE, selectable path
// to the megakernel (`spectral_transport.wgsl`'s `transport_main`), not a replacement.
// Ray state lives in storage buffers instead of per-thread local variables, and one
// bounce of every still-alive ray is a SEPARATE dispatch, so a divergent material
// (uniaxial/biaxial boundary solves, up to 7 exit-event polyhedron probes) does not
// keep every thread in a warp/wavefront resident through the whole bounce loop --
// threads that die early free their warp slot for the next `wavefront_generate` chunk's
// rays instead of idling until their warp's slowest ray finishes. See
// `renderer::gpu::frame`'s module doc comment ("Wavefront pipeline") for the buffer
// layout, per-bounce traffic estimate, and when to prefer this over the megakernel;
// `renderer::gpu_backend`/`renderer::gpu::frame::GpuPipelineKind` select it, default
// `Megakernel` until measured.
//
// # Kernel sequence (driven by `renderer::gpu::frame`'s host dispatch loop)
//
// 1. `wavefront_generate`: one invocation per (pixel, sample) tuple in the chunk (same
//    domain as the megakernel's own dispatch). Draws the exact same RNG sequence as the
//    megakernel's prologue (`transport_generate_ray`, `shaders/transport_bounce.wgsl`)
//    and writes the ray's initial state into the SoA buffers below, plus seeds
//    `active_ray_indices` with the identity permutation (every ray starts alive).
// 2. `wavefront_bounce`, called once per bounce round over the CURRENT active list:
//    each invocation loads one ray's state, recomputes that ray's material-derived
//    hoisted constants (`transport_hoist_ray_constants`) fresh from its own stored
//    `lambdas` -- see that function's own doc comment for why this is cheap-enough and
//    bit-identical -- calls the SAME `transport_bounce_step`
//    (`shaders/transport_bounce.wgsl`) the megakernel's loop calls, and either writes
//    the ray's state back (still alive) or finalizes it via `transport_finalize_ray`
//    (miss or Russian-roulette death) and clears its alive flag.
// 3. `wavefront_compact_scan` + `wavefront_compact_scatter`: order-preserving stream
//    compaction of the active list by alive flag -- see their own doc comments for the
//    two-pass block-prefix-sum this uses and why it stays deterministic.
// 4. Repeat 2-3 until the active count reaches zero or `params.max_bounces` bounce
//    rounds have run.
// 5. `wavefront_finalize_survivors`, over whatever remains of the active list after
//    step 4's loop ends by exhausting the bounce budget (NOT by every ray dying): mirrors
//    the megakernel's own unconditional tail -- a ray that survives every bounce without
//    ever escaping or dying still gets exactly one `out_xyz` write, with whatever
//    `radiance` it accumulated (typically from NEE contributions along the way, or zero).
// 6. `reduce_xyz.wgsl`'s existing `reduce_xyz_main`, UNCHANGED: every finalize site
//    above (`wavefront_bounce`'s death branch, `wavefront_finalize_survivors`) writes
//    into the SAME `out_xyz`/`out_radiance`/`out_lambdas`/`out_path_pdf`/`out_compat`
//    bindings (inherited from `transport_bounce.wgsl`'s shared prelude, bindings 4-7/9)
//    the megakernel writes into, at the SAME `idx` a ray would have used in the
//    megakernel (`pixel_in_chunk * camera.num_samples + sample_in_chunk`) -- so no
//    wavefront-specific output buffer or reduction kernel is needed at all.
//
// # Determinism
//
// Mirrors `spectral_transport.wgsl`'s own header comment: a ray's result depends only
// on (pixel, sample, seed, bounce) -- never on which OTHER rays share its dispatch,
// which workgroup it lands in, or what order `wavefront_bounce`/compaction process rays
// in. `transport_bounce_step` reads/writes only ITS OWN invocation's local copy of one
// ray's state (loaded from, and written back to, that ray's own fixed slot,
// `ray_idx`, in the SoA buffers below) -- no cross-thread communication, no atomics
// anywhere in the physics itself. `out_xyz[ray_idx]` is written exactly once, by
// whichever kernel invocation finalizes that ray, so which bounce round or which
// workgroup did it never affects the value written. `wavefront_compact_*`'s prefix sum
// is likewise a PURE function of the (fixed-order) alive-flag array, not of scheduling,
// so re-running the whole pipeline on identical input reproduces identical active-list
// orderings and identical `out_xyz` contents every time -- see
// `renderer::gpu::transport_check::run_wavefront_determinism` (or
// `determinism_check.rs`) for the two-runs-identical check, and
// `run_pipeline_equivalence` for the megakernel-vs-wavefront bit-identity check.
//
// # Ray-state SoA buffer layout
//
// Bindings 0-14 (camera/params/material/planes/output buffers/facet finishes/HDR
// environment) are `transport_bounce.wgsl`'s shared scene bindings, unchanged --
// EXACTLY the megakernel's own binding shapes, so `renderer::gpu::frame`'s host code
// builds this pipeline's scene bind group with the same buffers/helper
// (`build_chunk_bind_group`-style) it already uses for the megakernel. Bindings 15+
// below are wavefront-only. Every per-ray array is flat: a scalar field is indexed by
// `ray_idx` directly, an 8-channel field by `ray_idx * 8u + k`. Sized for
// `chunk_pixels * spp` rays (`chunk_rays`, `wf_params.chunk_rays`).
//
// Deliberately NOT stored per-ray (see `transport_hoist_ray_constants`'s own doc
// comment): the material-derived hoisted constants (`c_axis`, `is_anisotropic`,
// `is_biaxial`, the biaxial axis frame, hero indices, the four per-channel
// dispersion/absorption arrays, the studio rig directions) -- `wavefront_bounce`
// recomputes them fresh every bounce from `ray_lambdas`/`material` instead, trading a
// modest amount of recomputed ALU (a handful of dispersion evaluations and band
// summations per bounce) for a meaningfully smaller ray-state buffer and simpler
// generate/bounce kernels. Revisit if profiling shows this recomputation dominates.
struct WavefrontParams {
    chunk_rays: u32,
    active_count: u32,
    bounce: u32,
    workgroup_count: u32,
}

@group(0) @binding(15) var<uniform> wf_params: WavefrontParams;

@group(0) @binding(16) var<storage, read_write> ray_origin: array<vec4<f32>>;
@group(0) @binding(17) var<storage, read_write> ray_dir: array<vec4<f32>>;
@group(0) @binding(18) var<storage, read_write> ray_k: array<vec4<f32>>;
@group(0) @binding(19) var<storage, read_write> ray_prev_plane_normal: array<vec4<f32>>;
// Bit 0 `inside_gem`, bit 1 `is_extraordinary`, bit 2 `have_prev_plane_normal`, bit 3
// `path_escaped`, bit 4 `alive` -- see the `RAY_FLAG_*` constants below.
@group(0) @binding(20) var<storage, read_write> ray_flags: array<u32>;
@group(0) @binding(21) var<storage, read_write> ray_stokes: array<vec4<f32>>;
@group(0) @binding(22) var<storage, read_write> ray_radiance: array<f32>;
@group(0) @binding(23) var<storage, read_write> ray_path_pdf: array<f32>;
@group(0) @binding(24) var<storage, read_write> ray_split_radiance: array<f32>;
@group(0) @binding(25) var<storage, read_write> ray_compat: array<u32>;
@group(0) @binding(26) var<storage, read_write> ray_pending_light_mis: array<f32>;
@group(0) @binding(27) var<storage, read_write> ray_lambdas: array<f32>;
@group(0) @binding(28) var<storage, read_write> ray_seed: array<u32>;

// The current bounce round's live ray indices, `wf_params.active_count` of them
// (`wavefront_generate` seeds this with the identity permutation `0..chunk_rays`; each
// `wavefront_compact_scatter` round replaces it with that round's survivors -- see
// `renderer::gpu::frame`'s host loop for the buffer-swap between rounds).
@group(0) @binding(29) var<storage, read_write> active_ray_indices: array<u32>;
// `wavefront_compact_scatter`'s destination for this round's survivors -- becomes
// `active_ray_indices` for the NEXT bounce round after the host swaps the two buffer
// bindings (or copies, on a backend where an in-place binding swap is inconvenient).
@group(0) @binding(30) var<storage, read_write> active_ray_indices_next: array<u32>;
// Per-(active-list-position) LOCAL exclusive prefix count of alive rays within that
// ray's own workgroup -- `wavefront_compact_scan`'s output, `wavefront_compact_scatter`'s
// input.
@group(0) @binding(31) var<storage, read_write> compact_local_offset: array<u32>;
// One entry per workgroup: that workgroup's total alive-ray count --
// `wavefront_compact_scan`'s output. The host reads this back, computes its EXCLUSIVE
// prefix sum on the CPU (a simple, deterministic sequential loop over
// `wf_params.workgroup_count` entries -- see `renderer::gpu::frame`'s host loop), and
// uploads the result as `compact_block_offset` below. The same readback also gives the
// host this round's total survivor count (the last workgroup's offset plus its own
// count) for the "stop early when the live count hits zero" check.
@group(0) @binding(32) var<storage, read_write> compact_block_alive_count: array<u32>;
// The host-computed exclusive prefix sum of `compact_block_alive_count` -- see that
// binding's own doc comment.
@group(0) @binding(33) var<storage, read> compact_block_offset: array<u32>;
// The interior direction `ray_pending_light_mis`'s phase pdf was evaluated
// at -- see `transport_bounce_step`'s own doc comment on its matching
// `pending_light_mis_dir` parameter (`shaders/transport_bounce.wgsl`). `vec4` (not
// `vec3`) for the same storage-buffer-alignment reason `ray_origin`/`ray_dir`/`ray_k`/
// `ray_prev_plane_normal` above already use; `.w` is always `0.0` and unread.
//
// Rust-side storage buffer: `renderer::gpu::frame::WavefrontRayBuffers::pending_light_mis_dir`
// -- `chunk_rays` `[f32; 4]` entries, allocated/reset alongside `pending_light_mis`
// (`WavefrontRayBuffers::new`), bound at binding 34 in both `wavefront_generate`'s and
// `wavefront_bounce`'s bind groups (`wavefront_generate_bindings`/
// `wavefront_bounce_bind_group`). The megakernel path (`spectral_transport.wgsl`,
// `GpuPipelineKind`'s default) carries this same state in a per-invocation local
// variable instead (see `spectral_transport.wgsl`'s own `pending_light_mis_dir`).
@group(0) @binding(34) var<storage, read_write> ray_pending_light_mis_dir: array<vec4<f32>>;
// The running total of every Henyey-Greenstein scattering-point NEE deposit for this ray
// so far, each already integrated to XYZ at ITS OWN moment's `path_pdf`/`compat` -- see
// `transport_bounce_step`'s matching `nee_xyz` parameter and `nee_contribution_hg_scatter`'s
// own doc comment (`shaders/transport_bounce/06_nee_sampling.wgsl`).
// `vec4` (not `vec3`) for the same storage-buffer-alignment reason
// `ray_pending_light_mis_dir` above already uses; `.w` is always `0.0` and unread.
//
// The Rust side is `WavefrontRayBuffers::nee_xyz` (`renderer::gpu::frame::scene_buffers`), zeroed
// alongside `pending_light_mis_dir` and bound at binding 35 in `wavefront_generate`'s and
// `wavefront_bounce`'s bind groups (`renderer::gpu::frame::bind_groups`). The megakernel
// path (`spectral_transport.wgsl`, `GpuPipelineKind`'s default) needs no such buffer --
// its whole ray lives in one dispatch, so `nee_xyz` is a plain per-invocation local
// variable there instead, exactly like `pending_light_mis_dir`'s own precedent.
@group(0) @binding(35) var<storage, read_write> ray_nee_xyz: array<vec4<f32>>;

const RAY_FLAG_INSIDE_GEM: u32 = 0x1u;
const RAY_FLAG_IS_EXTRAORDINARY: u32 = 0x2u;
const RAY_FLAG_HAVE_PREV_PLANE_NORMAL: u32 = 0x4u;
const RAY_FLAG_PATH_ESCAPED: u32 = 0x8u;
const RAY_FLAG_ALIVE: u32 = 0x10u;

// ---------------------------------------------------------------------------------
// 1. wavefront_generate
// ---------------------------------------------------------------------------------

@compute @workgroup_size(64)
fn wavefront_generate(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx >= wf_params.chunk_rays) {
        return;
    }

    // Identical RNG draws to the megakernel's own prologue -- see
    // `transport_generate_ray`'s doc comment (`shaders/transport_bounce.wgsl`).
    let gen = transport_generate_ray(idx);

    ray_origin[idx] = vec4<f32>(gen.origin, 0.0);
    ray_dir[idx] = vec4<f32>(gen.dir, 0.0);
    ray_k[idx] = vec4<f32>(gen.dir, 0.0);
    ray_prev_plane_normal[idx] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
    // Every ray starts outside the gem, ordinary-mode, no previous plane normal, not
    // escaped -- ALIVE is the only bit set, mirroring `transport_main`'s own initial
    // `inside_gem = false; is_extraordinary = false; have_prev_plane_normal = false;
    // path_escaped = false;`.
    ray_flags[idx] = RAY_FLAG_ALIVE;
    for (var k: u32 = 0u; k < NUM_CHANNELS; k = k + 1u) {
        ray_stokes[idx * 8u + k] = vec4<f32>(1.0, 0.0, 0.0, 0.0);
        ray_radiance[idx * 8u + k] = 0.0;
        ray_path_pdf[idx * 8u + k] = 1.0;
        ray_split_radiance[idx * 8u + k] = 0.0;
        ray_compat[idx * 8u + k] = 0xFFu;
        ray_lambdas[idx * 8u + k] = gen.lambdas[k];
    }
    ray_pending_light_mis[idx] = 0.0;
    ray_pending_light_mis_dir[idx] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
    ray_nee_xyz[idx] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
    ray_seed[idx] = gen.seed0;

    // Identity permutation: every ray starts alive, in ascending `ray_idx` order --
    // see this file's header comment on why compaction order matters (determinism
    // auditability) even though a ray's own result never depends on it.
    active_ray_indices[idx] = idx;
}

// ---------------------------------------------------------------------------------
// 2. wavefront_bounce
// ---------------------------------------------------------------------------------

@compute @workgroup_size(64)
fn wavefront_bounce(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_index) local_index: u32,
) {
    // The workgroup-shared plane cache applies here exactly as in the
    // megakernel -- `transport_bounce_step` (via `intersect_ray`,
    // `shaders/transport_bounce.wgsl`) reads `planes_shared`/`planes` the same way
    // regardless of which entry point filled it, so this dispatch fills it itself,
    // before the out-of-range early return below, for the same
    // every-invocation-reaches-the-barrier reason `transport_main` does.
    let plane_count = arrayLength(&planes);
    if (plane_count <= PLANES_SHARED_CAPACITY) {
        var fill_i = local_index;
        loop {
            if (fill_i >= plane_count) {
                break;
            }
            planes_shared[fill_i] = planes[fill_i];
            fill_i = fill_i + 64u;
        }
    }
    workgroupBarrier();

    let active_pos = gid.x;
    if (active_pos >= wf_params.active_count) {
        return;
    }
    let ray_idx = active_ray_indices[active_pos];

    var lambdas: array<f32, 8>;
    for (var k: u32 = 0u; k < NUM_CHANNELS; k = k + 1u) {
        lambdas[k] = ray_lambdas[ray_idx * 8u + k];
    }
    let seed0 = ray_seed[ray_idx];
    // See `transport_hoist_ray_constants`'s own doc comment: recomputed fresh every
    // bounce round from this ray's own stored `lambdas`, rather than stored.
    let rc = transport_hoist_ray_constants(lambdas);
    // The megakernel's `observer = -gen.dir` (`spectral_transport.wgsl`): recomputed from
    // this ray's own deterministic generation rather than stored, like `rc`.
    let observer = -transport_generate_ray(ray_idx).dir;

    var current_origin = ray_origin[ray_idx].xyz;
    var current_dir = ray_dir[ray_idx].xyz;
    var current_k = ray_k[ray_idx].xyz;
    var prev_plane_normal = ray_prev_plane_normal[ray_idx].xyz;
    let flags = ray_flags[ray_idx];
    var inside_gem = (flags & RAY_FLAG_INSIDE_GEM) != 0u;
    var is_extraordinary = (flags & RAY_FLAG_IS_EXTRAORDINARY) != 0u;
    var have_prev_plane_normal = (flags & RAY_FLAG_HAVE_PREV_PLANE_NORMAL) != 0u;
    var path_escaped = (flags & RAY_FLAG_PATH_ESCAPED) != 0u;

    var stokes: array<vec4<f32>, 8>;
    var radiance: array<f32, 8>;
    var path_pdf: array<f32, 8>;
    var split_radiance: array<f32, 8>;
    var compat: array<u32, 8>;
    for (var k: u32 = 0u; k < NUM_CHANNELS; k = k + 1u) {
        stokes[k] = ray_stokes[ray_idx * 8u + k];
        radiance[k] = ray_radiance[ray_idx * 8u + k];
        path_pdf[k] = ray_path_pdf[ray_idx * 8u + k];
        split_radiance[k] = ray_split_radiance[ray_idx * 8u + k];
        compat[k] = ray_compat[ray_idx * 8u + k];
    }
    var pending_light_mis = ray_pending_light_mis[ray_idx];
    var pending_light_mis_dir = ray_pending_light_mis_dir[ray_idx].xyz;
    var nee_xyz = ray_nee_xyz[ray_idx].xyz;

    // One bounce of the SAME shared body `transport_main`'s loop calls -- see
    // `transport_bounce_step`'s own doc comment (`shaders/transport_bounce.wgsl`).
    let status = transport_bounce_step(
        wf_params.bounce, seed0, &lambdas, rc.c_axis, rc.birefringence_delta, rc.is_anisotropic, rc.is_biaxial,
        rc.biax_ax0, rc.biax_ax1, rc.biax_ax2, rc.n_o_hero_seed, rc.n_beta_hero, rc.n_alpha_hero, rc.n_gamma_hero,
        rc.n_e_hero_seed, rc.n_o_hoisted, rc.alpha_o_hoisted, rc.alpha_e_hoisted, rc.alpha_beta_hoisted,
        rc.studio_key_dir, rc.studio_fill_dir, rc.studio_sin_lp, observer,
        &stokes, &radiance, &path_pdf, &current_origin, &current_dir, &current_k,
        &inside_gem, &is_extraordinary, &prev_plane_normal, &have_prev_plane_normal,
        &split_radiance, &compat, &nee_xyz, &path_escaped, &pending_light_mis, &pending_light_mis_dir,
    );

    if (status == BOUNCE_STATUS_TERMINATE) {
        // Mirrors `transport_main`'s own unconditional tail call at the natural end of
        // its bounce loop -- writes `out_xyz`/the debug buffers at `ray_idx`, the same
        // slot the megakernel would have used for this (pixel, sample) tuple.
        transport_finalize_ray(ray_idx, lambdas, &radiance, split_radiance, path_pdf, compat, path_escaped, nee_xyz);
        ray_flags[ray_idx] = flags & ~RAY_FLAG_ALIVE;
        return;
    }

    // Still alive: write the mutated state back for the next bounce round.
    ray_origin[ray_idx] = vec4<f32>(current_origin, 0.0);
    ray_dir[ray_idx] = vec4<f32>(current_dir, 0.0);
    ray_k[ray_idx] = vec4<f32>(current_k, 0.0);
    ray_prev_plane_normal[ray_idx] = vec4<f32>(prev_plane_normal, 0.0);
    var new_flags = RAY_FLAG_ALIVE;
    if (inside_gem) {
        new_flags = new_flags | RAY_FLAG_INSIDE_GEM;
    }
    if (is_extraordinary) {
        new_flags = new_flags | RAY_FLAG_IS_EXTRAORDINARY;
    }
    if (have_prev_plane_normal) {
        new_flags = new_flags | RAY_FLAG_HAVE_PREV_PLANE_NORMAL;
    }
    if (path_escaped) {
        new_flags = new_flags | RAY_FLAG_PATH_ESCAPED;
    }
    ray_flags[ray_idx] = new_flags;
    for (var k: u32 = 0u; k < NUM_CHANNELS; k = k + 1u) {
        ray_stokes[ray_idx * 8u + k] = stokes[k];
        ray_radiance[ray_idx * 8u + k] = radiance[k];
        ray_path_pdf[ray_idx * 8u + k] = path_pdf[k];
        ray_split_radiance[ray_idx * 8u + k] = split_radiance[k];
        ray_compat[ray_idx * 8u + k] = compat[k];
    }
    ray_pending_light_mis[ray_idx] = pending_light_mis;
    ray_pending_light_mis_dir[ray_idx] = vec4<f32>(pending_light_mis_dir, 0.0);
    ray_nee_xyz[ray_idx] = vec4<f32>(nee_xyz, 0.0);
}

// ---------------------------------------------------------------------------------
// 3. wavefront_compact: order-preserving stream compaction of `active_ray_indices` by
//    each ray's current alive flag, via a deterministic two-pass block prefix sum.
//    `wavefront_compact_scan` computes, per workgroup, each ray's LOCAL exclusive
//    prefix count of alive predecessors within that same workgroup plus the
//    workgroup's own total; the host then computes the EXCLUSIVE prefix sum of every
//    workgroup's total (a trivial sequential CPU loop over
//    `wf_params.workgroup_count` values) and uploads it as `compact_block_offset`;
//    `wavefront_compact_scatter` then writes each alive ray to
//    `compact_block_offset[workgroup] + compact_local_offset[active_pos]`. Two
//    dispatches plus one small CPU pass rather than a single-dispatch global atomic
//    counter, deliberately: an atomic counter gives a valid COUNT but a
//    scheduling-dependent (non-reproducible) ORDER, and reproducible ordering is worth
//    keeping even though (see this file's header comment) no ray's OWN result depends
//    on it -- it keeps the wavefront pipeline's intermediate state, not just its final
//    `out_xyz`, auditable and diffable across runs.
// ---------------------------------------------------------------------------------

var<workgroup> compact_wg_alive: array<u32, 64>;
var<workgroup> compact_wg_offset: array<u32, 64>;
var<workgroup> compact_wg_total: u32;

@compute @workgroup_size(64)
fn wavefront_compact_scan(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_index) local_index: u32,
    @builtin(workgroup_id) wg_id: vec3<u32>,
) {
    let active_pos = gid.x;
    var alive: u32 = 0u;
    if (active_pos < wf_params.active_count) {
        let ray_idx = active_ray_indices[active_pos];
        alive = (ray_flags[ray_idx] >> 4u) & 1u;
    }
    compact_wg_alive[local_index] = alive;
    workgroupBarrier();

    // A single lane does the workgroup's own sequential exclusive prefix sum --
    // correctness-first: only 64 entries, and this scan runs at most once per bounce
    // round, not per bounce PER RAY. A parallel (Hillis-Steele) scan is a documented
    // future optimization if profiling ever shows this serialized step matters.
    if (local_index == 0u) {
        var running: u32 = 0u;
        for (var i: u32 = 0u; i < 64u; i = i + 1u) {
            compact_wg_offset[i] = running;
            running = running + compact_wg_alive[i];
        }
        compact_wg_total = running;
    }
    workgroupBarrier();

    if (active_pos < wf_params.active_count) {
        compact_local_offset[active_pos] = compact_wg_offset[local_index];
    }
    if (local_index == 0u) {
        compact_block_alive_count[wg_id.x] = compact_wg_total;
    }
}

@compute @workgroup_size(64)
fn wavefront_compact_scatter(@builtin(global_invocation_id) gid: vec3<u32>) {
    let active_pos = gid.x;
    if (active_pos >= wf_params.active_count) {
        return;
    }
    let ray_idx = active_ray_indices[active_pos];
    let alive = (ray_flags[ray_idx] >> 4u) & 1u;
    if (alive == 0u) {
        return;
    }
    let workgroup_idx = active_pos / 64u;
    let dest = compact_block_offset[workgroup_idx] + compact_local_offset[active_pos];
    active_ray_indices_next[dest] = ray_idx;
}

// ---------------------------------------------------------------------------------
// 5. wavefront_finalize_survivors -- see this file's header comment, step 5. Mirrors
//    `transport_main`'s unconditional tail for a ray that survives every bounce
//    without ever hitting `transport_bounce_step`'s `BOUNCE_STATUS_TERMINATE` --
//    `wf_params.active_count`/`active_ray_indices` here are whatever remained after
//    the host's bounce-round loop stopped because it ran out of bounce budget, NOT
//    because every ray died (that path is already fully finalized inside
//    `wavefront_bounce` itself).
// ---------------------------------------------------------------------------------

@compute @workgroup_size(64)
fn wavefront_finalize_survivors(@builtin(global_invocation_id) gid: vec3<u32>) {
    let active_pos = gid.x;
    if (active_pos >= wf_params.active_count) {
        return;
    }
    let ray_idx = active_ray_indices[active_pos];

    var radiance: array<f32, 8>;
    var path_pdf: array<f32, 8>;
    var split_radiance: array<f32, 8>;
    var compat: array<u32, 8>;
    var lambdas: array<f32, 8>;
    for (var k: u32 = 0u; k < NUM_CHANNELS; k = k + 1u) {
        radiance[k] = ray_radiance[ray_idx * 8u + k];
        path_pdf[k] = ray_path_pdf[ray_idx * 8u + k];
        split_radiance[k] = ray_split_radiance[ray_idx * 8u + k];
        compat[k] = ray_compat[ray_idx * 8u + k];
        lambdas[k] = ray_lambdas[ray_idx * 8u + k];
    }
    let path_escaped = (ray_flags[ray_idx] & RAY_FLAG_PATH_ESCAPED) != 0u;
    let nee_xyz = ray_nee_xyz[ray_idx].xyz;

    transport_finalize_ray(ray_idx, lambdas, &radiance, split_radiance, path_pdf, compat, path_escaped, nee_xyz);
    ray_flags[ray_idx] = ray_flags[ray_idx] & ~RAY_FLAG_ALIVE;
}