oxicuda_sparse/
handle.rs

1//! Sparse handle management.
2//!
3//! [`SparseHandle`] is the central object for all sparse operations, analogous
4//! to `cusparseHandle_t` in cuSPARSE. It owns a CUDA stream, a [`BlasHandle`]
5//! for dense sub-operations, a [`PtxCache`] for kernel caching, and caches the
6//! target SM version.
7//!
8//! # Example
9//!
10//! ```rust,no_run
11//! # use std::sync::Arc;
12//! # use oxicuda_driver::Context;
13//! # use oxicuda_sparse::handle::SparseHandle;
14//! # fn main() -> Result<(), oxicuda_sparse::error::SparseError> {
15//! # let ctx: Arc<Context> = unimplemented!();
16//! let handle = SparseHandle::new(&ctx)?;
17//! # Ok(())
18//! # }
19//! ```
20
21use std::sync::Arc;
22
23use oxicuda_blas::BlasHandle;
24use oxicuda_driver::{Context, Stream};
25use oxicuda_ptx::arch::SmVersion;
26use oxicuda_ptx::cache::PtxCache;
27
28use crate::error::{SparseError, SparseResult};
29
30/// Central handle for sparse matrix operations.
31///
32/// Every sparse routine requires a `SparseHandle`. The handle binds operations
33/// to a specific CUDA context and stream, maintains a BLAS sub-handle for
34/// dense operations (e.g. SpMM output accumulation), caches compiled PTX
35/// kernels, and stores the device's SM version for kernel dispatch.
36///
37/// # Thread safety
38///
39/// `SparseHandle` is `Send` but **not** `Sync`. Each thread should create its
40/// own handle (possibly sharing the same [`Arc<Context>`]).
41pub struct SparseHandle {
42    /// The CUDA context this handle is bound to.
43    context: Arc<Context>,
44    /// The stream on which sparse kernels are launched.
45    stream: Stream,
46    /// Sub-handle for BLAS operations.
47    blas_handle: BlasHandle,
48    /// On-disk cache for compiled PTX modules.
49    ptx_cache: PtxCache,
50    /// SM architecture of the device, cached for kernel selection.
51    sm_version: SmVersion,
52}
53
54impl SparseHandle {
55    /// Creates a new sparse handle with a freshly-allocated default stream.
56    ///
57    /// The device's compute capability is queried once and cached.
58    ///
59    /// # Errors
60    ///
61    /// Returns [`SparseError::Cuda`] if stream creation or device query fails.
62    pub fn new(ctx: &Arc<Context>) -> SparseResult<Self> {
63        let stream = Stream::new(ctx)?;
64        Self::build(ctx, stream)
65    }
66
67    /// Creates a new sparse handle bound to an existing stream.
68    ///
69    /// # Errors
70    ///
71    /// Same as [`new`](Self::new) except stream creation cannot fail.
72    pub fn with_stream(ctx: &Arc<Context>, stream: Stream) -> SparseResult<Self> {
73        Self::build(ctx, stream)
74    }
75
76    /// Shared construction logic.
77    fn build(ctx: &Arc<Context>, stream: Stream) -> SparseResult<Self> {
78        let device = ctx.device();
79        let (major, minor) = device.compute_capability()?;
80        let sm_version = SmVersion::from_compute_capability(major, minor).ok_or_else(|| {
81            SparseError::InternalError(format!("unsupported compute capability: {major}.{minor}"))
82        })?;
83
84        let blas_stream = Stream::new(ctx)?;
85        let blas_handle = BlasHandle::with_stream(ctx, blas_stream)?;
86        let ptx_cache = PtxCache::new()?;
87
88        Ok(Self {
89            context: Arc::clone(ctx),
90            stream,
91            blas_handle,
92            ptx_cache,
93            sm_version,
94        })
95    }
96
97    // -- Accessors -----------------------------------------------------------
98
99    /// Returns a reference to the CUDA context.
100    #[inline]
101    pub fn context(&self) -> &Arc<Context> {
102        &self.context
103    }
104
105    /// Returns a reference to the stream used for kernel launches.
106    #[inline]
107    pub fn stream(&self) -> &Stream {
108        &self.stream
109    }
110
111    /// Returns a reference to the internal BLAS handle.
112    #[inline]
113    pub fn blas_handle(&self) -> &BlasHandle {
114        &self.blas_handle
115    }
116
117    /// Returns a mutable reference to the internal BLAS handle.
118    #[inline]
119    pub fn blas_handle_mut(&mut self) -> &mut BlasHandle {
120        &mut self.blas_handle
121    }
122
123    /// Returns the SM version of the bound device.
124    #[inline]
125    pub fn sm_version(&self) -> SmVersion {
126        self.sm_version
127    }
128
129    /// Returns a reference to the PTX cache.
130    #[inline]
131    pub fn ptx_cache(&self) -> &PtxCache {
132        &self.ptx_cache
133    }
134
135    // -- Mutators ------------------------------------------------------------
136
137    /// Replaces the stream used for subsequent sparse operations.
138    pub fn set_stream(&mut self, stream: Stream) {
139        self.stream = stream;
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use oxicuda_ptx::arch::SmVersion;
146
147    #[test]
148    fn sm_version_is_copy() {
149        let v = SmVersion::Sm80;
150        let v2 = v;
151        assert_eq!(v, v2);
152    }
153}
oxicuda_sparse/handle.rs

oxicuda_sparse/
handle.rs
