user0/rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove the need to call clang for std::offload usages

Browse source
This commit is contained in:
Manuel Drehwald 2025-12-09 07:50:30 -08:00
parent 0ac9e59d8f
commit dfef2e96fe
8 changed files with 149 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

32
compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
View file

@ -43,8 +43,10 @@
// available. As such, we only try to build it in the first place, if
// llvm.offload is enabled.
#ifdef OFFLOAD
#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/Object/OffloadBinary.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
#endif
// for raw `write` in the bad-alloc handler
@ -174,12 +176,13 @@ static Error writeFile(StringRef Filename, StringRef Data) {
// --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) {
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, "module.bc");
auto MB = llvm::MemoryBuffer::getMemBufferCopy(Storage, "device.bc");
SmallVector<char, 1024> BinaryData;
raw_svector_ostream OS2(BinaryData);
@ -188,19 +191,38 @@ extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM) {
ImageBinary.TheImageKind = object::IMG_Bitcode;
ImageBinary.Image = std::move(MB);
ImageBinary.TheOffloadKind = object::OFK_OpenMP;
ImageBinary.StringData["triple"] = TM.getTargetTriple().str();
ImageBinary.StringData["arch"] = TM.getTargetCPU();
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("host.out",
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;
}
extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn) {
llvm::Function *oldFn = llvm::unwrap<llvm::Function>(OldFn);
llvm::Function *newFn = llvm::unwrap<llvm::Function>(NewFn);