ruccl/tensor_device/
reduction.rs

1use super::{Primitive, TensorBuffer, TensorDevice, TensorDeviceError, TensorElement, TensorReductionLaunch};
2use super::storage::{checked_length, checked_range};
3use crate::rank::ReductionOperation;
4use ruda_tensor::{Backend, DType, get_device_settings};
5use std::sync::Arc;
6
7impl<B: Backend> TensorDevice<B> {
8    pub fn reduction_kernel<T: TensorElement>(
9        &self,
10        operation: ReductionOperation,
11    ) -> Result<ReductionOperation, TensorDeviceError> {
12        self.validate_type::<T>()?;
13        if T::dtype() != DType::I32 && matches!(operation,
14            ReductionOperation::BitAnd | ReductionOperation::BitOr | ReductionOperation::BitXor)
15        {
16            return Err(TensorDeviceError::InvalidOperation("bitwise reductions require I32 elements"));
17        }
18        Ok(operation)
19    }
20
21    pub fn prepare<T: TensorElement>(
22        &self,
23        length: usize,
24        destination_offset: usize,
25    ) -> Result<TensorReductionLaunch, TensorDeviceError> {
26        let end = destination_offset.checked_add(length)
27            .ok_or(TensorDeviceError::InvalidBuffer("collective reduction range overflow"))?;
28        checked_length::<T>(end)?;
29        Ok(TensorReductionLaunch { length, destination_offset })
30    }
31
32    pub fn reduce<T: TensorElement>(
33        &self,
34        operation: ReductionOperation,
35        launch: &TensorReductionLaunch,
36        source: &TensorBuffer<B, T>,
37        destination: &TensorBuffer<B, T>,
38    ) -> Result<(), TensorDeviceError> {
39        self.validate_buffer(source)?;
40        self.validate_buffer(destination)?;
41        self.reduction_kernel::<T>(operation)?;
42        let source_range = checked_range(source.length, 0, launch.length)?;
43        let target_range = checked_range(destination.length, launch.destination_offset, launch.length)?;
44        if launch.length == 0 {
45            return Ok(());
46        }
47        let apply = |current: Primitive<B>, incoming: Primitive<B>| {
48            if target_range.start == 0 && target_range.end == destination.length {
49                self.reduce_values(operation, current, incoming)
50            } else {
51                let target = current.clone().slice(target_range.clone());
52                let reduced = self.reduce_values(operation, target, incoming);
53                current.assign(target_range.clone(), reduced)
54            }
55        };
56        if Arc::ptr_eq(&source.value, &destination.value) {
57            destination.update(|current| {
58                let incoming = current.clone().slice(source_range);
59                apply(current, incoming)
60            })
61        } else {
62            // Drop the source lock before acquiring the destination lock.
63            let mut incoming = source.snapshot()?;
64            if launch.length != source.length {
65                incoming = incoming.slice(source_range);
66            }
67            destination.update(|current| apply(current, incoming))
68        }
69    }
70
71    fn reduce_values(&self, operation: ReductionOperation, destination: Primitive<B>, source: Primitive<B>) -> Primitive<B> {
72        use ReductionOperation::*;
73        let bool_dtype = get_device_settings::<B>(&self.device).bool_dtype;
74        match (destination, source) {
75            (Primitive::Float(destination), Primitive::Float(source)) => Primitive::Float(match operation {
76                Sum => B::float_add(destination, source),
77                Product => B::float_mul(destination, source),
78                // Preserve the destination for unordered/equal comparisons,
79                // including NaNs and opposite signed zero, like native kernels.
80                Minimum | Maximum => {
81                    let mask = if operation == Minimum {
82                        B::float_lower(source.clone(), destination.clone(), bool_dtype)
83                    } else {
84                        B::float_lower(destination.clone(), source.clone(), bool_dtype)
85                    };
86                    B::float_mask_where(destination, mask, source)
87                }
88                BitAnd | BitOr | BitXor => unreachable!("validated floating reduction"),
89            }),
90            (Primitive::Int(destination), Primitive::Int(source)) => Primitive::Int(match operation {
91                Sum => B::int_add(destination, source),
92                Product => B::int_mul(destination, source),
93                Minimum | Maximum => {
94                    let mask = if operation == Minimum {
95                        B::int_lower(source.clone(), destination.clone(), bool_dtype)
96                    } else {
97                        B::int_lower(destination.clone(), source.clone(), bool_dtype)
98                    };
99                    B::int_mask_where(destination, mask, source)
100                }
101                BitAnd => B::bitwise_and(destination, source),
102                BitOr => B::bitwise_or(destination, source),
103                BitXor => B::bitwise_xor(destination, source),
104            }),
105            _ => unreachable!("typed collective buffer storage kind"),
106        }
107    }
108}
ruccl/tensor_device/reduction.rs

ruccl/tensor_device/
reduction.rs
