Add support for generics in offloadand remove no_mangle attribute

Cargo.lock

@@ -4532,12 +4532,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::{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,23 +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,
-    };
-
-    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)
     };
@@ -185,12 +174,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;
@@ -1850,7 +1852,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,
@@ -534,11 +606,19 @@ fn symbol_export_level(tcx: TyCtxt<'_>, sym_def_id: DefId) -> SymbolExportLevel
     // are not considered for export
     let codegen_fn_attrs = tcx.codegen_fn_attrs(sym_def_id);
     let is_extern = codegen_fn_attrs.contains_extern_indicator();
+    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
+        )
+    });
+    let is_offload = is_device_offload && is_offload_kernel(codegen_fn_attrs);
     let std_internal =
         codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::RUSTC_STD_INTERNAL_SYMBOL);
     let eii = codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::EXTERNALLY_IMPLEMENTABLE_ITEM);
 
-    if is_extern && !std_internal && !eii {
+    if (is_extern && !std_internal && !eii) || is_offload {
         let target = &tcx.sess.target.llvm_target;
         // WebAssembly cannot export data symbols, so reduce their export level
         // FIXME(jdonszelmann) don't do a substring match here.

compiler/rustc_codegen_ssa/src/back/write.rs

@@ -459,6 +459,29 @@ pub(crate) fn start_async_codegen<B: WriteBackendMethods>(
     }
 }
 
+/// Create an `OngoingCodegen` that has no coordinator thread and will finish
+/// immediately when joined. This is used for the offload host-metadata pass,
+/// that only need to run the monomorphization collector.
+pub(crate) fn empty_ongoing_codegen<B: WriteBackendMethods>(
+    backend: B,
+    tcx: TyCtxt<'_>,
+) -> OngoingCodegen<B> {
+    let (coordinator_send, _) = channel::<Message<B>>();
+    let (codegen_worker_send, codegen_worker_receive) = channel();
+    drop(codegen_worker_send);
+
+    let (shared_emitter, shared_emitter_main) = SharedEmitter::new();
+    drop(shared_emitter);
+
+    OngoingCodegen {
+        backend,
+        codegen_worker_receive,
+        shared_emitter_main,
+        coordinator: Coordinator { sender: coordinator_send, future: None, phantom: PhantomData },
+        output_filenames: Arc::clone(tcx.output_filenames(())),
+    }
+}
+
 fn copy_all_cgu_workproducts_to_incr_comp_cache_dir(
     sess: &Session,
     compiled_modules: &CompiledModules,
@@ -2091,7 +2114,15 @@ pub struct Coordinator<B: WriteBackendMethods> {
 
 impl<B: WriteBackendMethods> Coordinator<B> {
     fn join(mut self) -> std::thread::Result<Result<MaybeLtoModules<B>, ()>> {
-        self.future.take().unwrap().join()
+        if let Some(future) = self.future.take() {
+            future.join()
+        } else {
+            // Used for passes that do not codegen anything (e.g. the offload host-metadata pass).
+            Ok(Ok(MaybeLtoModules::NoLto(CompiledModules {
+                modules: vec![],
+                allocator_module: None,
+            })))
+        }
     }
 }
 

compiler/rustc_codegen_ssa/src/base.rs

@@ -41,8 +41,9 @@ use tracing::{debug, info};
 use crate::assert_module_sources::CguReuse;
 use crate::back::link::are_upstream_rust_objects_already_included;
 use crate::back::write::{
-    ComputedLtoType, OngoingCodegen, compute_per_cgu_lto_type, start_async_codegen,
-    submit_codegened_module_to_llvm, submit_post_lto_module_to_llvm, submit_pre_lto_module_to_llvm,
+    ComputedLtoType, OngoingCodegen, compute_per_cgu_lto_type, empty_ongoing_codegen,
+    start_async_codegen, submit_codegened_module_to_llvm, submit_post_lto_module_to_llvm,
+    submit_pre_lto_module_to_llvm,
 };
 use crate::common::{self, IntPredicate, RealPredicate, TypeKind};
 use crate::meth::load_vtable;
@@ -701,6 +702,21 @@ pub fn codegen_crate<
         tcx.dcx().emit_fatal(diagnostics::CpuRequired);
     }
 
+    // A `HostMetadata` pass only exists to collect the set of generic kernel instantiations
+    // required by the host and write the offload manifest.
+    let is_host_metadata = tcx
+        .sess
+        .opts
+        .unstable_opts
+        .offload
+        .iter()
+        .any(|o| matches!(o, rustc_session::config::Offload::HostMetadata(_)));
+
+    if is_host_metadata {
+        let _ = tcx.collect_and_partition_mono_items(());
+        return empty_ongoing_codegen(backend, tcx);
+    }
+
     if let Some(target_cpu) = &tcx.sess.opts.cg.target_cpu
         && tcx.sess.target.unsupported_cpus.contains(&target_cpu.into())
     {

compiler/rustc_middle/src/mono.rs

@@ -160,6 +160,12 @@ impl<'tcx> MonoItem<'tcx> {
             return InstantiationMode::GloballyShared { may_conflict: false };
         }
 
+        // Offload kernels are looked up by symbol name at runtime by the host.
+        // They must be emitted exactly once with external linkage.
+        if codegen_fn_attrs.flags.intersects(CodegenFnAttrFlags::OFFLOAD_KERNEL) {
+            return InstantiationMode::GloballyShared { may_conflict: false };
+        }
+
         // This is technically a heuristic even though it's in the "not a heuristic" part of
         // instantiation mode selection.
         // It is surely possible to untangle this; the root problem is that the way we instantiate