ruccl/in_process/collective/
exchange.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    pub fn broadcast(
10        &self,
11        buffer: &DistributedBuffer<T, D>,
12        root: usize,
13    ) -> Result<CollectiveStats, E> {
14        self.validate_buffer(buffer)?;
15        let root_values = self.download_rank(buffer, root)?;
16        for rank in 0..self.world_size() {
17            if rank != root {
18                <D as InProcessDevice<T>>::copy_to_device(
19                    &self.contexts[rank],
20                    &buffer.buffers[rank],
21                    &root_values,
22                )?;
23            }
24        }
25        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
26        stats.steps = u32::from(self.world_size() > 1);
27        stats.transferred_bytes = transferred_bytes(
28            self.world_size().saturating_sub(1),
29            buffer.length_per_rank,
30            D::ELEMENT_SIZE,
31        )?;
32        Ok(stats)
33    }
34
35    pub fn all_gather(
36        &self,
37        input: &DistributedBuffer<T, D>,
38    ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E> {
39        self.validate_buffer(input)?;
40        let output_length = input
41            .length_per_rank
42            .checked_mul(self.world_size())
43            .ok_or(InProcessError::Overflow("all-gather output length"))?;
44        let output = self.allocate(output_length)?;
45        for rank in 0..self.world_size() {
46            let local = self.download_rank(input, rank)?;
47            <D as InProcessDevice<T>>::copy_to_device_at(
48                &self.contexts[rank],
49                &output.buffers[rank],
50                rank * input.length_per_rank,
51                &local,
52            )?;
53        }
54        for step in 0..self.world_size().saturating_sub(1) {
55            let staged = (0..self.world_size())
56                .map(|rank| {
57                    let send_rank = (rank + self.world_size() - step) % self.world_size();
58                    <D as InProcessDevice<T>>::copy_from_device_at(
59                        &self.contexts[rank],
60                        &output.buffers[rank],
61                        send_rank * input.length_per_rank,
62                        input.length_per_rank,
63                    )
64                })
65                .collect::<Result<Vec<_>, _>>()?;
66            for rank in 0..self.world_size() {
67                let source_rank = (rank + self.world_size() - 1) % self.world_size();
68                let receive_rank = (rank + self.world_size() - step - 1) % self.world_size();
69                <D as InProcessDevice<T>>::copy_to_device_at(
70                    &self.contexts[rank],
71                    &output.buffers[rank],
72                    receive_rank * input.length_per_rank,
73                    &staged[source_rank],
74                )?;
75            }
76        }
77        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Ring);
78        stats.steps = self.world_size().saturating_sub(1) as u32;
79        stats.transferred_bytes = transferred_bytes(
80            self.world_size()
81                .saturating_mul(self.world_size().saturating_sub(1)),
82            input.length_per_rank,
83            D::ELEMENT_SIZE,
84        )?;
85        Ok((output, stats))
86    }
87
88    pub fn gather(
89        &self,
90        input: &DistributedBuffer<T, D>,
91        root: usize,
92    ) -> Result<(RootedBuffer<T, D>, CollectiveStats), E> {
93        self.validate_buffer(input)?;
94        let output_length = input
95            .length_per_rank
96            .checked_mul(self.world_size())
97            .ok_or(InProcessError::Overflow("gather output length"))?;
98        let output = self.allocate_rooted(root, output_length)?;
99        let root_context = self.context(root)?;
100        for rank in 0..self.world_size() {
101            let values = self.download_rank(input, rank)?;
102            <D as InProcessDevice<T>>::copy_to_device_at(
103                root_context,
104                &output.buffer,
105                rank * input.length_per_rank,
106                &values,
107            )?;
108        }
109        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
110        stats.steps = u32::from(self.world_size() > 1);
111        stats.transferred_bytes = transferred_bytes(
112            self.world_size().saturating_sub(1),
113            input.length_per_rank,
114            D::ELEMENT_SIZE,
115        )?;
116        Ok((output, stats))
117    }
118
119    pub fn scatter(
120        &self,
121        input: &RootedBuffer<T, D>,
122    ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E> {
123        self.validate_rooted_buffer(input)?;
124        if !input.len().is_multiple_of(self.world_size()) {
125            return Err(InProcessError::InvalidLength(
126                "scatter input length must be divisible by world size",
127            )
128            .into());
129        }
130        let output_length = input.len() / self.world_size();
131        let output = self.allocate(output_length)?;
132        let root_context = self.context(input.root)?;
133        for rank in 0..self.world_size() {
134            let values = <D as InProcessDevice<T>>::copy_from_device_at(
135                root_context,
136                &input.buffer,
137                rank * output_length,
138                output_length,
139            )?;
140            <D as InProcessDevice<T>>::copy_to_device(
141                &self.contexts[rank],
142                &output.buffers[rank],
143                &values,
144            )?;
145        }
146        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
147        stats.steps = u32::from(self.world_size() > 1);
148        stats.transferred_bytes = transferred_bytes(
149            self.world_size().saturating_sub(1),
150            output_length,
151            D::ELEMENT_SIZE,
152        )?;
153        Ok((output, stats))
154    }
155
156    /// Tagged point-to-point transfer between two rank-local buffers. This is
157    /// synchronous on the host-staged transport; `tag` is retained in the API
158    /// because network transports match send and receive by that value.
159    pub fn send_recv(
160        &self,
161        source: &DistributedBuffer<T, D>,
162        destination: &DistributedBuffer<T, D>,
163        transfer: PointToPointTransfer,
164    ) -> Result<CollectiveStats, E> {
165        self.validate_buffer(source)?;
166        self.validate_buffer(destination)?;
167        let source_context = self.context(transfer.source_rank)?;
168        let destination_context = self.context(transfer.destination_rank)?;
169        if transfer.source_range.start > transfer.source_range.end
170            || transfer.source_range.end > source.length_per_rank
171        {
172            return Err(InProcessError::InvalidLength(
173                "send source range is outside the rank buffer",
174            )
175            .into());
176        }
177        let length = transfer.source_range.len();
178        let destination_end = transfer
179            .destination_offset
180            .checked_add(length)
181            .ok_or(InProcessError::Overflow("send destination range"))?;
182        if destination_end > destination.length_per_rank {
183            return Err(InProcessError::InvalidLength(
184                "receive destination range is outside the rank buffer",
185            )
186            .into());
187        }
188        let values = <D as InProcessDevice<T>>::copy_from_device_at(
189            source_context,
190            &source.buffers[transfer.source_rank],
191            transfer.source_range.start,
192            length,
193        )?;
194        <D as InProcessDevice<T>>::copy_to_device_at(
195            destination_context,
196            &destination.buffers[transfer.destination_rank],
197            transfer.destination_offset,
198            &values,
199        )?;
200        let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
201        stats.steps = u32::from(length != 0);
202        stats.transferred_bytes = transferred_bytes(1, length, D::ELEMENT_SIZE)?;
203        Ok(stats)
204    }
205}
ruccl/in_process/collective/exchange.rs

ruccl/in_process/collective/
exchange.rs
