ruccl/in_process/collective/
all_to_all.rs

1use super::*;
2
3impl<T, D, E> InProcessCollective<'_, T, D, E>
4where
5    T: Copy + Send + Sync + 'static,
6    D: InProcessDevice<T>,
7    E: From<D::Error> + From<InProcessError>,
8{
9    /// Equal-count all-to-all. Each input rank stores destination shards in
10    /// rank order; each output rank stores source shards in rank order.
11    pub fn all_to_all(
12        &self,
13        input: &DistributedBuffer<T, D>,
14    ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E> {
15        self.validate_buffer(input)?;
16        if !input.length_per_rank.is_multiple_of(self.world_size()) {
17            return Err(InProcessError::InvalidLength(
18                "all-to-all rank length must be divisible by world size",
19            )
20            .into());
21        }
22        let shard_length = input.length_per_rank / self.world_size();
23        let output = self.allocate(input.length_per_rank)?;
24        for source in 0..self.world_size() {
25            for destination in 0..self.world_size() {
26                let values = <D as InProcessDevice<T>>::copy_from_device_at(
27                    &self.contexts[source],
28                    &input.buffers[source],
29                    destination * shard_length,
30                    shard_length,
31                )?;
32                <D as InProcessDevice<T>>::copy_to_device_at(
33                    &self.contexts[destination],
34                    &output.buffers[destination],
35                    source * shard_length,
36                    &values,
37                )?;
38            }
39        }
40        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
41        stats.steps = self.world_size().saturating_sub(1) as u32;
42        stats.transferred_bytes = transferred_bytes(
43            self.world_size()
44                .saturating_mul(self.world_size().saturating_sub(1)),
45            shard_length,
46            D::ELEMENT_SIZE,
47        )?;
48        Ok((output, stats))
49    }
50
51    pub fn all_to_all_v(
52        &self,
53        input: &DistributedBuffer<T, D>,
54        send_counts: &[Vec<usize>],
55    ) -> Result<(VariableDistributedBuffer<T, D>, CollectiveStats), E> {
56        self.validate_buffer(input)?;
57        if send_counts.len() != self.world_size()
58            || send_counts
59                .iter()
60                .any(|counts| counts.len() != self.world_size())
61        {
62            return Err(InProcessError::InvalidLength(
63                "all-to-all-v send_counts must be a world_size square matrix",
64            )
65            .into());
66        }
67        for counts in send_counts {
68            let total = counts
69                .iter()
70                .try_fold(0_usize, |sum, count| sum.checked_add(*count))
71                .ok_or(InProcessError::Overflow("all-to-all-v send count"))?;
72            if total != input.length_per_rank {
73                return Err(InProcessError::InvalidLength(
74                    "each all-to-all-v send-count row must sum to the input rank length",
75                )
76                .into());
77            }
78        }
79        let receive_lengths = (0..self.world_size())
80            .map(|destination| {
81                send_counts.iter().try_fold(0_usize, |total, counts| {
82                    total.checked_add(counts[destination])
83                })
84            })
85            .collect::<Option<Vec<_>>>()
86            .ok_or(InProcessError::Overflow("all-to-all-v receive count"))?;
87        let output = self.allocate_variable(&receive_lengths)?;
88        let mut receive_offsets = vec![0_usize; self.world_size()];
89        for (source, counts) in send_counts.iter().enumerate() {
90            let values = self.download_rank(input, source)?;
91            let mut send_offset = 0_usize;
92            for destination in 0..self.world_size() {
93                let count = counts[destination];
94                let send_end = send_offset
95                    .checked_add(count)
96                    .ok_or(InProcessError::Overflow("all-to-all-v source range"))?;
97                if count != 0 {
98                    let destination_buffer = output.buffers[destination]
99                        .as_ref()
100                        .expect("non-empty all-to-all-v receive has a device buffer");
101                    <D as InProcessDevice<T>>::copy_to_device_at(
102                        &self.contexts[destination],
103                        destination_buffer,
104                        receive_offsets[destination],
105                        &values[send_offset..send_end],
106                    )?;
107                }
108                receive_offsets[destination] = receive_offsets[destination]
109                    .checked_add(count)
110                    .ok_or(InProcessError::Overflow("all-to-all-v destination range"))?;
111                send_offset = send_end;
112            }
113        }
114        let remote_elements = send_counts
115            .iter()
116            .enumerate()
117            .try_fold(0_usize, |total, (source, counts)| {
118                counts
119                    .iter()
120                    .enumerate()
121                    .try_fold(total, |total, (destination, count)| {
122                        if source == destination {
123                            Some(total)
124                        } else {
125                            total.checked_add(*count)
126                        }
127                    })
128            })
129            .ok_or(InProcessError::Overflow(
130                "all-to-all-v transferred elements",
131            ))?;
132        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
133        stats.steps = self.world_size().saturating_sub(1) as u32;
134        stats.transferred_bytes = transferred_bytes(1, remote_elements, D::ELEMENT_SIZE)?;
135        Ok((output, stats))
136    }
137}
ruccl/in_process/collective/all_to_all.rs

ruccl/in_process/collective/
all_to_all.rs
