Skip to content

Commit 4bb6ec0

Browse files
borsbors
committed
Auto merge of #136823 - matthiaskrgr:rollup-vp20mk1, r=matthiaskrgr
Rollup of 6 pull requests Successful merges: - #136419 (adding autodiff tests) - #136628 (ci: upgrade to crosstool-ng 1.27.0) - #136681 (resolve `llvm-config` path properly on cross builds) - #136714 (Update `compiler-builtins` to 0.1.146) - #136731 (rustc_middle: parallel: TyCtxt: remove "unsafe impl DynSend/DynSync") - #136791 (Disable DWARF in linker options for i686-unknown-uefi) Failed merges: - #136767 (improve host/cross target checking) r? `@ghost` `@rustbot` modify labels: rollup
2 parents 4b293d9 + 4fff320 commit 4bb6ec0

File tree

24 files changed

+291
-115
lines changed

24 files changed

+291
-115
lines changed

Diff for: ‎compiler/rustc_ast/src/expand/autodiff_attrs.rs

+6-18
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,6 @@ pub enum DiffMode {
3030
Forward,
3131
/// The target function, to be created using reverse mode AD.
3232
Reverse,
33-
/// The target function, to be created using forward mode AD.
34-
/// This target function will also be used as a source for higher order derivatives,
35-
/// so compute it before all Forward/Reverse targets and optimize it through llvm.
36-
ForwardFirst,
37-
/// The target function, to be created using reverse mode AD.
38-
/// This target function will also be used as a source for higher order derivatives,
39-
/// so compute it before all Forward/Reverse targets and optimize it through llvm.
40-
ReverseFirst,
4133
}
4234

