tenferro_gpu/cuda/
raw.rs

1//! The owner-scoped raw session an extension crate runs its vendor calls in.
2//!
3//! The session exposes exactly what a vendor library needs — the runtime's stream, device
4//! addresses, and scratch — and nothing that would re-enter the runtime's admission: transfers run
5//! on the runtime before the scope opens, because their admission is the lock the scope holds.
6//!
7//! A caller's tensor is safe to borrow for as long as the vendor work runs: dropping a buffer
8//! releases it with `cuMemFreeAsync` on the same stream, so the free is ordered after any call the
9//! session enqueued. `retain_tensor` therefore only proves the tensor belongs to this runtime.
10
11use tenferro_tensor::{StorageBuffer, TensorScalar, TypedTensor};
12
13use super::engine::CudaEngine;
14use super::storage::CudaStorage;
15
16/// A handle to the runtime's own stream, valid for the session's life.
17#[derive(Debug, Clone, Copy)]
18#[doc(hidden)]
19pub struct StreamRef {
20    handle: u64,
21}
22
23impl StreamRef {
24    /// The raw `CUstream`, as the value a vendor binding takes.
25    #[must_use]
26    #[doc(hidden)]
27    pub fn raw_handle(self) -> u64 {
28        self.handle
29    }
30}
31
32/// A device address the session handed out, valid until the scope's work completes.
33#[derive(Debug, Clone, Copy)]
34#[doc(hidden)]
35pub struct DeviceBytes<'s> {
36    pointer: u64,
37    bytes: usize,
38    _session: std::marker::PhantomData<&'s ()>,
39}
40
41impl DeviceBytes<'_> {
42    /// Run `f` with the device address.
43    #[doc(hidden)]
44    pub fn with_ptr<R>(&self, f: impl FnOnce(*mut std::ffi::c_void) -> R) -> R {
45        f(self.pointer as usize as *mut std::ffi::c_void)
46    }
47
48    /// Number of bytes at the address.
49    #[must_use]
50    #[doc(hidden)]
51    pub fn byte_len(&self) -> usize {
52        self.bytes
53    }
54}
55
56/// A read-only device view of a caller's tensor.
57#[derive(Debug, Clone, Copy)]
58#[doc(hidden)]
59pub struct TensorRef<'a, T> {
60    pointer: u64,
61    bytes: usize,
62    _tensor: std::marker::PhantomData<&'a T>,
63}
64
65impl<T> TensorRef<'_, T> {
66    /// The tensor's device address, as a vendor binding takes it.
67    #[doc(hidden)]
68    pub fn raw_ptr(&self) -> *mut std::ffi::c_void {
69        self.pointer as usize as *mut std::ffi::c_void
70    }
71
72    /// Run `f` with the tensor's device address.
73    #[doc(hidden)]
74    pub fn with_ptr<R>(&self, f: impl FnOnce(*mut std::ffi::c_void) -> R) -> R {
75        f(self.pointer as usize as *mut std::ffi::c_void)
76    }
77
78    /// Number of bytes the tensor spans.
79    #[must_use]
80    #[doc(hidden)]
81    pub fn byte_len(&self) -> usize {
82        self.bytes
83    }
84}
85
86/// A writable device view of a caller's tensor.
87#[derive(Debug, Clone, Copy)]
88#[doc(hidden)]
89pub struct TensorMut<'a, T> {
90    pointer: u64,
91    bytes: usize,
92    _tensor: std::marker::PhantomData<&'a mut T>,
93}
94
95impl<T> TensorMut<'_, T> {
96    /// The tensor's device address, as a vendor binding takes it.
97    #[doc(hidden)]
98    pub fn raw_ptr(&self) -> *mut std::ffi::c_void {
99        self.pointer as usize as *mut std::ffi::c_void
100    }
101
102    /// Run `f` with the tensor's device address.
103    #[doc(hidden)]
104    pub fn with_ptr<R>(&self, f: impl FnOnce(*mut std::ffi::c_void) -> R) -> R {
105        f(self.pointer as usize as *mut std::ffi::c_void)
106    }
107
108    /// Number of bytes the tensor spans.
109    #[must_use]
110    #[doc(hidden)]
111    pub fn byte_len(&self) -> usize {
112        self.bytes
113    }
114}
115
116/// One raw scope, handed to an extension's closure.
117#[doc(hidden)]
118pub struct Session<'s> {
119    runtime: super::CudaRuntime,
120    scope: t4a_cuda_runtime::RawScope<'s>,
121    /// The runtime's extension caches, lent by the engine: a vendor handle is created once per
122    /// engine and reused by every later session. The cell is per session and the session is
123    /// thread-bound, so a shared borrow during a vendor call costs nothing and keeps `resource`
124    /// callable while the session is read.
125    caches: &'s std::cell::RefCell<tenferro_runtime::ExtensionCacheStore>,
126}
127
128impl<'s> Session<'s> {
129    /// The runtime's stream, for a vendor library to bind its plan to.
130    #[must_use]
131    #[doc(hidden)]
132    pub fn stream(&self) -> StreamRef {
133        StreamRef {
134            handle: self.scope.stream(),
135        }
136    }
137
138    /// Zero-filled scratch in this runtime's domain.
139    ///
140    /// # Errors
141    ///
142    /// [`crate::Error::RuntimeStateSource`] when the allocation or its zero fill fails.
143    #[doc(hidden)]
144    pub fn alloc_bytes(&self, nbytes: usize, op: &'static str) -> crate::Result<DeviceBytes<'s>> {
145        let buffer = self
146            .scope
147            .alloc_zeroed(nbytes, 256)
148            .map_err(|error| crate::Error::runtime_state_source(op, error))?;
149        let pointer = self
150            .scope
151            .register(&buffer)
152            .map_err(|error| crate::Error::runtime_state_source(op, error))?;
153        // The allocation lives as long as this handle: it moves into the scope's keepalives so the
154        // device address stays valid until the scope's completion, whatever the caller does.
155        self.scope
156            .retain(Box::new(buffer), nbytes)
157            .map_err(|error| crate::Error::runtime_state_source(op, error))?;
158        Ok(DeviceBytes {
159            pointer,
160            bytes: nbytes,
161            _session: std::marker::PhantomData,
162        })
163    }
164
165    /// Copy `bytes` between two device addresses on this session's stream.
166    ///
167    /// # Errors
168    ///
169    /// [`crate::Error::RuntimeState`] when the driver refuses the copy.
170    #[doc(hidden)]
171    /// # Safety
172    ///
173    /// `destination` and `source` must be device addresses this session's engine owns, `bytes`
174    /// must not exceed either allocation, and the ranges must not overlap unless
175    /// `cuMemcpyDtoDAsync` allows it. The copy is enqueued on this session's stream, so the
176    /// allocations must outlive the point where that work retires.
177    pub unsafe fn copy_bytes(
178        &self,
179        destination: *mut std::ffi::c_void,
180        source: *const std::ffi::c_void,
181        bytes: usize,
182        op: &'static str,
183    ) -> crate::Result<()> {
184        let status = unsafe {
185            cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
186                destination as usize as cudarc::driver::sys::CUdeviceptr,
187                source as usize as cudarc::driver::sys::CUdeviceptr,
188                bytes,
189                self.scope.stream() as usize as *mut cudarc::driver::sys::CUstream_st,
190            )
191        };
192        if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
193            return Err(crate::Error::runtime_state(
194                op,
195                format!("cuMemcpyDtoDAsync failed with {status:?}"),
196            ));
197        }
198        Ok(())
199    }
200
201    /// Copy host `bytes` into fresh scratch of this runtime, on this session's stream.
202    ///
203    /// The host source is pageable and the copy is asynchronous, so the bytes are moved into the
204    /// scope's keepalives: the driver may still be reading them when the caller returns.
205    ///
206    /// # Errors
207    ///
208    /// [`crate::Error::RuntimeState`] when the allocation or the copy is refused.
209    #[doc(hidden)]
210    pub fn upload_bytes(&self, bytes: &[u8], op: &'static str) -> crate::Result<DeviceBytes<'s>> {
211        let device = self.alloc_bytes(bytes.len(), op)?;
212        let owned = bytes.to_vec();
213        let source = owned.as_ptr() as usize;
214        device.with_ptr(|destination| {
215            let status = unsafe {
216                cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
217                    destination as usize as cudarc::driver::sys::CUdeviceptr,
218                    source as *const std::ffi::c_void,
219                    owned.len(),
220                    self.scope.stream() as usize as *mut cudarc::driver::sys::CUstream_st,
221                )
222            };
223            if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
224                return Err(crate::Error::runtime_state(
225                    op,
226                    format!("cuMemcpyHtoDAsync failed with {status:?}"),
227                ));
228            }
229            Ok(())
230        })?;
231        // The driver may still be reading the host source, so it outlives this call through the
232        // scope's completion.
233        self.scope
234            .retain(Box::new(owned), 0)
235            .map_err(|error| crate::Error::runtime_state_source(op, error))?;
236        Ok(device)
237    }
238
239    /// Copy a device tensor into a host tensor of the same shape.
240    ///
241    /// # Errors
242    ///
243    /// [`crate::Error::RuntimeState`] when the tensor is not this runtime's allocation, and the
244    /// runtime's own failure when the copy or its barrier fails.
245    #[doc(hidden)]
246    pub fn download_tensor<T: TensorScalar>(
247        &self,
248        tensor: &TypedTensor<T>,
249        op: &'static str,
250    ) -> crate::Result<TypedTensor<T>> {
251        let (pointer, bytes) = self.address(tensor, op)?;
252        let staging = PinnedStaging::new(bytes, op)?;
253        let status = unsafe {
254            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
255                staging.pointer,
256                pointer as cudarc::driver::sys::CUdeviceptr,
257                bytes,
258                self.scope.stream() as usize as *mut cudarc::driver::sys::CUstream_st,
259            )
260        };
261        if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
262            return Err(crate::Error::runtime_state(
263                op,
264                format!("cuMemcpyDtoHAsync failed with {status:?}"),
265            ));
266        }
267        self.synchronize()?;
268        let values: Vec<T> = staging
269            .as_slice::<T>(bytes / std::mem::size_of::<T>().max(1))
270            .to_vec();
271        TypedTensor::<T>::from_vec_col_major(tensor.shape().to_vec(), values)
272            .map_err(|error| crate::Error::runtime_state(op, error.to_string()))
273    }
274
275    /// Zero-filled device storage of `shape` for `T`, as a tensor of this runtime.
276    ///
277    /// # Errors
278    ///
279    /// [`crate::Error::RuntimeState`] when the allocation fails or the shape does not describe a
280    /// column-major tensor of `T`.
281    #[doc(hidden)]
282    pub fn alloc_output<T: TensorScalar>(&self, shape: &[usize]) -> crate::Result<TypedTensor<T>> {
283        const OP: &str = "cuda raw alloc_output";
284        let count = shape.iter().try_fold(1usize, |product, &extent| {
285            product
286                .checked_mul(extent)
287                .ok_or_else(|| crate::Error::runtime_state(OP, "the element count overflows usize"))
288        })?;
289        let bytes = count
290            .checked_mul(std::mem::size_of::<T>().max(1))
291            .ok_or_else(|| crate::Error::runtime_state(OP, "the byte length overflows usize"))?;
292        let buffer = self
293            .scope
294            .alloc_zeroed(bytes, 256)
295            .map_err(|error| crate::Error::runtime_state_source(OP, error))?;
296        // The returned tensor owns the allocation, so its own drop frees it — in this runtime's
297        // stream order, after any vendor call the session enqueued.
298        let storage = super::storage::fresh_storage(&self.runtime, buffer);
299        TypedTensor::<T>::from_buffer_col_major(
300            shape.to_vec(),
301            tenferro_tensor::StorageBuffer::Backend(Box::new(storage)),
302            super::transfer::placement(&self.runtime),
303        )
304        .map_err(|error| crate::Error::runtime_state(OP, error.to_string()))
305    }
306
307    /// A read-only device view of `tensor`.
308    ///
309    /// # Errors
310    ///
311    /// [`crate::Error::RuntimeState`] when the tensor is not this runtime's allocation.
312    #[doc(hidden)]
313    pub fn tensor<T: TensorScalar>(
314        &self,
315        tensor: &TypedTensor<T>,
316    ) -> crate::Result<TensorRef<'_, T>> {
317        let (pointer, bytes) = self.address(tensor, "cuda raw tensor")?;
318        Ok(TensorRef {
319            pointer,
320            bytes,
321            _tensor: std::marker::PhantomData,
322        })
323    }
324
325    /// A writable device view of `tensor`.
326    ///
327    /// # Errors
328    ///
329    /// [`crate::Error::RuntimeState`] when the tensor is not this runtime's allocation.
330    #[doc(hidden)]
331    pub fn tensor_mut<T: TensorScalar>(
332        &self,
333        tensor: &TypedTensor<T>,
334    ) -> crate::Result<TensorMut<'_, T>> {
335        let (pointer, bytes) = self.address(tensor, "cuda raw tensor mut")?;
336        Ok(TensorMut {
337            pointer,
338            bytes,
339            _tensor: std::marker::PhantomData,
340        })
341    }
342
343    /// Prove `tensor` is one of this runtime's allocations, for a vendor call that reads it.
344    ///
345    /// # Errors
346    ///
347    /// [`crate::Error::RuntimeState`] when the tensor is not this runtime's allocation.
348    #[doc(hidden)]
349    /// # Safety
350    ///
351    /// The returned address stays valid only while `tensor`'s allocation is resident on this
352    /// engine; the caller must not outlive it.
353    pub unsafe fn retain_tensor<T: TensorScalar>(
354        &self,
355        tensor: &TypedTensor<T>,
356        op: &'static str,
357    ) -> crate::Result<DeviceBytes<'static>> {
358        let (pointer, bytes) = self.address(tensor, op)?;
359        Ok(DeviceBytes {
360            pointer,
361            bytes,
362            _session: std::marker::PhantomData,
363        })
364    }
365
366    /// This session's vendor resource of type `T`, created by `init` on first use.
367    ///
368    /// One slot per type: the cache key carries the type name, so two resources of different types
369    /// in one engine cannot collide, and the same handle is reused by every later session.
370    ///
371    /// # Examples
372    ///
373    /// ```no_run
374    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
375    /// use tenferro_gpu::cuda::{gpu_available, CudaDeviceId, CudaEngine};
376    ///
377    /// # #[cfg(feature = "cuda")]
378    /// if gpu_available() {
379    ///     let mut engine = CudaEngine::open(CudaDeviceId::from_ordinal(0))?;
380    ///     // The guard borrows the session, so it stays inside the closure.
381    ///     let total = engine.with_raw("resource", |session| {
382    ///         let counter = session.resource(|| Ok::<u64, tenferro_tensor::Error>(7))?;
383    ///         Ok(*counter + 1)
384    ///     })?;
385    ///     assert_eq!(total, 8);
386    /// }
387    /// # Ok(())
388    /// # }
389    /// ```
390    ///
391    /// # Panics
392    ///
393    /// Only if storing the initialised resource succeeded and it is then missing from the cache,
394    /// which is an internal invariant of the store above.
395    ///
396    /// # Errors
397    ///
398    /// The initialiser's failure, and the engine's when the cache is poisoned.
399    pub fn resource<T: Send + Sync + 'static>(
400        &self,
401        init: impl FnOnce() -> crate::Result<T>,
402    ) -> crate::Result<CudaResourceGuard<'_, T>> {
403        let key = tenferro_runtime::ExtensionCacheKey::new(
404            "tenferro.cuda.session",
405            std::any::type_name::<T>(),
406            0,
407        );
408        let mut store = self.caches.borrow_mut();
409        if store.get::<T>(&key).is_none() {
410            let value = init()?;
411            store.put(key, value, 0);
412        }
413        let entry = std::cell::RefMut::map(store, |store| {
414            store
415                .get_mut::<T>(&key)
416                .expect("the resource was stored above")
417        });
418        Ok(CudaResourceGuard { entry })
419    }
420
421    /// Wait for the work enqueued on this session's stream.
422    ///
423    /// # Errors
424    ///
425    /// [`crate::Error::RuntimeStateSource`] when the stream cannot be synchronised.
426    #[doc(hidden)]
427    pub fn synchronize(&self) -> crate::Result<()> {
428        let stream = self.scope.stream() as usize as *mut cudarc::driver::sys::CUstream_st;
429        // SAFETY: the scope hands out this runtime's own stream; a barrier on it neither reads nor
430        // writes an operand.
431        let status = unsafe { cudarc::driver::sys::cuStreamSynchronize(stream) };
432        if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
433            return Err(crate::Error::runtime_state(
434                "cuda raw synchronize",
435                format!("cuStreamSynchronize failed with {status:?}"),
436            ));
437        }
438        Ok(())
439    }
440
441    fn address<T: TensorScalar>(
442        &self,
443        tensor: &TypedTensor<T>,
444        op: &'static str,
445    ) -> crate::Result<(u64, usize)> {
446        let storage = match tensor.buffer() {
447            StorageBuffer::Backend(buffer) => buffer
448                .as_ref()
449                .as_any()
450                .downcast_ref::<CudaStorage>()
451                .ok_or_else(|| {
452                    crate::Error::runtime_state(op, "expected a CUDA allocation of this backend")
453                })?,
454            StorageBuffer::Host(_) => {
455                return Err(crate::Error::runtime_state(
456                    op,
457                    "expected a device tensor, not host storage",
458                ))
459            }
460        };
461        if storage.domain() != self.runtime.allocation_domain_id() {
462            return Err(crate::Error::runtime_state(
463                op,
464                "the tensor belongs to another runtime",
465            ));
466        }
467        let pointer = self
468            .scope
469            .register(storage.allocation())
470            .map_err(|error| crate::Error::runtime_state_source(op, error))?;
471        Ok((pointer, storage.bytes()))
472    }
473}
474
475/// Marks a [`NativeSessionRef`](tenferro_tensor::NativeSessionRef) as this engine's, so an
476/// extension can reach the concrete engine through `&mut dyn Engine`.
477#[derive(Debug, Clone, Copy)]
478pub(crate) struct CudaEngineMarker;
479
480/// Visit this engine through the tensor crate's engine trait, as an extension does.
481///
482/// Returns `None` when `engine` is not this backend's engine.
483#[doc(hidden)]
484pub fn with_cuda_engine<B, R>(engine: &mut B, f: impl FnOnce(&mut CudaEngine) -> R) -> Option<R>
485where
486    B: tenferro_tensor::Engine + ?Sized,
487{
488    let pointer = engine
489        .native_engine()?
490        .into_marked_ptr::<CudaEngineMarker>()?;
491    // SAFETY: only this crate's `native_engine` creates a token carrying this marker, and it points
492    // that token at a live `CudaEngine`; the `&mut B` keeps the exclusive borrow for the visit.
493    Some(f(unsafe { pointer.cast::<CudaEngine>().as_mut() }))
494}
495
496impl CudaEngine {
497    /// Run `f` inside this engine's raw scope.
498    ///
499    /// Every vendor call the closure makes must be enqueued on [`Session::stream`]; a buffer the
500    /// caller drops afterwards is freed in that stream's order, so the call still sees it.
501    ///
502    /// # Errors
503    ///
504    /// The closure's own failure, and the runtime's when the scope cannot be opened.
505    #[doc(hidden)]
506    pub fn with_raw<R>(
507        &mut self,
508        _op: &'static str,
509        f: impl FnOnce(&mut Session<'_>) -> crate::Result<R>,
510    ) -> crate::Result<R> {
511        // The runtime handle is cloned first so the engine's caches can be borrowed mutably for the
512        // same call.
513        let runtime = self.runtime().clone();
514        let caches = &self.extension_caches;
515        // SAFETY: the closure only reads the addresses the session registers and enqueues its work
516        // on this scope's stream.
517        unsafe {
518            runtime
519                .with_raw_locked(Vec::new(), |scope| {
520                    let mut session = Session {
521                        runtime: runtime.clone(),
522                        scope,
523                        caches,
524                    };
525                    f(&mut session)
526                })
527                .and_then(|result| result)
528        }
529    }
530}
531
532/// A pinned host buffer used as the staging area of one download.
533struct PinnedStaging {
534    pointer: *mut std::ffi::c_void,
535    bytes: usize,
536}
537
538impl PinnedStaging {
539    fn new(bytes: usize, op: &'static str) -> crate::Result<PinnedStaging> {
540        let mut pointer: *mut std::ffi::c_void = std::ptr::null_mut();
541        let status = unsafe {
542            cudarc::driver::sys::cuMemHostAlloc(
543                &mut pointer,
544                bytes.max(1),
545                cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE,
546            )
547        };
548        if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
549            return Err(crate::Error::runtime_state(
550                op,
551                format!("cuMemHostAlloc failed with {status:?}"),
552            ));
553        }
554        Ok(PinnedStaging { pointer, bytes })
555    }
556
557    fn as_slice<T>(&self, count: usize) -> &[T] {
558        // SAFETY: the allocation holds `bytes` pinned bytes, the copy above filled them after a
559        // stream barrier, and `T` is plain data whose count fits the allocation.
560        unsafe { std::slice::from_raw_parts(self.pointer as *const T, count) }
561    }
562}
563
564impl Drop for PinnedStaging {
565    fn drop(&mut self) {
566        // SAFETY: the pointer came from `cuMemHostAlloc` and is freed once here.
567        unsafe {
568            let _ = cudarc::driver::sys::cuMemFreeHost(self.pointer);
569        }
570        self.bytes = 0;
571    }
572}
573
574/// A vendor resource lent from the engine's extension caches for the session's life.
575///
576/// # Examples
577///
578/// ```no_run
579/// use tenferro_gpu::cuda::raw::CudaResourceGuard;
580///
581/// fn handle_lifetime(guard: &CudaResourceGuard<'_, u64>) -> u64 {
582///     // The guard dereferences to the cached resource and releases it when dropped.
583///     **guard + 1
584/// }
585///
586/// # fn main() {}
587/// ```
588pub struct CudaResourceGuard<'s, T> {
589    entry: std::cell::RefMut<'s, T>,
590}
591
592impl<T> std::ops::Deref for CudaResourceGuard<'_, T> {
593    type Target = T;
594
595    fn deref(&self) -> &T {
596        &self.entry
597    }
598}
599
600impl<T> std::ops::DerefMut for CudaResourceGuard<'_, T> {
601    fn deref_mut(&mut self) -> &mut T {
602        &mut self.entry
603    }
604}
605
606impl<T: std::fmt::Debug> std::fmt::Debug for CudaResourceGuard<'_, T> {
607    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
608        self.entry.fmt(formatter)
609    }
610}
tenferro_gpu/cuda/raw.rs

tenferro_gpu/cuda/
raw.rs
