Rollup merge of #162309 - sgasho:offload-clang-linker-wrapper, r=ZuseZ4

offload: automate manual clang-linker-wrapper step

automate manual clang-linker-wrapper step from https://rustc-dev-guide.rust-lang.org/offload/usage.html
extract bitcode from device.bin and then wraps it into the host module.
needs some refactoring.

r? @ZuseZ4

Author: JonathanBrouwer

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -612,6 +612,8 @@ pub(crate) unsafe fn llvm_optimize(
     let pgo_use_path = get_pgo_use_path(config);
     let pgo_sample_use_path = get_pgo_sample_use_path(config);
     let is_lto = opt_stage == llvm::OptStage::ThinLTO || opt_stage == llvm::OptStage::FatLTO;
+    let is_final_stage =
+        !matches!(opt_stage, llvm::OptStage::PreLinkFatLTO | llvm::OptStage::PreLinkThinLTO);
     let instr_profile_output_path = get_instr_profile_output_path(config);
     let sanitize_dataflow_abilist: Vec<_> = config
         .sanitizer_dataflow_abilist
@@ -840,7 +842,7 @@ pub(crate) unsafe fn llvm_optimize(
     // 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 !cgcx.target_is_like_gpu && is_final_stage {
         if let Some(device_path) = config
             .offload
             .iter()
@@ -866,10 +868,11 @@ pub(crate) unsafe fn llvm_optimize(
             // 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())
+                llvm::RustOffloadWrapper::get_instance().llvm_rust_offload_embed_buffer_in_module(
+                    module.module_llvm.llmod(),
+                    device_bin_c.as_c_str(),
+                )
             };
             if !ok {
                 dcx.emit_err(crate::diagnostics::OffloadEmbedFailed);
@@ -878,7 +881,7 @@ pub(crate) unsafe fn llvm_optimize(
                 dcx,
                 module.module_llvm.tm.raw(),
                 config.no_builtins,
-                llmod2,
+                module.module_llvm.llmod(),
                 &out_obj,
                 None,
                 llvm::FileType::ObjectFile,
@@ -888,6 +891,16 @@ pub(crate) unsafe fn llvm_optimize(
             // 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.
+
+            let ok = unsafe {
+                llvm::RustOffloadWrapper::get_instance().llvm_rust_offload_wrap_images(
+                    module.module_llvm.llmod(),
+                    device_bin_c.as_c_str(),
+                )
+            };
+            if !ok {
+                dcx.emit_err(crate::diagnostics::OffloadWrapImagesFailed);
+            }
         }
     }
     result.into_result().unwrap_or_else(|()| llvm_err(dcx, LlvmError::RunLlvmPasses))

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -57,80 +57,6 @@ impl<'ll> OffloadGlobals<'ll> {
     }
 }
 
-// We need to register offload before using it. We also should unregister it once we are done, for
-// good measures. Previously we have done so before and after each individual offload intrinsic
-// call, but that comes at a performance cost. The repeated (un)register calls might also confuse
-// the LLVM ompOpt pass, which tries to move operations to a better location. The easiest solution,
-// which we copy from clang, is to just have those two calls once, in the global ctor/dtor section
-// of the final binary.
-pub(crate) fn register_offload<'ll>(cx: &CodegenCx<'ll, '_>) {
-    // First we check quickly whether we already have done our setup, in which case we return early.
-    // Shouldn't be needed for correctness.
-    let register_lib_name = "__tgt_register_lib";
-    if cx.get_function(register_lib_name).is_some() {
-        return;
-    }
-
-    let reg_lib_decl = cx.type_func(&[cx.type_ptr()], cx.type_void());
-    let register_lib = declare_offload_fn(&cx, register_lib_name, reg_lib_decl);
-    let unregister_lib = declare_offload_fn(&cx, "__tgt_unregister_lib", reg_lib_decl);
-
-    let ptr_null = cx.const_null(cx.type_ptr());
-    let const_struct = cx.const_struct(&[cx.get_const_i32(0), ptr_null, ptr_null, ptr_null], false);
-    let omp_descriptor =
-        add_global(cx, ".omp_offloading.descriptor", const_struct, InternalLinkage);
-    // @.omp_offloading.descriptor = internal constant %__tgt_bin_desc { i32 1, ptr @.omp_offloading.device_images, ptr @__start_llvm_offload_entries, ptr @__stop_llvm_offload_entries }
-    // @.omp_offloading.descriptor = internal constant %__tgt_bin_desc { i32 0, ptr null, ptr null, ptr null }
-
-    let atexit = cx.type_func(&[cx.type_ptr()], cx.type_i32());
-    let atexit_fn = declare_offload_fn(cx, "atexit", atexit);
-
-    // FIXME(offload): Drop this, once we fully automated our offload compilation pipeline, since
-    // LLVM will initialize them for us if it sees gpu kernels being registered.
-    let init_ty = cx.type_func(&[], cx.type_void());
-    let init_rtls = declare_offload_fn(cx, "__tgt_init_all_rtls", init_ty);
-
-    let desc_ty = cx.type_func(&[], cx.type_void());
-    let reg_name = ".omp_offloading.descriptor_reg";
-    let unreg_name = ".omp_offloading.descriptor_unreg";
-    let desc_reg_fn = declare_offload_fn(cx, reg_name, desc_ty);
-    let desc_unreg_fn = declare_offload_fn(cx, unreg_name, desc_ty);
-    llvm::set_linkage(desc_reg_fn, InternalLinkage);
-    llvm::set_linkage(desc_unreg_fn, InternalLinkage);
-    llvm::set_section(desc_reg_fn, c".text.startup");
-    llvm::set_section(desc_unreg_fn, c".text.startup");
-
-    // define internal void @.omp_offloading.descriptor_reg() section ".text.startup" {
-    // entry:
-    //   call void @__tgt_register_lib(ptr @.omp_offloading.descriptor)
-    //   call void @__tgt_init_all_rtls()
-    //   %0 = call i32 @atexit(ptr @.omp_offloading.descriptor_unreg)
-    //   ret void
-    // }
-    let bb = Builder::append_block(cx, desc_reg_fn, "entry");
-    let mut a = Builder::build(cx, bb);
-    a.call(reg_lib_decl, None, None, register_lib, &[omp_descriptor], None, None);
-    a.call(init_ty, None, None, init_rtls, &[], None, None);
-    a.call(atexit, None, None, atexit_fn, &[desc_unreg_fn], None, None);
-    a.ret_void();
-
-    // define internal void @.omp_offloading.descriptor_unreg() section ".text.startup" {
-    // entry:
-    //   call void @__tgt_unregister_lib(ptr @.omp_offloading.descriptor)
-    //   ret void
-    // }
-    let bb = Builder::append_block(cx, desc_unreg_fn, "entry");
-    let mut a = Builder::build(cx, bb);
-    a.call(reg_lib_decl, None, None, unregister_lib, &[omp_descriptor], None, None);
-    a.ret_void();
-
-    // @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 101, ptr @.omp_offloading.descriptor_reg, ptr null }]
-    let args = vec![cx.get_const_i32(101), desc_reg_fn, ptr_null];
-    let const_struct = cx.const_struct(&args, false);
-    let arr = cx.const_array(cx.val_ty(const_struct), &[const_struct]);
-    add_global(cx, "llvm.global_ctors", arr, AppendingLinkage);
-}
-
 pub(crate) struct OffloadKernelDims<'ll> {
     num_workgroups: &'ll Value,
     threads_per_block: &'ll Value,

compiler/rustc_codegen_llvm/src/context.rs

@@ -918,11 +918,6 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         self.get_const_int(self.type_i8(), n)
     }
 
