ruccl/in_process/
owner.rs

1use super::collective::InProcessCollective;
2use super::device::InProcessDevice;
3use super::error::InProcessError;
4use crate::rank::{CollectiveAlgorithm, CollectiveStats, CollectiveTransport};
5use std::marker::PhantomData;
6use std::sync::atomic::{AtomicU64, Ordering};
7
8static NEXT_COMMUNICATOR_ID: AtomicU64 = AtomicU64::new(1);
9
10#[derive(Debug, Clone)]
11pub struct Communicator<C> {
12    id: u64,
13    contexts: Vec<C>,
14}
15
16impl<C> Communicator<C> {
17    pub fn new<K, E>(
18        contexts: Vec<C>,
19        initialize: impl FnOnce(&C) -> Result<K, E>,
20    ) -> Result<(Self, K), E>
21    where
22        E: From<InProcessError>,
23    {
24        let Some(first) = contexts.first() else {
25            return Err(InProcessError::EmptyWorld.into());
26        };
27        let kernels = initialize(first)?;
28        Ok((
29            Self {
30                id: NEXT_COMMUNICATOR_ID.fetch_add(1, Ordering::Relaxed),
31                contexts,
32            },
33            kernels,
34        ))
35    }
36
37    pub fn world_size(&self) -> usize {
38        self.contexts.len()
39    }
40
41    pub const fn transport(&self) -> CollectiveTransport {
42        CollectiveTransport::HostStaged
43    }
44
45    pub fn collective<T, D, E>(&self) -> InProcessCollective<'_, T, D, E>
46    where
47        D: InProcessDevice<T, Context = C>,
48        E: From<D::Error> + From<InProcessError>,
49    {
50        InProcessCollective {
51            id: self.id,
52            contexts: &self.contexts,
53            marker: PhantomData,
54        }
55    }
56
57    pub fn barrier(&self) -> Result<CollectiveStats, InProcessError> {
58        if self.world_size() == 0 {
59            return Err(InProcessError::EmptyWorld);
60        }
61        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
62        stats.steps = u32::from(self.world_size() > 1);
63        Ok(stats)
64    }
65}
ruccl/in_process/owner.rs

ruccl/in_process/
owner.rs
