xlog_cuda/
cuda_compat.rs

1use std::ffi::c_void;
2
3use cudarc::driver::{self, SyncOnDrop};
4
5pub use cudarc::driver::{
6    sys, CudaSlice, CudaStream, CudaView, CudaViewMut, DevicePtr, DevicePtrMut, DeviceRepr,
7    DeviceSlice, DriverError, LaunchConfig, ValidAsZeroBits,
8};
9
10pub use crate::device::CudaFunction;
11
12mod sealed {
13    pub trait KernelScalarSealed {}
14
15    macro_rules! impl_kernel_scalar_sealed {
16        ($($ty:ty),* $(,)?) => {
17            $(impl KernelScalarSealed for $ty {})*
18        };
19    }
20
21    impl_kernel_scalar_sealed!(
22        bool, i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize, f32, f64
23    );
24}
25
26/// Stable host-side storage for a kernel argument pointer.
27pub trait KernelParamStorage {
28    fn as_kernel_param(&self) -> *mut c_void;
29}
30
31#[derive(Debug)]
32pub struct ScalarParamStorage<T>(T);
33
34impl<T> KernelParamStorage for ScalarParamStorage<T> {
35    fn as_kernel_param(&self) -> *mut c_void {
36        (&self.0 as *const T).cast_mut().cast()
37    }
38}
39
40#[derive(Debug)]
41pub struct DeviceParamStorage<'a> {
42    ptr: driver::sys::CUdeviceptr,
43    _sync: Option<SyncOnDrop<'a>>,
44}
45
46impl<'a> DeviceParamStorage<'a> {
47    pub fn synced(ptr: driver::sys::CUdeviceptr, sync: SyncOnDrop<'a>) -> Self {
48        Self {
49            ptr,
50            _sync: Some(sync),
51        }
52    }
53
54    pub fn unsynced(ptr: driver::sys::CUdeviceptr) -> Self {
55        Self { ptr, _sync: None }
56    }
57}
58
59impl KernelParamStorage for DeviceParamStorage<'_> {
60    fn as_kernel_param(&self) -> *mut c_void {
61        (&self.ptr as *const driver::sys::CUdeviceptr)
62            .cast_mut()
63            .cast()
64    }
65}
66
67/// Backwards-compatible `as_kernel_param()` helper for manual raw launch lists.
68pub trait AsKernelParam {
69    fn as_kernel_param(&self) -> *mut c_void;
70}
71
72/// Convert a launch argument into storage that lives until `cuLaunchKernel` runs.
73pub trait IntoKernelParamStorage {
74    type Storage: KernelParamStorage;
75
76    fn into_kernel_param_storage(self) -> Self::Storage;
77}
78
79/// Scalar kernel parameters that can be copied directly into launch storage.
80pub trait KernelScalar:
81    sealed::KernelScalarSealed
82    + cudarc::driver::DeviceRepr
83    + Copy
84    + 'static
85    + AsKernelParam
86    + IntoKernelParamStorage
87{
88}
89
90macro_rules! impl_kernel_scalar {
91    ($($ty:ty),* $(,)?) => {
92        $(
93            impl KernelScalar for $ty {}
94
95            impl AsKernelParam for $ty {
96                fn as_kernel_param(&self) -> *mut c_void {
97                    (self as *const $ty).cast_mut().cast()
98                }
99            }
100
101            impl AsKernelParam for &$ty {
102                fn as_kernel_param(&self) -> *mut c_void {
103                    (*self as *const $ty).cast_mut().cast()
104                }
105            }
106
107            impl AsKernelParam for &mut $ty {
108                fn as_kernel_param(&self) -> *mut c_void {
109                    (*self as *const $ty).cast_mut().cast()
110                }
111            }
112
113            impl IntoKernelParamStorage for $ty {
114                type Storage = ScalarParamStorage<$ty>;
115
116                fn into_kernel_param_storage(self) -> Self::Storage {
117                    ScalarParamStorage(self)
118                }
119            }
120        )*
121    };
122}
123
124impl_kernel_scalar!(bool, i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize, f32, f64);
125
126impl<'a, T> IntoKernelParamStorage for &'a CudaSlice<T> {
127    type Storage = DeviceParamStorage<'a>;
128
129    fn into_kernel_param_storage(self) -> Self::Storage {
130        let stream = self.stream();
131        let (ptr, sync) = cudarc::driver::DevicePtr::device_ptr(self, stream);
132        DeviceParamStorage::synced(ptr, sync)
133    }
134}
135
136impl<T> IntoKernelParamStorage for &mut CudaSlice<T> {
137    type Storage = DeviceParamStorage<'static>;
138
139    fn into_kernel_param_storage(self) -> Self::Storage {
140        let stream = self.stream().clone();
141        let (ptr, sync) = cudarc::driver::DevicePtrMut::device_ptr_mut(self, &stream);
142        std::mem::forget(sync);
143        DeviceParamStorage::unsynced(ptr)
144    }
145}
146
147impl<'a, 'b, T> IntoKernelParamStorage for &'a CudaView<'b, T> {
148    type Storage = DeviceParamStorage<'a>;
149
150    fn into_kernel_param_storage(self) -> Self::Storage {
151        let stream = self.stream();
152        let (ptr, sync) = cudarc::driver::DevicePtr::device_ptr(self, stream);
153        DeviceParamStorage::synced(ptr, sync)
154    }
155}
156
157impl<'a, 'b, T> IntoKernelParamStorage for &'a CudaViewMut<'b, T> {
158    type Storage = DeviceParamStorage<'a>;
159
160    fn into_kernel_param_storage(self) -> Self::Storage {
161        let stream = self.stream();
162        let (ptr, sync) = cudarc::driver::DevicePtr::device_ptr(self, stream);
163        DeviceParamStorage::synced(ptr, sync)
164    }
165}
166
167impl<'a, 'b, T> IntoKernelParamStorage for &'a mut CudaViewMut<'b, T> {
168    type Storage = DeviceParamStorage<'static>;
169
170    fn into_kernel_param_storage(self) -> Self::Storage {
171        let stream = self.stream().clone();
172        let (ptr, sync) = cudarc::driver::DevicePtrMut::device_ptr_mut(self, &stream);
173        std::mem::forget(sync);
174        DeviceParamStorage::unsynced(ptr)
175    }
176}
177
178/// Old cudarc-style launch trait reimplemented on top of CUDA 13-compatible
179/// raw kernel launches.
180///
181/// # Safety
182/// Implementors must preserve CUDA's launch semantics and must not let kernel
183/// parameter storage or referenced device memory expire before the launch is
184/// enqueued on the target stream.
185pub unsafe trait LaunchAsync<Params> {
186    /// Launch a kernel on the function's default stream.
187    ///
188    /// # Safety
189    /// `params` must match the underlying CUDA kernel ABI exactly, and all
190    /// referenced device pointers must stay valid until the launch is enqueued.
191    unsafe fn launch(self, cfg: LaunchConfig, params: Params) -> Result<(), DriverError>;
192
193    /// Launch a kernel on an explicit CUDA stream.
194    ///
195    /// # Safety
196    /// The caller must uphold the same ABI and lifetime guarantees as `launch`
197    /// and must ensure `stream` is valid for the target device.
198    unsafe fn launch_on_stream(
199        self,
200        stream: &CudaStream,
201        cfg: LaunchConfig,
202        params: Params,
203    ) -> Result<(), DriverError>;
204
205    /// Launch a cooperative kernel.
206    ///
207    /// # Safety
208    /// The caller must uphold the same ABI and lifetime guarantees as `launch`
209    /// and must also ensure the kernel/configuration satisfies CUDA cooperative
210    /// launch requirements.
211    unsafe fn launch_cooperative(
212        self,
213        cfg: LaunchConfig,
214        params: Params,
215    ) -> Result<(), DriverError>;
216
217    /// Launch a cooperative kernel on an explicit CUDA stream.
218    ///
219    /// # Safety
220    /// The caller must uphold the same ABI, lifetime, and cooperative launch
221    /// guarantees as [`Self::launch_cooperative`] and ensure `stream` belongs
222    /// to the function's CUDA context.
223    unsafe fn launch_cooperative_on_stream(
224        self,
225        stream: &CudaStream,
226        cfg: LaunchConfig,
227        params: Params,
228    ) -> Result<(), DriverError>;
229}
230
231unsafe impl LaunchAsync<&mut [*mut c_void]> for CudaFunction {
232    unsafe fn launch(
233        self,
234        cfg: LaunchConfig,
235        params: &mut [*mut c_void],
236    ) -> Result<(), DriverError> {
237        self.launch_raw(cfg, params)
238    }
239
240    unsafe fn launch_on_stream(
241        self,
242        stream: &CudaStream,
243        cfg: LaunchConfig,
244        params: &mut [*mut c_void],
245    ) -> Result<(), DriverError> {
246        self.launch_raw_on_stream(stream, cfg, params)
247    }
248
249    unsafe fn launch_cooperative(
250        self,
251        cfg: LaunchConfig,
252        params: &mut [*mut c_void],
253    ) -> Result<(), DriverError> {
254        self.launch_raw_cooperative(cfg, params)
255    }
256
257    unsafe fn launch_cooperative_on_stream(
258        self,
259        stream: &CudaStream,
260        cfg: LaunchConfig,
261        params: &mut [*mut c_void],
262    ) -> Result<(), DriverError> {
263        self.launch_raw_cooperative_on_stream(stream, cfg, params)
264    }
265}
266
267unsafe impl LaunchAsync<&mut Vec<*mut c_void>> for CudaFunction {
268    unsafe fn launch(
269        self,
270        cfg: LaunchConfig,
271        params: &mut Vec<*mut c_void>,
272    ) -> Result<(), DriverError> {
273        self.launch_raw(cfg, params)
274    }
275
276    unsafe fn launch_on_stream(
277        self,
278        stream: &CudaStream,
279        cfg: LaunchConfig,
280        params: &mut Vec<*mut c_void>,
281    ) -> Result<(), DriverError> {
282        self.launch_raw_on_stream(stream, cfg, params)
283    }
284
285    unsafe fn launch_cooperative(
286        self,
287        cfg: LaunchConfig,
288        params: &mut Vec<*mut c_void>,
289    ) -> Result<(), DriverError> {
290        self.launch_raw_cooperative(cfg, params)
291    }
292
293    unsafe fn launch_cooperative_on_stream(
294        self,
295        stream: &CudaStream,
296        cfg: LaunchConfig,
297        params: &mut Vec<*mut c_void>,
298    ) -> Result<(), DriverError> {
299        self.launch_raw_cooperative_on_stream(stream, cfg, params)
300    }
301}
302
303macro_rules! impl_launch_tuple {
304    ([$($var:ident),*], [$($idx:tt),*]) => {
305        #[allow(non_snake_case)]
306        unsafe impl<$($var: IntoKernelParamStorage),*> LaunchAsync<($($var,)*)> for CudaFunction {
307            unsafe fn launch(
308                self,
309                cfg: LaunchConfig,
310                params: ($($var,)*),
311            ) -> Result<(), DriverError> {
312                let ($($var,)*) = params;
313                $(let $var = $var.into_kernel_param_storage();)*
314                let mut raw = [$( $var.as_kernel_param(), )*];
315                self.launch_raw(cfg, &mut raw)
316            }
317
318            unsafe fn launch_on_stream(
319                self,
320                stream: &CudaStream,
321                cfg: LaunchConfig,
322                params: ($($var,)*),
323            ) -> Result<(), DriverError> {
324                let ($($var,)*) = params;
325                $(let $var = $var.into_kernel_param_storage();)*
326                let mut raw = [$( $var.as_kernel_param(), )*];
327                self.launch_raw_on_stream(stream, cfg, &mut raw)
328            }
329
330            unsafe fn launch_cooperative(
331                self,
332                cfg: LaunchConfig,
333                params: ($($var,)*),
334            ) -> Result<(), DriverError> {
335                let ($($var,)*) = params;
336                $(let $var = $var.into_kernel_param_storage();)*
337                let mut raw = [$( $var.as_kernel_param(), )*];
338                self.launch_raw_cooperative(cfg, &mut raw)
339            }
340
341            unsafe fn launch_cooperative_on_stream(
342                self,
343                stream: &CudaStream,
344                cfg: LaunchConfig,
345                params: ($($var,)*),
346            ) -> Result<(), DriverError> {
347                let ($($var,)*) = params;
348                $(let $var = $var.into_kernel_param_storage();)*
349                let mut raw = [$( $var.as_kernel_param(), )*];
350                self.launch_raw_cooperative_on_stream(stream, cfg, &mut raw)
351            }
352        }
353    };
354}
355
356impl_launch_tuple!([A], [0]);
357impl_launch_tuple!([A, B], [0, 1]);
358impl_launch_tuple!([A, B, C], [0, 1, 2]);
359impl_launch_tuple!([A, B, C, D], [0, 1, 2, 3]);
360impl_launch_tuple!([A, B, C, D, E], [0, 1, 2, 3, 4]);
361impl_launch_tuple!([A, B, C, D, E, F], [0, 1, 2, 3, 4, 5]);
362impl_launch_tuple!([A, B, C, D, E, F, G], [0, 1, 2, 3, 4, 5, 6]);
363impl_launch_tuple!([A, B, C, D, E, F, G, H], [0, 1, 2, 3, 4, 5, 6, 7]);
364impl_launch_tuple!([A, B, C, D, E, F, G, H, I], [0, 1, 2, 3, 4, 5, 6, 7, 8]);
365impl_launch_tuple!(
366    [A, B, C, D, E, F, G, H, I, J],
367    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
368);
369impl_launch_tuple!(
370    [A, B, C, D, E, F, G, H, I, J, K],
371    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
372);
373impl_launch_tuple!(
374    [A, B, C, D, E, F, G, H, I, J, K, L],
375    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
376);
377impl_launch_tuple!(
378    [A, B, C, D, E, F, G, H, I, J, K, L, M],
379    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
380);
xlog_cuda/cuda_compat.rs

xlog_cuda/
cuda_compat.rs
