vortx 0.2.0

Cross-platform GPU tensor library with Rust.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

use crate::shapes::{GGML_IDS, TensorLayout};
use crate::tensor::{AsTensorRef, Tensor, TensorRef};
use bytemuck::NoUninit;
use khal::backend::{Buffer, DeviceValue, GpuBuffer, GpuBufferSlice, GpuBufferSliceMut};

/// A view over a mutable vortx.
///
/// This is a view over an entier tensor, or only part of it, with a shape that doesn't necessarily
/// match the original tensor's shape.
pub struct TensorMut<'a, T: DeviceValue> {
    layout: TensorLayout,
    buffer: &'a mut GpuBuffer<T>,
}

pub trait AsTensorMut<T: DeviceValue> {
    fn as_tensor_mut(&mut self) -> TensorMut<'_, T>;
}

impl<'a, T: DeviceValue> From<&'a mut Tensor<T>> for TensorMut<'a, T> {
    fn from(val: &'a mut Tensor<T>) -> TensorMut<'a, T> {
        val.as_view_mut()
    }
}

impl<T: DeviceValue> AsTensorMut<T> for Tensor<T> {
    #[inline]
    fn as_tensor_mut(&mut self) -> TensorMut<'_, T> {
        self.as_view_mut()
    }
}

impl<T: DeviceValue> AsTensorMut<T> for &mut Tensor<T> {
    #[inline]
    fn as_tensor_mut(&mut self) -> TensorMut<'_, T> {
        self.as_view_mut()
    }
}

impl<'a, T: DeviceValue> AsTensorMut<T> for TensorMut<'a, T> {
    #[inline]
    fn as_tensor_mut(&mut self) -> TensorMut<'_, T> {
        TensorMut {
            layout: self.layout,
            buffer: &mut *self.buffer,
        }
    }
}

impl<'a, 'b, T: DeviceValue> AsTensorMut<T> for &'b mut TensorMut<'a, T> {
    #[inline]
    fn as_tensor_mut(&mut self) -> TensorMut<'_, T> {
        TensorMut {
            layout: self.layout,
            buffer: &mut *self.buffer,
        }
    }
}

impl<'a, T: DeviceValue> AsTensorRef<T> for TensorMut<'a, T> {
    #[inline]
    fn as_tensor_ref(&self) -> TensorRef<'_, T> {
        self.as_ref()
    }
}

impl<'a, 'b, T: DeviceValue> AsTensorRef<T> for &'b TensorMut<'a, T> {
    #[inline]
    fn as_tensor_ref(&self) -> TensorRef<'_, T> {
        self.as_ref()
    }
}

impl<'a, T: DeviceValue> TensorMut<'a, T> {
    pub(crate) fn new(layout: TensorLayout, buffer: &'a mut GpuBuffer<T>) -> TensorMut<'a, T> {
        TensorMut { layout, buffer }
    }

    pub(crate) fn contiguous(dims: &[u32], buffer: &'a mut GpuBuffer<T>) -> TensorMut<'a, T> {
        Self::new(TensorLayout::contiguous(dims), buffer)
    }

    /// Converts this mutable view into an immutable view.
    pub fn as_ref(&self) -> TensorRef<'_, T> {
        TensorRef::new(self.layout, &*self.buffer)
    }

    /// Checks if this tensor is contiguous in memory.
    pub fn is_contiguous(&self) -> bool {
        self.layout.is_contiguous()
    }

    /// Checks if `self` contains the same number oof elements and matches exactly the layout of
    /// its underlying `GpuTensor`.
    pub fn is_entire_tensor(&self) -> bool
    where
        T: NoUninit,
    {
        self.as_ref().is_entire_tensor()
    }

    pub fn rank(&self) -> u32 {
        self.layout().rank
    }

    /// The view's shape.
    pub fn layout(&self) -> TensorLayout {
        self.layout
    }

    /// The view's buffer.
    pub fn buffer(&self) -> GpuBufferSlice<'_, T> {
        self.buffer.slice(self.layout.offset as usize..)
    }

    /// The view's buffer.
    pub fn buffer_mut(&mut self) -> GpuBufferSliceMut<'_, T> {
        self.buffer.slice_mut(self.layout.offset as usize..)
    }

    /// The view's underlying buffer without any offset.
    pub fn raw_buffer(&mut self) -> &mut GpuBuffer<T> {
        self.buffer
    }

    /// Is this view empty?
    pub fn is_empty(&self) -> bool {
        self.layout.is_empty()
    }

    /// The number of elements in this view.
    pub fn len(&self) -> u64 {
        self.layout.len()
    }

    /// Size of this tensor along the dimension `i`.
    pub fn size(&self, i: usize) -> u32 {
        self.layout.size[i]
    }

    /// Size of this tensor along the dimension `i`.
    pub fn size_ggml(&self, i: usize) -> u32 {
        self.layout.size[GGML_IDS[i]]
    }

    /// Stride of this tensor along the dimension `i`.
    pub fn stride(&self, i: usize) -> u32 {
        self.layout.stride[i]
    }

    /// Stride of this tensor along the dimension `i`.
    pub fn stride_ggml(&self, i: usize) -> u32 {
        self.layout.stride[GGML_IDS[i]]
    }

    fn with_layout(self, layout: TensorLayout) -> Self {
        Self {
            layout,
            buffer: self.buffer,
        }
    }

    super::tensor_macro::impl_layout_modifiers! {}
}