Skip to content

offload: automate manual clang-linker-wrapper step #162309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rust-bors merged 4 commits into from
Sep 8, 2026
Merged

offload: automate manual clang-linker-wrapper step #162309

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c321d90
offload: automate manual clang-linker-wrapper step
sgasho Sep 4, 2026
18ada5c
compile bitcode in device.bin beforehand in order to avoid LLVM JIT, …
sgasho Sep 6, 2026
f9017b1
add rpath
sgasho Sep 6, 2026
ff680e2
compile device image without adding extra binaries
sgasho Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 18 additions & 5 deletions compiler/rustc_codegen_llvm/src/back/write.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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);
Expand All @@ -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,
Expand All @@ -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))
Expand Down
74 changes: 0 additions & 74 deletions compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
5 changes: 0 additions & 5 deletions compiler/rustc_codegen_llvm/src/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -883,11 +883,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(
Expand Down
4 changes: 4 additions & 0 deletions compiler/rustc_codegen_llvm/src/diagnostics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
5 changes: 1 addition & 4 deletions compiler/rustc_codegen_llvm/src/intrinsic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -1880,7 +1878,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(
Expand Down
82 changes: 49 additions & 33 deletions compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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,
}
Expand Down Expand Up @@ -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 =
Expand All @@ -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* ")
),
})
}
}
8 changes: 8 additions & 0 deletions compiler/rustc_codegen_ssa/src/back/link.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

@ZuseZ4 ZuseZ4 Sep 4, 2026 •
edited by rustbot
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks to your last PR, we should ship these in our sysroot. Is it guaranteed to be on our default search path? If not, we might want to add an rpath or -L, or?

View changes since the review

@sgasho sgasho Sep 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, looking at walk_native_lib_search_dirs, it is on our default linker search path, so we don't need an additional -L, but we need an rpath for runtime. I also checked this on a UofT server.

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,
Expand Down
Loading
Loading