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

j2k_cuda_runtime/
execution.rs
