Skip to main content

CollectiveDriver

cubecl_server::command

Trait CollectiveDriver 

Source 
pub trait CollectiveDriver: Driver {
    type Communicator;
    type UniqueId: Copy;
    type DataType: Copy;
    type CommStream: Copy;

    // Required methods
    fn group_id(id: &CommunicationId) -> Result<Self::UniqueId, ServerError>;
    fn join(
        id: Self::UniqueId,
        ranks: usize,
        rank: usize,
    ) -> Result<Self::Communicator, ServerError>;
    fn data_type(
        dtype: ElemType,
        size: u64,
    ) -> Result<(Self::DataType, usize), ServerError>;
    fn all_reduce(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        dtype: Self::DataType,
        count: usize,
        op: ReduceOperation,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn send(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dtype: Self::DataType,
        count: usize,
        peer: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn recv(
        comm: &Self::Communicator,
        dst: &DeviceResource<Self>,
        dtype: Self::DataType,
        count: usize,
        peer: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
}
Expand description

A driver that can run collectives across devices.

Required Associated Types§

Source

type Communicator

This device’s membership of one group, joined once and kept.

Source

type UniqueId: Copy

The identifier every rank of a group joins under.

Source

type DataType: Copy

How the driver names an element type.

Source

type CommStream: Copy

The stream collectives are issued on, kept apart from the compute streams so a collective never blocks one.

Required Methods§

Source

fn group_id(id: &CommunicationId) -> Result<Self::UniqueId, ServerError>

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.

Source

fn join( id: Self::UniqueId, ranks: usize, rank: usize, ) -> Result<Self::Communicator, 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.

Source

fn data_type( dtype: ElemType, size: u64, ) -> Result<(Self::DataType, usize), ServerError>

How the driver names dtype, and how many elements size bytes hold.

§Errors

An element type the driver has no name for. 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.

Source

fn all_reduce( comm: &Self::Communicator, src: &DeviceResource<Self>, dst: &DeviceResource<Self>, dtype: Self::DataType, count: usize, op: ReduceOperation, stream: Self::CommStream, ) -> Result<(), ServerError>

Reduce src across the group into dst on every rank.

§Errors

The driver’s refusal to enqueue the reduction.

Source

fn send( comm: &Self::Communicator, src: &DeviceResource<Self>, dtype: Self::DataType, count: usize, peer: usize, stream: Self::CommStream, ) -> Result<(), ServerError>

Send src to peer.

§Errors

The driver’s refusal to enqueue the send.

Source

fn recv( comm: &Self::Communicator, dst: &DeviceResource<Self>, dtype: Self::DataType, count: usize, peer: usize, stream: Self::CommStream, ) -> Result<(), ServerError>

Receive into dst from peer.

§Errors

The driver’s refusal to enqueue the receive.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§