feat: support nested functions with #[offload_kernel]

compiler/rustc_builtin_macros/src/diagnostics.rs

@@ -233,6 +233,13 @@ mod autodiff {
     }
 }
 
+#[derive(Diagnostic)]
+#[diag("offload_kernel must be applied to function")]
+pub(crate) struct OffloadKernelInvalidApplication {
+    #[primary_span]
+    pub(crate) span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag("cannot concatenate {$lit_kind} literals")]
 pub(crate) struct ConcatBytesInvalid {

compiler/rustc_builtin_macros/src/offload.rs

@@ -1,6 +1,6 @@
-use rustc_ast::ast;
 use rustc_ast::token::{Delimiter, IdentKind, Token, TokenKind};
 use rustc_ast::tokenstream::{DelimSpan, Spacing, TokenStream, TokenTree};
+use rustc_ast::{DUMMY_NODE_ID, ast};
 use rustc_expand::base::{Annotatable, ExtCtxt};
 use rustc_session::config::Offload;
 use rustc_span::{DUMMY_SP, Ident, Span, sym};
@@ -28,6 +28,15 @@ fn extract_fn(
             }
             _ => None,
         },
+        Annotatable::Stmt(stmt) => match &stmt.kind {
+            ast::StmtKind::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,
+        },
         _ => None,
     }
 }
