burn_dispatch/ops/
collective.rs

1use alloc::vec::Vec;
2use burn_backend::{
3    DeviceOps, TensorMetadata,
4    distributed::ReduceOperation,
5    ops::{CollectiveOps, EmulatedCollectives},
6    tensor::{BoolTensor, FloatTensor, IntTensor},
7};
8use burn_backend_extension::backend_dispatch;
9
10use crate::{Dispatch, DispatchDevice, DispatchTensor};
11
12// Members on one backend run its collective; members spanning backends are emulated, since only
13// dispatch's own transfers cross them. A broadcast is always emulated, as its devices may span them.
14impl CollectiveOps<Self> for Dispatch {
15    fn float_all_reduce(
16        tensors: Vec<FloatTensor<Self>>,
17        op: ReduceOperation,
18    ) -> Vec<FloatTensor<Self>> {
19        match Self::on_one_backend(&tensors) {
20            true => Self::float_all_reduce_dispatch(tensors, op),
21            false => EmulatedCollectives::<Self>::float_all_reduce(tensors, op),
22        }
23    }
24
25    fn float_all_gather(tensors: Vec<FloatTensor<Self>>, dim: usize) -> Vec<FloatTensor<Self>> {
26        match Self::on_one_backend(&tensors) {
27            true => Self::float_all_gather_dispatch(tensors, dim),
28            false => EmulatedCollectives::<Self>::float_all_gather(tensors, dim),
29        }
30    }
31
32    fn float_reduce_scatter(
33        tensors: Vec<FloatTensor<Self>>,
34        dim: usize,
35        op: ReduceOperation,
36    ) -> Vec<FloatTensor<Self>> {
37        match Self::on_one_backend(&tensors) {
38            true => Self::float_reduce_scatter_dispatch(tensors, dim, op),
39            false => EmulatedCollectives::<Self>::float_reduce_scatter(tensors, dim, op),
40        }
41    }
42
43    fn float_all_to_all(
44        tensors: Vec<FloatTensor<Self>>,
45        split_dim: usize,
46        concat_dim: usize,
47    ) -> Vec<FloatTensor<Self>> {
48        match Self::on_one_backend(&tensors) {
49            true => Self::float_all_to_all_dispatch(tensors, split_dim, concat_dim),
50            false => EmulatedCollectives::<Self>::float_all_to_all(tensors, split_dim, concat_dim),
51        }
52    }
53
54    fn int_all_reduce(tensors: Vec<IntTensor<Self>>, op: ReduceOperation) -> Vec<IntTensor<Self>> {
55        match Self::on_one_backend(&tensors) {
56            true => Self::int_all_reduce_dispatch(tensors, op),
57            false => EmulatedCollectives::<Self>::int_all_reduce(tensors, op),
58        }
59    }
60
61    fn int_all_gather(tensors: Vec<IntTensor<Self>>, dim: usize) -> Vec<IntTensor<Self>> {
62        match Self::on_one_backend(&tensors) {
63            true => Self::int_all_gather_dispatch(tensors, dim),
64            false => EmulatedCollectives::<Self>::int_all_gather(tensors, dim),
65        }
66    }
67
68    fn int_reduce_scatter(
69        tensors: Vec<IntTensor<Self>>,
70        dim: usize,
71        op: ReduceOperation,
72    ) -> Vec<IntTensor<Self>> {
73        match Self::on_one_backend(&tensors) {
74            true => Self::int_reduce_scatter_dispatch(tensors, dim, op),
75            false => EmulatedCollectives::<Self>::int_reduce_scatter(tensors, dim, op),
76        }
77    }
78
79    fn int_all_to_all(
80        tensors: Vec<IntTensor<Self>>,
81        split_dim: usize,
82        concat_dim: usize,
83    ) -> Vec<IntTensor<Self>> {
84        match Self::on_one_backend(&tensors) {
85            true => Self::int_all_to_all_dispatch(tensors, split_dim, concat_dim),
86            false => EmulatedCollectives::<Self>::int_all_to_all(tensors, split_dim, concat_dim),
87        }
88    }
89
90    fn bool_all_gather(tensors: Vec<BoolTensor<Self>>, dim: usize) -> Vec<BoolTensor<Self>> {
91        match Self::on_one_backend(&tensors) {
92            true => Self::bool_all_gather_dispatch(tensors, dim),
93            false => EmulatedCollectives::<Self>::bool_all_gather(tensors, dim),
94        }
95    }
96
97    fn bool_all_to_all(
98        tensors: Vec<BoolTensor<Self>>,
99        split_dim: usize,
100        concat_dim: usize,
101    ) -> Vec<BoolTensor<Self>> {
102        match Self::on_one_backend(&tensors) {
103            true => Self::bool_all_to_all_dispatch(tensors, split_dim, concat_dim),
104            false => EmulatedCollectives::<Self>::bool_all_to_all(tensors, split_dim, concat_dim),
105        }
106    }
107}
108
109#[backend_dispatch]
110impl Dispatch {
111    #[backend_dispatch(skip)]
112    fn on_one_backend(tensors: &[DispatchTensor]) -> bool {
113        let backend = |tensor: &DispatchTensor| {
114            DispatchDevice::decode_type_id(tensor.device().id().type_id).0
115        };
116        !tensors.is_empty()
117            && tensors
118                .windows(2)
119                .all(|pair| backend(&pair[0]) == backend(&pair[1]))
120    }
121
122    fn float_all_reduce_dispatch(
123        tensors: Vec<FloatTensor<Self>>,
124        op: ReduceOperation,
125    ) -> Vec<FloatTensor<Self>> {
126        B::float_all_reduce(tensors, op)
127    }
128
129    fn float_all_gather_dispatch(
130        tensors: Vec<FloatTensor<Self>>,
131        dim: usize,
132    ) -> Vec<FloatTensor<Self>> {
133        B::float_all_gather(tensors, dim)
134    }
135
136    fn float_reduce_scatter_dispatch(
137        tensors: Vec<FloatTensor<Self>>,
138        dim: usize,
139        op: ReduceOperation,
140    ) -> Vec<FloatTensor<Self>> {
141        B::float_reduce_scatter(tensors, dim, op)
142    }
143
144    fn float_all_to_all_dispatch(
145        tensors: Vec<FloatTensor<Self>>,
146        split_dim: usize,
147        concat_dim: usize,
148    ) -> Vec<FloatTensor<Self>> {
149        B::float_all_to_all(tensors, split_dim, concat_dim)
150    }
151
152    fn int_all_reduce_dispatch(
153        tensors: Vec<IntTensor<Self>>,
154        op: ReduceOperation,
155    ) -> Vec<IntTensor<Self>> {
156        B::int_all_reduce(tensors, op)
157    }
158
159    fn int_all_gather_dispatch(tensors: Vec<IntTensor<Self>>, dim: usize) -> Vec<IntTensor<Self>> {
160        B::int_all_gather(tensors, dim)
161    }
162
163    fn int_reduce_scatter_dispatch(
164        tensors: Vec<IntTensor<Self>>,
165        dim: usize,
166        op: ReduceOperation,
167    ) -> Vec<IntTensor<Self>> {
168        B::int_reduce_scatter(tensors, dim, op)
169    }
170
171    fn int_all_to_all_dispatch(
172        tensors: Vec<IntTensor<Self>>,
173        split_dim: usize,
174        concat_dim: usize,
175    ) -> Vec<IntTensor<Self>> {
176        B::int_all_to_all(tensors, split_dim, concat_dim)
177    }
178
179    fn bool_all_gather_dispatch(
180        tensors: Vec<BoolTensor<Self>>,
181        dim: usize,
182    ) -> Vec<BoolTensor<Self>> {
183        B::bool_all_gather(tensors, dim)
184    }
185
186    fn bool_all_to_all_dispatch(
187        tensors: Vec<BoolTensor<Self>>,
188        split_dim: usize,
189        concat_dim: usize,
190    ) -> Vec<BoolTensor<Self>> {
191        B::bool_all_to_all(tensors, split_dim, concat_dim)
192    }
193}
burn_dispatch/ops/collective.rs

burn_dispatch/ops/
collective.rs
