diff --git a/compiler/rustc_codegen_llvm/src/back/mod.rs b/compiler/rustc_codegen_llvm/src/back/mod.rs index 6cb89f80ab89a..1a6253bdba712 100644 --- a/compiler/rustc_codegen_llvm/src/back/mod.rs +++ b/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; diff --git a/compiler/rustc_codegen_llvm/src/back/offload.rs b/compiler/rustc_codegen_llvm/src/back/offload.rs new file mode 100644 index 0000000000000..7c1a657961839 --- /dev/null +++ b/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=`. +pub(crate) fn finalize_host_module( + cgcx: &CodegenContext, + prof: &SelfProfilerRef, + dcx: DiagCtxtHandle<'_>, + module: &ModuleCodegen, +) { + 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. +} diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs index 66b51d9184a79..12b8c63feab60 100644 --- a/compiler/rustc_codegen_llvm/src/back/write.rs +++ b/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)) } diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs index 552a91ffee071..04b170c6483fe 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/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( diff --git a/compiler/rustc_session/src/session.rs b/compiler/rustc_session/src/session.rs index 7db30b38f9be4..f1f58274ee134 100644 --- a/compiler/rustc_session/src/session.rs +++ b/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); } diff --git a/tests/run-make/offload-device-single-cgu/device.rs b/tests/run-make/offload-device-single-cgu/device.rs new file mode 100644 index 0000000000000..2defed6693dbc --- /dev/null +++ b/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 }; + } +} diff --git a/tests/run-make/offload-device-single-cgu/rmake.rs b/tests/run-make/offload-device-single-cgu/rmake.rs new file mode 100644 index 0000000000000..9adb7adcc3389 --- /dev/null +++ b/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`"); + } +}