ruccl/in_process/collective/
mod.rs

1use super::buffer::{DistributedBuffer, RootedBuffer, VariableDistributedBuffer};
2use super::device::InProcessDevice;
3use super::error::InProcessError;
4use crate::rank::{CollectiveAlgorithm, CollectiveStats, CollectiveTransport};
5use std::marker::PhantomData;
6use std::ops::Range;
7
8mod all_to_all;
9mod exchange;
10mod memory;
11mod reduction;
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct PointToPointTransfer {
15    pub source_rank: usize,
16    pub source_range: Range<usize>,
17    pub destination_rank: usize,
18    pub destination_offset: usize,
19    pub tag: u64,
20}
21
22pub struct InProcessCollective<'a, T, D: InProcessDevice<T>, E> {
23    pub(super) id: u64,
24    pub(super) contexts: &'a [D::Context],
25    pub(super) marker: PhantomData<fn() -> (T, E)>,
26}
27
28impl<T, D, E> InProcessCollective<'_, T, D, E>
29where
30    T: Copy + Send + Sync + 'static,
31    D: InProcessDevice<T>,
32    E: From<D::Error> + From<InProcessError>,
33{
34    pub fn world_size(&self) -> usize {
35        self.contexts.len()
36    }
37
38    pub const fn transport(&self) -> CollectiveTransport {
39        CollectiveTransport::HostStaged
40    }
41
42    fn context(&self, rank: usize) -> Result<&D::Context, E> {
43        self.contexts
44            .get(rank)
45            .ok_or(InProcessError::RankOutOfRange {
46                rank,
47                world_size: self.world_size(),
48            })
49            .map_err(Into::into)
50    }
51
52    fn validate_buffer(&self, buffer: &DistributedBuffer<T, D>) -> Result<(), E> {
53        if buffer.communicator_id != self.id || buffer.world_size() != self.world_size() {
54            return Err(InProcessError::WrongCommunicator.into());
55        }
56        Ok(())
57    }
58
59    fn validate_rooted_buffer(&self, buffer: &RootedBuffer<T, D>) -> Result<(), E> {
60        if buffer.communicator_id != self.id {
61            return Err(InProcessError::WrongCommunicator.into());
62        }
63        self.context(buffer.root)?;
64        Ok(())
65    }
66
67    fn validate_variable_buffer(&self, buffer: &VariableDistributedBuffer<T, D>) -> Result<(), E> {
68        if buffer.communicator_id != self.id || buffer.world_size() != self.world_size() {
69            return Err(InProcessError::WrongCommunicator.into());
70        }
71        Ok(())
72    }
73}
74
75fn ring_chunks(length: u32, world_size: usize) -> Result<Vec<Range<usize>>, InProcessError> {
76    if world_size == 0 {
77        return Err(InProcessError::EmptyWorld);
78    }
79    let length = length as usize;
80    let base = length / world_size;
81    let remainder = length % world_size;
82    let mut start = 0;
83    Ok((0..world_size)
84        .map(|rank| {
85            let chunk_length = base + usize::from(rank < remainder);
86            let range = start..start + chunk_length;
87            start += chunk_length;
88            range
89        })
90        .collect())
91}
92
93fn transferred_bytes(
94    transfers: usize,
95    elements: usize,
96    element_bytes: usize,
97) -> Result<u64, InProcessError> {
98    let bytes = transfers
99        .checked_mul(elements)
100        .and_then(|value| value.checked_mul(element_bytes))
101        .ok_or(InProcessError::Overflow("transferred bytes"))?;
102    u64::try_from(bytes).map_err(|_| InProcessError::Overflow("transferred bytes"))
103}
ruccl/in_process/collective/mod.rs

ruccl/in_process/collective/
mod.rs
