gpu_kernel/
safe_kernel_arg.rs

1use core::marker::PhantomData;
2
3#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
4use crate::{GpuBox, LaunchConfig};
5
6#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
7macro_rules! safe_kernel_arg_impl {
8    ($($ty:ty),*) => {
9        $(
10            #[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
11            unsafe impl SafeKernelArg for $ty {
12                type Output = Self;
13
14                fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
15                    self
16                }
17            }
18        )*
19    };
20}
21
22/// Marker trait for types that are safe to pass to GPU kernels.
23///
24/// The `Output` type that is passed to the GPU can be the same as the type
25/// the trait is implemented for or it can be different.
26/// This is useful to e.g. allow an allocated `Vec<T>` to be passed as a `&[T]`
27/// but not allow passing a slice directly as it might not point to memory that
28/// is readable by the GPU.
29///
30/// # Safety
31///
32/// An implementor guarantees that a GPU kernel receiving the output can freely
33/// use it in safe code.
34#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
35#[diagnostic::on_unimplemented(
36    message = "`SafeKernelArg` is not implemented for `{Self}`",
37    label = "`{Self}` is passed to a kernel here",
38    note = "All kernel arguments must implement `SafeKernelArg` or the kernel must be marked as `unsafe`"
39)]
40pub unsafe trait SafeKernelArg {
41    /// The type that is passed to the GPU.
42    type Output;
43
44    /// Convert into the actual GPU argument.
45    ///
46    /// May panic if necessary constraints are violated.
47    fn into_kernel_arg(self, launch_config: &LaunchConfig) -> Self::Output;
48}
49
50/// Safely pass a list to a kernel and let every thread mutably access one element of the list.
51///
52/// The size of the list needs to be equal to the number of launched threads otherwise launching the
53/// kernel panics.
54///
55/// # Example
56///
57/// ```no_run
58/// use gpu_kernel::{kernel, ThreadIndexedSlice};
59///
60/// gpu_kernel::kernel_lib!();
61///
62/// #[kernel]
63/// fn kernel(elem: ThreadIndexedSlice<'_, i32>) {
64///     // Every thread writes two into the element assigned to this thread
65///     *elem.get_mut() = 2;
66/// }
67///
68/// fn main() {
69///     // Needs the same length as the number of threads launched
70///     let mut data = vec![0; 10];
71///     kernel.launch(
72///         gpu_kernel::LaunchConfig::new()
73///             .threads_per_workgroup([data.len() as u32, 1, 1])
74///             .workgroups([1, 1, 1]),
75///         &mut data,
76///     );
77/// }
78/// ```
79// Needs repr transparent to be passed as a pointer. Structs would be passed by reference.
80#[repr(transparent)]
81pub struct ThreadIndexedSlice<'a, T> {
82    ptr: *mut T,
83    phantom: PhantomData<&'a mut T>,
84}
85
86// SAFETY: These primitive types have the same layout in the CPU and GPU calling convention.
87#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
88safe_kernel_arg_impl!(bool, u8, i8, u16, i16, u32, i32, u64, i64, f32, f64);
89
90// SAFETY: A pointer has the same layout in the CPU and GPU calling convention.
91// It might not point to GPU accessible memory, but that is fine as it is not
92// safely dereferenceable.
93#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
94unsafe impl<T> SafeKernelArg for *const T {
95    type Output = Self;
96
97    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
98        self
99    }
100}
101
102// SAFETY: See *const T
103#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
104unsafe impl<T> SafeKernelArg for *mut T {
105    type Output = Self;
106
107    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
108        self
109    }
110}
111
112// SAFETY: When using the allocator, heap memory is visible to the GPU, so the
113// slice is readable if `T` has the same layout on the GPU.
114#[cfg(all(
115    feature = "amd-allocator",
116    not(any(target_arch = "amdgpu", target_arch = "nvptx64"))
117))]
118unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a Vec<T> {
119    type Output = &'a [T];
120
121    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
122        self.as_slice()
123    }
124}
125
126// SAFETY: See Vec<T>
127#[cfg(all(
128    feature = "amd-allocator",
129    not(any(target_arch = "amdgpu", target_arch = "nvptx64"))
130))]
131unsafe impl<'a> SafeKernelArg for &'a String {
132    type Output = &'a str;
133
134    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
135        self.as_str()
136    }
137}
138
139// SAFETY: See Vec<T>
140#[cfg(all(
141    feature = "amd-allocator",
142    not(any(target_arch = "amdgpu", target_arch = "nvptx64"))
143))]
144unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a Box<T> {
145    type Output = &'a T;
146
147    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
148        self.as_ref()
149    }
150}
151
152// SAFETY: See Vec<T>
153#[cfg(all(
154    feature = "amd-allocator",
155    not(any(target_arch = "amdgpu", target_arch = "nvptx64"))
156))]
157unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a Box<[T]> {
158    type Output = &'a [T];
159
160    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
161        self.as_ref()
162    }
163}
164
165// SAFETY: See Vec<T>
166#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
167unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a GpuBox<T> {
168    type Output = &'a T;
169
170    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
171        self.as_ref()
172    }
173}
174
175// SAFETY: See Vec<T>
176#[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
177unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a GpuBox<[T]> {
178    type Output = &'a [T];
179
180    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
181        self.as_ref()
182    }
183}
184
185// SAFETY: See Vec<T>
186#[cfg(all(
187    feature = "amd-allocator",
188    not(any(target_arch = "amdgpu", target_arch = "nvptx64"))
189))]
190unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a std::sync::Arc<T> {
191    type Output = &'a T;
192
193    fn into_kernel_arg(self, _: &LaunchConfig) -> Self::Output {
194        std::ops::Deref::deref(self)
195    }
196}
197
198/// Implement SafeKernelArg<Output = ThreadIndexedSlice<T>> for a list type
199macro_rules! safe_kernel_arg_list_impl {
200    ($ty:ty: $len:expr; $ptr:expr) => {
201        #[cfg(not(any(target_arch = "amdgpu", target_arch = "nvptx64")))]
202        unsafe impl<'a, T: SafeKernelArg<Output = T>> SafeKernelArg for &'a mut $ty {
203            type Output = ThreadIndexedSlice<'a, T>;
204
205            fn into_kernel_arg(self, launch_config: &LaunchConfig) -> ThreadIndexedSlice<'a, T> {
206                // assert that vector is long enough for launched threads
207                let launch_size = launch_config
208                    .threads_per_workgroup
209                    .unwrap()
210                    .iter()
211                    .map(|i| *i as usize)
212                    .product::<usize>()
213                    * launch_config
214                        .workgroups
215                        .unwrap()
216                        .iter()
217                        .map(|i| *i as usize)
218                        .product::<usize>();
219                assert!(
220                    $len(self) >= launch_size as usize,
221                    "Passed vector is not large enough for the number of launched threads. Expected at least {launch_size} but got {}",
222                    $len(self)
223                );
224                ThreadIndexedSlice {
225                    ptr: $ptr(self),
226                    phantom: PhantomData,
227                }
228            }
229        }
230    };
231}
232
233// SAFETY: See Vec<T>
234#[cfg(feature = "amd-allocator")]
235safe_kernel_arg_list_impl!(Vec<T>: |v: &[_]| v.len(); |v: &mut [_]| v.as_mut_ptr());
236#[cfg(feature = "amd-allocator")]
237safe_kernel_arg_list_impl!(Box<[T]>: |v: &[_]| v.len(); |v: &mut [_]| v.as_mut_ptr());
238safe_kernel_arg_list_impl!(GpuBox<[T]>: |v: &[_]| v.len(); |v: &mut [_]| v.as_mut_ptr());
239
240#[cfg(any(target_arch = "amdgpu", target_arch = "nvptx64"))]
241fn thread_id() -> usize {
242    use crate::intrinsics::*;
243    let dispatch = crate::intrinsics::dispatch_ptr();
244
245    // Compute size as ((z * dimY) + y) * dimX + x
246    let mut id =
247        workitem_id_z() as usize + dispatch.workgroup_size_z as usize * workgroup_id_z() as usize;
248    id *= dispatch.grid_size_y as usize;
249    id += workitem_id_y() as usize + dispatch.workgroup_size_y as usize * workgroup_id_y() as usize;
250    id *= dispatch.grid_size_x as usize;
251    id += workitem_id_x() as usize + dispatch.workgroup_size_x as usize * workgroup_id_x() as usize;
252    id
253}
254
255impl<'a, T> ThreadIndexedSlice<'a, T> {
256    /// Constructs a `ThreadIndexedSlice` from a raw base pointer.
257    ///
258    /// # Safety
259    ///
260    /// - `ptr` must point to at least number of GPU threads consecutive properly initialized values of type T.
261    /// - No constant or mutable reference to the data must exist for the lifetime of this struct.
262    pub unsafe fn from_ptr(ptr: *mut T) -> Self {
263        Self {
264            ptr,
265            phantom: PhantomData,
266        }
267    }
268
269    /// Returns the base pointer of the wrapped list.
270    pub fn as_mut_base_ptr(&mut self) -> *mut T {
271        self.ptr
272    }
273
274    /// Get a reference to the element for the current thread index.
275    #[cfg(any(doc, target_arch = "amdgpu", target_arch = "nvptx64"))]
276    pub fn get(&self) -> &T {
277        unsafe { &*self.ptr.add(thread_id()) }
278    }
279
280    /// Get a mutable reference to the element for the current thread index.
281    #[cfg(any(doc, target_arch = "amdgpu", target_arch = "nvptx64"))]
282    pub fn get_mut(&mut self) -> &mut T {
283        unsafe { &mut *self.ptr.add(thread_id()) }
284    }
285}
gpu_kernel/safe_kernel_arg.rs

gpu_kernel/
safe_kernel_arg.rs
