From 9d291cbaab5e7c0f857ff19186187c99069def77 Mon Sep 17 00:00:00 2001
From: Manuel Drehwald <git@manuel.drehwald.info>
Date: Fri, 4 Sep 2026 00:35:20 -0400
Subject: [PATCH] Simplify mono entry-point computation for GPU Offload

`std::offload` previously compiled too much code for the GPU.
We first have a Host (CPU) compilation pass, in which we walk all
`offload` intrinsics that launch a Kernel. We create a Manifest, and
if we ever launch a generic kernel, than we add the kernel along with
it's needed instantiations. If the launched kernel wasn't generic, we
would skip it and not add it to the Manifest. During the following
Device (GPU) compilation pass, we'd then add add all those generic
kernels to the mono collector and force their instantiations. We'd also
walk all functions and check if any non-generic functions would have a
`rustc_offload_kernel` attribute, in which case we also add them to our
collector. We'd also codegen all the other mono roots.

After this PR, we simply write all launched Kernels into the Manifest,
including the non-generic ones. When compiling for the GPU Device, we
now drop all other mono roots and only add the ones from our Manifest.

This makes compiling core under build-std extremely fast. To make core
usage more reliable, `-Zoffload=Device` now also implies
`-Zcross-crate-inline-threshold=always`, so functions will always be
available. Our more aggressive mono roots computation makes it cheap.
---
 compiler/rustc_monomorphize/src/collector.rs  | 52 ++------------
 compiler/rustc_session/src/config.rs          |  7 ++
 .../offload-device-manifest-roots/dep.rs      | 13 ++++
 .../offload-device-manifest-roots/host.rs     | 33 +++++++++
 .../offload-device-manifest-roots/rmake.rs    | 70 +++++++++++++++++++
 5 files changed, 127 insertions(+), 48 deletions(-)
 create mode 100644 tests/run-make/offload-device-manifest-roots/dep.rs
 create mode 100644 tests/run-make/offload-device-manifest-roots/host.rs
 create mode 100644 tests/run-make/offload-device-manifest-roots/rmake.rs

diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index 5d7820df622b7..86830343f1ea2 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/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.
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
index efbffa8c0e486..ee6a8fde002a3 100644
--- a/compiler/rustc_session/src/config.rs
+++ b/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
diff --git a/tests/run-make/offload-device-manifest-roots/dep.rs b/tests/run-make/offload-device-manifest-roots/dep.rs
new file mode 100644
index 0000000000000..335f4a456776f
--- /dev/null
+++ b/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;
+}
diff --git a/tests/run-make/offload-device-manifest-roots/host.rs b/tests/run-make/offload-device-manifest-roots/host.rs
new file mode 100644
index 0000000000000..2d3d590b9a2f4
--- /dev/null
+++ b/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,),
+    }
+}
diff --git a/tests/run-make/offload-device-manifest-roots/rmake.rs b/tests/run-make/offload-device-manifest-roots/rmake.rs
new file mode 100644
index 0000000000000..a6c5c38bc5d54
--- /dev/null
+++ b/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");
+}