j2k_cuda_runtime/
execution.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3pub(crate) mod completion;
4mod events;
5mod memory_ops;
6mod queued;
7
8pub(crate) use completion::{select_uncertain_completion_error, CudaSynchronizationOutcome};
9pub(crate) use events::elapsed_event_us_ceil;
10#[cfg(test)]
11pub(crate) use events::CudaEvent;
12pub use queued::{
13    CudaExecutionStats, CudaKernelBatchOutput, CudaKernelContiguousBatchOutput, CudaKernelOutput,
14    CudaPooledKernelOutput, CudaQueuedExecution,
15};
16
17#[cfg(test)]
18use crate::context::{CudaKernelModule, CudaKernelName};
19use crate::{
20    context::CudaContext,
21    driver::{CuDevicePtr, CuFunction},
22    error::CudaError,
23    kernels::{self, copy_u8_launch_geometry},
24    memory::CudaDeviceBuffer,
25};
26use std::{ffi::c_void, ops::Range};
27
28#[derive(Clone, Copy, Debug, Eq, PartialEq)]
29pub(crate) enum CudaLaunchMode {
30    Sync,
31    Async,
32}
33
34impl CudaLaunchMode {
35    pub(crate) fn from_synchronize(synchronize: bool) -> Self {
36        if synchronize {
37            Self::Sync
38        } else {
39            Self::Async
40        }
41    }
42}
43
44/// Marker for values that can be passed by address to a CUDA kernel launch.
45///
46/// # Safety
47///
48/// Implementors must have a stable, CUDA-compatible by-value representation for
49/// the duration of `cuLaunchKernel`.
50pub(crate) unsafe trait CudaKernelParam {}
51
52// SAFETY: `CuDevicePtr` is the raw CUDA device pointer value expected by kernels.
53unsafe impl CudaKernelParam for CuDevicePtr {}
54// SAFETY: CUDA kernels receive these scalar ABI types by value via parameter pointers.
55unsafe impl CudaKernelParam for u32 {}
56// SAFETY: CUDA kernels receive these scalar ABI types by value via parameter pointers.
57unsafe impl CudaKernelParam for i32 {}
58// SAFETY: CUDA kernels receive these scalar ABI types by value via parameter pointers.
59unsafe impl CudaKernelParam for f32 {}
60
61pub(crate) fn cuda_kernel_param<T>(value: &mut T) -> *mut c_void
62where
63    T: CudaKernelParam,
64{
65    std::ptr::from_mut(value).cast::<c_void>()
66}
67
68impl CudaContext {
69    #[doc(hidden)]
70    /// Copy host bytes through a CUDA copy kernel and return device output.
71    pub fn copy_with_kernel(&self, bytes: &[u8]) -> Result<CudaKernelOutput, CudaError> {
72        let staging = self.upload(bytes)?;
73        let output = self.copy_device_to_device_with_kernel(&staging)?;
74        let copy_dispatches = usize::from(!bytes.is_empty());
75        Ok(CudaKernelOutput {
76            buffer: output,
77            execution: CudaExecutionStats {
78                kernel_dispatches: copy_dispatches,
79                copy_kernel_dispatches: copy_dispatches,
80                decode_kernel_dispatches: 0,
81                hardware_decode: false,
82            },
83        })
84    }
85
86    #[cfg(all(test, feature = "cuda-oxide-copy-u8", j2k_cuda_oxide_copy_u8_built))]
87    pub(crate) fn copy_with_cuda_oxide_kernel(
88        &self,
89        bytes: &[u8],
90    ) -> Result<CudaKernelOutput, CudaError> {
91        let staging = self.upload(bytes)?;
92        let output = self.copy_device_to_device_with_cuda_oxide_kernel(&staging)?;
93        let copy_dispatches = usize::from(!bytes.is_empty());
94        Ok(CudaKernelOutput {
95            buffer: output,
96            execution: CudaExecutionStats {
97                kernel_dispatches: copy_dispatches,
98                copy_kernel_dispatches: copy_dispatches,
99                decode_kernel_dispatches: 0,
100                hardware_decode: false,
101            },
102        })
103    }
104
105    pub(crate) fn launch_kernel(
106        &self,
107        function: CuFunction,
108        geometry: kernels::CudaLaunchGeometry,
109        params: &mut [*mut c_void],
110    ) -> Result<(), CudaError> {
111        self.launch_kernel_async(function, geometry, params)?;
112        // SAFETY: `function` was loaded from a live module in this context, and
113        // the kernel was launched on the current context; synchronize waits for
114        // completion before callers inspect outputs.
115        self.synchronize()
116    }
117
118    pub(crate) fn launch_kernel_async(
119        &self,
120        function: CuFunction,
121        geometry: kernels::CudaLaunchGeometry,
122        params: &mut [*mut c_void],
123    ) -> Result<(), CudaError> {
124        if !geometry.is_valid() {
125            return Err(CudaError::InvalidArgument {
126                message: format!(
127                    "CUDA launch geometry exceeds static limits: grid {:?}, block {:?}",
128                    geometry.grid(),
129                    geometry.block()
130                ),
131            });
132        }
133        let (grid_x, grid_y, grid_z) = geometry.grid();
134        let (block_x, block_y, block_z) = geometry.block();
135        self.inner.with_current_resource_operation(|| {
136            // SAFETY: `function` was loaded from a live module in this context,
137            // and `params` contains kernel argument pointers valid for the
138            // launch call. The context lifetime gate is held and this context
139            // is current for the calling thread.
140            let launch_status = unsafe {
141                (self.inner.driver.cu_launch_kernel)(
142                    function,
143                    grid_x,
144                    grid_y,
145                    grid_z,
146                    block_x,
147                    block_y,
148                    block_z,
149                    0,
150                    std::ptr::null_mut(),
151                    params.as_mut_ptr(),
152                    std::ptr::null_mut(),
153                )
154            };
155            self.inner.driver.check("cuLaunchKernel", launch_status)
156        })
157    }
158
159    pub(crate) fn copy_device_to_device_with_kernel(
160        &self,
161        src: &CudaDeviceBuffer,
162    ) -> Result<CudaDeviceBuffer, CudaError> {
163        self.copy_device_range_to_device_with_kernel(src, 0..src.byte_len())
164    }
165
166    pub(crate) fn copy_device_range_to_device_with_kernel(
167        &self,
168        src: &CudaDeviceBuffer,
169        range: Range<usize>,
170    ) -> Result<CudaDeviceBuffer, CudaError> {
171        self.copy_device_range_to_device_with_copy_u8_loader(src, range, |context| {
172            context.inner.cuda_oxide_copy_u8_kernel_function()
173        })
174    }
175
176    #[cfg(all(test, feature = "cuda-oxide-copy-u8", j2k_cuda_oxide_copy_u8_built))]
177    pub(crate) fn copy_device_to_device_with_cuda_oxide_kernel(
178        &self,
179        src: &CudaDeviceBuffer,
180    ) -> Result<CudaDeviceBuffer, CudaError> {
181        self.copy_device_range_to_device_with_copy_u8_loader(src, 0..src.byte_len(), |context| {
182            context.inner.cuda_oxide_copy_u8_kernel_function()
183        })
184    }
185
186    fn copy_device_range_to_device_with_copy_u8_loader(
187        &self,
188        src: &CudaDeviceBuffer,
189        range: Range<usize>,
190        load_function: impl FnOnce(&Self) -> Result<CuFunction, CudaError>,
191    ) -> Result<CudaDeviceBuffer, CudaError> {
192        if !src.is_owned_by(self) {
193            return Err(CudaError::InvalidArgument {
194                message: "CUDA copy source must belong to the launch context".to_string(),
195            });
196        }
197        if range.start > range.end {
198            return Err(CudaError::InvalidArgument {
199                message: "CUDA copy range start must not exceed its end".to_string(),
200            });
201        }
202        if range.end > src.byte_len() {
203            return Err(CudaError::OutputTooSmall {
204                required: range.end,
205                have: src.byte_len(),
206            });
207        }
208        let byte_len = range.end - range.start;
209        if byte_len == 0 {
210            self.inner.set_current()?;
211            return self.allocate(0);
212        }
213        let geometry =
214            copy_u8_launch_geometry(byte_len).ok_or(CudaError::LengthTooLarge { len: byte_len })?;
215        self.inner.set_current()?;
216        let dst = self.allocate(byte_len)?;
217
218        let source_offset = u64::try_from(range.start)
219            .map_err(|_| CudaError::LengthTooLarge { len: range.start })?;
220        let src_ptr = src
221            .device_ptr()
222            .checked_add(source_offset)
223            .ok_or(CudaError::LengthTooLarge { len: range.end })?;
224        let function = load_function(self)?;
225        let mut dst_ptr = dst.device_ptr();
226        let mut src_ptr = src_ptr;
227        let mut len =
228            u64::try_from(byte_len).map_err(|_| CudaError::LengthTooLarge { len: byte_len })?;
229        let mut params = cuda_kernel_params!(dst_ptr, src_ptr, len);
230
231        self.launch_kernel(function, geometry, &mut params)?;
232
233        Ok(dst)
234    }
235
236    /// Synchronize all work submitted to this CUDA context.
237    pub fn synchronize(&self) -> Result<(), CudaError> {
238        self.synchronize_for_resource_release().into_result()
239    }
240
241    pub(crate) fn synchronize_for_resource_release(&self) -> CudaSynchronizationOutcome {
242        let result = self.inner.with_current_completion_operation(|| {
243            // SAFETY: the context lifetime gate is held and this CUDA context
244            // is current for the calling thread.
245            let status = unsafe { (self.inner.driver.cu_ctx_synchronize)() };
246            self.inner.driver.check("cuCtxSynchronize", status)
247        });
248        match result {
249            Ok(()) => CudaSynchronizationOutcome::Completed,
250            Err(error) => {
251                // The CUDA API may return both precondition failures and fatal
252                // asynchronous errors here. Neither is sufficient evidence
253                // that host-side resource release is safe.
254                CudaSynchronizationOutcome::CompletionUncertain(error)
255            }
256        }
257    }
258
259    /// Synchronize before returning `error`; if synchronization itself fails,
260    /// return that completion failure instead.
261    pub(crate) fn synchronize_then_error<T>(&self, error: CudaError) -> Result<T, CudaError> {
262        if self.inner.resource_lifetimes_poisoned() {
263            // A synchronous operation may already have surfaced the driver
264            // error that poisoned this context. Do not replace that primary
265            // diagnostic with the generic follow-up poison sentinel.
266            return Err(select_uncertain_completion_error(error, None));
267        }
268        match self.synchronize() {
269            Ok(()) => Err(error),
270            Err(completion_error) => Err(select_uncertain_completion_error(
271                error,
272                Some(completion_error),
273            )),
274        }
275    }
276
277    /// Preload a bundled CUDA kernel module and return its metadata handle.
278    #[cfg(test)]
279    pub(crate) fn preload_kernel_module(
280        &self,
281        kernel: CudaKernelName,
282    ) -> Result<CudaKernelModule, CudaError> {
283        let _ = self.inner.cuda_oxide_kernel_function(kernel.kernel())?;
284        Ok(CudaKernelModule {
285            entrypoint: kernel.entrypoint(),
286        })
287    }
288}
j2k_cuda_runtime/execution.rs

j2k_cuda_runtime/
execution.rs
