OneCAD-Simulation/crates/lbm-gpu/src/lbm.wgsl

// The lattice Boltzmann step of fieldcraft_lbm::flow on the GPU: the same D3Q19 two-relaxation-
// time collision with Guo forcing, pull streaming with the per-link rules the CPU prepared
// (bounce-back with Ladd's moving-wall term, specular reflection, open placeholders), the Zou–He
// open boundaries, and the momentum exchange per force group. Populations are stored direction-
// major. The default AA pattern uses one population buffer (Bailey et al., ICPP 2009;
// Wittmann et al., Comput. Math. Appl. 65 (2013), 924–935). Even steps read a cell and
// write its opposite directions; odd steps read opposite directions at the upstream
// neighbours and write each outgoing direction at the downstream neighbour. Each thread
// writes exactly the addresses it read, so distinct cells never race. Physical coordinates
// wrap; the existing per-link rules still supply walls, reflection and open placeholders.
// A gather pass saves those boundary values before a scatter pass overrides their targets:
// overwriting in the gather itself could destroy another link's post-collision source.
// The optional A–B diagnostic path retains the original two-buffer collide/pull scheme.
//
// Single precision, with every population stored as its deviation from the weight, h_q = f_q − w_q
// (Krüger et al. 2017, §13.3.4: the deviations are small, so f32 keeps their digits instead of
// the weight's): ρ = 1 + Σ h, momentum Σ h c (Σ w c = 0), and every rule below is written for h.
// The CPU scheme (double precision, fixed sums) is the reference; see the crate's tests for the
// tolerances.
//
// With heat, a D3Q7 scalar lattice (the rest direction and the six axes of D3Q19, weights ¼ and
// ⅛) rides on the flow in the buffers t and t2, stored as deviations from w_q T₀ (T₀ the initial
// temperature) so that f32 keeps the digits of the differences: θ = T − T₀ = Σ t. It collides
// (BGK) with the velocity the momentum collision left in `moments`, streams with the links' rules
// read through the group's thermal condition (anti-bounce-back for a held temperature, bounce-
// back for an adiabatic face, a flux added to the bounced population), takes its inlet
// temperatures and outlet neighbours, and gives each group the heat its links exchanged (in
// contrib.w and forces.w). The temperature feeds back as the buoyant acceleration −β g θ.

override AA_MODE: bool = true;

const Q: u32 = 19u;
const QT: u32 = 7u;
const WG: u32 = 256u;
const KIND_PULL: u32 = 0u;
const KIND_BOUNCE: u32 = 1u;
const KIND_REFLECT: u32 = 2u;
const KIND_OPEN: u32 = 3u;

struct Params {
    n: u32,
    nx: u32,
    ny: u32,
    nz: u32,
    groups_x: u32,
    links: u32,
    open: u32,
    groups: u32,
    omega_plus: f32,
    omega_minus: f32,
    s_plus: f32,
    s_minus: f32,
    // The body acceleration (xyz) and whether there is one (w).
    acc: vec4<f32>,
    // The buoyancy −β g in lattice units (xyz) and whether there is any (w).
    buoy: vec4<f32>,
    // The scalar lattice: its relaxation rate (x), T₀ (y), whether there is heat (z).
    therm: vec4<f32>,
    // Current AA orientation (x: 0 normal, 1 opposite directions at upstream neighbours).
    phase: vec4<u32>,
};

