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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use super::*;

impl<T, D> DeviceCollective<'_, T, D>
where
    D: RankDevice<T>,
    D::Error: From<RankError> + From<NetworkError> + From<TopologyError> + From<WorkError>,
{
    pub fn reduce_concatenated(
        &self,
        payload: &[u8],
        output: &D::Buffer,
        length: usize,
        function: &D::Kernel,
    ) -> Result<u32, D::Error> {
        let rank_bytes = length
            .checked_mul(D::ELEMENT_SIZE)
            .ok_or(RankError::Overflow("TCP reduction rank bytes"))?;
        let expected = rank_bytes
            .checked_mul(self.world_size() as usize)
            .ok_or(RankError::Overflow("TCP reduction response bytes"))?;
        if payload.len() != expected {
            return Err(NetworkError::InvalidConfiguration(format!(
                "TCP reduction response has {} bytes, expected {expected}",
                payload.len()
            ))
            .into());
        }
        let first = D::decode(&payload[..rank_bytes])?;
        RankDevice::<T>::copy_to_device(self.execution, output, &first)?;
        if length == 0 {
            return Ok(0);
        }
        if self.world_size() == 1 {
            return Ok(0);
        }
        let scratch = RankDevice::<T>::alloc(self.execution, length)?;
        let launch = RankDevice::<T>::prepare_reduction(self.execution, length, 0)?;
        for source in 1..self.world_size() as usize {
            let start = source * rank_bytes;
            let values = D::decode(&payload[start..start + rank_bytes])?;
            RankDevice::<T>::copy_to_device(self.execution, &scratch, &values)?;
            RankDevice::<T>::launch_reduction(self.execution, function, &launch, &scratch, output)?;
        }
        Ok(self.world_size() - 1)
    }

    pub fn reduce_values_into(
        &self,
        destination: &D::Buffer,
        destination_offset: usize,
        values: &[T],
        function: &D::Kernel,
    ) -> Result<u32, D::Error> {
        let bytes = D::encode(values);
        self.reduce_bytes_into(
            destination,
            destination_offset,
            values.len(),
            &bytes,
            function,
        )
    }

    pub fn reduce_bytes_into(
        &self,
        destination: &D::Buffer,
        destination_offset: usize,
        element_count: usize,
        values: &[u8],
        function: &D::Kernel,
    ) -> Result<u32, D::Error> {
        let expected_bytes = element_count
            .checked_mul(D::ELEMENT_SIZE)
            .ok_or(RankError::Overflow("reduction input byte length"))?;
        if values.len() != expected_bytes {
            return Err(RankError::InvalidLength(
                "reduction input does not match its element count",
            )
            .into());
        }
        if element_count == 0 {
            return Ok(0);
        }
        let end = destination_offset
            .checked_add(element_count)
            .ok_or(RankError::Overflow(
                "hierarchical reduction destination range",
            ))?;
        if end > self.execution.buffer_len(destination) {
            return Err(RankError::InvalidLength(
                "hierarchical reduction destination range is out of bounds",
            )
            .into());
        }
        let scratch = RankDevice::<T>::alloc(self.execution, element_count)?;
        RankDevice::<T>::copy_bytes_to_device_at(self.execution, &scratch, 0, values)?;
        let launch =
            RankDevice::<T>::prepare_reduction(self.execution, element_count, destination_offset)?;
        RankDevice::<T>::launch_reduction(
            self.execution,
            function,
            &launch,
            &scratch,
            destination,
        )?;
        Ok(1)
    }
}