Improve Offload artifact creation

compiler/rustc_codegen_llvm/src/back/mod.rs

@@ -1,5 +1,6 @@
 pub(crate) mod archive;
 pub(crate) mod lto;
+pub(crate) mod offload;
 pub(crate) mod owned_target_machine;
 mod profiling;
 pub(crate) mod write;

compiler/rustc_codegen_llvm/src/back/offload.rs

@@ -0,0 +1,83 @@
+//! We take in a `device.bin` (from a previous device pass) and our current fat-lto LLVM (host)
+//! module.  We clone the host module, embed the device code in it, and write it out as a
+//! `host.o` object file. The `clang-linker-wrapper` then produces the final binary.
+
+use std::path::PathBuf;
+
+use rustc_codegen_ssa::ModuleCodegen;
+use rustc_codegen_ssa::back::write::CodegenContext;
+use rustc_data_structures::profiling::SelfProfilerRef;
+use rustc_errors::DiagCtxtHandle;
+use rustc_fs_util::path_to_c_string;
+use rustc_session::config::{self, OutputType};
+
+use crate::back::write::write_output_file;
+use crate::{ModuleLlvm, llvm};
+
+/// Embed the device image into the host module and write the result out as `host.o`.
+///
+/// Does nothing unless this is a host compilation with `-Zoffload=Host=<path>`.
+pub(crate) fn finalize_host_module(
+    cgcx: &CodegenContext,
+    prof: &SelfProfilerRef,
+    dcx: DiagCtxtHandle<'_>,
+    module: &ModuleCodegen<ModuleLlvm>,
+) {
+    if cgcx.target_is_like_gpu {
+        return;
+    }
+
+    let config = &cgcx.module_config;
+    let Some(device_path) = config
+        .offload
+        .iter()
+        .find_map(|o| if let config::Offload::Host(path) = o { Some(path) } else { None })
+    else {
+        return;
+    };
+
+    // This assumes that we previously compiled our kernels for a gpu target, which created a
+    // `device.bin` artifact. The user is supposed to provide us with a path to this artifact, we
+    // don't need any other artifacts from the previous run. We will embed this artifact into our
+    // LLVM-IR host module, to create a `host.o` ObjectFile, which we will write to disk.
+    // The last, not yet automated step uses the `clang-linker-wrapper` to process `host.o`.
+    let device_pathbuf = PathBuf::from(device_path);
+    if device_pathbuf.is_relative() {
+        dcx.emit_err(crate::diagnostics::OffloadWithoutAbsPath);
+    } else if device_pathbuf.file_name().and_then(|n| n.to_str()).is_some_and(|n| n != "device.bin")
+    {
+        dcx.emit_err(crate::diagnostics::OffloadWrongFileName);
+    } else if !device_pathbuf.exists() {
+        dcx.emit_err(crate::diagnostics::OffloadNonexistingPath);
+    }
+    let host_path = cgcx.output_filenames.path(OutputType::Object);
+    let host_dir = host_path.parent().unwrap();
+    let out_obj = host_dir.join("host.o");
+    let device_bin_c = path_to_c_string(device_pathbuf.as_path());
+
+    // Finalize host: lib.bc + device.bin -> host.o (host TM)
+    // We create a full clone of our LLVM host module, since we will embed the device IR
+    // into it, and this might break caching or incremental compilation otherwise.
+    let llmod2 = llvm::LLVMCloneModule(module.module_llvm.llmod());
+    let ok = unsafe {
+        llvm::RustOffloadWrapper::get_instance()
+            .llvm_rust_offload_embed_buffer_in_module(llmod2, device_bin_c.as_c_str())
+    };
+    if !ok {
+        dcx.emit_err(crate::diagnostics::OffloadEmbedFailed);
+    }
+    write_output_file(
+        dcx,
+        module.module_llvm.tm.raw(),
+        config.no_builtins,
+        llmod2,
+        &out_obj,
+        None,
+        llvm::FileType::ObjectFile,
+        prof,
+        true,
+    );
+    // We ignore cgcx.save_temps here and unconditionally always keep our `device.bin` artifact.
+    // Otherwise, recompiling the host code would fail since we deleted that device artifact
+    // in the previous host compilation, which would be confusing at best.
+}

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -50,7 +50,7 @@ pub(crate) fn llvm_err<'a>(dcx: DiagCtxtHandle<'_>, err: LlvmError<'a>) -> ! {
     }
 }
 
