Add support for generics in offloadand remove no_mangle attribute

Cargo.lock

@@ -4501,12 +4501,14 @@ name = "rustc_monomorphize"
 version = "0.0.0"
 dependencies = [
  "rustc_abi",
+ "rustc_ast",
  "rustc_data_structures",
  "rustc_errors",
  "rustc_hir",
  "rustc_index",
  "rustc_macros",
  "rustc_middle",
+ "rustc_serialize",
  "rustc_session",
  "rustc_span",
  "rustc_target",

compiler/rustc_builtin_macros/src/offload.rs

@@ -1,6 +1,6 @@
+use rustc_ast::ast;
 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::{DUMMY_SP, Ident, Span, sym};
@@ -9,7 +9,12 @@ use thin_vec::thin_vec;
 use crate::diagnostics;
 
 fn compile_for_device(ecx: &mut ExtCtxt<'_>) -> bool {
-    ecx.sess.opts.unstable_opts.offload.contains(&Offload::Device)
+    ecx.sess
+        .opts
+        .unstable_opts
+        .offload
+        .iter()
+        .any(|o| matches!(o, Offload::Device | Offload::DeviceWithManifest(_)))
 }
 
 fn outer_normal_attr(normal: &Box<ast::NormalAttr>, id: ast::AttrId, span: Span) -> ast::Attribute {
@@ -45,7 +50,6 @@ fn extract_fn(
 /// This expands to the host-side function:
 ///
 /// ```
-/// #[unsafe(no_mangle)]
 /// #[inline(never)]
 /// fn foo(_: &[f32], _: &[f32], _: *mut f32) {
 ///     ::core::panicking::panic("not implemented")
@@ -56,7 +60,6 @@ fn extract_fn(
 ///
 /// ```
 /// #[rustc_offload_kernel]
-/// #[unsafe(no_mangle)]
 /// unsafe extern "gpu-kernel" fn foo(a: &[f32], b: &[f32], c: *mut f32) {
 ///     *c = a[0] + b[0];
 /// }
@@ -110,24 +113,9 @@ pub(crate) fn expand_kernel(
         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::AttrArgs::Empty,
-        span,
-    };
-
-    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),
-        );
+        let mut item =
+            ecx.item(span, thin_vec![rustc_offload_kernel.clone()], ast::ItemKind::Fn(device_fn));
         item.vis = vis.clone();
         Annotatable::Item(item)
     };
@@ -187,12 +175,12 @@ pub(crate) fn expand_kernel(
     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));
+        let mut item = ecx.item(
+            span,
+            thin_vec![rustc_offload_kernel, inline_never],
+            ast::ItemKind::Fn(host_fn),
+        );
         item.vis = vis.clone();
         Annotatable::Item(item)
     };

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -735,7 +735,12 @@ pub(crate) unsafe fn llvm_optimize(
         llvm::set_value_name(new_fn, &name);
     }
 
