tract_gpu/tensor/
mod.rs

1#![allow(clippy::missing_safety_doc)]
2#![allow(clippy::missing_transmute_annotations)]
3
4mod lazy;
5mod owned;
6mod view;
7
8pub use lazy::*;
9pub use owned::*;
10pub use view::*;
11
12use num_traits::AsPrimitive;
13use std::ffi::c_void;
14use std::fmt::Display;
15use tract_core::internal::*;
16use tract_data::itertools::Itertools;
17
18use crate::device::{DeviceBuffer, get_context};
19use crate::utils::MAX_COPY_RANK;
20use crate::utils::merge_copy_axes;
21
22/// A GPU tensor: either a buffer of its own, or a window into someone else's.
23#[derive(Debug, Clone, Hash, PartialEq, Eq)]
24pub enum DeviceTensor {
25    Owned(Box<dyn OwnedDeviceTensor>),
26    View(DeviceView),
27}
28
29impl DeviceTensor {
30    pub const SUPPORTED_DT: [DatumType; 11] = [
31        DatumType::Bool,
32        DatumType::F32,
33        DatumType::F16,
34        DatumType::I8,
35        DatumType::U8,
36        DatumType::I16,
37        DatumType::U16,
38        DatumType::I32,
39        DatumType::U32,
40        DatumType::I64,
41        DatumType::U64,
42    ];
43
44    pub fn tname(dt: DatumType) -> TractResult<&'static str> {
45        Ok(match dt {
46            DatumType::F32 => "f32",
47            DatumType::F16 => "f16",
48            DatumType::U8 => "u8",
49            DatumType::U16 => "u16",
50            DatumType::U32 => "u32",
51            DatumType::U64 => "u64",
52            DatumType::I8 => "i8",
53            DatumType::I16 => "i16",
54            DatumType::I32 => "i32",
55            DatumType::I64 => "i64",
56            DatumType::Bool => "bool",
57            _ => bail!("Unsupported dt {:?} for GPU Tensor", dt),
58        })
59    }
60
61    /// Create an uninitialized DeviceTensor
62    pub fn uninitialized_dt(dt: DatumType, shape: &[usize]) -> TractResult<DeviceTensor> {
63        Ok(DeviceTensor::Owned(get_context()?.uninitialized_device_tensor(shape, dt)?))
64    }
65
66    pub fn uninitialized<T: Datum>(shape: &[usize]) -> TractResult<DeviceTensor> {
67        Self::uninitialized_dt(T::datum_type(), shape)
68    }
69
70    pub fn uninitialized_exotic(exotic_fact: Box<dyn ExoticFact>) -> TractResult<DeviceTensor> {
71        Ok(DeviceTensor::Owned(get_context()?.uninitialized_device_exotic_tensor(exotic_fact)?))
72    }
73    // Create a device tensor with a given shape and a slice of elements. The data is copied and aligned to size of T.
74    pub fn from_shape<T: Copy + Datum>(shape: &[usize], data: &[T]) -> TractResult<DeviceTensor> {
75        Tensor::from_shape(shape, data)?.into_device()
76    }
77
78    pub fn is_supported_dt(dt: DatumType) -> bool {
79        Self::SUPPORTED_DT.contains(&dt)
80    }
81
82    /// Get the datum type of the tensor.
83    #[inline]
84    pub fn datum_type(&self) -> DatumType {
85        match self {
86            Self::Owned(owned) => owned.datum_type(),
87            Self::View(view) => view.datum_type(),
88        }
89    }
90
91    /// Get the number of dimensions (or axes) of the tensor.
92    #[inline]
93    pub fn rank(&self) -> usize {
94        self.shape().len()
95    }
96
97    /// Get the shape of the tensor.
98    #[inline]
99    pub fn shape(&self) -> &[usize] {
100        match self {
101            Self::Owned(t) => t.shape(),
102            Self::View(t) => t.shape(),
103        }
104    }
105
106    /// Get the number of values in the tensor.
107    #[inline]
108    #[allow(clippy::len_without_is_empty)]
109    pub fn len(&self) -> usize {
110        match self {
111            Self::Owned(t) => t.len(),
112            Self::View(t) => t.len(),
113        }
114    }
115
116    /// Get the strides of the tensor.
117    #[inline]
118    pub fn strides(&self) -> &[isize] {
119        match self {
120            Self::Owned(t) => t.strides(),
121            Self::View(t) => t.strides(),
122        }
123    }
124
125    /// Get underlying inner buffer.
126    pub fn device_buffer(&self) -> &dyn DeviceBuffer {
127        match self {
128            Self::Owned(t) => t.device_buffer(),
129            Self::View(t) => t.device_buffer(),
130        }
131    }
132
133    /// Get underlying inner buffer offset
134    pub fn buffer_offset<I: Copy + 'static>(&self) -> I
135    where
136        usize: AsPrimitive<I>,
137    {
138        match self {
139            Self::Owned(_) => 0.as_(),
140            Self::View(t) => t.buffer_offset(),
141        }
142    }
143
144    pub fn device_buffer_ptr(&self) -> *const c_void {
145        match self {
146            Self::Owned(t) => t.device_buffer().ptr(),
147            Self::View(t) => t.device_buffer().ptr(),
148        }
149    }
150
151    /// Returns short description of the inner tensor.
152    pub fn description(&self) -> String {
153        format!("|{},{:?}|", self.shape().iter().join(","), self.datum_type(),)
154    }
155
156    /// Reshaped tensor with given shape.
157    pub fn reshaped(&self, shape: TVec<usize>) -> TractResult<Self> {
158        match self {
159            Self::Owned(t) => Ok(t.reshaped(shape)?),
160            Self::View(t) => Ok(Self::View(t.reshaped(shape)?)),
161        }
162    }
163
164    pub fn restrided(&self, strides: TVec<isize>) -> TractResult<Self> {
165        match self {
166            Self::Owned(t) => Ok(t.restrided(strides)?),
167            Self::View(t) => Ok(Self::View(t.restrided(strides)?)),
168        }
169    }
170
171    /// Convert device tensor to a Tensor backed by device storage.
172    ///
173    /// The resulting tensor carries the real datum type and shape from the
174    /// device tensor (e.g. F32 / \[2,3\]), rather than an exotic scalar wrapper.
175    pub fn into_tensor(self) -> Tensor {
176        let dt = self.datum_type();
177        let shape: TVec<usize> = self.shape().into();
178        Tensor::from_storage(dt, &shape, self)
179    }
180
181    /// The first `len` along `axis`, as a view keeping this tensor's strides.
182    ///
183    /// What a buffer holding spare capacity hands out: the live prefix costs
184    /// nothing to describe, and the tail it leaves behind is simply not part of
185    /// the shape. A [`DeviceView`] is turn-scoped, so this is for a reader
186    /// downstream of the producer, never for the application.
187    ///
188    /// The result is strided wherever `axis` is not the outermost: a consumer
189    /// has to read shape and strides rather than assume a packed tensor.
190    pub fn prefix_window(&self, axis: usize, len: usize) -> TractResult<DeviceTensor> {
191        ensure!(axis < self.rank(), "axis {axis} out of rank {}", self.rank());
192        ensure!(
193            len <= self.shape()[axis],
194            "window of {len} on axis {axis} of extent {}",
195            self.shape()[axis]
196        );
197        ensure!(!self.is_exotic(), "an exotic tensor carries no shape to window");
198        let mut shape: TVec<usize> = self.shape().into();
199        shape[axis] = len;
200        let (buffer, offset_bytes) = match self {
201            Self::Owned(owned) => (Arc::new(tract_core::dyn_clone::clone_box(&**owned)), 0),
202            Self::View(view) => (Arc::clone(&view.buffer), view.offset_bytes),
203        };
204        Ok(DeviceView {
205            buffer,
206            dt: self.datum_type(),
207            len: shape.iter().product(),
208            shape,
209            strides: self.strides().into(),
210            offset_bytes,
211            exotic_fact: None,
212        }
213        .into())
214    }
215
216    /// Slice `[start, end)` along `axis` by a copy made on the device, or `None`
217    /// where the backends' `copy_nd` kernels cannot serve the case and the
218    /// caller has to fall back.
219    ///
220    /// `dt` and `shape` are the owning tensor's, and storage whose geometry has
221    /// drifted from them is refused rather than sliced on stale dimensions.
222    ///
223    /// A copy rather than a view: a view would alias what it was cut from, which
224    /// outlives neither an arena's turn nor a later write to the source. The
225    /// result owns its buffer, and only the source of the copy is strided, which
226    /// is what lets any axis be sliced and not just a dense one.
227    ///
228    /// The copy is ordered on the stream behind whatever produced the source and
229    /// waits for nothing of its own; a caller that needs the bytes now is the
230    /// one that synchronizes.
231    pub(crate) fn slice_on_device(
232        &self,
233        dt: DatumType,
234        shape: &[usize],
235        axis: usize,
236        start: usize,
237        end: usize,
238    ) -> TractResult<Option<DeviceTensor>> {
239        ensure!(
240            axis < shape.len() && start < end && end <= shape[axis],
241            "Invalid slicing range {start}..{end} on axis {axis} of {shape:?}"
242        );
243        let mut sliced: TVec<usize> = shape.into();
244        sliced[axis] = end - start;
245        let servable = dt == self.datum_type()
246            && shape == self.shape()
247            && !self.is_exotic()
248            && Self::is_supported_dt(dt)
249            && self.strides().iter().all(|s| *s > 0)
250            && (1..=MAX_COPY_RANK).contains(
251                &merge_copy_axes(self.strides(), &sliced, &Tensor::natural_strides(&sliced))
252                    .shape
253                    .len(),
254            );
255        if !servable {
256            return Ok(None);
257        }
258        let output = DeviceTensor::uninitialized_dt(dt, &sliced)?;
259        get_context()?.assign_slice(&output, 0..end - start, self, start..end, axis)?;
260        Ok(Some(output))
261    }
262
263    /// Synchronize the GPU Tensor by completing all current
264    /// commands on GPU and returns the inner tensor.
265    pub fn to_host(&self) -> TractResult<Arc<Tensor>> {
266        get_context()?.synchronize()?;
267
268        Ok(match self {
269            Self::Owned(o) => o.to_host()?,
270            Self::View(v) => v.to_host()?.into(),
271        })
272    }
273}
274
275impl Display for DeviceTensor {
276    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
277        match self {
278            Self::Owned(o) => o.fmt(f),
279            Self::View(v) => {
280                let content =
281                    v.to_host().unwrap().dump(false).unwrap_or_else(|e| format!("Error : {e:?}"));
282                write!(f, "View: {{ {content} }}")
283            }
284        }
285    }
286}
287
288pub trait IntoDevice<T> {
289    fn into_device(self) -> TractResult<T>;
290}
291
292impl IntoDevice<DeviceTensor> for Tensor {
293    fn into_device(self) -> TractResult<DeviceTensor> {
294        Ok(DeviceTensor::Owned(get_context()?.tensor_to_device(self.into_tvalue())?))
295    }
296}
297
298impl IntoDevice<DeviceTensor> for Arc<Tensor> {
299    fn into_device(self) -> TractResult<DeviceTensor> {
300        Ok(DeviceTensor::Owned(get_context()?.tensor_to_device(self.into_tvalue())?))
301    }
302}
303
304impl TensorStorage for DeviceTensor {
305    fn byte_len(&self) -> usize {
306        self.len() * self.datum_type().size_of()
307    }
308
309    fn is_empty(&self) -> bool {
310        self.byte_len() == 0
311    }
312
313    fn deep_clone(&self) -> Box<dyn TensorStorage> {
314        Box::new(self.clone())
315    }
316
317    fn as_plain_ram(&self) -> Option<&PlainStorage> {
318        None
319    }
320
321    fn as_plain_ram_mut(&mut self) -> Option<&mut PlainStorage> {
322        None
323    }
324
325    fn into_plain_ram(self: Box<Self>) -> Option<PlainStorage> {
326        None
327    }
328
329    fn dyn_hash(&self, _state: &mut dyn std::hash::Hasher) {
330        // no meaningful hash for device memory
331    }
332
333    fn is_exotic(&self) -> bool {
334        match self {
335            Self::Owned(owned) => owned.exotic_fact().is_some(),
336            Self::View(view) => view.exotic_fact().is_some(),
337        }
338    }
339
340    fn in_ram(&self) -> bool {
341        false
342    }
343
344    fn slice(
345        &self,
346        dt: DatumType,
347        shape: &[usize],
348        axis: usize,
349        start: usize,
350        end: usize,
351    ) -> TractResult<Option<Tensor>> {
352        Ok(self.slice_on_device(dt, shape, axis, start, end)?.map(|t| t.into_tensor()))
353    }
354
355    fn exotic_fact(&self, _shape: &[usize]) -> TractResult<Option<Box<dyn ExoticFact>>> {
356        bail!(
357            "DeviceTensor cannot reconstruct a DeviceFact: origin (FromHost/FromDevice) is not carried by storage"
358        )
359    }
360}
361
362impl From<DeviceView> for DeviceTensor {
363    fn from(view: DeviceView) -> Self {
364        Self::View(view)
365    }
366}
367
368pub trait DeviceTensorExt {
369    fn to_device_tensor(&self) -> TractResult<&DeviceTensor>;
370    fn as_device_tensor(&self) -> Option<&DeviceTensor>;
371}
372
373impl DeviceTensorExt for Tensor {
374    fn to_device_tensor(&self) -> TractResult<&DeviceTensor> {
375        self.try_storage_as::<DeviceTensor>()
376    }
377
378    fn as_device_tensor(&self) -> Option<&DeviceTensor> {
379        self.storage_as::<DeviceTensor>()
380    }
381}
tract_gpu/tensor/mod.rs

tract_gpu/tensor/
mod.rs
