cubecl_runtime/memory_management/
handle.rs

1use crate::memory_management::MemoryHandle;
2use alloc::{
3    sync::{Arc, Weak},
4    vec::Vec,
5};
6use core::sync::atomic::{AtomicU64, Ordering};
7
8/// Managed Memory handle
9#[derive(Debug)]
10pub struct ManagedMemoryHandle {
11    descriptor: Arc<ManagedMemoryDescriptor>,
12    // Holds only the reference counts of the handle.
13    handle_count: Arc<()>,
14}
15
16/// Binding of a memory handle
17#[derive(Debug)]
18pub struct ManagedMemoryBinding {
19    descriptor: Arc<ManagedMemoryDescriptor>,
20}
21
22/// A [`ManagedMemoryBinding`] that does not keep its memory reserved: what a
23/// record of an allocation holds when holding the allocation itself would
24/// change what the pools may reuse.
25#[derive(Debug, Clone)]
26pub struct WeakMemoryBinding {
27    descriptor: Weak<ManagedMemoryDescriptor>,
28}
29
30/// A list of bindings that are shared across multiple streams.
31#[derive(Debug, Default)]
32pub struct SharedMemoryBindings {
33    /// The bindings.
34    pub bindings: Vec<ManagedMemoryBinding>,
35}
36
37impl Clone for ManagedMemoryHandle {
38    fn clone(&self) -> Self {
39        Self {
40            descriptor: self.descriptor.clone(),
41            handle_count: self.handle_count.clone(),
42        }
43    }
44}
45
46/// Managed memory descriptor.
47///
48/// Multiple handles share the same descriptor via `Arc`, yet the memory
49/// management system needs to update the location after creation (e.g. during
50/// `reserve` / `bind`). The location is packed into an atomic for that.
51///
52/// All mutation happens on the device thread, so `Relaxed` is all the ordering
53/// it needs — and on the targets that matter that is a plain load and store.
54/// Being atomic rather than a `Cell` is what makes the descriptor `Sync` by
55/// construction: these methods are public so the pools in `cubecl-server` can
56/// reach them, and anything public can be called from any thread.
57#[doc(hidden)]
58pub struct ManagedMemoryDescriptor {
59    #[doc(hidden)]
60    pub id: ManagedMemoryId,
61    location: AtomicU64,
62}
63
64impl core::fmt::Debug for ManagedMemoryDescriptor {
65    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
66        f.debug_struct("ManagedMemoryDescriptor")
67            .field("id", &self.id)
68            .field("location", &self.location())
69            .finish()
70    }
71}
72
73#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
74/// Managed memory unique identifier.
75pub struct ManagedMemoryId {
76    #[doc(hidden)]
77    pub value: usize,
78}
79
80impl PartialEq for ManagedMemoryDescriptor {
81    fn eq(&self, other: &Self) -> bool {
82        self.id == other.id
83    }
84}
85
86impl Eq for ManagedMemoryDescriptor {}
87
88#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
89/// Defines where the [`ManagedMemoryId`] is located.
90#[doc(hidden)]
91pub struct MemoryLocation {
92    /// The memory pool index in the global memory management.
93    pub pool: u8,
94    /// The memory page index in a memory pool.
95    pub page: u16,
96    /// The memory slice index in a memory page.
97    pub slice: u32,
98    /// Whether the memory location is known/initialized.
99    pub init: u8,
100}
101
102impl ManagedMemoryDescriptor {
103    /// Update the memory location for the given [`ManagedMemoryId`].
104    #[doc(hidden)]
105    pub fn update_location(&self, location: MemoryLocation) {
106        self.location.store(location.to_bits(), Ordering::Relaxed);
107    }
108
109    /// Update only the slice position for the given [`ManagedMemoryId`].
110    #[doc(hidden)]
111    pub fn update_slice(&self, slice: u32) {
112        self.modify(|location| MemoryLocation { slice, ..location });
113    }
114
115    /// Update only the memory page position for the given [`ManagedMemoryId`].
116    #[doc(hidden)]
117    pub fn update_page(&self, page: u16) {
118        self.modify(|location| MemoryLocation { page, ..location });
119    }
120
121    /// Retrieves the current location.
122    #[doc(hidden)]
123    pub fn location(&self) -> MemoryLocation {
124        MemoryLocation::from_bits(self.location.load(Ordering::Relaxed))
125    }
126
127    /// Whether a reservation ever gave this memory a location.
128    ///
129    /// A handle is minted before its memory exists, and the reservation that
130    /// follows can fail: a full device, a size no pool accepts.
131    #[doc(hidden)]
132    pub fn is_allocated(&self) -> bool {
133        self.location().init != 0
134    }
135
136    #[doc(hidden)]
137    pub fn slice(&self) -> usize {
138        self.location().slice as usize
139    }
140
141    #[doc(hidden)]
142    pub fn page(&self) -> usize {
143        self.location().page as usize
144    }
145
146    fn modify(&self, update: impl Fn(MemoryLocation) -> MemoryLocation) {
147        // Never `Err`: the closure always has an update to make.
148        let _ = self
149            .location
150            .try_update(Ordering::Relaxed, Ordering::Relaxed, |bits| {
151                Some(update(MemoryLocation::from_bits(bits)).to_bits())
152            });
153    }
154}
155
156impl MemoryLocation {
157    /// The location packed into one word, so it can live in an atomic: pool
158    /// in the low byte, then page, then slice, then the init flag on top.
159    fn to_bits(self) -> u64 {
160        self.pool as u64
161            | (self.page as u64) << 8
162            | (self.slice as u64) << 24
163            | (self.init as u64) << 56
164    }
165
166    fn from_bits(bits: u64) -> Self {
167        Self {
168            pool: bits as u8,
169            page: (bits >> 8) as u16,
170            slice: (bits >> 24) as u32,
171            init: (bits >> 56) as u8,
172        }
173    }
174
175    /// Creates a new memory location.
176    #[doc(hidden)]
177    pub fn new(pool: u8, page: u16, slice: u32) -> Self {
178        Self {
179            pool,
180            page,
181            slice,
182            init: 1,
183        }
184    }
185
186    /// Creates a new uninitialized memory location.
187    #[doc(hidden)]
188    pub fn uninit() -> Self {
189        Self {
190            pool: 0,
191            page: 0,
192            slice: 0,
193            init: 0,
194        }
195    }
196}
197
198impl ManagedMemoryHandle {
199    /// Creates a new managed memory handle.
200    pub fn new() -> Self {
201        let value = Self::gen_id();
202
203        Self {
204            descriptor: Arc::new(ManagedMemoryDescriptor {
205                id: ManagedMemoryId { value },
206                location: AtomicU64::new(MemoryLocation::uninit().to_bits()),
207            }),
208            handle_count: Arc::new(()),
209        }
210    }
211
212    /// Retrieves the descriptor for the current handle.
213    #[doc(hidden)]
214    pub fn descriptor(&self) -> &ManagedMemoryDescriptor {
215        &self.descriptor
216    }
217
218    /// Return whether the current handle can be modified in-place.
219    pub fn can_mut(&self) -> bool {
220        Arc::strong_count(&self.handle_count) <= 2
221    }
222
223    /// Whether a [binding](Self::binding) is out: something holds the
224    /// allocation's address, not only its handle.
225    ///
226    /// Handles and bindings share the descriptor, and only handles count in
227    /// `handle_count`, so the difference is the bindings.
228    pub fn is_bound(&self) -> bool {
229        Arc::strong_count(&self.descriptor) > Arc::strong_count(&self.handle_count)
230    }
231
232    /// Return whether the current handle is free.
233    pub fn is_free(&self) -> bool {
234        Arc::strong_count(&self.descriptor) <= 1
235    }
236
237    /// Returns the binding for the current handle.
238    pub fn binding(self) -> ManagedMemoryBinding {
239        ManagedMemoryBinding {
240            descriptor: self.descriptor.clone(),
241        }
242    }
243
244    fn gen_id() -> usize {
245        static COUNTER: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0);
246        let value = COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
247        if value == usize::MAX {
248            core::panic!("Memory ID overflowed");
249        }
250        value
251    }
252}
253
254impl ManagedMemoryBinding {
255    /// Retrieves the descriptor for the current binding.
256    #[doc(hidden)]
257    pub fn descriptor(&self) -> &ManagedMemoryDescriptor {
258        &self.descriptor
259    }
260
261    /// The id of the memory this binding is bound to, stable for as long as the
262    /// allocation lives and never reused by a later one.
263    pub fn id(&self) -> ManagedMemoryId {
264        self.descriptor.id
265    }
266
267    /// This binding, without keeping its memory reserved.
268    pub fn downgrade(&self) -> WeakMemoryBinding {
269        WeakMemoryBinding {
270            descriptor: Arc::downgrade(&self.descriptor),
271        }
272    }
273}
274
275impl WeakMemoryBinding {
276    /// Whether any handle or binding still owns the descriptor, including the pool's reference.
277    /// Unlike [`upgrade`](Self::upgrade), this also counts the sole owner of unallocated memory.
278    pub fn is_alive(&self) -> bool {
279        self.descriptor.strong_count() != 0
280    }
281
282    /// The binding back, while the memory it names is allocated: held by
283    /// something other than the pool that carved it, which keeps a reference
284    /// of its own whether the slice is free or not (see
285    /// [`ManagedMemoryHandle::is_free`]). `None` once every owner let it go.
286    pub fn upgrade(&self) -> Option<ManagedMemoryBinding> {
287        if self.descriptor.strong_count() <= 1 {
288            return None;
289        }
290        let descriptor = self.descriptor.upgrade()?;
291        Some(ManagedMemoryBinding { descriptor })
292    }
293}
294
295impl Default for ManagedMemoryHandle {
296    fn default() -> Self {
297        Self::new()
298    }
299}
300
301impl Clone for ManagedMemoryBinding {
302    fn clone(&self) -> Self {
303        Self {
304            descriptor: self.descriptor.clone(),
305        }
306    }
307}
308
309impl MemoryHandle<ManagedMemoryBinding> for ManagedMemoryHandle {
310    fn can_mut(&self) -> bool {
311        self.can_mut()
312    }
313
314    fn binding(self) -> ManagedMemoryBinding {
315        self.binding()
316    }
317}
318
319impl SharedMemoryBindings {
320    /// Clears the shared bindings list.
321    pub fn clear(&mut self) {
322        self.bindings.clear();
323    }
324
325    /// Returns true if the shared bindings list is empty.
326    pub fn is_empty(&self) -> bool {
327        self.bindings.is_empty()
328    }
329
330    /// Push a memory binding to the list of shared bindings.
331    pub fn push(&mut self, binding: ManagedMemoryBinding) {
332        self.bindings.push(binding)
333    }
334}
335
336impl cubecl_common::pool::Reclaim for SharedMemoryBindings {
337    fn reclaim(&mut self) {
338        self.clear();
339    }
340}
341
342/// Calculates a best-effort heuristic for the alignment of row-aligned tensors.
343/// Prefers contiguous alignments for unit dimensions, 16-byte minimum alignment for non-unit,
344/// scaling with input size up to `buffer_align`.
345pub fn optimal_align(shape: usize, elem_size: usize, buffer_align: usize) -> usize {
346    if shape == 1 {
347        elem_size
348    } else {
349        (shape * elem_size)
350            .next_power_of_two()
351            .clamp(16, buffer_align)
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    #[test]
360    fn test_memory_id_mutability() {
361        let handle1 = ManagedMemoryHandle::new();
362        handle1.descriptor().update_slice(4);
363        assert_eq!(handle1.descriptor().slice(), 4);
364
365        let handle2 = ManagedMemoryHandle::new();
366        handle2
367            .clone()
368            .descriptor()
369            .update_location(handle1.descriptor().location());
370        assert_eq!(handle2.descriptor().slice(), 4);
371    }
372
373    #[test]
374    fn test_location_visible_through_shared_arc() {
375        let handle = ManagedMemoryHandle::new();
376        let handle2 = handle.clone();
377
378        let location = MemoryLocation::new(1, 2, 3);
379        handle.descriptor().update_location(location);
380
381        assert_eq!(handle2.descriptor().location().pool, 1);
382        assert_eq!(handle2.descriptor().location().page, 2);
383        assert_eq!(handle2.descriptor().location().slice, 3);
384        assert_eq!(handle2.descriptor().location().init, 1);
385
386        handle.descriptor().update_slice(42);
387        assert_eq!(handle2.descriptor().slice(), 42);
388    }
389
390    /// Every field gets its own bits in the packed word, so none bleeds into
391    /// its neighbour even at its widest.
392    #[test]
393    fn a_location_survives_packing_at_every_extreme() {
394        let fields = |location: MemoryLocation| {
395            (location.pool, location.page, location.slice, location.init)
396        };
397
398        for location in [
399            MemoryLocation::uninit(),
400            MemoryLocation::new(1, 2, 3),
401            MemoryLocation::new(u8::MAX, u16::MAX, u32::MAX),
402            MemoryLocation {
403                init: u8::MAX,
404                ..MemoryLocation::uninit()
405            },
406        ] {
407            let packed = MemoryLocation::from_bits(location.to_bits());
408
409            assert_eq!(fields(packed), fields(location));
410        }
411    }
412}
cubecl_runtime/memory_management/handle.rs

cubecl_runtime/memory_management/
handle.rs