@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read> fluid: array<u32>;
@group(0) @binding(2) var<storage, read_write> f: array<f32>;
@group(0) @binding(3) var<storage, read_write> g: array<f32>;
// Per link: (q | kind << 8 | q2 << 16, from, group, 0) and the moving-wall term.
@group(0) @binding(5) var<storage, read> link_u: array<vec4<u32>>;
@group(0) @binding(6) var<storage, read> link_f: array<f32>;
// Per open cell: (cell, axis, kind: 0 velocity / 1 density, 0) and (target xyz or density in x, inward sign in w).
@group(0) @binding(7) var<storage, read> open_u: array<vec4<u32>>;
@group(0) @binding(8) var<storage, read> open_f: array<vec4<f32>>;
// Per cell: velocity (with the half force) and density, from the moments kernel.
@group(0) @binding(9) var<storage, read_write> moments: array<vec4<f32>>;
// [0]: the largest squared lattice speed of the chunk (as f32 bits); [1]: a density was not positive.
@group(0) @binding(10) var<storage, read_write> stats: array<atomic<u32>>;
// Per link, the momentum it transferred to its group in the last step.
@group(0) @binding(11) var<storage, read_write> contrib: array<vec4<f32>>;
// Link indices ordered by group, and each group's range in that order (groups + 1 entries).
@group(0) @binding(12) var<storage, read> by_group: array<u32>;
@group(0) @binding(13) var<storage, read> group_start: array<u32>;
@group(0) @binding(14) var<storage, read_write> forces: array<vec4<f32>>;
// Per workgroup, the mass of its cells.
@group(0) @binding(15) var<storage, read_write> partials: array<f32>;
// The scalar populations remain A–B (t[q n + cell]; see the notes above).
@group(0) @binding(16) var<storage, read_write> t: array<f32>;
@group(0) @binding(17) var<storage, read_write> t2: array<f32>;
// Per force group: its thermal condition (x: 0 adiabatic, 1 temperature, 2 flux) and the value
// (y: the held temperature − T₀, or the flux as a temperature per step).
@group(0) @binding(18) var<storage, read> group_t: array<vec4<f32>>;
// Per open cell: whether it holds a temperature (x) and that temperature − T₀ (y).
@group(0) @binding(19) var<storage, read> open_t: array<vec4<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 {
        case 0u: { return vec3<i32>(0, 0, 0); }
        case 1u: { return vec3<i32>(1, 0, 0); }
        case 2u: { return vec3<i32>(-1, 0, 0); }
        case 3u: { return vec3<i32>(0, 1, 0); }
        case 4u: { return vec3<i32>(0, -1, 0); }
        case 5u: { return vec3<i32>(0, 0, 1); }
        case 6u: { return vec3<i32>(0, 0, -1); }
        case 7u: { return vec3<i32>(1, 1, 0); }
        case 8u: { return vec3<i32>(-1, -1, 0); }
        case 9u: { return vec3<i32>(1, -1, 0); }
        case 10u: { return vec3<i32>(-1, 1, 0); }
        case 11u: { return vec3<i32>(1, 0, 1); }
        case 12u: { return vec3<i32>(-1, 0, -1); }
        case 13u: { return vec3<i32>(1, 0, -1); }
        case 14u: { return vec3<i32>(-1, 0, 1); }
        case 15u: { return vec3<i32>(0, 1, 1); }
        case 16u: { return vec3<i32>(0, -1, -1); }
        case 17u: { return vec3<i32>(0, 1, -1); }
        default: { return vec3<i32>(0, -1, 1); }
    }
}

fn opp(q: u32) -> u32 {
    if q == 0u { return 0u; }
    if (q & 1u) == 1u { return q + 1u; }
    return q - 1u;
}

fn w_of(q: u32) -> f32 {
    if q == 0u { return 1.0 / 3.0; }
    if q < 7u { return 1.0 / 18.0; }
    return 1.0 / 36.0;
}

fn wt_of(q: u32) -> f32 {
    if q == 0u { return 0.25; }
    return 0.125;
}

fn has_heat() -> bool {
    return params.therm.z > 0.5;
}

// The temperature of a cell above T₀: the sum of its scalar populations.
fn theta_of(cell: u32) -> f32 {
    var s = 0.0;
    for (var q = 0u; q < QT; q = q + 1u) {
        s = s + t[q * params.n + cell];
    }
    return s;
}

