4545#ifdef OFFLOAD
4646#include " llvm/Object/OffloadBinary.h"
4747#include " llvm/Target/TargetMachine.h"
48+ #include " llvm/Transforms/Utils/ModuleUtils.h"
49+ #include " llvm/Analysis/VectorUtils.h"
50+ #include " llvm/ADT/SmallString.h"
51+ #include " llvm/IR/DerivedTypes.h"
52+ #include " llvm/IR/Function.h"
53+ #include " llvm/IR/IRBuilder.h"
54+ #include " llvm/IR/MDBuilder.h"
55+ #include " llvm/IR/Module.h"
56+ #include " llvm/Support/Casting.h"
57+ #include " llvm/Support/MD5.h"
58+ #include " llvm/Support/raw_ostream.h"
59+ #include " llvm/Support/xxhash.h"
4860#endif
4961
5062// for raw `write` in the bad-alloc handler
@@ -174,12 +186,12 @@ static Error writeFile(StringRef Filename, StringRef Data) {
174186// --image=file=device.bc,triple=amdgcn-amd-amdhsa,arch=gfx90a,kind=openmp
175187// The input module is the rust code compiled for a gpu target like amdgpu.
176188// Based on clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
177- extern " C" bool LLVMRustBundleImages (LLVMModuleRef M, TargetMachine &TM) {
189+ extern " C" bool LLVMRustBundleImages (LLVMModuleRef M, TargetMachine &TM, const char *HostOutPath ) {
178190 std::string Storage;
179191 llvm::raw_string_ostream OS1 (Storage);
180192 llvm::WriteBitcodeToFile (*unwrap (M), OS1);
181193 OS1.flush ();
182- auto MB = llvm::MemoryBuffer::getMemBufferCopy (Storage, " module .bc" );
194+ auto MB = llvm::MemoryBuffer::getMemBufferCopy (Storage, " device .bc" );
183195
184196 SmallVector<char , 1024 > BinaryData;
185197 raw_svector_ostream OS2 (BinaryData);
@@ -188,19 +200,103 @@ extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM) {
188200 ImageBinary.TheImageKind = object::IMG_Bitcode;
189201 ImageBinary.Image = std::move (MB);
190202 ImageBinary.TheOffloadKind = object::OFK_OpenMP;
191- ImageBinary.StringData [" triple" ] = TM.getTargetTriple ().str ();
192- ImageBinary.StringData [" arch" ] = TM.getTargetCPU ();
203+
204+
205+ std::string TripleStr = TM.getTargetTriple ().str ();
206+ llvm::StringRef CPURef = TM.getTargetCPU ();
207+ ImageBinary.StringData [" triple" ] = TripleStr;
208+ ImageBinary.StringData [" arch" ] = CPURef;
193209 llvm::SmallString<0 > Buffer = OffloadBinary::write (ImageBinary);
194210 if (Buffer.size () % OffloadBinary::getAlignment () != 0 )
195211 // Offload binary has invalid size alignment
196212 return false ;
197213 OS2 << Buffer;
198- if (Error E = writeFile (" host.out " ,
214+ if (Error E = writeFile (HostOutPath ,
199215 StringRef (BinaryData.begin (), BinaryData.size ())))
200216 return false ;
201217 return true ;
202218}
203219
220+ #include " llvm/Bitcode/BitcodeReader.h"
221+ Expected<std::unique_ptr<Module>>
222+ loadHostModuleFromBitcode (LLVMContext &Ctx, StringRef LibBCPath) {
223+ auto MBOrErr = MemoryBuffer::getFile (LibBCPath);
224+ if (!MBOrErr)
225+ return errorCodeToError (MBOrErr.getError ());
226+
227+ MemoryBufferRef Ref = (*MBOrErr)->getMemBufferRef ();
228+ return parseBitcodeFile (Ref, Ctx);
229+ }
230+
231+ extern " C" void embedBufferInModule (Module &M, MemoryBufferRef Buf) {
232+ StringRef SectionName = " .llvm.offloading" ;
233+ Align Alignment = Align (8 );
234+ // Embed the memory buffer into the module.
235+ Constant *ModuleConstant = ConstantDataArray::get (
236+ M.getContext (), ArrayRef (Buf.getBufferStart (), Buf.getBufferSize ()));
237+ GlobalVariable *GV = new GlobalVariable (
238+ M, ModuleConstant->getType (), true , GlobalValue::PrivateLinkage,
239+ ModuleConstant, " llvm.embedded.object" );
240+ GV->setSection (SectionName);
241+ GV->setAlignment (Alignment);
242+
243+ LLVMContext &Ctx = M.getContext ();
244+ NamedMDNode *MD = M.getOrInsertNamedMetadata (" llvm.embedded.objects" );
245+ Metadata *MDVals[] = {ConstantAsMetadata::get (GV),
246+ MDString::get (Ctx, SectionName)};
247+
248+ MD->addOperand (llvm::MDNode::get (Ctx, MDVals));
249+ GV->setMetadata (LLVMContext::MD_exclude, llvm::MDNode::get (Ctx, {}));
250+
251+ appendToCompilerUsed (M, GV);
252+ }
253+
254+ Error embedHostOutIntoHostModule (Module &HostM, StringRef HostOutPath) {
255+ llvm::errs () << " embedHostOutIntoHostModule step 1:\n " ;
256+ auto MBOrErr = MemoryBuffer::getFile (HostOutPath);
257+ llvm::errs () << " embedHostOutIntoHostModule step 2:\n " ;
258+ if (!MBOrErr)
259+ return errorCodeToError (MBOrErr.getError ());
260+
261+ llvm::errs () << " embedHostOutIntoHostModule step 3:\n " ;
262+ MemoryBufferRef Buf = (*MBOrErr)->getMemBufferRef ();
263+ llvm::errs () << " embedHostOutIntoHostModule step 4:\n " ;
264+ embedBufferInModule (HostM, Buf);
265+ return Error::success ();
266+ }
267+
268+ #include " llvm/Support/TargetSelect.h"
269+ #include " llvm/Target/TargetMachine.h"
270+ #include " llvm/Target/TargetOptions.h"
271+ #include " llvm/IR/LegacyPassManager.h"
272+ // #include "llvm/Support/Host.h"
273+ // #include "llvm/Support/TargetRegistry.h"
274+ #include " llvm/MC/TargetRegistry.h"
275+ #include " llvm/Support/raw_ostream.h"
276+ #include " llvm/Support/FileSystem.h"
277+ #include " llvm/Support/CodeGen.h" // <-- new
278+
279+ Error emitHostObjectWithTM (Module &HostM,
280+ TargetMachine &TM,
281+ StringRef OutObjPath) {
282+ // Make sure module matches the TM
283+ // HostM.setDataLayout(TM.createDataLayout());
284+ // HostM.setTargetTriple(TM.getTargetTriple().str());
285+
286+ legacy::PassManager PM;
287+ std::error_code EC;
288+ raw_fd_ostream OS (OutObjPath, EC, sys::fs::OF_None);
289+ if (EC)
290+ return errorCodeToError (EC);
291+
292+ if (TM.addPassesToEmitFile (PM, OS, nullptr , llvm::CodeGenFileType::ObjectFile))
293+ return createStringError (inconvertibleErrorCode (),
294+ " TargetMachine can't emit a file of this type" );
295+
296+ PM.run (HostM);
297+ return Error::success ();
298+ }
299+
204300extern " C" void LLVMRustOffloadMapper (LLVMValueRef OldFn, LLVMValueRef NewFn) {
205301 llvm::Function *oldFn = llvm::unwrap<llvm::Function>(OldFn);
206302 llvm::Function *newFn = llvm::unwrap<llvm::Function>(NewFn);
@@ -222,6 +318,65 @@ extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn) {
222318}
223319#endif
224320
321+ // Create a host TargetMachine with HARDCODED triple/CPU
322+ static std::unique_ptr<TargetMachine> createHostTargetMachine () {
323+ static bool Initialized = false ;
324+ if (!Initialized) {
325+ InitializeAllTargets ();
326+ InitializeAllTargetMCs ();
327+ InitializeAllAsmPrinters ();
328+ InitializeAllAsmParsers ();
329+ Initialized = true ;
330+ }
331+
332+ // Hardcoded host triple + CPU (adapt if your CI/host differs)
333+ std::string TripleStr = " x86_64-unknown-linux-gnu" ;
334+ std::string CPU = " x86-64" ; // OK for X86
335+
336+ std::string Err;
337+ const Target *T = TargetRegistry::lookupTarget (TripleStr, Err);
338+ if (!T) {
339+ // Could log Err here
340+ return nullptr ;
341+ }
342+
343+ TargetOptions Opts;
344+ auto RM = std::optional<Reloc::Model>(Reloc::PIC_);
345+
346+ std::unique_ptr<TargetMachine> TM (
347+ T->createTargetMachine (TripleStr, CPU, /* Features*/ " " , Opts, RM));
348+
349+ return TM;
350+ }
351+
352+ // Top-level entry: host finalize in second rustc invocation
353+ // lib.bc (from first rustc) + host.out (from LLVMRustBundleImages) => host.offload.o
354+ extern " C" bool LLVMRustFinalizeOffload (const char *LibBCPath,
355+ const char *HostOutPath,
356+ const char *OutObjPath) {
357+ LLVMContext Ctx;
358+
359+ // 1. Load host lib.bc
360+ auto ModOrErr = loadHostModuleFromBitcode (Ctx, LibBCPath);
361+ if (!ModOrErr)
362+ return !errorToBool (ModOrErr.takeError ());
363+ std::unique_ptr<Module> HostM = std::move (*ModOrErr);
364+
365+ // 2. Embed host.out
366+ if (Error E = embedHostOutIntoHostModule (*HostM, HostOutPath))
367+ return !errorToBool (std::move (E));
368+
369+ // 3. Create host TM and emit host object
370+ auto HostTM = createHostTargetMachine ();
371+ if (!HostTM)
372+ return false ;
373+
374+ if (Error E = emitHostObjectWithTM (*HostM, *HostTM, OutObjPath))
375+ return !errorToBool (std::move (E));
376+
377+ return true ;
378+ }
379+
225380extern " C" LLVMValueRef LLVMRustGetNamedValue (LLVMModuleRef M, const char *Name,
226381 size_t NameLen) {
227382 return wrap (unwrap (M)->getNamedValue (StringRef (Name, NameLen)));
0 commit comments