1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
//! Shared host-side scaffolding for every cudarc-backed module under
//! `src/gpu/*` and `src/solver/gpu/*`.
//!
//! Before this module existed, each device backend (`bms_flex`,
//! `survival_flex`, `polya_gamma`, `reml_trace`, ...) carried its own
//! near-identical copy of two patterns:
//!
//! 1. A power-of-two bucketed free list of reusable f64 device slices
//! (the per-backend `DeviceArena`).
//! 2. A `OnceLock<Result<{module: Arc<CudaModule>}, GpuError>>` that
//! NVRTC-compiled one source string the first time the backend
//! dispatched and cached the resulting module for the process lifetime.
//!
//! Both are now provided here so every cudarc backend points at the same
//! implementation. The migration is atomic: no per-backend `DeviceArena`
//! type, no per-backend ad-hoc OnceLock, no transitional shim.
#[cfg(target_os = "linux")]
pub use linux::{DeviceArena, KeyedPtxModuleCache, PtxModuleCache, compile_ptx_arch};
#[cfg(target_os = "linux")]
mod linux {
use super::super::gpu_error::GpuError;
use crate::gpu_error::GpuResultExt;
use cudarc::driver::{CudaContext, CudaModule, CudaSlice, CudaStream};
use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex};
/// Power-of-two bucketed free list of f64 device slices.
///
/// Allocations round the requested element count up to the next
/// `usize::next_power_of_two`. On drop the slab is handed back to the
/// arena under the same bucket via [`DeviceArena::release`]. Held under
/// a `Mutex` by every backend that uses it because large-scale fits
/// dispatch from multiple rayon workers; the mutex is only held during
/// `alloc` / `release`, never across kernel launches.
#[derive(Default)]
pub struct DeviceArena {
free: HashMap<usize, Vec<CudaSlice<f64>>>,
}
impl DeviceArena {
#[inline]
pub fn bucket_of(elements: usize) -> usize {
elements.max(1).next_power_of_two()
}
/// Allocate a device slice of at least `elements` f64s. Returns the
/// bucket size actually allocated so the caller can release into the
/// same bucket on drop. `label` is woven into the error message if
/// the underlying `alloc_zeros` fails so failures stay attributable
/// to the originating backend (matching the pre-extraction wording).
pub fn alloc(
&mut self,
stream: &Arc<CudaStream>,
elements: usize,
label: &'static str,
) -> Result<(usize, CudaSlice<f64>), GpuError> {
let bucket = Self::bucket_of(elements);
if let Some(bucket_vec) = self.free.get_mut(&bucket)
&& let Some(slot) = bucket_vec.pop()
{
return Ok((bucket, slot));
}
let fresh = stream
.alloc_zeros::<f64>(bucket)
.gpu_ctx_with(|err| format!("{label} arena alloc_zeros<{bucket}>: {err}"))?;
Ok((bucket, fresh))
}
pub fn release(&mut self, bucket: usize, slab: CudaSlice<f64>) {
self.free.entry(bucket).or_default().push(slab);
}
}
/// Process-wide NVRTC module cache for a single PTX source string.
///
/// The first call to [`PtxModuleCache::get_or_compile`] compiles the
/// source via `cudarc::nvrtc::compile_ptx`, loads the module on the
/// supplied context, and stores the resulting `Arc<CudaModule>`.
/// Subsequent calls return the cached module without recompiling.
///
/// The `label` is woven into the error message so the originating
/// backend stays identifiable in logs; the wording matches each
/// caller's previous bespoke `format!` so existing log assertions
/// continue to hold.
#[derive(Default)]
pub struct PtxModuleCache {
module: std::sync::OnceLock<Arc<CudaModule>>,
}
impl PtxModuleCache {
pub const fn new() -> Self {
Self {
module: std::sync::OnceLock::new(),
}
}
pub fn get(&self) -> Option<&Arc<CudaModule>> {
self.module.get()
}
/// Compile `source` and load it on `ctx` the first time; return
/// the cached `Arc<CudaModule>` on every subsequent call.
pub fn get_or_compile(
&self,
ctx: &Arc<CudaContext>,
label: &'static str,
source: &str,
) -> Result<&Arc<CudaModule>, GpuError> {
self.get_or_load(ctx, label, || {
compile_ptx_with_opts(source, nvrtc_compile_options()?)
.map_err(|err| {
// The historical silent-CPU class: an NVRTC failure here
// is swallowed by callers' `.ok()?`; count it so telemetry
// can distinguish "device engaged" from "silently
// declined".
crate::profile::telemetry_record_cpu_fallback(format!(
"{label} NVRTC compile failed: {err}"
));
err
})
.gpu_ctx_with(|err| format!("{label} NVRTC compile failed: {err}"))
})
}
/// Load the PTX that `compile` produces on `ctx` the first time; return
/// the cached `Arc<CudaModule>` on every subsequent call. This is the
/// entry for a kernel whose NVRTC options differ from the shared ones
/// (a family that must disable FMA contraction, say); everything else
/// goes through [`Self::get_or_compile`].
pub fn get_or_load<F>(
&self,
ctx: &Arc<CudaContext>,
label: &'static str,
compile: F,
) -> Result<&Arc<CudaModule>, GpuError>
where
F: FnOnce() -> Result<cudarc::nvrtc::Ptx, GpuError>,
{
if let Some(existing) = self.module.get() {
return Ok(existing);
}
let ptx = compile()?;
let module = ctx
.load_module(ptx)
.gpu_ctx_with(|err| format!("{label} module load failed: {err}"))?;
if self.module.set(module).is_err() {
// A concurrent compile of the same label won the race; its
// module is already in the slot and is the one we return.
log::debug!("{label} module slot already populated by a concurrent compile");
}
Ok(self
.module
.get()
.expect("module slot populated immediately after set"))
}
}
/// A per-key family of compiled modules: one [`CudaModule`] per value of a
/// kernel-shape parameter that is baked into the source (a row width `P`,
/// say), each compiled with the shared device-keyed NVRTC options on first
/// use and shared thereafter. The keyed twin of [`PtxModuleCache`].
pub struct KeyedPtxModuleCache<K> {
modules: Mutex<HashMap<K, Arc<CudaModule>>>,
}
impl<K: Eq + std::hash::Hash + Copy + std::fmt::Display> KeyedPtxModuleCache<K> {
pub fn new() -> Self {
Self {
modules: Mutex::new(HashMap::new()),
}
}
/// The module for `key`, compiling `source(key)` and loading it on
/// `ctx` the first time that key is seen.
pub fn get_or_compile<S>(
&self,
ctx: &Arc<CudaContext>,
key: K,
label: &'static str,
source: S,
) -> Result<Arc<CudaModule>, GpuError>
where
S: FnOnce(K) -> String,
{
if let Ok(guard) = self.modules.lock() {
if let Some(module) = guard.get(&key) {
return Ok(Arc::clone(module));
}
}
let ptx = compile_ptx_arch(source(key))
.gpu_ctx_with(|err| format!("{label} NVRTC compile failed (key={key}): {err}"))?;
let module = ctx
.load_module(ptx)
.gpu_ctx_with(|err| format!("{label} module load failed (key={key}): {err}"))?;
if let Ok(mut guard) = self.modules.lock() {
// A concurrent compile of the same key may have won the race;
// its module is the one every later caller sees.
return Ok(Arc::clone(guard.entry(key).or_insert(module)));
}
Ok(module)
}
}
impl<K: Eq + std::hash::Hash + Copy + std::fmt::Display> Default for KeyedPtxModuleCache<K> {
fn default() -> Self {
Self::new()
}
}
/// Compile a kernel source string to PTX with the SAME device-keyed NVRTC
/// options [`PtxModuleCache::get_or_compile`] uses — crucially the
/// `--gpu-architecture` pin (#1551), without which NVRTC defaults below
/// `sm_60` and rejects `atomicAdd(double*, double)`. Call sites that compile
/// via the bare `cudarc::nvrtc::compile_ptx` (no options) MUST route through
/// this instead when their kernel uses double atomics, or the device path
/// silently falls back to the CPU.
pub fn compile_ptx_arch<S: AsRef<str>>(source: S) -> Result<cudarc::nvrtc::Ptx, GpuError> {
compile_ptx_with_opts(source.as_ref(), nvrtc_compile_options()?)
.gpu_ctx_with(|err| std::format!("NVRTC compile failed: {err}"))
}
fn nvrtc_compile_options() -> Result<CompileOptions, GpuError> {
let mut opts = CompileOptions::default();
opts.include_paths = nvrtc_include_paths();
// GPU↔CPU PARITY: disable FMA contraction. NVRTC's default is
// `--fmad=true`, which fuses `a*b + c` into a single fused multiply-add
// (ONE rounding). The CPU oracle computes `a*b` then `+ c` as two
// SEPARATELY-rounded f64 ops. For shallow kernels the gap is ~1 ULP;
// for deep derivative towers (the survival/SAE seeded jets, whose
// Hessian + contracted third/fourth channels chain dozens of mul/add
// steps) the per-op FMA divergence accumulates to ~5e-8 — enough to
// blow a 1e-9 parity gate on a real device (measured on a V100,
// compute 7.0: survival_rowjet device-vs-CPU max abs diff 5.09e-8).
// `--use_fast_math` was already off, but that does NOT imply fmad off
// (use_fast_math only ADDS fmad=true; the converse default is still
// on). Pinning fmad=false makes every shared-options kernel
// bit-comparable to the separately-rounded CPU path. `Option::None`
// would defer to NVRTC's `true` default, so we set it explicitly.
opts.fmad = Some(false);
// #1551: pin the NVRTC virtual arch to the selected device's compute
// capability. Without it NVRTC defaults below sm_60, where the
// `atomicAdd(double*, double)` overload is absent — so kernels using
// double atomics (the SAE arrow/Schur PCG kernels) fail to compile and
// the device path silently falls back to the CPU (SAE ran at 0% GPU).
// `arch` is `Option<&'static str>`; `nvrtc_arch()` returns a static
// `compute_NN` for the device's real capability.
if let Some(runtime) = crate::device_runtime::GpuRuntime::resolve(crate::global_policy())? {
opts.arch = Some(runtime.selected_device().capability.nvrtc_arch());
}
Ok(opts)
}
fn nvrtc_include_paths() -> Vec<String> {
let mut paths = Vec::new();
push_existing_include_path(&mut paths, Path::new("/usr/local/cuda/include"));
push_existing_include_path(&mut paths, Path::new("/usr/include"));
push_existing_include_path(&mut paths, Path::new("/usr/include/x86_64-linux-gnu"));
push_gcc_include_paths(&mut paths, Path::new("/usr/lib/gcc/x86_64-linux-gnu"));
paths
}
fn push_gcc_include_paths(paths: &mut Vec<String>, root: &Path) {
let Ok(entries) = std::fs::read_dir(root) else {
return;
};
for entry in entries.flatten() {
push_existing_include_path(paths, &entry.path().join("include"));
}
}
fn push_existing_include_path(paths: &mut Vec<String>, path: &Path) {
if !path.is_dir() {
return;
}
let display = path.to_string_lossy().into_owned();
if !paths.iter().any(|existing| existing == &display) {
paths.push(display);
}
}
}