bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

update autodiff flags

Browse source
This commit is contained in:
Manuel Drehwald 2025-02-21 21:51:20 -05:00
parent 161a4bf6ff
commit e2d250c3f6
11 changed files with 204 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

1
compiler/rustc_codegen_llvm/messages.ftl
View file

@ -1,3 +1,4 @@
codegen_llvm_autodiff_without_enable = using the autodiff feature requires -Z autodiff=Enable
codegen_llvm_autodiff_without_lto = using the autodiff feature requires using fat-lto
codegen_llvm_copy_bitcode = failed to copy bitcode to object file: {$err}

87
compiler/rustc_codegen_llvm/src/back/lto.rs
View file

@ -586,6 +586,42 @@ fn thin_lto(
}
}
fn enable_autodiff_settings(ad: &[config::AutoDiff], module: &mut ModuleCodegen<ModuleLlvm>) {
for &val in ad {
match val {
config::AutoDiff::PrintModBefore => {
unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
}
config::AutoDiff::PrintPerf => {
llvm::set_print_perf(true);
}
config::AutoDiff::PrintAA => {
llvm::set_print_activity(true);
}
config::AutoDiff::PrintTA => {
llvm::set_print_type(true);
}
config::AutoDiff::Inline => {
llvm::set_inline(true);
}
config::AutoDiff::LooseTypes => {
llvm::set_loose_types(false);
}
config::AutoDiff::PrintSteps => {
llvm::set_print(true);
}
// We handle this below
config::AutoDiff::PrintModAfter => {}
// This is required and already checked
config::AutoDiff::Enable => {}
}
}
// This helps with handling enums for now.
llvm::set_strict_aliasing(false);
// FIXME(ZuseZ4): Test this, since it was added a long time ago.
llvm::set_rust_rules(true);
}
pub(crate) fn run_pass_manager(
cgcx: &CodegenContext<LlvmCodegenBackend>,
dcx: DiagCtxtHandle<'_>,
@ -604,34 +640,37 @@ pub(crate) fn run_pass_manager(
let opt_stage = if thin { llvm::OptStage::ThinLTO } else { llvm::OptStage::FatLTO };
let opt_level = config.opt_level.unwrap_or(config::OptLevel::No);
// If this rustc version was build with enzyme/autodiff enabled, and if users applied the
// `#[autodiff]` macro at least once, then we will later call llvm_optimize a second time.
debug!("running llvm pm opt pipeline");
unsafe {
write::llvm_optimize(
cgcx,
dcx,
module,
config,
opt_level,
opt_stage,
write::AutodiffStage::DuringAD,
)?;
// The PostAD behavior is the same that we would have if no autodiff was used.
// It will run the default optimization pipeline. If AD is enabled we select
// the DuringAD stage, which will disable vectorization and loop unrolling, and
// schedule two autodiff optimization + differentiation passes.
// We then run the llvm_optimize function a second time, to optimize the code which we generated
// in the enzyme differentiation pass.
let enable_ad = config.autodiff.contains(&config::AutoDiff::Enable);
let stage =
if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD };
if enable_ad {
enable_autodiff_settings(&config.autodiff, module);
}
// FIXME(ZuseZ4): Make this more granular
if cfg!(llvm_enzyme) && !thin {
unsafe {
write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
}
if cfg!(llvm_enzyme) && enable_ad {
let opt_stage = llvm::OptStage::FatLTO;
let stage = write::AutodiffStage::PostAD;
unsafe {
write::llvm_optimize(
cgcx,
dcx,
module,
config,
opt_level,
llvm::OptStage::FatLTO,
write::AutodiffStage::PostAD,
)?;
write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
}
// This is the final IR, so people should be able to inspect the optimized autodiff output.
if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
}
}
debug!("lto done");
Ok(())
}

13
compiler/rustc_codegen_llvm/src/builder/autodiff.rs
View file

@ -10,7 +10,7 @@ use crate::back::write::llvm_err;
use crate::builder::SBuilder;
use crate::context::SimpleCx;
use crate::declare::declare_simple_fn;
use crate::errors::LlvmError;
use crate::errors::{AutoDiffWithoutEnable, LlvmError};
use crate::llvm::AttributePlace::Function;
use crate::llvm::{Metadata, True};
use crate::value::Value;
@ -46,9 +46,6 @@ fn generate_enzyme_call<'ll>(
let output = attrs.ret_activity;
// We have to pick the name depending on whether we want forward or reverse mode autodiff.
// FIXME(ZuseZ4): The new pass based approach should not need the {Forward/Reverse}First method anymore, since
// it will handle higher-order derivatives correctly automatically (in theory). Currently
// higher-order derivatives fail, so we should debug that before adjusting this code.
let mut ad_name: String = match attrs.mode {
DiffMode::Forward => "__enzyme_fwddiff",
DiffMode::Reverse => "__enzyme_autodiff",
@ -291,6 +288,14 @@ pub(crate) fn differentiate<'ll>(
let diag_handler = cgcx.create_dcx();
let cx = SimpleCx { llmod: module.module_llvm.llmod(), llcx: module.module_llvm.llcx };
// First of all, did the user try to use autodiff without using the -Zautodiff=Enable flag?
if !diff_items.is_empty()
&& !cgcx.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable)
{
let dcx = cgcx.create_dcx();
return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutEnable));
}
// Before dumping the module, we want all the TypeTrees to become part of the module.
for item in diff_items.iter() {
let name = item.source.clone();

5
compiler/rustc_codegen_llvm/src/errors.rs
View file

@ -92,9 +92,12 @@ impl<G: EmissionGuarantee> Diagnostic<'_, G> for ParseTargetMachineConfig<'_> {
#[derive(Diagnostic)]
#[diag(codegen_llvm_autodiff_without_lto)]
#[note]
pub(crate) struct AutoDiffWithoutLTO;
#[derive(Diagnostic)]
#[diag(codegen_llvm_autodiff_without_enable)]
pub(crate) struct AutoDiffWithoutEnable;
#[derive(Diagnostic)]
#[diag(codegen_llvm_lto_disallowed)]
pub(crate) struct LtoDisallowed;

94
compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs
View file

@ -35,3 +35,97 @@ pub enum LLVMRustVerifierFailureAction {
LLVMPrintMessageAction = 1,
LLVMReturnStatusAction = 2,
}
#[cfg(llvm_enzyme)]
pub use self::Enzyme_AD::*;
#[cfg(llvm_enzyme)]
pub mod Enzyme_AD {
use libc::c_void;
extern "C" {
pub fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8);
}
extern "C" {
static mut EnzymePrintPerf: c_void;
static mut EnzymePrintActivity: c_void;
static mut EnzymePrintType: c_void;
static mut EnzymePrint: c_void;
static mut EnzymeStrictAliasing: c_void;
static mut looseTypeAnalysis: c_void;
static mut EnzymeInline: c_void;
static mut RustTypeRules: c_void;
}
pub fn set_print_perf(print: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintPerf), print as u8);
}
}
pub fn set_print_activity(print: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintActivity), print as u8);
}
}
pub fn set_print_type(print: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintType), print as u8);
}
}
pub fn set_print(print: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrint), print as u8);
}
}
pub fn set_strict_aliasing(strict: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeStrictAliasing), strict as u8);
}
}
pub fn set_loose_types(loose: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(looseTypeAnalysis), loose as u8);
}
}
pub fn set_inline(val: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeInline), val as u8);
}
}
pub fn set_rust_rules(val: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(RustTypeRules), val as u8);
}
}
}
#[cfg(not(llvm_enzyme))]
pub use self::Fallback_AD::*;
#[cfg(not(llvm_enzyme))]
pub mod Fallback_AD {
#![allow(unused_variables)]
pub fn set_inline(val: bool) {
unimplemented!()
}
pub fn set_print_perf(print: bool) {
unimplemented!()
}
pub fn set_print_activity(print: bool) {
unimplemented!()
}
pub fn set_print_type(print: bool) {
unimplemented!()
}
pub fn set_print(print: bool) {
unimplemented!()
}
pub fn set_strict_aliasing(strict: bool) {
unimplemented!()
}
pub fn set_loose_types(loose: bool) {
unimplemented!()
}
pub fn set_rust_rules(val: bool) {
unimplemented!()
}
}

