lbm-gpu: stream in place with the AA pattern

The collide and stream passes each moved the whole population field
through two buffers (304 bytes per cell update). The AA pattern (Bailey
et al. 2009; Wittmann et al. 2013) keeps one buffer: even steps collide
each cell in place with its directions swapped, odd steps read a cell's
incoming populations at its neighbours and write the outgoing ones where
the next step reads them, so each thread writes the addresses it read
and streaming happens inside the collision (152 bytes). A cell's
coordinates are computed once and neighbours wrap with a compare and an
add. The per-link wall rules gather their sources before a scatter
writes them, so no link overwrites another's post-collision value, and
the momentum exchange is taken in the gather of the step whose forces
are read back.

Options::streaming picks the layout; Auto, what a solve uses, streams
in place, which measured faster on every device tried: 1.5x on an RTX
4090, 1.85x on the GB10 and 1.45x on an M4 Max at periodic 128^3.
Options::profile times each kernel with timestamp queries for the
throughput test, which now runs both layouts. The verification tests
run in both layouts, with two new ones comparing fields, mass, forces
and heat with the CPU after odd and even step counts across submission
boundaries.

Part of #562.

crates/lbm-gpu/src/lbm.wgsl

@@ -2,8 +2,15 @@
 // 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 (f[q n + cell]) in two buffers, A–B: a step collides `f` in place, streams `f` into `g`
-// and applies the open boundaries to `g`; the next step binds the buffers the other way round.
+// 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
@@ -20,6 +27,8 @@
 // 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;
@@ -47,6 +56,8 @@ struct Params {
     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;
@@ -71,7 +82,7 @@ struct Params {
 @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, A–B like f and g (t[q n + cell]; see the notes above).
+// 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
@@ -164,6 +175,60 @@ 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)
@@ -183,11 +248,14 @@ fn collide(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg:
 
 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 = f[q * n + cell];
+        let v = population(q, cell, xyz);
         fq[q] = v;
         drho = drho + v;
         j = j + vec3<f32>(c_of(q)) * v;
@@ -213,7 +281,7 @@ fn collide_cell(cell: u32, lid: u32) {
     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[cell] = out0;
+    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));
@@ -232,8 +300,8 @@ fn collide_cell(cell: u32, lid: u32) {
             s_plus = sp * w * (9.0 * cu * cf - 3.0 * uf);
             s_minus = sm * w * 3.0 * cf;
         }
-        f[q * n + cell] = fq[q] - relax_plus - relax_minus + s_plus + s_minus;
-        f[o * n + cell] = fq[o] - relax_plus + relax_minus + s_plus - s_minus;
+        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;
     }
 }
 
