metrale_gpu_runtime/
registry.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: [`MetraleRegistry`]: one model's kernel modules, loaded into the
4//! process CUDA context, with raw driver-API helpers for launches, copies and
5//! device symbols.
6//!
7//! Owner: gpu-runtime.
8//! Invariants:
9//! - Each raw handle in `raw_modules` is a view of the cudarc module stored under
10//!   the same name in `modules`; `unload_raw` drains both maps together.
11//! - A registry owns neither the context nor the stream; both belong to the
12//!   process-scoped [`CudaHost`].
13
14use std::collections::HashMap;
15use std::ffi::{CString, c_void};
16use std::sync::{Arc, OnceLock};
17
18use cudarc::driver::{CudaContext, CudaFunction, CudaModule, CudaStream, LaunchConfig};
19use cudarc::nvrtc::Ptx;
20
21pub use crate::cuda_host::{CudaHost, host, release};
22use metrale_core::error::{MetraleError, Result};
23
24unsafe extern "C" {
25    fn cuModuleGetFunction(hfunc: *mut *mut c_void, hmod: *mut c_void, name: *const i8) -> i32;
26    fn cuLaunchKernel(
27        f: *mut c_void,
28        gridDimX: u32,
29        gridDimY: u32,
30        gridDimZ: u32,
31        blockDimX: u32,
32        blockDimY: u32,
33        blockDimZ: u32,
34        sharedMemBytes: u32,
35        hStream: *mut c_void,
36        kernelParams: *mut *mut c_void,
37        extra: *mut *mut c_void,
38    ) -> i32;
39    fn cuFuncSetAttribute(hfunc: *mut c_void, attrib: i32, value: i32) -> i32;
40    fn cuGetErrorName(error: i32, pStr: *mut *const i8) -> i32;
41    fn cuGetErrorString(error: i32, pStr: *mut *const i8) -> i32;
42    fn cuModuleGetGlobal_v2(
43        dptr: *mut u64,
44        bytes: *mut usize,
45        hmod: *mut c_void,
46        name: *const i8,
47    ) -> i32;
48    fn cuMemcpyHtoDAsync_v2(dst: u64, src: *const c_void, bytes: usize, stream: u64) -> i32;
49    fn cuMemcpyDtoHAsync_v2(dst: *mut c_void, src: u64, bytes: usize, stream: u64) -> i32;
50    fn cuStreamSynchronize(stream: u64) -> i32;
51}
52
53/// 2026-09-25: Format a CUresult as `"<NAME> (<status>): <description>"`, with
54/// `CUDA_UNKNOWN` and `(no message)` when the driver has no name or string for it.
55pub fn cuda_error_text(status: i32) -> String {
56    use std::ffi::CStr;
57    let mut name_ptr: *const i8 = std::ptr::null();
58    let mut msg_ptr: *const i8 = std::ptr::null();
59    let name = unsafe {
60        if cuGetErrorName(status, &mut name_ptr) == 0 && !name_ptr.is_null() {
61            CStr::from_ptr(name_ptr as *const std::os::raw::c_char)
62                .to_string_lossy()
63                .into_owned()
64        } else {
65            "CUDA_UNKNOWN".to_string()
66        }
67    };
68    let msg = unsafe {
69        if cuGetErrorString(status, &mut msg_ptr) == 0 && !msg_ptr.is_null() {
70            CStr::from_ptr(msg_ptr as *const std::os::raw::c_char)
71                .to_string_lossy()
72                .into_owned()
73        } else {
74            "(no message)".to_string()
75        }
76    };
77    format!("{name} ({status}): {msg}")
78}
79
80/// 2026-09-25: `CUDA_ERROR_DEINITIALIZED`: the driver has already shut down. The
81/// free paths treat it as nothing to do (see [`is_teardown_noop`]).
82pub const CUDA_ERROR_DEINITIALIZED: i32 = 4;
83
84/// 2026-09-25: Whether a CUresult means the context is already gone, so there is
85/// nothing to free: `CUDA_ERROR_DEINITIALIZED`, `CUDA_ERROR_INVALID_CONTEXT` (201)
86/// or `CUDA_ERROR_CONTEXT_IS_DESTROYED` (709).
87pub fn is_teardown_noop(status: i32) -> bool {
88    matches!(status, CUDA_ERROR_DEINITIALIZED | 201 | 709)
89}
90
91/// 2026-09-25: A raw `CUfunction` handle from [`MetraleRegistry::raw_function_cached`].
92#[derive(Clone, Copy)]
93pub struct RawCudaFunc(pub *mut c_void);
94// 2026-09-25: SAFETY: the handle is an opaque pointer that is only copied and
95// passed to the driver, never dereferenced in Rust. It is valid only while the
96// module it came from is loaded.
97unsafe impl Send for RawCudaFunc {}
98unsafe impl Sync for RawCudaFunc {}
99
100/// 2026-09-25: The kernel modules of one loaded model.
101///
102/// The blob set is the model's own (the caller passes the modules for this
103/// checkpoint). Obtain one with [`MetraleRegistry::load`] and pass the
104/// `Arc<MetraleRegistry>` along; there is no global accessor. Dropping the last
105/// handle drops the registry's module handles.
106pub struct MetraleRegistry {
107    host: Arc<CudaHost>,
108    modules: HashMap<&'static str, Arc<CudaModule>>,
109    /// 2026-09-25: Raw `CUmodule` views of `modules`, for `cuModuleGetFunction` and
110    /// `cuModuleGetGlobal_v2`.
111    raw_modules: HashMap<&'static str, *mut c_void>,
112}
113
114impl Drop for MetraleRegistry {
115    /// 2026-09-25: Drains both module maps (see `unload_raw`); cudarc unloads each
116    /// module when its last `Arc<CudaModule>` goes.
117    fn drop(&mut self) {
118        let failures = self.unload_raw();
119        if !failures.is_empty() {
120            eprintln!(
121                "metrale: {} module(s) failed to unload: {}",
122                failures.len(),
123                failures.join("; ")
124            );
125        }
126    }
127}
128
129// 2026-09-25: SAFETY: `raw_modules` holds raw pointers. They are written only in
130// `init` and in `unload_raw`, which takes `&mut self`; shared access only reads
131// them and passes them to the driver.
132unsafe impl Send for MetraleRegistry {}
133unsafe impl Sync for MetraleRegistry {}
134
135impl MetraleRegistry {
136    /// 2026-09-25: Load this model's kernel modules into the process CUDA context.
137    ///
138    /// Each call produces a fresh, independent module set; it shares only the
139    /// context and stream with a previously loaded model.
140    pub fn load(
141        ordinal: usize,
142        kernel_blobs: &[(&'static str, &'static [u8])],
143    ) -> Result<Arc<Self>> {
144        Ok(Arc::new(Self::init(host(ordinal)?, kernel_blobs)?))
145    }
146
147    /// 2026-09-25: The process CUDA host this registry's modules live in.
148    pub fn host(&self) -> &Arc<CudaHost> {
149        &self.host
150    }
151
152    pub fn ctx(&self) -> &Arc<CudaContext> {
153        &self.host.ctx
154    }
155
156    pub fn stream(&self) -> &Arc<CudaStream> {
157        &self.host.stream
158    }
159
160    /// 2026-09-25: Module names this registry loaded.
161    pub fn module_names(&self) -> impl Iterator<Item = &'static str> + '_ {
162        self.modules.keys().copied()
163    }
164
165    fn init(
166        host: Arc<CudaHost>,
167        kernel_blobs: &[(&'static str, &'static [u8])],
168    ) -> Result<MetraleRegistry> {
169        let ctx = &host.ctx;
170
171        let mut modules = HashMap::new();
172        let mut raw_modules = HashMap::new();
173        for &(name, blob) in kernel_blobs {
174            let is_binary = blob.starts_with(b"\x7fELF")
175                || blob.starts_with(b"__CLANG_OFFLOAD_BUNDLE__")
176                || std::str::from_utf8(&blob[..blob.len().min(64)]).is_err();
177
178            let ptx = if is_binary {
179                Ptx::from_binary(blob.to_vec())
180            } else {
181                let src = std::str::from_utf8(blob).map_err(|e| {
182                    MetraleError::ModuleLoad(format!("{name}: PTX not valid UTF-8: {e}"))
183                })?;
184                Ptx::from_src(src)
185            };
186            let module = ctx
187                .load_module(ptx)
188                .map_err(|e| MetraleError::ModuleLoad(format!("{name}: {e}")))?;
189
190            // 2026-09-25: The raw handle is a view of this same module, not a second
191            // load. It stays valid while `modules` holds the `Arc<CudaModule>` beside
192            // it; `unload_raw` drains both maps together.
193            raw_modules.insert(name, module.cu_module_raw() as *mut c_void);
194            modules.insert(name, module);
195        }
196
197        Ok(MetraleRegistry {
198            host,
199            modules,
200            raw_modules,
201        })
202    }
203
204    /// 2026-09-25: Look up `func_name` in `module_name` through cudarc. Not memoized;
205    /// see [`Self::function_cached`].
206    pub fn function(&self, module_name: &str, func_name: &str) -> Result<CudaFunction> {
207        let module = self.modules.get(module_name).ok_or_else(|| {
208            MetraleError::ModuleLoad(format!("Module '{module_name}' not loaded"))
209        })?;
210        module
211            .load_function(func_name)
212            .map_err(|e| MetraleError::ModuleLoad(format!("{module_name}::{func_name}: {e}")))
213    }
214
215    /// 2026-09-25: [`Self::function`], memoized in the caller's `OnceLock`.
216    pub fn function_cached(
217        &self,
218        cache: &OnceLock<CudaFunction>,
219        module_name: &str,
220        func_name: &str,
221    ) -> Result<CudaFunction> {
222        if let Some(f) = cache.get() {
223            return Ok(f.clone());
224        }
225        let func = self.function(module_name, func_name)?;
226        let _ = cache.set(func.clone());
227        Ok(func)
228    }
229
230    /// 2026-09-25: Whether a module of this name was loaded for this run.
231    pub fn has_module(&self, module_name: &str) -> bool {
232        self.raw_modules.contains_key(module_name)
233    }
234
235    /// 2026-09-25: Look up `func_name` through the raw driver API
236    /// (`cuModuleGetFunction`), memoized in the caller's `OnceLock`.
237    pub fn raw_function_cached(
238        &self,
239        cache: &OnceLock<RawCudaFunc>,
240        module_name: &str,
241        func_name: &str,
242    ) -> Result<RawCudaFunc> {
243        if let Some(f) = cache.get() {
244            return Ok(*f);
245        }
246        let raw_mod = self.raw_modules.get(module_name).ok_or_else(|| {
247            MetraleError::ModuleLoad(format!("Module '{module_name}' not loaded"))
248        })?;
249        let c_name = CString::new(func_name).map_err(|e| {
250            MetraleError::ModuleLoad(format!("{module_name}::{func_name}: CString: {e}"))
251        })?;
252        let mut func: *mut c_void = std::ptr::null_mut();
253        let status =
254            // 2026-09-25: `.cast()` because `c_char` is `i8` on x86_64 and `u8` on
255            // aarch64; `as *const i8` would be an `unnecessary_cast` on x86_64.
256            unsafe { cuModuleGetFunction(&mut func, *raw_mod, c_name.as_ptr().cast()) };
257        if status != 0 {
258            return Err(MetraleError::ModuleLoad(format!(
259                "{module_name}::{func_name}: cuModuleGetFunction failed: {}",
260                cuda_error_text(status)
261            )));
262        }
263        let raw = RawCudaFunc(func);
264        let _ = cache.set(raw);
265        Ok(raw)
266    }
267
268    /// 2026-09-25: Drain both module maps, raw views first, so no raw handle
269    /// outlives its module. Nothing is unloaded here: cudarc unloads each module
270    /// when its last `Arc<CudaModule>` drops. Idempotent; always returns an empty
271    /// list.
272    pub(crate) fn unload_raw(&mut self) -> Vec<String> {
273        self.raw_modules.drain().for_each(drop);
274        self.modules.drain().for_each(drop);
275        Vec::new()
276    }
277
278    /// 2026-09-25: The raw `CUstream` of the process stream held by [`CudaHost`].
279    pub fn raw_stream(&self) -> u64 {
280        self.host.stream.cu_stream() as u64
281    }
282
283    /// 2026-09-25: Resolve a `__device__` symbol of a loaded module to its device
284    /// pointer and byte length, for drivers that read or write device globals
285    /// without a kernel (the InnerQ calibration state). `symbol` is the
286    /// linker-visible name, so C++ namespace symbols are Itanium-mangled
287    /// (`_ZN7tq_plus14d_innerq_scaleE`).
288    pub fn device_symbol(&self, module_name: &str, symbol: &str) -> Result<(u64, usize)> {
289        let raw_mod = self.raw_modules.get(module_name).ok_or_else(|| {
290            MetraleError::ModuleLoad(format!("Module '{module_name}' not loaded"))
291        })?;
292        let c_sym = CString::new(symbol).map_err(|e| {
293            MetraleError::ModuleLoad(format!("{module_name}::{symbol}: CString: {e}"))
294        })?;
295        let mut dptr: u64 = 0;
296        let mut bytes: usize = 0;
297        let status =
298            unsafe { cuModuleGetGlobal_v2(&mut dptr, &mut bytes, *raw_mod, c_sym.as_ptr().cast()) };
299        if status != 0 {
300            return Err(MetraleError::ModuleLoad(format!(
301                "{module_name}::{symbol}: cuModuleGetGlobal_v2 failed: {}",
302                cuda_error_text(status)
303            )));
304        }
305        Ok((dptr, bytes))
306    }
307
308    /// 2026-09-25: Async H2D copy into a device pointer.
309    ///
310    /// # Safety
311    /// `dst` must be a valid device pointer, and the bytes at `src` must outlive
312    /// the copy (keep the host buffer until the next sync on `stream`).
313    pub unsafe fn copy_h2d_async(
314        &self,
315        dst: u64,
316        src: *const c_void,
317        bytes: usize,
318        stream: u64,
319    ) -> Result<()> {
320        let status = unsafe { cuMemcpyHtoDAsync_v2(dst, src, bytes, stream) };
321        if status != 0 {
322            return Err(MetraleError::KernelLaunch(format!(
323                "cuMemcpyHtoDAsync_v2 failed: {}",
324                cuda_error_text(status)
325            )));
326        }
327        Ok(())
328    }
329
330    /// 2026-09-25: Async D2H copy from a device pointer.
331    ///
332    /// # Safety
333    /// `dst` must stay alive until `stream` is synchronised.
334    pub unsafe fn copy_d2h_async(
335        &self,
336        dst: *mut c_void,
337        src: u64,
338        bytes: usize,
339        stream: u64,
340    ) -> Result<()> {
341        let status = unsafe { cuMemcpyDtoHAsync_v2(dst, src, bytes, stream) };
342        if status != 0 {
343            return Err(MetraleError::KernelLaunch(format!(
344                "cuMemcpyDtoHAsync_v2 failed: {}",
345                cuda_error_text(status)
346            )));
347        }
348        Ok(())
349    }
350
351    /// 2026-09-25: Block the calling thread until all prior work on `stream` completes.
352    pub fn stream_synchronize(&self, stream: u64) -> Result<()> {
353        let status = unsafe { cuStreamSynchronize(stream) };
354        if status != 0 {
355            return Err(MetraleError::KernelLaunch(format!(
356                "cuStreamSynchronize failed: {}",
357                cuda_error_text(status)
358            )));
359        }
360        Ok(())
361    }
362
363    /// 2026-09-25: Launch a kernel on the raw CUDA stream `stream_ptr`.
364    ///
365    /// # Safety
366    /// - `kernel_params` must contain valid pointers to arguments matching the kernel signature.
367    /// - `stream_ptr` must be a valid CUstream handle, or 0 for the legacy default stream.
368    /// - `raw_func` must be a valid CUfunction obtained from `raw_function_cached`.
369    pub unsafe fn launch_on_stream(
370        &self,
371        raw_func: RawCudaFunc,
372        cfg: LaunchConfig,
373        stream_ptr: u64,
374        kernel_params: &mut [*mut c_void],
375    ) -> Result<()> {
376        // 2026-09-25: The caller's stream is used as given; the registry's own
377        // stream is never substituted.
378        let stream = stream_ptr;
379        // 2026-09-25: Opt in to more than 48 KiB of dynamic shared memory when asked.
380        if cfg.shared_mem_bytes > 48 * 1024 {
381            const CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES: i32 = 8;
382            let attr_status = unsafe {
383                cuFuncSetAttribute(
384                    raw_func.0,
385                    CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
386                    cfg.shared_mem_bytes as i32,
387                )
388            };
389            if attr_status != 0 {
390                return Err(MetraleError::KernelLaunch(format!(
391                    "cuFuncSetAttribute(MAX_DYNAMIC_SHARED={}) failed: {}",
392                    cfg.shared_mem_bytes,
393                    cuda_error_text(attr_status)
394                )));
395            }
396        }
397        let status = unsafe {
398            cuLaunchKernel(
399                raw_func.0,
400                cfg.grid_dim.0,
401                cfg.grid_dim.1,
402                cfg.grid_dim.2,
403                cfg.block_dim.0,
404                cfg.block_dim.1,
405                cfg.block_dim.2,
406                cfg.shared_mem_bytes,
407                stream as *mut c_void,
408                kernel_params.as_mut_ptr(),
409                std::ptr::null_mut(),
410            )
411        };
412        if status != 0 {
413            return Err(MetraleError::KernelLaunch(format!(
414                "cuLaunchKernel failed: {} (grid=[{},{},{}], block=[{},{},{}], shared_mem={})",
415                cuda_error_text(status),
416                cfg.grid_dim.0,
417                cfg.grid_dim.1,
418                cfg.grid_dim.2,
419                cfg.block_dim.0,
420                cfg.block_dim.1,
421                cfg.block_dim.2,
422                cfg.shared_mem_bytes
423            )));
424        }
425        Ok(())
426    }
427}
metrale_gpu_runtime/registry.rs

metrale_gpu_runtime/
registry.rs
