burn_dispatch/ops/
distributed.rs

1use alloc::vec::Vec;
2
3use burn_backend::{
4    distributed::{DistributedConfig, DistributedOps, DistributedParams, TensorRef},
5    tensor::FloatTensor,
6};
7
8use crate::{Dispatch, DispatchDevice};
9
10macro_rules! dispatch_distributed_devices_arms {
11    (
12        $device:expr,
13        $devices:expr,
14        |$inner_devices:ident| $body:expr;
15        $([$Backend:ident, $cfg:meta]),*
16    ) => {
17        match $device {
18            // Autodiff arm first
19            #[cfg(feature = "autodiff")]
20            $crate::DispatchDevice::Autodiff(inner) => {
21                let inner_devices = $devices
22                    .iter()
23                    .map(|d| {
24                         match &d {
25                            #[cfg(feature = "autodiff")]
26                            $crate::DispatchDevice::Autodiff(d_inner) => *d_inner.inner.clone(),
27                            _ => unreachable!("All devices are expected to be of the same variant."),
28                        }
29                    })
30                    .collect::<Vec<_>>();
31                let inner_devices = inner_devices.as_slice();
32                // Recursively dispatch on inner
33                dispatch_distributed_devices_arms!(
34                    @autodiff
35                    &**inner,
36                    &*inner_devices,
37                    |$inner_devices| $body;
38                    $([$Backend, $cfg]),*
39                )
40            },
41            $(
42                #[cfg($cfg)]
43                $crate::DispatchDevice::$Backend(_) => {
44                    type B = $crate::backends::$Backend;
45                    let $inner_devices = $devices
46                        .iter()
47                        .map(|d| {
48                            // In a build with only one dispatch variant compiled in, the
49                            // pattern is irrefutable and the else clause is dead — which is
50                            // the point: it only guards the multi-backend builds.
51                            #[allow(irrefutable_let_patterns)]
52                            let DispatchDevice::$Backend(dev) = d else {
53                                unreachable!("All devices are expected to be of the same variant.")
54                            };
55                            dev.clone()
56                        })
57                        .collect::<Vec<_>>();
58                    $body
59                }
60            )*
61            #[allow(unreachable_patterns)]
62            other => panic!("Distributed operations are not supported for device {other:?}"),
63        }
64    };
65    (
66        @autodiff
67        $device:expr,
68        $devices:expr,
69        |$inner_devices:ident| $body:expr;
70        $([$Backend:ident, $cfg:meta]),*
71    ) => {
72        match $device {
73            $(
74                #[cfg($cfg)]
75                $crate::DispatchDevice::$Backend(_) => {
76                    type B = $crate::backends::Autodiff<$crate::backends::$Backend>;
77                    let $inner_devices = $devices
78                        .iter()
79                        .map(|d| {
80                            // In a build with only one dispatch variant compiled in, the
81                            // pattern is irrefutable and the else clause is dead — which is
82                            // the point: it only guards the multi-backend builds.
83                            #[allow(irrefutable_let_patterns)]
84                            let DispatchDevice::$Backend(dev) = d else {
85                                unreachable!("All devices are expected to be of the same variant.")
86                            };
87                            dev.clone()
88                        })
89                        .collect::<Vec<_>>();
90                    $body
91                }
92            )*
93            $crate::DispatchDevice::Autodiff(_) => panic!("Autodiff should not wrap an autodiff device."),
94            #[allow(unreachable_patterns)]
95            other => panic!("Distributed operations are not supported for device {other:?}"),
96        }
97    };
98}
99
100/// Dispatches an operation body based on the provided devices.
101macro_rules! dispatch_distributed_devices {
102    ($device:expr, $devices:expr, |$inner_devices:ident| $body:expr) => {
103        distributed_backend_list!(
104            dispatch_distributed_devices_arms,
105            $device,
106            $devices,
107            |$inner_devices| $body
108        )
109    };
110}
111
112// In builds without a backend the catalog marks distributed, the dispatch arms are all cfg'd out,
113// leaving only a diverging fallback, so the captured arguments go unused.
114#[allow(unused_variables, unreachable_code)]
115impl DistributedOps<Self> for Dispatch {
116    fn start_communication_server(devices: &[DispatchDevice], config: DistributedConfig) {
117        if !devices.is_empty() {
118            let first = &devices[0];
119            dispatch_distributed_devices!(first, devices, |inner_devices| {
120                B::start_communication_server(&inner_devices, config)
121            });
122        }
123    }
124
125    fn close_communication_server(device: &DispatchDevice) {
126        dispatch_device!(@distributed device, |device| {
127            B::close_communication_server(device)
128        })
129    }
130
131    fn register_sync_parameters(
132        device: &DispatchDevice,
133        sharded_param_ids: Vec<DistributedParams>,
134    ) {
135        dispatch_device!(@distributed device, |device| B::register_sync_parameters(
136            device,
137            sharded_param_ids,
138        ))
139    }
140
141    fn submit_sync_collective(device: &DispatchDevice) {
142        dispatch_device!(@distributed device, |device| B::submit_sync_collective(device))
143    }
144
145    fn submit_gradient_sync(_tensor: TensorRef<Self>, _distributed_params: DistributedParams) {
146        unimplemented!()
147    }
148
149    unsafe fn comm_device(_tensor: &TensorRef<Self>) -> DispatchDevice {
150        unimplemented!()
151    }
152
153    unsafe fn float_from_ref(_tensor: &TensorRef<Self>) -> FloatTensor<Self> {
154        unimplemented!()
155    }
156}
burn_dispatch/ops/distributed.rs

burn_dispatch/ops/
distributed.rs