-fn write_output_file<'ll>(
+pub(crate) fn write_output_file<'ll>(
     dcx: DiagCtxtHandle<'_>,
     target: &'ll llvm::TargetMachine,
     no_builtins: bool,
@@ -835,61 +835,8 @@ pub(crate) unsafe fn llvm_optimize(
         }
     }
 
-    // This assumes that we previously compiled our kernels for a gpu target, which created a
-    // `device.bin` artifact. The user is supposed to provide us with a path to this artifact, we
-    // don't need any other artifacts from the previous run. We will embed this artifact into our
-    // LLVM-IR host module, to create a `host.o` ObjectFile, which we will write to disk.
-    // The last, not yet automated steps uses the `clang-linker-wrapper` to process `host.o`.
-    if !cgcx.target_is_like_gpu {
-        if let Some(device_path) = config
-            .offload
-            .iter()
-            .find_map(|o| if let config::Offload::Host(path) = o { Some(path) } else { None })
-        {
-            let device_pathbuf = PathBuf::from(device_path);
-            if device_pathbuf.is_relative() {
-                dcx.emit_err(crate::diagnostics::OffloadWithoutAbsPath);
-            } else if device_pathbuf
-                .file_name()
-                .and_then(|n| n.to_str())
-                .is_some_and(|n| n != "device.bin")
-            {
-                dcx.emit_err(crate::diagnostics::OffloadWrongFileName);
-            } else if !device_pathbuf.exists() {
-                dcx.emit_err(crate::diagnostics::OffloadNonexistingPath);
-            }
-            let host_path = cgcx.output_filenames.path(OutputType::Object);
-            let host_dir = host_path.parent().unwrap();
-            let out_obj = host_dir.join("host.o");
-            let device_bin_c = path_to_c_string(device_pathbuf.as_path());
-
-            // 2) Finalize host: lib.bc + device.bin -> host.o (host TM)
-            // We create a full clone of our LLVM host module, since we will embed the device IR
-            // into it, and this might break caching or incremental compilation otherwise.
-            let llmod2 = llvm::LLVMCloneModule(module.module_llvm.llmod());
-            let ok = unsafe {
-                llvm::RustOffloadWrapper::get_instance()
-                    .llvm_rust_offload_embed_buffer_in_module(llmod2, device_bin_c.as_c_str())
-            };
-            if !ok {
-                dcx.emit_err(crate::diagnostics::OffloadEmbedFailed);
-            }
-            write_output_file(
-                dcx,
-                module.module_llvm.tm.raw(),
-                config.no_builtins,
-                llmod2,
-                &out_obj,
-                None,
-                llvm::FileType::ObjectFile,
-                prof,
-                true,
-            );
-            // We ignore cgcx.save_temps here and unconditionally always keep our `device.bin` artifact.
-            // Otherwise, recompiling the host code would fail since we deleted that device artifact
-            // in the previous host compilation, which would be confusing at best.
-        }
-    }
+    // The host side runs in `crate::back::offload::finalize_host_module` instead: it needs the
+    // merged fat LTO module, which only exists once we get past this per-CGU pass.
     result.into_result().unwrap_or_else(|()| llvm_err(dcx, LlvmError::RunLlvmPasses))
 }
 

compiler/rustc_codegen_llvm/src/lib.rs

@@ -157,6 +157,8 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         let dcx = dcx.handle();
         back::lto::run_pass_manager(cgcx, &sess.prof, dcx, &mut module, false);
 
