lbm: a Backend trait steps the flow, with the CPU scheme as its reference

run() drives a Backend in chunks of steps (to the next frame, history sample
or progress check) instead of the CPU Flow directly, so a device backend
registered by an application can step the same prepared flow: Flow keeps
building the grid, links and open cells and exposes them read-only, and
implements the trait itself. solve_fluid_on picks the CPU or the registered
device backend; solve_fluid stays the CPU and is unchanged in its results.

Part of #529.

crates/lbm/src/backend.rs

@@ -0,0 +1,91 @@
+//! What steps a prepared flow: the CPU scheme in [`crate::flow::Flow`], or a device backend an
+//! application registers (a GPU crate above this layer implements the trait and hands a factory
+//! to [`register_backend`]; `solve.fluid` picks it when asked). Every backend takes its grid,
+//! links and open cells from the [`Flow`] that prepared them, so voxelisation and every boundary
+//! rule exist once, here.
+//!
+//! The CPU result is the reference: a device backend is deterministic on its device and agrees
+//! with the CPU within a stated tolerance (its tests say which).
+
+use std::sync::OnceLock;
+
+use crate::flow::{Flow, Group};
+use crate::setup::Setup;
+use crate::{Error, Result};
+
+/// The density and velocity of every cell (lattice units; zero in solid cells).
+#[derive(Debug, Clone, PartialEq)]
+pub struct Fields {
+    pub density: Vec<f64>,
+    pub velocity: Vec<[f64; 3]>,
+}
+
+/// A stepper of one prepared flow.
+pub trait Backend {
+    /// What steps the flow, for reports: `cpu`, `gpu: <adapter>`.
+    fn name(&self) -> String;
+    /// Advance `steps` time steps. Fails with [`Error::Unstable`] when a density stops being a
+    /// positive number or a lattice speed exceeds [`crate::setup::MAX_LATTICE_SPEED`] (a device
+    /// backend may notice at the end of the chunk rather than the step).
+    fn advance(&mut self, steps: u64) -> Result<()>;
+    fn steps_done(&self) -> u64;
+    /// The density and velocity (with the half force) of every cell.
+    fn fields(&mut self) -> Result<Fields>;
+    /// The total mass of the fluid cells (the sum of every population).
+    fn mass(&mut self) -> Result<f64>;
+    /// The momentum each force group received in the last step (lattice units: force per step).
+    fn forces(&mut self) -> Result<Vec<[f64; 3]>>;
+    fn groups(&self) -> &[Group];
+    fn fluid_cells(&self) -> usize;
+    fn is_fluid(&self, cell: usize) -> bool;
+    /// Obstacles that cover no cell centre.
+    fn empty_obstacles(&self) -> &[String];
+}
+
+/// Makes a device backend from a prepared flow (its initial state included) and its setup.
+pub type BackendFactory = dyn Fn(&Flow, &Setup) -> Result<Box<dyn Backend>> + Send + Sync;
+
+static FACTORY: OnceLock<Box<BackendFactory>> = OnceLock::new();
+
+/// Register the process's device backend (an application does this once at start; a second
+/// registration is ignored).
+pub fn register_backend(factory: Box<BackendFactory>) {
+    let _ = FACTORY.set(factory);
+}
+
+/// The registered device backend's factory, if an application registered one.
+pub fn backend_factory() -> Option<&'static BackendFactory> {
+    FACTORY.get().map(|b| b.as_ref())
+}
+
+/// Which backend a solve asks for.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum Choice {
+    /// The CPU scheme: the reference, bit for bit reproducible.
+    #[default]
+    Cpu,
+    /// The registered device backend; an error when none is registered.
+    Device,
+}
+
+impl Choice {
+    /// `cpu` or `gpu`, exactly.
+    pub fn parse(s: &str) -> Option<Choice> {
+        match s {
+            "cpu" => Some(Choice::Cpu),
+            "gpu" => Some(Choice::Device),
+            _ => None,
+        }
+    }
+}
+
+/// The backend `choice` names for `flow`.
+pub fn backend(choice: Choice, flow: Flow, setup: &Setup) -> Result<Box<dyn Backend>> {
+    match choice {
+        Choice::Cpu => Ok(Box::new(flow)),
+        Choice::Device => match backend_factory() {
+            Some(make) => make(&flow, setup),
+            None => Err(Error::Invalid("no GPU backend is available in this build: solve on the CPU (`backend: \"cpu\"`)".into())),
+        },
+    }
+}

