core/offload/
mod.rs

1// offload module
2#[unstable(feature = "gpu_offload", issue = "131513")]
3pub use crate::macros::builtin::offload_kernel;
4#[unstable(feature = "gpu_offload", issue = "131513")]
5pub use crate::offload;
6
7/// Launches a kernel on an offload device (e.g., a GPU).
8///
9/// This macro is an interface over the `offload` intrinsic. The kernel itself must be defined
10/// using the [`offload_kernel`] macro.
11///
12/// The following named arguments are accepted:
13///
14/// - `kernel`: The kernel function to offload. Must be a function item. (required)
15/// - `args`: A tuple of arguments forwarded to `kernel`. (required)
16/// - `workgroup_dim`: A 3D size specifying the number of workgroups to launch.
17///   Defaults to `[1, 1, 1]`.
18/// - `thread_dim`: A 3D size specifying the number of threads per workgroup.
19///   Defaults to `[1, 1, 1]`.
20/// - `dyn_cache`: The amount of dynamic shared memory, in bytes, to allocate for the kernel.
21///   Defaults to `0`.
22/// - `device`: The index of the device to offload to. Must be `>= 0`. If omitted, the
23///   default device is used. Use [`crate::intrinsics::offload_get_num_devices`] to discover
24///   which device ids are valid.
25///
26/// Each argument may only be specified once.
27///
28/// # Examples
29///
30/// ```rust,ignore (offload requires a -Z flag)
31/// let mut x = [0.0f64; 256];
32/// core::offload::offload! {
33///     kernel = kernel,
34///     workgroup_dim = [256, 1, 1],
35///     args = (&mut x as *mut [f64; 256],),
36/// }
37/// ```
38#[macro_export]
39#[unstable(feature = "gpu_offload", issue = "131513")]
40#[allow_internal_unstable(core_intrinsics)]
41macro_rules! offload {
42    ( $($field:ident = $val:expr),* $(,)? ) => {
43        $crate::offload!(@munch
44            [ $($field = $val),* ];
45            kernel = NONE;
46            workgroup_dim = ([1, 1, 1]);
47            thread_dim = ([1, 1, 1]);
48            dyn_cache = (0);
49            device = NONE;
50            args = NONE
51        )
52    };
53
54    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
55        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = (SOME $val); workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = $a)
56    };
57    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = (SOME $old:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
58        compile_error!("duplicate field `kernel`")
59    };
60    (@munch [workgroup_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = ([1, 1, 1]); thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
61        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = (SOME $val); thread_dim = $t; dyn_cache = $d; device = $device; args = $a)
62    };
63    (@munch [workgroup_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = (SOME $old:expr); thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
64        compile_error!("duplicate field `workgroup_dim`")
65    };
66    (@munch [thread_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = ([1, 1, 1]); dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
67        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = (SOME $val); dyn_cache = $d; device = $device; args = $a)
68    };
69    (@munch [thread_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = (SOME $old:expr); dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
70        compile_error!("duplicate field `thread_dim`")
71    };
72    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (0); device = $device:tt; args = $a:tt) => {
73        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = (SOME $val); device = $device; args = $a)
74    };
75    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = (SOME $old:expr); device = $device:tt; args = $a:tt) => {
76        compile_error!("duplicate field `dyn_cache`")
77    };
78    (@munch [device = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = NONE; args = $a:tt) => {
79        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = (SOME $val); args = $a)
80    };
81    (@munch [device = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = (SOME $old:expr); args = $a:tt) => {
82        compile_error!("duplicate field `device`")
83    };
84    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
85        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; device = $device; args = (SOME $val))
86    };
87    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME $old:expr)) => {
88        compile_error!("duplicate field `args`")
89    };
90
91    (@munch [$invalid:ident = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
92        compile_error!(concat!("unknown field `", stringify!($invalid), "`"))
93    };
94
95    (@munch []; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = $a:tt) => {
96        compile_error!("missing `kernel`")
97    };
98    (@munch []; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = NONE) => {
99        compile_error!("missing `args`")
100    };
101    (@munch []; kernel = (SOME $kernel:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; device = $device:tt; args = (SOME $args:expr)) => {
102        $crate::intrinsics::offload::<_, _, ()>(
103            $kernel,
104            $crate::offload!(@value $w),
105            $crate::offload!(@value $t),
106            $crate::offload!(@value $d),
107            $crate::offload!(@device $device),
108            $args,
109        )
110    };
111
112    (@value (SOME $val:expr)) => { $val };
113    (@value ($val:expr)) => { $val };
114
115    // if `device` is omitted (`NONE), we use the OpenMP default device (`-1`)
116    (@device NONE) => { -1 };
117    (@device (SOME $val:expr)) => { {
118        const { $crate::assert!($val >= 0, "offload device must be non-negative; omit `device` to use the default device") };
119        let device: i32 = $val;
120        $crate::assert!(
121            device < $crate::intrinsics::offload_get_num_devices(),
122            "offload device {} is not available",
123            device,
124        );
125        device
126    } };
127}
core/offload/mod.rs

core/offload/
mod.rs
