Trait CollectiveDriver 

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 elements(
        dtype: ElemType,
        size: u64,
    ) -> Result<Elements<Self::DataType>, ServerError>;
    fn all_reduce(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        op: ReduceOperation,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn all_gather(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn reduce_scatter(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        op: ReduceOperation,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn broadcast(
        comm: &Self::Communicator,
        src: Option<&DeviceResource<Self>>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        root: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn all_to_all(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        ranks: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn send(
        comm: &Self::Communicator,
        src: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        peer: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
    fn recv(
        comm: &Self::Communicator,
        dst: &DeviceResource<Self>,
        elements: Elements<Self::DataType>,
        peer: usize,
        stream: Self::CommStream,
    ) -> Result<(), ServerError>;
}

A driver that can run collectives across devices.

Required Associated Types

type Communicator

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

type UniqueId: Copy

The identifier every rank of a group joins under.

type DataType: Copy

How the driver names an element type.

type CommStream: Copy

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

Required Methods

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.

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.

fn elements( dtype: ElemType, size: u64, ) -> Result<Elements<Self::DataType>, 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 all_reduce( comm: &Self::Communicator, src: &DeviceResource<Self>, dst: &DeviceResource<Self>, elements: Elements<Self::DataType>, 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.

fn all_gather( comm: &Self::Communicator, src: &DeviceResource<Self>, dst: &DeviceResource<Self>, elements: Elements<Self::DataType>, 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 reduce_scatter( comm: &Self::Communicator, src: &DeviceResource<Self>, dst: &DeviceResource<Self>, elements: Elements<Self::DataType>, op: ReduceOperation, 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 broadcast( comm: &Self::Communicator, src: Option<&DeviceResource<Self>>, dst: &DeviceResource<Self>, elements: Elements<Self::DataType>, root: usize, 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 all_to_all( comm: &Self::Communicator, src: &DeviceResource<Self>, dst: &DeviceResource<Self>, elements: Elements<Self::DataType>, ranks: 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 send( comm: &Self::Communicator, src: &DeviceResource<Self>, elements: Elements<Self::DataType>, peer: usize, stream: Self::CommStream, ) -> Result<(), ServerError>

Send src to peer.

Errors

The driver’s refusal to enqueue the send.

fn recv( comm: &Self::Communicator, dst: &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.

Dyn Compatibility

This trait is not dyn compatible.

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

Implementors