ruccl/tensor_device/
mod.rs

1//! Shared tensor-backend adapters for host-staged rank and in-process collectives.
2//!
3//! Transfers use the existing collective transports. Reductions execute through
4//! the selected backend, including its normal runtime, fusion and copy-on-write
5//! rules. This module does not provide a peer-memory or NCCL ABI implementation.
6
7mod device;
8mod element;
9mod error;
10mod reduction;
11mod storage;
12mod communicator;
13mod tensor;
14
15pub use element::TensorElement;
16pub use error::TensorDeviceError;
17
18use ruda_tensor::{Backend, DType, TensorMetadata};
19use std::marker::PhantomData;
20use std::sync::{Arc, Mutex};
21
22/// Execution context shared by both collective device interfaces.
23#[derive(Clone, Debug)]
24pub struct TensorDevice<B: Backend> {
25    device: B::Device,
26}
27
28/// A typed, one-dimensional collective buffer on a tensor backend.
29///
30/// Clones refer to the same collective state. Exported tensor primitives are
31/// snapshots: later collective writes use the backend's copy-on-write semantics
32/// and do not modify previously exported tensors.
33#[derive(Clone, Debug)]
34pub struct TensorBuffer<B: Backend, T: TensorElement> {
35    value: Arc<Mutex<Option<Primitive<B>>>>,
36    device: B::Device,
37    length: usize,
38    element: PhantomData<T>,
39}
40
41#[derive(Clone, Debug)]
42enum Primitive<B: Backend> {
43    Float(B::FloatTensorPrimitive),
44    Int(B::IntTensorPrimitive),
45}
46
47/// Checked launch coordinates, in elements rather than bytes.
48#[derive(Clone, Copy, Debug)]
49pub struct TensorReductionLaunch {
50    length: usize,
51    destination_offset: usize,
52}
53
54impl<B: Backend> TensorDevice<B> {
55    pub fn new(device: B::Device) -> Self {
56        Self { device }
57    }
58
59    pub fn device(&self) -> &B::Device {
60        &self.device
61    }
62
63    /// Import a rank-one floating tensor without downloading its values.
64    pub fn import_float<T: TensorElement>(
65        &self,
66        value: B::FloatTensorPrimitive,
67    ) -> Result<TensorBuffer<B, T>, TensorDeviceError> {
68        if T::dtype() == DType::I32 {
69            return Err(TensorDeviceError::InvalidBuffer("integer elements require import_int"));
70        }
71        self.validate_type::<T>()?;
72        self.validate_import::<T>(&value, &B::float_device(&value))?;
73        let length = value.shape()[0];
74        Ok(self.wrap(Primitive::Float(value), length))
75    }
76
77    /// Import a rank-one integer tensor without downloading its values.
78    pub fn import_int<T: TensorElement>(
79        &self,
80        value: B::IntTensorPrimitive,
81    ) -> Result<TensorBuffer<B, T>, TensorDeviceError> {
82        if T::dtype() != DType::I32 {
83            return Err(TensorDeviceError::InvalidBuffer("floating elements require import_float"));
84        }
85        self.validate_type::<T>()?;
86        self.validate_import::<T>(&value, &B::int_device(&value))?;
87        let length = value.shape()[0];
88        Ok(self.wrap(Primitive::Int(value), length))
89    }
90
91    pub fn upload<T: TensorElement>(
92        &self,
93        values: &[T],
94    ) -> Result<TensorBuffer<B, T>, TensorDeviceError> {
95        self.validate_type::<T>()?;
96        storage::checked_length::<T>(values.len())?;
97        Ok(self.wrap(self.from_values(values), values.len()))
98    }
99
100    /// Wait for work already submitted to the selected tensor backend.
101    pub fn synchronize(&self) -> Result<(), TensorDeviceError> {
102        B::sync(&self.device).map_err(Into::into)
103    }
104
105    fn wrap<T: TensorElement>(&self, value: Primitive<B>, length: usize) -> TensorBuffer<B, T> {
106        TensorBuffer {
107            value: Arc::new(Mutex::new(Some(value))),
108            device: self.device.clone(),
109            length,
110            element: PhantomData,
111        }
112    }
113
114    fn validate_type<T: TensorElement>(&self) -> Result<(), TensorDeviceError> {
115        if !B::supports_dtype(&self.device, T::dtype()) {
116            return Err(TensorDeviceError::UnsupportedDType(T::dtype()));
117        }
118        Ok(())
119    }
120
121    fn validate_import<T: TensorElement>(
122        &self,
123        value: &impl TensorMetadata,
124        device: &B::Device,
125    ) -> Result<(), TensorDeviceError> {
126        if value.dtype() != T::dtype() {
127            return Err(TensorDeviceError::DTypeMismatch {
128                expected: T::dtype(),
129                actual: value.dtype(),
130            });
131        }
132        if device != &self.device {
133            return Err(TensorDeviceError::DeviceMismatch);
134        }
135        let shape = value.shape();
136        if shape.num_dims() != 1 {
137            return Err(TensorDeviceError::InvalidBuffer("collective tensors must have rank one"));
138        }
139        storage::checked_length::<T>(shape[0])
140    }
141
142    fn validate_buffer<T: TensorElement>(
143        &self,
144        buffer: &TensorBuffer<B, T>,
145    ) -> Result<(), TensorDeviceError> {
146        if buffer.device != self.device {
147            return Err(TensorDeviceError::DeviceMismatch);
148        }
149        Ok(())
150    }
151}
152
153impl<B: Backend, T: TensorElement> TensorBuffer<B, T> {
154    pub fn len(&self) -> usize {
155        self.length
156    }
157
158    pub fn is_empty(&self) -> bool {
159        self.length == 0
160    }
161
162    pub fn dtype(&self) -> DType {
163        T::dtype()
164    }
165
166    pub fn device(&self) -> &B::Device {
167        &self.device
168    }
169
170    /// Export the current floating tensor without a host transfer.
171    pub fn float_tensor(&self) -> Result<B::FloatTensorPrimitive, TensorDeviceError> {
172        match self.snapshot()? {
173            Primitive::Float(value) => Ok(value),
174            Primitive::Int(_) => Err(TensorDeviceError::InvalidBuffer("expected a floating tensor")),
175        }
176    }
177
178    /// Export the current integer tensor without a host transfer.
179    pub fn int_tensor(&self) -> Result<B::IntTensorPrimitive, TensorDeviceError> {
180        match self.snapshot()? {
181            Primitive::Int(value) => Ok(value),
182            Primitive::Float(_) => Err(TensorDeviceError::InvalidBuffer("expected an integer tensor")),
183        }
184    }
185
186    fn snapshot(&self) -> Result<Primitive<B>, TensorDeviceError> {
187        self.value.lock().map_err(|_| TensorDeviceError::Poisoned)?
188            .as_ref().cloned().ok_or(TensorDeviceError::Poisoned)
189    }
190
191    fn update(&self, update: impl FnOnce(Primitive<B>) -> Primitive<B>) -> Result<(), TensorDeviceError> {
192        let mut value = self.value.lock().map_err(|_| TensorDeviceError::Poisoned)?;
193        let previous = value.take().ok_or(TensorDeviceError::Poisoned)?;
194        *value = Some(update(previous));
195        Ok(())
196    }
197}
ruccl/tensor_device/mod.rs

ruccl/tensor_device/
mod.rs