+        back::offload::finalize_host_module(cgcx, &sess.prof, dcx, &module);
+
         back::write::codegen(cgcx, &sess.prof, shared_emitter, module, &cgcx.module_config)
     }
     fn run_thin_lto(

compiler/rustc_session/src/session.rs

@@ -1056,6 +1056,15 @@ impl Session {
     /// Returns the number of codegen units that should be used for this
     /// compilation
     pub fn codegen_units(&self) -> CodegenUnits {
+        // The Device pass needs to bundle all Kernels into a `device.bin` artifact. We could
+        // compute it at the end of the fat-lto pass when all Modules are combined, but some
+        // build configs like rlib currently don't use fat-lto. We can either restrict the
+        // configs in which the Device pass can run, or we reduce CGU to 1 for the Device pass.
+        // This is the easier solution for now, especially given that kernels tend to be smaller.
+        if self.opts.unstable_opts.offload.iter().any(|o| matches!(o, config::Offload::Device(_))) {
+            return CodegenUnits::Default(1);
+        }
+
         if let Some(n) = self.opts.cli_forced_codegen_units {
             return CodegenUnits::User(n);
         }

tests/run-make/offload-device-single-cgu/device.rs

@@ -0,0 +1,23 @@
+#![feature(abi_gpu_kernel, rustc_attrs, no_core)]
+#![no_core]
+#![crate_type = "rlib"]
+
+extern crate minicore;
+
+// Partitioning assigns items to codegen units by module, so with `-Ccodegen-units=2` these two
+// kernels would land in separate CGUs.
+pub mod first {
+    #[unsafe(no_mangle)]
+    #[rustc_offload_kernel]
+    pub unsafe extern "gpu-kernel" fn kernel_in_first_module(x: *mut f32, k: f32) {
+        unsafe { *x = k };
+    }
+}
+
+pub mod second {
+    #[unsafe(no_mangle)]
+    #[rustc_offload_kernel]
+    pub unsafe extern "gpu-kernel" fn kernel_in_second_module(x: *mut f32, k: f32) {
+        unsafe { *x = k };
+    }
+}

tests/run-make/offload-device-single-cgu/rmake.rs

@@ -0,0 +1,42 @@
+// The device pass bundles each codegen unit's module into `device.bin`, and every CGU writes the
+// same file, so a crate split into several CGUs would keep only the kernels of whichever CGU was
+// written last. `-Zoffload=Device` therefore forces a single CGU, even over an explicit
+// `-Ccodegen-units`. Check that kernels from two modules, which partitioning would otherwise put
+// into two CGUs, both end up in the bundle. If we were to enforce fat-lto also for the device, we
+// could move the artifact creation to the linker, but enforcing a single CGU is easier for now.
+//
+// The crate is emitted as an rlib on purpose: `--emit=obj` together with `-o` already resets the
+// CGU count to one, which would hide the problem.
+
+//@ needs-offload
+//@ needs-llvm-components: amdgpu
+
+use run_make_support::{rfs, rustc, rustc_minicore};
+
+fn contains(haystack: &[u8], needle: &str) -> bool {
+    haystack.windows(needle.len()).any(|w| w == needle.as_bytes())
+}
+
+fn main() {
+    rustc_minicore()
+        .target("amdgcn-amd-amdhsa")
+        .target_cpu("gfx90a")
+        .output("libminicore.rlib")
+        .run();
+
+    rustc()
+        .input("device.rs")
+        .target("amdgcn-amd-amdhsa")
+        .target_cpu("gfx90a")
+        .arg("-Zunstable-options")
+        .arg("-Zoffload=Device")
+        .codegen_units(2)
+        .extern_("minicore", "libminicore.rlib")
+        .run();
+
+    // The bundle holds the module as bitcode, whose string table keeps symbol names verbatim.
+    let device_bin = rfs::read("device.bin");
+    for kernel in ["kernel_in_first_module", "kernel_in_second_module"] {
+        assert!(contains(&device_bin, kernel), "`{kernel}` is missing from `device.bin`");
+    }
+}