crates/lbm/src/flow.rs

@@ -45,9 +45,10 @@ use crate::setup::{MAX_LATTICE_SPEED, Obstacle, Setup};
 use crate::voxel::{Grid, TriangleTree, voxelise};
 use crate::{Error, Result};
 
-/// Where a population of a cell next to a wall or box face comes from.
+/// Where a population of a cell next to a wall or box face comes from (every other population
+/// streams from the neighbour at `−c_q`). Public so a device backend can apply the same rules.
 #[derive(Debug, Clone, Copy, PartialEq)]
-enum Rule {
+pub enum Rule {
     /// The same direction from another cell (a neighbour or a periodic image).
     Pull { from: u32 },
     /// Bounce-back at a wall of force group `group`; `du` is the moving-wall term.
@@ -59,25 +60,26 @@ enum Rule {
 }
 
 /// A population, direction `q` of `cell`, and where it comes from.
-#[derive(Debug, Clone, Copy)]
-struct Link {
-    cell: u32,
-    q: u8,
-    rule: Rule,
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct Link {
+    pub cell: u32,
+    pub q: u8,
+    pub rule: Rule,
 }
 
 /// An inlet or outlet cell.
-#[derive(Debug, Clone, Copy)]
-struct OpenCell {
-    cell: u32,
-    axis: usize,
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct OpenCell {
+    pub cell: u32,
+    pub axis: usize,
     /// +1 at a min face (the inward normal points up the axis), −1 at a max face.
-    inward: f64,
-    target: Target,
+    pub inward: f64,
+    pub target: Target,
 }
 
-#[derive(Debug, Clone, Copy)]
-enum Target {
+/// What an open cell holds.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum Target {
     /// The velocity (lattice units).
     Velocity([f64; 3]),
     /// The density.
@@ -377,6 +379,36 @@ impl Flow {
         self.steps_done
     }
 
+    /// The links that do not stream from the neighbour at `−c_q`, in cell order.
+    pub fn links(&self) -> &[Link] {
+        &self.links
+    }
+
+    /// The inlet and outlet cells.
+    pub fn open_cells(&self) -> &[OpenCell] {
+        &self.open
+    }
+
+    /// Per direction, the index offset of the neighbour a population streams from.
+    pub fn offsets(&self) -> [isize; Q] {
+        self.offset
+    }
+
+    /// `1/τ⁺` and `1/τ⁻`.
+    pub fn omegas(&self) -> (f64, f64) {
+        (self.omega_plus, self.omega_minus)
+    }
+
+    /// The body acceleration in lattice units.
+    pub fn acceleration(&self) -> [f64; 3] {
+        self.acceleration
+    }
+
+    /// Population `q` of `cell` (zero in a solid cell).
+    pub fn population(&self, q: usize, cell: usize) -> f64 {
+        self.f.get(q * self.stride + cell).copied().unwrap_or(0.0)
+    }
+
     /// Set every fluid cell to the equilibrium of `state(cell) = (density, velocity)` in lattice
     /// units (velocity as defined with the half force).
     pub fn set_equilibrium(&mut self, state: impl Fn(usize) -> (f64, [f64; 3])) {
@@ -642,3 +674,57 @@ impl Flow {
         }
     }
 }
+
+impl crate::backend::Backend for Flow {
+    fn name(&self) -> String {
+        "cpu".into()
+    }
+
+    fn advance(&mut self, steps: u64) -> Result<()> {
+        for _ in 0..steps {
+            self.step()?;
+        }
+        Ok(())
+    }
+
+    fn steps_done(&self) -> u64 {
+        self.steps_done
+    }
+
+    fn fields(&mut self) -> Result<crate::backend::Fields> {
+        let n = self.n;
+        let mut density = filled(n, 0.0, "density field")?;
+        let mut velocity = filled(n, [0.0; 3], "velocity field")?;
+        for cell in 0..n {
+            if let (Some((rho, u)), Some(d), Some(v)) = (self.moments(cell), density.get_mut(cell), velocity.get_mut(cell)) {
+                *d = rho;
+                *v = u;
+            }
+        }
+        Ok(crate::backend::Fields { density, velocity })
+    }
+
+    fn mass(&mut self) -> Result<f64> {
+        Ok(Flow::mass(self))
+    }
+
+    fn forces(&mut self) -> Result<Vec<[f64; 3]>> {
+        Ok(self.forces.clone())
+    }
+
+    fn groups(&self) -> &[Group] {
+        &self.groups
+    }
+
+    fn fluid_cells(&self) -> usize {
+        self.fluid_cells
+    }
+
+    fn is_fluid(&self, cell: usize) -> bool {
+        Flow::is_fluid(self, cell)
+    }
+
+    fn empty_obstacles(&self) -> &[String] {
+        &self.empty_obstacles
+    }
+}

crates/lbm/src/lib.rs

@@ -48,15 +48,17 @@
 // Population arrays are indexed by cell and direction counts fixed when they are allocated.
 #![allow(clippy::needless_range_loop)]
 
+pub mod backend;
 pub mod flow;
 pub mod lattice;
 pub mod setup;
 pub mod solve;
 pub mod voxel;
 
-pub use flow::{Flow, Group};
+pub use backend::{Backend, Choice, Fields, register_backend};
+pub use flow::{Flow, Group, Link, OpenCell, Rule, Target};
 pub use setup::{Obstacle, Setup, fluid_step};
-pub use solve::{FluidResult, FluidStats, ForceHistory, Frame, run, solve_fluid};
+pub use solve::{FluidResult, FluidStats, ForceHistory, Frame, run, solve_fluid, solve_fluid_on};
 pub use voxel::Grid;
 
 #[derive(Debug, Clone, PartialEq, thiserror::Error)]

crates/lbm/src/solve.rs

@@ -2,6 +2,7 @@
 
 use fieldcraft_model::Model;
 
+use crate::backend::{Backend, Choice, Fields};
 use crate::flow::Flow;
 use crate::setup::{Obstacle, Setup};
 use crate::voxel::Grid;
@@ -83,39 +84,44 @@ impl FluidResult {
 }
 
 /// Solve the fluid step named `step` (or the model's only fluid step) with the obstacles' surfaces
-/// prepared by the caller.
+/// prepared by the caller, on the CPU.
 pub fn solve_fluid(model: &Model, step: Option<&str>, obstacles: &[Obstacle], limits: &Limits) -> Result<FluidResult> {
-    let setup = Setup::new(model, step, obstacles, limits)?;
-    let mut flow = Flow::new(&setup, obstacles)?;
-    run(&setup, &mut flow, limits)
+    solve_fluid_on(model, step, obstacles, limits, Choice::Cpu)
 }
 
-fn frame(setup: &Setup, flow: &Flow, step: u64) -> Result<Frame> {
+/// [`solve_fluid`] on the backend `choice` names (the registered device backend, or the CPU).
+pub fn solve_fluid_on(model: &Model, step: Option<&str>, obstacles: &[Obstacle], limits: &Limits, choice: Choice) -> Result<FluidResult> {
+    let setup = Setup::new(model, step, obstacles, limits)?;
+    let flow = Flow::new(&setup, obstacles)?;
+    let mut backend = crate::backend::backend(choice, flow, &setup)?;
+    run(&setup, backend.as_mut(), limits)
+}
+
+fn frame(setup: &Setup, flow: &mut dyn Backend, step: u64) -> Result<Frame> {
     let n = setup.grid.len();
     let mut velocity = Vec::new();
     let mut pressure = Vec::new();
     velocity.try_reserve_exact(n).map_err(|_| Error::Limit("not enough memory for an output frame".into()))?;
     pressure.try_reserve_exact(n).map_err(|_| Error::Limit("not enough memory for an output frame".into()))?;
     let vs = setup.velocity_scale();
+    let Fields { density, velocity: u } = flow.fields()?;
     for cell in 0..n {
-        match flow.moments(cell) {
-            Some((rho, u)) => {
-                velocity.push(u.map(|c| (c * vs) as f32));
-                pressure.push(setup.gauge_pressure(rho) as f32);
-            }
-            None => {
-                velocity.push([0.0; 3]);
-                pressure.push(0.0);
-            }
+        if flow.is_fluid(cell) {
+            velocity.push(u.get(cell).copied().unwrap_or([0.0; 3]).map(|c| (c * vs) as f32));
+            pressure.push(setup.gauge_pressure(density.get(cell).copied().unwrap_or(1.0)) as f32);
+        } else {
+            velocity.push([0.0; 3]);
+            pressure.push(0.0);
         }
     }
-    let mass = flow.mass() * setup.density * setup.grid.spacing.powi(3);
+    let mass = flow.mass()? * setup.density * setup.grid.spacing.powi(3);
     Ok(Frame { time: step as f64 * setup.time_step, step, velocity, pressure, mass })
 }
 
-/// Run a prepared flow (fresh from [`Flow::new`], possibly with its own initial state) through
-/// the setup's steps, recording frames and force histories.
-pub fn run(setup: &Setup, flow: &mut Flow, limits: &Limits) -> Result<FluidResult> {
+/// Run a prepared flow (fresh from [`Flow::new`], possibly with its own initial state, on any
+/// backend) through the setup's steps, recording frames and force histories. The steps go to
+/// the backend in chunks, up to the next frame, history sample or progress check.
+pub fn run(setup: &Setup, flow: &mut dyn Backend, limits: &Limits) -> Result<FluidResult> {
     let n = setup.grid.len();
     let total = setup.steps;
     let frames = u64::from(setup.frames).min(total).max(1);
@@ -135,21 +141,30 @@ pub fn run(setup: &Setup, flow: &mut Flow, limits: &Limits) -> Result<FluidResul
     let mut history_times = Vec::new();
     let mut forces: Vec<Vec<[f64; 3]>> = vec![Vec::new(); flow.groups().len()];
     let mut next_frame = 1u64;
+    // Frame j (1..=frames) at step round(j · total / frames).
+    let frame_step = |j: u64| (j * total + frames / 2) / frames;
     fieldcraft_model::progress::stage("Fluid steps", Some(total as f64), "");
-    for k in 1..=total {
+    let mut k = 0u64;
+    while k < total {
         fieldcraft_model::progress::check()?;
-        flow.step()?;
-        if k % 64 == 0 {
+        // Up to the next event: a history sample, a frame, a progress count.
+        let mut next = total.min(k + 64 - k % 64);
+        next = next.min((k / every + 1) * every);
+        if next_frame <= frames {
+            next = next.min(frame_step(next_frame).max(k + 1));
+        }
+        flow.advance(next - k)?;
+        k = next;
+        if k.is_multiple_of(64) || k == total {
             fieldcraft_model::progress::count(k as f64);
         }
-        if k % every == 0 || k == total {
+        if k.is_multiple_of(every) || k == total {
             history_times.push(k as f64 * setup.time_step);
-            for (h, f) in forces.iter_mut().zip(flow.forces()) {
+            for (h, f) in forces.iter_mut().zip(flow.forces()?) {
                 h.push(f.map(|c| c * fs));
             }
         }
-        // Frame j (1..=frames) at step round(j · total / frames).
-        while next_frame <= frames && k == (next_frame * total + frames / 2) / frames {
+        while next_frame <= frames && k == frame_step(next_frame) {
             out_frames.push(frame(setup, flow, k)?);
             next_frame += 1;
         }
@@ -158,7 +173,8 @@ pub fn run(setup: &Setup, flow: &mut Flow, limits: &Limits) -> Result<FluidResul
     let mut warnings = setup.warnings.clone();
     // The weakly compressible scheme turns pressure differences into density differences of
     // Δp / (ρ c_s²) in lattice units; large ones show the time step is too long for the flow.
-    let spread = (0..n).filter_map(|c| flow.moments(c)).map(|(rho, _)| (rho - 1.0).abs()).fold(0.0, f64::max);
+    let end = flow.fields()?;
+    let spread = (0..n).filter(|c| flow.is_fluid(*c)).filter_map(|c| end.density.get(c)).map(|rho| (rho - 1.0).abs()).fold(0.0, f64::max);
     if spread > COMPRESSIBILITY_WARNING {
         warnings.push(format!(
             "the density varies by {:.1}% at the end: the weakly compressible error is visible; a shorter time step reduces it (quadratically)",
@@ -168,6 +184,10 @@ pub fn run(setup: &Setup, flow: &mut Flow, limits: &Limits) -> Result<FluidResul
     for body in flow.empty_obstacles() {
         warnings.push(format!("obstacle `{body}` covers no cell centre: it is smaller than the spacing or outside the box"));
     }
+    if flow.name() != "cpu" {
+        warnings
+            .push(format!("solved on {}: single precision, deterministic on this device, within the backend's tolerance of the CPU", flow.name()));
+    }
     let solid: Vec<bool> = (0..n).map(|c| !flow.is_fluid(c)).collect();
     Ok(FluidResult {
         step: setup.step.clone(),