Struct ProcessGroup 

pub struct ProcessGroup { /* private fields */ }

A group of processes that can communicate with each other.

Implementations

impl ProcessGroup

pub fn new(backend: Arc<dyn Backend>) -> ProcessGroup

Creates a new process group with all ranks.

pub fn with_ranks(backend: Arc<dyn Backend>, ranks: Vec<usize>) -> ProcessGroup

Creates a process group with specific ranks.

pub fn mock() -> ProcessGroup

Creates a mock process group for testing.

pub fn backend(&self) -> &dyn Backend

Returns the backend.

pub fn rank(&self) -> usize

Returns the rank of this process.

pub fn world_size(&self) -> usize

Returns the world size.

pub fn size(&self) -> usize

Returns the number of processes in this group.

pub fn ranks(&self) -> &[usize]

Returns the ranks in this group.

pub fn contains(&self, rank: usize) -> bool

Checks if this process is part of the group.

pub fn barrier(&self)

Synchronizes all processes in the group.

pub fn all_reduce_tensor(&self, tensor: &mut Tensor<f32>, op: ReduceOp)

Performs all-reduce on a tensor.

pub fn broadcast_tensor(&self, tensor: &mut Tensor<f32>, src: usize)

Broadcasts a tensor from a source rank.

pub fn all_gather_tensor(&self, send_tensor: &Tensor<f32>) -> Tensor<f32>

Performs all-gather on tensors.

pub fn reduce_scatter_tensor( &self, send_tensor: &Tensor<f32>, op: ReduceOp, ) -> Tensor<f32>

Performs reduce-scatter on a tensor.

pub fn send_tensor(&self, tensor: &mut Tensor<f32>, dst: usize)

Sends a tensor to a destination rank.

pub fn recv_tensor(&self, src: usize, shape: &[usize]) -> Tensor<f32>

Receives a tensor from a source rank.

Trait Implementations

impl Clone for ProcessGroup

fn clone(&self) -> ProcessGroup

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.