@@ -68,7 +77,7 @@ pub(crate) fn expand_kernel(
     let dcx = ecx.sess.dcx();
 
     let Some((vis, sig, ident, generics, body)) = extract_fn(&item) else {
-        dcx.emit_err(diagnostics::AutoDiffInvalidApplication { span: item.span() });
+        dcx.emit_err(diagnostics::OffloadKernelInvalidApplication { span: item.span() });
         return vec![item];
     };
 
@@ -109,10 +118,21 @@ pub(crate) fn expand_kernel(
     );
 
     let device_item = {
-        let mut item =
+        let mut device_item_ecx =
             ecx.item(span, thin_vec![rustc_offload_kernel.clone()], ast::ItemKind::Fn(device_fn));
-        item.vis = vis.clone();
-        Annotatable::Item(item)
+        device_item_ecx.vis = vis.clone();
+
+        match &item {
+            Annotatable::Item(_) => Annotatable::Item(device_item_ecx),
+            Annotatable::Stmt(_) => Annotatable::Stmt(Box::new(ast::Stmt {
+                id: DUMMY_NODE_ID,
+                kind: ast::StmtKind::Item(device_item_ecx),
+                span,
+            })),
+            _ => {
+                unreachable!("item kind checked previously")
+            }
+        }
     };
 
     // unimplemented! body
@@ -171,13 +191,24 @@ pub(crate) fn expand_kernel(
     let inline_never = outer_normal_attr(&inline_never_attr, new_id, span);
 
     let host_item = {
-        let mut item = ecx.item(
+        let mut host_item_ecx = ecx.item(
             span,
             thin_vec![rustc_offload_kernel, inline_never],
             ast::ItemKind::Fn(host_fn),
         );
-        item.vis = vis;
-        Annotatable::Item(item)
+        host_item_ecx.vis = vis;
+
+        match &item {
+            Annotatable::Item(_) => Annotatable::Item(host_item_ecx),
+            Annotatable::Stmt(_) => Annotatable::Stmt(Box::new(ast::Stmt {
+                id: DUMMY_NODE_ID,
+                kind: ast::StmtKind::Item(host_item_ecx),
+                span,
+            })),
+            _ => {
+                unreachable!("item kind checked previously")
+            }
+        }
     };
 
     if compile_for_device(ecx) { vec![device_item] } else { vec![host_item] }

tests/pretty/offload/offload_kernel_nested.device.pp

@@ -0,0 +1,25 @@
+#![feature(prelude_import)]
+#![no_std]
+//@ only-nightly
+//@ revisions: host device
+
+//@ pretty-mode:expanded
+//@ pretty-compare-only
+//@[host] pp-exact:offload_kernel_nested.host.pp
+//@[device] pp-exact:offload_kernel_nested.device.pp
+
+//@[device] compile-flags: -Zunstable-options -Zoffload=Device
+
+#![feature(gpu_offload)]
+extern crate std;
+#[prelude_import]
+use ::std::prelude::rust_2015::*;
+
+use std::offload::offload_kernel;
+
+fn kernel() {
+    #[rustc_offload_kernel]
+    unsafe extern "gpu-kernel" fn inner_kernel() {}
+}
+
+fn main() {}

tests/pretty/offload/offload_kernel_nested.host.pp

@@ -0,0 +1,28 @@
+#![feature(prelude_import)]
+#![no_std]
+//@ only-nightly
+//@ revisions: host device
+
+//@ pretty-mode:expanded
+//@ pretty-compare-only
+//@[host] pp-exact:offload_kernel_nested.host.pp
+//@[device] pp-exact:offload_kernel_nested.device.pp
+
+//@[device] compile-flags: -Zunstable-options -Zoffload=Device
+
+#![feature(gpu_offload)]
+extern crate std;
+#[prelude_import]
+use ::std::prelude::rust_2015::*;
+
+use std::offload::offload_kernel;
+
+fn kernel() {
+    #[rustc_offload_kernel]
+    #[inline(never)]
+    fn inner_kernel() {
+
+        ::core::panicking::panic("not implemented")
+    }
+}
+fn main() {}

tests/pretty/offload/offload_kernel_nested.rs

@@ -0,0 +1,20 @@
+//@ only-nightly
+//@ revisions: host device
+
+//@ pretty-mode:expanded
+//@ pretty-compare-only
+//@[host] pp-exact:offload_kernel_nested.host.pp
+//@[device] pp-exact:offload_kernel_nested.device.pp
+
+//@[device] compile-flags: -Zunstable-options -Zoffload=Device
+
+#![feature(gpu_offload)]
+
+use std::offload::offload_kernel;
+
+fn kernel() {
+    #[offload_kernel]
+    fn inner_kernel() {}
+}
+
+fn main() {}

tests/ui/offload/offload_kernel_illegal.rs

@@ -0,0 +1,21 @@
+#![feature(gpu_offload)]
+
+fn dummy() {
+    #[core::offload::offload_kernel]
+    //~^ ERROR macro attributes on statements are unstable
+    let mut x = 5;
+    //~^ ERROR offload_kernel must be applied to function
+
+    #[core::offload::offload_kernel]
+    x = x + 3;
+    //~^^  ERROR attributes on expressions are experimental [E0658]
+    //~|   ERROR macro attributes on expressions are unstable
+    //~^^^ ERROR offload_kernel must be applied to function
+
+    #[core::offload::offload_kernel]
+    //~^ ERROR macro attributes on statements are unstable
+    let add_one_v2 = |x: u32| -> u32 { x + 1 };
+    //~^ ERROR offload_kernel must be applied to function
+}
+
+fn main() {}

tests/ui/offload/offload_kernel_illegal.stderr

@@ -0,0 +1,61 @@
+error[E0658]: attributes on expressions are experimental
+  --> $DIR/offload_kernel_illegal.rs:9:5
+   |
+LL |     #[core::offload::offload_kernel]
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: see issue #15701 <https://github.com/rust-lang/rust/issues/15701> for more information
+   = help: add `#![feature(stmt_expr_attributes)]` to the crate attributes to enable
+   = note: this compiler was built on YYYY-MM-DD; consider upgrading it if it is out of date
+
+error[E0658]: macro attributes on statements are unstable
+  --> $DIR/offload_kernel_illegal.rs:4:5
+   |
+LL |     #[core::offload::offload_kernel]
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: see issue #54727 <https://github.com/rust-lang/rust/issues/54727> for more information
+   = help: add `#![feature(proc_macro_hygiene)]` to the crate attributes to enable
+   = note: this compiler was built on YYYY-MM-DD; consider upgrading it if it is out of date
+
+error: offload_kernel must be applied to function
+  --> $DIR/offload_kernel_illegal.rs:6:5
+   |
+LL |     let mut x = 5;
+   |     ^^^^^^^^^^^^^^
+
+error[E0658]: macro attributes on expressions are unstable
+  --> $DIR/offload_kernel_illegal.rs:9:5
+   |
+LL |     #[core::offload::offload_kernel]
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: see issue #54727 <https://github.com/rust-lang/rust/issues/54727> for more information
+   = help: add `#![feature(proc_macro_hygiene)]` to the crate attributes to enable
+   = note: this compiler was built on YYYY-MM-DD; consider upgrading it if it is out of date
+
+error: offload_kernel must be applied to function
+  --> $DIR/offload_kernel_illegal.rs:10:5
+   |
+LL |     x = x + 3;
+   |     ^
+
+error[E0658]: macro attributes on statements are unstable
+  --> $DIR/offload_kernel_illegal.rs:15:5
+   |
+LL |     #[core::offload::offload_kernel]
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: see issue #54727 <https://github.com/rust-lang/rust/issues/54727> for more information
+   = help: add `#![feature(proc_macro_hygiene)]` to the crate attributes to enable
+   = note: this compiler was built on YYYY-MM-DD; consider upgrading it if it is out of date
+
+error: offload_kernel must be applied to function
+  --> $DIR/offload_kernel_illegal.rs:17:5
+   |
+LL |     let add_one_v2 = |x: u32| -> u32 { x + 1 };
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to 7 previous errors
+
+For more information about this error, try `rustc --explain E0658`.