cubecl-server 0.12.0-pre.1

//! Collectives across devices: the groups this device has joined, and what a
//! driver has to supply to run one.
//!
//! The bookkeeping is the same whichever library is underneath. A group is
//! named by the devices in it, every rank joins under one identifier the group
//! agrees on, and a device's rank is its position in the sorted list — get
//! that wrong on one rank and the whole group hangs. That is what lives here.
//!
//! What does not is [`CollectiveDriver`]: agreeing the identifier, joining,
//! naming an element type, and the operations themselves.

use super::{DeviceResource, Driver};
use crate::server::{CommunicationId, ReduceOperation, ServerError};
use alloc::{format, vec, vec::Vec};
use cubecl_common::device::DeviceId;
use cubecl_environment::backtrace::BackTrace;
use cubecl_environment::collections::HashMap;
use cubecl_ir::ElemType;

/// A driver that can run collectives across devices.
pub trait CollectiveDriver: Driver {
    /// This device's membership of one group, joined once and kept.
    type Communicator;
    /// The identifier every rank of a group joins under.
    type UniqueId: Copy;
    /// How the driver names an element type.
    type DataType: Copy;
    /// The stream collectives are issued on, kept apart from the compute
    /// streams so a collective never blocks one.
    type CommStream: Copy;

    /// The identifier the group `id` names joins under.
    ///
    /// Minted by whichever rank asks first and remembered for the rest, so
    /// this is process-wide state the driver keeps: the servers of two devices
    /// in one group are two objects that have to agree on one answer.
    ///
    /// # Errors
    ///
    /// The driver's refusal to mint one, which stops the group forming at all.
    fn group_id(id: &CommunicationId) -> Result<Self::UniqueId, ServerError>;

    /// Join the group `id` names as rank `rank` of `ranks`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to join, which every other rank sees as this one
    /// never arriving.
    fn join(
        id: Self::UniqueId,
        ranks: usize,
        rank: usize,
    ) -> Result<Self::Communicator, ServerError>;

    /// The elements of `dtype` that `size` bytes hold, as the driver names them.
    ///
    /// # Errors
    ///
    /// An element type the driver has no name for, and a size that is not whole elements.
    /// Reported rather than fatal: a collective is one operation among many, and
    /// refusing it is not a reason to take the process down. The caller can pick
    /// another type, or another way to move the tensor.
    fn elements(dtype: ElemType, size: u64) -> Result<Elements<Self::DataType>, ServerError>;

