henad_models/gpu_sir/
mod.rs

1//! The SIR epidemic as a [`GpuGridModel`], with the rules of [`crate::sir`] and the grid in a storage buffer.
2//!
3//! The CPU model draws from one `xorshift64` stream along each row, and WGSL has no 64-bit integers to run it.
4//! Each cell instead keeps a `pcg_hash` state in a second buffer, `rng`, and the step advances it one round beside the
5//! cell. The engine ping-pongs both buffers together, and display and reduce read `state` alone.
6//!
7//! [`seed_cells`] repeats the draws of the CPU model's `init`, and the two backends agree at tick 0. Their random
8//! streams differ after tick 0.
9
10use henad_compute::cpu::grid_engine::{GRID_INIT_SEED, grid_init_rng};
11use henad_core::action::ActionDescriptor;
12use henad_core::authoring::model::binding::BindingDecl;
13use henad_core::authoring::model::gpu_grid_model::{GpuGridAction, GpuGridModel};
14use henad_core::authoring::primitives::rng::{mix_seed, pcg_hash, xorshift64};
15use henad_core::helpers::{extract_f32, extract_u32, f32_param, u32_param};
16use henad_core::params::{ParamDescriptor, ParamValue};
17use henad_core::view::{StatDescriptor, StatValue};
18
19use crate::shader_bindings::gpu_sir::seed_outbreak::Params as ActionParams;
20use crate::shader_bindings::gpu_sir::step::Params as StepParams;
21use crate::sir::PALETTE;
22
23/// Domain separator of the `rng` buffer's seed, so the buffer and the cell seeding stream do not start correlated.
24const RNG_INIT_SEED: u64 = GRID_INIT_SEED ^ 0x5EED_5EED_5EED_5EED;
25
26// The list repeats the ids and order the CPU engine composes for `SirGridModel`, so one parameter vector drives
27// either backend.
28henad_core::params! {
29    const PARAM_WIDTH = u32_param("grid_width", "Grid Width", DEFAULT_DIM, 1, 16_384);
30    const PARAM_HEIGHT = u32_param("grid_height", "Grid Height", DEFAULT_DIM, 1, 16_384);
31    const PARAM_INFECTION_RATE =
32        f32_param("infection_rate", "Infection Rate", DEFAULT_INFECTION_RATE, 0.0, 1.0, Some(0.01));
33    const PARAM_RECOVERY_RATE =
34        f32_param("recovery_rate", "Recovery Rate", DEFAULT_RECOVERY_RATE, 0.0, 1.0, Some(0.01));
35    const PARAM_INITIAL_INFECTED_PCT = f32_param(
36        "initial_infected_pct",
37        "Initial Infected",
38        DEFAULT_INITIAL_INFECTED_PCT,
39        0.0,
40        1.0,
41        Some(0.001),
42    )
43    .percent();
44}
45
46const DEFAULT_DIM: u32 = 1024;
47const DEFAULT_INFECTION_RATE: f32 = 0.3;
48const DEFAULT_RECOVERY_RATE: f32 = 0.05;
49const DEFAULT_INITIAL_INFECTED_PCT: f32 = 0.01;
50
51// Cell states, taken from the shader that defines them rather than mirrored.
52use crate::shader_bindings::gpu_sir::step::{I as CELL_I, R as CELL_R, S as CELL_S};
53
54/// Returns the initial cell states, one `u32` per cell, each infected with probability `initial_infected_pct`.
55///
56/// The cells are drawn from `rng` as `SirGridModel::init` draws them. With `rng` from `grid_init_rng`, the grid
57/// matches the CPU model's bit for bit.
58pub fn seed_cells(width: u32, height: u32, initial_infected_pct: f32, mut rng: u64) -> Vec<u32> {
59    let threshold = (initial_infected_pct * u32::MAX as f32) as u32;
60    let mut cells = vec![0u32; (width as usize) * (height as usize)];
61    for cell in &mut cells {
62        rng = xorshift64(rng);
63        *cell = u32::from(((rng >> 32) as u32) < threshold);
64    }
65    cells
66}
67
68/// Returns the initial `pcg_hash` state of every cell, hashed from `seed` and the cell's index.
69fn seed_rng_states(width: u32, height: u32, seed: u64) -> Vec<u32> {
70    let seed32 = (seed ^ (seed >> 32)) as u32;
71    (0..(width as usize) * (height as usize))
72        .map(|idx| pcg_hash(seed32 ^ idx as u32))
73        .collect()
74}
75
76/// The SIR epidemic as a [`GpuGridModel`], with a random number state per cell.
77#[derive(Debug)]
78pub struct GpuSir;
79
80impl GpuGridModel for GpuSir {
81    const NAME: &'static str = "SIR Epidemic (GPU)";
82    const ID: &'static str = "gpu_sir";
83    const DESCRIPTION: &'static str =
84        "Classic SIR compartmental model on a toroidal grid with Moore neighborhood, stepped entirely on the GPU";
85    const PALETTE: &'static [[u8; 4]] = &PALETTE;
86    const STATS: &'static [StatDescriptor] = &[
87        StatDescriptor::new("Susceptible", PALETTE[0]),
88        StatDescriptor::new("Infected", PALETTE[1]),
89        StatDescriptor::new("Recovered", PALETTE[2]),
90    ];
91
92    const BUFFERS: &'static [&'static str] = &["state", "rng"];
93
94    const STEP_BINDINGS: &'static [BindingDecl] = crate::binding_decls::bindings::GPU_SIR_STEP;
95    const DISPLAY_BINDINGS: &'static [BindingDecl] = crate::binding_decls::bindings::GPU_SIR_DISPLAY;
96    const REDUCE_BINDINGS: &'static [BindingDecl] = crate::binding_decls::bindings::GPU_SIR_REDUCE;
97
98    const STEP_SHADER: &'static str = crate::shader_bindings::gpu_sir::step::SHADER_STRING;
99    const DISPLAY_SHADER: &'static str = crate::shader_bindings::gpu_sir::display::SHADER_STRING;
100    const REDUCE_SHADER: &'static str = crate::shader_bindings::gpu_sir::reduce::SHADER_STRING;
101
102    const ACTIONS: &'static [GpuGridAction] = &[GpuGridAction {
103        desc: ActionDescriptor::new("seed_outbreak", "Seed outbreak"),
104        shader: crate::shader_bindings::gpu_sir::seed_outbreak::SHADER_STRING,
105        bindings: crate::binding_decls::bindings::GPU_SIR_SEED_OUTBREAK,
106    }];
107
108    fn param_descriptors() -> Vec<ParamDescriptor> {
109        descriptors()
110    }
111
112    fn dims(params: &[ParamValue]) -> (u32, u32) {
113        (
114            extract_u32(params, PARAM_WIDTH, DEFAULT_DIM),
115            extract_u32(params, PARAM_HEIGHT, DEFAULT_DIM),
116        )
117    }
118
119    fn seed_buffers(width: u32, height: u32, params: &[ParamValue], seed: Option<u64>) -> Vec<Vec<u32>> {
120        let initial_infected_pct = extract_f32(params, PARAM_INITIAL_INFECTED_PCT, DEFAULT_INITIAL_INFECTED_PCT);
121        let cells = grid_init_rng(seed);
122        let rng = seed.map_or(RNG_INIT_SEED, |s| mix_seed(s ^ RNG_INIT_SEED));
123        vec![
124            seed_cells(width, height, initial_infected_pct, cells),
125            seed_rng_states(width, height, rng),
126        ]
127    }
128
129    fn step_params_bytes(width: u32, height: u32, params: &[ParamValue]) -> Vec<u8> {
130        bytemuck::bytes_of(&StepParams {
131            width,
132            height,
133            infection_rate: extract_f32(params, PARAM_INFECTION_RATE, DEFAULT_INFECTION_RATE),
134            recovery_rate: extract_f32(params, PARAM_RECOVERY_RATE, DEFAULT_RECOVERY_RATE),
135        })
136        .to_vec()
137    }
138
139    fn action_params_bytes(_action: usize, width: u32, height: u32, params: &[ParamValue], seed: u32) -> Vec<u8> {
140        let pct = extract_f32(params, PARAM_INITIAL_INFECTED_PCT, DEFAULT_INITIAL_INFECTED_PCT);
141        bytemuck::bytes_of(&ActionParams {
142            width,
143            height,
144            threshold: (pct * u32::MAX as f32) as u32,
145            seed,
146        })
147        .to_vec()
148    }
149
150    fn stats(counts: &[u32]) -> Vec<StatValue> {
151        vec![
152            StatValue::Scalar(f64::from(counts[CELL_S as usize])),
153            StatValue::Scalar(f64::from(counts[CELL_I as usize])),
154            StatValue::Scalar(f64::from(counts[CELL_R as usize])),
155        ]
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use henad_compute::cpu::grid_engine::{GRID_PARAM_BASE, GridModelState};
163    use henad_compute::gpu::GpuContext;
164    use henad_compute::gpu::grid_engine::GpuGridState;
165    use henad_compute::gpu::sim_thread::GpuSimState as _;
166    use henad_core::authoring::model::grid_model::GridModel as _;
167    use henad_core::grid::Grid2D;
168    use henad_core::model::SimState as _;
169    use henad_core::view::StatEntry;
170    use henad_explore::testing::{TestDeviceRequest, headless_test_device};
171
172    use crate::sir::SirGridModel;
173
174    type State = GpuGridState<GpuSir>;
175
176    pub(super) fn headless_context() -> Option<GpuContext> {
177        headless_test_device(&TestDeviceRequest::baseline())
178    }
179
180    pub(super) fn params(
181        width: u32,
182        height: u32,
183        infection_rate: f32,
184        recovery_rate: f32,
185        initial_infected_pct: f32,
186    ) -> Vec<ParamValue> {
187        vec![
188            ParamValue::U32(width),
189            ParamValue::U32(height),
190            ParamValue::F32(infection_rate),
191            ParamValue::F32(recovery_rate),
192            ParamValue::F32(initial_infected_pct),
193        ]
194    }
195
196    fn sir_counts(state: &State) -> (u64, u64, u64) {
197        let stats = state.stats();
198        let scalar = |entry: &StatEntry| match &entry.value {
199            StatValue::Scalar(v) => *v as u64,
200            other => panic!("expected a scalar stat, got {other:?}"),
201        };
202        (scalar(&stats[0]), scalar(&stats[1]), scalar(&stats[2]))
203    }
204
205    /// Runs the display, reduce and readback passes as the sim thread's one-shot snapshot does.
206    fn refresh_stats(ctx: &GpuContext, state: &mut State) {
207        let mut encoder = ctx
208            .device
209            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
210        state.encode_snapshot_passes(&mut encoder);
211        ctx.queue.submit(Some(encoder.finish()));
212        state.begin_stats_readback();
213        state.poll_stats_readback(&ctx.device, true);
214    }
215
216    fn step_once(ctx: &GpuContext, state: &mut State) {
217        let mut encoder = ctx
218            .device
219            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
220        state.encode_steps(&mut encoder, 1, None);
221        ctx.queue.submit(Some(encoder.finish()));
222    }
223
224    /// The total of S, I and R holds exactly at every tick, whatever the per-cell RNG streams do.
225    #[test]
226    fn population_conserved_over_many_ticks() {
227        let Some(ctx) = headless_context() else {
228            log::warn!("skipping population_conserved_over_many_ticks: no adapter");
229            return;
230        };
231
232        let (width, height) = (64u32, 64u32);
233        let mut state = State::new(&ctx, &params(width, height, 0.3, 0.05, 0.1));
234        let total = u64::from(width) * u64::from(height);
235
236        for tick in 0..50 {
237            refresh_stats(&ctx, &mut state);
238            let (s, i, r) = sir_counts(&state);
239            assert_eq!(s + i + r, total, "S+I+R must equal total population at tick {tick}");
240            step_once(&ctx, &mut state);
241        }
242    }
243
244    /// With `infection_rate == 0`, S can never lose a member.
245    #[test]
246    fn zero_infection_rate_freezes_susceptible_count() {
247        let Some(ctx) = headless_context() else {
248            log::warn!("skipping zero_infection_rate_freezes_susceptible_count: no adapter");
249            return;
250        };
251
252        let (width, height) = (64u32, 64u32);
253        let mut state = State::new(&ctx, &params(width, height, 0.0, 0.05, 0.2));
254
255        refresh_stats(&ctx, &mut state);
256        let (initial_s, _, _) = sir_counts(&state);
257
258        for tick in 0..20 {
259            step_once(&ctx, &mut state);
260            refresh_stats(&ctx, &mut state);
261            let (s, _, _) = sir_counts(&state);
262            assert_eq!(
263                s, initial_s,
264                "susceptible count must not change at tick {tick} when infection_rate is 0"
265            );
266        }
267    }
268
269    /// With `recovery_rate == 0`, I can never lose a member, so at any positive infection rate I never decreases.
270    #[test]
271    fn zero_recovery_rate_keeps_infected_non_decreasing() {
272        let Some(ctx) = headless_context() else {
273            log::warn!("skipping zero_recovery_rate_keeps_infected_non_decreasing: no adapter");
274            return;
275        };
276
277        let (width, height) = (64u32, 64u32);
278        let mut state = State::new(&ctx, &params(width, height, 0.5, 0.0, 0.1));
279
280        refresh_stats(&ctx, &mut state);
281        let (_, mut prev_i, _) = sir_counts(&state);
282
283        for tick in 0..20 {
284            step_once(&ctx, &mut state);
285            refresh_stats(&ctx, &mut state);
286            let (_, i, r) = sir_counts(&state);
287            assert!(
288                i >= prev_i,
289                "infected count must not decrease at tick {tick} when recovery_rate is 0"
290            );
291            assert_eq!(r, 0, "no cell can reach R when recovery_rate is 0 (tick {tick})");
292            prev_i = i;
293        }
294    }
295
296    /// The port repeats `SirGridModel::init` rather than calling it, so every cell is compared.
297    #[test]
298    fn the_seeded_grid_matches_the_cpu_init() {
299        let (width, height) = (37u32, 23u32);
300        let p = params(width, height, 0.3, 0.05, 0.2);
301        for seed in [None, Some(7)] {
302            let mut grid = Grid2D::new(width, height);
303            SirGridModel::init(&mut grid, &p[GRID_PARAM_BASE..], &mut grid_init_rng(seed));
304            let cpu: Vec<u32> = grid.current().iter().copied().map(u32::from).collect();
305            assert_eq!(
306                GpuSir::seed_buffers(width, height, &p, seed)[0],
307                cpu,
308                "the seeded state differs from the CPU grid for seed {seed:?}"
309            );
310        }
311    }
312
313    /// Initial seeding uses the same PRNG, traversal order, and threshold as the CPU model, so the
314    /// tick-0 compartment counts (before any RNG-dependent step) must match exactly.
315    #[test]
316    fn initial_seed_matches_cpu_model() {
317        let Some(ctx) = headless_context() else {
318            log::warn!("skipping initial_seed_matches_cpu_model: no adapter");
319            return;
320        };
321
322        let (width, height) = (64u32, 64u32);
323        let p = params(width, height, 0.3, 0.05, 0.2);
324
325        let mut gpu = State::new(&ctx, &p);
326        let cpu = GridModelState::<SirGridModel>::from_params(&p);
327
328        refresh_stats(&ctx, &mut gpu);
329        let (gpu_s, gpu_i, gpu_r) = sir_counts(&gpu);
330
331        let cpu_scalar = |entry: &StatEntry| match &entry.value {
332            StatValue::Scalar(v) => *v as u64,
333            other => panic!("expected a scalar stat, got {other:?}"),
334        };
335        let cpu_stats = cpu.stats();
336        assert_eq!(
337            gpu_s,
338            cpu_scalar(&cpu_stats[0]),
339            "initial susceptible count must match the CPU model"
340        );
341        assert_eq!(
342            gpu_i,
343            cpu_scalar(&cpu_stats[1]),
344            "initial infected count must match the CPU model"
345        );
346        assert_eq!(
347            gpu_r,
348            cpu_scalar(&cpu_stats[2]),
349            "initial recovered count must be 0 on both backends"
350        );
351    }
352
353    #[test]
354    fn population_is_total_cells() {
355        let Some(ctx) = headless_context() else {
356            log::warn!("skipping population_is_total_cells: no adapter");
357            return;
358        };
359
360        let (width, height) = (32u32, 16u32);
361        let state = State::new(&ctx, &params(width, height, 0.3, 0.05, 0.1));
362        assert_eq!(state.population(), u64::from(width) * u64::from(height));
363    }
364}
365
366/// Tests that drive the real [`henad_compute::gpu::GpuSimThread`], similar to
367/// `gpu_game_of_life`'s `runner_tests`.
368#[cfg(test)]
369#[cfg(not(target_arch = "wasm32"))]
370mod runner_tests {
371    use std::time::{Duration, Instant};
372
373    use henad_compute::gpu::grid_engine::GpuGridState;
374    use henad_compute::gpu::sim_thread::{GpuBatchSettings, GpuSimThread};
375    use henad_compute::snapshot::{Snapshot, SnapshotView};
376
377    use super::GpuSir;
378    use super::tests::{headless_context, params};
379    use henad_compute::entry::ModelState;
380    use henad_core::view::StatValue;
381
382    fn wait_for(thread: &mut GpuSimThread, timeout: Duration, pred: impl Fn(&Snapshot) -> bool) -> Option<Snapshot> {
383        let deadline = Instant::now() + timeout;
384        while Instant::now() < deadline {
385            if let Some(snap) = thread.take_snapshot()
386                && pred(&snap)
387            {
388                return Some(snap);
389            }
390            std::thread::sleep(Duration::from_millis(2));
391        }
392        None
393    }
394
395    fn sir_counts(snap: &Snapshot) -> (u64, u64, u64) {
396        let scalar = |entry: &henad_core::view::StatEntry| match &entry.value {
397            StatValue::Scalar(v) => *v as u64,
398            other => panic!("expected a scalar stat, got {other:?}"),
399        };
400        (scalar(&snap.stats[0]), scalar(&snap.stats[1]), scalar(&snap.stats[2]))
401    }
402
403    #[test]
404    fn gpu_thread_publishes_gpu_snapshots_and_conserves_population() {
405        let Some(ctx) = headless_context() else {
406            log::warn!("skipping gpu_thread_publishes_gpu_snapshots_and_conserves_population: no adapter");
407            return;
408        };
409
410        let (width, height) = (128u32, 128u32);
411        let state = GpuGridState::<GpuSir>::new(&ctx, &params(width, height, 0.3, 0.05, 0.1));
412        let mut thread = GpuSimThread::new(ctx, Box::new(state), GpuBatchSettings::default(), None);
413
414        let initial = wait_for(&mut thread, Duration::from_secs(5), |_| true)
415            .expect("the GPU thread must publish an initial snapshot before Play");
416        assert!(matches!(initial.view, SnapshotView::Gpu(_)));
417        let total = initial.population;
418        let (s, i, r) = sir_counts(&initial);
419        assert_eq!(s + i + r, total, "initial snapshot must already conserve population");
420
421        thread.play();
422        let running = wait_for(&mut thread, Duration::from_secs(5), |snap| snap.tick > 0)
423            .expect("playing must advance the tick counter and publish fresh snapshots");
424        assert!(matches!(running.view, SnapshotView::Gpu(_)));
425        let (s, i, r) = sir_counts(&running);
426        assert_eq!(s + i + r, total, "population must stay conserved while running");
427
428        thread.pause();
429        drop(thread);
430    }
431
432    #[test]
433    fn registry_with_gpu_context_offers_a_drivable_gpu_sir_model() {
434        let Some(ctx) = headless_context() else {
435            log::warn!("skipping registry_with_gpu_context_offers_a_drivable_gpu_sir_model: no adapter");
436            return;
437        };
438
439        let models = crate::example_models();
440        let entry = models
441            .lookup("gpu_sir", Some(&ctx))
442            .expect("a GPU context must make the GPU SIR model selectable");
443
444        let built = entry
445            .build(&params(32, 32, 0.3, 0.05, 0.1), None, Some(&ctx))
446            .unwrap_or_else(|fault| panic!("the GPU SIR entry's factory failed to build: {fault}"));
447        let ModelState::Gpu(mut state) = built else {
448            panic!("the GPU SIR entry's factory must yield ModelState::Gpu, not ModelState::Cpu");
449        };
450
451        let mut encoder = ctx
452            .device
453            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
454        state.encode_steps(&mut encoder, 4, None);
455        ctx.queue.submit(Some(encoder.finish()));
456        assert_eq!(state.tick(), 4, "the registry-built state must be steppable");
457    }
458}
henad_models/gpu_sir/mod.rs

henad_models/gpu_sir/
mod.rs