-    if cgcx.target_is_like_gpu && config.offload.contains(&config::Offload::Device) {
+    if cgcx.target_is_like_gpu
+        && config
+            .offload
+            .iter()
+            .any(|o| matches!(o, config::Offload::Device | config::Offload::DeviceWithManifest(_)))
+    {
         let cx =
             SimpleCx::new(module.module_llvm.llmod(), module.module_llvm.llcx, cgcx.pointer_size);
         for func in cx.get_functions() {
@@ -806,7 +811,12 @@ pub(crate) unsafe fn llvm_optimize(
         )
     };
 
-    if cgcx.target_is_like_gpu && config.offload.contains(&config::Offload::Device) {
+    if cgcx.target_is_like_gpu
+        && config
+            .offload
+            .iter()
+            .any(|o| matches!(o, config::Offload::Device | config::Offload::DeviceWithManifest(_)))
+    {
         let device_path = cgcx.output_filenames.path(OutputType::Object);
         let device_dir = device_path.parent().unwrap();
         let device_out = device_dir.join("device.bin");

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -26,7 +26,9 @@ use rustc_session::config::CrateType;
 use rustc_session::diagnostics::feature_err;
 use rustc_session::lint::builtin::DEPRECATED_LLVM_INTRINSIC;
 use rustc_span::{ErrorGuaranteed, Span, Symbol, sym};
-use rustc_symbol_mangling::{mangle_internal_symbol, symbol_name_for_instance_in_crate};
+use rustc_symbol_mangling::{
+    mangle_internal_symbol, mangle_offload_export, symbol_name_for_instance_in_crate,
+};
 use rustc_target::callconv::PassMode;
 use rustc_target::spec::Arch;
 use tracing::debug;
@@ -1851,7 +1853,7 @@ fn codegen_offload<'ll, 'tcx>(
         _ => panic!("unparsable"),
     };
     let args = get_args_from_tuple(bx, args[4], fn_target);
-    let target_symbol = symbol_name_for_instance_in_crate(tcx, fn_target, LOCAL_CRATE);
+    let target_symbol = mangle_offload_export(tcx, fn_target);
 
     let sig = tcx.fn_sig(fn_target.def_id()).skip_binder();
     let sig = tcx.instantiate_bound_regions_with_erased(sig);

compiler/rustc_codegen_ssa/src/back/symbol_export.rs

@@ -2,6 +2,7 @@ use std::collections::hash_map::Entry::*;
 
 use rustc_abi::{CanonAbi, X86Call};
 use rustc_ast::expand::allocator::{AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, global_fn_name};
+use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::unord::UnordMap;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LOCAL_CRATE, LocalDefId};
@@ -17,7 +18,7 @@ use rustc_middle::ty::{
 use rustc_middle::util::Providers;
 use rustc_session::config::CrateType;
 use rustc_span::Span;
-use rustc_symbol_mangling::mangle_internal_symbol;
+use rustc_symbol_mangling::{is_offload_kernel, mangle_internal_symbol};
 use rustc_target::spec::{Arch, Os, TlsModel};
 use tracing::debug;
 
@@ -225,6 +226,51 @@ fn exported_non_generic_symbols_provider_local<'tcx>(
         ));
     }
 
+    let is_device_offload = tcx.sess.opts.unstable_opts.offload.iter().any(|o| {
+        matches!(
+            o,
+            rustc_session::config::Offload::DeviceWithManifest(_)
+                | rustc_session::config::Offload::Device
+        )
+    });
+    if is_device_offload {
+        let crate_items = tcx.hir_crate_items(());
+        let mut seen: rustc_data_structures::fx::FxHashSet<DefId> = symbols
+            .iter()
+            .filter_map(|(s, _)| match s {
+                ExportedSymbol::NonGeneric(d) => Some(*d),
+                _ => None,
+            })
+            .collect();
+
+        let mut try_emit_offload_kernel = |def_id: DefId, seen: &mut FxHashSet<DefId>| {
+            if !matches!(tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+                return;
+            }
+            if !tcx.generics_of(def_id).requires_monomorphization(tcx)
+                && is_offload_kernel(tcx.codegen_fn_attrs(def_id))
+                && seen.insert(def_id)
+            {
+                symbols.push((
+                    ExportedSymbol::NonGeneric(def_id),
+                    SymbolExportInfo {
+                        level: SymbolExportLevel::C,
+                        kind: SymbolExportKind::Text,
+                        used: false,
+                        rustc_std_internal_symbol: false,
+                    },
+                ));
+            }
+        };
+
+        for id in crate_items.free_items() {
+            try_emit_offload_kernel(id.owner_id.to_def_id(), &mut seen);
+        }
+        for id in crate_items.impl_items() {
+            try_emit_offload_kernel(id.owner_id.to_def_id(), &mut seen);
+        }
+    }
+
     // Sort so we get a stable incr. comp. hash.
     symbols.sort_by_cached_key(|s| s.0.symbol_name_for_local_instance(tcx));
 
@@ -241,7 +287,16 @@ fn exported_generic_symbols_provider_local<'tcx>(
 
     let mut symbols: Vec<_> = vec![];
 
-    if tcx.local_crate_exports_generics() {
+    let export_generics = tcx.local_crate_exports_generics();
+    let is_device_offload = tcx.sess.opts.unstable_opts.offload.iter().any(|o| {
+        matches!(
+            o,
+            rustc_session::config::Offload::DeviceWithManifest(_)
+                | rustc_session::config::Offload::Device
+        )
+    });
+
+    if export_generics || is_device_offload {
         use rustc_hir::attrs::Linkage;
         use rustc_middle::mono::{MonoItem, Visibility};
         use rustc_middle::ty::InstanceKind;
@@ -287,6 +342,14 @@ fn exported_generic_symbols_provider_local<'tcx>(
                     })
             };
 
