Skip to content

Commit c1e5961

Browse files
committed
gpu host code generation
1 parent 3dc0c3b commit c1e5961

File tree

3 files changed

+621
-0
lines changed

3 files changed

+621
-0
lines changed

compiler/rustc_codegen_llvm/src/back/lto.rs

Lines changed: 8 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,13 @@ 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+
dbg!(&enable_gpu);
673+
let cx =
674+
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
675+
crate::builder::gpu_offload::handle_gpu_code(cgcx, &cx);
676+
}
677+
670678
if cfg!(llvm_enzyme) && enable_ad && !thin {
671679
let cx =
672680
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);

compiler/rustc_codegen_llvm/src/builder.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use std::ops::Deref;
33
use std::{iter, ptr};
44

55
pub(crate) mod autodiff;
6+
pub(crate) mod gpu_offload;
67

78
use libc::{c_char, c_uint, size_t};
89
use rustc_abi as abi;

0 commit comments

Comments
 (0)