ruccl/tensor_device/
tensor.rs

1use super::{TensorDevice, TensorDeviceError, TensorElement, storage};
2use crate::ReduceOperation;
3use crate::rank::{ReductionOperation, communicator::RankCommunicator};
4use ruda_tensor::{Backend, DType, Shape, TensorMetadata, bf16, f16};
5
6impl<B: Backend> RankCommunicator<TensorDevice<B>> {
7    /// Reduce a floating tensor through this rank's configured transport.
8    ///
9    /// All ranks must submit the same tensor shapes, dtypes and operations in
10    /// the same order. The backend reshapes non-contiguous inputs according to
11    /// its normal tensor semantics; output shape, dtype and device are retained.
12    /// Transfers are host-staged, while arithmetic uses the tensor backend.
13    pub fn all_reduce_float(
14        &self,
15        value: B::FloatTensorPrimitive,
16        operation: ReduceOperation,
17    ) -> Result<B::FloatTensorPrimitive, TensorDeviceError> {
18        match value.dtype() {
19            DType::F32 => self.reduce_float::<f32>(value, operation),
20            DType::F16 => self.reduce_float::<f16>(value, operation),
21            DType::BF16 => self.reduce_float::<bf16>(value, operation),
22            dtype => Err(TensorDeviceError::UnsupportedDType(dtype)),
23        }
24    }
25
26    fn reduce_float<T: TensorElement>(
27        &self,
28        value: B::FloatTensorPrimitive,
29        operation: ReduceOperation,
30    ) -> Result<B::FloatTensorPrimitive, TensorDeviceError> {
31        let execution = self.execution();
32        execution.validate_type::<T>()?;
33        if &B::float_device(&value) != execution.device() {
34            return Err(TensorDeviceError::DeviceMismatch);
35        }
36        let shape = value.shape();
37        let length = shape.iter().try_fold(1_usize, |length, dim| {
38            length.checked_mul(*dim).ok_or(TensorDeviceError::InvalidBuffer(
39                "collective tensor element count overflow",
40            ))
41        })?;
42        storage::checked_length::<T>(length)?;
43        let value = B::float_reshape(value, Shape::new([length]));
44        let buffer = execution.import_float::<T>(value)?;
45        self.tensor_collective::<T>().all_reduce(
46            &buffer,
47            ReductionOperation::Sum,
48            &ReductionOperation::Sum,
49        )?;
50        let mut value = buffer.float_tensor()?;
51        if operation == ReduceOperation::Mean {
52            value = B::float_div_scalar(value, (self.world_size() as f32).into());
53        }
54        Ok(B::float_reshape(value, shape))
55    }
56}
ruccl/tensor_device/tensor.rs

ruccl/tensor_device/
tensor.rs
