Skip to content

Commit 0fdf25f

Browse files
committed
gpu host code generation
1 parent fabc472 commit 0fdf25f

File tree

4 files changed

+478
-7
lines changed

4 files changed

+478
-7
lines changed

compiler/rustc_codegen_llvm/src/back/lto.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -657,6 +657,7 @@ pub(crate) fn run_pass_manager(
657657
// We then run the llvm_optimize function a second time, to optimize the code which we generated
658658
// in the enzyme differentiation pass.
659659
let enable_ad = config.autodiff.contains(&config::AutoDiff::Enable);
660+
let enable_gpu = config.offload.contains(&config::Offload::Enable);
660661
let stage = if thin {
661662
write::AutodiffStage::PreAD
662663
} else {
@@ -671,6 +672,12 @@ pub(crate) fn run_pass_manager(
671672
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
672673
}
673674

675+
if cfg!(llvm_enzyme) && enable_gpu && !thin {
676+
let cx =
677+
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
678+
crate::builder::gpu_offload::handle_gpu_code(cgcx, &cx);
679+
}
680+
674681
if cfg!(llvm_enzyme) && enable_ad && !thin {
675682
let cx =
676683
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);

0 commit comments

Comments
 (0)