Struct InProcessCollective 

pub struct InProcessCollective<'a, T, D: InProcessDevice<T>, E> { /* private fields */ }

Implementations

impl<T, D, E> InProcessCollective<'_, T, D, E>

where T: Copy + Send + Sync + 'static, D: InProcessDevice<T>, E: From<D::Error> + From<InProcessError>,

pub fn all_to_all( &self, input: &DistributedBuffer<T, D>, ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E>

Equal-count all-to-all. Each input rank stores destination shards in rank order; each output rank stores source shards in rank order.

pub fn all_to_all_v( &self, input: &DistributedBuffer<T, D>, send_counts: &[Vec<usize>], ) -> Result<(VariableDistributedBuffer<T, D>, CollectiveStats), E>

impl<T, D, E> InProcessCollective<'_, T, D, E>

where T: Copy + Send + Sync + 'static, D: InProcessDevice<T>, E: From<D::Error> + From<InProcessError>,

pub fn broadcast( &self, buffer: &DistributedBuffer<T, D>, root: usize, ) -> Result<CollectiveStats, E>

pub fn all_gather( &self, input: &DistributedBuffer<T, D>, ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E>

pub fn gather( &self, input: &DistributedBuffer<T, D>, root: usize, ) -> Result<(RootedBuffer<T, D>, CollectiveStats), E>

pub fn scatter( &self, input: &RootedBuffer<T, D>, ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E>

pub fn send_recv( &self, source: &DistributedBuffer<T, D>, destination: &DistributedBuffer<T, D>, transfer: PointToPointTransfer, ) -> Result<CollectiveStats, E>

Tagged point-to-point transfer between two rank-local buffers. This is synchronous on the host-staged transport; tag is retained in the API because network transports match send and receive by that value.

impl<B, T, E> InProcessCollective<'_, T, TensorDevice<B>, E>

where B: Backend, T: TensorElement, E: From<TensorDeviceError> + From<InProcessError>,

pub fn import_buffers( &self, buffers: Vec<TensorBuffer<B, T>>, ) -> Result<DistributedBuffer<T, TensorDevice<B>>, E>

Attach existing rank tensors without downloading and re-uploading values.

impl<T, D, E> InProcessCollective<'_, T, D, E>

where T: Copy + Send + Sync + 'static, D: InProcessDevice<T>, E: From<D::Error> + From<InProcessError>,

pub fn allocate( &self, length_per_rank: usize, ) -> Result<DistributedBuffer<T, D>, E>

pub fn allocate_rooted( &self, root: usize, length: usize, ) -> Result<RootedBuffer<T, D>, E>

pub fn allocate_variable( &self, lengths: &[usize], ) -> Result<VariableDistributedBuffer<T, D>, E>

pub fn upload_variable_rank( &self, buffer: &VariableDistributedBuffer<T, D>, rank: usize, values: &[T], ) -> Result<(), E>

pub fn download_variable_rank( &self, buffer: &VariableDistributedBuffer<T, D>, rank: usize, ) -> Result<Vec<T>, E>

pub fn upload_root( &self, buffer: &RootedBuffer<T, D>, values: &[T], ) -> Result<(), E>

pub fn download_root(&self, buffer: &RootedBuffer<T, D>) -> Result<Vec<T>, E>

pub fn upload_rank( &self, buffer: &DistributedBuffer<T, D>, rank: usize, values: &[T], ) -> Result<(), E>

pub fn download_rank( &self, buffer: &DistributedBuffer<T, D>, rank: usize, ) -> Result<Vec<T>, E>

impl<T, D, E> InProcessCollective<'_, T, D, E>

where T: Copy + Send + Sync + 'static, D: InProcessDevice<T>, E: From<D::Error> + From<InProcessError>,

pub fn reduce_sum( &self, input: &DistributedBuffer<T, D>, root: usize, function: &D::Kernel, ) -> Result<(RootedBuffer<T, D>, CollectiveStats), E>

pub fn all_reduce_sum( &self, buffer: &DistributedBuffer<T, D>, function: &D::Kernel, ) -> Result<CollectiveStats, E>

pub fn reduce_scatter_sum( &self, input: &DistributedBuffer<T, D>, function: &D::Kernel, ) -> Result<(DistributedBuffer<T, D>, CollectiveStats), E>

impl<T, D, E> InProcessCollective<'_, T, D, E>

where T: Copy + Send + Sync + 'static, D: InProcessDevice<T>, E: From<D::Error> + From<InProcessError>,

pub fn world_size(&self) -> usize

pub const fn transport(&self) -> CollectiveTransport