Skip to content

Commit 34fab8c

Browse files
committed
gpu host code generation
1 parent f89744f commit 34fab8c

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
@@ -653,6 +653,7 @@ pub(crate) fn run_pass_manager(
653653
// We then run the llvm_optimize function a second time, to optimize the code which we generated
654654
// in the enzyme differentiation pass.
655655
let enable_ad = config.autodiff.contains(&config::AutoDiff::Enable);
656+
let enable_gpu = config.offload.contains(&config::Offload::Enable);
656657
let stage = if thin {
657658
write::AutodiffStage::PreAD
658659
} else {
@@ -667,6 +668,12 @@ pub(crate) fn run_pass_manager(
667668
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
668669
}
669670

671+
if cfg!(llvm_enzyme) && enable_gpu && !thin {
672+
let cx =
673+
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
674+
crate::builder::gpu_offload::handle_gpu_code(cgcx, &cx);
675+
}
676+
670677
if cfg!(llvm_enzyme) && enable_ad && !thin {
671678
let cx =
672679
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);

0 commit comments

Comments
 (0)