cuest_cudarc/
context.rs

1//! The cuEST context: the top-level handle every other object is derived
2//! from.
3//!
4//! A [`CuestContext`] wraps a user-supplied [`Arc<CudaStream>`]. The stream
5//! (and its cuBLAS/cuSolver resources) is bound to the cuEST handle once at
6//! creation — cuEST has no per-call streams — so device buffers passed to
7//! compute calls must be allocated on this stream.
8
9use crate::error::{status_check, CuestError};
10use crate::parameters::CuestParameters;
11use cudarc::driver::CudaStream;
12use cuest_ffi::ffi;
13use std::ffi::{c_int, c_void, CString};
14use std::path::{Path, PathBuf};
15use std::ptr;
16use std::sync::Arc;
17
18/// Options for [`CuestContext::with_options`].
19///
20/// The defaults (all fields `None`) reproduce cuEST's default handle
21/// parameters. For JIT-related options and semantics, see the cuEST
22/// documentation.
23#[derive(Debug, Default, Clone)]
24pub struct CuestHandleOptions {
25    /// JIT kernel cache directory.
26    ///
27    /// `None` uses the cuEST default
28    /// (`~/.cuest_cache/cuest-cuda<N>-v<version>/`), which requires the
29    /// `HOME` environment variable to be set; set this explicitly in
30    /// sandboxes, CI, and tests. cuEST rejects unsafe or unwritable
31    /// directories.
32    pub jit_cache_dir: Option<PathBuf>,
33    /// Number of parallel JIT precompile worker threads (>= 1, clamped to
34    /// 256 by cuEST). cuEST default: 16.
35    pub jit_compile_threads: Option<i32>,
36    /// Maximum number of Gauss-Hermite quadrature points. cuEST default: 20.
37    pub max_gauss_hermite: Option<u64>,
38    /// Highest angular momentum for solid harmonics. cuEST default: 10.
39    pub max_l_solid_harmonic: Option<u64>,
40    /// Maximum number of Rys quadrature points (0 selects the largest
41    /// available). cuEST default: 0.
42    pub max_rys: Option<u64>,
43}
44
45/// cuEST math mode (see [`CuestContext::math_mode`]).
46#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
47#[non_exhaustive]
48pub enum CuestMathMode {
49    /// Math mode is heuristically chosen per function for maximum efficiency;
50    /// some results may use mixed precision arithmetic.
51    Default,
52    /// Native FP64 compute is used everywhere.
53    NativeFp64,
54}
55
56impl CuestMathMode {
57    fn to_raw(self) -> ffi::cuestMathMode_t {
58        match self {
59            Self::Default => ffi::CUEST_DEFAULT_MATH_MODE,
60            Self::NativeFp64 => ffi::CUEST_NATIVE_FP64_MATH_MODE,
61        }
62    }
63
64    fn from_raw(raw: ffi::cuestMathMode_t) -> Result<Self, CuestError> {
65        match raw {
66            ffi::CUEST_DEFAULT_MATH_MODE => Ok(Self::Default),
67            ffi::CUEST_NATIVE_FP64_MATH_MODE => Ok(Self::NativeFp64),
68            _ => Err(CuestError::Rust(format!("unknown cuestMathMode_t value {}", raw.0))),
69        }
70    }
71}
72
73/// Shared context state; the last [`Arc`] reference destroys the cuEST handle.
74pub(crate) struct ContextInner {
75    handle: ffi::cuestHandle_t,
76    stream: Arc<CudaStream>,
77    version: (u32, u32, u32),
78    /// Kept alive defensively; cuEST is not documented to copy the string.
79    _jit_cache_dir: Option<CString>,
80}
81
82impl ContextInner {
83    pub(crate) fn handle(&self) -> ffi::cuestHandle_t {
84        self.handle
85    }
86
87    pub(crate) fn stream(&self) -> &Arc<CudaStream> {
88        &self.stream
89    }
90
91    pub(crate) fn version(&self) -> (u32, u32, u32) {
92        self.version
93    }
94
95    /// Verify a device buffer was allocated on this context's stream
96    /// (`buffer.stream()`, obtained from any `DeviceSlice`).
97    ///
98    /// cuEST has a single stream bound at handle creation; using buffers from
99    /// another stream would race with the cuEST kernels. Streams are compared
100    /// by their underlying CUDA stream and context (`CudaStream::eq`), not by
101    /// `Arc` identity — cudarc hands out fresh `Arc`s for e.g. the default
102    /// stream on each call.
103    pub(crate) fn check_stream(
104        &self,
105        stream: &Arc<CudaStream>,
106        arg: &str,
107    ) -> Result<(), CuestError> {
108        if self.stream.as_ref() == stream.as_ref() {
109            Ok(())
110        } else {
111            Err(CuestError::Rust(format!(
112                "`{arg}` is allocated on a different CUDA stream than the cuEST context; \
113                 allocate it on the context stream"
114            )))
115        }
116    }
117}
118
119impl Drop for ContextInner {
120    fn drop(&mut self) {
121        // cuEST requires the handle's device to be current when destroying.
122        // Drop cannot report errors; failures only leak GPU resources.
123        let _ = self.stream.context().bind_to_thread();
124        let _ = unsafe { ffi::cuestDestroy(self.handle) };
125    }
126}
127
128/// A cuEST context: the library handle bound to one CUDA stream.
129///
130/// This is the root object of the API; all other objects are created from it
131/// (directly or transitively) and keep it alive.
132pub struct CuestContext {
133    inner: Arc<ContextInner>,
134}
135
136impl CuestContext {
137    /// Create a context on `stream` with default handle parameters.
138    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, CuestError> {
139        Self::with_options(stream, &CuestHandleOptions::default())
140    }
141
142    /// Create a context on `stream` with explicit handle parameters.
143    ///
144    /// The stream's device is made current on the calling thread (the
145    /// `cudaSetDevice` equivalent) before the handle is created, as cuEST
146    /// requires.
147    pub fn with_options(
148        stream: &Arc<CudaStream>,
149        options: &CuestHandleOptions,
150    ) -> Result<Self, CuestError> {
151        stream.context().bind_to_thread()?;
152
153        let params = CuestParameters::create(ffi::CUEST_HANDLE_PARAMETERS)?;
154        let attr = |a: ffi::cuestHandleParametersAttributes_t| a.0 as c_int;
155        // Bind the stream to the handle; cuBLAS/cuSolver handles are left
156        // NULL so cuEST allocates its own on this stream.
157        params.configure_ptr(
158            attr(ffi::CUEST_HANDLE_PARAMETERS_CUDASTREAM),
159            stream.cu_stream() as *mut c_void,
160        )?;
161        if let Some(threads) = options.jit_compile_threads {
162            params
163                .configure_i32(attr(ffi::CUEST_HANDLE_PARAMETERS_JIT_COMPILE_THREADS), threads)?;
164        }
165        if let Some(value) = options.max_gauss_hermite {
166            params.configure_u64(attr(ffi::CUEST_HANDLE_PARAMETERS_MAX_GAUSS_HERMITE), value)?;
167        }
168        if let Some(value) = options.max_l_solid_harmonic {
169            params.configure_u64(attr(ffi::CUEST_HANDLE_PARAMETERS_MAX_L_SOLID_HARMONIC), value)?;
170        }
171        if let Some(value) = options.max_rys {
172            params.configure_u64(attr(ffi::CUEST_HANDLE_PARAMETERS_MAX_RYS), value)?;
173        }
174        let jit_cache_dir = options.jit_cache_dir.as_deref().map(cstr_from_path).transpose()?;
175        if let Some(dir) = &jit_cache_dir {
176            params.configure_cstr(attr(ffi::CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR), dir)?;
177        }
178
179        let mut handle: ffi::cuestHandle_t = ptr::null_mut();
180        status_check(unsafe { ffi::cuestCreate(params.raw(), &mut handle) }, "cuestCreate")?;
181
182        let version = match query_version(handle) {
183            Ok(version) => version,
184            Err(err) => {
185                let _ = unsafe { ffi::cuestDestroy(handle) };
186                return Err(err);
187            },
188        };
189        let needed = (ffi::CUEST_VER_MAJOR, ffi::CUEST_VER_MINOR, ffi::CUEST_VER_PATCH);
190        if (version.0, version.1) != (needed.0, needed.1) {
191            let _ = unsafe { ffi::cuestDestroy(handle) };
192            return Err(CuestError::UnsupportedVersion { needed, found: version });
193        }
194
195        Ok(Self {
196            inner: Arc::new(ContextInner {
197                handle,
198                stream: stream.clone(),
199                version,
200                _jit_cache_dir: jit_cache_dir,
201            }),
202        })
203    }
204
205    /// The CUDA stream this context's handle is bound to.
206    pub fn stream(&self) -> &Arc<CudaStream> {
207        self.inner.stream()
208    }
209
210    /// cuEST library version (major, minor, patch) reported at creation.
211    pub fn version(&self) -> (u32, u32, u32) {
212        self.inner.version()
213    }
214
215    /// The math mode of this handle.
216    pub fn math_mode(&self) -> Result<CuestMathMode, CuestError> {
217        let mut mode = CuestMathMode::Default.to_raw();
218        status_check(
219            unsafe { ffi::cuestGetMathMode(self.inner.handle(), &mut mode) },
220            "cuestGetMathMode",
221        )?;
222        CuestMathMode::from_raw(mode)
223    }
224
225    /// Set the math mode of this handle.
226    pub fn set_math_mode(&self, mode: CuestMathMode) -> Result<(), CuestError> {
227        status_check(
228            unsafe { ffi::cuestSetMathMode(self.inner.handle(), mode.to_raw()) },
229            "cuestSetMathMode",
230        )
231    }
232
233    /// The compute capability (major, minor) used for heuristics.
234    pub fn compute_capability_target(&self) -> Result<(u32, u32), CuestError> {
235        let (mut major, mut minor) = (0u32, 0u32);
236        status_check(
237            unsafe {
238                ffi::cuestGetComputeCapabilityTarget(self.inner.handle(), &mut major, &mut minor)
239            },
240            "cuestGetComputeCapabilityTarget",
241        )?;
242        Ok((major, minor))
243    }
244
245    /// Override the compute capability used for heuristics (not for kernel
246    /// execution or JIT compilation); see the cuEST documentation.
247    pub fn set_compute_capability_target(&self, major: u32, minor: u32) -> Result<(), CuestError> {
248        status_check(
249            unsafe { ffi::cuestSetComputeCapabilityTarget(self.inner.handle(), major, minor) },
250            "cuestSetComputeCapabilityTarget",
251        )
252    }
253
254    /// The raw cuEST handle (escape hatch for API surface not wrapped yet).
255    pub fn raw_handle(&self) -> ffi::cuestHandle_t {
256        self.inner.handle()
257    }
258
259    pub(crate) fn inner(&self) -> &Arc<ContextInner> {
260        &self.inner
261    }
262}
263
264impl Clone for CuestContext {
265    fn clone(&self) -> Self {
266        Self { inner: self.inner.clone() }
267    }
268}
269
270impl std::fmt::Debug for CuestContext {
271    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        f.debug_struct("CuestContext").field("version", &self.version()).finish_non_exhaustive()
273    }
274}
275
276// SAFETY: see the "Thread safety" section of the crate documentation: the
277// handle may be moved between threads, but a handle (and everything derived
278// from it) must be used from one thread at a time — cuEST documents no
279// internal synchronization.
280unsafe impl Send for CuestContext {}
281
282fn query_version(handle: ffi::cuestHandle_t) -> Result<(u32, u32, u32), CuestError> {
283    let (mut major, mut minor, mut patch) = (0u32, 0u32, 0u32);
284    status_check(unsafe { ffi::cuestGetMajorVersion(handle, &mut major) }, "cuestGetMajorVersion")?;
285    status_check(unsafe { ffi::cuestGetMinorVersion(handle, &mut minor) }, "cuestGetMinorVersion")?;
286    status_check(unsafe { ffi::cuestGetPatchVersion(handle, &mut patch) }, "cuestGetPatchVersion")?;
287    Ok((major, minor, patch))
288}
289
290fn cstr_from_path(path: &Path) -> Result<CString, CuestError> {
291    let string = path.to_str().ok_or_else(|| {
292        CuestError::Rust(format!("JIT cache directory path is not valid UTF-8: {}", path.display()))
293    })?;
294    CString::new(string).map_err(|_| {
295        CuestError::Rust(format!("JIT cache directory path contains NUL: {}", path.display()))
296    })
297}
cuest_cudarc/context.rs

cuest_cudarc/
context.rs
