Skip to content

Commit b2a20e4

Browse files
ZuseZ4ZuseZ4
committed
wip, works
1 parent 33cea95 commit b2a20e4

File tree

3 files changed

+107
-9
lines changed

3 files changed

+107
-9
lines changed

‎compiler/rustc_codegen_llvm/src/back/write.rs‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -767,11 +767,36 @@ pub(crate) unsafe fn llvm_optimize(
767767
};
768768

769769
if cgcx.target_is_like_gpu && config.offload.contains(&config::Offload::Enable) {
770+
let lib_bc_c = CString::new("/p/lustre1/drehwald1/prog/offload/r/lib.bc").unwrap();
771+
let host_out_c = CString::new("/p/lustre1/drehwald1/prog/offload/r/host.out").unwrap();
772+
let out_obj_c = CString::new("/p/lustre1/drehwald1/prog/offload/r/host.o").unwrap();
773+
770774
unsafe {
771-
llvm::LLVMRustBundleImages(module.module_llvm.llmod(), module.module_llvm.tm.raw());
775+
llvm::LLVMRustBundleImages(
776+
module.module_llvm.llmod(),
777+
module.module_llvm.tm.raw(),
778+
host_out_c.as_ptr(),
779+
);
772780
}
773-
}
781+
unsafe {
782+
// 1) Bundle device module into offload image host.out (device TM)
783+
let ok = llvm::LLVMRustBundleImages(
784+
module.module_llvm.llmod(),
785+
module.module_llvm.tm.raw(),
786+
host_out_c.as_ptr(),
787+
);
788+
assert!(ok, "LLVMRustBundleImages (device -> host.out) failed");
774789

790+
// 2) Finalize host: lib.bc + host.out -> host.offload.o (host TM created in C++)
791+
let ok = llvm::LLVMRustFinalizeOffload(
792+
lib_bc_c.as_ptr(),
793+
host_out_c.as_ptr(),
794+
out_obj_c.as_ptr(),
795+
);
796+
assert!(ok, "LLVMRustFinalizeOffload (host finalize) failed");
797+
}
798+
dbg!("done");
799+
}
775800
result.into_result().unwrap_or_else(|()| llvm_err(dcx, LlvmError::RunLlvmPasses))
776801
}
777802

‎compiler/rustc_codegen_llvm/src/llvm/ffi.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1642,7 +1642,16 @@ unsafe extern "C" {
16421642
) -> &'a Value;
16431643

16441644
/// Processes the module and writes it in an offload compatible way into a "host.out" file.
1645-
pub(crate) fn LLVMRustBundleImages<'a>(M: &'a Module, TM: &'a TargetMachine) -> bool;
1645+
pub(crate) fn LLVMRustBundleImages<'a>(
1646+
M: &'a Module,
1647+
TM: &'a TargetMachine,
1648+
host_out: *const c_char,
1649+
) -> bool;
1650+
pub(crate) fn LLVMRustFinalizeOffload(
1651+
lib_bc_path: *const c_char,
1652+
host_out_path: *const c_char,
1653+
out_obj_path: *const c_char,
1654+
) -> bool;
16461655

16471656
/// Writes a module to the specified path. Returns 0 on success.
16481657
pub(crate) fn LLVMWriteBitcodeToFile(M: &Module, Path: *const c_char) -> c_int;
@@ -2028,6 +2037,7 @@ unsafe extern "C" {
20282037
) -> &Attribute;
20292038

20302039
// Operations on functions
2040+
20312041
pub(crate) fn LLVMRustOffloadMapper<'a>(Fn: &'a Value, Fn: &'a Value);
20322042
pub(crate) fn LLVMRustGetOrInsertFunction<'a>(
20332043
M: &'a Module,

‎compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp‎

