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

ruccl/tensor_device/
mod.rs
