burn-dispatch 0.22.1

use alloc::vec::Vec;

use burn_backend::{
    DeviceId,
    distributed::{
        CollectiveTensor, DistributedConfig, DistributedOps, DistributedParams, ReduceOperation,
        TensorRef,
    },
    tensor::FloatTensor,
};

use crate::{Dispatch, DispatchDevice};

macro_rules! dispatch_distributed_devices_arms {
    (
        $device:expr,
        $devices:expr,
        |$inner_devices:ident| $body:expr;
        $([$Backend:ident, $cfg:meta]),*
    ) => {
        match $device {
            // Autodiff arm first
            #[cfg(feature = "autodiff")]
            $crate::DispatchDevice::Autodiff(inner) => {
                let inner_devices = $devices
                    .iter()
                    .map(|d| {
                         match &d {
                            #[cfg(feature = "autodiff")]
                            $crate::DispatchDevice::Autodiff(d_inner) => *d_inner.inner.clone(),
                            _ => unreachable!("All devices are expected to be of the same variant."),
                        }
                    })
                    .collect::<Vec<_>>();
                let inner_devices = inner_devices.as_slice();
                // Recursively dispatch on inner
                dispatch_distributed_devices_arms!(
                    @autodiff
                    &**inner,
                    &*inner_devices,
                    |$inner_devices| $body;
                    $([$Backend, $cfg]),*
                )
            },
            $(
                #[cfg($cfg)]
                $crate::DispatchDevice::$Backend(_) => {
                    type B = $crate::backends::$Backend;
                    let $inner_devices = $devices
                        .iter()
                        .map(|d| {
                            // In a build with only one dispatch variant compiled in, the
                            // pattern is irrefutable and the else clause is dead — which is
                            // the point: it only guards the multi-backend builds.
                            #[allow(irrefutable_let_patterns)]
                            let DispatchDevice::$Backend(dev) = d else {
                                unreachable!("All devices are expected to be of the same variant.")
                            };
                            dev.clone()
                        })
                        .collect::<Vec<_>>();
                    $body
                }
            )*
            #[allow(unreachable_patterns)]
            other => panic!("Distributed operations are not supported for device {other:?}"),
        }
    };
    (
        @autodiff
        $device:expr,
        $devices:expr,
        |$inner_devices:ident| $body:expr;
        $([$Backend:ident, $cfg:meta]),*
    ) => {
        match $device {
            $(
                #[cfg($cfg)]
                $crate::DispatchDevice::$Backend(_) => {
                    type B = $crate::backends::Autodiff<$crate::backends::$Backend>;
                    let $inner_devices = $devices
                        .iter()
                        .map(|d| {
                            // In a build with only one dispatch variant compiled in, the
                            // pattern is irrefutable and the else clause is dead — which is
                            // the point: it only guards the multi-backend builds.
                            #[allow(irrefutable_let_patterns)]
                            let DispatchDevice::$Backend(dev) = d else {
                                unreachable!("All devices are expected to be of the same variant.")
                            };
                            dev.clone()
                        })
                        .collect::<Vec<_>>();
                    $body
                }
            )*
            $crate::DispatchDevice::Autodiff(_) => panic!("Autodiff should not wrap an autodiff device."),
            #[allow(unreachable_patterns)]
            other => panic!("Distributed operations are not supported for device {other:?}"),
        }
    };
}

/// Dispatches an operation body based on the provided devices.
macro_rules! dispatch_distributed_devices {
    ($device:expr, $devices:expr, |$inner_devices:ident| $body:expr) => {
        distributed_backend_list!(
            dispatch_distributed_devices_arms,
            $device,
            $devices,
            |$inner_devices| $body
        )
    };
}

