ruCCL 0.21.1

Ruda collective communication algorithms and orchestration.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

//! # Common Tensor Map for Local Collective Operations
use ruda_tensor::{
    Backend, TensorMetadata,
    tensor::{Device, FloatTensor},
};
use ruda_core::tensor::Shape;
use std::collections::HashMap;

use crate::PeerId;

pub type CollectiveTensorMap<B> = HashMap<PeerId, FloatTensor<B>>;

pub type PeerDeviceMap<B> = HashMap<PeerId, Device<B>>;

/// Get the shape of the tensors. They should all have the same shape, otherwise None is returned.
pub fn get_common_shape<B: Backend>(tensors: &CollectiveTensorMap<B>) -> Option<Shape> {
    let mut it = tensors.values();
    if let Some(first) = it.next() {
        let shape = first.shape();
        for tensor in it {
            if tensor.shape() != shape {
                return None;
            }
        }
        return Some(shape);
    }
    None
}

/// Get the `{ peer_id -> device }` mapping for the given tensors.
pub fn get_peer_devices<B: Backend>(tensors: &CollectiveTensorMap<B>) -> PeerDeviceMap<B> {
    tensors
        .iter()
        .map(|(id, tensor)| (*id, B::float_device(tensor)))
        .collect()
}