paulistrings/engine/
transport.rs

1//! The [`Transport`] trait, the [`Payload`] wire seam, the exchange's [`ChunkMap`], and the one-partition [`SoloTransport`].
2
3use crate::collectives::sealed::Sealed;
4use crate::collectives::Collectives;
5use crate::engine::coset::CosetOrder;
6
7/// The chunks a layer's bulk transfer is cut into, over the positions of the receiver's [`CosetOrder`].
8///
9/// Chunk edges always fall on coset boundaries, so a coset task waits on exactly one chunk.
10#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
11pub struct ChunkMap {
12    positions: u32,
13    /// `log2` of the coset size: coset `c` owns positions `c << coset_dimension .. (c + 1) << coset_dimension`.
14    coset_dimension: u32,
15    cosets: u32,
16    /// Clamped to `1..=cosets`.
17    chunks: u32,
18    /// `chunks` before the clamp, which every partition must have asked for alike; only the MPI header carries it.
19    #[cfg(feature = "mpi")]
20    requested: u32,
21}
22
23impl ChunkMap {
24    /// `positions` positions in cosets of `2^coset_dimension`, cut into at most `chunks` pieces; panics unless `positions` is a power of two and `chunks > 0`.
25    pub(crate) fn new(positions: usize, coset_dimension: usize, chunks: usize) -> Self {
26        assert!(
27            positions.is_power_of_two(),
28            "ChunkMap: {positions} positions is not a power of two",
29        );
30        assert!(chunks > 0, "ChunkMap: a layer needs at least one chunk");
31        let coset_dimension = coset_dimension as u32;
32        let positions = positions as u32;
33        debug_assert!(
34            (1u32 << coset_dimension) <= positions,
35            "ChunkMap: a coset of {} buckets does not fit in {positions}",
36            1u32 << coset_dimension,
37        );
38        let cosets = positions >> coset_dimension;
39        Self {
40            positions,
41            coset_dimension,
42            cosets,
43            chunks: (chunks as u32).min(cosets).max(1),
44            #[cfg(feature = "mpi")]
45            requested: chunks as u32,
46        }
47    }
48
49    /// The chunking of `order`'s positions.
50    pub(crate) fn of_order(order: &CosetOrder, chunks: usize) -> Self {
51        Self::new(order.positions(), order.dimension(), chunks)
52    }
53
54    /// Positions, i.e. buckets.
55    pub(crate) fn positions(&self) -> usize {
56        self.positions as usize
57    }
58
59    pub(crate) fn chunks(&self) -> usize {
60        self.chunks as usize
61    }
62
63    /// The chunk count this map was asked for, before the clamp to the coset count.
64    #[cfg(feature = "mpi")]
65    pub(crate) fn requested_chunks(&self) -> u32 {
66        self.requested
67    }
68
69    /// The first position of chunk `k`, a multiple of the coset size; `bound(chunks())` is the position count.
70    pub(crate) fn bound(&self, k: usize) -> u32 {
71        assert!(k <= self.chunks(), "ChunkMap: chunk {k} is out of range");
72        let coset = (k as u64 * u64::from(self.cosets)).div_ceil(u64::from(self.chunks)) as u32;
73        coset << self.coset_dimension
74    }
75
76    /// The `k` with `bound(k) <= position < bound(k + 1)`.
77    pub(crate) fn chunk_of_position(&self, position: u32) -> usize {
78        let coset = u64::from(position >> self.coset_dimension);
79        ((coset * u64::from(self.chunks)) / u64::from(self.cosets)) as usize
80    }
81}
82
83/// Something a [`Transport`] can move between partitions as zero-copy byte views of its own columns.
84///
85/// The in-process transport moves the typed value and calls none of these methods.
86pub trait Payload: Default + Send + 'static {
87    /// Borrowed byte views of this payload's columns, in wire order.
88    fn byte_parts(&self) -> Vec<&[u8]>;
89
90    /// Reshape for an incoming message of part lengths `lens` (growing, never shrinking, the storage) and hand out one mutable byte view per part, in `byte_parts` order.
91    ///
92    /// Panics if `lens` is not a shape this payload can take.
93    fn receive_into(&mut self, lens: &[usize]) -> Vec<&mut [u8]>;
94
95    /// Check, once the bytes have arrived and before the engine reads them, what `receive_into` could not; panics on an inconsistent payload.
96    fn finish_receive(&mut self);
97
98    /// The parts a two-phase transport must deliver before the body runs, in `byte_parts` order; every part by default.
99    fn early_parts(&self) -> Vec<&[u8]> {
100        self.byte_parts()
101    }
102
103    /// The parts after `early_parts`, each cut at `map`'s chunks: `bulk_parts()[i][k]` is part `i`'s slice for chunk `k`; none by default.
104    ///
105    /// Must cut exactly as `bulk_receive_into` does on the receiver, or the two sides post different message sizes.
106    fn bulk_parts(&self, map: &ChunkMap) -> Vec<Vec<&[u8]>> {
107        let _ = map;
108        Vec::new()
109    }
110
111    /// `receive_into` sizing every part from `lens` but handing out only the `early_parts` views.
112    fn early_receive_into(&mut self, lens: &[usize]) -> Vec<&mut [u8]> {
113        self.receive_into(lens)
114    }
115
116    /// The receive-side mirror of `bulk_parts`, called once the early parts have arrived.
117    fn bulk_receive_into(&mut self, map: &ChunkMap) -> Vec<Vec<&mut [u8]>> {
118        let _ = map;
119        Vec::new()
120    }
121}
122
123/// What a coset task waits on before it reads a received row of its chunk.
124///
125/// Called from every Rayon worker at once, so an MPI implementation must serialize itself (`MPI_THREAD_SERIALIZED`).
126pub trait ChunkWait: Sync {
127    /// Block until every row of chunk `k` has arrived; safe to call concurrently, repeatedly and out of order.
128    fn wait_chunk(&self, k: usize);
129}
130
131/// The [`ChunkWait`] of a transport whose exchange already completed.
132pub(super) struct AlreadyHere;
133
134impl ChunkWait for AlreadyHere {
135    fn wait_chunk(&self, _k: usize) {}
136}
137
138/// The per-layer all-to-all exchange between the partitions of a group.
139///
140/// Sealed through [`Collectives`]; not object-safe, so the layer loop monomorphizes over it.
141pub trait Transport: Collectives {
142    /// Send `send[q]` to partition `q`, run `body` on what the partners sent here, and return the received payloads with `body`'s result.
143    ///
144    /// Two-phase: the `early_parts` have arrived before `body` starts, and `body` calls `ChunkWait::wait_chunk` before reading a chunk's bulk rows.
145    /// `send.len()` must be the group size with `None` in the own slot and for any partner with nothing to send; every partition passes the same `map`.
146    /// `spare` is a payload pool in unspecified state, possibly empty: the transport takes receive payloads from it and returns finished send payloads to it.
147    fn exchange_layer<P, F, R>(
148        &self,
149        send: Vec<Option<P>>,
150        spare: &mut Vec<P>,
151        map: &ChunkMap,
152        body: F,
153    ) -> (Vec<Option<P>>, R)
154    where
155        P: Payload,
156        F: FnOnce(&[Option<P>], &dyn ChunkWait) -> R;
157
158    /// The blocking [`exchange_layer`](Self::exchange_layer) under an empty `ChunkMap`, for payloads whose every part is early.
159    fn exchange<P: Payload>(&self, send: Vec<Option<P>>, spare: &mut Vec<P>) -> Vec<Option<P>> {
160        self.exchange_layer(send, spare, &ChunkMap::default(), |_, _| ()).0
161    }
162
163    /// Collect every partition's `parts` on rank 0: `Some(v)` there with `v[q]` partition `q`'s parts, `None` elsewhere; one collective.
164    ///
165    /// A transport that infers its partner set from the `Some` positions, as the MPI one does, must override this asymmetric call.
166    fn gather_to_root(&self, parts: Vec<&[u8]>) -> Option<Vec<Vec<Vec<u8>>>> {
167        let size = self.size() as usize;
168        let this_rank = self.rank() as usize;
169        let mine = ByteParts(parts.iter().map(|p| p.to_vec()).collect());
170
171        if this_rank != ROOT {
172            let mut send: Vec<Option<ByteParts>> = (0..size).map(|_| None).collect();
173            send[ROOT] = Some(mine);
174            self.exchange(send, &mut Vec::new());
175            return None;
176        }
177
178        let mut received = self
179            .exchange((0..size).map(|_| None).collect::<Vec<Option<ByteParts>>>(), &mut Vec::new());
180        let mut mine = Some(mine);
181        Some(
182            (0..size)
183                .map(|q| {
184                    let got = if q == ROOT { mine.take() } else { received[q].take() };
185                    got.unwrap_or_else(|| {
186                        panic!(
187                            "gather_to_root: partition {q} sent nothing (it must send its \
188                                parts, empty or not)"
189                        )
190                    })
191                    .0
192                })
193                .collect(),
194        )
195    }
196}
197
198/// The group of one that unpartitioned propagation runs on: no reduction or exchange has anyone to talk to.
199pub(crate) struct SoloTransport;
200
201impl Sealed for SoloTransport {}
202
203impl Collectives for SoloTransport {
204    fn rank(&self) -> u32 {
205        0
206    }
207    fn size(&self) -> u32 {
208        1
209    }
210    fn allreduce_max_u8(&self, v: u8) -> u8 {
211        v
212    }
213    fn allreduce_sum_u64(&self, _buffer: &mut [u64]) {}
214    fn allreduce_sum_f64(&self, _buffer: &mut [f64]) {}
215    fn barrier(&self) {}
216}
217
218impl Transport for SoloTransport {
219    fn exchange_layer<P, F, R>(
220        &self,
221        send: Vec<Option<P>>,
222        _spare: &mut Vec<P>,
223        _map: &ChunkMap,
224        body: F,
225    ) -> (Vec<Option<P>>, R)
226    where
227        P: Payload,
228        F: FnOnce(&[Option<P>], &dyn ChunkWait) -> R,
229    {
230        debug_assert!(send.iter().all(Option::is_none), "a group of one sends to no one");
231        let result = body(&send, &AlreadyHere);
232        (send, result)
233    }
234}
235
236/// The partition every gather lands on.
237pub(super) const ROOT: usize = 0;
238
239/// A [`Payload`] that is already bytes, the default gather's wire form.
240#[derive(Default)]
241struct ByteParts(Vec<Vec<u8>>);
242
243impl Payload for ByteParts {
244    fn byte_parts(&self) -> Vec<&[u8]> {
245        self.0.iter().map(Vec::as_slice).collect()
246    }
247
248    fn receive_into(&mut self, lens: &[usize]) -> Vec<&mut [u8]> {
249        self.0.resize_with(lens.len(), Vec::new);
250        for (part, &len) in self.0.iter_mut().zip(lens) {
251            part.clear();
252            part.resize(len, 0);
253        }
254        self.0.iter_mut().map(Vec::as_mut_slice).collect()
255    }
256
257    fn finish_receive(&mut self) {}
258}
259
260#[cfg(test)]
261mod tests;
paulistrings/engine/transport.rs

paulistrings/engine/
transport.rs
