teeny_cuda/device/
context.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::marker::PhantomData;
18
19use teeny_core::device::context::{Context, DeviceInfo};
20
21use crate::{
22    cuda,
23    device::{CudaDevice, CudaDeviceInfo},
24    errors::{Error, Result},
25};
26
27/// The CUDA [`Context`]: entry point for listing/opening devices.
28pub struct Cuda<'a> {
29    _unused: PhantomData<&'a ()>,
30}
31
32impl<'a> Cuda<'a> {
33    /// Initializes the CUDA driver. Errors if no CUDA-capable device is available.
34    pub fn try_new() -> Result<Self> {
35        Self::is_available().and_then(|is_available| {
36            if !is_available {
37                return Err(Error::CudaNotAvailable.into());
38            }
39
40            let flags = 0;
41            let status = unsafe { cuda::cuInit(flags) };
42            if status != cuda::cudaError_enum_CUDA_SUCCESS {
43                return Err(Error::from_cuda_error(status).into());
44            }
45
46            Ok(Self {
47                _unused: PhantomData,
48            })
49        })
50    }
51
52    /// Whether at least one CUDA-capable device is present.
53    pub fn is_available() -> Result<bool> {
54        let mut device_count = 0;
55        let err = unsafe { cuda::cudaGetDeviceCount(&mut device_count) };
56        if err != cuda::cudaError_enum_CUDA_SUCCESS {
57            return Err(Error::from_cuda_error(err).into());
58        }
59
60        Ok(device_count > 0)
61    }
62}
63
64impl<'a> Context<'a> for Cuda<'a> {
65    type Device = CudaDevice<'a>;
66    type DeviceInfo = CudaDeviceInfo;
67
68    fn list_devices(&self) -> Result<Vec<Self::DeviceInfo>> {
69        let mut devices = Vec::new();
70        let mut device_count = 0;
71        let err = unsafe { cuda::cudaGetDeviceCount(&mut device_count) };
72        if err != cuda::cudaError_enum_CUDA_SUCCESS {
73            return Err(Error::from_cuda_error(err).into());
74        }
75
76        for id in 0..device_count {
77            let mut props = cuda::cudaDeviceProp::default();
78            #[cfg(cuda_props_v2)]
79            let err = unsafe { cuda::cudaGetDeviceProperties_v2(&mut props, id) };
80            #[cfg(not(cuda_props_v2))]
81            let err = unsafe { cuda::cudaGetDeviceProperties(&mut props, id) };
82            if err != cuda::cudaError_enum_CUDA_SUCCESS {
83                return Err(Error::from_cuda_error(err).into());
84            }
85
86            let device_info = CudaDeviceInfo::new(id, props);
87            devices.push(device_info);
88        }
89
90        Ok(devices)
91    }
92
93    fn device(&self, id: &<Self::DeviceInfo as DeviceInfo>::Id) -> Result<Self::Device> {
94        CudaDevice::try_new(*id)
95    }
96}
teeny_cuda/device/context.rs

teeny_cuda/device/
context.rs