@@ -266,7 +334,17 @@ fn stream(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg:
 // write the same population.
 @compute @workgroup_size(256)
 fn stream_links(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wg: vec3<u32>) {
-    let l = cell_of(lid, wg);
+    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];
@@ -276,16 +354,30 @@ fn stream_links(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id
     let cell = lu.w;
     var v = 0.0;
     if kind == KIND_PULL {
-        v = f[q * n + lu.y];
+        v = post(q, lu.y);
     } else if kind == KIND_BOUNCE {
-        v = f[opp(q) * n + cell] + link_f[l];
+        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 = f[q2 * n + lu.y];
+        v = post(q2, lu.y);
     } else {
-        v = f[opp(q) * n + cell];
+        v = post(opp(q), cell);
     }
-    g[q * n + cell] = v;
+    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;
@@ -301,6 +393,7 @@ fn open_boundaries(@builtin(local_invocation_index) lid: u32, @builtin(workgroup
     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
@@ -308,7 +401,7 @@ fn open_boundaries(@builtin(local_invocation_index) lid: u32, @builtin(workgroup
     var tangential = 0.0;
     var outgoing = 0.0;
     for (var q = 0u; q < Q; q = q + 1u) {
-        let v = g[q * n + cell];
+        let v = streamed(q, cell, xyz);
         let cn = f32(c_of(q)[axis]) * sign;
         if cn == 0.0 {
             tangential = tangential + v;
@@ -333,24 +426,37 @@ fn open_boundaries(@builtin(local_invocation_index) lid: u32, @builtin(workgroup
         let c = c_of(q);
         if f32(c[axis]) * sign > 0.0 {
             let cu = dot(vec3<f32>(c), bare);
-            g[q * n + cell] = g[opp(q) * n + cell] + 6.0 * w_of(q) * rho * cu;
+            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 + g[q * n + cell] * f32(c_of(q)[t]);
+            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 {
-                g[q * n + cell] = g[q * n + cell] - 0.5 * m * f32(c[t]);
+                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)
@@ -366,16 +472,9 @@ fn link_forces(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id)
     // 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 = vec3<f32>(0.0);
+    var v = contrib[l].xyz;
+    if !AA_MODE { v = momentum_exchange(l, lu, q, kind, q2); }
     var heat = 0.0;
-    if kind == KIND_BOUNCE {
-        let out = f[opp(q) * n + cell];
-        let back = out + link_f[l];
-        v = -vec3<f32>(c_of(q)) * (out + back);
-    } else if kind == KIND_REFLECT {
-        let out = f[q2 * n + lu.y];
-        v = vec3<f32>(c_of(q2)) * out - vec3<f32>(c_of(q)) * out;
-    }
     // 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);
@@ -525,8 +624,9 @@ fn cell_moments(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id
     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 = f[q * n + cell];
+        let v = population(q, cell, xyz);
         drho = drho + v;
         j = j + vec3<f32>(c_of(q)) * v;
     }

crates/lbm-gpu/src/lib.rs

@@ -18,7 +18,7 @@
 #[cfg(not(target_arch = "wasm32"))]
 mod native;
 #[cfg(not(target_arch = "wasm32"))]
-pub use native::{GpuFlow, register};
+pub use native::{GpuFlow, KernelTime, Options, Streaming, register};
 
 /// In the browser there is no device backend yet (WebGPU later): registering is a no-op and
 /// `solve.fluid` with `backend: "gpu"` reports that none is available.

crates/lbm-gpu/src/native.rs

@@ -31,6 +31,9 @@ const MAX_CHUNK: u64 = 64;
 /// How long one wait for the device may take before the solve gives up on it.
 const DEVICE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
 
+/// A bounded chunk's dispatches, shared by the normal and timestamped paths.
+type Dispatch = (usize, &'static str, (u32, u32));
+
 /// A vector of `len` copies of `value`, or [`Error::Limit`] when memory runs out (the host
 /// arrays are sized by the grid: 16 M cells is over a gigabyte of populations).
 fn filled<T: Clone>(len: usize, value: T, what: &str) -> Result<Vec<T>> {
@@ -72,13 +75,48 @@ pub fn register() {
     fieldcraft_lbm::register_backend(Box::new(|flow, _setup| GpuFlow::from_env(flow).map(|g| Box::new(g) as Box<dyn Backend>)));
 }
 
+/// How the momentum populations stream (see `lbm.wgsl`).
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum Streaming {
+    /// In place: measured faster on every device tried (RTX 4090, GB10, M4 Max; #563).
+    #[default]
+    Auto,
+    /// One population buffer, collided and streamed in place (the AA pattern).
+    InPlace,
+    /// Two population buffers: collide in place, then pull into the other (A–B).
+    TwoBuffers,
+}
+
+/// How a [`GpuFlow`] runs; the defaults are what a solve uses.
+#[derive(Clone, Copy, Debug, Default)]
+pub struct Options {
+    pub streaming: Streaming,
+    /// Time every kernel with timestamp queries ([`GpuFlow::kernel_times`]). Each dispatch then
+    /// gets its own pass and every chunk waits for its timestamps: for measurements, not solves.
+    pub profile: bool,
+}
+
+/// GPU time per kernel, summed over the dispatches since the flow was created (with
+/// [`Options::profile`]).
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct KernelTime {
+    pub dispatches: u64,
+    pub seconds: f64,
+}
+
 /// A prepared flow on the device.
 pub struct GpuFlow {
     device: wgpu::Device,
     queue: wgpu::Queue,
+    /// With [`Options::profile`]: the GPU time of each kernel so far.
+    kernel_times: Option<std::cell::RefCell<std::collections::BTreeMap<&'static str, KernelTime>>>,
+    /// How long creating the compute pipelines (compiling the shader) took.
+    pipeline_seconds: f64,
+    aa: bool,
     adapter: String,
     pipelines: HashMap<&'static str, wgpu::ComputePipeline>,
-    /// Bind groups with (f, g) = (buffer 0, buffer 1) and the other way round.
+    /// AB swaps populations; AA keeps one population buffer and changes its orientation.
+    /// Both modes swap the separate scalar lattice buffers.
     bind: [wgpu::BindGroup; 2],
     /// Which bind group's `f` holds the state.
     current: usize,
@@ -95,7 +133,7 @@ pub struct GpuFlow {
     forces: wgpu::Buffer,
     partials: wgpu::Buffer,
     partial_count: usize,
-    /// The scalar populations, A–B like the flow's, with heat; and T₀ they are shifted by.
+    /// The scalar populations remain A–B with heat; and T₀ they are shifted by.
     thermal: Option<[wgpu::Buffer; 2]>,
     initial_temperature: f64,
     fluid: Vec<bool>,
@@ -121,10 +159,20 @@ impl GpuFlow {
     /// `flow` (its grid, rules and initial state: everything the step needs) on the device
     /// `policy` picks.
     pub fn new(flow: &Flow, policy: Policy) -> Result<GpuFlow> {
+        GpuFlow::with_options(flow, policy, Options::default())
+    }
+
+    /// [`GpuFlow::new`] with `options`.
+    pub fn with_options(flow: &Flow, policy: Policy, options: Options) -> Result<GpuFlow> {
         let Some((adapter, tier)) = select::adapter(policy).map_err(gpu_error)? else {
             return Err(Error::Invalid(format!("{}=cpu: no GPU device; solve with `backend: \"cpu\"`", select::ENV)));
         };
         let info = adapter.get_info();
+        let profile = options.profile;
+        let features = if profile { wgpu::Features::TIMESTAMP_QUERY } else { wgpu::Features::empty() };
+        if !adapter.features().contains(features) {
+            return Err(gpu_error("profiling needs timestamp queries, which this adapter does not offer"));
+        }
         let have = adapter.limits();
         let mut limits = wgpu::Limits::default().using_resolution(have.clone());
         limits.max_buffer_size = have.max_buffer_size;
@@ -132,13 +180,17 @@ impl GpuFlow {
         limits.max_storage_buffers_per_shader_stage = have.max_storage_buffers_per_shader_stage.max(limits.max_storage_buffers_per_shader_stage);
         let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
             label: Some("fieldcraft lbm"),
-            required_features: wgpu::Features::empty(),
+            required_features: features,
             required_limits: limits,
             ..Default::default()
         }))
         .map_err(gpu_error)?;
         let adapter_name = format!("{}{}", info.name, if tier == select::Tier::Software { " (software)" } else { "" });
 
+        let aa = match options.streaming {
+            Streaming::Auto | Streaming::InPlace => true,
+            Streaming::TwoBuffers => false,
+        };
         let n = flow.grid().len();
         let elem = 4u64;
         let pop_bytes = elem * (Q * n) as u64;
@@ -328,13 +380,18 @@ impl GpuFlow {
                 mapped_at_creation: false,
             })
         };
-        let params = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
-            label: Some("lbm params"),
-            contents: &params,
-            usage: wgpu::BufferUsages::UNIFORM,
+        let params = [0u32, 1].map(|phase| {
+            let mut bytes = params.clone();
+            bytes.extend_from_slice(&bytes_u32(&[phase, 0, 0, 0]));
+            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+                label: Some("lbm params"),
+                contents: &bytes,
+                usage: wgpu::BufferUsages::UNIFORM,
+            })
         });
         let fluid_buf = storage("lbm fluid", &bytes_u32(&fluid_u32));
-        let pops = [storage("lbm f", &bytes_f32(&f)), empty("lbm g", pop_bytes)];
+        // AA reuses binding g for boundary scratch only, never another full population field.
+        let pops = [storage("lbm f", &bytes_f32(&f)), empty("lbm g / AA boundary scratch", if aa { 4 * links.len() as u64 } else { pop_bytes })];
         let link_u = storage("lbm links", &bytes_u32(&link_u));
         let link_f = storage("lbm link du", &bytes_f32(&link_f));
         let open_u = storage("lbm open cells", &bytes_u32(&open_u));
@@ -394,10 +451,10 @@ impl GpuFlow {
                 label: Some("fieldcraft lbm"),
                 layout: &layout,
                 entries: &[
-                    wgpu::BindGroupEntry { binding: 0, resource: params.as_entire_binding() },
+                    wgpu::BindGroupEntry { binding: 0, resource: params[k].as_entire_binding() },
                     wgpu::BindGroupEntry { binding: 1, resource: fluid_buf.as_entire_binding() },
-                    wgpu::BindGroupEntry { binding: 2, resource: pops[k].as_entire_binding() },
-                    wgpu::BindGroupEntry { binding: 3, resource: pops[1 - k].as_entire_binding() },
+                    wgpu::BindGroupEntry { binding: 2, resource: pops[if aa { 0 } else { k }].as_entire_binding() },
+                    wgpu::BindGroupEntry { binding: 3, resource: pops[if aa { 1 } else { 1 - k }].as_entire_binding() },
                     wgpu::BindGroupEntry { binding: 5, resource: link_u.as_entire_binding() },
                     wgpu::BindGroupEntry { binding: 6, resource: link_f.as_entire_binding() },
                     wgpu::BindGroupEntry { binding: 7, resource: open_u.as_entire_binding() },
@@ -422,10 +479,13 @@ impl GpuFlow {
             immediate_size: 0,
         });
         let mut pipelines = HashMap::new();
+        let pipeline_start = std::time::Instant::now();
         for name in [
             "collide",
             "stream",
             "stream_links",
+            "stream_links_exchange",
+            "scatter_links",
             "open_boundaries",
             "link_forces",
             "group_forces",
@@ -442,7 +502,7 @@ impl GpuFlow {
                 layout: Some(&pl),
                 module: &module,
                 entry_point: Some(name),
-                compilation_options: Default::default(),
+                compilation_options: wgpu::PipelineCompilationOptions { constants: &[("AA_MODE", if aa { 1.0 } else { 0.0 })], ..Default::default() },
                 cache: None,
             });
             pipelines.insert(name, p);
@@ -452,9 +512,13 @@ impl GpuFlow {
                 return Err(gpu_error(e));
             }
         }
+        let pipeline_seconds = pipeline_start.elapsed().as_secs_f64();
         Ok(GpuFlow {
             device,
             queue,
+            kernel_times: profile.then(Default::default),
+            pipeline_seconds,
+            aa,
             adapter: adapter_name,
             pipelines,
             bind,
@@ -489,6 +553,21 @@ impl GpuFlow {
         &self.adapter
     }
 
+    /// Whether the populations stream in place (see [`Streaming`]).
+    pub fn streams_in_place(&self) -> bool {
+        self.aa
+    }
+
+    /// How long creating the compute pipelines took, in seconds.
+    pub fn pipeline_seconds(&self) -> f64 {
+        self.pipeline_seconds
+    }
+
+    /// With [`Options::profile`], the GPU time of each kernel so far, by kernel name.
+    pub fn kernel_times(&self) -> Option<std::collections::BTreeMap<&'static str, KernelTime>> {
+        self.kernel_times.as_ref().map(|t| t.borrow().clone())
+    }
+
     fn pipeline(&self, name: &str) -> Result<&wgpu::ComputePipeline> {
         self.pipelines.get(name).ok_or_else(|| Error::Invalid(format!("GPU: no kernel {name}")))
     }
@@ -500,6 +579,57 @@ impl GpuFlow {
         Ok(())
     }
 
+    /// Timestamped passes are diagnostic: their extra pass boundaries and readback are not
+    /// used for throughput measurements. Normally every chunk remains one compute pass.
+    fn submit(&self, label: &str, commands: &[Dispatch]) -> Result<()> {
+        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some(label) });
+        let Some(times) = &self.kernel_times else {
+            {
+                let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some(label), timestamp_writes: None });
+                for &(bind, name, groups) in commands {
+                    self.dispatch(&mut pass, bind, name, groups)?;
+                }
+            }
+            self.queue.submit([enc.finish()]);
+            return Ok(());
+        };
+        let count = u32::try_from(commands.len() * 2).map_err(|_| gpu_error("too many timestamp queries"))?;
+        if count == 0 || count > wgpu::QUERY_SET_MAX_QUERIES {
+            return Err(gpu_error("timestamp query count is outside the device limit"));
+        }
+        let query = self.device.create_query_set(&wgpu::QuerySetDescriptor { label: Some(label), ty: wgpu::QueryType::Timestamp, count });
+        let result = self.device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some("lbm timestamps"),
+            size: u64::from(count) * 8,
+            usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
+            mapped_at_creation: false,
+        });
+        for (i, &(bind, name, groups)) in commands.iter().enumerate() {
+            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
+                label: Some(name),
+                timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
+                    query_set: &query,
+                    beginning_of_pass_write_index: Some(2 * i as u32),
+                    end_of_pass_write_index: Some(2 * i as u32 + 1),
+                }),
+            });
+            self.dispatch(&mut pass, bind, name, groups)?;
+        }
+        enc.resolve_query_set(&query, 0..count, &result, 0);
+        self.queue.submit([enc.finish()]);
+        let bytes = self.read(&result, u64::from(count) * 8)?;
+        let period = f64::from(self.queue.get_timestamp_period());
+        let mut totals = times.borrow_mut();
+        for ((_, name, _), pair) in commands.iter().zip(bytes.as_chunks::<16>().0) {
+            let [begin, end] = pair.as_chunks::<8>().0 else { return Err(gpu_error("invalid timestamp pair")) };
+            let ticks = u64::from_le_bytes(*end).checked_sub(u64::from_le_bytes(*begin)).ok_or_else(|| gpu_error("invalid GPU timestamps"))?;
+            let total = totals.entry(name).or_default();
+            total.dispatches += 1;
+            total.seconds += ticks as f64 * period * 1e-9;
+        }
+        Ok(())
+    }
+
     /// Copy `size` bytes of `src` back to the host.
     fn read(&self, src: &wgpu::Buffer, size: u64) -> Result<Vec<u8>> {
         let size = size.max(16);
@@ -551,13 +681,7 @@ impl GpuFlow {
 
     /// The moments of the current state into the moments buffer.
     fn compute_moments(&self) -> Result<()> {
-        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("lbm moments") });
-        {
-            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("lbm moments"), timestamp_writes: None });
-            self.dispatch(&mut pass, self.current, "cell_moments", self.cells)?;
-        }
-        self.queue.submit([enc.finish()]);
-        Ok(())
+        self.submit("lbm moments", &[(self.current, "cell_moments", self.cells)])
     }
 }
 