    /// Reduce `src` across the group into `dst` on every rank.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the reduction.
    fn all_reduce(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        op: ReduceOperation,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;

    /// Write every rank's `src` of `elements`, concatenated in rank order, to `dst`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the gather.
    fn all_gather(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;

    /// Reduce every rank's `src` and write this rank's chunk of `elements` to `dst`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the reduction.
    fn reduce_scatter(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        op: ReduceOperation,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;

    /// Write rank `root`'s `src` of `elements` to `dst` on every rank. Only `root` has a `src`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the broadcast.
    fn broadcast(
        comm: &Self::Communicator,
        src: Option<&DeviceResource<Self>>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        root: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;

    /// Send chunk `r` of `src` to rank `r` and write the chunk rank `r` sends to chunk `r` of
    /// `dst`, each of the `ranks` chunks holding `elements`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the exchange.
    fn all_to_all(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        ranks: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;

    /// Send `src` to `peer`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the send.
    fn send(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        peer: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;

    /// Receive into `dst` from `peer`.
    ///
    /// # Errors
    ///
    /// The driver's refusal to enqueue the receive.
    fn recv(
        comm: &Self::Communicator,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        peer: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
}

/// A buffer's elements, with their type as the driver names it.
#[derive(Debug, Clone, Copy)]
pub struct Elements<T> {
    /// The driver's name for the element type.
    pub data_type: T,
    /// How many elements.
    pub count: usize,
}

/// The groups this device has joined.
///
/// One per server, unlike the identifiers behind
/// [`group_id`](CollectiveDriver::group_id): an identifier is what every rank
/// of a group agrees on, a communicator is one rank's membership of it.
pub struct Collectives<D: CollectiveDriver> {
    /// This device, whose position in a sorted group is its rank.
    device: DeviceId,
    joined: HashMap<CommunicationId, D::Communicator>,
    /// The groups whose part here failed, each with the failure that broke it.
    broken: HashMap<CommunicationId, ServerError>,
}

impl<D: CollectiveDriver> core::fmt::Debug for Collectives<D> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        // The communicators are opaque driver handles, so which groups this
        // device is in is the whole of what there is to say.
        f.debug_struct("Collectives")
            .field("device", &self.device)
            .field("joined", &self.joined.len())
            .field("broken", &self.broken.len())
            .finish()
    }
}

impl<D: CollectiveDriver> Collectives<D> {
    /// A device that has joined nothing yet.
    pub fn new(device: DeviceId) -> Self {
        Self {
            device,
            joined: HashMap::default(),
            broken: HashMap::default(),
        }
    }

    /// Records that this device's part of a collective over `group` failed with `cause`. Its
    /// peers wait in that collective, and another one over the group would pair with theirs.
    pub fn mark_broken(&mut self, group: &CommunicationId, cause: ServerError) {
        self.broken.entry(group.clone()).or_insert(cause);
    }

    /// Refuses a collective over `group` once a part of one here has failed.
    ///
    /// # Errors
    ///
    /// [`ServerError::Generic`] naming the failure that broke the group.
    pub fn ensure_intact(&self, group: &CommunicationId) -> Result<(), ServerError> {
        match self.broken.get(group) {
            None => Ok(()),
            Some(cause) => Err(ServerError::Generic {
                reason: format!(
                    "an earlier collective over this group failed on this device, so its peers \
                     wait in that one and this device runs no other over the group: {cause}"
                ),
                backtrace: BackTrace::capture(),
            }),
        }
    }

    /// Join the group over `devices`, and answer with the id it is known by.
    ///
    /// A group already joined is joined once: `Ok(None)` says there was
    /// nothing to do, so the caller does not announce a membership twice.
    ///
    /// # Errors
    ///
    /// [`ServerError::Generic`] when this device is not among `devices` — it
    /// would have no rank in a group it is not in — and whatever the driver
    /// says about agreeing an identifier or joining.
    pub fn join(&mut self, devices: Vec<DeviceId>) -> Result<Option<CommunicationId>, ServerError> {
        self.join_as(CommunicationId::from(devices.clone()), devices)
    }

    /// The communicator for transfers between this device and `peer`, joined on first use.
    ///
    /// # Errors
    ///
    /// Whatever the driver says about agreeing an identifier or joining.
    pub fn transfers_with(&mut self, peer: DeviceId) -> Result<&D::Communicator, ServerError> {
        let id = CommunicationId::transfers_between(self.device, peer);
        self.join_as(id.clone(), vec![self.device, peer])?;
        self.get(&id)
    }

    fn join_as(
        &mut self,
        id: CommunicationId,
        devices: Vec<DeviceId>,
    ) -> Result<Option<CommunicationId>, ServerError> {
        if self.joined.contains_key(&id) {
            return Ok(None);
        }
        // Sorted, because a device's rank is its position and every rank has
        // to compute the same one. Two ranks disagreeing does not fail: it
        // hangs, with each waiting for a peer that is answering to a different
        // number.
        let mut devices = devices;
        devices.sort();
        // Two ranks on one device would each wait in its join for the other.
        if let Some(pair) = devices.windows(2).find(|pair| pair[0] == pair[1]) {
            return Err(ServerError::Generic {
                reason: format!(
                    "{:?} appears twice among the devices of a group, which need one rank each",
                    pair[0]
                ),
                backtrace: BackTrace::capture(),
            });
        }
        let rank = rank_in(self.device, &devices).ok_or_else(|| ServerError::Generic {
            reason: format!(
                "this device ({:?}) is not among the {} the group was formed over, \
                     so it has no rank in it",
                self.device,
                devices.len()
            ),
            backtrace: BackTrace::capture(),
        })?;

        let comm = D::join(D::group_id(&id)?, devices.len(), rank)?;
        self.joined.insert(id.clone(), comm);
        Ok(Some(id))
    }

    /// This device's membership of the group over `devices`.
    ///
    /// # Errors
    ///
    /// [`ServerError::Generic`] when this device never joined that group, which
    /// is a missing [`join`](Self::join) rather than anything the device did.
    pub fn get(&self, devices: &CommunicationId) -> Result<&D::Communicator, ServerError> {
        self.joined
            .get(devices)
            .ok_or_else(|| ServerError::Generic {
                reason: "no communicator for this group; it has to be joined first".into(),
                backtrace: BackTrace::capture(),
            })
    }

    /// The rank of `device` in the group over `devices`.
    ///
    /// # Errors
    ///
    /// [`ServerError::Generic`] when `device` is not among `devices`.
    pub fn rank(&self, device: DeviceId, devices: &[DeviceId]) -> Result<usize, ServerError> {
        let mut sorted = devices.to_vec();
        sorted.sort();
        rank_in(device, &sorted).ok_or_else(|| ServerError::Generic {
            reason: format!("{device:?} is not among the devices of the group"),
            backtrace: BackTrace::capture(),
        })
    }

    /// The rank of the one device in `devices` that is not this one.
    ///
    /// For the two-device operations — a send has exactly one peer, and its
    /// rank is whichever position this device does not occupy.
    ///
    /// # Errors
    ///
    /// [`ServerError::Generic`] when every device in the pair is this one, so
    /// there is no peer to name.
    pub fn peer_rank(&self, devices: &[DeviceId]) -> Result<usize, ServerError> {
        peer_of(self.device, devices).ok_or_else(|| ServerError::Generic {
            reason: format!(
                "every device in the pair is this one ({:?}), so there is no peer",
                self.device
            ),
            backtrace: BackTrace::capture(),
        })
    }
}

/// `device`'s position in `devices`, which is its rank.
///
/// Free of the driver because a rank is arithmetic over device ids and
/// nothing else — which is also what lets it be checked without one.
fn rank_in(device: DeviceId, devices: &[DeviceId]) -> Option<usize> {
    devices.iter().position(|id| id.index_id == device.index_id)
}

/// The position of the one device in `devices` that is not `device`.
fn peer_of(device: DeviceId, devices: &[DeviceId]) -> Option<usize> {
    devices.iter().position(|id| id.index_id != device.index_id)
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::vec;

    /// A device is its index, since that is all a rank is computed from.
    fn device(index: u16) -> DeviceId {
        DeviceId {
            type_id: 0,
            index_id: index,
        }
    }

    /// A rank is a position in the sorted group, so every member computes the
    /// same one whatever order it was handed the devices in.
    ///
    /// The property worth defending above every other here, because breaking
    /// it does not fail: two ranks that disagree hang, each waiting for a peer
    /// that is answering to a different number.
    #[test]
    fn a_rank_does_not_depend_on_the_order_the_group_was_given_in() {
        let group = [device(7), device(2), device(5)];
        for given in [
            vec![group[0], group[1], group[2]],
            vec![group[2], group[1], group[0]],
            vec![group[1], group[2], group[0]],
        ] {
            let mut sorted = given.clone();
            sorted.sort();
            for (expected, member) in sorted.iter().enumerate() {
                assert_eq!(
                    rank_in(*member, &sorted),
                    Some(expected),
                    "{member:?} in {given:?}"
                );
            }
        }
    }

    /// A device outside the group has no rank in it, rather than silently
    /// taking someone else's.
    #[test]
    fn a_device_outside_the_group_has_no_rank() {
        assert_eq!(rank_in(device(9), &[device(1), device(2)]), None);
    }

    /// The peer of a pair is whichever member is not this device.
    #[test]
    fn the_peer_of_a_pair_is_the_other_member() {
        let pair = [device(3), device(8)];
        assert_eq!(peer_of(device(3), &pair), Some(1));
        assert_eq!(peer_of(device(8), &pair), Some(0));
    }

    /// A pair of one device has no peer, rather than a rank that would send a
    /// transfer to the device it came from.
    #[test]
    fn a_pair_of_one_device_has_no_peer() {
        assert_eq!(peer_of(device(4), &[device(4), device(4)]), None);
    }

    /// Ranks are compared by device index alone, so two devices of different
    /// kinds at the same index are the same member of a group.
    ///
    /// Not obviously right — it is what the collective has always done — but
    /// worth pinning, because a group spanning two device types would silently
    /// give both the same rank.
    #[test]
    fn a_rank_is_the_device_index_and_nothing_else() {
        let other_kind = DeviceId {
            type_id: 1,
            index_id: 2,
        };
        assert_eq!(rank_in(other_kind, &[device(1), device(2)]), Some(1));
    }
}