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::{Primitive, TensorBuffer, TensorDevice, TensorDeviceError, TensorElement, TensorReductionLaunch};
use super::storage::{checked_length, checked_range};
use crate::rank::ReductionOperation;
use ruda_tensor::{Backend, DType, get_device_settings};
use std::sync::Arc;

impl<B: Backend> TensorDevice<B> {
    pub fn reduction_kernel<T: TensorElement>(
        &self,
        operation: ReductionOperation,
    ) -> Result<ReductionOperation, TensorDeviceError> {
        self.validate_type::<T>()?;
        if T::dtype() != DType::I32 && matches!(operation,
            ReductionOperation::BitAnd | ReductionOperation::BitOr | ReductionOperation::BitXor)
        {
            return Err(TensorDeviceError::InvalidOperation("bitwise reductions require I32 elements"));
        }
        Ok(operation)
    }

    pub fn prepare<T: TensorElement>(
        &self,
        length: usize,
        destination_offset: usize,
    ) -> Result<TensorReductionLaunch, TensorDeviceError> {
        let end = destination_offset.checked_add(length)
            .ok_or(TensorDeviceError::InvalidBuffer("collective reduction range overflow"))?;
        checked_length::<T>(end)?;
        Ok(TensorReductionLaunch { length, destination_offset })
    }

    pub fn reduce<T: TensorElement>(
        &self,
        operation: ReductionOperation,
        launch: &TensorReductionLaunch,
        source: &TensorBuffer<B, T>,
        destination: &TensorBuffer<B, T>,
    ) -> Result<(), TensorDeviceError> {
        self.validate_buffer(source)?;
        self.validate_buffer(destination)?;
        self.reduction_kernel::<T>(operation)?;
        let source_range = checked_range(source.length, 0, launch.length)?;
        let target_range = checked_range(destination.length, launch.destination_offset, launch.length)?;
        if launch.length == 0 {
            return Ok(());
        }
        let apply = |current: Primitive<B>, incoming: Primitive<B>| {
            if target_range.start == 0 && target_range.end == destination.length {
                self.reduce_values(operation, current, incoming)
            } else {
                let target = current.clone().slice(target_range.clone());
                let reduced = self.reduce_values(operation, target, incoming);
                current.assign(target_range.clone(), reduced)
            }
        };
        if Arc::ptr_eq(&source.value, &destination.value) {
            destination.update(|current| {
                let incoming = current.clone().slice(source_range);
                apply(current, incoming)
            })
        } else {
            // Drop the source lock before acquiring the destination lock.
            let mut incoming = source.snapshot()?;
            if launch.length != source.length {
                incoming = incoming.slice(source_range);
            }
            destination.update(|current| apply(current, incoming))
        }
    }

    fn reduce_values(&self, operation: ReductionOperation, destination: Primitive<B>, source: Primitive<B>) -> Primitive<B> {
        use ReductionOperation::*;
        let bool_dtype = get_device_settings::<B>(&self.device).bool_dtype;
        match (destination, source) {
            (Primitive::Float(destination), Primitive::Float(source)) => Primitive::Float(match operation {
                Sum => B::float_add(destination, source),
                Product => B::float_mul(destination, source),
                // Preserve the destination for unordered/equal comparisons,
                // including NaNs and opposite signed zero, like native kernels.
                Minimum | Maximum => {
                    let mask = if operation == Minimum {
                        B::float_lower(source.clone(), destination.clone(), bool_dtype)
                    } else {
                        B::float_lower(destination.clone(), source.clone(), bool_dtype)
                    };
                    B::float_mask_where(destination, mask, source)
                }
                BitAnd | BitOr | BitXor => unreachable!("validated floating reduction"),
            }),
            (Primitive::Int(destination), Primitive::Int(source)) => Primitive::Int(match operation {
                Sum => B::int_add(destination, source),
                Product => B::int_mul(destination, source),
                Minimum | Maximum => {
                    let mask = if operation == Minimum {
                        B::int_lower(source.clone(), destination.clone(), bool_dtype)
                    } else {
                        B::int_lower(destination.clone(), source.clone(), bool_dtype)
                    };
                    B::int_mask_where(destination, mask, source)
                }
                BitAnd => B::bitwise_and(destination, source),
                BitOr => B::bitwise_or(destination, source),
                BitXor => B::bitwise_xor(destination, source),
            }),
            _ => unreachable!("typed collective buffer storage kind"),
        }
    }
}