Skip to main content

dynlink_cuda/
api.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use 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        // Note: on error we will leak the memory allocated in CString::new().
91        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}