1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
//! Rank conventions shared by every distributed entry point.
//!
//! Frame ownership is a pure function of `(frame index, n_frames,
//! n_ranks)`: a contiguous block partition with the remainder going to
//! the lowest ranks. This part has no MPI dependency, so serial code and
//! tests can reason about the same partition the MPI paths use.
//!
//! With the `mpi` feature, [`Comm`] wraps a communicator handed in by
//! the caller. The library never calls `MPI_Init` or `MPI_Finalize` and
//! never falls back to the world communicator: every entry point takes
//! a communicator, either as a Rust [`Comm`] or as a Fortran `MPI_Fint`
//! handle over the C ABI (see `include/readcon-core-mpi.h`).
/// Half-open frame range `[start, end)` owned by `rank`.
///
/// Returns an empty range at `n_frames` when `rank >= n_ranks` or
/// `n_ranks == 0`.
pub fn block_range(n_frames: usize, n_ranks: usize, rank: usize) -> std::ops::Range<usize> {
if n_ranks == 0 || rank >= n_ranks {
return n_frames..n_frames;
}
let base = n_frames / n_ranks;
let rem = n_frames % n_ranks;
let start = rank * base + rank.min(rem);
let len = base + usize::from(rank < rem);
start..start + len
}
/// Rank that owns frame `i` under [`block_range`].
pub fn block_owner(n_frames: usize, n_ranks: usize, i: usize) -> Option<usize> {
if n_ranks == 0 || i >= n_frames {
return None;
}
let base = n_frames / n_ranks;
let rem = n_frames % n_ranks;
let big = rem * (base + 1);
Some(if i < big {
i / (base + 1)
} else {
rem + (i - big) / base
})
}
/// C ABI for [`block_range`]; writes `[*start, *end)`. Returns 0, or -1
/// on a null pointer.
///
/// # Safety
/// `start` and `end` must be null or valid for writes.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_block_range(
n_frames: u64,
n_ranks: u32,
rank: u32,
start: *mut u64,
end: *mut u64,
) -> i32 {
if start.is_null() || end.is_null() {
return -1;
}
let r = block_range(n_frames as usize, n_ranks as usize, rank as usize);
unsafe {
*start = r.start as u64;
*end = r.end as u64;
}
0
}
#[cfg(feature = "mpi")]
pub use with_mpi::*;
#[cfg(feature = "mpi")]
mod with_mpi {
use ::mpi::ffi;
use ::mpi::raw::FromRaw;
use ::mpi::topology::SimpleCommunicator;
use ::mpi::traits::*;
/// The caller's communicator, duplicated so library traffic never
/// matches the caller's own messages. The duplicate is freed on drop.
pub struct Comm {
inner: SimpleCommunicator,
}
impl Comm {
/// Duplicate any rsmpi communicator.
pub fn dup_of<C: Communicator>(c: &C) -> Self {
Comm {
inner: c.duplicate(),
}
}
/// Duplicate a raw `MPI_Comm` owned by the caller.
///
/// # Safety
/// MPI must be initialized and `raw` must be a live intra-communicator.
pub unsafe fn dup_raw(raw: ffi::MPI_Comm) -> Result<Self, i32> {
let mut out: ffi::MPI_Comm = unsafe { std::mem::zeroed() };
let rc = unsafe { ffi::MPI_Comm_dup(raw, &mut out) };
if rc != 0 {
return Err(rc);
}
Ok(Comm {
inner: unsafe { SimpleCommunicator::from_raw(out) },
})
}
/// Duplicate a Fortran `MPI_Fint` communicator handle.
///
/// # Safety
/// MPI must be initialized and `f` must name a live intra-communicator.
pub unsafe fn dup_fint(f: ffi::RSMPI_Fint) -> Result<Self, i32> {
unsafe { Self::dup_raw(ffi::RSMPI_Comm_f2c(f)) }
}
pub fn rank(&self) -> usize {
self.inner.rank() as usize
}
pub fn size(&self) -> usize {
self.inner.size() as usize
}
/// The wrapped communicator.
pub fn raw(&self) -> &SimpleCommunicator {
&self.inner
}
/// This rank's frames under [`super::block_range`].
pub fn my_frames(&self, n_frames: usize) -> std::ops::Range<usize> {
super::block_range(n_frames, self.size(), self.rank())
}
/// Broadcast a byte blob from `root`. `data` is read on `root` and
/// ignored elsewhere; every rank returns the root's bytes.
pub fn bcast_bytes(&self, root: usize, data: &[u8]) -> Vec<u8> {
let r = self.inner.process_at_rank(root as i32);
let mut n: u64 = if self.rank() == root {
data.len() as u64
} else {
0
};
r.broadcast_into(&mut n);
let mut buf = if self.rank() == root {
data.to_vec()
} else {
vec![0u8; n as usize]
};
if n > 0 {
r.broadcast_into(&mut buf[..]);
}
buf
}
}
/// Size and rank of a Fortran communicator; returns 0 or the MPI error.
///
/// # Safety
/// MPI must be initialized; `size` and `rank` must be valid for writes.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_mpi_comm_info_f(
comm_f: ffi::RSMPI_Fint,
size: *mut i32,
rank: *mut i32,
) -> i32 {
if size.is_null() || rank.is_null() {
return -1;
}
match unsafe { Comm::dup_fint(comm_f) } {
Ok(c) => {
unsafe {
*size = c.size() as i32;
*rank = c.rank() as i32;
}
0
}
Err(e) => e,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn block_partition_covers_and_matches_owner() {
for n in 0..40 {
for p in 1..9 {
let mut next = 0;
for r in 0..p {
let range = block_range(n, p, r);
assert_eq!(range.start, next);
for i in range.clone() {
assert_eq!(block_owner(n, p, i), Some(r));
}
next = range.end;
}
assert_eq!(next, n);
assert!(block_range(n, p, p).is_empty());
}
}
}
}