xlog_cuda/device_runtime/
direct.rs

1//! Default-device-stream allocation through the shared raw storage owner.
2//!
3//! Allocation/free use the same explicit mode selected from context capability.
4//! Failed initialization retains actual storage; reclamation proves physical
5//! release before refunding bytes. Caller stream IDs are block metadata, not
6//! proof of arbitrary pool routing. Sanitizer qualification remains separate.
7
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex};
10
11use crate::memory::RawDeviceAllocation;
12
13use super::resource::{
14    AllocationAccounting, AllocationRequest, BlockState, DeviceBlock, DeviceMemoryResource,
15    Generation, ResourceBudgetSnapshot, ResourceError, ResourceResult,
16};
17use crate::CudaDevice;
18
19/// Default-device-stream allocation adaptor. The live map retains actual raw
20/// owners while callers hold opaque [`DeviceBlock`] identities. Deallocation
21/// removes an exact identity and queues its actual owner after unlocking the
22/// map. Reclamation waits for shared leases; an unknown physical-free outcome
23/// retains the owner and its outstanding byte charge for a later reap.
24pub struct DirectCudaResource {
25    device: Arc<CudaDevice>,
26    device_ordinal: u32,
27    /// Live allocation identities and their physical owners.
28    live: Mutex<HashMap<u64, (super::resource::BlockId, Arc<RawDeviceAllocation>)>>,
29    /// Logical detach may leave graph or operation leases alive. The existing
30    /// resource reap contract retires their actual owner after the last lease.
31    pending: Mutex<Vec<Option<Arc<RawDeviceAllocation>>>>,
32    /// All physical bytes, including retained initialization/free failures.
33    /// Only the raw owner's proven release callback decrements this counter.
34    bytes_outstanding: Arc<AllocationAccounting>,
35}
36
37impl DirectCudaResource {
38    /// Construct a resource bound to `device`. `device_ordinal` is the
39    /// CUDA ordinal for logging / multi-device disambiguation.
40    pub fn new(device: Arc<CudaDevice>, device_ordinal: u32) -> Self {
41        Self {
42            device,
43            device_ordinal,
44            live: Mutex::new(HashMap::new()),
45            pending: Mutex::new(Vec::new()),
46            bytes_outstanding: Arc::default(),
47        }
48    }
49
50    /// Borrow the device handle. Tests and downstream resources use
51    /// this to launch kernels against the same device this resource
52    /// allocates on.
53    pub fn device(&self) -> &Arc<CudaDevice> {
54        &self.device
55    }
56}
57
58impl DeviceMemoryResource for DirectCudaResource {
59    fn materialize(&self, request: AllocationRequest) -> ResourceResult<DeviceBlock> {
60        let AllocationRequest {
61            bytes,
62            stream,
63            tag,
64            reclamation,
65            ..
66        } = request;
67        if bytes == 0 {
68            // Zero-byte allocations are not legal in CUDA; surface as
69            // a contract error rather than calling cuMemAlloc(0).
70            return Err(ResourceError::Driver(
71                "DirectCudaResource: zero-byte allocation not supported".to_string(),
72            ));
73        }
74
75        reclamation.attach_resource(Arc::clone(&self.bytes_outstanding), bytes)?;
76        let allocation = crate::memory::RawDeviceAllocation::allocate(
77            Arc::clone(self.device.inner().stream()),
78            Arc::clone(self.device.inner().allocation_stream()),
79            bytes,
80            None,
81            Arc::clone(&reclamation),
82        )?;
83        let ptr = allocation.ptr();
84        let generation = Generation::next();
85        let identity = super::resource::BlockId {
86            ptr,
87            generation,
88            alloc_stream: stream,
89            device_ordinal: self.device_ordinal,
90        };
91        let mut live = self.live.lock().expect("direct allocation map poisoned");
92        if live.contains_key(&ptr) {
93            return Err(
94                ResourceError::Driver(format!("allocation pointer collision: {ptr:#x}"))
95                    .retaining(bytes, reclamation),
96            );
97        }
98        live.insert(ptr, (identity, allocation));
99        Ok(DeviceBlock {
100            ptr,
101            device_ordinal: self.device_ordinal,
102            alloc_stream: stream,
103            bytes,
104            align: std::mem::align_of::<u8>(),
105            tag,
106            generation,
107            state: BlockState::Live,
108        })
109    }
110
111    fn deallocate(&self, block: DeviceBlock) -> ResourceResult<()> {
112        let mut live = self.live.lock().expect("direct allocation map poisoned");
113        let (identity, allocation) = live.get(&block.ptr).ok_or(ResourceError::UseAfterFree {
114            generation: block.generation,
115        })?;
116        super::resource::BlockId::from_block(&block).validate_allocation(
117            block.bytes,
118            *identity,
119            allocation.len(),
120        )?;
121        let (_, allocation) = live.remove(&block.ptr).expect("validated live owner");
122        drop(live);
123        self.pending
124            .lock()
125            .expect("direct pending allocation queue poisoned")
126            .push(Some(allocation));
127        self.reap_pending()
128    }
129
130    fn reap_pending(&self) -> ResourceResult<()> {
131        let mut pending =
132            crate::memory::PendingReclamationBatch::take(&self.pending, |queue, owners| {
133                owners.retain(Option::is_some);
134                queue.append(owners);
135            });
136        crate::memory::reap_raw_allocations(&mut pending)
137    }
138
139    fn access_dependencies(
140        &self,
141        block: super::resource::BlockId,
142        bytes: usize,
143    ) -> ResourceResult<Option<Arc<super::resource::DeviceAccessDependencies>>> {
144        let live = self.live.lock().expect("direct allocation map poisoned");
145        let (identity, allocation) = live.get(&block.ptr).ok_or(ResourceError::UseAfterFree {
146            generation: block.generation,
147        })?;
148        block.validate_allocation(bytes, *identity, allocation.len())?;
149        Ok(Some(allocation.dependencies()))
150    }
151
152    fn device_ordinal(&self) -> u32 {
153        self.device_ordinal
154    }
155
156    fn bytes_outstanding(&self) -> usize {
157        self.bytes_outstanding.snapshot().0
158    }
159
160    fn allocation_accounting(&self) -> Arc<AllocationAccounting> {
161        Arc::clone(&self.bytes_outstanding)
162    }
163
164    fn budget_snapshot(&self) -> Option<ResourceBudgetSnapshot> {
165        None
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::super::resource::{AllocTag, StreamId};
172    use super::*;
173
174    fn try_device() -> Option<Arc<CudaDevice>> {
175        CudaDevice::new(0).ok().map(Arc::new)
176    }
177
178    #[test]
179    fn allocate_then_deallocate_round_trips() {
180        let Some(device) = try_device() else {
181            eprintln!("Skipping: no CUDA device");
182            return;
183        };
184        let r = DirectCudaResource::new(device, 0);
185        assert_eq!(r.bytes_outstanding(), 0);
186
187        let block = r
188            .allocate(4096, StreamId::DEFAULT, AllocTag::UNTAGGED)
189            .expect("alloc");
190        assert_eq!(block.bytes, 4096);
191        assert_eq!(block.state, BlockState::Live);
192        assert_eq!(r.bytes_outstanding(), 4096);
193
194        r.deallocate(block).expect("dealloc");
195        assert_eq!(r.bytes_outstanding(), 0);
196    }
197
198    #[test]
199    fn zero_byte_allocate_rejects() {
200        let Some(device) = try_device() else {
201            return;
202        };
203        let r = DirectCudaResource::new(device, 0);
204        let err = r.allocate(0, StreamId::DEFAULT, AllocTag::UNTAGGED);
205        assert!(matches!(err, Err(ResourceError::Driver(_))));
206        assert_eq!(r.bytes_outstanding(), 0);
207    }
208
209    #[test]
210    fn deallocate_unknown_block_returns_use_after_free() {
211        let Some(device) = try_device() else {
212            return;
213        };
214        let r = DirectCudaResource::new(device, 0);
215        let bogus = DeviceBlock {
216            ptr: 0xdead_beef,
217            device_ordinal: 0,
218            alloc_stream: StreamId::DEFAULT,
219            bytes: 16,
220            align: 1,
221            tag: AllocTag::UNTAGGED,
222            generation: Generation::next(),
223            state: BlockState::Live,
224        };
225        assert!(matches!(
226            r.deallocate(bogus),
227            Err(ResourceError::UseAfterFree { .. })
228        ));
229    }
230
231    /// Locks the contract that DirectCudaResource does NOT
232    /// silently accept `record_block_use`. If a caller (e.g. the
233    /// future xlog launch builder) calls record_block_use against
234    /// a runtime built around DirectCudaResource, the call must
235    /// fail loudly with StreamMisuse — not return Ok and quietly
236    /// fail to track anything. False safety here would let
237    /// downstream code queue cross-stream kernels and drop
238    /// blocks while the cross-stream use was never recorded,
239    /// reproducing exactly the use-after-free this whole layer
240    /// exists to prevent.
241    ///
242    /// Implementation note: DirectCudaResource inherits the
243    /// trait's default `record_block_use` impl (which returns
244    /// `StreamMisuse`). It does NOT override. If a future change
245    /// adds a real override, it must make the override
246    /// genuinely track cross-stream uses (similar to
247    /// AsyncCudaResource's implementation) — anything else
248    /// regresses this contract.
249    #[test]
250    fn record_block_use_rejected_with_stream_misuse() {
251        let Some(device) = try_device() else {
252            return;
253        };
254        let r = DirectCudaResource::new(device, 0);
255        let block = r
256            .allocate(64, StreamId::DEFAULT, AllocTag::UNTAGGED)
257            .expect("alloc");
258        let err = r.record_block_use(&block, StreamId::DEFAULT);
259        match err {
260            Err(ResourceError::StreamMisuse(msg)) => {
261                assert!(
262                    msg.contains("unsupported"),
263                    "expected 'unsupported' in StreamMisuse message, got {:?}",
264                    msg
265                );
266            }
267            other => panic!(
268                "DirectCudaResource::record_block_use must return StreamMisuse \
269                 to surface unsupported cross-stream tracking; got {:?}",
270                other
271            ),
272        }
273        // The block stays live — a failed record_block_use must
274        // NOT have removed the entry or dropped the slice.
275        assert_eq!(r.bytes_outstanding(), 64);
276        r.deallocate(block).expect("dealloc still works");
277    }
278}
xlog_cuda/device_runtime/direct.rs

xlog_cuda/device_runtime/
direct.rs