// The body acceleration of a cell: the constant one plus the buoyant one.
fn acc_of(cell: u32) -> vec3<f32> {
    var a = params.acc.xyz;
    if params.buoy.w > 0.5 { a = a + params.buoy.xyz * theta_of(cell); }
    return a;
}

fn forced() -> bool {
    return params.acc.w > 0.5 || params.buoy.w > 0.5;
}

// The scalar equilibrium w_q T (1 + 4 c_q·u), shifted by w_q T₀ like the populations.
fn eq_t(q: u32, theta: f32, u: vec3<f32>) -> f32 {
    let cu = 4.0 * dot(vec3<f32>(c_of(q)), u);
    return wt_of(q) * (theta * (1.0 + cu) + params.therm.y * cu);
}

fn cell_of(lid: u32, wg: vec3<u32>) -> u32 {
    return (wg.y * params.groups_x + wg.x) * WG + lid;
}

fn is_fluid(cell: u32) -> bool {
    return fluid[cell] != 0u;
}

// A cell's grid coordinates (integer divisions once per cell, not once per direction).
fn coords(cell: u32) -> vec3<i32> {
    let plane = params.nx * params.ny;
    return vec3<i32>(i32(cell % params.nx), i32((cell % plane) / params.nx), i32(cell / plane));
}

// The cell at `p`, wrapped periodically. A lattice step moves at most one cell per axis, so a
// compare and an add per axis replace the remainder.
fn wrapped(p: vec3<i32>) -> u32 {
    let size = vec3<i32>(i32(params.nx), i32(params.ny), i32(params.nz));
    var w = select(p, p + size, p < vec3<i32>(0));
    w = select(w, w - size, w >= size);
    return u32((w.z * size.y + w.y) * size.x + w.x);
}

// Where population q of the cell at `xyz` is stored in the given AA orientation: in place when
// even, reversed at the upstream neighbour when odd.
fn population_index(q: u32, cell: u32, xyz: vec3<i32>, phase: u32) -> u32 {
    if phase == 0u { return q * params.n + cell; }
    return opp(q) * params.n + wrapped(xyz - c_of(q));
}

fn population(q: u32, cell: u32, xyz: vec3<i32>) -> f32 {
    if AA_MODE { return f[population_index(q, cell, xyz, params.phase.x)]; }
    return f[q * params.n + cell];
}

// Collision's output at its source cell, before the boundary rules.
fn post_index(q: u32, cell: u32, xyz: vec3<i32>) -> u32 {
    if !AA_MODE { return q * params.n + cell; }
    if params.phase.x == 0u { return opp(q) * params.n + cell; }
    return q * params.n + wrapped(xyz + c_of(q));
}

fn post(q: u32, cell: u32) -> f32 {
    return f[post_index(q, cell, aa_coords(cell))];
}

fn streamed(q: u32, cell: u32, xyz: vec3<i32>) -> f32 {
    if AA_MODE { return f[population_index(q, cell, xyz, params.phase.x ^ 1u)]; }
    return g[q * params.n + cell];
}

fn store_streamed(q: u32, cell: u32, xyz: vec3<i32>, v: f32) {
    if AA_MODE { f[population_index(q, cell, xyz, params.phase.x ^ 1u)] = v; }
    else { g[q * params.n + cell] = v; }
}

// The coordinates the AA addressing needs; none in the two-buffer layout.
fn aa_coords(cell: u32) -> vec3<i32> {
    if AA_MODE { return coords(cell); }
    return vec3<i32>(0);
}

