cuda-async 0.0.1-alpha

Safe Async CUDA support via Async Rust.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

use crate::device_context::with_deallocator_stream;
use crate::launch::{AsyncKernelLaunch, KernelArgument};
use cuda_core::free_async;
use cuda_core::sys::CUdeviceptr;
use std::marker::PhantomData;

#[derive(Debug, Copy, Clone)]
pub struct DevicePointer<T> {
    dtype: PhantomData<T>,
    pub dptr: CUdeviceptr,
}

unsafe impl<T> Send for DevicePointer<T> {}

impl<T> DevicePointer<T> {
    pub fn cu_deviceptr(&self) -> CUdeviceptr {
        self.dptr
    }
}

impl<T: Send + Sized> KernelArgument for DevicePointer<T> {
    fn push_arg(self, launcher: &mut AsyncKernelLaunch) {
        launcher.push_arg(Box::new(self.cu_deviceptr()));
    }
}

#[derive(Debug)]
pub struct DeviceBox<T: Send + ?Sized> {
    _device_id: usize,
    dtype: PhantomData<T>,
    cudptr: CUdeviceptr,
    len: usize,
}

unsafe impl<T: Send + ?Sized> Send for DeviceBox<T> {}
unsafe impl<T: Send + ?Sized> Sync for DeviceBox<T> {}

impl<T: Send + ?Sized> Drop for DeviceBox<T> {
    fn drop(&mut self) {
        unsafe {
            // Safety: The CUDA driver is guaranteed to complete any queued async operations.
            with_deallocator_stream(self._device_id, |stream| {
                free_async(self.cudptr, stream);
            })
            .expect(
                format!(
                    "Failed to free device pointer on device_id={}",
                    self._device_id
                )
                .as_str(),
            )
        }
    }
}

impl<DType: Send + Sized> DeviceBox<[DType]> {
    pub unsafe fn from_raw_parts(dptr: CUdeviceptr, len: usize, device_id: usize) -> Self {
        Self {
            dtype: PhantomData,
            cudptr: dptr,
            len,
            _device_id: device_id,
        }
    }
    pub fn len(&self) -> usize {
        self.len
    }
    pub fn cu_deviceptr(&self) -> CUdeviceptr {
        self.cudptr
    }
    pub fn device_id(&self) -> usize {
        self._device_id
    }
    pub fn device_pointer(&self) -> DevicePointer<DType> {
        DevicePointer::<DType> {
            dtype: PhantomData,
            dptr: self.cudptr,
        }
    }
}