macro_rules! dispatch_distributed_float_arms {
    ($tensor:expr, |$inner:ident| $body:expr; $([$Backend:ident, $cfg:meta]),*) => {{
        let autodiff = $tensor.autodiff;
        match $tensor.kind {
            #[cfg(feature = "autodiff")]
            $crate::DispatchTensorKind::Autodiff(inner) => match *inner {
                $(
                    #[cfg($cfg)]
                    $crate::DispatchTensorKind::$Backend($inner) => {
                        let $crate::DispatchAutodiffContext::Enabled(checkpointing) = autodiff else {
                            panic!("an autodiff float primitive must have an enabled autodiff context")
                        };
                        with_autodiff_backend!($Backend, checkpointing, |B| {
                            let $inner = $inner.autodiff();
                            $crate::DispatchTensor {
                                kind: $crate::DispatchTensorKind::Autodiff(alloc::boxed::Box::new(
                                    $crate::DispatchTensorKind::$Backend(
                                        $crate::BackendTensor::Autodiff($body),
                                    ),
                                )),
                                autodiff,
                            }
                        })
                    }
                )*
                #[allow(unreachable_patterns)]
                other => panic!("Distributed operations are not supported for tensor kind {other:?}"),
            },
            $(
                #[cfg($cfg)]
                $crate::DispatchTensorKind::$Backend($inner) => {
                    assert_eq!(
                        autodiff,
                        $crate::DispatchAutodiffContext::Disabled,
                        "an enabled float tensor must use an autodiff primitive",
                    );
                    type B = $crate::backends::$Backend;
                    let $inner = $inner.float();
                    $crate::DispatchTensor {
                        kind: $crate::DispatchTensorKind::$Backend(
                            $crate::BackendTensor::Float($body),
                        ),
                        autodiff,
                    }
                }
            )*
            #[allow(unreachable_patterns)]
            other => panic!("Distributed operations are not supported for tensor kind {other:?}"),
        }
    }};
}

macro_rules! dispatch_distributed_float {
    ($tensor:expr, |$inner:ident| $body:expr) => {
        distributed_backend_list!(dispatch_distributed_float_arms, $tensor, |$inner| $body)
    };
}

// In builds without a collective-capable backend (Cuda/Remote), the distributed dispatch arms
// are all cfg'd out, leaving only a diverging fallback — so the captured arguments and trailing
// expressions are intentionally unused/unreachable.
#[allow(unused_variables, unreachable_code)]
impl DistributedOps<Self> for Dispatch {
    fn start_communication_server(devices: &[DispatchDevice], config: DistributedConfig) {
        if !devices.is_empty() {
            let first = &devices[0];
            dispatch_distributed_devices!(first, devices, |inner_devices| {
                B::start_communication_server(&inner_devices, config)
            });
        }
    }

    fn close_communication_server(device: &DispatchDevice) {
        dispatch_device!(@distributed device, |device| {
            B::close_communication_server(device)
        })
    }

    fn register_sync_parameters(
        device: &DispatchDevice,
        sharded_param_ids: Vec<DistributedParams>,
    ) {
        dispatch_device!(@distributed device, |device| B::register_sync_parameters(
            device,
            sharded_param_ids,
        ))
    }

    fn submit_sync_collective(device: &DispatchDevice) {
        dispatch_device!(@distributed device, |device| B::submit_sync_collective(device))
    }

    fn submit_gradient_sync(_tensor: TensorRef<Self>, _distributed_params: DistributedParams) {
        unimplemented!()
    }

    fn all_reduce(
        tensor: FloatTensor<Self>,
        op: ReduceOperation,
        device_ids: Vec<DeviceId>,
    ) -> CollectiveTensor<Self> {
        // Safety: we call `assume_resolved` only to wrap it in a new `CollectiveTensor`.
        // Explicit type: the distributed dispatch only emits arms for collective-capable
        // backends (Cuda, Remote), so a build with none of them leaves only the diverging
        // fallback and the match would otherwise infer `!`.
        let tensor: FloatTensor<Self> = dispatch_distributed_float!(tensor, |tensor| {
            let collective_tensor = B::all_reduce(tensor, op, device_ids);
            unsafe { collective_tensor.assume_resolved() }
        });
        CollectiveTensor::new(tensor)
    }

    fn sync_collective(device: &DispatchDevice) {
        dispatch_device!(@distributed device, |device| B::sync_collective(device))
    }

    unsafe fn comm_device(_tensor: &TensorRef<Self>) -> DispatchDevice {
        unimplemented!()
    }

    unsafe fn float_from_ref(_tensor: &TensorRef<Self>) -> FloatTensor<Self> {
        unimplemented!()
    }
}