lbm-gpu: one atomic per workgroup for the chunk's largest speed

Every cell did an atomicMax on the same word; a workgroup reduction first
leaves one per 256 cells. Same result, less contention on a single
address.

Part of #529.
This commit is contained in:
Christian Svensson authored and Irakli Kourtis committed 2026-10-17 07:23:40 -04:00
1 parent 48cc80bef5
commit 1e0c21d088
1 file changed
+17 -3
+17 -3
View File
@@ -59,6 +59,9 @@ struct Params {
@group(0) @binding(15) var<storage, read_write> partials: array<f32>;
var<workgroup> red: array<vec4<f32>, 256>;
// The largest squared speed of a workgroup's cells (collide): one atomic per workgroup, not
// per cell, which serialise on a single address.
var<workgroup> speed: array<f32, 256>;
fn c_of(q: u32) -> vec3<i32> {
switch q {
@@ -109,8 +112,19 @@ fn is_fluid(cell: u32) -> bool {
@compute @workgroup_size(256)
fn collide(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
let cell = cell_of(lid, wg);
if cell >= params.n { return; }
if !is_fluid(cell) { return; }
speed[lid] = 0.0;
if cell < params.n && is_fluid(cell) {
collide_cell(cell, lid);
}
workgroupBarrier();
for (var k = 128u; k > 0u; k = k >> 1u) {
if lid < k { speed[lid] = max(speed[lid], speed[lid + k]); }
workgroupBarrier();
}
if lid == 0u { atomicMax(&stats[0], bitcast<u32>(speed[0])); }
}
fn collide_cell(cell: u32, lid: u32) {
let n = params.n;
var fq: array<f32, 19>;
var drho = 0.0;
@@ -130,7 +144,7 @@ fn collide(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg:
let force = rho * a;
let u = (j + 0.5 * force) / rho;
let uu = dot(u, u);
atomicMax(&stats[0], bitcast<u32>(uu));
speed[lid] = uu;
let uf = dot(u, force);
let forced = params.acc.w > 0.5;
let wp = params.omega_plus;