gam_gpu/backend_probe.rs
1//! Shared CUDA backend-probe contract for every cudarc-backed module under
2//! `src/gpu/*`.
3//!
4//! Before this module existed, every GPU backend (`bms_flex`,
5//! `survival_flex`, `cubic_bspline_moments`, `cubic_cell`, `pirls_row`,
6//! `sphere`, ...) carried its own near-identical `probe_linux` prologue:
7//!
8//! 1. Resolve the process-wide [`crate::GpuRuntime`] losslessly. Typed hardware
9//! absence becomes a labelled `DriverLibraryUnavailable`; probe faults
10//! keep their original [`crate::GpuError`] variant.
11//! 2. Read the runtime's selected device ordinal.
12//! 3. Create (or reuse) the per-ordinal `CudaContext` or fail with a
13//! `DriverCallFailed { reason: "<module> backend: failed to create
14//! CUDA context for device N" }`.
15//! 4. Open the context's default `CudaStream`.
16//! 5. Carry the device's compute capability alongside the handles.
17//!
18//! Those five steps are identical apart from the per-module label that gets
19//! woven into the two error messages. Drift between copies meant error
20//! wording, capability handling, context reuse, and stream choice could
21//! diverge module to module. This module hosts the single contract: each
22//! backend now calls [`probe_cuda_backend`] with its label and keeps only
23//! its module caches and optional eager-compilation step.
24//!
25//! The migration is atomic: no backend re-implements the prologue, and
26//! there is no transitional shim.
27
28// `CudaBackendParts` is re-exported alongside the probe entry points: sibling
29// crates (`gam-terms`, `gam-models`, ...) call `probe_cuda_backend` and receive
30// a `CudaBackendParts` value (without ever naming the type), so `probe_cuda_backend`'s
31// public return type must itself be reachable or `-D warnings` rejects the leak
32// (`private_interfaces`). It carries no fields a caller can misuse out of context.
33#[cfg(target_os = "linux")]
34pub use linux::{
35 CudaBackendContext, CudaBackendParts, probe_backend_with_compile, probe_cuda_backend,
36};
37
38#[cfg(target_os = "linux")]
39mod linux {
40 use crate::device::GpuCapability;
41 use crate::device_cache::{DeviceArena, PtxModuleCache};
42 use crate::device_runtime::{GpuAvailabilityRef, GpuRuntime, cuda_context_for};
43 use crate::gpu_error::GpuError;
44 use cudarc::driver::{CudaContext, CudaStream};
45 use std::sync::{Arc, Mutex};
46
47 /// The handles every cudarc backend shares once the probe succeeds:
48 /// a context on the runtime's selected device, that context's default
49 /// stream, and the device's compute capability. Module-specific
50 /// backends layer their own caches and optional eager compilation on
51 /// top of these.
52 #[derive(Debug)]
53 pub struct CudaBackendParts {
54 pub ctx: Arc<CudaContext>,
55 pub stream: Arc<CudaStream>,
56 pub capability: GpuCapability,
57 }
58
59 /// Probe the process-wide CUDA backend for the calling module.
60 ///
61 /// Resolves the global [`GpuRuntime`], creates (or reuses) the
62 /// [`CudaContext`] for its selected device, opens that context's
63 /// default stream, and returns the trio bundled in [`CudaBackendParts`].
64 /// `label` names the calling module (e.g. `"bms_flex"`) and is woven
65 /// into both failure messages so the uniform contract still attributes
66 /// errors to their originating backend.
67 pub fn probe_cuda_backend(label: &'static str) -> Result<CudaBackendParts, GpuError> {
68 let runtime = match GpuRuntime::availability()? {
69 GpuAvailabilityRef::Available(runtime) => runtime,
70 GpuAvailabilityRef::Absent(reason) => {
71 return Err(GpuError::DriverLibraryUnavailable {
72 reason: format!("{label} backend: {reason}"),
73 });
74 }
75 };
76 let ordinal = runtime.selected_device().ordinal;
77 let ctx = cuda_context_for(ordinal).ok_or_else(|| {
78 gpu_err!("{label} backend: failed to create CUDA context for device {ordinal}")
79 })?;
80 let stream = ctx.default_stream();
81 let capability = runtime.selected_device().capability.clone();
82 Ok(CudaBackendParts {
83 ctx,
84 stream,
85 capability,
86 })
87 }
88
89 /// Probe the CUDA backend for `label` and run a backend-specific build
90 /// step on the resolved handles.
91 ///
92 /// This is [`probe_cuda_backend`] plus the one piece that genuinely
93 /// differs between backends: the NVRTC compile (and any per-backend cache
94 /// construction). The runtime resolution, context creation, and stream
95 /// selection — together with their uniform, label-attributed error
96 /// messages — live in the shared probe; `build` receives the resolved
97 /// [`CudaBackendParts`] (so it can clone the `Arc<CudaContext>` /
98 /// `Arc<CudaStream>` it needs) and returns the backend's own state `T`.
99 pub fn probe_backend_with_compile<F, T>(label: &'static str, build: F) -> Result<T, GpuError>
100 where
101 F: FnOnce(&CudaBackendParts) -> Result<T, GpuError>,
102 {
103 let parts = probe_cuda_backend(label)?;
104 build(&parts)
105 }
106
107 /// The process-wide device handles every cudarc backend stores after a
108 /// successful probe: the [`CudaContext`], its default [`CudaStream`], the
109 /// lazily NVRTC-compiled [`PtxModuleCache`], and the bucketed
110 /// [`DeviceArena`] of reusable f64 device buffers (held under a `Mutex`
111 /// because large-scale fits dispatch from multiple rayon worker threads; the
112 /// mutex is only held during `alloc` / `release`, not across kernel
113 /// launches). Module-specific backends (`bms_flex`, `survival_flex`, …)
114 /// wrap one of these as their `inner` context so the host-side
115 /// scaffolding (arena pooling, module cache, mutex around alloc) is
116 /// uniform instead of duplicated per backend.
117 pub struct CudaBackendContext {
118 pub ctx: Arc<CudaContext>,
119 pub stream: Arc<CudaStream>,
120 pub module: PtxModuleCache,
121 pub arena: Mutex<DeviceArena>,
122 }
123
124 impl CudaBackendContext {
125 /// Build the stored context from a fresh [`CudaBackendParts`] probe
126 /// result: adopt its context and stream, start an empty module cache
127 /// (the backend's eager-compile step fills it), and an empty device
128 /// arena. The probe's compute `capability` is consumed by the probe
129 /// path itself and is not retained here.
130 pub fn from_parts(parts: CudaBackendParts) -> Self {
131 CudaBackendContext {
132 ctx: parts.ctx,
133 stream: parts.stream,
134 module: PtxModuleCache::new(),
135 arena: Mutex::new(DeviceArena::default()),
136 }
137 }
138 }
139}
140
141#[cfg(all(test, target_os = "linux"))]
142mod tests {
143 use super::probe_cuda_backend;
144 use crate::device_runtime::{GpuAvailabilityRef, GpuRuntime};
145 use crate::gpu_error::GpuError;
146
147 /// Parity: every backend's probe must agree with the shared contract on
148 /// the same device. On a host with no CUDA runtime, the shared probe
149 /// must return the uniform `DriverLibraryUnavailable` carrying the
150 /// caller's label; on a host with a runtime, the probe must resolve the
151 /// *same* selected-device ordinal and compute capability the runtime
152 /// advertises, with a context bound to that ordinal and a usable
153 /// default stream. This is the regression guard that keeps the six
154 /// migrated backends (`bms_flex`, `survival_flex`,
155 /// `cubic_bspline_moments`, `cubic_cell`, `pirls_row`, `sphere`) routed
156 /// through one prologue instead of drifting copies.
157 #[test]
158 fn shared_probe_matches_runtime_device_and_labels_errors() {
159 match GpuRuntime::availability() {
160 Ok(GpuAvailabilityRef::Absent(absence)) => {
161 // No runtime: the shared probe must fail uniformly and
162 // attribute the failure to the supplied label.
163 match probe_cuda_backend("bms_flex") {
164 Err(GpuError::DriverLibraryUnavailable { reason }) => {
165 assert_eq!(
166 reason,
167 format!("bms_flex backend: {absence}"),
168 "shared probe must emit the uniform no-runtime message"
169 );
170 }
171 other => panic!(
172 "expected DriverLibraryUnavailable on a host without a CUDA runtime, \
173 got {other:?}"
174 ),
175 }
176 }
177 Ok(GpuAvailabilityRef::Available(runtime)) => {
178 // Runtime present: every label resolves the same selected
179 // device and the same compute capability the runtime
180 // advertises, and the context binds to that ordinal.
181 let expected_ordinal = runtime.selected_device().ordinal;
182 let expected_capability = &runtime.selected_device().capability;
183 for label in [
184 "bms_flex",
185 "survival_flex",
186 "cubic_bspline_moments",
187 "cubic_cell",
188 "pirls_row",
189 "sphere",
190 ] {
191 let parts = probe_cuda_backend(label)
192 .unwrap_or_else(|err| panic!("probe for {label} must succeed: {err:?}"));
193 assert_eq!(
194 parts.ctx.ordinal(),
195 expected_ordinal,
196 "{label}: context must bind the runtime's selected device ordinal"
197 );
198 assert_eq!(
199 &parts.capability, expected_capability,
200 "{label}: probe capability must match the runtime's selected device"
201 );
202 parts
203 .stream
204 .synchronize()
205 .unwrap_or_else(|err| panic!("{label}: default stream must sync: {err:?}"));
206 }
207 }
208 Err(error) => panic!("GPU probe fault must fail this backend contract test: {error}"),
209 }
210 }
211}