3
compiler/rustc_codegen_ssa/src/back/write.rs
View file

@ -405,7 +405,8 @@ fn generate_lto_work<B: ExtraBackendMethods>(
B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
let config = cgcx.config(ModuleKind::Regular);
module = unsafe { module.autodiff(cgcx, autodiff, config).unwrap() };
module =
unsafe { module.autodiff(cgcx, autodiff, config).unwrap_or_else(|e| e.raise()) };
}
// We are adding a single work item, so the cost doesn't matter.
vec![(WorkItem::LTO(module), 0)]

2
compiler/rustc_interface/src/tests.rs
View file

@ -759,7 +759,7 @@ fn test_unstable_options_tracking_hash() {
tracked!(allow_features, Some(vec![String::from("lang_items")]));
tracked!(always_encode_mir, true);
tracked!(assume_incomplete_release, true);
tracked!(autodiff, vec![AutoDiff::Print]);
tracked!(autodiff, vec![AutoDiff::Enable]);
tracked!(binary_dep_depinfo, true);
tracked!(box_noalias, false);
tracked!(

25
compiler/rustc_session/src/config.rs
View file

@ -198,33 +198,26 @@ pub enum CoverageLevel {
/// The different settings that the `-Z autodiff` flag can have.
#[derive(Clone, Copy, PartialEq, Hash, Debug)]
pub enum AutoDiff {
/// Enable the autodiff opt pipeline
Enable,
/// Print TypeAnalysis information
PrintTA,
/// Print ActivityAnalysis Information
PrintAA,
/// Print Performance Warnings from Enzyme
PrintPerf,
/// Combines the three print flags above.
Print,
/// Print intermediate IR generation steps
PrintSteps,
/// Print the whole module, before running opts.
PrintModBefore,
/// Print the whole module just before we pass it to Enzyme.
/// For Debug purpose, prefer the OPT flag below
PrintModAfterOpts,
/// Print the module after Enzyme differentiated everything.
PrintModAfterEnzyme,
PrintModAfter,
/// Enzyme's loose type debug helper (can cause incorrect gradients)
/// Enzyme's loose type debug helper (can cause incorrect gradients!!)
/// Usable in cases where Enzyme errors with `can not deduce type of X`.
LooseTypes,
/// More flags
NoModOptAfter,
/// Tell Enzyme to run LLVM Opts on each function it generated. By default off,
/// since we already optimize the whole module after Enzyme is done.
EnableFncOpt,
NoVecUnroll,
RuntimeActivity,
/// Runs Enzyme specific Inlining
/// Runs Enzyme's aggressive inlining
Inline,
}

39
compiler/rustc_session/src/options.rs
View file

@ -707,7 +707,7 @@ mod desc {
pub(crate) const parse_list: &str = "a space-separated list of strings";
pub(crate) const parse_list_with_polarity: &str =
"a comma-separated list of strings, with elements beginning with + or -";
pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Print`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfterOpts`, `PrintModAfterEnzyme`, `LooseTypes`, `NoModOptAfter`, `EnableFncOpt`, `NoVecUnroll`, `Inline`";
pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `LooseTypes`, `Inline`";
pub(crate) const parse_comma_list: &str = "a comma-separated list of strings";
pub(crate) const parse_opt_comma_list: &str = parse_comma_list;
pub(crate) const parse_number: &str = "a number";
@ -1348,17 +1348,14 @@ pub mod parse {
v.sort_unstable();
for &val in v.iter() {
let variant = match val {
"Enable" => AutoDiff::Enable,
"PrintTA" => AutoDiff::PrintTA,
"PrintAA" => AutoDiff::PrintAA,
"PrintPerf" => AutoDiff::PrintPerf,
"Print" => AutoDiff::Print,
"PrintSteps" => AutoDiff::PrintSteps,
"PrintModBefore" => AutoDiff::PrintModBefore,
"PrintModAfterOpts" => AutoDiff::PrintModAfterOpts,
"PrintModAfterEnzyme" => AutoDiff::PrintModAfterEnzyme,
"PrintModAfter" => AutoDiff::PrintModAfter,
"LooseTypes" => AutoDiff::LooseTypes,
"NoModOptAfter" => AutoDiff::NoModOptAfter,
"EnableFncOpt" => AutoDiff::EnableFncOpt,
"NoVecUnroll" => AutoDiff::NoVecUnroll,
"Inline" => AutoDiff::Inline,
_ => {
// FIXME(ZuseZ4): print an error saying which value is not recognized
@ -2081,21 +2078,19 @@ options! {
assume_incomplete_release: bool = (false, parse_bool, [TRACKED],
"make cfg(version) treat the current version as incomplete (default: no)"),
autodiff: Vec<crate::config::AutoDiff> = (Vec::new(), parse_autodiff, [TRACKED],
"a list of optional autodiff flags to enable
Optional extra settings:
`=PrintTA`
`=PrintAA`
`=PrintPerf`
`=Print`
`=PrintModBefore`
`=PrintModAfterOpts`
`=PrintModAfterEnzyme`
`=LooseTypes`
`=NoModOptAfter`
`=EnableFncOpt`
`=NoVecUnroll`
`=Inline`
Multiple options can be combined with commas."),
"a list of autodiff flags to enable
Mandatory setting:
`=Enable`
Optional extra settings:
`=PrintTA`
`=PrintAA`
`=PrintPerf`
`=PrintSteps`
`=PrintModBefore`
`=PrintModAfter`
`=LooseTypes`
`=Inline`
Multiple options can be combined with commas."),
#[rustc_lint_opt_deny_field_access("use `Session::binary_dep_depinfo` instead of this field")]
binary_dep_depinfo: bool = (false, parse_bool, [TRACKED],
"include artifacts (sysroot, crate dependencies) used during compilation in dep-info \

9
src/doc/unstable-book/src/compiler-flags/autodiff.md
View file

@ -8,16 +8,13 @@ This feature allows you to differentiate functions using automatic differentiati
Set the `-Zautodiff=<options>` compiler flag to adjust the behaviour of the autodiff feature.
Multiple options can be separated with a comma. Valid options are:
`Enable` - Required flag to enable autodiff
`PrintTA` - print Type Analysis Information
`PrintAA` - print Activity Analysis Information
`PrintPerf` - print Performance Warnings from Enzyme
`Print` - prints all intermediate transformations
`PrintSteps` - prints all intermediate transformations
`PrintModBefore` - print the whole module, before running opts
`PrintModAfterOpts` - print the whole module just before we pass it to Enzyme
`PrintModAfterEnzyme` - print the module after Enzyme differentiated everything
`PrintModAfter` - print the module after Enzyme differentiated everything
`LooseTypes` - Enzyme's loose type debug helper (can cause incorrect gradients)
`Inline` - runs Enzyme specific Inlining
`NoModOptAfter` - do not optimize the module after Enzyme is done
`EnableFncOpt` - tell Enzyme to run LLVM Opts on each function it generated
`NoVecUnroll` - do not unroll vectorized loops
`RuntimeActivity` - allow specifying activity at runtime

2
tests/codegen/autodiff.rs
View file

@ -1,4 +1,4 @@
//@ compile-flags: -C opt-level=3 -Clto=fat
//@ compile-flags: -Zautodiff=Enable -C opt-level=3 -Clto=fat
//@ no-prefer-dynamic
//@ needs-enzyme
#![feature(autodiff)]