burn_cubecl/ops/
distributed.rs

1use burn_backend::distributed::DistributedOps;
2
3use crate::{CubeBackend, CubeRuntime};
4
5#[cfg(feature = "std")]
6use crate::ops::numeric::{self, zeros_client};
7#[cfg(feature = "std")]
8use burn_backend::{
9    DeviceId, TensorMetadata,
10    cubecl::dtype_to_elem_type,
11    distributed::{CollectiveTensor, ReduceOperation},
12    tensor::{Device, FloatTensor},
13};
14
15impl<R: CubeRuntime> DistributedOps<Self> for CubeBackend<R> {
16    #[cfg(feature = "std")]
17    fn all_reduce(
18        tensor: FloatTensor<Self>,
19        op: ReduceOperation,
20        device_ids: Vec<DeviceId>,
21    ) -> CollectiveTensor<Self> {
22        let device = &tensor.device.clone();
23        let out_tensor = if tensor.handle.can_mut() && tensor.is_contiguous() {
24            tensor
25        } else {
26            let zeros_tensor = zeros_client::<R>(
27                tensor.client.clone(),
28                device.clone(),
29                tensor.shape(),
30                tensor.dtype(),
31            );
32            numeric::add(zeros_tensor, tensor)
33        };
34
35        let op = match op {
36            ReduceOperation::Sum => cubecl::server::ReduceOperation::Sum,
37            ReduceOperation::Mean => cubecl::server::ReduceOperation::Mean,
38        };
39
40        let mut client = R::client(device);
41        client.all_reduce(
42            out_tensor.handle.clone(),
43            out_tensor.handle.clone(),
44            dtype_to_elem_type(out_tensor.dtype),
45            device_ids.clone(),
46            op,
47        );
48        CollectiveTensor::new(out_tensor)
49    }
50
51    #[cfg(feature = "std")]
52    fn sync_collective(device: &Device<Self>) {
53        let client = R::client(device);
54        client.sync_collective();
55    }
56}
burn_cubecl/ops/distributed.rs

burn_cubecl/ops/
distributed.rs
