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 within an already admitted memory operation, retaining its capture pin.
187    ///
188    /// # Safety
189    /// Parameters must match the kernel ABI and all accessed memory must be
190    /// covered by `enqueue` or a live enclosing admission on the same stream.
191    unsafe fn launch_in(
192        self,
193        enqueue: &crate::launch::CudaEnqueue<'_>,
194        cfg: LaunchConfig,
195        params: Params,
196    ) -> Result<(), DriverError>;
197
198    /// Launch cooperatively within an already admitted memory operation.
199    ///
200    /// # Safety
201    /// The requirements of `launch_in` and CUDA cooperative launch apply.
202    unsafe fn launch_cooperative_in(
203        self,
204        enqueue: &crate::launch::CudaEnqueue<'_>,
205        cfg: LaunchConfig,
206        params: Params,
207    ) -> Result<(), DriverError>;
208    /// Launch a kernel on the function's default stream.
209    ///
210    /// # Safety
211    /// `params` must match the underlying CUDA kernel ABI exactly, and all
212    /// referenced device pointers must stay valid until the launch is enqueued.
213    unsafe fn launch(self, cfg: LaunchConfig, params: Params) -> Result<(), DriverError>;
214
215    /// Launch a kernel on an explicit CUDA stream.
216    ///
217    /// # Safety
218    /// The caller must uphold the same ABI and lifetime guarantees as `launch`
219    /// and must ensure `stream` is valid for the target device.
220    unsafe fn launch_on_stream(
221        self,
222        stream: &CudaStream,
223        cfg: LaunchConfig,
224        params: Params,
225    ) -> Result<(), DriverError>;
226
227    /// Launch a cooperative kernel.
228    ///
229    /// # Safety
230    /// The caller must uphold the same ABI and lifetime guarantees as `launch`
231    /// and must also ensure the kernel/configuration satisfies CUDA cooperative
232    /// launch requirements.
233    unsafe fn launch_cooperative(
234        self,
235        cfg: LaunchConfig,
236        params: Params,
237    ) -> Result<(), DriverError>;
238
239    /// Launch a cooperative kernel on an explicit CUDA stream.
240    ///
241    /// # Safety
242    /// The caller must uphold the same ABI, lifetime, and cooperative launch
243    /// guarantees as [`Self::launch_cooperative`] and ensure `stream` belongs
244    /// to the function's CUDA context.
245    unsafe fn launch_cooperative_on_stream(
246        self,
247        stream: &CudaStream,
248        cfg: LaunchConfig,
249        params: Params,
250    ) -> Result<(), DriverError>;
251}
252
253unsafe impl LaunchAsync<&mut [*mut c_void]> for CudaFunction {
254    unsafe fn launch_in(
255        self,
256        enqueue: &crate::launch::CudaEnqueue<'_>,
257        cfg: LaunchConfig,
258        params: &mut [*mut c_void],
259    ) -> Result<(), DriverError> {
260        self.launch_raw_in(enqueue, cfg, params, false)
261    }
262
263    unsafe fn launch_cooperative_in(
264        self,
265        enqueue: &crate::launch::CudaEnqueue<'_>,
266        cfg: LaunchConfig,
267        params: &mut [*mut c_void],
268    ) -> Result<(), DriverError> {
269        self.launch_raw_in(enqueue, cfg, params, true)
270    }
271    unsafe fn launch(
272        self,
273        cfg: LaunchConfig,
274        params: &mut [*mut c_void],
275    ) -> Result<(), DriverError> {
276        self.launch_raw(cfg, params)
277    }
278
279    unsafe fn launch_on_stream(
280        self,
281        stream: &CudaStream,
282        cfg: LaunchConfig,
283        params: &mut [*mut c_void],
284    ) -> Result<(), DriverError> {
285        self.launch_raw_on_stream(stream, cfg, params)
286    }
287
288    unsafe fn launch_cooperative(
289        self,
290        cfg: LaunchConfig,
291        params: &mut [*mut c_void],
292    ) -> Result<(), DriverError> {
293        self.launch_raw_cooperative(cfg, params)
294    }
295
296    unsafe fn launch_cooperative_on_stream(
297        self,
298        stream: &CudaStream,
299        cfg: LaunchConfig,
300        params: &mut [*mut c_void],
301    ) -> Result<(), DriverError> {
302        self.launch_raw_cooperative_on_stream(stream, cfg, params)
303    }
304}
305
306unsafe impl LaunchAsync<&mut Vec<*mut c_void>> for CudaFunction {
307    unsafe fn launch_in(
308        self,
309        enqueue: &crate::launch::CudaEnqueue<'_>,
310        cfg: LaunchConfig,
311        params: &mut Vec<*mut c_void>,
312    ) -> Result<(), DriverError> {
313        self.launch_raw_in(enqueue, cfg, params, false)
314    }
315
316    unsafe fn launch_cooperative_in(
317        self,
318        enqueue: &crate::launch::CudaEnqueue<'_>,
319        cfg: LaunchConfig,
320        params: &mut Vec<*mut c_void>,
321    ) -> Result<(), DriverError> {
322        self.launch_raw_in(enqueue, cfg, params, true)
323    }
324    unsafe fn launch(
325        self,
326        cfg: LaunchConfig,
327        params: &mut Vec<*mut c_void>,
328    ) -> Result<(), DriverError> {
329        self.launch_raw(cfg, params)
330    }
331
332    unsafe fn launch_on_stream(
333        self,
334        stream: &CudaStream,
335        cfg: LaunchConfig,
336        params: &mut Vec<*mut c_void>,
337    ) -> Result<(), DriverError> {
338        self.launch_raw_on_stream(stream, cfg, params)
339    }
340
341    unsafe fn launch_cooperative(
342        self,
343        cfg: LaunchConfig,
344        params: &mut Vec<*mut c_void>,
345    ) -> Result<(), DriverError> {
346        self.launch_raw_cooperative(cfg, params)
347    }
348
349    unsafe fn launch_cooperative_on_stream(
350        self,
351        stream: &CudaStream,
352        cfg: LaunchConfig,
353        params: &mut Vec<*mut c_void>,
354    ) -> Result<(), DriverError> {
355        self.launch_raw_cooperative_on_stream(stream, cfg, params)
356    }
357}
358
359macro_rules! impl_launch_tuple {
360    ([$($var:ident),*], [$($idx:tt),*]) => {
361        #[expect(
362            non_snake_case,
363            reason = "tuple type parameters are also destructured as bindings by this arity-generating macro"
364        )]
365        unsafe impl<$($var: IntoKernelParamStorage),*> LaunchAsync<($($var,)*)> for CudaFunction {
366            unsafe fn launch_in(self, enqueue: &crate::launch::CudaEnqueue<'_>, cfg: LaunchConfig, params: ($($var,)*)) -> Result<(), DriverError> {
367                let ($($var,)*) = params;
368                $(let $var = $var.into_kernel_param_storage();)*
369                let mut raw = [$( $var.as_kernel_param(), )*];
370                self.launch_raw_in(enqueue, cfg, &mut raw, false)
371            }
372
373            unsafe fn launch_cooperative_in(self, enqueue: &crate::launch::CudaEnqueue<'_>, cfg: LaunchConfig, params: ($($var,)*)) -> Result<(), DriverError> {
374                let ($($var,)*) = params;
375                $(let $var = $var.into_kernel_param_storage();)*
376                let mut raw = [$( $var.as_kernel_param(), )*];
377                self.launch_raw_in(enqueue, cfg, &mut raw, true)
378            }
379            unsafe fn launch(
380                self,
381                cfg: LaunchConfig,
382                params: ($($var,)*),
383            ) -> Result<(), DriverError> {
384                let ($($var,)*) = params;
385                $(let $var = $var.into_kernel_param_storage();)*
386                let mut raw = [$( $var.as_kernel_param(), )*];
387                self.launch_raw(cfg, &mut raw)
388            }
389
390            unsafe fn launch_on_stream(
391                self,
392                stream: &CudaStream,
393                cfg: LaunchConfig,
394                params: ($($var,)*),
395            ) -> Result<(), DriverError> {
396                let ($($var,)*) = params;
397                $(let $var = $var.into_kernel_param_storage();)*
398                let mut raw = [$( $var.as_kernel_param(), )*];
399                self.launch_raw_on_stream(stream, cfg, &mut raw)
400            }
401
402            unsafe fn launch_cooperative(
403                self,
404                cfg: LaunchConfig,
405                params: ($($var,)*),
406            ) -> Result<(), DriverError> {
407                let ($($var,)*) = params;
408                $(let $var = $var.into_kernel_param_storage();)*
409                let mut raw = [$( $var.as_kernel_param(), )*];
410                self.launch_raw_cooperative(cfg, &mut raw)
411            }
412
413            unsafe fn launch_cooperative_on_stream(
414                self,
415                stream: &CudaStream,
416                cfg: LaunchConfig,
417                params: ($($var,)*),
418            ) -> Result<(), DriverError> {
419                let ($($var,)*) = params;
420                $(let $var = $var.into_kernel_param_storage();)*
421                let mut raw = [$( $var.as_kernel_param(), )*];
422                self.launch_raw_cooperative_on_stream(stream, cfg, &mut raw)
423            }
424        }
425    };
426}
427
428impl_launch_tuple!([A], [0]);
429impl_launch_tuple!([A, B], [0, 1]);
430impl_launch_tuple!([A, B, C], [0, 1, 2]);
431impl_launch_tuple!([A, B, C, D], [0, 1, 2, 3]);
432impl_launch_tuple!([A, B, C, D, E], [0, 1, 2, 3, 4]);
433impl_launch_tuple!([A, B, C, D, E, F], [0, 1, 2, 3, 4, 5]);
434impl_launch_tuple!([A, B, C, D, E, F, G], [0, 1, 2, 3, 4, 5, 6]);
435impl_launch_tuple!([A, B, C, D, E, F, G, H], [0, 1, 2, 3, 4, 5, 6, 7]);
436impl_launch_tuple!([A, B, C, D, E, F, G, H, I], [0, 1, 2, 3, 4, 5, 6, 7, 8]);
437impl_launch_tuple!(
438    [A, B, C, D, E, F, G, H, I, J],
439    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
440);
441impl_launch_tuple!(
442    [A, B, C, D, E, F, G, H, I, J, K],
443    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
444);
445impl_launch_tuple!(
446    [A, B, C, D, E, F, G, H, I, J, K, L],
447    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
448);
449impl_launch_tuple!(
450    [A, B, C, D, E, F, G, H, I, J, K, L, M],
451    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
452);
xlog_cuda/cuda_compat.rs

xlog_cuda/
cuda_compat.rs
