metrale_gpu_runtime/
pinned_hosts.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Which host buffers are page-locked, so the CUDA backend's
4//! `copy_h2d_async` can keep its promise that the caller may drop the source on
5//! return.
6//!
7//! From pageable memory the driver stages the bytes before the call returns; from
8//! page-locked memory the DMA engine reads the caller's pages after it returns.
9//! `MetraleCudaBackend::alloc_host_pinned` registers each buffer it allocates and
10//! `free_host_pinned` unregisters it before the free. `copy_h2d_async` asks
11//! [`is_pinned`] and, for a pinned source, synchronises the stream before
12//! returning and logs a warning once per process.
13//!
14//! Owner: gpu-runtime.
15//! Invariants:
16//! - Only memory from `alloc_host_pinned` is tracked. Page-locked memory from
17//!   another route, such as the `cudaHostAlloc` workspace in `flashinfer.rs`, is
18//!   not seen.
19//! - Entries stay sorted by base, with at most one entry per base.
20
21use std::sync::RwLock;
22
23/// 2026-09-25: `(base, len)` of each live page-locked host region, sorted by base.
24static PINNED: RwLock<Vec<(usize, usize)>> = RwLock::new(Vec::new());
25
26/// 2026-09-25: Record a page-locked region. A repeated base replaces the stored
27/// length; a null pointer or a zero length is ignored.
28pub fn register(ptr: *const u8, bytes: usize) {
29    if ptr.is_null() || bytes == 0 {
30        return;
31    }
32    let base = ptr as usize;
33    let mut g = match PINNED.write() {
34        Ok(g) => g,
35        // 2026-09-25: A poisoned lock only means an earlier holder panicked;
36        // recover the table rather than propagate the panic.
37        Err(e) => e.into_inner(),
38    };
39    match g.binary_search_by_key(&base, |&(b, _)| b) {
40        Ok(i) => g[i].1 = bytes,
41        Err(i) => g.insert(i, (base, bytes)),
42    }
43}
44
45/// 2026-09-25: Forget a region. `free_host_pinned` calls it before the free, so a
46/// reused address is not reported as still pinned.
47pub fn unregister(ptr: *const u8) {
48    if ptr.is_null() {
49        return;
50    }
51    let base = ptr as usize;
52    let mut g = match PINNED.write() {
53        Ok(g) => g,
54        Err(e) => e.into_inner(),
55    };
56    if let Ok(i) = g.binary_search_by_key(&base, |&(b, _)| b) {
57        g.remove(i);
58    }
59}
60
61/// 2026-09-25: Whether all of `src` lies inside one live registered region. An
62/// empty slice is never pinned.
63pub fn is_pinned(src: &[u8]) -> bool {
64    if src.is_empty() {
65        return false;
66    }
67    let start = src.as_ptr() as usize;
68    let g = match PINNED.read() {
69        Ok(g) => g,
70        Err(e) => e.into_inner(),
71    };
72    // 2026-09-25: Largest base <= start, then a containment test, so a sub-slice
73    // of a registered region matches too.
74    match g.binary_search_by_key(&start, |&(b, _)| b) {
75        Ok(i) => start + src.len() <= g[i].0 + g[i].1,
76        Err(0) => false,
77        Err(i) => {
78            let (b, len) = g[i - 1];
79            start < b + len && start + src.len() <= b + len
80        }
81    }
82}
83
84/// 2026-09-25: Number of live registered regions.
85pub fn live_count() -> usize {
86    match PINNED.read() {
87        Ok(g) => g.len(),
88        Err(e) => e.into_inner().len(),
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    /// 2026-09-25: Serialises this module's tests, since `PINNED` is process-global.
97    static SERIAL: std::sync::Mutex<()> = std::sync::Mutex::new(());
98
99    #[test]
100    fn a_pageable_buffer_is_not_pinned() {
101        let _s = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
102        let v = vec![0u8; 128];
103        assert!(!is_pinned(&v));
104    }
105
106    #[test]
107    fn a_registered_region_and_its_subslices_are_pinned() {
108        let _s = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
109        let mut v = vec![0u8; 256];
110        register(v.as_ptr(), v.len());
111
112        assert!(is_pinned(&v), "the whole region");
113        // 2026-09-25: A base-address-only lookup would miss these sub-slices.
114        assert!(is_pinned(&v[..64]), "prefix");
115        assert!(is_pinned(&v[64..128]), "interior");
116        assert!(is_pinned(&v[255..]), "last byte");
117
118        unregister(v.as_ptr());
119        assert!(!is_pinned(&v), "unregistered again");
120        v[0] = 1;
121        assert_eq!(v[0], 1);
122    }
123
124    #[test]
125    fn a_neighbouring_pageable_buffer_is_not_caught() {
126        let _s = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
127        let pinned = vec![0u8; 128];
128        let other = vec![0u8; 128];
129        register(pinned.as_ptr(), pinned.len());
130        assert!(is_pinned(&pinned));
131        assert!(!is_pinned(&other), "a different allocation must not match");
132        unregister(pinned.as_ptr());
133    }
134
135    #[test]
136    fn empty_and_null_are_handled() {
137        let _s = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
138        let before = live_count();
139        register(std::ptr::null(), 16);
140        register(0x1000 as *const u8, 0);
141        assert_eq!(live_count(), before, "neither is a real region");
142        unregister(std::ptr::null());
143        assert!(!is_pinned(&[]));
144    }
145
146    #[test]
147    fn re_registering_the_same_base_updates_the_length() {
148        let _s = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
149        let v = vec![0u8; 256];
150        let before = live_count();
151        register(v.as_ptr(), 64);
152        register(v.as_ptr(), 256);
153        assert_eq!(live_count(), before + 1, "one entry, not two");
154        assert!(is_pinned(&v[..256]), "the updated length is in effect");
155        unregister(v.as_ptr());
156    }
157}
metrale_gpu_runtime/pinned_hosts.rs

metrale_gpu_runtime/
pinned_hosts.rs
