offload: remove device.bin and pass device bitcode directly

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -818,28 +818,8 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
         )
     };
 
-    if cgcx.target_is_like_gpu
-        && config.offload.iter().any(|o| matches!(o, config::Offload::Device(_)))
-    {
-        let device_path = cgcx.output_filenames.path(OutputType::Object);
-        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());
-        // 1) Bundle device module into offload image device.bin (device TM)
-        let ok = unsafe {
-            llvm::RustOffloadWrapper::get_instance().llvm_rust_bundle_images(
-                llmod,
-                tm,
-                device_out_c.as_c_str(),
-            )
-        };
-        if !ok || !device_out.exists() {
-            dcx.emit_err(crate::diagnostics::OffloadBundleImagesFailed);
-        }
-    }
-
     // 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
+    // `<crate>.bc` 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.
     if !cgcx.target_is_like_gpu && is_final_stage {
         if let Some(device_path) = config
@@ -850,19 +830,15 @@ pub(crate) unsafe fn llvm_optimize<'ll>(
             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")
-            {
+            } else if device_pathbuf.extension().and_then(|e| e.to_str()) != Some("bc") {
                 dcx.emit_err(crate::diagnostics::OffloadWrongFileName);
             } else if !device_pathbuf.exists() {
                 dcx.emit_err(crate::diagnostics::OffloadNonexistingPath);
             }
-            let device_bin_c = path_to_c_string(device_pathbuf.as_path());
+            let device_bitcode_c = path_to_c_string(device_pathbuf.as_path());
             let ok = unsafe {
                 llvm::RustOffloadWrapper::get_instance()
-                    .llvm_rust_offload_wrap_images(llmod, device_bin_c.as_c_str())
+                    .llvm_rust_offload_wrap_images(llmod, device_bitcode_c.as_c_str())
             };
             if !ok {
                 dcx.emit_err(crate::diagnostics::OffloadWrapImagesFailed);

compiler/rustc_codegen_llvm/src/diagnostics.rs

@@ -73,7 +73,7 @@ pub(crate) struct RustOffloadComponentMissing {
 
 #[derive(Diagnostic)]
 #[diag(
-    "using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/device.bin>"
+    "using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/<crate>.bc>"
 )]
 pub(crate) struct OffloadWithoutEnable;
 
@@ -82,25 +82,21 @@ pub(crate) struct OffloadWithoutEnable;
 pub(crate) struct OffloadWithoutFatLTO;
 
 #[derive(Diagnostic)]
-#[diag("using the `-Z offload=Host=/absolute/path/to/device.bin` flag requires an absolute path")]
+#[diag("using the `-Z offload=Host=/absolute/path/to/<crate>.bc` flag requires an absolute path")]
 pub(crate) struct OffloadWithoutAbsPath;
 
 #[derive(Diagnostic)]
 #[diag(
-    "using the `-Z offload=Host=/absolute/path/to/device.bin` flag must point to a `device.bin` file"
+    "using the `-Z offload=Host=/absolute/path/to/<crate>.bc` flag must point to a .bc file produced by the device compilation (--emit=llvm-bc)"
 )]
 pub(crate) struct OffloadWrongFileName;
 
 #[derive(Diagnostic)]
 #[diag(
-    "the given path/file to `device.bin` does not exist. Did you forget to run the device compilation first?"
+    "the given path/file to `<crate>.bc` does not exist. Did you forget to run the device compilation first?"
 )]
 pub(crate) struct OffloadNonexistingPath;
 
-#[derive(Diagnostic)]
-#[diag("call to BundleImages failed, `device.bin` was not created")]
-pub(crate) struct OffloadBundleImagesFailed;
-
 #[derive(Diagnostic)]
 #[diag("call to WrapImages failed, device image was not wrapped into the host module")]
 pub(crate) struct OffloadWrapImagesFailed;

compiler/rustc_codegen_llvm/src/llvm/offload_ffi.rs

@@ -2,9 +2,8 @@ use std::ffi::{CStr, c_char};
 use std::path::PathBuf;
 use std::sync::OnceLock;
 
-use super::ffi::{Module, TargetMachine, Value};
+use super::ffi::{Module, Value};
 
