wip offload frontend

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,159 @@
+use crate::errors;
+use rustc_ast::{ForeignMod, ast};
+use rustc_expand::base::Annotatable;
+use rustc_expand::base::ExtCtxt;
+use rustc_session::config::Offload;
+use rustc_span::{DUMMY_SP, Ident};
+use rustc_span::{Span, sym};
+use thin_vec::thin_vec;
+
+/*
+```
+#[offload_kernel]
+fn foo(..args) {
+    // body
+}
+```
+
+expands to:
+```
+#[cfg(host)]
+unsafe extern "C" {
+    pub fn foo(..args)
+}
+
+#[cfg(device)]
+#[rustc_offload_kernel]
+unsafe extern "gpu-kernel" fn foo(args) {
+    // body
+}
+```
+*/
+fn is_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(box ast::Fn { sig, ident, generics, body, .. }) => {
+                Some((iitem.vis.clone(), sig.clone(), *ident, generics.clone(), body.clone()))
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+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,
+    );
+
+    let device_item = {
+        let mut item =
+            ecx.item(span, thin_vec![rustc_offload_kernel], ast::ItemKind::Fn(device_fn));
+        item.vis = vis.clone();
+        Annotatable::Item(item)
+    };
+
+    // host function
+    let host_fn = Box::new(ast::Fn {
+        defaultness: ast::Defaultness::Implicit,
+        sig: sig.clone(),
+        ident,
+        generics: generics.clone(),
+        contract: None,
+        body: None,
+        define_opaque: None,
+        eii_impls: Default::default(),
+    });
+
+    let foreign_fn = ast::ForeignItem {
+        attrs: Default::default(),
+        id: ast::DUMMY_NODE_ID,
+        span,
+        vis: vis.clone(),
+        kind: ast::ForeignItemKind::Fn(host_fn),
+        tokens: None,
+    };
+
+    let extern_c_lit = ast::StrLit {
+        symbol: sym::C,
+        suffix: None,
+        symbol_unescaped: sym::C,
+        style: ast::StrStyle::Cooked,
+        span,
+    };
+
+    let foreign_mod = ForeignMod {
+        abi: Some(extern_c_lit),
+        safety: ast::Safety::Unsafe(span),
+        items: thin_vec![Box::new(foreign_fn)],
+        extern_span: DUMMY_SP,
+    };
+
+    let host_item = {
+        let mut item = ecx.item(span, thin_vec![], ast::ItemKind::ForeignMod(foreign_mod));
+        item.vis = vis;
+        Annotatable::Item(item)
+    };
+
+    if is_device(ecx) { vec![device_item] } else { vec![host_item] }
+}

compiler/rustc_span/src/symbol.rs

@@ -1034,6 +1034,7 @@ symbols! {
         global_asm,
         global_registration,
         globs,
+        gpu_kernel,
         gpu_launch_sized_workgroup_mem,
         gt,
         guard,
@@ -1172,6 +1173,7 @@ symbols! {
         linkonce,
         linkonce_odr,
         lint_reasons,
+        linux,
         literal,
         little, big,
         load,
@@ -1416,6 +1418,7 @@ symbols! {
         of,
         off,
         offload,
+        offload_kernel,
         offset,
         offset_of,
         offset_of_enum,

library/core/src/lib.rs

@@ -235,6 +235,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

@@ -881,6 +881,14 @@ macro_rules! todo {
     };
 }
 
+/// TODO(offload): add docs
+#[macro_export]
+#[unstable(feature = "gpu_offload", issue = "131513")]
+#[allow_internal_unstable(core_intrinsics)]
+macro_rules! offload {
+    () => {};
+}
+
 /// Definitions of built-in macros.
 ///
 /// Most of the macro properties (stability, visibility, etc.) are taken from the source code here,
@@ -1627,6 +1635,15 @@ pub(crate) mod builtin {
         /* compiler built-in */
     }
 
+    /// TODO(Sa4dUs): add docs
+    #[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/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

@@ -286,6 +286,7 @@
 #![feature(f16)]
 #![feature(f128)]
 #![feature(ffi_const)]
+#![feature(gpu_offload)]
 #![feature(intra_doc_pointers)]
 #![feature(lang_items)]
 #![feature(link_cfg)]
@@ -658,6 +659,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.

tests/pretty/offload/offload_kernel.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.host.pp
+//@[device] pp-exact:offload_kernel.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;
+
+#[rustc_offload_kernel]
+unsafe extern "gpu_kernel" fn foo(a: &[f32], b: &[f32], c: &mut [f32]) {
+    *c[0] = a[0] + b[0];
+}
+
+fn main() {}

tests/pretty/offload/offload_kernel.host.pp

@@ -0,0 +1,24 @@
+#![feature(prelude_import)]
+#![no_std]
+//@ only-nightly
+//@ revisions: host device
+
+//@ pretty-mode:expanded
+//@ pretty-compare-only
+//@[host] pp-exact:offload_kernel.host.pp
+//@[device] pp-exact:offload_kernel.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;
+
+unsafe extern "C" {
+    fn foo(a: &[f32], b: &[f32], c: &mut [f32]);
+}
+
+fn main() {}

tests/pretty/offload/offload_kernel.rs

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