Simplify mono entry-point computation for GPU Offload

compiler/rustc_monomorphize/src/collector.rs

@@ -1484,7 +1484,9 @@ fn collect_roots(tcx: TyCtxt<'_>, mode: MonoItemCollectionStrategy) -> Vec<MonoI
     debug!("collecting roots");
     let mut roots = MonoItems::new();
 
-    // Read the manifest and add the recorded kernel instantiations as roots so they are codegened.
+    // By design, the `HostMetadata` contains all the kernels launched by the host. For generic
+    // kernels, it also includes the used instantiations. These are the *only* roots for offload
+    // Device compilation, which makes it very cheap.
     if let Some(manifest_path) = tcx.sess.opts.unstable_opts.offload.iter().find_map(|o| {
         if let rustc_session::config::Offload::Device(p) = o
             && !p.is_empty()
@@ -1510,9 +1512,7 @@ fn collect_roots(tcx: TyCtxt<'_>, mode: MonoItemCollectionStrategy) -> Vec<MonoI
                 });
             }
         }
-    }
-
-    {
+    } else {
         let entry_fn = tcx.entry_fn(());
 
         debug!("collect_roots: entry_fn = {:?}", entry_fn);
@@ -1536,50 +1536,6 @@ fn collect_roots(tcx: TyCtxt<'_>, mode: MonoItemCollectionStrategy) -> Vec<MonoI
         collector.push_extra_entry_roots();
     }
 
-    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 crate_items = tcx.hir_crate_items(());
-        for id in crate_items.free_items() {
-            if !matches!(tcx.def_kind(id.owner_id), DefKind::Fn | DefKind::AssocFn) {
-                continue;
-            }
-            let def_id = id.owner_id.to_def_id();
-            if !tcx.generics_of(def_id).requires_monomorphization(tcx)
-                && tcx.codegen_fn_attrs(def_id).flags.intersects(CodegenFnAttrFlags::OFFLOAD_KERNEL)
-            {
-                roots.push(dummy_spanned(MonoItem::Fn(Instance::mono(tcx, def_id))));
-            }
-        }
-        for id in crate_items.impl_items() {
-            if !matches!(tcx.def_kind(id.owner_id), DefKind::Fn | DefKind::AssocFn) {
-                continue;
-            }
-            let def_id = id.owner_id.to_def_id();
-            if !tcx.generics_of(def_id).requires_monomorphization(tcx)
-                && tcx.codegen_fn_attrs(def_id).flags.intersects(CodegenFnAttrFlags::OFFLOAD_KERNEL)
-            {
-                roots.push(dummy_spanned(MonoItem::Fn(Instance::mono(tcx, def_id))));
-            }
-        }
-        for id in crate_items.trait_items() {
-            if !matches!(tcx.def_kind(id.owner_id), DefKind::Fn | DefKind::AssocFn) {
-                continue;
-            }
-            let def_id = id.owner_id.to_def_id();
-            if !tcx.generics_of(def_id).requires_monomorphization(tcx)
-                && tcx.codegen_fn_attrs(def_id).flags.intersects(CodegenFnAttrFlags::OFFLOAD_KERNEL)
-            {
-                roots.push(dummy_spanned(MonoItem::Fn(Instance::mono(tcx, def_id))));
-            }
-        }
-    }
-
     // We can only codegen items that are instantiable - items all of
     // whose predicates hold. Luckily, items that aren't instantiable
     // can't actually be used, so we can just skip codegenning them.

compiler/rustc_session/src/config.rs

@@ -2854,6 +2854,13 @@ pub fn build_session_options(early_dcx: &mut EarlyDiagCtxt, matches: &getopts::M
         )
     }
 
+    // Under Offload the device pass codegens only what launched kernels reach, and the device image
+    // is linked from that crate's bitcode alone. With this threshold a GPU kernel can still
+    // reach into `core` and its dependencies without them having to codegen anything themselves.
+    if unstable_opts.offload.iter().any(|o| matches!(o, Offload::Device(_))) {
+        unstable_opts.cross_crate_inline_threshold = InliningThreshold::Always;
+    }
+
     let target_triple = parse_target_triple(early_dcx, matches);
 
     // Ensure `-Z unstable-options` is required when using the unstable `-C link-self-contained` and

tests/run-make/offload-device-manifest-roots/dep.rs

@@ -0,0 +1,13 @@
+#![crate_type = "rlib"]
+
+// This tests two things.
+//
+// First, if we directly compile this for the device, then helper should not end up in the LLVM-IR,
+// despite being pub. This makes sure that our mono collector overwrite works.
+//
+// Second, when we compile host.rs, this file becomes a dependency. In that case its MIR should be
+// available, since our Device pass forces `InliningThreshold::Always`.
+#[inline(never)]
+pub fn helper(x: &mut f32) {
+    *x = 1.0;
+}

