core/offload/
mod.rs

1// offload module
2use core::ptr::NonNull;
3
4#[unstable(feature = "gpu_offload", issue = "131513")]
5pub use crate::macros::builtin::offload_kernel;
6#[unstable(feature = "gpu_offload", issue = "131513")]
7pub use crate::offload;
8
9/// Launches a kernel on an offload device (e.g., a GPU).
10///
11/// This macro is an interface over the `offload` intrinsic. The kernel itself must be defined
12/// using the [`offload_kernel`] macro.
13///
14/// The following named arguments are accepted:
15///
16/// - `kernel`: The kernel function to offload. Must be a function item. (required)
17/// - `args`: A tuple of arguments forwarded to `kernel`. (required)
18/// - `workgroup_dim`: A 3D size specifying the number of workgroups to launch.
19///   Defaults to `[1, 1, 1]`.
20/// - `thread_dim`: A 3D size specifying the number of threads per workgroup.
21///   Defaults to `[1, 1, 1]`.
22/// - `dyn_cache`: The amount of dynamic shared memory, in bytes, to allocate for the kernel.
23///   Defaults to `0`.
24/// - `device`: The index of the device to offload to. Must be `>= 0`. If omitted, the
25///   default device is used. Use [`crate::intrinsics::offload_get_num_devices`] to discover
26///   which device ids are valid.
27///
28/// Each argument may only be specified once.
29///
30/// # Examples
31///
32/// ```rust,ignore (offload requires a -Z flag)
33/// let mut x = [0.0f64; 256];
34/// core::offload::offload! {
35///     kernel = kernel,
36///     workgroup_dim = [256, 1, 1],
37///     args = (&mut x as *mut [f64; 256],),
38/// }
39/// ```
40#[macro_export]
41#[unstable(feature = "gpu_offload", issue = "131513")]
42#[allow_internal_unstable(core_intrinsics, gpu_offload)]
43macro_rules! offload {
44    (@munch [kernel = $val:expr $(, $($rest:tt)*)?]; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
45        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = (SOME $val); workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = $a)
46    };
47    (@munch [kernel = $val:expr $(, $($rest:tt)*)?]; kernel = (SOME $old:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
48        compile_error!("duplicate field `kernel`")
49    };
50    (@munch [workgroup_dim = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = ([1, 1, 1]); thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
51        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = (SOME $val); thread_dim = $t; dyn_cache = $d; device = $device; args = $a)
52    };
53    (@munch [workgroup_dim = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = (SOME $old:expr); thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
54        compile_error!("duplicate field `workgroup_dim`")
55    };
56    (@munch [thread_dim = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = ([1, 1, 1]); dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
57        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = (SOME $val); dyn_cache = $d; device = $device; args = $a)
58    };
59    (@munch [thread_dim = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = (SOME $old:expr); dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
60        compile_error!("duplicate field `thread_dim`")
61    };
62    (@munch [dyn_cache = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (0); device = $device:tt; args = $a:tt) => {
63        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = (SOME $val); device = $device; args = $a)
64    };
65    (@munch [dyn_cache = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (SOME $old:expr); device = $device:tt; args = $a:tt) => {
66        compile_error!("duplicate field `dyn_cache`")
67    };
68    (@munch [device = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = NONE; args = $a:tt) => {
69        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = (SOME $val); args = $a)
70    };
71    (@munch [device = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = (SOME $old:expr); args = $a:tt) => {
72        compile_error!("duplicate field `device`")
73    };
74    (@munch [args = ($($arg:expr),* $(,)?) $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
75        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = (SOME_TUPLE ($($arg),*)))
76    };
77    (@munch [args = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
78        $crate::offload!(@munch [ $($($rest)*)? ]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = (SOME $val))
79    };
80    (@munch [args = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME $old:expr)) => {
81        compile_error!("duplicate field `args`")
82    };
83    (@munch [args = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME_TUPLE $old:tt)) => {
84        compile_error!("duplicate field `args`")
85    };
86
87    (@munch [$invalid:ident = $val:expr $(, $($rest:tt)*)?]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
88        compile_error!(concat!("unknown field `", stringify!($invalid), "`"))
89    };
90
91    (@munch []; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
92        compile_error!("missing `kernel`")
93    };
94    (@munch []; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
95        compile_error!("missing `args`")
96    };
97    (@munch []; kernel = (SOME $kernel:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME_TUPLE ($($arg:expr),* $(,)?))) => {{
98        #[allow(unused_imports)]
99        use $crate::offload::{DefaultLaunchCheck as _, RegionLaunchCheck as _};
100
101        let __offload_grid = $crate::offload!(@value $w);
102        let __offload_block = $crate::offload!(@value $t);
103
104        $crate::intrinsics::offload::<_, _, ()>(
105            $kernel,
106            __offload_grid,
107            __offload_block,
108            $crate::offload!(@value $d),
109            $crate::offload!(@device $device),
110            ($({
111                let __offload_arg = $arg;
112                (&__offload_arg).__offload_check_launch(__offload_grid, __offload_block);
113                __offload_arg
114            },)*),
115        )
116    }};
117    // Accepting an arbitrary expression here would let callers bypass the per-argument launch
118    // performed for `args = (..)` by binding the tuple to avariable first.
119    // FIXME(offload): run those checks on arbitrary expressions
120    (@munch []; kernel = (SOME $kernel:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME $args:expr)) => {
121        compile_error!(
122            "`args` must be a tuple literal"
123        )
124    };
125
126    (@value (SOME $val:expr)) => { $val };
127    (@value ($val:expr)) => { $val };
128
129    // if `device` is omitted (`NONE), we use the OpenMP default device (`-1`)
130    (@device NONE) => { -1 };
131    (@device (SOME $val:expr)) => { {
132        const { $crate::assert!($val >= 0, "offload device must be non-negative; omit `device` to use the default device") };
133        let device: i32 = $val;
134        $crate::assert!(
135            device < $crate::intrinsics::offload_get_num_devices(),
136            "offload device {} is not available",
137            device,
138        );
139        device
140    } };
141
142    ( $($tt:tt)* ) => {
143        $crate::offload!(@munch
144            [ $($tt)* ];
145            kernel = NONE;
146            workgroup_dim = ([1, 1, 1]);
147            thread_dim = ([1, 1, 1]);
148            dyn_cache = (0);
149            device = NONE;
150            args = NONE
151        )
152    };
153}
154
155// Region & Partitioning Strategy
156
157/// Error returned by [`PartitioningStrategy::check_launch`] when a launch
158/// configuration is not compatible with a partitioning strategy.
159#[derive(Debug)]
160#[unstable(feature = "gpu_offload", issue = "131513")]
161pub struct LaunchError {
162    message: &'static str,
163}
164
165impl LaunchError {
166    /// Creates a new [`LaunchError`] with the given message.
167    #[unstable(feature = "gpu_offload", issue = "131513")]
168    pub const fn new(message: &'static str) -> Self {
169        Self { message }
170    }
171}
172
173#[unstable(feature = "gpu_offload", issue = "131513")]
174impl core::fmt::Display for LaunchError {
175    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
176        f.write_str(self.message)
177    }
178}
179
180/// Defines how execution units access memory regions.
181///
182/// # Safety
183///
184/// Implementations must guarantee that generated views are disjoint.
185///
186/// Implementations must also implement [`check_launch`] correctly since it is soundness-critical.
187///
188/// [`check_launch`]: PartitioningStrategy::check_launch
189#[unstable(feature = "gpu_offload", issue = "131513")]
190pub unsafe trait PartitioningStrategy {
191    /// Read-only view type for the partitioned memory region.
192    type View<'a, T: 'a>;
193
194    /// Mutable view type for the partitioned memory region.
195    type ViewMut<'a, T: 'a>;
196
197    /// Returns the execution index of the current unit.
198    fn index() -> usize;
199
200    /// Returns a read-only view of the region for the current execution context.
201    ///
202    /// # Safety
203    ///
204    /// `ptr` must point to `len` valid, initialized elements of type `T`.
205    /// The memory must stay valid for lifetime `'a`.
206    unsafe fn get<'a, T>(ptr: NonNull<T>, len: usize) -> Option<Self::View<'a, T>>;
207
208    /// Returns a mutable view of the region for the current execution context.
209    ///
210    /// # Safety
211    ///
212    /// `ptr` must point to `len` valid, initialized elements of type `T`.
213    /// The memory must stay valid for lifetime `'a`.
214    /// The returned view must be disjoint from all other active views.
215    unsafe fn get_mut<'a, T>(ptr: NonNull<T>, len: usize) -> Option<Self::ViewMut<'a, T>>;
216
217    /// Checks that a kernel using this strategy can be launched with `len`
218    /// elements on a `grid` by `block` launch configuration.
219    ///
220    /// This is called automatically before the kernel is launched.
221    ///
222    /// # Soundness
223    ///
224    /// This check is soundness-critical. It may only return `Ok` unconditionally
225    /// if and only if a launch is valid for arbitrary `len`, `grid` and `block` values.
226    /// If there is any configuration for which this strategy would produce incorrect
227    /// results or cause undefined behavior, it must return a [`LaunchError`] instead.
228    fn check_launch(len: usize, grid: [u32; 3], block: [u32; 3]) -> Result<(), LaunchError>;
229}
230
231/// A memory region bound to a partitioning strategy.
232#[derive(Debug)]
233#[unstable(feature = "gpu_offload", issue = "131513")]
234#[rustc_diagnostic_item = "offload_region"]
235pub struct Region<'a, T, S: PartitioningStrategy> {
236    ptr: NonNull<T>,
237    len: usize,
238    _marker: core::marker::PhantomData<(&'a mut [T], S)>,
239}
240
241/// Raw representation used to build a [`Region`] from common aggregate types.
242#[derive(Debug)]
243#[unstable(feature = "gpu_offload", issue = "131513")]
244pub struct RawRegion<'a, T> {
245    ptr: NonNull<T>,
246    len: usize,
247    _marker: core::marker::PhantomData<&'a mut [T]>,
248}
249
250impl<'a, T> From<&'a mut [T]> for RawRegion<'a, T> {
251    fn from(data: &'a mut [T]) -> Self {
252        // SAFETY: `data.as_mut_ptr()` is non-null, because it is derived from a reference.
253        let ptr = unsafe { NonNull::new_unchecked(data.as_mut_ptr()) };
254        Self { ptr, len: data.len(), _marker: core::marker::PhantomData }
255    }
256}
257
258impl<'a, T, const N: usize> From<&'a mut [T; N]> for RawRegion<'a, T> {
259    fn from(data: &'a mut [T; N]) -> Self {
260        // SAFETY: `data.as_mut_ptr()` is non-null, because it is derived from a reference.
261        let ptr = unsafe { NonNull::new_unchecked(data.as_mut_ptr()) };
262        Self { ptr, len: N, _marker: core::marker::PhantomData }
263    }
264}
265
266#[unstable(feature = "gpu_offload", issue = "131513")]
267impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> {
268    /// Creates a new partitioned region from data convertible into a [`RawRegion`].
269    pub fn new<D>(data: D) -> Self
270    where
271        D: Into<RawRegion<'a, T>>,
272    {
273        let raw = data.into();
274        Self { ptr: raw.ptr, len: raw.len, _marker: core::marker::PhantomData }
275    }
276
277    /// Returns a read-only view for the current execution context.
278    pub fn get(&self) -> Option<S::View<'_, T>> {
279        // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`.
280        unsafe { S::get(self.ptr, self.len) }
281    }
282
283    /// Returns a mutable view for the current execution context.
284    pub fn get_mut(&mut self) -> Option<S::ViewMut<'_, T>> {
285        // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`.
286        // The strategy guarantees that the returned view is disjoint.
287        unsafe { S::get_mut(self.ptr, self.len) }
288    }
289
290    /// Checks that this region can be launched with the given configuration.
291    ///
292    /// Called automatically by the [`offload!`] macro before launching a kernel.
293    pub fn check_launch(&self, grid: [u32; 3], block: [u32; 3]) -> Result<(), LaunchError> {
294        S::check_launch(self.len, grid, block)
295    }
296
297    /// Reborrows the region, producing a new region that aliases the same memory with
298    /// the lifetime of the borrow.
299    pub fn reborrow(&mut self) -> Region<'_, T, S> {
300        Region { ptr: self.ptr, len: self.len, _marker: core::marker::PhantomData }
301    }
302}
303
304// Launch validation helpers.
305
306/// Implementation detail of [`offload!`]: validates a [`Region`] argument against
307/// the launch configuration of an offload call.
308#[doc(hidden)]
309#[unstable(feature = "gpu_offload", issue = "131513")]
310pub trait RegionLaunchCheck {
311    /// Panics if the region's partitioning strategy rejects the launch.
312    fn __offload_check_launch(&self, grid: [u32; 3], block: [u32; 3]);
313}
314
315#[unstable(feature = "gpu_offload", issue = "131513")]
316impl<'a, T, S: PartitioningStrategy> RegionLaunchCheck for Region<'a, T, S> {
317    fn __offload_check_launch(&self, grid: [u32; 3], block: [u32; 3]) {
318        if let Err(err) = self.check_launch(grid, block) {
319            panic!(
320                "offload launch is not supported by the region's partitioning strategy: {}",
321                err,
322            );
323        }
324    }
325}
326
327/// no-op launch check for kernel arguments that are not [`Region`]s.
328#[doc(hidden)]
329#[unstable(feature = "gpu_offload", issue = "131513")]
330pub trait DefaultLaunchCheck {
331    /// Does nothing.
332    fn __offload_check_launch(&self, grid: [u32; 3], block: [u32; 3]);
333}
334
335#[unstable(feature = "gpu_offload", issue = "131513")]
336impl<T: ?Sized> DefaultLaunchCheck for &T {
337    fn __offload_check_launch(&self, _grid: [u32; 3], _block: [u32; 3]) {}
338}
core/offload/mod.rs

core/offload/
mod.rs