Lines changed: 69 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ static Error writeFile(StringRef Filename, StringRef Data) {
177177
// --image=file=device.bc,triple=amdgcn-amd-amdhsa,arch=gfx90a,kind=openmp
178178
// The input module is the rust code compiled for a gpu target like amdgpu.
179179
// Based on clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
180-
extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM) {
180+
extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM, const char *HostOutPath) {
181181
std::string Storage;
182182
llvm::raw_string_ostream OS1(Storage);
183183
llvm::WriteBitcodeToFile(*unwrap(M), OS1);
@@ -202,16 +202,16 @@ extern "C" bool LLVMRustBundleImages(LLVMModuleRef M, TargetMachine &TM) {
202202
// Offload binary has invalid size alignment
203203
return false;
204204
OS2 << Buffer;
205-
if (Error E = writeFile("host.out",
205+
if (Error E = writeFile(HostOutPath,
206206
StringRef(BinaryData.begin(), BinaryData.size())))
207207
return false;
208208
return true;
209209
}
210210

211211
#include "llvm/Bitcode/BitcodeReader.h"
212212
Expected<std::unique_ptr<Module>>
213-
loadHostModuleFromBitcode(LLVMContext &Ctx) {
214-
auto MBOrErr = MemoryBuffer::getFile("/g/todo");
213+
loadHostModuleFromBitcode(LLVMContext &Ctx, StringRef LibBCPath) {
214+
auto MBOrErr = MemoryBuffer::getFile(LibBCPath);
215215
if (!MBOrErr)
216216
return errorCodeToError(MBOrErr.getError());
217217

@@ -243,11 +243,15 @@ extern "C" void embedBufferInModule(Module &M, MemoryBufferRef Buf) {
243243
}
244244

245245
Error embedHostOutIntoHostModule(Module &HostM, StringRef HostOutPath) {
246+
llvm::errs() << "embedHostOutIntoHostModule step 1:\n";
246247
auto MBOrErr = MemoryBuffer::getFile(HostOutPath);
248+
llvm::errs() << "embedHostOutIntoHostModule step 2:\n";
247249
if (!MBOrErr)
248250
return errorCodeToError(MBOrErr.getError());
249251

252+
llvm::errs() << "embedHostOutIntoHostModule step 3:\n";
250253
MemoryBufferRef Buf = (*MBOrErr)->getMemBufferRef();
254+
llvm::errs() << "embedHostOutIntoHostModule step 4:\n";
251255
embedBufferInModule(HostM, Buf);
252256
return Error::success();
253257
}
@@ -267,8 +271,8 @@ Error emitHostObjectWithTM(Module &HostM,
267271
TargetMachine &TM,
268272
StringRef OutObjPath) {
269273
// Make sure module matches the TM
270-
HostM.setDataLayout(TM.createDataLayout());
271-
HostM.setTargetTriple(TM.getTargetTriple().str());
274+
//HostM.setDataLayout(TM.createDataLayout());
275+
//HostM.setTargetTriple(TM.getTargetTriple().str());
272276

273277
legacy::PassManager PM;
274278
std::error_code EC;
@@ -304,6 +308,65 @@ extern "C" void LLVMRustOffloadMapper(LLVMValueRef OldFn, LLVMValueRef NewFn) {
304308
returns);
305309
}
306310

311+
// Create a host TargetMachine with HARDCODED triple/CPU
312+
static std::unique_ptr<TargetMachine> createHostTargetMachine() {
313+
static bool Initialized = false;
314+
if (!Initialized) {
315+
InitializeAllTargets();
316+
InitializeAllTargetMCs();
317+
InitializeAllAsmPrinters();
318+
InitializeAllAsmParsers();
319+
Initialized = true;
320+
}
321+
322+
// Hardcoded host triple + CPU (adapt if your CI/host differs)
323+
std::string TripleStr = "x86_64-unknown-linux-gnu";
324+
std::string CPU = "x86-64"; // OK for X86
325+
326+
std::string Err;
327+
const Target *T = TargetRegistry::lookupTarget(TripleStr, Err);
328+
if (!T) {
329+
// Could log Err here
330+
return nullptr;
331+
}
332+
333+
TargetOptions Opts;
334+
auto RM = std::optional<Reloc::Model>(Reloc::PIC_);
335+
336+
std::unique_ptr<TargetMachine> TM(
337+
T->createTargetMachine(TripleStr, CPU, /*Features*/"", Opts, RM));
338+
339+
return TM;
340+
}
341+
342+
// Top-level entry: host finalize in second rustc invocation
343+
// lib.bc (from first rustc) + host.out (from LLVMRustBundleImages) => host.offload.o
344+
extern "C" bool LLVMRustFinalizeOffload(const char *LibBCPath,
345+
const char *HostOutPath,
346+
const char *OutObjPath) {
347+
LLVMContext Ctx;
348+
349+
// 1. Load host lib.bc
350+
auto ModOrErr = loadHostModuleFromBitcode(Ctx, LibBCPath);
351+
if (!ModOrErr)
352+
return !errorToBool(ModOrErr.takeError());
353+
std::unique_ptr<Module> HostM = std::move(*ModOrErr);
354+
355+
// 2. Embed host.out
356+
if (Error E = embedHostOutIntoHostModule(*HostM, HostOutPath))
357+
return !errorToBool(std::move(E));
358+
359+
// 3. Create host TM and emit host object
360+
auto HostTM = createHostTargetMachine();
361+
if (!HostTM)
362+
return false;
363+
364+
if (Error E = emitHostObjectWithTM(*HostM, *HostTM, OutObjPath))
365+
return !errorToBool(std::move(E));
366+
367+
return true;
368+
}
369+
307370
extern "C" LLVMValueRef LLVMRustGetNamedValue(LLVMModuleRef M, const char *Name,
308371
size_t NameLen) {
309372
return wrap(unwrap(M)->getNamedValue(StringRef(Name, NameLen)));

0 commit comments

Comments
 (0)