tests/run-make/offload-device-manifest-roots/host.rs

@@ -0,0 +1,33 @@
+#![feature(gpu_offload, rustc_attrs)]
+#![allow(internal_features, dead_code)]
+
+extern crate dep;
+
+// Launched by `main`, so the manifest lists it and the device pass compiles it, together with
+// everything it reaches.
+#[rustc_offload_kernel]
+fn launched(x: &mut f32) {
+    dep::helper(x);
+}
+
+// Never launched, so it is not in the manifest and the device pass drops it, despite attribute.
+#[rustc_offload_kernel]
+fn dormant(x: &mut f32) {
+    *x = 2.0;
+}
+
+// Public, but not reachable from any launched kernel.
+pub fn plain_pub(x: &mut f32) {
+    *x = 3.0;
+}
+
+// Offload previously kept the normal mono roots, so we'd need to add various `#[cfg(...)]`
+// attributes to functions like main that shouldn't end up on the Device. This tests that our new
+// mono logic keeps working and correctly disregards this function during device compilation.
+fn main() {
+    let mut x = 0.0f32;
+    core::offload::offload! {
+        kernel = launched,
+        args = (&mut x,),
+    }
+}

tests/run-make/offload-device-manifest-roots/rmake.rs

@@ -0,0 +1,70 @@
+// With a manifest, the device pass compiles exactly the kernels the host launches and whatever
+// they reach, nothing else. That includes reaching into dependencies: their functions are made
+// available for instantiation by `-Zoffload=Device` itself, so the dependencies do not have to
+// codegen anything for the device, and the kernel crate ends up self-contained.
+
+//@ needs-offload
+
+use std::path::Path;
+
+use run_make_support::{rfs, rustc};
+
+fn defines(ll: &str, name: &str) -> bool {
+    ll.lines().any(|line| line.starts_with("define") && line.contains(name))
+}
+
+fn main() {
+    // Host build of the dependency, so the metadata pass can resolve `dep::helper`.
+    rustc().input("dep.rs").run();
+
+    // Pass 1: Our manifest now includes both generic and non-generic kernels that were launched.
+    rustc()
+        .input("host.rs")
+        .extern_("dep", "libdep.rlib")
+        .arg("-Zunstable-options")
+        .arg("-Zoffload=HostMetadata=kernels.manifest")
+        .arg("-Csymbol-mangling-version=v0")
+        .arg("-Clto=fat")
+        .emit("metadata")
+        .run();
+
+    // Pass 2a: the dependency for the device. It reads a manifest naming kernels in a crate it
+    // has never heard of, has no launched kernels of its own, and so has nothing to codegen.
+    // FIXME(offload): In the future, we should add better errorhandling here. It's fine to not find
+    // the kernels mentioned in the manifest if this is just a dep. However, if the Manifest entry
+    // names this crate and the path to the Kernel does not resolve, then we should error. The
+    // decoder should be able to tell the difference between both cases.
+    rustc()
+        .input("dep.rs")
+        .arg("-Zunstable-options")
+        .arg("-Zoffload=Device=kernels.manifest")
+        .arg("-Csymbol-mangling-version=v0")
+        .codegen_units(1)
+        .emit("link,llvm-ir")
+        .out_dir("device")
+        .run();
+    if Path::new("device/dep.ll").exists() {
+        let dep_ll = rfs::read_to_string("device/dep.ll");
+        assert!(!defines(&dep_ll, "helper"), "`dep::helper` was codegened in its own crate");
+    }
+
+    // Pass 2b: the kernel crate for the device.
+    rustc()
+        .input("host.rs")
+        .extern_("dep", "device/libdep.rlib")
+        .arg("-Zunstable-options")
+        .arg("-Zoffload=Device=kernels.manifest")
+        .arg("-Csymbol-mangling-version=v0")
+        .arg("-Clto=fat")
+        .codegen_units(1)
+        .emit("llvm-ir")
+        .out_dir("device")
+        .run();
+    let ll = rfs::read_to_string("device/host.ll");
+    assert!(defines(&ll, "launched"), "the launched kernel is missing");
+    assert!(defines(&ll, "helper"), "`dep::helper` was not instantiated in the kernel crate");
+    assert!(!defines(&ll, "dormant"), "an unlaunched kernel was compiled");
+    assert!(!defines(&ll, "plain_pub"), "a function no kernel reaches was compiled");
+    // The entry point would be a root under the usual rules; on the device it is not launched.
+    assert!(!defines(&ll, "main"), "the host entry point was compiled for the device");
+}