-    pub(crate) fn get_function(&self, name: &str) -> Option<&'ll Value> {
-        let name = SmallCStr::new(name);
-        unsafe { llvm::LLVMGetNamedFunction((**self).borrow().llmod, name.as_ptr()) }
-    }
-
     pub(crate) fn get_md_kind_id(&self, name: &str) -> llvm::MetadataKindId {
         unsafe {
             llvm::LLVMGetMDKindIDInContext(

compiler/rustc_codegen_llvm/src/diagnostics.rs

@@ -107,6 +107,10 @@ pub(crate) struct OffloadBundleImagesFailed;
 #[diag("call to EmbedBufferInModule failed, `host.o` was not created")]
 pub(crate) struct OffloadEmbedFailed;
 
+#[derive(Diagnostic)]
+#[diag("call to WrapImages failed, device image was not wrapped into the host module")]
+pub(crate) struct OffloadWrapImagesFailed;
+
 #[derive(Diagnostic)]
 #[diag("failed to get bitcode from object file for LTO ({$err})")]
 pub(crate) struct LtoBitcodeFromRlib {

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -36,9 +36,7 @@ use tracing::debug;
 use crate::abi::FnAbiLlvmExt;
 use crate::builder::Builder;
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
-use crate::builder::gpu_offload::{
-    self, OffloadKernelDims, declare_omp_get_num_devices, register_offload,
-};
+use crate::builder::gpu_offload::{self, OffloadKernelDims, declare_omp_get_num_devices};
 use crate::context::CodegenCx;
 use crate::declare::declare_raw_fn;
 use crate::diagnostics::{
@@ -1900,7 +1898,6 @@ fn codegen_offload<'ll, 'tcx>(
             return;
         }
     };
-    register_offload(cx);
     let offload_data =
         gpu_offload::gen_define_handling(&cx, &metadata, target_symbol, offload_globals);
     gpu_offload::gen_call_handling(

compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs

@@ -1,12 +1,16 @@
 use std::ffi::{CStr, c_char};
+use std::path::PathBuf;
 use std::sync::OnceLock;
 
 use super::ffi::{Module, TargetMachine, Value};
 
 type LLVMRustBundleImagesFn = unsafe extern "C" fn(&Module, &TargetMachine, *const c_char) -> bool;
 type LLVMRustOffloadEmbedBufferInModuleFn = unsafe extern "C" fn(&Module, *const c_char) -> bool;
 type LLVMRustOffloadMapperFn = unsafe extern "C" fn(&Value, &Value, *const &Value);
+type LLVMRustOffloadWrapImagesFn =
+    unsafe extern "C" fn(&Module, *const c_char, *const c_char) -> bool;
 
+use rustc_fs_util::path_to_c_string;
 use rustc_session::config::host_tuple;
 use rustc_session::filesearch;
 
@@ -16,6 +20,8 @@ pub(crate) struct RustOffloadWrapper {
     LLVMRustBundleImages: LLVMRustBundleImagesFn,
     LLVMRustOffloadEmbedBufferInModule: LLVMRustOffloadEmbedBufferInModuleFn,
     LLVMRustOffloadMapper: LLVMRustOffloadMapperFn,
+    LLVMRustOffloadWrapImages: LLVMRustOffloadWrapImagesFn,
+    lld_path: Option<PathBuf>,
     // Keep the dynamic library loaded while the function pointers are used.
     _lib: libloading::Library,
 }
@@ -71,10 +77,21 @@ impl RustOffloadWrapper {
         unsafe { (self.LLVMRustOffloadMapper)(v1, v2, vs.as_ptr()) }
     }
 
+    pub(crate) unsafe fn llvm_rust_offload_wrap_images(
+        &self,
+        host_m: &Module,
+        device_bin_path: &CStr,
+    ) -> bool {
+        let lld_c = self.lld_path.as_deref().map(path_to_c_string).unwrap_or_default();
+        unsafe {
+            (self.LLVMRustOffloadWrapImages)(host_m, lld_c.as_ptr(), device_bin_path.as_ptr())
+        }
+    }
+
     fn call_dynamic(
         sysroot: &rustc_session::config::Sysroot,
     ) -> Result<Self, RustOffloadLibraryError> {
-        let rust_offload_path = Self::get_rust_offload_path(sysroot)?;
+        let (rust_offload_path, lld_path) = Self::get_offload_and_lld_paths(sysroot)?;
         let lib = unsafe { libloading::Library::new(rust_offload_path)? };
 
         let llvm_rust_bundle_images =
@@ -86,48 +103,47 @@ impl RustOffloadWrapper {
         };
         let llvm_rust_offload_wrapper =
             *unsafe { lib.get::<LLVMRustOffloadMapperFn>(b"LLVMRustOffloadMapper\0")? };
+        let llvm_rust_offload_wrap_images =
+            *unsafe { lib.get::<LLVMRustOffloadWrapImagesFn>(b"LLVMRustOffloadWrapImages\0")? };
 
         Ok(Self {
             LLVMRustBundleImages: llvm_rust_bundle_images,
             LLVMRustOffloadEmbedBufferInModule: llvm_rust_offload_embed_buffer_in_module,
             LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
+            LLVMRustOffloadWrapImages: llvm_rust_offload_wrap_images,
+            lld_path,
             _lib: lib,
         })
     }
 
-    fn get_rust_offload_path(
+    fn get_offload_and_lld_paths(
         sysroot: &rustc_session::config::Sysroot,
-    ) -> Result<String, RustOffloadLibraryError> {
+    ) -> Result<(PathBuf, Option<PathBuf>), RustOffloadLibraryError> {
         let llvm_version_major = unsafe { LLVMRustVersionMajor() };
-
-        let path_buf = sysroot
-            .all_paths()
-            .find_map(|p| {
-                let candidate = filesearch::make_target_lib_path(p, host_tuple())
-                    .join(format!("libRustOffload-{}", llvm_version_major))
-                    .with_extension(std::env::consts::DLL_EXTENSION);
-
-                candidate.exists().then_some(candidate)
-            })
-            .ok_or_else(|| {
-                let candidates = sysroot
-                    .all_paths()
-                    .map(|p| p.join("lib").display().to_string())
-                    .collect::<Vec<String>>()
-                    .join("\n* ");
-                RustOffloadLibraryError::NotFound {
-                    err: format!(
-                        "failed to find a `libRustOffload-{llvm_version_major}` \
-                    in the sysroot candidates:\n* {candidates}"
-                    ),
-                }
-            })?;
-
-        Ok(path_buf
-            .to_str()
-            .ok_or_else(|| RustOffloadLibraryError::LoadFailed {
-                err: format!("invalid UTF-8 in path: {}", path_buf.display()),
-            })?
-            .to_string())
+        let mut searched = Vec::new();
+
+        for root in sysroot.all_paths() {
+            let rust_offload_path = filesearch::make_target_lib_path(root, host_tuple())
+                .join(format!("libRustOffload-{llvm_version_major}"))
+                .with_extension(std::env::consts::DLL_EXTENSION);
+
+            if !rust_offload_path.is_file() {
+                searched.push(rust_offload_path);
+                continue;
+            }
+
+            let lld_path = filesearch::make_target_bin_path(root, host_tuple())
+                .join(format!("rust-lld{}", std::env::consts::EXE_SUFFIX));
+            let lld_path = lld_path.is_file().then_some(lld_path);
+
+            return Ok((rust_offload_path, lld_path));
+        }
+
+        Err(RustOffloadLibraryError::NotFound {
+            err: format!(
+                "could not find libRustOffload-{llvm_version_major} in the sysroot candidates:\n* {}",
+                searched.iter().map(|p| p.display().to_string()).collect::<Vec<_>>().join("\n* ")
+            ),
+        })
     }
 }

compiler/rustc_codegen_ssa/src/back/link.rs

@@ -3029,6 +3029,14 @@ fn linker_with_args(
         link_output_kind,
     );
 
+    if sess.opts.unstable_opts.offload.iter().any(|o| matches!(o, config::Offload::Host(_))) {
+        cmd.link_dylib_by_name("omptarget", false, true);
+        cmd.link_dylib_by_name("omp", false, true);
+        cmd.link_args(["-z", "nostart-stop-gc"]);
+        cmd.link_arg("-rpath");
+        cmd.link_arg(std::path::absolute(&*sess.target_tlib_path.dir).unwrap());
+    }
+
     // Upstream rust crates and their non-dynamic native libraries.
     add_upstream_rust_crates(
         cmd,