// Two-relaxation-time collision with Guo forcing, in place in f (Ginzburg et al. 2008; Guo et
// al. 2002; Krüger et al. 2017 ch. 10), and the chunk's largest speed.
@compute @workgroup_size(256)
fn collide(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let cell = cell_of(lid, wg);
    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;
    // Only the odd AA orientation addresses neighbours.
    var xyz = vec3<i32>(0);
    if AA_MODE && params.phase.x == 1u { xyz = coords(cell); }
    var fq: array<f32, 19>;
    var drho = 0.0;
    var j = vec3<f32>(0.0);
    for (var q = 0u; q < Q; q = q + 1u) {
        let v = population(q, cell, xyz);
        fq[q] = v;
        drho = drho + v;
        j = j + vec3<f32>(c_of(q)) * v;
    }
    let rho = 1.0 + drho;
    if !(rho > 0.0) || rho != rho || abs(rho) > 1e30 {
        atomicStore(&stats[1], 1u);
        return;
    }
    let a = acc_of(cell);
    let force = rho * a;
    let u = (j + 0.5 * force) / rho;
    let uu = dot(u, u);
    speed[lid] = uu;
    if has_heat() { moments[cell] = vec4<f32>(u, rho); }
    let uf = dot(u, force);
    let forced = forced();
    let wp = params.omega_plus;
    let wm = params.omega_minus;
    let sp = params.s_plus;
    let sm = params.s_minus;
    // Rest population: symmetric only. Equilibria are shifted by the weight like the populations.
    let eq0 = w_of(0u) * (drho - 1.5 * rho * uu);
    var out0 = fq[0] - wp * (fq[0] - eq0);
    if forced { out0 = out0 + sp * w_of(0u) * (-3.0 * uf); }
    f[post_index(0u, cell, xyz)] = out0;
    for (var q = 1u; q < Q; q = q + 2u) {
        let o = q + 1u;
        let cv = vec3<f32>(c_of(q));
        let cu = dot(cv, u);
        let w = w_of(q);
        let eq_plus = w * (drho + rho * (4.5 * cu * cu - 1.5 * uu));
        let eq_minus = w * rho * 3.0 * cu;
        let plus = 0.5 * (fq[q] + fq[o]);
        let minus = 0.5 * (fq[q] - fq[o]);
        let relax_plus = wp * (plus - eq_plus);
        let relax_minus = wm * (minus - eq_minus);
        var s_plus = 0.0;
        var s_minus = 0.0;
        if forced {
            let cf = dot(cv, force);
            s_plus = sp * w * (9.0 * cu * cf - 3.0 * uf);
            s_minus = sm * w * 3.0 * cf;
        }
        f[post_index(q, cell, xyz)] = fq[q] - relax_plus - relax_minus + s_plus + s_minus;
        f[post_index(o, cell, xyz)] = fq[o] - relax_plus + relax_minus + s_plus - s_minus;
    }
}

// Pull streaming of f into g, first pass: every population from the neighbour at −c_q
// (wrapped; a wall or an open face overrides it in the second pass).
@compute @workgroup_size(256)
fn stream(@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; }
    let n = params.n;
    let nx = i32(params.nx);
    let ny = i32(params.ny);
    let nz = i32(params.nz);
    let i = i32(cell % params.nx);
    let jj = i32((cell / params.nx) % params.ny);
    let k = i32(cell / (params.nx * params.ny));
    for (var q = 0u; q < Q; q = q + 1u) {
        let c = c_of(q);
        let si = (i - c.x + nx) % nx;
        let sj = (jj - c.y + ny) % ny;
        let sk = (k - c.z + nz) % nz;
        let src = u32((sk * ny + sj) * nx + si);
        g[q * n + cell] = f[q * n + src];
    }
}

// Second pass, one thread per link: the populations with a rule (a periodic image, a wall, a
// slip face, an open face). Every (cell, direction) has at most one link, so no two threads
// write the same population.
@compute @workgroup_size(256)
fn stream_links(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    gather_link(cell_of(lid, wg), false);
}

// The same on the step whose wall forces are read back (a chunk's last): the AA layout
// overwrites the post-collision sources in the scatter, so the momentum exchange is taken here.
@compute @workgroup_size(256)
fn stream_links_exchange(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    gather_link(cell_of(lid, wg), true);
}