-type LLVMRustBundleImagesFn = unsafe extern "C" fn(&Module, &TargetMachine, *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;
@@ -16,7 +15,6 @@ use rustc_session::filesearch;
 use crate::llvm;
 
 pub(crate) struct RustOffloadWrapper {
-    LLVMRustBundleImages: LLVMRustBundleImagesFn,
     LLVMRustOffloadMapper: LLVMRustOffloadMapperFn,
     LLVMRustOffloadWrapImages: LLVMRustOffloadWrapImagesFn,
     lld_path: Option<PathBuf>,
@@ -54,27 +52,18 @@ impl RustOffloadWrapper {
             .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_wrapper(&self, v1: &Value, v2: &Value, vs: &[&Value]) {
         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,
+        device_bitcode_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())
+            (self.LLVMRustOffloadWrapImages)(host_m, lld_c.as_ptr(), device_bitcode_path.as_ptr())
         }
     }
 
@@ -84,15 +73,12 @@ impl RustOffloadWrapper {
         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 =
-            *unsafe { lib.get::<LLVMRustBundleImagesFn>(b"LLVMRustBundleImages\0")? };
         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,
             LLVMRustOffloadMapper: llvm_rust_offload_wrapper,
             LLVMRustOffloadWrapImages: llvm_rust_offload_wrap_images,
             lld_path,

compiler/rustc_llvm/llvm-wrapper/offload/OffloadWrapper.cpp

@@ -1,21 +1,18 @@
 #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/CodeGen.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileOutputBuffer.h"
@@ -36,7 +33,6 @@
 #include <utility>
 
 using namespace llvm;
-using namespace llvm::object;
 
 static Error writeFile(StringRef Filename, StringRef Data) {
   Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
@@ -50,43 +46,6 @@ static Error writeFile(StringRef Filename, StringRef Data) {
   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<char, 1024> 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;
-}
-
 // 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
@@ -122,22 +81,6 @@ extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn,
   B.CreateBr(&clonedEntry);
 }
 
-static Error extractImages(StringRef DeviceBinPath,
-                           SmallVectorImpl<OffloadFile> &Binaries) {
-  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
-      MemoryBuffer::getFile(DeviceBinPath);
-  if (std::error_code EC = BufOrErr.getError())
-    return createFileError(DeviceBinPath, EC);
-  std::unique_ptr<MemoryBuffer> 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")
@@ -248,15 +191,23 @@ linkWithRustLld(StringRef Obj, StringRef LldPath) {
 }
 
 static Expected<std::unique_ptr<MemoryBuffer>>
-compileDeviceImage(const OffloadBinary &Input, const char *LldPath) {
-  const Triple DeviceTriple(Input.getTriple());
-  const StringRef Arch = Input.getArch();
-
+compileDeviceImage(MemoryBufferRef Input, const char *LldPath) {
   LLVMContext Ctx;
-  Expected<std::unique_ptr<Module>> ImageObjOrError =
-      parseBitcodeFile(MemoryBufferRef(Input.getImage(), "device.bc"), Ctx);
-  if (!ImageObjOrError)
-    return ImageObjOrError.takeError();
+  Expected<std::unique_ptr<Module>> ModOrErr = parseBitcodeFile(Input, Ctx);
+  if (!ModOrErr)
+    return ModOrErr.takeError();
+  Module &M = **ModOrErr;
+
+  const Triple &DeviceTriple = M.getTargetTriple();
+  StringRef Arch;
+  for (Function &F : M)
+    if (F.hasFnAttribute("target-cpu")) {
+      Arch = F.getFnAttribute("target-cpu").getValueAsString();
+      break;
+    }
+
+  if (Arch.empty())
+    return createStringError("no target-cpu attribute in device bitcode");
 
   std::string ErrorMsg;
   const Target *DeviceTarget =
@@ -282,7 +233,7 @@ compileDeviceImage(const OffloadBinary &Input, const char *LldPath) {
                              DeviceTriple.str().c_str(),
                              IsNvptx ? "assembly" : "object");
 
-  PM.run(**ImageObjOrError);
+  PM.run(**ModOrErr);
 
   if (IsNvptx)
     return assembleWithPtxas(Emitted, Arch);
@@ -299,25 +250,20 @@ compileDeviceImage(const OffloadBinary &Input, const char *LldPath) {
 
 extern "C" bool LLVMRustOffloadWrapImages(LLVMModuleRef HostMRef,
                                           const char *LldPath,
-                                          const char *DeviceBinPath) {
+                                          const char *DeviceBitcodePath) {
   Module &M = *unwrap(HostMRef);
   if (!hasOffloadEntries(M))
     return true;
 
-  SmallVector<OffloadFile> 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();
+  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
+      MemoryBuffer::getFile(DeviceBitcodePath);
+  if (std::error_code EC = BufOrErr.getError())
+    return reportAndFailWrappingImages(createFileError(DeviceBitcodePath, EC),
+                                       "device bitcode");
+  std::unique_ptr<MemoryBuffer> Buf = std::move(*BufOrErr);
 
-  auto ImageOrErr = compileDeviceImage(Input, LldPath);
+  Expected<std::unique_ptr<MemoryBuffer>> ImageOrErr =
+      compileDeviceImage(Buf->getMemBufferRef(), LldPath);
   if (!ImageOrErr)
     return reportAndFailWrappingImages(ImageOrErr.takeError(),
                                        "device compile");

tests/ui/offload/check_config.fail.stderr

@@ -1,4 +1,4 @@
-error: using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/device.bin>
+error: using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/<crate>.bc>
 
 error: using the offload feature requires -C lto=fat
 

tests/ui/offload/check_config.rs

@@ -6,7 +6,7 @@
 //@[pass] compile-flags: -Zunstable-options -Zoffload=Device -Clto=fat --emit=metadata
 //@[fail] compile-flags: -Clto=thin
 
-//[fail]~? ERROR: using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/device.bin>
+//[fail]~? ERROR: using the offload feature requires -Z offload=<Device or Host=/absolute/path/to/<crate>.bc>
 //[fail]~? ERROR: using the offload feature requires -C lto=fat
 
 #![feature(gpu_offload)]