diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs index b8952ffc6bf81..90b2cab5b63e2 100644 --- a/compiler/rustc_codegen_llvm/src/back/write.rs +++ b/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)) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index d20a73e8e6825..e2ec20226e3ce 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/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, diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs index 853c4bfc9ca3f..ecad6e5629172 100644 --- a/compiler/rustc_codegen_llvm/src/context.rs +++ b/compiler/rustc_codegen_llvm/src/context.rs @@ -883,11 +883,6 @@ impl<'ll, CX: Borrow>> 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( diff --git a/compiler/rustc_codegen_llvm/src/diagnostics.rs b/compiler/rustc_codegen_llvm/src/diagnostics.rs index fb43b36fe39b9..70a14288aec0c 100644 --- a/compiler/rustc_codegen_llvm/src/diagnostics.rs +++ b/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 { diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 58957b46964c7..1f36a9543df72 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/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::{ @@ -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( diff --git a/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs index 46d9320248a9b..7ecf450ab1dba 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs @@ -1,4 +1,5 @@ use std::ffi::{CStr, c_char}; +use std::path::PathBuf; use std::sync::OnceLock; use super::ffi::{Module, TargetMachine, Value}; @@ -6,7 +7,10 @@ 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, // 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 { - 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::(b"LLVMRustOffloadMapper\0")? }; + let llvm_rust_offload_wrap_images = + *unsafe { lib.get::(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 { + ) -> Result<(PathBuf, Option), 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::>() - .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::>().join("\n* ") + ), + }) } } diff --git a/compiler/rustc_codegen_ssa/src/back/link.rs b/compiler/rustc_codegen_ssa/src/back/link.rs index 25003e071beb7..abe41f8ae8e7c 100644 --- a/compiler/rustc_codegen_ssa/src/back/link.rs +++ b/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, diff --git a/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp index 8c18f2453e9d8..bed54da0a7045 100644 --- a/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp +++ b/compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp @@ -1,18 +1,41 @@ #include "../SuppressLLVMWarnings.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/MapVector.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/ScopeExit.h" +#include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/Bitcode/BitcodeReader.h" #include "llvm/Bitcode/BitcodeWriter.h" +#include "llvm/Frontend/Offloading/OffloadWrapper.h" +#include "llvm/Frontend/Offloading/Utility.h" #include "llvm/IR/IRBuilder.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/MC/TargetRegistry.h" #include "llvm/Object/OffloadBinary.h" -#include "llvm/Support/CBindingWrapping.h" +#include "llvm/Support/CodeGen.h" +#include "llvm/Support/Error.h" #include "llvm/Support/FileOutputBuffer.h" +#include "llvm/Support/FileSystem.h" #include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/MemoryBufferRef.h" +#include "llvm/Support/Program.h" +#include "llvm/Support/raw_ostream.h" #include "llvm/Target/TargetMachine.h" +#include "llvm/Target/TargetOptions.h" +#include "llvm/TargetParser/Triple.h" #include "llvm/Transforms/Utils/Cloning.h" #include "llvm/Transforms/Utils/ModuleUtils.h" #include "llvm/Transforms/Utils/ValueMapper.h" +#include +#include +#include +#include + using namespace llvm; using namespace llvm::object; @@ -115,3 +138,214 @@ extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn, IRBuilder<> B(&entry); B.CreateBr(&clonedEntry); } + +static Error extractImages(StringRef DeviceBinPath, + SmallVectorImpl &Binaries) { + ErrorOr> BufOrErr = + MemoryBuffer::getFile(DeviceBinPath); + if (std::error_code EC = BufOrErr.getError()) + return createFileError(DeviceBinPath, EC); + std::unique_ptr Buf = std::move(*BufOrErr); + + if (!isAddrAligned(Align(OffloadBinary::getAlignment()), + Buf->getBufferStart())) + Buf = MemoryBuffer::getMemBufferCopy(Buf->getBuffer(), + Buf->getBufferIdentifier()); + + return extractOffloadBinaries(*Buf, Binaries); +} + +static bool hasOffloadEntries(Module &M) { + for (GlobalVariable &GV : M.globals()) + if (GV.hasSection() && GV.getSection() == "llvm_offload_entries") + return true; + return false; +} + +static bool reportAndFailWrappingImages(Error E, const char *What) { + handleAllErrors(std::move(E), [&](const ErrorInfoBase &EI) { + errs() << "LLVMRustOffloadWrapImages: " << What << ": " << EI.message() + << "\n"; + }); + return false; +} + +static Expected> +assembleWithPtxas(StringRef Ptx, StringRef Arch) { + const ErrorOr Ptxas = sys::findProgramByName("ptxas"); + if (!Ptxas) + return createStringError(Ptxas.getError(), "ptxas not found in PATH"); + + SmallString<128> PtxFilePath; + if (std::error_code E = + sys::fs::createTemporaryFile("rust-offload", "ptx", PtxFilePath)) + return errorCodeToError(E); + + SmallString<128> CubinFilePath; + if (std::error_code E = + sys::fs::createTemporaryFile("rust-offload", "cubin", CubinFilePath)) + return errorCodeToError(E); + + auto Cleanup = scope_exit([&] { + if (std::error_code E = sys::fs::remove(PtxFilePath)) + (void)reportAndFailWrappingImages( + errorCodeToError(E), "assembleWithPtxas: PtxFilePath cleanup"); + if (std::error_code E = sys::fs::remove(CubinFilePath)) + (void)reportAndFailWrappingImages( + errorCodeToError(E), "assembleWithPtxas: CubinFilePath cleanup"); + }); + + if (Error E = writeFile(PtxFilePath, Ptx)) + return std::move(E); + + const StringRef Args[] = { + *Ptxas, "-m64", "-O3", "--gpu-name", + Arch, "--output-file", CubinFilePath, PtxFilePath, + }; + + std::string ErrorMsg; + const int Status = + sys::ExecuteAndWait(*Ptxas, Args, std::nullopt, {}, 0, 0, &ErrorMsg); + + if (Status != 0) + return createStringError("assembleWithPtxas: status %d: %s", Status, + ErrorMsg.c_str()); + + ErrorOr> CubinOrError = + MemoryBuffer::getFileAsStream(CubinFilePath); + if (!CubinOrError) + return errorCodeToError(CubinOrError.getError()); + + return std::move(*CubinOrError); +} + +static Expected> +linkWithRustLld(StringRef Obj, StringRef LldPath) { + SmallString<128> ObjFilePath; + if (std::error_code E = + sys::fs::createTemporaryFile("rust-offload", "o", ObjFilePath)) + return errorCodeToError(E); + + SmallString<128> SoFilePath; + if (std::error_code E = + sys::fs::createTemporaryFile("rust-offload", "so", SoFilePath)) + return errorCodeToError(E); + + auto Cleanup = scope_exit([&] { + if (std::error_code E = sys::fs::remove(ObjFilePath)) + (void)reportAndFailWrappingImages(errorCodeToError(E), + "linkWithRustLld: ObjFilePath cleanup"); + if (std::error_code E = sys::fs::remove(SoFilePath)) + (void)reportAndFailWrappingImages(errorCodeToError(E), + "linkWithRustLld: SoFilePath cleanup"); + }); + + if (Error E = writeFile(ObjFilePath, Obj)) + return std::move(E); + + const StringRef Args[] = { + LldPath, "-flavor", "gnu", "-shared", + "--no-undefined", "-o", SoFilePath, ObjFilePath, + }; + + std::string ErrorMsg; + const int Status = + sys::ExecuteAndWait(LldPath, Args, std::nullopt, {}, 0, 0, &ErrorMsg); + + if (Status != 0) + return createStringError("linkWithRustLld: status %d: %s", Status, + ErrorMsg.c_str()); + + ErrorOr> ElfOrError = + MemoryBuffer::getFileAsStream(SoFilePath); + if (!ElfOrError) + return errorCodeToError(ElfOrError.getError()); + + return std::move(*ElfOrError); +} + +static Expected> +compileDeviceImage(const OffloadBinary &Input, const char *LldPath) { + const Triple DeviceTriple(Input.getTriple()); + const StringRef Arch = Input.getArch(); + + LLVMContext Ctx; + Expected> ImageObjOrError = + parseBitcodeFile(MemoryBufferRef(Input.getImage(), "device.bc"), Ctx); + if (!ImageObjOrError) + return ImageObjOrError.takeError(); + + std::string ErrorMsg; + const Target *DeviceTarget = + TargetRegistry::lookupTarget(DeviceTriple, ErrorMsg); + if (!DeviceTarget) + return createStringError(ErrorMsg); + + std::unique_ptr TM(DeviceTarget->createTargetMachine( + DeviceTriple, Arch, /*Features=*/"", TargetOptions(), Reloc::PIC_)); + if (!TM) + return createStringError("createTargetMachine failed for %s", + DeviceTriple.str().c_str()); + + const bool IsNvptx = DeviceTriple.isNVPTX(); + + legacy::PassManager PM; + SmallString<0> Emitted; + raw_svector_ostream OS(Emitted); + const CodeGenFileType FileType = + IsNvptx ? CodeGenFileType::AssemblyFile : CodeGenFileType::ObjectFile; + if (TM->addPassesToEmitFile(PM, OS, nullptr, FileType)) + return createStringError("target %s cannot emit %s", + DeviceTriple.str().c_str(), + IsNvptx ? "assembly" : "object"); + + PM.run(**ImageObjOrError); + + if (IsNvptx) + return assembleWithPtxas(Emitted, Arch); + if (DeviceTriple.isAMDGPU()) { + if (!LldPath || !*LldPath) + return createStringError("rust-lld path was not provided for %s", + DeviceTriple.str().c_str()); + return linkWithRustLld(Emitted, LldPath); + } + + return createStringError("unsupported offload target %s", + DeviceTriple.str().c_str()); +} + +extern "C" bool LLVMRustOffloadWrapImages(LLVMModuleRef HostMRef, + const char *LldPath, + const char *DeviceBinPath) { + Module &M = *unwrap(HostMRef); + if (!hasOffloadEntries(M)) + return true; + + SmallVector Binaries; + if (Error E = extractImages(DeviceBinPath, Binaries)) + return reportAndFailWrappingImages(std::move(E), "extract"); + + // LLVMRustBundleImages writes exactly one device image + if (Binaries.size() != 1) + return reportAndFailWrappingImages( + createStringError("expected exactly one device image, found %zu", + Binaries.size()), + "extract"); + + const OffloadBinary &Input = *Binaries.front().getBinary(); + + auto ImageOrErr = compileDeviceImage(Input, LldPath); + if (!ImageOrErr) + return reportAndFailWrappingImages(ImageOrErr.takeError(), + "device compile"); + + StringRef ImageBuf = (*ImageOrErr)->getBuffer(); + ArrayRef Image(ImageBuf.data(), ImageBuf.size()); + + if (Error E = offloading::wrapOpenMPBinaries( + M, {Image}, offloading::getOffloadEntryArray(M), /*Suffix=*/"", + /*Relocatable=*/ + false)) + return reportAndFailWrappingImages(std::move(E), "wrap"); + return true; +}