ruccl/tensor_device/
storage.rs

1use super::{Primitive, TensorBuffer, TensorDevice, TensorDeviceError, TensorElement};
2use ruda_tensor::{Backend, DType, FloatDType, IntDType, Slice, TensorData, read_sync};
3use std::ops::Range;
4
5pub(super) fn checked_length<T: TensorElement>(length: usize) -> Result<(), TensorDeviceError> {
6    let bytes = length.checked_mul(std::mem::size_of::<T>())
7        .ok_or(TensorDeviceError::InvalidBuffer("collective buffer byte length overflow"))?;
8    if bytes > isize::MAX as usize {
9        return Err(TensorDeviceError::InvalidBuffer("collective buffer exceeds addressable size"));
10    }
11    Ok(())
12}
13
14pub(super) fn checked_range(
15    total: usize,
16    offset: usize,
17    length: usize,
18) -> Result<Range<usize>, TensorDeviceError> {
19    let end = offset.checked_add(length)
20        .ok_or(TensorDeviceError::InvalidBuffer("collective element range overflow"))?;
21    if end > total {
22        return Err(TensorDeviceError::InvalidBuffer("collective element range out of bounds"));
23    }
24    Ok(offset..end)
25}
26
27impl<B: Backend> Primitive<B> {
28    pub(super) fn slice(self, range: Range<usize>) -> Self {
29        let slice = Slice::from(range);
30        match self {
31            Self::Float(value) => Self::Float(B::float_slice(value, &[slice])),
32            Self::Int(value) => Self::Int(B::int_slice(value, &[slice])),
33        }
34    }
35
36    pub(super) fn assign(self, range: Range<usize>, value: Self) -> Self {
37        let slice = Slice::from(range);
38        match (self, value) {
39            (Self::Float(tensor), Self::Float(value)) => {
40                Self::Float(B::float_slice_assign(tensor, &[slice], value))
41            }
42            (Self::Int(tensor), Self::Int(value)) => {
43                Self::Int(B::int_slice_assign(tensor, &[slice], value))
44            }
45            _ => unreachable!("typed collective buffer storage kind"),
46        }
47    }
48}
49
50impl<B: Backend> TensorDevice<B> {
51    pub fn allocate<T: TensorElement>(
52        &self,
53        length: usize,
54    ) -> Result<TensorBuffer<B, T>, TensorDeviceError> {
55        self.validate_type::<T>()?;
56        checked_length::<T>(length)?;
57        let shape = [length].into();
58        let value = match T::dtype() {
59            DType::F32 => Primitive::Float(B::float_empty(shape, &self.device, FloatDType::F32)),
60            DType::F16 => Primitive::Float(B::float_empty(shape, &self.device, FloatDType::F16)),
61            DType::BF16 => Primitive::Float(B::float_empty(shape, &self.device, FloatDType::BF16)),
62            DType::I32 => Primitive::Int(B::int_empty(shape, &self.device, IntDType::I32)),
63            DType::I64 => Primitive::Int(B::int_empty(shape, &self.device, IntDType::I64)),
64            DType::U8 => Primitive::Int(B::int_empty(shape, &self.device, IntDType::U8)),
65            DType::U32 => Primitive::Int(B::int_empty(shape, &self.device, IntDType::U32)),
66            _ => unreachable!("sealed collective element type"),
67        };
68        Ok(self.wrap(value, length))
69    }
70
71    pub(super) fn from_values<T: TensorElement>(&self, values: &[T]) -> Primitive<B> {
72        let data = TensorData::new(values.to_vec(), [values.len()]);
73        if matches!(T::dtype(), DType::U8 | DType::U32 | DType::I32 | DType::I64) {
74            Primitive::Int(B::int_from_data(data, &self.device))
75        } else {
76            Primitive::Float(B::float_from_data(data, &self.device))
77        }
78    }
79
80    pub fn write<T: TensorElement>(
81        &self,
82        buffer: &TensorBuffer<B, T>,
83        values: &[T],
84    ) -> Result<(), TensorDeviceError> {
85        if values.len() != buffer.length {
86            return Err(TensorDeviceError::InvalidBuffer("collective full write length mismatch"));
87        }
88        self.write_at(buffer, 0, values)
89    }
90
91    pub fn write_at<T: TensorElement>(
92        &self,
93        buffer: &TensorBuffer<B, T>,
94        offset: usize,
95        values: &[T],
96    ) -> Result<(), TensorDeviceError> {
97        self.validate_buffer(buffer)?;
98        let range = checked_range(buffer.length, offset, values.len())?;
99        if range.is_empty() {
100            return Ok(());
101        }
102        let incoming = self.from_values(values);
103        buffer.update(|current| {
104            if range.start == 0 && range.end == buffer.length {
105                incoming
106            } else {
107                current.assign(range, incoming)
108            }
109        })
110    }
111
112    pub fn download<T: TensorElement>(
113        &self,
114        buffer: &TensorBuffer<B, T>,
115    ) -> Result<Vec<T>, TensorDeviceError> {
116        self.download_at(buffer, 0, buffer.length)
117    }
118
119    pub fn download_at<T: TensorElement>(
120        &self,
121        buffer: &TensorBuffer<B, T>,
122        offset: usize,
123        length: usize,
124    ) -> Result<Vec<T>, TensorDeviceError> {
125        self.validate_buffer(buffer)?;
126        let range = checked_range(buffer.length, offset, length)?;
127        if range.is_empty() {
128            return Ok(Vec::new());
129        }
130        let mut value = buffer.snapshot()?;
131        if range.start != 0 || range.end != buffer.length {
132            value = value.slice(range);
133        }
134        let data = match value {
135            Primitive::Float(value) => read_sync(B::float_into_data(value))?,
136            Primitive::Int(value) => read_sync(B::int_into_data(value))?,
137        };
138        data.into_vec::<T>().map_err(|error| TensorDeviceError::Data(format!("{error:?}")))
139    }
140}
ruccl/tensor_device/storage.rs

ruccl/tensor_device/
storage.rs
