From 162cba5202c77d583a9cfb1de3b6b3c0d5cf9fbb Mon Sep 17 00:00:00 2001 From: sgasho Date: Thu, 6 Aug 2026 23:16:59 +0000 Subject: [PATCH] dlopen Offload --- compiler/rustc_codegen_llvm/src/back/write.rs | 28 ++-- .../src/builder/gpu_offload.rs | 2 +- .../rustc_codegen_llvm/src/diagnostics.rs | 13 ++ compiler/rustc_codegen_llvm/src/lib.rs | 20 +++ compiler/rustc_codegen_llvm/src/llvm/ffi.rs | 57 -------- compiler/rustc_codegen_llvm/src/llvm/mod.rs | 2 + .../src/llvm/offload_ffi.rs | 133 ++++++++++++++++++ .../rustc_llvm/llvm-wrapper/RustWrapper.cpp | 105 -------------- .../llvm-wrapper/offload/CMakeLists.txt | 27 ++++ .../llvm-wrapper/offload/OffloadWrapper.cpp | 117 +++++++++++++++ src/bootstrap/src/core/build_steps/compile.rs | 8 ++ src/bootstrap/src/core/build_steps/llvm.rs | 92 +++++++++++- src/bootstrap/src/core/builder/mod.rs | 1 + src/bootstrap/src/lib.rs | 6 +- 14 files changed, 435 insertions(+), 176 deletions(-) create mode 100644 compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs create mode 100644 compiler/rustc_llvm/llvm-wrapper/offload/CMakeLists.txt create mode 100644 compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs index edf52e67b434b..6aaefbe82ec2b 100644 --- a/compiler/rustc_codegen_llvm/src/back/write.rs +++ b/compiler/rustc_codegen_llvm/src/back/write.rs @@ -720,7 +720,11 @@ pub(crate) unsafe fn llvm_optimize( // Here we map the old arguments to the new arguments, with an offset of 1 to make sure // that we don't use the newly added `%dyn_ptr`. unsafe { - llvm::LLVMRustOffloadMapper(old_fn, new_fn, old_args_rebuilt.as_ptr()); + llvm::RustOffloadWrapper::get_instance().llvm_rust_offload_wrapper( + old_fn, + new_fn, + old_args_rebuilt.as_slice(), + ); } llvm::set_linkage(new_fn, llvm::get_linkage(old_fn)); @@ -814,16 +818,16 @@ pub(crate) unsafe fn llvm_optimize( let device_dir = device_path.parent().unwrap(); let device_out = device_dir.join("device.bin"); let device_out_c = path_to_c_string(device_out.as_path()); - unsafe { - // 1) Bundle device module into offload image device.bin (device TM) - let ok = llvm::LLVMRustBundleImages( + // 1) Bundle device module into offload image device.bin (device TM) + let ok = unsafe { + llvm::RustOffloadWrapper::get_instance().llvm_rust_bundle_images( module.module_llvm.llmod(), module.module_llvm.tm.raw(), - device_out_c.as_ptr(), - ); - if !ok || !device_out.exists() { - dcx.emit_err(crate::diagnostics::OffloadBundleImagesFailed); - } + device_out_c.as_c_str(), + ) + }; + if !ok || !device_out.exists() { + dcx.emit_err(crate::diagnostics::OffloadBundleImagesFailed); } } @@ -859,8 +863,10 @@ pub(crate) unsafe fn llvm_optimize( // 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::LLVMRustOffloadEmbedBufferInModule(llmod2, device_bin_c.as_ptr()) }; + 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); } diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 0b009321802cf..3d0bb6fcc48fd 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -296,7 +296,7 @@ struct KernelArgsTy { impl KernelArgsTy { const OFFLOAD_VERSION: u64 = 3; - const FLAGS: u64 = 0; + const FLAGS: u64 = 1 << 6; // Enable StrictBlocksAndThreads const TRIPCOUNT: u64 = 0; fn new_decl<'ll>(cx: &CodegenCx<'ll, '_>) -> &'ll Type { let kernel_arguments_ty = cx.type_named_struct("struct.__tgt_kernel_arguments"); diff --git a/compiler/rustc_codegen_llvm/src/diagnostics.rs b/compiler/rustc_codegen_llvm/src/diagnostics.rs index ea29683b9d289..54f8ffbb881da 100644 --- a/compiler/rustc_codegen_llvm/src/diagnostics.rs +++ b/compiler/rustc_codegen_llvm/src/diagnostics.rs @@ -60,6 +60,19 @@ pub(crate) struct AutoDiffWithoutLto; #[diag("using the autodiff feature requires -Z autodiff=Enable")] pub(crate) struct AutoDiffWithoutEnable; +#[derive(Diagnostic)] +#[diag("failed to load our rust offload backend: {$err}")] +pub(crate) struct RustOffloadComponentUnavailable { + pub err: String, +} + +#[derive(Diagnostic)] +#[diag("rust offload backend not found in the sysroot: {$err}")] +#[note("it will be distributed via rustup in the future")] +pub(crate) struct RustOffloadComponentMissing { + pub err: String, +} + #[derive(Diagnostic)] #[diag( "using the offload feature requires -Z offload=" diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs index 3ec0495956c4c..fe39fc6b3fca0 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/compiler/rustc_codegen_llvm/src/lib.rs @@ -373,6 +373,26 @@ impl CodegenBackend for LlvmCodegenBackend { } fn codegen_crate<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Box { + use rustc_session::config::Offload; + + if tcx.sess.opts.unstable_opts.offload.contains(&Offload::Device) + || tcx.sess.opts.unstable_opts.offload.iter().any(|o| matches!(o, Offload::Host(_))) + { + match llvm::RustOffloadWrapper::get_or_init(&tcx.sess.opts.sysroot) { + Ok(_) => {} + Err(llvm::RustOffloadLibraryError::NotFound { err }) => { + tcx.sess + .dcx() + .emit_fatal(crate::diagnostics::RustOffloadComponentMissing { err }); + } + Err(llvm::RustOffloadLibraryError::LoadFailed { err }) => { + tcx.sess + .dcx() + .emit_fatal(crate::diagnostics::RustOffloadComponentUnavailable { err }); + } + } + } + Box::new(rustc_codegen_ssa::base::codegen_crate(LlvmCodegenBackend(()), tcx)) } diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs index 4cc5d326bdc9e..1a60b59a93525 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs @@ -1713,63 +1713,6 @@ unsafe extern "C" { ) -> &'a Value; } -#[cfg(feature = "llvm_offload")] -pub(crate) use self::Offload::*; - -#[cfg(feature = "llvm_offload")] -mod Offload { - use super::*; - unsafe extern "C" { - /// Processes the module and writes it in an offload compatible way into a "device.bin" file. - pub(crate) fn LLVMRustBundleImages<'a>( - M: &'a Module, - TM: &'a TargetMachine, - device_bin: *const c_char, - ) -> bool; - pub(crate) unsafe fn LLVMRustOffloadEmbedBufferInModule<'a>( - _M: &'a Module, - _device_bin: *const c_char, - ) -> bool; - pub(crate) fn LLVMRustOffloadMapper<'a>( - OldFn: &'a Value, - NewFn: &'a Value, - RebuiltArgs: *const &Value, - ); - } -} - -#[cfg(not(feature = "llvm_offload"))] -pub(crate) use self::Offload_fallback::*; - -#[cfg(not(feature = "llvm_offload"))] -mod Offload_fallback { - use super::*; - /// Processes the module and writes it in an offload compatible way into a "device.bin" file. - /// Marked as unsafe to match the real offload wrapper which is unsafe due to FFI. - #[allow(unused_unsafe)] - pub(crate) unsafe fn LLVMRustBundleImages<'a>( - _M: &'a Module, - _TM: &'a TargetMachine, - _device_bin: *const c_char, - ) -> bool { - unimplemented!("This rustc version was not built with LLVM Offload support!"); - } - pub(crate) unsafe fn LLVMRustOffloadEmbedBufferInModule<'a>( - _M: &'a Module, - _device_bin: *const c_char, - ) -> bool { - unimplemented!("This rustc version was not built with LLVM Offload support!"); - } - #[allow(unused_unsafe)] - pub(crate) unsafe fn LLVMRustOffloadMapper<'a>( - _OldFn: &'a Value, - _NewFn: &'a Value, - _RebuiltArgs: *const &Value, - ) { - unimplemented!("This rustc version was not built with LLVM Offload support!"); - } -} - // FFI bindings for `DIBuilder` functions in the LLVM-C API. // Try to keep these in the same order as in `llvm/include/llvm-c/DebugInfo.h`. // diff --git a/compiler/rustc_codegen_llvm/src/llvm/mod.rs b/compiler/rustc_codegen_llvm/src/llvm/mod.rs index a2d17e93b4996..eb7a529c0b198 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/mod.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/mod.rs @@ -21,8 +21,10 @@ pub(crate) mod diagnostic; pub(crate) mod enzyme_ffi; mod ffi; mod metadata_kind; +pub(crate) mod offload_ffi; pub(crate) use self::enzyme_ffi::*; +pub(crate) use self::offload_ffi::*; impl LLVMRustResult { pub(crate) fn into_result(self) -> Result<(), ()> { diff --git a/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs new file mode 100644 index 0000000000000..46d9320248a9b --- /dev/null +++ b/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs @@ -0,0 +1,133 @@ +use std::ffi::{CStr, c_char}; +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); + +use rustc_session::config::host_tuple; +use rustc_session::filesearch; + +use crate::llvm::LLVMRustVersionMajor; + +pub(crate) struct RustOffloadWrapper { + LLVMRustBundleImages: LLVMRustBundleImagesFn, + LLVMRustOffloadEmbedBufferInModule: LLVMRustOffloadEmbedBufferInModuleFn, + LLVMRustOffloadMapper: LLVMRustOffloadMapperFn, + // Keep the dynamic library loaded while the function pointers are used. + _lib: libloading::Library, +} + +#[derive(Debug)] +pub(crate) enum RustOffloadLibraryError { + NotFound { err: String }, + LoadFailed { err: String }, +} + +impl From for RustOffloadLibraryError { + fn from(err: libloading::Error) -> Self { + Self::LoadFailed { err: format!("{err:?}") } + } +} + +static OFFLOAD_INSTANCE: OnceLock = OnceLock::new(); + +impl RustOffloadWrapper { + pub(crate) fn get_or_init( + sysroot: &rustc_session::config::Sysroot, + ) -> Result<&'static RustOffloadWrapper, RustOffloadLibraryError> { + OFFLOAD_INSTANCE.get_or_try_init(|| { + let w = Self::call_dynamic(sysroot)?; + Ok(w) + }) + } + + pub(crate) fn get_instance() -> &'static RustOffloadWrapper { + OFFLOAD_INSTANCE + .get() + .expect("RustOffloadWrapper not initialized. Call get_or_init with sysroot first.") + } + + pub(crate) unsafe fn llvm_rust_bundle_images( + &self, + m: &Module, + tm: &TargetMachine, + c: &CStr, + ) -> bool { + unsafe { (self.LLVMRustBundleImages)(m, tm, c.as_ptr()) } + } + + pub(crate) unsafe fn llvm_rust_offload_embed_buffer_in_module( + &self, + m: &Module, + i: &CStr, + ) -> bool { + unsafe { (self.LLVMRustOffloadEmbedBufferInModule)(m, i.as_ptr()) } + } + + pub(crate) unsafe fn llvm_rust_offload_wrapper(&self, v1: &Value, v2: &Value, vs: &[&Value]) { + unsafe { (self.LLVMRustOffloadMapper)(v1, v2, vs.as_ptr()) } + } + + fn call_dynamic( + sysroot: &rustc_session::config::Sysroot, + ) -> Result { + let rust_offload_path = Self::get_rust_offload_path(sysroot)?; + let lib = unsafe { libloading::Library::new(rust_offload_path)? }; + + let llvm_rust_bundle_images = + *unsafe { lib.get::(b"LLVMRustBundleImages\0")? }; + let llvm_rust_offload_embed_buffer_in_module = *unsafe { + lib.get::( + b"LLVMRustOffloadEmbedBufferInModule\0", + )? + }; + let llvm_rust_offload_wrapper = + *unsafe { lib.get::(b"LLVMRustOffloadMapper\0")? }; + + Ok(Self { + LLVMRustBundleImages: llvm_rust_bundle_images, + LLVMRustOffloadEmbedBufferInModule: llvm_rust_offload_embed_buffer_in_module, + LLVMRustOffloadMapper: llvm_rust_offload_wrapper, + _lib: lib, + }) + } + + fn get_rust_offload_path( + sysroot: &rustc_session::config::Sysroot, + ) -> Result { + 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::>() + .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()) + } +} diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp index f500041a12d8b..983a506bd4ac6 100644 --- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp +++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp @@ -164,111 +164,6 @@ extern "C" bool LLVMRustIsCall(LLVMValueRef V) { return llvm::isa(llvm::unwrap(V)); } -// Some of the functions here rely on LLVM modules that may not always be -// available. As such, we only try to build it in the first place, if -// llvm.offload is enabled. -#ifdef OFFLOAD -static Error writeFile(StringRef Filename, StringRef Data) { - Expected> OutputOrErr = - FileOutputBuffer::create(Filename, Data.size()); - if (!OutputOrErr) - return OutputOrErr.takeError(); - std::unique_ptr Output = std::move(*OutputOrErr); - llvm::copy(Data, Output->getBufferStart()); - if (Error E = Output->commit()) - return E; - return Error::success(); -} - -// This is the first of many steps in creating a binary using llvm offload, -// to run code on the gpu. Concrete, it replaces the following binary use: -// clang-offload-packager -o device.bin -// --image=file=device.bc,triple=amdgcn-amd-amdhsa,arch=gfx90a,kind=openmp -// The input module is the rust code compiled for a gpu target like amdgpu. -// Based on clang/tools/clang-offload-packager/ClangOffloadPackager.cpp -extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM, - const char *HostOutPath) { - std::string Storage; - llvm::raw_string_ostream OS1(Storage); - llvm::WriteBitcodeToFile(*unwrap(M), OS1); - OS1.flush(); - auto MB = llvm::MemoryBuffer::getMemBufferCopy(Storage, "device.bc"); - - SmallVector BinaryData; - raw_svector_ostream OS2(BinaryData); - - OffloadBinary::OffloadingImage ImageBinary{}; - ImageBinary.TheImageKind = object::IMG_Bitcode; - ImageBinary.Image = std::move(MB); - ImageBinary.TheOffloadKind = object::OFK_OpenMP; - - std::string TripleStr = TM.getTargetTriple().str(); - llvm::StringRef CPURef = TM.getTargetCPU(); - ImageBinary.StringData["triple"] = TripleStr; - ImageBinary.StringData["arch"] = CPURef; - llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary); - if (Buffer.size() % OffloadBinary::getAlignment() != 0) - // Offload binary has invalid size alignment - return false; - OS2 << Buffer; - if (Error E = writeFile(HostOutPath, - StringRef(BinaryData.begin(), BinaryData.size()))) - return false; - return true; -} - -extern "C" bool LLVMRustOffloadEmbedBufferInModule(LLVMModuleRef HostM, - const char *HostOutPath) { - auto MBOrErr = MemoryBuffer::getFile(HostOutPath); - if (!MBOrErr) { - auto E = MBOrErr.getError(); - auto _B = errorCodeToError(E); - return false; - } - MemoryBufferRef Buf = (*MBOrErr)->getMemBufferRef(); - Module *M = unwrap(HostM); - StringRef SectionName = ".llvm.offloading"; - Align Alignment = Align(8); - llvm::embedBufferInModule(*M, Buf, SectionName, Alignment); - return true; -} - -// Clone OldFn into NewFn, remapping its arguments to RebuiltArgs. -// Each arg of OldFn is replaced with the corresponding value in RebuiltArgs. -// For scalars, RebuiltArgs contains the value cast and/or truncated to the -// original type. -extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn, - const LLVMValueRef *RebuiltArgs) { - llvm::Function *oldFn = llvm::unwrap(OldFn); - llvm::Function *newFn = llvm::unwrap(NewFn); - - // Map old arguments to new arguments. We skip the first dyn_ptr argument, - // since it can't be used directly by user code. - llvm::ValueToValueMapTy vmap; - auto newArgIt = newFn->arg_begin(); - newArgIt->setName("dyn_ptr"); - - unsigned i = 0; - for (auto &oldArg : oldFn->args()) { - vmap[&oldArg] = unwrap(RebuiltArgs[i++]); - } - - llvm::SmallVector returns; - llvm::CloneFunctionInto(newFn, oldFn, vmap, - llvm::CloneFunctionChangeType::LocalChangesOnly, - returns); - - BasicBlock &entry = newFn->getEntryBlock(); - BasicBlock &clonedEntry = *std::next(newFn->begin()); - - if (entry.getTerminator()) - entry.getTerminator()->eraseFromParent(); - - IRBuilder<> B(&entry); - B.CreateBr(&clonedEntry); -} -#endif - extern "C" LLVMValueRef LLVMRustGetNamedValue(LLVMModuleRef M, const char *Name, size_t NameLen) { return wrap(unwrap(M)->getNamedValue(StringRef(Name, NameLen))); diff --git a/compiler/rustc_llvm/llvm-wrapper/offload/CMakeLists.txt b/compiler/rustc_llvm/llvm-wrapper/offload/CMakeLists.txt new file mode 100644 index 0000000000000..37c747a902d87 --- /dev/null +++ b/compiler/rustc_llvm/llvm-wrapper/offload/CMakeLists.txt @@ -0,0 +1,27 @@ +cmake_minimum_required(VERSION 3.20) +project(RustOffload LANGUAGES CXX) + +find_package(LLVM CONFIG REQUIRED) + +add_library(RustOffload-${LLVM_VERSION_MAJOR} SHARED + OffloadWrapper.cpp +) + +target_include_directories(RustOffload-${LLVM_VERSION_MAJOR} PRIVATE + ${LLVM_INCLUDE_DIRS} +) + +target_link_libraries(RustOffload-${LLVM_VERSION_MAJOR} PRIVATE + LLVM +) + +if(NOT LLVM_ENABLE_RTTI) + target_compile_options( + RustOffload-${LLVM_VERSION_MAJOR} + PRIVATE -fno-rtti + ) +endif() + +install(TARGETS RustOffload-${LLVM_VERSION_MAJOR} + LIBRARY DESTINATION lib +) diff --git a/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp new file mode 100644 index 0000000000000..8c18f2453e9d8 --- /dev/null +++ b/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp @@ -0,0 +1,117 @@ +#include "../SuppressLLVMWarnings.h" + +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Bitcode/BitcodeWriter.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/Object/OffloadBinary.h" +#include "llvm/Support/CBindingWrapping.h" +#include "llvm/Support/FileOutputBuffer.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Transforms/Utils/Cloning.h" +#include "llvm/Transforms/Utils/ModuleUtils.h" +#include "llvm/Transforms/Utils/ValueMapper.h" + +using namespace llvm; +using namespace llvm::object; + +static Error writeFile(StringRef Filename, StringRef Data) { + Expected> OutputOrErr = + FileOutputBuffer::create(Filename, Data.size()); + if (!OutputOrErr) + return OutputOrErr.takeError(); + std::unique_ptr Output = std::move(*OutputOrErr); + llvm::copy(Data, Output->getBufferStart()); + if (Error E = Output->commit()) + return E; + return Error::success(); +} + +// This is the first of many steps in creating a binary using llvm offload, +// to run code on the gpu. Concrete, it replaces the following binary use: +// clang-offload-packager -o device.bin +// --image=file=device.bc,triple=amdgcn-amd-amdhsa,arch=gfx90a,kind=openmp +// The input module is the rust code compiled for a gpu target like amdgpu. +// Based on clang/tools/clang-offload-packager/ClangOffloadPackager.cpp +extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM, + const char *HostOutPath) { + std::string Storage; + llvm::raw_string_ostream OS1(Storage); + llvm::WriteBitcodeToFile(*unwrap(M), OS1); + OS1.flush(); + auto MB = llvm::MemoryBuffer::getMemBufferCopy(Storage, "device.bc"); + + SmallVector BinaryData; + raw_svector_ostream OS2(BinaryData); + + OffloadBinary::OffloadingImage ImageBinary{}; + ImageBinary.TheImageKind = object::IMG_Bitcode; + ImageBinary.Image = std::move(MB); + ImageBinary.TheOffloadKind = object::OFK_OpenMP; + + std::string TripleStr = TM.getTargetTriple().str(); + llvm::StringRef CPURef = TM.getTargetCPU(); + ImageBinary.StringData["triple"] = TripleStr; + ImageBinary.StringData["arch"] = CPURef; + llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary); + if (Buffer.size() % OffloadBinary::getAlignment() != 0) + // Offload binary has invalid size alignment + return false; + OS2 << Buffer; + if (Error E = writeFile(HostOutPath, + StringRef(BinaryData.begin(), BinaryData.size()))) + return false; + return true; +} + +extern "C" bool LLVMRustOffloadEmbedBufferInModule(LLVMModuleRef HostM, + const char *HostOutPath) { + auto MBOrErr = MemoryBuffer::getFile(HostOutPath); + if (!MBOrErr) { + auto E = MBOrErr.getError(); + auto _B = errorCodeToError(E); + return false; + } + MemoryBufferRef Buf = (*MBOrErr)->getMemBufferRef(); + Module *M = unwrap(HostM); + StringRef SectionName = ".llvm.offloading"; + Align Alignment = Align(8); + llvm::embedBufferInModule(*M, Buf, SectionName, Alignment); + return true; +} + +// Clone OldFn into NewFn, remapping its arguments to RebuiltArgs. +// Each arg of OldFn is replaced with the corresponding value in RebuiltArgs. +// For scalars, RebuiltArgs contains the value cast and/or truncated to the +// original type. +extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn, + const LLVMValueRef *RebuiltArgs) { + llvm::Function *oldFn = llvm::unwrap(OldFn); + llvm::Function *newFn = llvm::unwrap(NewFn); + + // Map old arguments to new arguments. We skip the first dyn_ptr argument, + // since it can't be used directly by user code. + llvm::ValueToValueMapTy vmap; + auto newArgIt = newFn->arg_begin(); + newArgIt->setName("dyn_ptr"); + + unsigned i = 0; + for (auto &oldArg : oldFn->args()) { + vmap[&oldArg] = unwrap(RebuiltArgs[i++]); + } + + llvm::SmallVector returns; + llvm::CloneFunctionInto(newFn, oldFn, vmap, + llvm::CloneFunctionChangeType::LocalChangesOnly, + returns); + + BasicBlock &entry = newFn->getEntryBlock(); + BasicBlock &clonedEntry = *std::next(newFn->begin()); + + if (entry.getTerminator()) + entry.getTerminator()->eraseFromParent(); + + IRBuilder<> B(&entry); + B.CreateBr(&clonedEntry); +} diff --git a/src/bootstrap/src/core/build_steps/compile.rs b/src/bootstrap/src/core/build_steps/compile.rs index 021a652a5ac50..25c3df3ae541e 100644 --- a/src/bootstrap/src/core/build_steps/compile.rs +++ b/src/bootstrap/src/core/build_steps/compile.rs @@ -2275,10 +2275,18 @@ impl CommandLineStep for Assemble { if builder.config.llvm_offload && !builder.config.dry_run() { debug!("`llvm_offload` requested"); + let rust_offload = builder.ensure(llvm::RustOffload { target: build_compiler.host }); let offload_install = builder.ensure(llvm::OmpOffload { target: build_compiler.host }); if let Some(_llvm_config) = builder.llvm_config(builder.config.host_target) { let target_libdir = builder.sysroot_target_libdir(target_compiler, target_compiler.host); + let rust_offload_dst_lib = target_libdir.join(rust_offload.rust_offload_filename()); + builder.copy_link( + &rust_offload.rust_offload_path(), + &rust_offload_dst_lib, + FileType::NativeLibrary, + ); + for p in offload_install.offload_paths() { let libname = p.file_name().unwrap(); let dst_lib = target_libdir.join(libname); diff --git a/src/bootstrap/src/core/build_steps/llvm.rs b/src/bootstrap/src/core/build_steps/llvm.rs index d3276cfb5371b..2e14082310350 100644 --- a/src/bootstrap/src/core/build_steps/llvm.rs +++ b/src/bootstrap/src/core/build_steps/llvm.rs @@ -942,6 +942,96 @@ fn get_var(var_base: &str, host: &str, target: &str) -> Option { .or_else(|| env::var_os(var_base)) } +#[derive(Clone)] +pub struct BuiltRustOffload { + /// Path to the rust offload dylib + offload: PathBuf, +} + +impl BuiltRustOffload { + pub fn rust_offload_path(&self) -> PathBuf { + self.offload.clone() + } + + pub fn rust_offload_filename(&self) -> String { + self.offload.file_name().unwrap().to_str().unwrap().to_owned() + } +} + +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub struct RustOffload { + pub target: TargetSelection, +} + +impl CommandLineStep for RustOffload { + type Output = BuiltRustOffload; + const IS_HOST: bool = true; + + fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> { + run.alias("rust-offload") + } + + fn make_run(run: RunConfig<'_>) { + run.builder.ensure(RustOffload { target: run.target }); + } + + fn run(self, builder: &Builder<'_>) -> Self::Output { + if builder.config.dry_run() { + return BuiltRustOffload { + offload: builder.config.tempdir().join("rust-offload-dry-run"), + }; + } + + let target = self.target; + + let LlvmResult { host_llvm_config, llvm_cmake_dir } = builder.ensure(Llvm { target }); + + let out_dir = builder.rust_offload_out(target); + + let llvm_version_major = llvm::get_llvm_version_major(builder, &host_llvm_config); + let lib_ext = std::env::consts::DLL_EXTENSION; + let lib_rust_offload = format!("libRustOffload-{llvm_version_major}"); + let build_dir = out_dir.join(libdir(target)); + let dylib = build_dir.join(&lib_rust_offload).with_extension(lib_ext); + + let mut cfg = + cmake::Config::new(builder.src.join("compiler/rustc_llvm/llvm-wrapper/offload/")); + + // Logic copied from `configure_llvm` + // ThinLTO is only available when building with LLVM, enabling LLD is required. + // Apple's linker ld64 supports ThinLTO out of the box though, so don't use LLD on Darwin. + let mut ldflags = LdFlags::default(); + if builder.config.llvm_thin_lto && !target.contains("apple") { + ldflags.push_all("-fuse-ld=lld"); + } + + configure_cmake(builder, target, &mut cfg, true, ldflags, CcFlags::default(), &[]); + + let profile = match (builder.config.llvm_optimize, builder.config.llvm_release_debuginfo) { + (false, _) => "Debug", + (true, false) => "Release", + (true, true) => "RelWithDebInfo", + }; + + cfg.out_dir(&out_dir) + .profile(profile) + .env("LLVM_CONFIG_REAL", &host_llvm_config) + .define("LLVM_DIR", llvm_cmake_dir); + + cfg.build(); + + if !dylib.exists() { + eprintln!( + "`{lib_rust_offload}` not found in `{}`. Either the build has failed or RustOffload was built with a wrong version of LLVM", + build_dir.display() + ); + exit!(1); + } + + BuiltRustOffload { offload: dylib } + } +} + #[derive(Clone)] pub struct BuiltOmpOffload { /// Path to the omp and offload dylibs. @@ -998,7 +1088,7 @@ impl CommandLineStep for OmpOffload { // Running cmake twice in the same folder is known to cause issues, like deleting existing // binaries. We therefore write our offload artifacts into it's own folder, instead of // using the llvm build dir. - let out_dir = builder.offload_out(target); + let out_dir = builder.omp_offload_out(target); let mut files = vec![]; let lib_ext = std::env::consts::DLL_EXTENSION; diff --git a/src/bootstrap/src/core/builder/mod.rs b/src/bootstrap/src/core/builder/mod.rs index 051e01a0a6666..22adbfd946965 100644 --- a/src/bootstrap/src/core/builder/mod.rs +++ b/src/bootstrap/src/core/builder/mod.rs @@ -823,6 +823,7 @@ impl<'a> Builder<'a> { tool::CargoMiri, llvm::Lld, llvm::Enzyme, + llvm::RustOffload, llvm::CrtBeginEnd, tool::RustdocGUITest, tool::OptimizedDist, diff --git a/src/bootstrap/src/lib.rs b/src/bootstrap/src/lib.rs index 7d119247b3bac..27a03b1616192 100644 --- a/src/bootstrap/src/lib.rs +++ b/src/bootstrap/src/lib.rs @@ -983,10 +983,14 @@ impl Build { self.out.join(&*target.triple).join("enzyme") } - fn offload_out(&self, target: TargetSelection) -> PathBuf { + fn omp_offload_out(&self, target: TargetSelection) -> PathBuf { self.out.join(&*target.triple).join("offload") } + fn rust_offload_out(&self, target: TargetSelection) -> PathBuf { + self.out.join(&*target.triple).join("rust-offload") + } + fn lld_out(&self, target: TargetSelection) -> PathBuf { self.out.join(target).join("lld") }