gridsim/
square_grid.rs

1#![allow(clippy::reversed_empty_ranges)]
2
3use crate::{Neumann, Sim};
4use ndarray::{par_azip, s, Array2, ArrayView2, ArrayViewMut2};
5use std::{
6    cell::UnsafeCell,
7    mem::{self, ManuallyDrop},
8};
9use itertools::Itertools;
10
11/// Represents the state of the simulation.
12#[derive(Clone, Debug)]
13pub struct SquareGrid<S>
14where
15    S: Sim<Neumann>,
16{
17    sim: S,
18    cells: Array2<S::Cell>,
19}
20
21impl<S> SquareGrid<S>
22where
23    S: Sim<Neumann>,
24    S::Cell: Send,
25{
26    /// Make a new grid with the given cells.
27    pub fn new(sim: S, mut original_cells: Array2<S::Cell>) -> Self {
28        let dims = original_cells.dim();
29        assert!(
30            dims.0 >= 1 && dims.1 >= 1,
31            "grid is empty, which isnt allowed"
32        );
33        let mut cells =
34            Array2::from_shape_simple_fn((dims.0 + 2, dims.1 + 2), || sim.cell_padding());
35        par_azip!((dest in &mut cells.slice_mut(s![1..-1, 1..-1]), cell in &mut original_cells) {
36            mem::swap(dest, cell);
37        });
38        Self { sim, cells }
39    }
40
41    /// Get view of cells on the grid.
42    pub fn cells(&self) -> ArrayView2<'_, S::Cell> {
43        self.cells.slice(s![1..-1, 1..-1])
44    }
45
46    /// Get mutable view of cells on the grid.
47    pub fn cells_mut(&mut self) -> ArrayViewMut2<'_, S::Cell> {
48        self.cells.slice_mut(s![1..-1, 1..-1])
49    }
50}
51
52#[cfg(feature = "use-rayon")]
53impl<S> SquareGrid<S>
54where
55    S: Sim<Neumann> + Sync,
56    S::Cell: Send + Sync,
57    S::Diff: Send + Sync,
58    S::Flow: Send,
59{
60    pub fn step_parallel(&mut self) {
61        let diffs = self.compute_diffs();
62        let flows = self.perform_egress(diffs.view());
63        self.perform_ingress(flows);
64    }
65
66    fn compute_diffs(&self) -> Array2<S::Diff> {
67        let mut diffs = Array2::from_shape_simple_fn(self.cells.dim(), || self.sim.diff_padding());
68        par_azip!((diff in diffs.slice_mut(s![1..-1, 1..-1]), cell in self.cells.windows((3, 3))) {
69            *diff = self.sim.compute(cell);
70        });
71        diffs
72    }
73
74    fn perform_egress(
75        &mut self,
76        diffs: ArrayView2<'_, S::Diff>,
77    ) -> Array2<ManuallyDrop<UnsafeCell<[S::Flow; 8]>>> {
78        let mut flows = Array2::from_shape_simple_fn(self.cells.dim(), || {
79            ManuallyDrop::new(UnsafeCell::new([
80                self.sim.flow_padding(),
81                self.sim.flow_padding(),
82                self.sim.flow_padding(),
83                self.sim.flow_padding(),
84                self.sim.flow_padding(),
85                self.sim.flow_padding(),
86                self.sim.flow_padding(),
87                self.sim.flow_padding(),
88            ]))
89        });
90        let sim = &self.sim;
91        par_azip!((flow in flows.slice_mut(s![1..-1, 1..-1]), cell in self.cells.slice_mut(s![1..-1, 1..-1]), diffs in diffs.windows((3, 3))) {
92            *flow.get_mut() = sim.egress(cell, diffs);
93        });
94
95        unsafe fn exchange_chunk<T>(chunk: ArrayViewMut2<'_, ManuallyDrop<UnsafeCell<[T; 8]>>>) {
96            let top_left = &mut *chunk[(0, 0)].get();
97            let top_right = &mut *chunk[(0, 1)].get();
98            let bottom_left = &mut *chunk[(1, 0)].get();
99            let bottom_right = &mut *chunk[(1, 1)].get();
100            // Left to right
101            mem::swap(&mut bottom_left[0], &mut bottom_right[4]);
102            // Top to bottom
103            mem::swap(&mut top_right[6], &mut bottom_right[2]);
104            // Across
105            mem::swap(&mut top_left[7], &mut bottom_right[3]);
106            mem::swap(&mut top_right[5], &mut bottom_left[1]);
107        }
108
109        // The flows need to be moved around to where they are consumed.
110        // We do this by splitting the grid into 2x2 chunks.
111        // We exchange the diagnal movements in each chunk, and then we also
112        // exchange two specific sides (but the specific sides can be chosen arbitrarily).
113        // This performs four swaps per chunk, exchanging eight flows.
114        // At this point we have only exchanged 1/4th of the total flows, but by sequentially
115        // performing this same operation offset by (0, 0), (0, 1), (1, 0), and (1, 1)
116        // we can actually exchange all flows in four simple parallel operations.
117        for (y, x) in (0..2).cartesian_product(0..2) {
118            par_azip!((chunk in flows.slice_mut(s![y.., x..]).exact_chunks_mut((2, 2))) {
119                unsafe { exchange_chunk(chunk); }
120            });
121        }
122
123        flows
124    }
125
126    fn perform_ingress(&mut self, mut flows: Array2<ManuallyDrop<UnsafeCell<[S::Flow; 8]>>>) {
127        let (h, w) = self.cells.dim();
128        let sim = &self.sim;
129        // At the end of this line, all of the manually drops MUST have been taken or dropped.
130        par_azip!((index (y, x), flow in &mut flows, cell in &mut self.cells) {
131            unsafe {
132                if (1..h-1).contains(&y) && (1..w-1).contains(&x) {
133                    // If its not part of the padding, we run the sim here.
134                    sim.ingress(cell, ManuallyDrop::take(flow).into_inner());
135                } else {
136                    // If this is part of the padding, we must manually drop.
137                    ManuallyDrop::drop(flow);
138                }
139            }
140        });
141    }
142}
gridsim/square_grid.rs

gridsim/
square_grid.rs
