1use std::mem::MaybeUninit;
5
6use crate::error::CudaError;
7use crate::ffi::*;
8use crate::load::SharedLibrary;
9
10macro_rules! api_call {
11 ($expr:expr_2021) => {{
12 let status = $expr;
13 if status != cudaError_enum::CUDA_SUCCESS {
14 return Err(CudaError::ErrCode { status });
15 }
16 }};
17}
18
19#[expect(non_snake_case)]
20pub struct LibCuda<'lib> {
21 cuInit: libloading::Symbol<'lib, extern "C" fn(::std::os::raw::c_uint) -> CUresult>,
22 cuDriverGetVersion:
23 libloading::Symbol<'lib, extern "C" fn(*mut ::std::os::raw::c_int) -> CUresult>,
24 cuDeviceGetCount:
25 libloading::Symbol<'lib, extern "C" fn(*mut ::std::os::raw::c_int) -> CUresult>,
26 cuDeviceGet:
27 libloading::Symbol<'lib, extern "C" fn(*mut CUdevice, ::std::os::raw::c_int) -> CUresult>,
28 cuDeviceGetName: libloading::Symbol<
29 'lib,
30 extern "C" fn(
31 name: *mut ::std::os::raw::c_char,
32 ::std::os::raw::c_int,
33 CUdevice,
34 ) -> CUresult,
35 >,
36 cuCtxCreate_v2: libloading::Symbol<
37 'lib,
38 extern "C" fn(*mut CUcontext, ::std::os::raw::c_uint, CUdevice) -> CUresult,
39 >,
40}
41
42impl LibCuda<'_> {
43 pub fn init(&self, flags: u32) -> Result<(), CudaError> {
44 api_call!((*self.cuInit)(flags));
45 Ok(())
46 }
47 pub fn driver_get_version(&self) -> Result<i32, CudaError> {
48 let mut value = 0;
49 api_call!((*self.cuDriverGetVersion)(&mut value));
50 Ok(value)
51 }
52 pub fn device_get_count(&self) -> Result<i32, CudaError> {
53 let mut value = 0;
54 api_call!((*self.cuDeviceGetCount)(&mut value));
55 Ok(value)
56 }
57 pub fn new_device(&self, i: i32) -> Result<CudaDevice<'_>, CudaError> {
58 let inner = MaybeUninit::zeroed();
59 let mut inner: CUdevice = unsafe { inner.assume_init() };
60 api_call!((*self.cuDeviceGet)(&mut inner, i));
61 Ok(CudaDevice {
62 parent: self,
63 inner,
64 })
65 }
66}
67
68pub struct CudaDevice<'a> {
69 parent: &'a LibCuda<'a>,
70 inner: CUdevice,
71}
72
73pub struct CudaContext<'a> {
74 _parent: &'a LibCuda<'a>,
75 inner: CUcontext,
76}
77
78impl CudaContext<'_> {
79 pub fn as_mut_void_ptr(&mut self) -> *mut std::ffi::c_void {
80 self.inner as *mut std::ffi::c_void
81 }
82}
83
84impl<'a> CudaDevice<'a> {
85 pub fn name(&self) -> Result<String, CudaError> {
86 const MAX_LEN: i32 = 255;
87 let value = std::ffi::CString::new(vec![b' '; MAX_LEN.try_into().unwrap()]).unwrap();
88 let raw = value.into_raw();
89 api_call!((*self.parent.cuDeviceGetName)(raw, MAX_LEN, self.inner));
90 let cs = unsafe { std::ffi::CString::from_raw(raw) };
92 let r = cs.into_string().unwrap();
93 Ok(r)
94 }
95 pub fn into_context(self) -> Result<CudaContext<'a>, CudaError> {
96 let context = MaybeUninit::zeroed();
97 let mut context: CUcontext = unsafe { context.assume_init() };
98 api_call!((*self.parent.cuCtxCreate_v2)(&mut context, 0, self.inner));
99 Ok(CudaContext {
100 _parent: self.parent,
101 inner: context,
102 })
103 }
104}
105
106macro_rules! get_func {
107 ($lib:expr_2021, $name:expr_2021) => {{
108 unsafe { $lib.library.get($name) }.map_err(|source| CudaError::NameFFIError {
109 name: String::from_utf8_lossy($name).to_string(),
110 source,
111 })?
112 }};
113}
114
115pub fn init(library: &SharedLibrary) -> Result<LibCuda<'_>, CudaError> {
116 let lib_cuda = LibCuda {
117 cuInit: get_func!(library, b"cuInit\0"),
118 cuDriverGetVersion: get_func!(library, b"cuDriverGetVersion\0"),
119 cuDeviceGetCount: get_func!(library, b"cuDeviceGetCount\0"),
120 cuDeviceGet: get_func!(library, b"cuDeviceGet\0"),
121 cuDeviceGetName: get_func!(library, b"cuDeviceGetName\0"),
122 cuCtxCreate_v2: get_func!(library, b"cuCtxCreate_v2\0"),
123 };
124
125 Ok(lib_cuda)
126}