cubecl_server/storage/
pinned.rs

1//! Page-locked host memory, and handing a slice of it out as [`Bytes`].
2//!
3//! Pinned pages are what a driver can DMA from without a bounce, so every
4//! backend that stages transfers through the host allocates them. What one
5//! looks like from the outside — a pointer, a length, and a binding that keeps
6//! the allocation alive while a caller holds a slice of it — is the same
7//! whichever driver page-locked it.
8
9use crate::memory_management::ManagedMemoryBinding;
10use cubecl_common::bytes::{AccessError, AccessPolicy, AllocationController, AllocationProperty};
11
12/// The alignment pinned allocations are handed out at.
13///
14/// A `u128`'s worth, which is the widest load a host-side copy will make of
15/// staged bytes.
16pub const PINNED_MEMORY_ALIGNMENT: usize = core::mem::size_of::<u128>();
17
18/// A range of page-locked host memory.
19#[derive(Debug)]
20pub struct PinnedMemoryResource {
21    /// Pointer to the pinned memory buffer.
22    pub ptr: *mut u8,
23    /// Size of the memory resource in bytes.
24    pub size: usize,
25}
26
27// SAFETY: the pointer is to page-locked host memory, which stays valid and
28// pinned whichever thread touches it; access is serialized by the device
29// handle above it.
30unsafe impl Send for PinnedMemoryResource {}
31
32/// Hands out a pinned allocation as [`Bytes`](cubecl_common::bytes::Bytes),
33/// keeping the allocation alive for as long as the bytes are.
34///
35/// The binding is held and never read: dropping it is what returns the pages
36/// to the pool, so a caller that still has the slice still has the memory.
37pub struct PinnedMemoryAllocController {
38    resource: PinnedMemoryResource,
39    /// The memory binding, kept alive until deallocation.
40    _binding: ManagedMemoryBinding,
41}
42
43impl PinnedMemoryAllocController {
44    /// A controller over the pinned allocation `binding` names, resolved to
45    /// `resource`.
46    pub fn init(binding: ManagedMemoryBinding, resource: PinnedMemoryResource) -> Self {
47        Self {
48            _binding: binding,
49            resource,
50        }
51    }
52}
53
54impl AllocationController for PinnedMemoryAllocController {
55    fn alloc_align(&self) -> usize {
56        PINNED_MEMORY_ALIGNMENT
57    }
58
59    fn property(&self) -> AllocationProperty {
60        AllocationProperty::Pinned
61    }
62
63    // Pinned host memory is always host-resident: the policy never forces a
64    // copy here.
65    unsafe fn memory_mut(
66        &mut self,
67        _policy: AccessPolicy,
68    ) -> Result<&mut [core::mem::MaybeUninit<u8>], AccessError> {
69        // A zero-size resource carries a NULL pointer — page-locking nothing
70        // succeeds without allocating — which `from_raw_parts_mut` rejects
71        // even for an empty slice. Hand out an aligned dangling pointer.
72        if self.resource.size == 0 {
73            return Ok(empty_pinned_slice_mut());
74        }
75        // SAFETY:
76        // - the pointer is valid while the binding is alive,
77        // - the resource was allocated with `size` bytes,
78        // - `MaybeUninit<u8>` has the same layout as `u8`,
79        // - the caller promises to write only initialized data into it.
80        Ok(unsafe {
81            core::slice::from_raw_parts_mut(
82                self.resource.ptr as *mut core::mem::MaybeUninit<u8>,
83                self.resource.size,
84            )
85        })
86    }
87
88    fn memory(&self, _policy: AccessPolicy) -> Result<&[core::mem::MaybeUninit<u8>], AccessError> {
89        // See `memory_mut`: a zero-size resource carries a NULL pointer.
90        if self.resource.size == 0 {
91            return Ok(empty_pinned_slice_mut());
92        }
93        // SAFETY: as `memory_mut`, without the write.
94        Ok(unsafe {
95            core::slice::from_raw_parts(
96                self.resource.ptr as *mut core::mem::MaybeUninit<u8>,
97                self.resource.size,
98            )
99        })
100    }
101}
102
103/// An empty slice whose dangling pointer still satisfies
104/// [`PINNED_MEMORY_ALIGNMENT`], matching what `alloc_align` advertises.
105fn empty_pinned_slice_mut<'a>() -> &'a mut [core::mem::MaybeUninit<u8>] {
106    // SAFETY: a dangling, well-aligned, non-null pointer is valid for a
107    // zero-length slice.
108    unsafe {
109        core::slice::from_raw_parts_mut(
110            core::ptr::without_provenance_mut(PINNED_MEMORY_ALIGNMENT),
111            0,
112        )
113    }
114}
cubecl_server/storage/pinned.rs

cubecl_server/storage/
pinned.rs
