burn_tensor/tensor/api/collective.rs
1use alloc::vec::Vec;
2
3use burn_std::distributed::ReduceOperation;
4
5use crate::{
6 AsIndex, Device, Tensor, check,
7 check::{TensorCheck, unwrap_dim_index},
8 kind::{Basic, Numeric},
9 ops::BridgeTensor,
10};
11
12impl<const D: usize, K: Numeric> Tensor<D, K> {
13 /// Reduces every member's tensor with `op`, and gives every member the result.
14 ///
15 /// Member `i` is position `i` of `tensors`, each on its own member's device and every one of
16 /// one shape and dtype. The outputs are in member order, each on its member's device. A device
17 /// may appear twice, which runs the collective over transfers rather than natively.
18 ///
19 /// A [`ReduceOperation::Mean`] over integers divides by the member count, truncating toward
20 /// zero.
21 ///
22 /// # Panics
23 ///
24 /// - If `tensors` is empty.
25 /// - If the members don't all have the same shape and dtype.
26 /// - If the members are quantized.
27 ///
28 /// # Example
29 ///
30 /// ```rust
31 /// use burn_tensor::{Tensor, distributed::ReduceOperation};
32 ///
33 /// let device = Default::default();
34 /// let members = vec![
35 /// Tensor::<1>::from_data([1.0, 2.0], &device),
36 /// Tensor::<1>::from_data([3.0, 4.0], &device),
37 /// ];
38 ///
39 /// let reduced = Tensor::all_reduce(members, ReduceOperation::Sum);
40 ///
41 /// assert_eq!(reduced[0].clone().try_into_vec_as::<f32>().unwrap(), vec![4.0, 6.0]);
42 /// assert_eq!(reduced[1].clone().try_into_vec_as::<f32>().unwrap(), vec![4.0, 6.0]);
43 /// ```
44 pub fn all_reduce(tensors: Vec<Self>, op: ReduceOperation) -> Vec<Self> {
45 check!(TensorCheck::collective("All Reduce", &tensors));
46
47 from_primitives(K::all_reduce(into_primitives(tensors), op))
48 }
49
50 /// Reduces every member's tensor with `op`, and gives member `i` chunk `i` along `dim` of the
51 /// result.
52 ///
53 /// Member `i` is position `i` of `tensors`, each on its own member's device and every one of
54 /// one shape and dtype. Chunks are as even as possible, with the first members holding one
55 /// more. The dimension supports negative indexing.
56 ///
57 /// A [`ReduceOperation::Mean`] over integers divides by the member count, truncating toward
58 /// zero.
59 ///
60 /// # Panics
61 ///
62 /// - If `tensors` is empty.
63 /// - If the members don't all have the same shape and dtype.
64 /// - If the members are quantized.
65 /// - If `dim` is outside the rank, or holds fewer elements than there are members.
66 ///
67 /// # Example
68 ///
69 /// ```rust
70 /// use burn_tensor::{Tensor, distributed::ReduceOperation};
71 ///
72 /// let device = Default::default();
73 /// let members = vec![
74 /// Tensor::<1>::from_data([1.0, 2.0], &device),
75 /// Tensor::<1>::from_data([3.0, 4.0], &device),
76 /// ];
77 ///
78 /// let scattered = Tensor::reduce_scatter(members, 0, ReduceOperation::Sum);
79 ///
80 /// assert_eq!(scattered[0].clone().try_into_vec_as::<f32>().unwrap(), vec![4.0]);
81 /// assert_eq!(scattered[1].clone().try_into_vec_as::<f32>().unwrap(), vec![6.0]);
82 /// ```
83 pub fn reduce_scatter(tensors: Vec<Self>, dim: impl AsIndex, op: ReduceOperation) -> Vec<Self> {
84 let dim = unwrap_dim_index(dim.try_dim_index(D), "Reduce Scatter");
85 check!(TensorCheck::collective("Reduce Scatter", &tensors));
86 check!(TensorCheck::collective_chunks(
87 "Reduce Scatter",
88 &tensors,
89 dim
90 ));
91
92 from_primitives(K::reduce_scatter(into_primitives(tensors), dim, op))
93 }
94}
95
96impl<const D: usize, K: Basic> Tensor<D, K> {
97 /// Gives every member every member's tensor, concatenated along `dim` in member order.
98 ///
99 /// Member `i` is position `i` of `tensors`, each on its own member's device and every one of
100 /// one shape and dtype. The outputs are in member order, each on its member's device. The
101 /// dimension supports negative indexing.
102 ///
103 /// # Panics
104 ///
105 /// - If `tensors` is empty.
106 /// - If the members don't all have the same shape and dtype.
107 /// - If the members are quantized.
108 /// - If `dim` is outside the rank.
109 ///
110 /// # Example
111 ///
112 /// ```rust
113 /// use burn_tensor::Tensor;
114 ///
115 /// let device = Default::default();
116 /// let members = vec![
117 /// Tensor::<1>::from_data([1.0, 2.0], &device),
118 /// Tensor::<1>::from_data([3.0, 4.0], &device),
119 /// ];
120 ///
121 /// let gathered = Tensor::all_gather(members, 0);
122 ///
123 /// assert_eq!(gathered[0].clone().try_into_vec_as::<f32>().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);
124 /// assert_eq!(gathered[1].clone().try_into_vec_as::<f32>().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);
125 /// ```
126 pub fn all_gather(tensors: Vec<Self>, dim: impl AsIndex) -> Vec<Self> {
127 let dim = unwrap_dim_index(dim.try_dim_index(D), "All Gather");
128 check!(TensorCheck::collective("All Gather", &tensors));
129
130 from_primitives(K::all_gather(into_primitives(tensors), dim))
131 }
132
133 /// Gives member `i` chunk `i` along `split_dim` of every member's tensor, concatenated along
134 /// `concat_dim` in member order.
135 ///
136 /// Member `i` is position `i` of `tensors`, each on its own member's device and every one of
137 /// one shape and dtype. Chunks are as even as possible, with the first members holding one
138 /// more. Both dimensions support negative indexing.
139 ///
140 /// # Panics
141 ///
142 /// - If `tensors` is empty.
143 /// - If the members don't all have the same shape and dtype.
144 /// - If the members are quantized.
145 /// - If either dimension is outside the rank, or `split_dim` holds fewer elements than there
146 /// are members.
147 ///
148 /// # Example
149 ///
150 /// ```rust
151 /// use burn_tensor::Tensor;
152 ///
153 /// let device = Default::default();
154 /// let members = vec![
155 /// Tensor::<1>::from_data([1.0, 2.0], &device),
156 /// Tensor::<1>::from_data([3.0, 4.0], &device),
157 /// ];
158 ///
159 /// let exchanged = Tensor::all_to_all(members, 0, 0);
160 ///
161 /// assert_eq!(exchanged[0].clone().try_into_vec_as::<f32>().unwrap(), vec![1.0, 3.0]);
162 /// assert_eq!(exchanged[1].clone().try_into_vec_as::<f32>().unwrap(), vec![2.0, 4.0]);
163 /// ```
164 pub fn all_to_all(
165 tensors: Vec<Self>,
166 split_dim: impl AsIndex,
167 concat_dim: impl AsIndex,
168 ) -> Vec<Self> {
169 let split_dim = unwrap_dim_index(split_dim.try_dim_index(D), "All to All");
170 let concat_dim = unwrap_dim_index(concat_dim.try_dim_index(D), "All to All");
171 check!(TensorCheck::collective("All to All", &tensors));
172 check!(TensorCheck::collective_chunks(
173 "All to All",
174 &tensors,
175 split_dim
176 ));
177
178 from_primitives(K::all_to_all(
179 into_primitives(tensors),
180 split_dim,
181 concat_dim,
182 ))
183 }
184
185 /// Gives each of `devices` the tensor, in the order of `devices`.
186 ///
187 /// To repeat a tensor over a shape rather than over devices, see [`Tensor::expand`].
188 ///
189 /// # Panics
190 ///
191 /// - If `devices` is empty.
192 /// - If the tensor is quantized.
193 ///
194 /// # Example
195 ///
196 /// ```rust
197 /// use burn_tensor::{Device, Tensor};
198 ///
199 /// let device = Device::default();
200 /// let tensor = Tensor::<1>::from_data([1.0, 2.0], &device);
201 ///
202 /// let members = Tensor::broadcast(tensor, &[device.clone(), device]);
203 ///
204 /// assert_eq!(members[0].clone().try_into_vec_as::<f32>().unwrap(), vec![1.0, 2.0]);
205 /// assert_eq!(members[1].clone().try_into_vec_as::<f32>().unwrap(), vec![1.0, 2.0]);
206 /// ```
207 pub fn broadcast(tensor: Self, devices: &[Device]) -> Vec<Self> {
208 check!(TensorCheck::collective(
209 "Broadcast",
210 core::slice::from_ref(&tensor)
211 ));
212 check!(TensorCheck::collective_devices("Broadcast", devices));
213
214 from_primitives(K::broadcast(tensor.primitive, devices))
215 }
216}
217
218fn into_primitives<const D: usize, K: Basic>(tensors: Vec<Tensor<D, K>>) -> Vec<BridgeTensor> {
219 tensors.into_iter().map(|tensor| tensor.primitive).collect()
220}
221
222fn from_primitives<const D: usize, K: Basic>(primitives: Vec<BridgeTensor>) -> Vec<Tensor<D, K>> {
223 primitives.into_iter().map(Tensor::new).collect()
224}