From 1e0c21d08883d288d60cedf976af7324d7c0e273 Mon Sep 17 00:00:00 2001
From: Christian Svensson <blue@cmd.nu>
Date: Sat, 17 Oct 2026 00:58:33 +0200
Subject: [PATCH] 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.
---
 crates/lbm-gpu/src/lbm.wgsl | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/crates/lbm-gpu/src/lbm.wgsl b/crates/lbm-gpu/src/lbm.wgsl
index 3313452..2bda46c 100644
--- a/crates/lbm-gpu/src/lbm.wgsl
+++ b/crates/lbm-gpu/src/lbm.wgsl
@@ -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;