@@ -565,41 +689,45 @@ impl GpuFlow {
     /// One submission of `steps` (≤ MAX_CHUNK) time steps, then the stability check.
     fn advance_chunk(&mut self, steps: u64) -> Result<()> {
         self.queue.write_buffer(&self.stats, 0, &bytes_u32(&[0, 0, 0, 0]));
-        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("lbm steps") });
-        {
-            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("lbm steps"), timestamp_writes: None });
-            for _ in 0..steps {
-                self.dispatch(&mut pass, self.current, "collide", self.cells)?;
-                self.dispatch(&mut pass, self.current, "stream", self.cells)?;
-                if self.link_count > 0 {
-                    self.dispatch(&mut pass, self.current, "stream_links", self.links_groups)?;
-                }
-                if self.open_count > 0 {
-                    self.dispatch(&mut pass, self.current, "open_boundaries", self.open_groups)?;
-                }
-                if self.thermal.is_some() {
-                    self.dispatch(&mut pass, self.current, "collide_thermal", self.cells)?;
-                    self.dispatch(&mut pass, self.current, "stream_thermal", self.cells)?;
-                    if self.link_count > 0 {
-                        self.dispatch(&mut pass, self.current, "stream_links_thermal", self.links_groups)?;
-                    }
-                    if self.open_count > 0 {
-                        self.dispatch(&mut pass, self.current, "open_thermal", self.open_groups)?;
-                    }
-                }
-                self.current ^= 1;
+        let mut commands = Vec::with_capacity(8 * MAX_CHUNK as usize + 3);
+        for step in 0..steps {
+            commands.push((self.current, "collide", self.cells));
+            if !self.aa {
+                commands.push((self.current, "stream", self.cells));
             }
-            // The momentum (and heat) exchange of the last step, from its post-collision
-            // populations.
-            if self.link_count > 0 && self.group_count > 0 {
-                self.dispatch(&mut pass, self.current ^ 1, "link_forces", self.links_groups)?;
-                self.dispatch(&mut pass, self.current ^ 1, "group_forces", self.force_groups)?;
+            if self.link_count > 0 {
+                // The last step's momentum exchange is the one read back.
+                let exchange = self.aa && step + 1 == steps;
+                commands.push((self.current, if exchange { "stream_links_exchange" } else { "stream_links" }, self.links_groups));
+                if self.aa {
+                    commands.push((self.current, "scatter_links", self.links_groups));
+                }
+            }
+            if self.open_count > 0 {
+                commands.push((self.current, "open_boundaries", self.open_groups));
             }
             if self.thermal.is_some() {
-                self.dispatch(&mut pass, self.current, "thermal_check", self.cells)?;
+                commands.push((self.current, "collide_thermal", self.cells));
+                commands.push((self.current, "stream_thermal", self.cells));
+                if self.link_count > 0 {
+                    commands.push((self.current, "stream_links_thermal", self.links_groups));
+                }
+                if self.open_count > 0 {
+                    commands.push((self.current, "open_thermal", self.open_groups));
+                }
             }
+            self.current ^= 1;
         }
-        self.queue.submit([enc.finish()]);
+        // The momentum (and heat) exchange of the last step, from its post-collision
+        // populations.
+        if self.link_count > 0 && self.group_count > 0 {
+            commands.push((self.current ^ 1, "link_forces", self.links_groups));
+            commands.push((self.current ^ 1, "group_forces", self.force_groups));
+        }
+        if self.thermal.is_some() {
+            commands.push((self.current, "thermal_check", self.cells));
+        }
+        self.submit("lbm steps", &commands)?;
         self.steps_done += steps;
         self.exchange_cache = None;
         let stats = self.read(&self.stats, 16)?;
@@ -624,7 +752,7 @@ impl GpuFlow {
 
 impl Backend for GpuFlow {
     fn name(&self) -> String {
-        format!("gpu: {}", self.adapter)
+        format!("gpu: {}, {} streaming", self.adapter, if self.aa { "in-place" } else { "two-buffer" })
     }
 
     fn advance(&mut self, steps: u64) -> Result<()> {
@@ -671,12 +799,7 @@ impl Backend for GpuFlow {
 
     fn mass(&mut self) -> Result<f64> {
         self.compute_moments()?;
-        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("lbm mass") });
-        {
-            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("lbm mass"), timestamp_writes: None });
-            self.dispatch(&mut pass, self.current, "mass", self.cells)?;
-        }
-        self.queue.submit([enc.finish()]);
+        self.submit("lbm mass", &[(self.current, "mass", self.cells)])?;
         let partials = self.read_f32(&self.partials, self.partial_count)?;
         // Summed in a fixed order on the host, so the result is deterministic on this device.
         Ok(partials.iter().map(|v| f64::from(*v)).sum())