+        let is_offload_instance = |mono_item: &MonoItem<'tcx>| {
+            if let MonoItem::Fn(instance) = mono_item {
+                is_offload_kernel(tcx.codegen_fn_attrs(instance.def_id()))
+            } else {
+                false
+            }
+        };
+
         // The symbols created in this loop are sorted below it
         #[allow(rustc::potential_query_instability)]
         for (mono_item, data) in cgus.iter().flat_map(|cgu| cgu.items().iter()) {
@@ -302,7 +365,9 @@ fn exported_generic_symbols_provider_local<'tcx>(
                 continue;
             }
 
-            if !tcx.sess.opts.share_generics() {
+            let item_is_offload = is_offload_instance(mono_item);
+
+            if !item_is_offload && !tcx.sess.opts.share_generics() {
                 if tcx.codegen_fn_attrs(mono_item.def_id()).inline
                     == rustc_hir::attrs::InlineAttr::Never
                 {
@@ -319,15 +384,22 @@ fn exported_generic_symbols_provider_local<'tcx>(
                 MonoItem::Fn(Instance { def: InstanceKind::Item(def), args }) => {
                     let has_generics = args.non_erasable_generics().next().is_some();
 
-                    let should_export =
-                        has_generics && is_instantiable_downstream(Some(def), &args);
+                    let should_export = if item_is_offload {
+                        has_generics
+                    } else {
+                        has_generics && is_instantiable_downstream(Some(def), &args)
+                    };
 
                     if should_export {
                         let symbol = ExportedSymbol::Generic(def, args);
                         symbols.push((
                             symbol,
                             SymbolExportInfo {
-                                level: SymbolExportLevel::Rust,
+                                level: if item_is_offload {
+                                    SymbolExportLevel::C
+                                } else {
+                                    SymbolExportLevel::Rust
+                                },
                                 kind: SymbolExportKind::Text,
                                 used: false,
                                 rustc_std_internal_symbol: false,

compiler/rustc_interface/src/passes.rs

@@ -1299,11 +1299,25 @@ pub(crate) fn start_codegen<'tcx>(
 
     let metadata = rustc_metadata::fs::encode_and_write_metadata(tcx);
 
+    let is_host_metadata = tcx
+        .sess
+        .opts
+        .unstable_opts
+        .offload
+        .iter()
+        .any(|o| matches!(o, rustc_session::config::Offload::HostMetadata(_)));
+
     let codegen = tcx.sess.time("codegen_crate", || {
-        if tcx.sess.opts.unstable_opts.no_codegen || !tcx.sess.opts.output_types.should_codegen() {
-            // Skip crate items and just output metadata in -Z no-codegen mode.
+        if tcx.sess.opts.unstable_opts.no_codegen
+            || !tcx.sess.opts.output_types.should_codegen()
+            || is_host_metadata
+        {
             tcx.sess.dcx().abort_if_errors();
 
+            if is_host_metadata {
+                rustc_monomorphize::write_host_metadata_offload_manifest(tcx);
+            }
+
             // Linker::link will skip join_codegen in case of a CodegenResults Any value.
             Box::new(CompiledModules { modules: vec![], allocator_module: None })
         } else {

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -295,6 +295,8 @@ impl CodegenFnAttrs {
             // note: for these we do also set a symbol name so technically also handled by the
             // condition below. However, I think that regardless these should be treated as extern.
             || self.flags.contains(CodegenFnAttrFlags::EXTERNALLY_IMPLEMENTABLE_ITEM)
+            // `#[rustc_offload_kernel]`: this item is an externally-launched kernel entry point.
+            || self.flags.contains(CodegenFnAttrFlags::OFFLOAD_KERNEL)
             || self.symbol_name.is_some()
             || match self.linkage {
                 // These are private, so make sure we don't try to consider

compiler/rustc_monomorphize/Cargo.toml

@@ -6,12 +6,14 @@ edition = "2024"
 [dependencies]
 # tidy-alphabetical-start
 rustc_abi = { path = "../rustc_abi" }
+rustc_ast = { path = "../rustc_ast" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
 rustc_hir = { path = "../rustc_hir" }
 rustc_index = { path = "../rustc_index" }
 rustc_macros = { path = "../rustc_macros" }
 rustc_middle = { path = "../rustc_middle" }
+rustc_serialize = { path = "../rustc_serialize" }
 rustc_session = { path = "../rustc_session" }
 rustc_span = { path = "../rustc_span" }
 rustc_target = { path = "../rustc_target" }