Rollup merge of rust-lang#156642 - Sa4dUs:offload-macro, r=ZuseZ4

`offload_kernel` macro expansion

r? @ZuseZ4

Author: JonathanBrouwer

compiler/rustc_builtin_macros/src/lib.rs

@@ -42,6 +42,7 @@ mod format_foreign;
 mod global_allocator;
 mod iter;
 mod log_syntax;
+mod offload;
 mod pattern_type;
 mod source_util;
 mod test;
@@ -116,6 +117,7 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
         eii_declaration: eii::eii_declaration,
         eii_shared_macro: eii::eii_shared_macro,
         global_allocator: global_allocator::expand,
+        offload_kernel: offload::expand_kernel,
         test: test::expand_test,
         test_case: test::expand_test_case,
         unsafe_eii: eii::unsafe_eii,

compiler/rustc_builtin_macros/src/offload.rs

@@ -0,0 +1,205 @@
+use rustc_ast::token::{Delimiter, Token, TokenKind};
+use rustc_ast::tokenstream::{DelimSpan, Spacing, TokenStream, TokenTree};
+use rustc_ast::{AttrItem, ast};
+use rustc_expand::base::{Annotatable, ExtCtxt};
+use rustc_session::config::Offload;
+use rustc_span::{Ident, Span, sym};
+use thin_vec::thin_vec;
+
+use crate::errors;
+
+fn compile_for_device(ecx: &mut ExtCtxt<'_>) -> bool {
+    ecx.sess.opts.unstable_opts.offload.contains(&Offload::Device)
+}
+
+fn outer_normal_attr(
+    kind: &Box<rustc_ast::NormalAttr>,
+    id: rustc_ast::AttrId,
+    span: Span,
+) -> rustc_ast::Attribute {
+    let style = rustc_ast::AttrStyle::Outer;
+    let kind = rustc_ast::AttrKind::Normal(kind.clone());
+    rustc_ast::Attribute { kind, id, style, span }
+}
+
+fn extract_fn(
+    item: &Annotatable,
+) -> Option<(ast::Visibility, ast::FnSig, Ident, ast::Generics, Option<Box<ast::Block>>)> {
+    match item {
+        Annotatable::Item(iitem) => match &iitem.kind {
+            ast::ItemKind::Fn(ast::Fn { sig, ident, generics, body, .. }) => {
+                Some((iitem.vis.clone(), sig.clone(), *ident, generics.clone(), body.clone()))
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+/// The `offload_kernel` macro expands the function into two separate definitions:
+/// one on the host to handle the call, and one on the device for executing the kernel.
+///
+/// ```
+/// #[offload_kernel]
+/// fn foo(a: &[f32], b: &[f32], c: *mut f32) {
+///     *c = a[0] + b[0];
+/// }
+/// ```
+///
+/// This expands to the host-side function:
+///
+/// ```
+/// #[unsafe(no_mangle)]
+/// #[inline(never)]
+/// fn foo(_: &[f32], _: &[f32], _: *mut f32) {
+///     ::core::panicking::panic("not implemented")
+/// }
+/// ```
+///
+/// And the device-side kernel:
+///
+/// ```
+/// #[rustc_offload_kernel]
+/// #[unsafe(no_mangle)]
+/// unsafe extern "gpu-kernel" fn foo(a: &[f32], b: &[f32], c: *mut f32) {
+///     *c = a[0] + b[0];
+/// }
+/// ```
+pub(crate) fn expand_kernel(
+    ecx: &mut ExtCtxt<'_>,
+    expand_span: Span,
+    _meta_item: &ast::MetaItem,
+    item: Annotatable,
+) -> Vec<Annotatable> {
+    let dcx = ecx.sess.dcx();
+
+    let Some((vis, sig, ident, generics, body)) = extract_fn(&item) else {
+        dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
+        return vec![item];
+    };
+
+    let span = ecx.with_def_site_ctxt(expand_span);
+
+    // device function
+    let mut device_fn = Box::new(ast::Fn {
+        defaultness: ast::Defaultness::Implicit,
+        sig: sig.clone(),
+        ident,
+        generics: generics.clone(),
+        contract: None,
+        body,
+        define_opaque: None,
+        eii_impls: Default::default(),
+    });
+
+    let extern_gpu_kernel = ast::Extern::from_abi(
+        Some(ast::StrLit {
+            symbol: sym::gpu_kernel,
+            suffix: None,
+            symbol_unescaped: sym::gpu_kernel,
+            style: ast::StrStyle::Cooked,
+            span,
+        }),
+        span,
+    );
+    device_fn.sig.header.ext = extern_gpu_kernel;
+    device_fn.sig.header.safety = ast::Safety::Unsafe(span);
+
+    // rustc_offload_kernel attr
+    let rustc_offload_kernel_attr =
+        Box::new(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::rustc_offload_kernel)));
+    let rustc_offload_kernel = outer_normal_attr(
+        &rustc_offload_kernel_attr,
+        ecx.sess.psess.attr_id_generator.mk_attr_id(),
+        span,
+    );
+
+    // unsafe(no_mangle) attr
+    let unsafe_item = AttrItem {
+        unsafety: ast::Safety::Unsafe(span),
+        path: ast::Path::from_ident(Ident::new(sym::no_mangle, span)),
+        args: ast::AttrItemKind::Unparsed(ast::AttrArgs::Empty),
+        tokens: None,
+    };
+
+    let no_mangle_attr = Box::new(ast::NormalAttr { item: unsafe_item, tokens: None });
+    let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
+    let unsafe_no_mangle = outer_normal_attr(&no_mangle_attr, new_id, span);
+
+    let device_item = {
+        let mut item = ecx.item(
+            span,
+            thin_vec![rustc_offload_kernel, unsafe_no_mangle],
+            ast::ItemKind::Fn(device_fn),
+        );
+        item.vis = vis.clone();
+        Annotatable::Item(item)
+    };
+
+    // unimplemented! body
+    let macro_expr = ecx.expr_macro_call(
+        span,
+        ecx.macro_call(
+            span,
+            ecx.path_global(
+                span,
+                [sym::std, sym::unimplemented].map(|s| Ident::new(s, span)).to_vec(),
+            ),
+            Delimiter::Parenthesis,
+            TokenStream::default(),
+        ),
+    );
+    let stmt = ecx.stmt_expr(macro_expr);
+    let body = ecx.block(span, thin_vec![stmt]);
+
+    // host function
+    let mut host_fn = Box::new(ast::Fn {
+        defaultness: ast::Defaultness::Implicit,
+        sig: sig.clone(),
+        ident,
+        generics: generics.clone(),
+        contract: None,
+        body: Some(body),
+        define_opaque: None,
+        eii_impls: Default::default(),
+    });
+
+    for param in host_fn.sig.decl.inputs.iter_mut() {
+        param.pat = Box::new(ecx.pat_wild(param.pat.span));
+    }
+
+    // inline(never) attr
+    let ts: Vec<TokenTree> = vec![TokenTree::Token(
+        Token::new(TokenKind::Ident(sym::never, false.into()), span),
+        Spacing::Joint,
+    )];
+
+    let never_arg = ast::DelimArgs {
+        dspan: DelimSpan::from_single(span),
+        delim: Delimiter::Parenthesis,
+        tokens: TokenStream::from_iter(ts),
+    };
+
+    let inline_item = ast::AttrItem {
+        unsafety: ast::Safety::Default,
+        path: ast::Path::from_ident(Ident::with_dummy_span(sym::inline)),
+        args: rustc_ast::ast::AttrItemKind::Unparsed(ast::AttrArgs::Delimited(never_arg)),
+        tokens: None,
+    };
+    let inline_never_attr = Box::new(ast::NormalAttr { item: inline_item, tokens: None });
+
+    let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
+    let inline_never = outer_normal_attr(&inline_never_attr, new_id, span);
+
+    let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
+    let unsafe_no_mangle = outer_normal_attr(&no_mangle_attr, new_id, span);
+
+    let host_item = {
+        let mut item =
+            ecx.item(span, thin_vec![unsafe_no_mangle, inline_never], ast::ItemKind::Fn(host_fn));
+        item.vis = vis.clone();
+        Annotatable::Item(item)
+    };
+
+    if compile_for_device(ecx) { vec![device_item] } else { vec![host_item] }
+}

compiler/rustc_span/src/symbol.rs

@@ -1040,6 +1040,7 @@ symbols! {
         global_asm,
         global_registration,
         globs,
+        gpu_kernel: "gpu-kernel",
         gpu_launch_sized_workgroup_mem,
         gt,
         guard,
@@ -1424,6 +1425,7 @@ symbols! {
         of,
         off,
         offload,
+        offload_kernel,
         offset,
         offset_of,
         offset_of_enum,
@@ -2163,6 +2165,7 @@ symbols! {
         underscore_imports,
         underscore_lifetimes,
         uniform_paths,
+        unimplemented,
         unit,
         universal_impl_trait,
         unix,

library/core/src/lib.rs

@@ -234,6 +234,10 @@ pub mod autodiff {
     pub use crate::macros::builtin::{autodiff_forward, autodiff_reverse};
 }
 
+#[unstable(feature = "gpu_offload", issue = "131513")]
+#[doc = include_str!("../../core/src/offload.md")]
+pub mod offload;
+
 #[unstable(feature = "contracts", issue = "128044")]
 pub mod contracts;
 

library/core/src/macros/mod.rs

@@ -1628,6 +1628,48 @@ pub(crate) mod builtin {
         /* compiler built-in */
     }
 
+    /// The `offload_kernel` macro is applied to a function to generate two separate
+    /// definitions: a host-side wrapper for dispatch and a device-side kernel.
+    ///
+    /// The macro does not perform the offload itself. It generates the necessary
+    /// code required by the compiler's offloading infrastructure.
+    ///
+    /// ### Usage example:
+    ///
+    /// ```rust,ignore (offload requires a -Z flag)
+    /// #[offload_kernel]
+    /// fn foo(a: &[f32], b: &[f32], c: *mut f32) {
+    ///     *c = a[0] + b[0];
+    /// }
+    /// ```
+    ///
+    /// This expands to the host-side function:
+    ///
+    /// ```rust,ignore (offload requires a -Z flag)
+    /// #[unsafe(no_mangle)]
+    /// #[inline(never)]
+    /// fn foo(_: &[f32], _: &[f32], _: *mut f32) {
+    ///     ::core::panicking::panic("not implemented")
+    /// }
+    /// ```
+    ///
+    /// And the device-side kernel:
+    ///
+    /// ```rust,ignore (offload requires a -Z flag)
+    /// #[rustc_offload_kernel]
+    /// #[unsafe(no_mangle)]
+    /// unsafe extern "gpu-kernel" fn foo(a: &[f32], b: &[f32], c: *mut f32) {
+    ///     *c = a[0] + b[0];
+    /// }
+    /// ```
+    #[unstable(feature = "gpu_offload", issue = "131513")]
+    #[allow_internal_unstable(rustc_attrs)]
+    #[allow_internal_unstable(core_intrinsics)]
+    #[rustc_builtin_macro]
+    pub macro offload_kernel($item:item) {
+        /* compiler built-in */
+    }
+
     /// Asserts that a boolean expression is `true` at runtime.
     ///
     /// This will invoke the [`panic!`] macro if the provided expression cannot be

library/core/src/offload.md

@@ -0,0 +1,39 @@
+This module provides support for gpu offloading. For technical details regarding the `offload_kernel`
+and how to use it, see their respective documentation.
+
+## General usage
+The `offload_kernel` macro can be applied to a function to generate the necessary code to launch a
+kernel on the target device.
+
+```rust,ignore (optional component)
+#[offload_kernel]
+fn kernel(x: *mut [f64; 256]) {
+    // SAFETY:
+    // calling our `arch` functions and dereferencing a raw pointer is unsafe
+    unsafe {
+        let n = (*x).len();
+        let i = (thread_idx_x() + block_idx_x() * block_dim_x()) as usize;
+        if i < n {
+            (*x)[i] = i as f64;
+        }
+    }
+}
+```
+
+To launch an offloaded kernel, the only current way is to use the `core::intrinsic::offload`
+intrinsic (note that intrinsics usage is discouraged outside the standard library). This
+allows you to specify grid and block dimensions and pass the required arguments to the device.
+
+```rust,ignore (optional component)
+let mut x = [0.0f64; 256];
+core::intrinsics::offload::<_, _, ()>(kernel, [256, 1, 1], [1, 1, 1], (&mut x as *mut [f64; 256],));
+```
+
+For precise information on the `offload` intrinsic, see its respective documentation.
+
+## Current limitations:
+
+- Usage is restricted to types supported by the current device-mapping implementation.
+- Generics and functions accepting dyn Trait are not supported.
+- Kernel execution is currently restricted to intrinsics usage, which is discouraged outside of the
+standard library.

library/core/src/offload/mod.rs

@@ -0,0 +1,5 @@
+// offload module
+#[unstable(feature = "gpu_offload", issue = "131513")]
+pub use crate::macros::builtin::offload_kernel;
+#[unstable(feature = "gpu_offload", issue = "131513")]
+pub use crate::offload;

library/std/src/lib.rs

@@ -289,6 +289,7 @@
 #![feature(f16)]
 #![feature(f128)]
 #![feature(ffi_const)]
+#![feature(gpu_offload)]
 #![feature(intra_doc_pointers)]
 #![feature(lang_items)]
 #![feature(link_cfg)]
@@ -663,6 +664,12 @@ pub mod autodiff {
     pub use core::autodiff::{autodiff_forward, autodiff_reverse};
 }
 
+#[unstable(feature = "gpu_offload", issue = "131513")]
+#[doc = include_str!("../../core/src/offload.md")]
+pub mod offload {
+    pub use core::offload::{offload, offload_kernel};
+}
+
 #[stable(feature = "futures_api", since = "1.36.0")]
 pub mod task {
     //! Types and Traits for working with asynchronous tasks.