mesh_sieve/algs/completion/
stack_completion.rs

1//! Complete the vertical‐stack arrows (mirror of section completion).
2//!
3//! This module provides routines for completing stack arrows in a distributed mesh,
4//! mirroring section completion logic. Communication uses explicit [`CommTag`]s and
5//! drains all send/receive handles before returning. Neighbor ranks are derived from
6//! the provided overlap graph, ensuring only true neighbors participate in the exchange.
7
8use std::collections::{BTreeSet, HashMap, HashSet};
9
10use crate::algs::wire::{WirePoint, cast_slice, cast_slice_mut};
11use bytemuck::{Pod, Zeroable};
12
13use crate::algs::communicator::{CommTag, StackCommTags, Wait};
14use crate::algs::completion::size_exchange::exchange_sizes_symmetric;
15use crate::mesh_error::MeshSieveError;
16use crate::topology::sieve::sieve_trait::Sieve;
17
18/// Trait for extracting rank from overlap payloads.
19pub trait HasRank {
20    /// Returns the MPI rank associated with this payload as a 32‑bit value.
21    fn rank_u32(&self) -> u32;
22}
23
24impl HasRank for crate::overlap::overlap::Remote {
25    #[inline]
26    fn rank_u32(&self) -> u32 {
27        u32::try_from(self.rank)
28            .expect("rank does not fit in u32; increase wire width or cap n_ranks")
29    }
30}
31
32/// Fixed width `(base, cap, payload)` triple used on the wire.
33#[repr(C)]
34#[derive(Copy, Clone, Zeroable)]
35struct WireTriple64<Pay>
36where
37    Pay: Copy + Pod + Zeroable,
38{
39    base_le: u64,
40    cap_le: u64,
41    pay: Pay,
42}
43
44impl<Pay: Copy + Pod + Zeroable> WireTriple64<Pay> {
45    fn new(base: u64, cap: u64, pay: Pay) -> Self {
46        Self {
47            base_le: base.to_le(),
48            cap_le: cap.to_le(),
49            pay,
50        }
51    }
52}
53
54unsafe impl<Pay: Copy + Pod + Zeroable> Pod for WireTriple64<Pay> {}
55
56/// Complete the stack by exchanging arrows with all true neighbor ranks.
57pub fn complete_stack_with_tags<P, Q, Pay, C, S, O, R>(
58    stack: &mut S,
59    overlap: &O,
60    comm: &C,
61    my_rank: usize,
62    n_ranks: usize,
63    tags: StackCommTags,
64) -> Result<(), MeshSieveError>
65where
66    P: WirePoint + Default + Eq + std::hash::Hash + Copy + Send + 'static,
67    Q: WirePoint + Default + Eq + std::hash::Hash + Copy + Send + 'static,
68    Pay: Copy + Pod + Zeroable + Default + PartialEq + Send + 'static,
69    C: crate::algs::communicator::Communicator + Sync,
70    S: crate::topology::stack::Stack<Point = P, CapPt = Q, VerticalPayload = Pay>,
71    O: Sieve<Point = P, Payload = R> + Sync,
72    R: HasRank + Copy + Send + 'static,
73{
74    if n_ranks == 0 {
75        return Err(MeshSieveError::CommError {
76            neighbor: my_rank,
77            source: "n_ranks must be > 0".into(),
78        });
79    }
80
81    // 1. Build owned links per neighbor
82    let mut nb_links: HashMap<usize, Vec<(P, Q, Pay)>> = HashMap::new();
83    for base in stack.base().base_points() {
84        let mut has_owned = false;
85        let mut owned_caps = Vec::new();
86        for (cap, pay) in stack.lift(base) {
87            if pay != Pay::default() {
88                has_owned = true;
89                owned_caps.push((cap, pay));
90            }
91        }
92        if !has_owned {
93            continue;
94        }
95        for (cap, pay) in owned_caps {
96            for (_dst, rem) in overlap.cone(base) {
97                let r = rem.rank_u32() as usize;
98                if r != my_rank {
99                    nb_links.entry(r).or_default().push((base, cap, pay));
100                }
101            }
102        }
103    }
104
105    // 2. Determine neighbor set and validate ranks.
106    //
107    // We still inspect the provided overlap to validate any referenced ranks,
108    // but for symmetric exchange we communicate with all ranks (except self).
109    // This ensures progress even if only one side seeded overlap structure.
110    let mut nb_seen: BTreeSet<usize> = BTreeSet::new();
111    for p in overlap.base_points() {
112        for (_dst, rem) in overlap.cone(p) {
113            nb_seen.insert(rem.rank_u32() as usize);
114        }
115    }
116    for &r in &nb_seen {
117        if r >= n_ranks {
118            return Err(MeshSieveError::CommError {
119                neighbor: r,
120                source: format!("rank {r} ≥ n_ranks {n_ranks}").into(),
121            });
122        }
123    }
124    // Symmetric neighbor set: all ranks except self.
125    let neighbors: Vec<usize> = (0..n_ranks).filter(|&r| r != my_rank).collect();
126    let all_neighbors: HashSet<usize> = neighbors.iter().copied().collect();
127
128    // 3. Build wire buffers per neighbor
129    let mut wires: HashMap<usize, Vec<WireTriple64<Pay>>> = HashMap::new();
130    for (&nbr, triples) in nb_links.iter() {
131        let mut buf = Vec::with_capacity(triples.len());
132        for &(b, c, p) in triples {
133            buf.push(WireTriple64::new(b.to_wire(), c.to_wire(), p));
134        }
135        wires.insert(nbr, buf);
136    }
137
138    // 4. Symmetric exchange of counts
139    let counts = exchange_sizes_symmetric(&wires, comm, tags.sizes, &all_neighbors)?;
140
141    // 5. Exchange payloads
142    let mut recv_data = Vec::new();
143    for &nbr in &neighbors {
144        let n = counts.get(&nbr).copied().unwrap_or(0) as usize;
145        let mut buf = vec![WireTriple64::<Pay>::zeroed(); n];
146        let h = comm.irecv_result(nbr, tags.data.as_u16(), cast_slice_mut(&mut buf))?;
147        recv_data.push((nbr, h, buf));
148    }
149
150    let mut pending_sends = Vec::new();
151    for &nbr in &neighbors {
152        let out = wires.get(&nbr).map_or(&[][..], |v| &v[..]);
153        pending_sends.push(comm.isend_result(nbr, tags.data.as_u16(), cast_slice(out))?);
154    }
155
156    let mut maybe_err: Option<MeshSieveError> = None;
157    for (nbr, h, mut buf) in recv_data {
158        match h.wait() {
159            Some(raw) if raw.len() == buf.len() * std::mem::size_of::<WireTriple64<Pay>>() => {
160                if maybe_err.is_none() {
161                    cast_slice_mut(&mut buf).copy_from_slice(&raw);
162                    for w in &buf {
163                        let b = P::from_wire(u64::from_le(w.base_le));
164                        let c = Q::from_wire(u64::from_le(w.cap_le));
165                        stack.add_arrow(b, c, w.pay)?;
166                    }
167                }
168            }
169            Some(raw) if maybe_err.is_none() => {
170                maybe_err = Some(MeshSieveError::CommError {
171                    neighbor: nbr,
172                    source: format!(
173                        "payload size mismatch: expected {}B, got {}B",
174                        buf.len() * std::mem::size_of::<WireTriple64<Pay>>(),
175                        raw.len()
176                    )
177                    .into(),
178                });
179            }
180            None if maybe_err.is_none() => {
181                maybe_err = Some(MeshSieveError::CommError {
182                    neighbor: nbr,
183                    source: "recv returned None".into(),
184                });
185            }
186            _ => {}
187        }
188    }
189
190    for s in pending_sends {
191        let _ = s.wait();
192    }
193
194    if let Some(e) = maybe_err {
195        Err(e)
196    } else {
197        Ok(())
198    }
199}
200
201/// Convenience wrapper using a legacy default base tag (0xC0DE).
202pub fn complete_stack<P, Q, Pay, C, S, O, R>(
203    stack: &mut S,
204    overlap: &O,
205    comm: &C,
206    my_rank: usize,
207    n_ranks: usize,
208) -> Result<(), MeshSieveError>
209where
210    P: WirePoint + Default + Eq + std::hash::Hash + Copy + Send + 'static,
211    Q: WirePoint + Default + Eq + std::hash::Hash + Copy + Send + 'static,
212    Pay: Copy + Pod + Zeroable + Default + PartialEq + Send + 'static,
213    C: crate::algs::communicator::Communicator + Sync,
214    S: crate::topology::stack::Stack<Point = P, CapPt = Q, VerticalPayload = Pay>,
215    O: Sieve<Point = P, Payload = R> + Sync,
216    R: HasRank + Copy + Send + 'static,
217{
218    // Legacy default tags keep ranks in sync for thread-local comms; use
219    // complete_stack_with_tags for concurrent or coordinated epochs.
220    let tags = StackCommTags::from_base(CommTag::new(0xC0DE));
221    complete_stack_with_tags::<P, Q, Pay, C, S, O, R>(stack, overlap, comm, my_rank, n_ranks, tags)
222}
mesh_sieve/algs/completion/stack_completion.rs

mesh_sieve/algs/completion/
stack_completion.rs
