cubecl_runtime/server/
handle.rs

1use cubecl_common::device::ServiceId;
2use cubecl_environment::stream::StreamId;
3use cubecl_zspace::{Shape, Strides};
4
5use crate::{
6    memory_management::{
7        ManagedMemoryBinding, ManagedMemoryHandle, ManagedMemoryId, WeakMemoryBinding,
8    },
9    server::{CopyDescriptor, TensorMapBinding},
10};
11
12/// How two handles' bytes relate.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum Overlap {
15    /// No byte in common.
16    Disjoint,
17    /// The same bytes.
18    Same,
19    /// Some bytes in common, but not all.
20    Partial,
21}
22
23/// Server handle containing the [memory handle](crate::server::Handle).
24pub struct Handle {
25    /// Memory handle.
26    pub memory: ManagedMemoryHandle,
27    /// The service whose memory this handle addresses. A client checks it
28    /// before handing the handle to its device: memory coordinates mean
29    /// nothing on another device, and reading them there is not an error
30    /// but garbage.
31    pub service: ServiceId,
32    /// Memory offset in bytes.
33    pub offset_start: Option<u64>,
34    /// Memory offset in bytes.
35    pub offset_end: Option<u64>,
36    /// The stream where the data was created.
37    pub stream: StreamId,
38    /// Length of the underlying buffer ignoring offsets
39    pub(crate) size: u64,
40}
41
42impl core::fmt::Debug for Handle {
43    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
44        f.debug_struct("Handle")
45            .field("id", &self.memory)
46            .field("service", &self.service)
47            .field("offset_start", &self.offset_start)
48            .field("offset_end", &self.offset_end)
49            .field("stream", &self.stream)
50            .field("size", &self.size)
51            .finish()
52    }
53}
54
55impl Clone for Handle {
56    fn clone(&self) -> Self {
57        Self {
58            memory: self.memory.clone(),
59            service: self.service,
60            offset_start: self.offset_start,
61            offset_end: self.offset_end,
62            stream: self.stream,
63            size: self.size,
64        }
65    }
66}
67
68impl Handle {
69    /// Creates a new handle of the given size.
70    pub fn from_memory(
71        id: ManagedMemoryHandle,
72        service: ServiceId,
73        stream: StreamId,
74        size: u64,
75    ) -> Self {
76        Self {
77            memory: id,
78            service,
79            offset_start: None,
80            offset_end: None,
81            stream,
82            size,
83        }
84    }
85    /// Creates a new handle of the given size.
86    pub fn new(service: ServiceId, stream: StreamId, size: u64) -> Self {
87        Self {
88            memory: ManagedMemoryHandle::new(),
89            service,
90            offset_start: None,
91            offset_end: None,
92            stream,
93            size,
94        }
95    }
96    /// Checks whether the handle can be mutated in-place without affecting other computation.
97    pub fn can_mut(&self) -> bool {
98        self.memory.can_mut()
99    }
100
101    /// Returns the [`BufferBinding`] corresponding to the current handle.
102    pub fn binding(self) -> BufferBinding {
103        BufferBinding {
104            memory: self.memory.binding(),
105            service: self.service,
106            offset_start: self.offset_start,
107            offset_end: self.offset_end,
108            stream: self.stream,
109            size: self.size,
110        }
111    }
112
113    /// Add to the current offset in bytes.
114    pub fn offset_start(mut self, offset: u64) -> Self {
115        if let Some(val) = &mut self.offset_start {
116            *val += offset;
117        } else {
118            self.offset_start = Some(offset);
119        }
120
121        self
122    }
123    /// Add to the current offset in bytes.
124    pub fn offset_end(mut self, offset: u64) -> Self {
125        if let Some(val) = &mut self.offset_end {
126            *val += offset;
127        } else {
128            self.offset_end = Some(offset);
129        }
130
131        self
132    }
133
134    /// Convert the [handle](Handle) into a [binding](Binding) with shape and stride metadata.
135    pub fn copy_descriptor(
136        self,
137        shape: Shape,
138        strides: Strides,
139        elem_size: usize,
140    ) -> CopyDescriptor {
141        CopyDescriptor {
142            shape,
143            strides,
144            elem_size,
145            handle: self.binding(),
146        }
147    }
148    /// How this handle's bytes relate to `other`'s.
149    pub fn overlap(&self, other: &Self) -> Overlap {
150        let bytes = |handle: &Self| {
151            handle.offset_start.unwrap_or(0)..handle.size - handle.offset_end.unwrap_or(0)
152        };
153        let mine = bytes(self);
154        let theirs = bytes(other);
155        if self.memory.descriptor().id != other.memory.descriptor().id
156            || mine.end <= theirs.start
157            || theirs.end <= mine.start
158        {
159            Overlap::Disjoint
160        } else if mine == theirs {
161            Overlap::Same
162        } else {
163            Overlap::Partial
164        }
165    }
166    /// Get the size of the handle, in bytes, accounting for offsets
167    pub fn size_in_used(&self) -> u64 {
168        self.size - self.offset_start.unwrap_or(0) - self.offset_end.unwrap_or(0)
169    }
170    /// Get the total size of the handle, in bytes.
171    pub fn size(&self) -> u64 {
172        self.size
173    }
174}
175
176/// A resource passed to a kernel function
177#[allow(clippy::large_enum_variant)]
178#[derive(Clone, Debug)]
179pub enum KernelResource {
180    /// Buffer resource
181    Buffer(BufferBinding),
182    /// Tensor map resource for CUDA
183    TensorMap(TensorMapBinding),
184}
185
186/// A buffer binding represents a [Handle] that is bound to managed memory.
187///
188/// The memory used is known by the compute server.
189/// A buffer binding is only valid after being initlized with [`super::Server::initialize_bindings`]
190///
191/// # Notes
192///
193/// A buffer binding is detached from a [`Handle`], meaning that is won't affect [`Handle::can_mut`].
194#[derive(Clone, Debug)]
195pub struct BufferBinding {
196    /// The id of the handle the binding is bound to.
197    pub memory: ManagedMemoryBinding,
198    /// The service whose memory this binding addresses; see [`Handle::service`].
199    pub service: ServiceId,
200    /// Memory offset in bytes.
201    pub offset_start: Option<u64>,
202    /// Memory offset in bytes.
203    pub offset_end: Option<u64>,
204    /// The stream where the data was created.
205    pub stream: StreamId,
206    /// Length of the underlying buffer ignoring offsets
207    pub size: u64,
208}
209
210/// A [`BufferBinding`] that does not keep its memory reserved: what a record
211/// of a buffer holds when holding the buffer would change what its pool may
212/// reuse.
213#[derive(Clone, Debug)]
214pub struct WeakBufferBinding {
215    memory: WeakMemoryBinding,
216    service: ServiceId,
217    offset_start: Option<u64>,
218    offset_end: Option<u64>,
219    stream: StreamId,
220    size: u64,
221}
222
223impl WeakBufferBinding {
224    /// The binding back, while its memory is allocated (see
225    /// [`WeakMemoryBinding::upgrade`]).
226    pub fn upgrade(&self) -> Option<BufferBinding> {
227        Some(BufferBinding {
228            memory: self.memory.upgrade()?,
229            service: self.service,
230            offset_start: self.offset_start,
231            offset_end: self.offset_end,
232            stream: self.stream,
233            size: self.size,
234        })
235    }
236}
237
238impl BufferBinding {
239    /// This binding, without keeping its memory reserved.
240    pub fn downgrade(&self) -> WeakBufferBinding {
241        WeakBufferBinding {
242            memory: self.memory.downgrade(),
243            service: self.service,
244            offset_start: self.offset_start,
245            offset_end: self.offset_end,
246            stream: self.stream,
247            size: self.size,
248        }
249    }
250
251    /// Get the size of the handle, in bytes, accounting for offsets
252    pub fn size_in_used(&self) -> u64 {
253        self.size - self.offset_start.unwrap_or(0) - self.offset_end.unwrap_or(0)
254    }
255
256    /// The byte range of the allocation this binding names: what the offsets
257    /// leave of the buffer. This is the region the taint bookkeeping claims
258    /// when work writing through this binding fails, and releases when work
259    /// writing through it lands.
260    pub fn range(&self) -> core::ops::Range<u64> {
261        self.offset_start.unwrap_or(0)..self.size - self.offset_end.unwrap_or(0)
262    }
263
264    /// The identity a claim on this binding is stored under: the allocation
265    /// and the byte range of it, exactly as [`range`](Self::range) computes
266    /// it.
267    ///
268    /// Anything that deduplicates or compares claims — a capture's write set,
269    /// above all — must key by this and nothing coarser. Two tensors carved
270    /// from one batched allocation share a memory id and nothing else; keyed
271    /// by the id alone, one sibling's claim swallows the others'.
272    pub fn claim_key(&self) -> (ManagedMemoryId, u64, u64) {
273        let range = self.range();
274        (self.memory.id(), range.start, range.end)
275    }
276    /// Get the total size of the handle, in bytes.
277    pub fn size(&self) -> u64 {
278        self.size
279    }
280}
cubecl_runtime/server/handle.rs

cubecl_runtime/server/
handle.rs
