xlog_cuda/
device.rs

1//! CUDA device management
2//!
3//! This module keeps XLOG's historical single-stream device abstraction while
4//! targeting cudarc's newer CUDA 13-capable context/stream APIs.
5
6use std::collections::BTreeMap;
7use std::ffi::{c_void, CString};
8use std::path::Path;
9use std::sync::{Arc, RwLock};
10
11use cudarc::driver::result::{self, DriverError};
12use cudarc::driver::{
13    sys, CudaContext as CudarcContext, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, DeviceRepr,
14    HostSlice, LaunchConfig, ValidAsZeroBits,
15};
16use cudarc::nvrtc::Ptx;
17use xlog_core::{Result, XlogError};
18
19#[derive(Debug)]
20struct LoadedModule {
21    cu_module: sys::CUmodule,
22    functions: BTreeMap<String, sys::CUfunction>,
23}
24
25unsafe impl Send for LoadedModule {}
26unsafe impl Sync for LoadedModule {}
27
28/// Kernel handle bound to XLOG's default CUDA stream.
29#[derive(Debug, Clone)]
30pub struct CudaFunction {
31    cu_function: sys::CUfunction,
32    context: Arc<CudarcContext>,
33    stream: Arc<CudaStream>,
34}
35
36impl CudaFunction {
37    pub(crate) unsafe fn launch_raw(
38        &self,
39        cfg: LaunchConfig,
40        params: &mut [*mut c_void],
41    ) -> std::result::Result<(), DriverError> {
42        self.context.bind_to_thread()?;
43        result::launch_kernel(
44            self.cu_function,
45            cfg.grid_dim,
46            cfg.block_dim,
47            cfg.shared_mem_bytes,
48            self.stream.cu_stream(),
49            params,
50        )
51    }
52
53    pub(crate) unsafe fn launch_raw_on_stream(
54        &self,
55        stream: &CudaStream,
56        cfg: LaunchConfig,
57        params: &mut [*mut c_void],
58    ) -> std::result::Result<(), DriverError> {
59        self.context.bind_to_thread()?;
60        result::launch_kernel(
61            self.cu_function,
62            cfg.grid_dim,
63            cfg.block_dim,
64            cfg.shared_mem_bytes,
65            stream.cu_stream(),
66            params,
67        )
68    }
69
70    pub(crate) unsafe fn launch_raw_cooperative(
71        &self,
72        cfg: LaunchConfig,
73        params: &mut [*mut c_void],
74    ) -> std::result::Result<(), DriverError> {
75        self.launch_raw_cooperative_on_stream(&self.stream, cfg, params)
76    }
77
78    pub(crate) unsafe fn launch_raw_cooperative_on_stream(
79        &self,
80        stream: &CudaStream,
81        cfg: LaunchConfig,
82        params: &mut [*mut c_void],
83    ) -> std::result::Result<(), DriverError> {
84        self.context.bind_to_thread()?;
85        result::launch_cooperative_kernel(
86            self.cu_function,
87            cfg.grid_dim,
88            cfg.block_dim,
89            cfg.shared_mem_bytes,
90            stream.cu_stream(),
91            params,
92        )
93    }
94
95    pub fn occupancy_available_dynamic_smem_per_block(
96        &self,
97        num_blocks: u32,
98        block_size: u32,
99    ) -> std::result::Result<usize, DriverError> {
100        let mut dynamic_smem_size: usize = 0;
101        unsafe {
102            sys::cuOccupancyAvailableDynamicSMemPerBlock(
103                &mut dynamic_smem_size,
104                self.cu_function,
105                num_blocks as std::ffi::c_int,
106                block_size as std::ffi::c_int,
107            )
108            .result()?
109        };
110        Ok(dynamic_smem_size)
111    }
112
113    pub fn occupancy_max_active_blocks_per_multiprocessor(
114        &self,
115        block_size: u32,
116        dynamic_smem_size: usize,
117        flags: Option<sys::CUoccupancy_flags_enum>,
118    ) -> std::result::Result<u32, DriverError> {
119        let mut num_blocks: std::ffi::c_int = 0;
120        let flags = flags.unwrap_or(sys::CUoccupancy_flags_enum::CU_OCCUPANCY_DEFAULT);
121        unsafe {
122            sys::cuOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
123                &mut num_blocks,
124                self.cu_function,
125                block_size as std::ffi::c_int,
126                dynamic_smem_size,
127                flags as std::ffi::c_uint,
128            )
129            .result()?
130        };
131        Ok(num_blocks as u32)
132    }
133
134    pub fn occupancy_max_active_clusters(
135        &self,
136        config: LaunchConfig,
137    ) -> std::result::Result<u32, DriverError> {
138        let mut num_clusters: std::ffi::c_int = 0;
139        let cfg = sys::CUlaunchConfig {
140            gridDimX: config.grid_dim.0,
141            gridDimY: config.grid_dim.1,
142            gridDimZ: config.grid_dim.2,
143            blockDimX: config.block_dim.0,
144            blockDimY: config.block_dim.1,
145            blockDimZ: config.block_dim.2,
146            sharedMemBytes: config.shared_mem_bytes,
147            hStream: self.stream.cu_stream(),
148            attrs: std::ptr::null_mut(),
149            numAttrs: 0,
150        };
151        unsafe {
152            sys::cuOccupancyMaxActiveClusters(&mut num_clusters, self.cu_function, &cfg).result()?
153        };
154        Ok(num_clusters as u32)
155    }
156
157    pub fn occupancy_max_potential_block_size(
158        &self,
159        block_size_to_dynamic_smem_size: extern "C" fn(block_size: std::ffi::c_int) -> usize,
160        dynamic_smem_size: usize,
161        block_size_limit: u32,
162        flags: Option<sys::CUoccupancy_flags_enum>,
163    ) -> std::result::Result<(u32, u32), DriverError> {
164        let mut min_grid_size: std::ffi::c_int = 0;
165        let mut block_size: std::ffi::c_int = 0;
166        let flags = flags.unwrap_or(sys::CUoccupancy_flags_enum::CU_OCCUPANCY_DEFAULT);
167        unsafe {
168            sys::cuOccupancyMaxPotentialBlockSizeWithFlags(
169                &mut min_grid_size,
170                &mut block_size,
171                self.cu_function,
172                Some(block_size_to_dynamic_smem_size),
173                dynamic_smem_size,
174                block_size_limit as std::ffi::c_int,
175                flags as std::ffi::c_uint,
176            )
177            .result()?
178        };
179        Ok((min_grid_size as u32, block_size as u32))
180    }
181
182    pub fn occupancy_max_potential_cluster_size(
183        &self,
184        config: LaunchConfig,
185    ) -> std::result::Result<u32, DriverError> {
186        let mut cluster_size: std::ffi::c_int = 0;
187        let cfg = sys::CUlaunchConfig {
188            gridDimX: config.grid_dim.0,
189            gridDimY: config.grid_dim.1,
190            gridDimZ: config.grid_dim.2,
191            blockDimX: config.block_dim.0,
192            blockDimY: config.block_dim.1,
193            blockDimZ: config.block_dim.2,
194            sharedMemBytes: config.shared_mem_bytes,
195            hStream: self.stream.cu_stream(),
196            attrs: std::ptr::null_mut(),
197            numAttrs: 0,
198        };
199        unsafe {
200            sys::cuOccupancyMaxPotentialClusterSize(&mut cluster_size, self.cu_function, &cfg)
201                .result()?
202        };
203        Ok(cluster_size as u32)
204    }
205
206    pub fn get_attribute(
207        &self,
208        attribute: sys::CUfunction_attribute_enum,
209    ) -> std::result::Result<i32, DriverError> {
210        self.context.bind_to_thread()?;
211        unsafe { result::function::get_function_attribute(self.cu_function, attribute) }
212    }
213
214    pub fn num_regs(&self) -> std::result::Result<i32, DriverError> {
215        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_NUM_REGS)
216    }
217
218    pub fn shared_size_bytes(&self) -> std::result::Result<i32, DriverError> {
219        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES)
220    }
221
222    pub fn const_size_bytes(&self) -> std::result::Result<i32, DriverError> {
223        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_CONST_SIZE_BYTES)
224    }
225
226    pub fn local_size_bytes(&self) -> std::result::Result<i32, DriverError> {
227        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES)
228    }
229
230    pub fn max_threads_per_block(&self) -> std::result::Result<i32, DriverError> {
231        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK)
232    }
233
234    pub fn ptx_version(&self) -> std::result::Result<i32, DriverError> {
235        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_PTX_VERSION)
236    }
237
238    pub fn binary_version(&self) -> std::result::Result<i32, DriverError> {
239        self.get_attribute(sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_BINARY_VERSION)
240    }
241
242    pub fn set_attribute(
243        &self,
244        attribute: sys::CUfunction_attribute_enum,
245        value: i32,
246    ) -> std::result::Result<(), DriverError> {
247        unsafe { result::function::set_function_attribute(self.cu_function, attribute, value) }
248    }
249
250    pub fn set_function_cache_config(
251        &self,
252        config: sys::CUfunc_cache,
253    ) -> std::result::Result<(), DriverError> {
254        unsafe { result::function::set_function_cache_config(self.cu_function, config) }
255    }
256}
257
258#[derive(Debug)]
259pub struct CudaDeviceInner {
260    context: Arc<CudarcContext>,
261    stream: Arc<CudaStream>,
262    modules: RwLock<BTreeMap<String, LoadedModule>>,
263}
264
265impl Drop for CudaDeviceInner {
266    fn drop(&mut self) {
267        let _ = self.context.bind_to_thread();
268        if let Ok(modules) = self.modules.get_mut() {
269            for module in modules.values() {
270                let _ = unsafe { result::module::unload(module.cu_module) };
271            }
272            modules.clear();
273        }
274    }
275}
276
277impl CudaDeviceInner {
278    fn insert_module(
279        &self,
280        module_name: &str,
281        cu_module: sys::CUmodule,
282        kernels: &[&str],
283    ) -> std::result::Result<(), DriverError> {
284        let mut functions = BTreeMap::new();
285        for &kernel in kernels {
286            let name_c = CString::new(kernel).unwrap();
287            let cu_function = unsafe { result::module::get_function(cu_module, name_c) }?;
288            functions.insert(kernel.to_string(), cu_function);
289        }
290        let module = LoadedModule {
291            cu_module,
292            functions,
293        };
294
295        let mut modules = self.modules.write().unwrap();
296        if let Some(prev) = modules.insert(module_name.to_string(), module) {
297            unsafe { result::module::unload(prev.cu_module) }?;
298        }
299        Ok(())
300    }
301
302    pub fn stream(&self) -> &Arc<CudaStream> {
303        &self.stream
304    }
305
306    pub fn has_func(&self, module_name: &str, func_name: &str) -> bool {
307        let modules = self.modules.read().unwrap();
308        modules
309            .get(module_name)
310            .is_some_and(|module| module.functions.contains_key(func_name))
311    }
312
313    pub fn get_func(&self, module_name: &str, func_name: &str) -> Option<CudaFunction> {
314        let modules = self.modules.read().unwrap();
315        let cu_function = modules
316            .get(module_name)
317            .and_then(|module| module.functions.get(func_name))
318            .copied()?;
319        Some(CudaFunction {
320            cu_function,
321            context: self.context.clone(),
322            stream: self.stream.clone(),
323        })
324    }
325
326    pub fn load_file(
327        &self,
328        path: &Path,
329        module_name: &str,
330        kernels: &[&str],
331    ) -> std::result::Result<(), DriverError> {
332        self.context.bind_to_thread()?;
333        let name_c = CString::new(path.to_string_lossy().as_bytes()).unwrap();
334        let cu_module = result::module::load(name_c)?;
335        self.insert_module(module_name, cu_module, kernels)
336    }
337
338    pub fn load_ptx(
339        &self,
340        ptx: Ptx,
341        module_name: &str,
342        kernels: &[&str],
343    ) -> std::result::Result<(), DriverError> {
344        self.context.bind_to_thread()?;
345        let cu_module = if let Some(bytes) = ptx.as_bytes() {
346            unsafe { result::module::load_data(bytes.as_ptr() as *const _) }?
347        } else {
348            let src = CString::new(ptx.to_src()).unwrap();
349            unsafe { result::module::load_data(src.as_ptr() as *const _) }?
350        };
351        self.insert_module(module_name, cu_module, kernels)
352    }
353
354    /// Allocate an uninitialized device slice on this device stream.
355    ///
356    /// # Safety
357    ///
358    /// The caller must initialize the returned allocation before any device or
359    /// host read observes its contents.
360    pub unsafe fn alloc<T: DeviceRepr>(
361        &self,
362        len: usize,
363    ) -> std::result::Result<CudaSlice<T>, DriverError> {
364        self.stream.alloc(len)
365    }
366
367    pub fn alloc_zeros<T: DeviceRepr + ValidAsZeroBits>(
368        &self,
369        len: usize,
370    ) -> std::result::Result<CudaSlice<T>, DriverError> {
371        self.stream.alloc_zeros(len)
372    }
373
374    pub fn memset_zeros<T: DeviceRepr + ValidAsZeroBits, Dst: DevicePtrMut<T>>(
375        &self,
376        dst: &mut Dst,
377    ) -> std::result::Result<(), DriverError> {
378        self.stream.memset_zeros(dst)?;
379        self.stream.synchronize()
380    }
381
382    pub fn htod_sync_copy_into<T: DeviceRepr, Dst: DevicePtrMut<T>, Src: HostSlice<T> + ?Sized>(
383        &self,
384        src: &Src,
385        dst: &mut Dst,
386    ) -> std::result::Result<(), DriverError> {
387        self.stream.memcpy_htod(src, dst)?;
388        self.stream.synchronize()
389    }
390
391    pub fn dtoh_sync_copy_into<T: DeviceRepr, Src: DevicePtr<T>, Dst: HostSlice<T> + ?Sized>(
392        &self,
393        src: &Src,
394        dst: &mut Dst,
395    ) -> std::result::Result<(), DriverError> {
396        self.stream.memcpy_dtoh(src, dst)?;
397        self.stream.synchronize()
398    }
399
400    pub fn htod_sync_copy<T: DeviceRepr, Src: HostSlice<T> + ?Sized>(
401        &self,
402        src: &Src,
403    ) -> std::result::Result<CudaSlice<T>, DriverError> {
404        let dst = self.stream.clone_htod(src)?;
405        self.stream.synchronize()?;
406        Ok(dst)
407    }
408
409    pub fn dtoh_sync_copy<T: DeviceRepr, Src: DevicePtr<T>>(
410        &self,
411        src: &Src,
412    ) -> std::result::Result<Vec<T>, DriverError> {
413        let dst = self.stream.clone_dtoh(src)?;
414        self.stream.synchronize()?;
415        Ok(dst)
416    }
417
418    pub fn dtod_copy<T, Src: DevicePtr<T>, Dst: DevicePtrMut<T>>(
419        &self,
420        src: &Src,
421        dst: &mut Dst,
422    ) -> std::result::Result<(), DriverError> {
423        self.stream.memcpy_dtod(src, dst)?;
424        self.stream.synchronize()
425    }
426
427    /// Enqueue a device-to-device copy without synchronizing the stream.
428    ///
429    /// Callers batching many copies must synchronize once after the last
430    /// enqueue and before reading any destination.
431    pub fn dtod_copy_async<T, Src: DevicePtr<T>, Dst: DevicePtrMut<T>>(
432        &self,
433        src: &Src,
434        dst: &mut Dst,
435    ) -> std::result::Result<(), DriverError> {
436        self.stream.memcpy_dtod(src, dst)
437    }
438
439    /// Wrap an existing CUDA device pointer in a typed cudarc slice.
440    ///
441    /// # Safety
442    ///
443    /// `cu_device_ptr` must point to a live allocation containing at least
444    /// `len * size_of::<T>()` bytes, and the resulting wrapper must not outlive
445    /// the allocation or alias another owner that will free it independently.
446    pub unsafe fn upgrade_device_ptr<T>(
447        &self,
448        cu_device_ptr: sys::CUdeviceptr,
449        len: usize,
450    ) -> CudaSlice<T> {
451        self.stream.upgrade_device_ptr(cu_device_ptr, len)
452    }
453
454    pub fn attribute(
455        &self,
456        attrib: sys::CUdevice_attribute,
457    ) -> std::result::Result<i32, DriverError> {
458        self.context.attribute(attrib)
459    }
460
461    pub fn synchronize(&self) -> std::result::Result<(), DriverError> {
462        self.stream.synchronize()
463    }
464
465    pub fn ordinal(&self) -> usize {
466        self.context.ordinal()
467    }
468}
469
470/// CUDA device wrapper for GPU operations.
471///
472/// This keeps XLOG's historical "device with a built-in default stream" API,
473/// but is backed by cudarc's newer `CudaContext` and `CudaStream`.
474pub struct CudaDevice {
475    device: Arc<CudaDeviceInner>,
476}
477
478impl CudaDevice {
479    /// Create a new CUDA device on the specified GPU ordinal.
480    pub fn new(ordinal: usize) -> Result<Self> {
481        let context = std::panic::catch_unwind(|| CudarcContext::new(ordinal))
482            .map_err(|_| {
483                XlogError::Kernel(format!(
484                    "Failed to create CUDA device {}: cudarc panicked during driver initialization",
485                    ordinal
486                ))
487            })?
488            .map_err(|e| {
489                XlogError::Kernel(format!("Failed to create CUDA device {}: {}", ordinal, e))
490            })?;
491
492        let stream = context.default_stream();
493        Ok(Self {
494            device: Arc::new(CudaDeviceInner {
495                context,
496                stream,
497                modules: RwLock::new(BTreeMap::new()),
498            }),
499        })
500    }
501
502    pub fn count() -> Result<i32> {
503        std::panic::catch_unwind(|| {
504            result::init()?;
505            result::device::get_count()
506        })
507        .map_err(|_| {
508            XlogError::Kernel(
509                "Failed to count CUDA devices: cudarc panicked during driver initialization"
510                    .to_string(),
511            )
512        })?
513        .map_err(|e| XlogError::Kernel(format!("Failed to count CUDA devices: {}", e)))
514    }
515
516    pub fn synchronize(&self) -> Result<()> {
517        self.device
518            .synchronize()
519            .map_err(|e| XlogError::Kernel(format!("Failed to synchronize device: {}", e)))
520    }
521
522    pub fn inner(&self) -> &Arc<CudaDeviceInner> {
523        &self.device
524    }
525
526    pub fn ordinal(&self) -> usize {
527        self.device.ordinal()
528    }
529}
530
531// Compile-time assertion: CudaDevice must be Send so pyxlog can use py.allow_threads().
532const _: () = {
533    fn _assert_send<T: Send>() {}
534    fn _check() {
535        _assert_send::<CudaDevice>();
536    }
537};
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542
543    #[test]
544    fn test_device_creation() {
545        let device = match CudaDevice::new(0) {
546            Ok(d) => d,
547            Err(e) => {
548                eprintln!("Skipping test: CUDA runtime unavailable: {}", e);
549                return;
550            }
551        };
552        drop(device);
553    }
554
555    #[test]
556    fn test_device_synchronize() {
557        let device = match CudaDevice::new(0) {
558            Ok(d) => d,
559            Err(e) => {
560                eprintln!("Skipping test: CUDA runtime unavailable: {}", e);
561                return;
562            }
563        };
564        let result = device.synchronize();
565        assert!(result.is_ok(), "Failed to synchronize: {:?}", result.err());
566    }
567
568    #[test]
569    fn test_device_ordinal() {
570        let device = match CudaDevice::new(0) {
571            Ok(d) => d,
572            Err(e) => {
573                eprintln!("Skipping test: CUDA runtime unavailable: {}", e);
574                return;
575            }
576        };
577        assert_eq!(device.ordinal(), 0);
578    }
579
580    #[test]
581    fn test_device_inner_access() {
582        let device = match CudaDevice::new(0) {
583            Ok(d) => d,
584            Err(e) => {
585                eprintln!("Skipping test: CUDA runtime unavailable: {}", e);
586                return;
587            }
588        };
589        let inner = device.inner();
590        assert_eq!(inner.ordinal(), 0);
591    }
592
593    #[test]
594    fn test_invalid_device_ordinal() {
595        let result = CudaDevice::new(9999);
596        assert!(result.is_err(), "Should fail with invalid ordinal");
597
598        if let Err(XlogError::Kernel(msg)) = result {
599            assert!(msg.contains("9999"), "Error should mention device ordinal");
600        } else {
601            panic!("Expected XlogError::Kernel");
602        }
603    }
604}
xlog_cuda/device.rs

xlog_cuda/
device.rs
