mpi/
window.rs

1//! One-sided (RMA) communication: [`Window`]s of memory that remote ranks can
2//! `put` into, `get` from, and `accumulate` into without the target making a
3//! matching call. This corresponds to MPI's one-sided chapter (`MPI_Win_*`),
4//! which released rsmpi does not currently expose.
5//!
6//! Each window runs on its own communication context; a lightweight
7//! active-message handler on every rank services incoming RMA requests against
8//! the local window memory (guarded by a mutex), so operations complete even
9//! while the target is busy computing. Operations are synchronous (they wait
10//! for the target to acknowledge), which makes [`Window::fence`] a simple
11//! epoch barrier.
12//!
13//! ```no_run
14//! use mpi::traits::*;
15//! use mpi::window::Window;
16//!
17//! let universe = mpi::initialize().unwrap();
18//! let world = universe.world();
19//! let mut win = Window::<i32>::allocate(4, &world);
20//! win.fence();
21//! if world.rank() == 0 {
22//!     win.put(1, 0, &[10, 20, 30, 40]); // write into rank 1's window
23//! }
24//! win.fence();
25//! ```
26
27use std::marker::PhantomData;
28use std::sync::{Arc, Mutex};
29
30use crate::collective::{CommunicatorCollectives, Operation, SystemOperation};
31use crate::datatype::Equivalence;
32use crate::topology::{Communicator, SimpleCommunicator};
33use crate::transport;
34use crate::Rank;
35
36// Request tags (on the window request context).
37const OP_PUT: i32 = 1;
38const OP_GET: i32 = 2;
39const OP_ACC: i32 = 3;
40// Reply tags (on the window reply context).
41const REP_ACK: i32 = 1;
42const REP_GET: i32 = 2;
43
44#[inline]
45fn as_bytes<T>(v: &[T]) -> &[u8] {
46    // SAFETY: callers use POD `Equivalence` types.
47    unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, std::mem::size_of_val(v)) }
48}
49
50#[inline]
51fn as_bytes_mut<T>(v: &mut [T]) -> &mut [u8] {
52    // SAFETY: as above.
53    unsafe { std::slice::from_raw_parts_mut(v.as_mut_ptr() as *mut u8, std::mem::size_of_val(v)) }
54}
55
56fn u64_le(bytes: &[u8]) -> u64 {
57    u64::from_le_bytes(bytes[..8].try_into().unwrap())
58}
59
60/// A window of memory exposed for one-sided access (`MPI_Win`).
61pub struct Window<T: Equivalence + Send + Sync> {
62    comm: SimpleCommunicator,
63    req_ctx: u32,
64    reply_ctx: u32,
65    mem: Arc<Mutex<Vec<T>>>,
66    elem: usize,
67    _t: PhantomData<T>,
68}
69
70impl<T: Equivalence + Send + Sync> Window<T> {
71    /// Collectively create a window of `count` elements (zero-initialized) on
72    /// every rank of `comm` (`MPI_Win_allocate`).
73    pub fn allocate<C: Communicator>(count: usize, comm: &C) -> Window<T>
74    where
75        T: Default + Clone,
76    {
77        Window::from_vec(vec![T::default(); count], comm)
78    }
79
80    /// Collectively expose an existing buffer as a window (`MPI_Win_create`).
81    pub fn from_vec<C: Communicator>(data: Vec<T>, comm: &C) -> Window<T> {
82        // A private context for this window's traffic, agreed by all members.
83        let dup = comm.duplicate();
84        let base = dup.comm_data().derive_context(0x5749_4E00); // "WIN\0"
85        let req_ctx = dup.comm_data().derive_context(base ^ 0x11);
86        let reply_ctx = dup.comm_data().derive_context(base ^ 0x22);
87
88        let elem = std::mem::size_of::<T>().max(1);
89        let mem = Arc::new(Mutex::new(data));
90
91        // Register the active-message handler that services incoming RMA ops.
92        let mem_h = Arc::clone(&mem);
93        let world_ranks: Vec<i32> = dup.comm_data().world_ranks.clone();
94        let my_rank = dup.rank();
95        let handler: transport::Handler =
96            Arc::new(move |source, tag, _count, datatype, payload| {
97                let rt = transport::runtime();
98                let reply_to = world_ranks[source as usize];
99                match tag {
100                    OP_PUT => {
101                        let disp = u64_le(&payload) as usize;
102                        let data = &payload[8..];
103                        {
104                            let mut m = mem_h.lock().unwrap();
105                            let off = disp * elem;
106                            as_bytes_mut(&mut m[..])[off..off + data.len()].copy_from_slice(data);
107                        }
108                        let _ = rt.send(reply_ctx, my_rank, reply_to, REP_ACK, 0, datatype, &[]);
109                    }
110                    OP_GET => {
111                        let disp = u64_le(&payload) as usize;
112                        let cnt = u64_le(&payload[8..]) as usize;
113                        let out = {
114                            let m = mem_h.lock().unwrap();
115                            let off = disp * elem;
116                            as_bytes(&m[..])[off..off + cnt * elem].to_vec()
117                        };
118                        let _ = rt.send(
119                            reply_ctx, my_rank, reply_to, REP_GET, cnt as u64, datatype, &out,
120                        );
121                    }
122                    OP_ACC => {
123                        let op = SystemOperation::from_code(payload[0]);
124                        let disp = u64_le(&payload[1..]) as usize;
125                        let data = &payload[9..];
126                        {
127                            let mut m = mem_h.lock().unwrap();
128                            let off = disp * elem;
129                            op.reduce_bytes(
130                                datatype,
131                                &mut as_bytes_mut(&mut m[..])[off..off + data.len()],
132                                data,
133                            );
134                        }
135                        let _ = rt.send(reply_ctx, my_rank, reply_to, REP_ACK, 0, datatype, &[]);
136                    }
137                    _ => {}
138                }
139            });
140        transport::runtime().register_handler(req_ctx, handler);
141
142        Window {
143            comm: dup,
144            req_ctx,
145            reply_ctx,
146            mem,
147            elem,
148            _t: PhantomData,
149        }
150    }
151
152    /// This rank in the window's communicator.
153    pub fn rank(&self) -> Rank {
154        self.comm.rank()
155    }
156
157    /// Number of ranks sharing the window.
158    pub fn size(&self) -> Rank {
159        self.comm.size()
160    }
161
162    /// Number of elements in the local window memory.
163    pub fn local_len(&self) -> usize {
164        self.mem.lock().unwrap().len()
165    }
166
167    /// Inspect the local window memory.
168    pub fn with_local<R>(&self, f: impl FnOnce(&[T]) -> R) -> R {
169        f(&self.mem.lock().unwrap())
170    }
171
172    /// Mutate the local window memory.
173    pub fn with_local_mut<R>(&self, f: impl FnOnce(&mut [T]) -> R) -> R {
174        f(&mut self.mem.lock().unwrap())
175    }
176
177    /// Write `data` into rank `target`'s window at element offset
178    /// `target_disp` (`MPI_Put`). Completes when the target acknowledges.
179    pub fn put(&self, target: Rank, target_disp: usize, data: &[T]) {
180        let mut payload = Vec::with_capacity(8 + data.len() * self.elem);
181        payload.extend_from_slice(&(target_disp as u64).to_le_bytes());
182        payload.extend_from_slice(as_bytes(data));
183        self.request(
184            target,
185            OP_PUT,
186            T::equivalent_datatype().id,
187            data.len() as u64,
188            &payload,
189        );
190        let _ = transport::runtime().recv(self.reply_ctx, target, REP_ACK);
191    }
192
193    /// Read `buf.len()` elements from rank `target`'s window at element offset
194    /// `target_disp` into `buf` (`MPI_Get`).
195    pub fn get(&self, target: Rank, target_disp: usize, buf: &mut [T]) {
196        let mut payload = Vec::with_capacity(16);
197        payload.extend_from_slice(&(target_disp as u64).to_le_bytes());
198        payload.extend_from_slice(&(buf.len() as u64).to_le_bytes());
199        self.request(
200            target,
201            OP_GET,
202            T::equivalent_datatype().id,
203            buf.len() as u64,
204            &payload,
205        );
206        let (_s, _t, _c, _d, reply) = transport::runtime().recv(self.reply_ctx, target, REP_GET);
207        let dst = as_bytes_mut(buf);
208        let n = dst.len().min(reply.len());
209        dst[..n].copy_from_slice(&reply[..n]);
210    }
211
212    /// Combine `data` into rank `target`'s window at element offset
213    /// `target_disp` using `op` (`MPI_Accumulate`).
214    pub fn accumulate(&self, target: Rank, target_disp: usize, data: &[T], op: SystemOperation) {
215        let mut payload = Vec::with_capacity(9 + data.len() * self.elem);
216        payload.push(op.to_code());
217        payload.extend_from_slice(&(target_disp as u64).to_le_bytes());
218        payload.extend_from_slice(as_bytes(data));
219        self.request(
220            target,
221            OP_ACC,
222            T::equivalent_datatype().id,
223            data.len() as u64,
224            &payload,
225        );
226        let _ = transport::runtime().recv(self.reply_ctx, target, REP_ACK);
227    }
228
229    fn request(&self, target: Rank, tag: i32, datatype: u32, count: u64, payload: &[u8]) {
230        transport::runtime()
231            .send(
232                self.req_ctx,
233                self.comm.rank(),
234                self.comm.comm_data().world_rank(target),
235                tag,
236                count,
237                datatype,
238                payload,
239            )
240            .expect("RMA request send failed");
241    }
242
243    /// Collective synchronization delimiting an access epoch (`MPI_Win_fence`).
244    /// Because operations are synchronous, this is a barrier over the window's
245    /// communicator.
246    pub fn fence(&self) {
247        self.comm.barrier();
248    }
249
250    /// Begin passive-target access to `_target` (`MPI_Win_lock`). Operations in
251    /// this implementation are already atomic per element run, so this is a
252    /// no-op provided for API completeness.
253    pub fn lock(&self, _target: Rank) {}
254
255    /// End passive-target access (`MPI_Win_unlock`).
256    pub fn unlock(&self, _target: Rank) {}
257}
258
259impl<T: Equivalence + Send + Sync> Drop for Window<T> {
260    fn drop(&mut self) {
261        if transport::is_initialized() {
262            transport::runtime().unregister_handler(self.req_ctx);
263        }
264    }
265}
mpi/window.rs

mpi/
window.rs