fn gather_link(l: u32, exchange: bool) {
    if l >= params.links { return; }
    let n = params.n;
    let lu = link_u[l];
    let q = lu.x & 0xffu;
    let kind = (lu.x >> 8u) & 0xffu;
    let q2 = (lu.x >> 16u) & 0xffu;
    let cell = lu.w;
    var v = 0.0;
    if kind == KIND_PULL {
        v = post(q, lu.y);
    } else if kind == KIND_BOUNCE {
        v = post(opp(q), cell) + link_f[l];
    } else if kind == KIND_REFLECT {
        // The mirrored direction has the same weight, so h reflects like f.
        v = post(q2, lu.y);
    } else {
        v = post(opp(q), cell);
    }
    if AA_MODE {
        g[l] = v;
        if exchange { contrib[l] = vec4<f32>(momentum_exchange(l, lu, q, kind, q2), 0.0); }
    } else {
        g[q * n + cell] = v;
    }
}

// The gather above must complete before any AA boundary population is overwritten.
@compute @workgroup_size(256)
fn scatter_links(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let l = cell_of(lid, wg);
    if l >= params.links { return; }
    let lu = link_u[l];
    store_streamed(lu.x & 0xffu, lu.w, aa_coords(lu.w), g[l]);
}

// Zou–He velocity and pressure boundaries on the streamed populations in g (Zou & He 1997;
// Hecht & Harting 2010 for the D3Q19 transverse correction).
@compute @workgroup_size(256)
fn open_boundaries(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let o = cell_of(lid, wg);
    if o >= params.open { return; }
    let n = params.n;
    let ou = open_u[o];
    let ot = open_f[o];
    let cell = ou.x;
    let axis = ou.y;
    let sign = ot.w;
    let a_half = 0.5 * acc_of(cell);
    let xyz = aa_coords(cell);
    // The sums of h over the tangential (c_n = 0) and outgoing (c_n < 0) directions. The sums of
    // f add the weights' (2/3 and 1/6 for D3Q19, so tangential + 2 outgoing gains exactly 1): the
    // weight is added after the deviations are summed, not before, which in single precision would
    // round away the deviations' last digits at low speed.
    var tangential = 0.0;
    var outgoing = 0.0;
    for (var q = 0u; q < Q; q = q + 1u) {
        let v = streamed(q, cell, xyz);
        let cn = f32(c_of(q)[axis]) * sign;
        if cn == 0.0 {
            tangential = tangential + v;
        } else if cn < 0.0 {
            outgoing = outgoing + v;
        }
    }
    var rho = 0.0;
    var bare = vec3<f32>(0.0);
    if ou.z == 0u {
        bare = ot.xyz - a_half;
        let un = sign * bare[axis];
        rho = (1.0 + (tangential + 2.0 * outgoing)) / (1.0 - un);
    } else {
        rho = ot.x;
        let un = ((rho - 1.0) - (tangential + 2.0 * outgoing)) / rho;
        // Not `bare[axis] = …`: DX12's shader compiler (FXC) refuses a write through a
        // dynamic index into a local vector (X3500).
        bare = select(-a_half, vec3<f32>(sign * un), vec3<u32>(0u, 1u, 2u) == vec3<u32>(axis));
    }
    for (var q = 1u; q < Q; q = q + 1u) {
        let c = c_of(q);
        if f32(c[axis]) * sign > 0.0 {
            let cu = dot(vec3<f32>(c), bare);
            store_streamed(q, cell, xyz, streamed(opp(q), cell, xyz) + 6.0 * w_of(q) * rho * cu);
        }
    }
    for (var t = 0u; t < 3u; t = t + 1u) {
        if t == axis { continue; }
        var m = -rho * bare[t];
        for (var q = 0u; q < Q; q = q + 1u) {
            m = m + streamed(q, cell, xyz) * f32(c_of(q)[t]);
        }
        for (var q = 1u; q < Q; q = q + 1u) {
            let c = c_of(q);
            if f32(c[axis]) * sign > 0.0 && c[t] != 0 {
                store_streamed(q, cell, xyz, streamed(q, cell, xyz) - 0.5 * m * f32(c[t]));
            }
        }
    }
}

