ruccl/in_process/collective/
reduction.rs1use 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 reduce_sum(
10 &self,
11 input: &DistributedBuffer<T, D>,
12 root: usize,
13 function: &D::Kernel,
14 ) -> Result<(RootedBuffer<T, D>, CollectiveStats), E> {
15 self.validate_buffer(input)?;
16 let root_context = self.context(root)?;
17 let output = self.allocate_rooted(root, input.length_per_rank)?;
18 let root_values = self.download_rank(input, root)?;
19 <D as InProcessDevice<T>>::copy_to_device(root_context, &output.buffer, &root_values)?;
20 if input.length_per_rank == 0 {
21 return Ok((output, CollectiveStats::new(CollectiveAlgorithm::Direct)));
22 }
23 let scratch = <D as InProcessDevice<T>>::alloc(root_context, input.length_per_rank)?;
24 let length = u32::try_from(input.length_per_rank)
25 .map_err(|_| InProcessError::InvalidLength("GX kernels use u32 element indices"))?;
26 let launch = <D as InProcessDevice<T>>::prepare_reduction(length, 0);
27 let mut stats = CollectiveStats::new(CollectiveAlgorithm::Direct);
28 for rank in 0..self.world_size() {
29 if rank == root {
30 continue;
31 }
32 let values = self.download_rank(input, rank)?;
33 <D as InProcessDevice<T>>::copy_to_device(root_context, &scratch, &values)?;
34 <D as InProcessDevice<T>>::launch_reduction(
35 root_context,
36 function,
37 &launch,
38 &scratch,
39 &output.buffer,
40 )?;
41 stats.steps += 1;
42 stats.reduction_kernel_launches += 1;
43 }
44 stats.transferred_bytes = transferred_bytes(
45 self.world_size().saturating_sub(1),
46 input.length_per_rank,
47 D::ELEMENT_SIZE,
48 )?;
49 Ok((output, stats))
50 }
51
52 pub fn all_reduce_sum(
53 &self,
54 buffer: &DistributedBuffer<T, D>,
55 function: &D::Kernel,
56 ) -> Result<CollectiveStats, E> {
57 self.validate_buffer(buffer)?;
58 let world_size = self.world_size();
59 let mut stats = CollectiveStats::new(CollectiveAlgorithm::Ring);
60 if world_size == 1 {
61 return Ok(stats);
62 }
63 let length_u32 = u32::try_from(buffer.length_per_rank)
64 .map_err(|_| InProcessError::InvalidLength("GX kernels use u32 element indices"))?;
65 let chunks = ring_chunks(length_u32, world_size)?;
66 let max_chunk = chunks.iter().map(|range| range.len()).max().unwrap_or(0);
67 let scratch = self
68 .contexts
69 .iter()
70 .map(|context| <D as InProcessDevice<T>>::alloc(context, max_chunk.max(1)))
71 .collect::<Result<Vec<_>, _>>()?;
72
73 for step in 0..world_size - 1 {
74 let staged = (0..world_size)
75 .map(|rank| {
76 let send_chunk = (rank + world_size - step - 1) % world_size;
77 let range = &chunks[send_chunk];
78 <D as InProcessDevice<T>>::copy_from_device_at(
79 &self.contexts[rank],
80 &buffer.buffers[rank],
81 range.start,
82 range.len(),
83 )
84 })
85 .collect::<Result<Vec<_>, _>>()?;
86 for (rank, scratch_buffer) in scratch.iter().enumerate() {
87 let source_rank = (rank + world_size - 1) % world_size;
88 let receive_chunk = (rank + world_size - step - 2) % world_size;
89 let range = &chunks[receive_chunk];
90 if range.is_empty() {
91 continue;
92 }
93 <D as InProcessDevice<T>>::copy_to_device_at(
94 &self.contexts[rank],
95 scratch_buffer,
96 0,
97 &staged[source_rank],
98 )?;
99 let launch = <D as InProcessDevice<T>>::prepare_reduction(
100 range.len() as u32,
101 range.start as u32,
102 );
103 <D as InProcessDevice<T>>::launch_reduction(
104 &self.contexts[rank],
105 function,
106 &launch,
107 scratch_buffer,
108 &buffer.buffers[rank],
109 )?;
110 stats.reduction_kernel_launches += 1;
111 }
112 stats.steps += 1;
113 }
114
115 for step in 0..world_size - 1 {
116 let staged = (0..world_size)
117 .map(|rank| {
118 let send_chunk = (rank + world_size - step) % world_size;
119 let range = &chunks[send_chunk];
120 <D as InProcessDevice<T>>::copy_from_device_at(
121 &self.contexts[rank],
122 &buffer.buffers[rank],
123 range.start,
124 range.len(),
125 )
126 })
127 .collect::<Result<Vec<_>, _>>()?;
128 for rank in 0..world_size {
129 let source_rank = (rank + world_size - 1) % world_size;
130 let receive_chunk = (rank + world_size - step - 1) % world_size;
131 let range = &chunks[receive_chunk];
132 if !range.is_empty() {
133 <D as InProcessDevice<T>>::copy_to_device_at(
134 &self.contexts[rank],
135 &buffer.buffers[rank],
136 range.start,
137 &staged[source_rank],
138 )?;
139 }
140 }
141 stats.steps += 1;
142 }
143
144 stats.transferred_bytes = transferred_bytes(
145 2 * (world_size - 1),
146 buffer.length_per_rank,
147 D::ELEMENT_SIZE,
148 )?;
149 Ok(stats)
150 }
151
152 pub fn reduce_scatter_sum(
153 &self,
154 input: &DistributedBuffer<T, D>,
155 function: &D::Kernel,
156 ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E> {
157 self.validate_buffer(input)?;
158 let world_size = self.world_size();
159 if !input.length_per_rank.is_multiple_of(world_size) {
160 return Err(InProcessError::InvalidLength(
161 "reduce-scatter input length must be divisible by world size",
162 )
163 .into());
164 }
165 let output_length = input.length_per_rank / world_size;
166 let length_u32 = u32::try_from(input.length_per_rank)
167 .map_err(|_| InProcessError::InvalidLength("GX kernels use u32 element indices"))?;
168 let chunks = ring_chunks(length_u32, world_size)?;
169 let working = self.allocate(input.length_per_rank)?;
170 for rank in 0..world_size {
171 let values = self.download_rank(input, rank)?;
172 <D as InProcessDevice<T>>::copy_to_device(
173 &self.contexts[rank],
174 &working.buffers[rank],
175 &values,
176 )?;
177 }
178 let scratch = self
179 .contexts
180 .iter()
181 .map(|context| <D as InProcessDevice<T>>::alloc(context, output_length))
182 .collect::<Result<Vec<_>, _>>()?;
183 let mut stats = CollectiveStats::new(CollectiveAlgorithm::Ring);
184 for step in 0..world_size.saturating_sub(1) {
185 let staged = (0..world_size)
186 .map(|rank| {
187 let send_chunk = (rank + world_size - step - 1) % world_size;
188 let range = &chunks[send_chunk];
189 <D as InProcessDevice<T>>::copy_from_device_at(
190 &self.contexts[rank],
191 &working.buffers[rank],
192 range.start,
193 range.len(),
194 )
195 })
196 .collect::<Result<Vec<_>, _>>()?;
197 for (rank, scratch_buffer) in scratch.iter().enumerate() {
198 let source_rank = (rank + world_size - 1) % world_size;
199 let receive_chunk = (rank + world_size - step - 2) % world_size;
200 let range = &chunks[receive_chunk];
201 if range.is_empty() {
202 continue;
203 }
204 <D as InProcessDevice<T>>::copy_to_device(
205 &self.contexts[rank],
206 scratch_buffer,
207 &staged[source_rank],
208 )?;
209 let launch = <D as InProcessDevice<T>>::prepare_reduction(
210 range.len() as u32,
211 range.start as u32,
212 );
213 <D as InProcessDevice<T>>::launch_reduction(
214 &self.contexts[rank],
215 function,
216 &launch,
217 scratch_buffer,
218 &working.buffers[rank],
219 )?;
220 stats.reduction_kernel_launches += 1;
221 }
222 stats.steps += 1;
223 }
224
225 let output = self.allocate(output_length)?;
226 for (rank, range) in chunks.iter().enumerate() {
227 let reduced = <D as InProcessDevice<T>>::copy_from_device_at(
228 &self.contexts[rank],
229 &working.buffers[rank],
230 range.start,
231 range.len(),
232 )?;
233 <D as InProcessDevice<T>>::copy_to_device(
234 &self.contexts[rank],
235 &output.buffers[rank],
236 &reduced,
237 )?;
238 }
239 stats.transferred_bytes = transferred_bytes(
240 world_size.saturating_sub(1),
241 input.length_per_rank,
242 D::ELEMENT_SIZE,
243 )?;
244 Ok((output, stats))
245 }
246}