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

use std::collections::HashMap;

use crate::{
    PeerId,
    local::tensor_map::{CollectiveTensorMap, PeerDeviceMap},
};
use ruda_tensor::Backend;

/// Broadcasts the tensor from one device in a map to all the others
#[cfg_attr(
    feature = "tracing",
    tracing::instrument(level = "trace", skip(devices, tensor))
)]
pub(crate) fn broadcast_centralized<B: Backend>(
    mut devices: PeerDeviceMap<B>,
    central: PeerId,
    tensor: B::FloatTensorPrimitive,
) -> CollectiveTensorMap<B> {
    let mut output = HashMap::new();

    devices
        .remove(&central)
        .expect("Central device id is in `devices`");
    for (dest, dest_device) in devices {
        let tensor = B::float_to_device(tensor.clone(), &dest_device);
        output.insert(dest, tensor);
    }
    output.insert(central, tensor);

    output
}