fn momentum_exchange(l: u32, lu: vec4<u32>, q: u32, kind: u32, q2: u32) -> vec3<f32> {
    var v = vec3<f32>(0.0);
    if kind == KIND_BOUNCE {
        let out = post(opp(q), lu.w);
        let back = out + link_f[l];
        v = -vec3<f32>(c_of(q)) * (out + back);
    } else if kind == KIND_REFLECT {
        let out = post(q2, lu.y);
        v = vec3<f32>(c_of(q2)) * out - vec3<f32>(c_of(q)) * out;
    }
    return v;
}

// The momentum each link transferred to its group in the last step, from the post-collision
// populations in f (Ladd 1994; Mei et al. 2002): c_out f* − c_in f'.
@compute @workgroup_size(256)
fn link_forces(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let l = cell_of(lid, wg);
    if l >= params.links { return; }
    let n = params.n;
    let lu = link_u[l];
    let q = lu.x & 0xffu;
    let kind = (lu.x >> 8u) & 0xffu;
    let q2 = (lu.x >> 16u) & 0xffu;
    let cell = lu.w;
    // The deviations only: the weights' share of the exchange (2 w_q c_q per bounced link, which
    // cancels between a wall's links in exact arithmetic but is 10⁴ times the shear they carry)
    // is added on the host in double precision, so the sum keeps the deviations' digits.
    var v = contrib[l].xyz;
    if !AA_MODE { v = momentum_exchange(l, lu, q, kind, q2); }
    var heat = 0.0;
    // The heat the link gave its wall: what left the fluid minus what came back.
    if has_heat() && q < QT && (kind == KIND_BOUNCE || kind == KIND_REFLECT) {
        heat = t[opp(q) * n + cell] - thermal_rule(lu, q, kind, cell);
    }
    contrib[l] = vec4<f32>(v, heat);
}

// ---- The scalar lattice ----

// BGK collision of the scalar lattice, in place in t, with the velocity of the momentum
// collision.
@compute @workgroup_size(256)
fn collide_thermal(@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; }
    let n = params.n;
    var tq: array<f32, 7>;
    var theta = 0.0;
    for (var q = 0u; q < QT; q = q + 1u) {
        let v = t[q * n + cell];
        tq[q] = v;
        theta = theta + v;
    }
    let u = moments[cell].xyz;
    let omega = params.therm.x;
    for (var q = 0u; q < QT; q = q + 1u) {
        t[q * n + cell] = tq[q] - omega * (tq[q] - eq_t(q, theta, u));
    }
}

// Pull streaming of t into t2, first pass (as `stream`).
@compute @workgroup_size(256)
fn stream_thermal(@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; }
    let n = params.n;
    let nx = i32(params.nx);
    let ny = i32(params.ny);
    let nz = i32(params.nz);
    let i = i32(cell % params.nx);
    let jj = i32((cell / params.nx) % params.ny);
    let k = i32(cell / (params.nx * params.ny));
    for (var q = 0u; q < QT; q = q + 1u) {
        let c = c_of(q);
        let si = (i - c.x + nx) % nx;
        let sj = (jj - c.y + ny) % ny;
        let sk = (k - c.z + nz) % nz;
        let src = u32((sk * ny + sj) * nx + si);
        t2[q * n + cell] = t[q * n + src];
    }
}

// The scalar population a link's rule gives, from t: a pull as the flow's; at a wall or a slip
// face the group's thermal condition on the bounced population; an open placeholder.
fn thermal_rule(lu: vec4<u32>, q: u32, kind: u32, cell: u32) -> f32 {
    let n = params.n;
    if kind == KIND_PULL { return t[q * n + lu.y]; }
    let back = t[opp(q) * n + cell];
    if kind == KIND_OPEN { return back; }
    let cond = group_t[lu.z];
    let k = u32(cond.x);
    if k == 1u { return -back + 2.0 * wt_of(q) * cond.y; }
    if k == 2u { return back + cond.y; }
    return back;
}

