|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! Kernel launches with stream-ordered device memory arguments. |
| 5 | +
|
| 6 | +use std::ffi::c_void; |
| 7 | + |
| 8 | +use crate::memory::AccessKind; |
| 9 | +use crate::memory::ArgAccess; |
| 10 | +use crate::memory::CudaSlice; |
| 11 | +use crate::memory::CudaView; |
| 12 | +use crate::memory::CudaViewMut; |
| 13 | +use crate::module::CudaFunction; |
| 14 | +use crate::repr::DeviceRepr; |
| 15 | +use crate::result::DriverError; |
| 16 | +use crate::stream::CudaEvent; |
| 17 | +use crate::stream::CudaStream; |
| 18 | +use crate::sys; |
| 19 | + |
| 20 | +/// Grid, block and dynamic shared memory of a launch. |
| 21 | +#[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| 22 | +pub struct LaunchConfig { |
| 23 | + /// Blocks per grid dimension. |
| 24 | + pub grid_dim: (u32, u32, u32), |
| 25 | + /// Threads per block dimension. |
| 26 | + pub block_dim: (u32, u32, u32), |
| 27 | + /// Dynamic shared memory per block, in bytes. |
| 28 | + pub shared_mem_bytes: u32, |
| 29 | +} |
| 30 | + |
| 31 | +impl LaunchConfig { |
| 32 | + /// One-dimensional blocks of 1024 threads covering `n` elements. |
| 33 | + pub fn for_num_elems(n: u32) -> Self { |
| 34 | + Self { |
| 35 | + grid_dim: (n.div_ceil(1024), 1, 1), |
| 36 | + block_dim: (1024, 1, 1), |
| 37 | + shared_mem_bytes: 0, |
| 38 | + } |
| 39 | + } |
| 40 | +} |
| 41 | + |
| 42 | +/// A kernel launch being assembled on a stream; create with [`CudaStream::launch_builder`]. |
| 43 | +/// |
| 44 | +/// Device memory arguments are ordered after earlier conflicting accesses and recorded as |
| 45 | +/// accesses of this launch. |
| 46 | +pub struct LaunchArgs<'a> { |
| 47 | + stream: &'a CudaStream, |
| 48 | + function: &'a CudaFunction, |
| 49 | + // `Vec`, unlike `SmallVec`, lets argument borrows end before the builder drops. |
| 50 | + args: Vec<*mut c_void>, |
| 51 | + accesses: Vec<ArgAccess<'a>>, |
| 52 | + timing: Option<sys::CUevent_flags>, |
| 53 | +} |
| 54 | + |
| 55 | +impl std::fmt::Debug for LaunchArgs<'_> { |
| 56 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 57 | + f.debug_struct("LaunchArgs") |
| 58 | + .field("args", &self.args.len()) |
| 59 | + .finish() |
| 60 | + } |
| 61 | +} |
| 62 | + |
| 63 | +impl CudaStream { |
| 64 | + /// Starts a launch of `function` on this stream. |
| 65 | + pub fn launch_builder<'a>(&'a self, function: &'a CudaFunction) -> LaunchArgs<'a> { |
| 66 | + LaunchArgs { |
| 67 | + stream: self, |
| 68 | + function, |
| 69 | + args: Vec::with_capacity(8), |
| 70 | + accesses: Vec::with_capacity(4), |
| 71 | + timing: None, |
| 72 | + } |
| 73 | + } |
| 74 | +} |
| 75 | + |
| 76 | +/// A value that can be passed as a kernel argument. |
| 77 | +/// |
| 78 | +/// # Safety |
| 79 | +/// |
| 80 | +/// Implementations must push a pointer to storage that lives as long as the [`LaunchArgs`] and |
| 81 | +/// holds the argument in the kernel's expected representation. |
| 82 | +pub unsafe trait PushKernelArg<T> { |
| 83 | + /// Appends the argument. |
| 84 | + fn arg(&mut self, arg: T) -> &mut Self; |
| 85 | +} |
| 86 | + |
| 87 | +// SAFETY: The reference outlives the launch and `T` has a device representation. |
| 88 | +unsafe impl<'a, 'b: 'a, T: DeviceRepr> PushKernelArg<&'b T> for LaunchArgs<'a> { |
| 89 | + #[inline] |
| 90 | + fn arg(&mut self, arg: &'b T) -> &mut Self { |
| 91 | + self.args.push((arg as *const T).cast_mut().cast()); |
| 92 | + self |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +// SAFETY: The device pointer lives in the slice, which outlives the launch. |
| 97 | +unsafe impl<'a, 'b: 'a, T> PushKernelArg<&'b CudaSlice<T>> for LaunchArgs<'a> { |
| 98 | + #[inline] |
| 99 | + fn arg(&mut self, arg: &'b CudaSlice<T>) -> &mut Self { |
| 100 | + let (access, pointer) = arg.arg_access(AccessKind::Read); |
| 101 | + self.accesses.push(access); |
| 102 | + self.args |
| 103 | + .push((pointer as *const sys::CUdeviceptr).cast_mut().cast()); |
| 104 | + self |
| 105 | + } |
| 106 | +} |
| 107 | + |
| 108 | +// SAFETY: The device pointer lives in the slice, which outlives the launch. |
| 109 | +unsafe impl<'a, 'b: 'a, T> PushKernelArg<&'b mut CudaSlice<T>> for LaunchArgs<'a> { |
| 110 | + #[inline] |
| 111 | + fn arg(&mut self, arg: &'b mut CudaSlice<T>) -> &mut Self { |
| 112 | + let (access, pointer) = arg.arg_access(AccessKind::Write); |
| 113 | + self.accesses.push(access); |
| 114 | + self.args |
| 115 | + .push((pointer as *const sys::CUdeviceptr).cast_mut().cast()); |
| 116 | + self |
| 117 | + } |
| 118 | +} |
| 119 | + |
| 120 | +// SAFETY: The device pointer lives in the view, which outlives the launch. |
| 121 | +unsafe impl<'a, 'b: 'a, 'c: 'b, T> PushKernelArg<&'b CudaView<'c, T>> for LaunchArgs<'a> { |
| 122 | + #[inline] |
| 123 | + fn arg(&mut self, arg: &'b CudaView<'c, T>) -> &mut Self { |
| 124 | + let (access, pointer) = arg.arg_access(); |
| 125 | + self.accesses.push(access); |
| 126 | + self.args |
| 127 | + .push((pointer as *const sys::CUdeviceptr).cast_mut().cast()); |
| 128 | + self |
| 129 | + } |
| 130 | +} |
| 131 | + |
| 132 | +// SAFETY: The device pointer lives in the view, which outlives the launch. |
| 133 | +unsafe impl<'a, 'b: 'a, 'c: 'b, T> PushKernelArg<&'b mut CudaViewMut<'c, T>> for LaunchArgs<'a> { |
| 134 | + #[inline] |
| 135 | + fn arg(&mut self, arg: &'b mut CudaViewMut<'c, T>) -> &mut Self { |
| 136 | + let (access, pointer) = arg.arg_access(); |
| 137 | + self.accesses.push(access); |
| 138 | + self.args |
| 139 | + .push((pointer as *const sys::CUdeviceptr).cast_mut().cast()); |
| 140 | + self |
| 141 | + } |
| 142 | +} |
| 143 | + |
| 144 | +impl LaunchArgs<'_> { |
| 145 | + /// Makes [`launch`](Self::launch) return events recorded before and after the kernel. |
| 146 | + pub fn record_kernel_launch(&mut self, flags: sys::CUevent_flags) -> &mut Self { |
| 147 | + self.timing = Some(flags); |
| 148 | + self |
| 149 | + } |
| 150 | + |
| 151 | + /// Enqueues the kernel after the arguments' earlier conflicting accesses. |
| 152 | + /// |
| 153 | + /// # Safety |
| 154 | + /// |
| 155 | + /// The arguments must match the kernel's parameters, and the kernel must access device |
| 156 | + /// memory only within the arguments, writing only those passed mutably. |
| 157 | + pub unsafe fn launch( |
| 158 | + &mut self, |
| 159 | + config: LaunchConfig, |
| 160 | + ) -> Result<Option<(CudaEvent, CudaEvent)>, DriverError> { |
| 161 | + self.stream.bind()?; |
| 162 | + for access in &self.accesses { |
| 163 | + access |
| 164 | + .access |
| 165 | + .before(access.owner, self.stream, access.kind)?; |
| 166 | + } |
| 167 | + let start = self |
| 168 | + .timing |
| 169 | + .map(|flags| self.stream.record_event(Some(flags))) |
| 170 | + .transpose()?; |
| 171 | + // SAFETY: Under the caller's guarantees; the argument pointers outlive the call. |
| 172 | + unsafe { |
| 173 | + sys::cuLaunchKernel( |
| 174 | + self.function.cu_function(), |
| 175 | + config.grid_dim.0, |
| 176 | + config.grid_dim.1, |
| 177 | + config.grid_dim.2, |
| 178 | + config.block_dim.0, |
| 179 | + config.block_dim.1, |
| 180 | + config.block_dim.2, |
| 181 | + config.shared_mem_bytes, |
| 182 | + self.stream.cu_stream(), |
| 183 | + self.args.as_mut_ptr(), |
| 184 | + std::ptr::null_mut(), |
| 185 | + ) |
| 186 | + } |
| 187 | + .result()?; |
| 188 | + // The kernel is already queued: a trailing-event failure must not skip access tracking. |
| 189 | + let end = self |
| 190 | + .timing |
| 191 | + .map(|flags| self.stream.record_event(Some(flags))) |
| 192 | + .transpose(); |
| 193 | + for access in &self.accesses { |
| 194 | + access |
| 195 | + .access |
| 196 | + .after(access.owner, self.stream, access.kind)?; |
| 197 | + } |
| 198 | + Ok(start.zip(end?)) |
| 199 | + } |
| 200 | +} |
0 commit comments