Skip to main content

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///
23/// Each argument may only be specified once.
24///
25/// # Examples
26///
27/// ```rust,ignore (offload requires a -Z flag)
28/// let mut x = [0.0f64; 256];
29/// core::offload::offload! {
30///     kernel = kernel,
31///     workgroup_dim = [256, 1, 1],
32///     args = (&mut x as *mut [f64; 256],),
33/// }
34/// ```
35#[macro_export]
36#[unstable(feature = "gpu_offload", issue = "131513")]
37#[allow_internal_unstable(core_intrinsics)]
38macro_rules! offload {
39    ( $($field:ident = $val:expr),* $(,)? ) => {
40        $crate::offload!(@munch
41            [ $($field = $val),* ];
42            kernel = NONE;
43            workgroup_dim = ([1, 1, 1]);
44            thread_dim = ([1, 1, 1]);
45            dyn_cache = (0);
46            args = NONE
47        )
48    };
49
50    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
51        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = (SOME $val); workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; args = $a)
52    };
53    (@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; args = $a:tt) => {
54        compile_error!("duplicate field `kernel`")
55    };
56    (@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; args = $a:tt) => {
57        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = (SOME $val); thread_dim = $t; dyn_cache = $d; args = $a)
58    };
59    (@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; args = $a:tt) => {
60        compile_error!("duplicate field `workgroup_dim`")
61    };
62    (@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; args = $a:tt) => {
63        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = (SOME $val); dyn_cache = $d; args = $a)
64    };
65    (@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; args = $a:tt) => {
66        compile_error!("duplicate field `thread_dim`")
67    };
68    (@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); args = $a:tt) => {
69        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = (SOME $val); args = $a)
70    };
71    (@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); args = $a:tt) => {
72        compile_error!("duplicate field `dyn_cache`")
73    };
74    (@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; args = NONE) => {
75        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; workgroup_dim = $w; thread_dim = $t; dyn_cache = $d; args = (SOME $val))
76    };
77    (@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; args = (SOME $old:expr)) => {
78        compile_error!("duplicate field `args`")
79    };
80
81    (@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; args = $a:tt) => {
82        compile_error!(concat!("unknown field `", stringify!($invalid), "`"))
83    };
84
85    (@munch []; kernel = NONE; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = $a:tt) => {
86        compile_error!("missing `kernel`")
87    };
88    (@munch []; kernel = $k:tt; workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = NONE) => {
89        compile_error!("missing `args`")
90    };
91    (@munch []; kernel = (SOME $kernel:expr); workgroup_dim = $w:tt; thread_dim = $t:tt; dyn_cache = $d:tt; args = (SOME $args:expr)) => {
92        $crate::intrinsics::offload::<_, _, ()>(
93            $kernel,
94            $crate::offload!(@value $w),
95            $crate::offload!(@value $t),
96            $crate::offload!(@value $d),
97            $args,
98        )
99    };
100
101    (@value (SOME $val:expr)) => { $val };
102    (@value ($val:expr)) => { $val };
103}