// Second pass, one thread per link of the first seven directions.
@compute @workgroup_size(256)
fn stream_links_thermal(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let l = cell_of(lid, wg);
    if l >= params.links { return; }
    let lu = link_u[l];
    let q = lu.x & 0xffu;
    if q >= QT { return; }
    let kind = (lu.x >> 8u) & 0xffu;
    let cell = lu.w;
    t2[q * params.n + cell] = thermal_rule(lu, q, kind, cell);
}

// Inlets hold their temperature (the equilibrium at the cell's velocity); outlets take the
// streamed populations of the neighbour one cell inward.
@compute @workgroup_size(256)
fn open_thermal(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let o = cell_of(lid, wg);
    if o >= params.open { return; }
    let n = params.n;
    let ou = open_u[o];
    let cell = ou.x;
    let held = open_t[o];
    if held.x > 0.5 {
        let u = moments[cell].xyz;
        for (var q = 0u; q < QT; q = q + 1u) {
            t2[q * n + cell] = eq_t(q, held.y, u);
        }
    } else {
        var step = 1;
        if ou.y == 1u { step = i32(params.nx); }
        if ou.y == 2u { step = i32(params.nx * params.ny); }
        let src = u32(i32(cell) + i32(open_f[o].w) * step);
        for (var q = 0u; q < QT; q = q + 1u) {
            t2[q * n + cell] = t2[q * n + src];
        }
    }
}

// Flags a temperature that stopped being a number.
@compute @workgroup_size(256)
fn thermal_check(@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; }
    let theta = theta_of(cell);
    if theta != theta || abs(theta) > 1e30 { atomicStore(&stats[2], 1u); }
}

// One workgroup per group: the sum of its links' contributions.
@compute @workgroup_size(256)
fn group_forces(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let group = wg.y * params.groups_x + wg.x;
    if group >= params.groups { return; }
    let lo = group_start[group];
    let hi = group_start[group + 1u];
    var s = vec4<f32>(0.0);
    for (var i = lo + lid; i < hi; i = i + WG) {
        s = s + contrib[by_group[i]];
    }
    red[lid] = s;
    workgroupBarrier();
    for (var k = 128u; k > 0u; k = k >> 1u) {
        if lid < k { red[lid] = red[lid] + red[lid + k]; }
        workgroupBarrier();
    }
    if lid == 0u { forces[group] = red[0]; }
}

// The density and velocity (with the half force) of every cell from f; zero in solid cells.
@compute @workgroup_size(256)
fn cell_moments(@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) {
        moments[cell] = vec4<f32>(0.0);
        return;
    }
    let n = params.n;
    var drho = 0.0;
    var j = vec3<f32>(0.0);
    let xyz = aa_coords(cell);
    for (var q = 0u; q < Q; q = q + 1u) {
        let v = population(q, cell, xyz);
        drho = drho + v;
        j = j + vec3<f32>(c_of(q)) * v;
    }
    let rho = 1.0 + drho;
    let u = (j + 0.5 * rho * acc_of(cell)) / rho;
    moments[cell] = vec4<f32>(u, rho);
}

// Per workgroup, the sum of its cells' densities (the mass, since every cell has unit volume).
@compute @workgroup_size(256)
fn mass(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
    let group = wg.y * params.groups_x + wg.x;
    let cell = group * WG + lid;
    var v = 0.0;
    if cell < params.n && is_fluid(cell) {
        v = moments[cell].w;
    }
    red[lid] = vec4<f32>(v, 0.0, 0.0, 0.0);
    workgroupBarrier();
    for (var k = 128u; k > 0u; k = k >> 1u) {
        if lid < k { red[lid] = red[lid] + red[lid + k]; }
        workgroupBarrier();
    }
    if lid == 0u { partials[group] = red[0].x; }
}