4335
/// Dual and Duplicated (and their Only variants) are getting lowered to the same Enzyme Activity.
@@ -92,10 +84,10 @@ pub struct AutoDiffAttrs {
9284

9385
impl DiffMode {
9486
pub fn is_rev(&self) -> bool {
95-
matches!(self, DiffMode::Reverse | DiffMode::ReverseFirst)
87+
matches!(self, DiffMode::Reverse)
9688
}
9789
pub fn is_fwd(&self) -> bool {
98-
matches!(self, DiffMode::Forward | DiffMode::ForwardFirst)
90+
matches!(self, DiffMode::Forward)
9991
}
10092
}
10193

@@ -106,8 +98,6 @@ impl Display for DiffMode {
10698
DiffMode::Source => write!(f, "Source"),
10799
DiffMode::Forward => write!(f, "Forward"),
108100
DiffMode::Reverse => write!(f, "Reverse"),
109-
DiffMode::ForwardFirst => write!(f, "ForwardFirst"),
110-
DiffMode::ReverseFirst => write!(f, "ReverseFirst"),
111101
}
112102
}
113103
}
@@ -125,12 +115,12 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
125115
match mode {
126116
DiffMode::Error => false,
127117
DiffMode::Source => false,
128-
DiffMode::Forward | DiffMode::ForwardFirst => {
118+
DiffMode::Forward => {
129119
activity == DiffActivity::Dual
130120
|| activity == DiffActivity::DualOnly
131121
|| activity == DiffActivity::Const
132122
}
133-
DiffMode::Reverse | DiffMode::ReverseFirst => {
123+
DiffMode::Reverse => {
134124
activity == DiffActivity::Const
135125
|| activity == DiffActivity::Active
136126
|| activity == DiffActivity::ActiveOnly
@@ -166,10 +156,10 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
166156
return match mode {
167157
DiffMode::Error => false,
168158
DiffMode::Source => false,
169-
DiffMode::Forward | DiffMode::ForwardFirst => {
159+
DiffMode::Forward => {
170160
matches!(activity, Dual | DualOnly | Const)
171161
}
172-
DiffMode::Reverse | DiffMode::ReverseFirst => {
162+
DiffMode::Reverse => {
173163
matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
174164
}
175165
};
@@ -200,8 +190,6 @@ impl FromStr for DiffMode {
200190
"Source" => Ok(DiffMode::Source),
201191
"Forward" => Ok(DiffMode::Forward),
202192
"Reverse" => Ok(DiffMode::Reverse),
203-
"ForwardFirst" => Ok(DiffMode::ForwardFirst),
204-
"ReverseFirst" => Ok(DiffMode::ReverseFirst),
205193
_ => Err(()),
206194
}
207195
}

Diff for: ‎compiler/rustc_codegen_cranelift/patches/0029-stdlib-Disable-f16-and-f128-in-compiler-builtins.patch

+2-2
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ index 7165c3e48af..968552ad435 100644
1616

1717
[dependencies]
1818
core = { path = "../core" }
19-
-compiler_builtins = { version = "=0.1.145", features = ['rustc-dep-of-std'] }
20-
+compiler_builtins = { version = "=0.1.145", features = ['rustc-dep-of-std', 'no-f16-f128'] }
19+
-compiler_builtins = { version = "=0.1.146", features = ['rustc-dep-of-std'] }
20+
+compiler_builtins = { version = "=0.1.146", features = ['rustc-dep-of-std', 'no-f16-f128'] }
2121

2222
[dev-dependencies]
2323
rand = { version = "0.8.5", default-features = false, features = ["alloc"] }

Diff for: ‎compiler/rustc_codegen_llvm/src/back/lto.rs

+23-2
Original file line numberDiff line numberDiff line change
@@ -606,10 +606,31 @@ pub(crate) fn run_pass_manager(
606606

607607
// If this rustc version was build with enzyme/autodiff enabled, and if users applied the
608608
// `#[autodiff]` macro at least once, then we will later call llvm_optimize a second time.
609-
let first_run = true;
610609
debug!("running llvm pm opt pipeline");
611610
unsafe {
612-
write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, first_run)?;
611+
write::llvm_optimize(
612+
cgcx,
613+
dcx,
614+
module,
615+
config,
616+
opt_level,
617+
opt_stage,
618+
write::AutodiffStage::DuringAD,
619+
)?;
620+
}
621+
// FIXME(ZuseZ4): Make this more granular
622+
if cfg!(llvm_enzyme) && !thin {
623+
unsafe {
624+
write::llvm_optimize(
625+
cgcx,
626+
dcx,
627+
module,
628+
config,
629+
opt_level,
630+
llvm::OptStage::FatLTO,
631+
write::AutodiffStage::PostAD,
632+
)?;
633+
}
613634
}
614635
debug!("lto done");
615636
Ok(())

Diff for: ‎compiler/rustc_codegen_llvm/src/back/write.rs

+26-15
Original file line numberDiff line numberDiff line change
@@ -530,14 +530,24 @@ fn get_instr_profile_output_path(config: &ModuleConfig) -> Option<CString> {
530530
config.instrument_coverage.then(|| c"default_%m_%p.profraw".to_owned())
531531
}
532532

533+
// PreAD will run llvm opts but disable size increasing opts (vectorization, loop unrolling)
534+
// DuringAD is the same as above, but also runs the enzyme opt and autodiff passes.
535+
// PostAD will run all opts, including size increasing opts.
536+
#[derive(Debug, Eq, PartialEq)]
537+
pub(crate) enum AutodiffStage {
538+
PreAD,
539+
DuringAD,
540+
PostAD,
541+
}
542+
533543
pub(crate) unsafe fn llvm_optimize(
534544
cgcx: &CodegenContext<LlvmCodegenBackend>,
535545
dcx: DiagCtxtHandle<'_>,
536546
module: &ModuleCodegen<ModuleLlvm>,
537547
config: &ModuleConfig,
538548
opt_level: config::OptLevel,
539549
opt_stage: llvm::OptStage,
540-
skip_size_increasing_opts: bool,
550+
autodiff_stage: AutodiffStage,
541551
) -> Result<(), FatalError> {
542552
// Enzyme:
543553
// The whole point of compiler based AD is to differentiate optimized IR instead of unoptimized
@@ -550,12 +560,16 @@ pub(crate) unsafe fn llvm_optimize(
550560
let unroll_loops;
551561
let vectorize_slp;
552562
let vectorize_loop;
563+
let run_enzyme = cfg!(llvm_enzyme) && autodiff_stage == AutodiffStage::DuringAD;
553564

554565
// When we build rustc with enzyme/autodiff support, we want to postpone size-increasing
555-
// optimizations until after differentiation. FIXME(ZuseZ4): Before shipping on nightly,
566+
// optimizations until after differentiation. Our pipeline is thus: (opt + enzyme), (full opt).
567+
// We therefore have two calls to llvm_optimize, if autodiff is used.
568+
//
569+
// FIXME(ZuseZ4): Before shipping on nightly,
556570
// we should make this more granular, or at least check that the user has at least one autodiff
557571
// call in their code, to justify altering the compilation pipeline.
558-
if skip_size_increasing_opts && cfg!(llvm_enzyme) {
572+
if cfg!(llvm_enzyme) && autodiff_stage != AutodiffStage::PostAD {
559573
unroll_loops = false;
560574
vectorize_slp = false;
561575
vectorize_loop = false;
@@ -565,7 +579,7 @@ pub(crate) unsafe fn llvm_optimize(
565579
vectorize_slp = config.vectorize_slp;
566580
vectorize_loop = config.vectorize_loop;
567581
}
568-
trace!(?unroll_loops, ?vectorize_slp, ?vectorize_loop);
582+
trace!(?unroll_loops, ?vectorize_slp, ?vectorize_loop, ?run_enzyme);
569583
let using_thin_buffers = opt_stage == llvm::OptStage::PreLinkThinLTO || config.bitcode_needed();
570584
let pgo_gen_path = get_pgo_gen_path(config);
571585
let pgo_use_path = get_pgo_use_path(config);
@@ -633,6 +647,7 @@ pub(crate) unsafe fn llvm_optimize(
633647
vectorize_loop,
634648
config.no_builtins,
635649
config.emit_lifetime_markers,
650+
run_enzyme,
636651
sanitizer_options.as_ref(),
637652
pgo_gen_path.as_ref().map_or(std::ptr::null(), |s| s.as_ptr()),
638653
pgo_use_path.as_ref().map_or(std::ptr::null(), |s| s.as_ptr()),
@@ -684,18 +699,14 @@ pub(crate) unsafe fn optimize(
684699
_ => llvm::OptStage::PreLinkNoLTO,
685700
};
686701

687-
// If we know that we will later run AD, then we disable vectorization and loop unrolling
688-
let skip_size_increasing_opts = cfg!(llvm_enzyme);
702+
// If we know that we will later run AD, then we disable vectorization and loop unrolling.
703+
// Otherwise we pretend AD is already done and run the normal opt pipeline (=PostAD).
704+
// FIXME(ZuseZ4): Make this more granular, only set PreAD if we actually have autodiff
705+
// usages, not just if we build rustc with autodiff support.
706+
let autodiff_stage =
707+
if cfg!(llvm_enzyme) { AutodiffStage::PreAD } else { AutodiffStage::PostAD };
689708
return unsafe {
690-
llvm_optimize(
691-
cgcx,
692-
dcx,
693-
module,
694-
config,
695-
opt_level,
696-
opt_stage,
697-
skip_size_increasing_opts,
698-
)
709+
llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, autodiff_stage)
699710
};
700711
}
701712
Ok(())

Diff for: ‎compiler/rustc_codegen_llvm/src/builder/autodiff.rs

+14-32
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,9 @@ use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivit
44
use rustc_codegen_ssa::ModuleCodegen;
55
use rustc_codegen_ssa::back::write::ModuleConfig;
66
use rustc_errors::FatalError;
7-
use rustc_session::config::Lto;
87
use tracing::{debug, trace};
98

10-
use crate::back::write::{llvm_err, llvm_optimize};
9+
use crate::back::write::llvm_err;
1110
use crate::builder::SBuilder;
1211
use crate::context::SimpleCx;
1312
use crate::declare::declare_simple_fn;
@@ -53,8 +52,6 @@ fn generate_enzyme_call<'ll>(
5352
let mut ad_name: String = match attrs.mode {
5453
DiffMode::Forward => "__enzyme_fwddiff",
5554
DiffMode::Reverse => "__enzyme_autodiff",
56-
DiffMode::ForwardFirst => "__enzyme_fwddiff",
57-
DiffMode::ReverseFirst => "__enzyme_autodiff",
5855
_ => panic!("logic bug in autodiff, unrecognized mode"),
5956
}
6057
.to_string();
@@ -153,7 +150,7 @@ fn generate_enzyme_call<'ll>(
153150
_ => {}
154151
}
155152

156-
trace!("matching autodiff arguments");
153+
debug!("matching autodiff arguments");
157154
// We now handle the issue that Rust level arguments not always match the llvm-ir level
158155
// arguments. A slice, `&[f32]`, for example, is represented as a pointer and a length on
159156
// llvm-ir level. The number of activities matches the number of Rust level arguments, so we
@@ -164,10 +161,10 @@ fn generate_enzyme_call<'ll>(
164161
let mut activity_pos = 0;
165162
let outer_args: Vec<&llvm::Value> = get_params(outer_fn);
166163
while activity_pos < inputs.len() {
167-
let activity = inputs[activity_pos as usize];
164+
let diff_activity = inputs[activity_pos as usize];
168165
// Duplicated arguments received a shadow argument, into which enzyme will write the
169166
// gradient.
170-
let (activity, duplicated): (&Metadata, bool) = match activity {
167+
let (activity, duplicated): (&Metadata, bool) = match diff_activity {
171168
DiffActivity::None => panic!("not a valid input activity"),
172169
DiffActivity::Const => (enzyme_const, false),
173170
DiffActivity::Active => (enzyme_out, false),
@@ -222,7 +219,15 @@ fn generate_enzyme_call<'ll>(
222219
// A duplicated pointer will have the following two outer_fn arguments:
223220
// (..., ptr, ptr, ...). We add the following llvm-ir to our __enzyme call:
224221
// (..., metadata! enzyme_dup, ptr, ptr, ...).
225-
assert!(llvm::LLVMRustGetTypeKind(next_outer_ty) == llvm::TypeKind::Pointer);
222+
if matches!(
223+
diff_activity,
224+
DiffActivity::Duplicated | DiffActivity::DuplicatedOnly
225+
) {
226+
assert!(
227+
llvm::LLVMRustGetTypeKind(next_outer_ty) == llvm::TypeKind::Pointer
228+
);
229+
}
230+
// In the case of Dual we don't have assumptions, e.g. f32 would be valid.
226231
args.push(next_outer_arg);
227232
outer_pos += 2;
228233
activity_pos += 1;
@@ -277,7 +282,7 @@ pub(crate) fn differentiate<'ll>(
277282
module: &'ll ModuleCodegen<ModuleLlvm>,
278283
cgcx: &CodegenContext<LlvmCodegenBackend>,
279284
diff_items: Vec<AutoDiffItem>,
280-
config: &ModuleConfig,
285+
_config: &ModuleConfig,
281286
) -> Result<(), FatalError> {
282287
for item in &diff_items {
283288
trace!("{}", item);
@@ -318,29 +323,6 @@ pub(crate) fn differentiate<'ll>(
318323

319324
// FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
320325

321-
if let Some(opt_level) = config.opt_level {
322-
let opt_stage = match cgcx.lto {
323-
Lto::Fat => llvm::OptStage::PreLinkFatLTO,
324-
Lto::Thin | Lto::ThinLocal => llvm::OptStage::PreLinkThinLTO,
325-
_ if cgcx.opts.cg.linker_plugin_lto.enabled() => llvm::OptStage::PreLinkThinLTO,
326-
_ => llvm::OptStage::PreLinkNoLTO,
327-
};
328-
// This is our second opt call, so now we run all opts,
329-
// to make sure we get the best performance.
330-
let skip_size_increasing_opts = false;
331-
trace!("running Module Optimization after differentiation");
332-
unsafe {
333-
llvm_optimize(
334-
cgcx,
335-
diag_handler.handle(),
336-
module,
337-
config,
338-
opt_level,
339-
opt_stage,
340-
skip_size_increasing_opts,
341-
)?
342-
};
343-
}
344326
trace!("done with differentiate()");
345327

346328
Ok(())

Diff for: ‎compiler/rustc_codegen_llvm/src/llvm/ffi.rs

+1
Original file line numberDiff line numberDiff line change
@@ -2382,6 +2382,7 @@ unsafe extern "C" {
23822382
LoopVectorize: bool,
23832383
DisableSimplifyLibCalls: bool,
23842384
EmitLifetimeMarkers: bool,
2385+
RunEnzyme: bool,
23852386
SanitizerOptions: Option<&SanitizerOptions>,
23862387
PGOGenPath: *const c_char,
23872388
PGOUsePath: *const c_char,

Diff for: ‎compiler/rustc_codegen_ssa/src/codegen_attrs.rs

-2
Original file line numberDiff line numberDiff line change
@@ -914,8 +914,6 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
914914
let mode = match mode.as_str() {
915915
"Forward" => DiffMode::Forward,
916916
"Reverse" => DiffMode::Reverse,
917-
"ForwardFirst" => DiffMode::ForwardFirst,
918-
"ReverseFirst" => DiffMode::ReverseFirst,
919917
_ => {
920918
span_bug!(mode.span, "rustc_autodiff attribute contains invalid mode");
921919
}

Diff for: ‎compiler/rustc_llvm/build.rs

+4
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,10 @@ fn main() {
193193
cfg.define(&flag, None);
194194
}
195195

196+
if tracked_env_var_os("LLVM_ENZYME").is_some() {
197+
cfg.define("ENZYME", None);
198+
}
199+
196200
if tracked_env_var_os("LLVM_RUSTLLVM").is_some() {
197201
cfg.define("LLVM_RUSTLLVM", None);
198202
}

Diff for: ‎compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp

+20-2
Original file line numberDiff line numberDiff line change
@@ -688,14 +688,20 @@ struct LLVMRustSanitizerOptions {
688688
bool SanitizeKernelAddressRecover;
689689
};
690690

691+
// This symbol won't be available or used when Enzyme is not enabled
692+
#ifdef ENZYME
693+
extern "C" void registerEnzyme(llvm::PassBuilder &PB);
694+
#endif
695+
691696
extern "C" LLVMRustResult LLVMRustOptimize(
692697
LLVMModuleRef ModuleRef, LLVMTargetMachineRef TMRef,
693698
LLVMRustPassBuilderOptLevel OptLevelRust, LLVMRustOptStage OptStage,
694699
bool IsLinkerPluginLTO, bool NoPrepopulatePasses, bool VerifyIR,
695700
bool LintIR, bool UseThinLTOBuffers, bool MergeFunctions, bool UnrollLoops,
696701
bool SLPVectorize, bool LoopVectorize, bool DisableSimplifyLibCalls,
697-
bool EmitLifetimeMarkers, LLVMRustSanitizerOptions *SanitizerOptions,
698-
const char *PGOGenPath, const char *PGOUsePath, bool InstrumentCoverage,
702+
bool EmitLifetimeMarkers, bool RunEnzyme,
703+
LLVMRustSanitizerOptions *SanitizerOptions, const char *PGOGenPath,
704+
const char *PGOUsePath, bool InstrumentCoverage,
699705
const char *InstrProfileOutput, const char *PGOSampleUsePath,
700706
bool DebugInfoForProfiling, void *LlvmSelfProfiler,
701707
LLVMRustSelfProfileBeforePassCallback BeforePassCallback,
@@ -1010,6 +1016,18 @@ extern "C" LLVMRustResult LLVMRustOptimize(
10101016
MPM.addPass(NameAnonGlobalPass());
10111017
}
10121018

1019+
// now load "-enzyme" pass:
1020+
#ifdef ENZYME
1021+
if (RunEnzyme) {
1022+
registerEnzyme(PB);
1023+
if (auto Err = PB.parsePassPipeline(MPM, "enzyme")) {
1024+
std::string ErrMsg = toString(std::move(Err));
1025+
LLVMRustSetLastError(ErrMsg.c_str());
1026+
return LLVMRustResult::Failure;
1027+
}
1028+
}
1029+
#endif
1030+
10131031
// Upgrade all calls to old intrinsics first.
10141032
for (Module::iterator I = TheModule->begin(), E = TheModule->end(); I != E;)
10151033
UpgradeCallsToIntrinsic(&*I++); // must be post-increment, as we remove

Diff for: ‎compiler/rustc_middle/src/ty/context.rs

-3
Original file line numberDiff line numberDiff line change
@@ -1301,9 +1301,6 @@ pub struct TyCtxt<'tcx> {
13011301
gcx: &'tcx GlobalCtxt<'tcx>,
13021302
}
13031303

1304-
// Explicitly implement `DynSync` and `DynSend` for `TyCtxt` to short circuit trait resolution.
1305-
unsafe impl DynSend for TyCtxt<'_> {}
1306-
unsafe impl DynSync for TyCtxt<'_> {}
13071304
fn _assert_tcx_fields() {
13081305
sync::assert_dyn_sync::<&'_ GlobalCtxt<'_>>();
13091306
sync::assert_dyn_send::<&'_ GlobalCtxt<'_>>();

0 commit comments

Comments
 (0)