Skip to content

Commit

Permalink
Auto merge of rust-lang#134245 - Zalathar:rollup-2dhkxya, r=Zalathar
Browse files Browse the repository at this point in the history
Rollup of 9 pull requests

Successful merges:

 - rust-lang#130060 (Autodiff Upstreaming - rustc_codegen_llvm changes)
 - rust-lang#132038 (Add lint rule for `#[deprecated]` on re-exports)
 - rust-lang#132150 (Fix powerpc64 big-endian FreeBSD ABI)
 - rust-lang#133633 (don't show the full linker args unless `--verbose` is passed)
 - rust-lang#133942 (Clarify how to use `black_box()`)
 - rust-lang#134081 (Try to evaluate constants in legacy mangling)
 - rust-lang#134192 (Remove `Lexer`'s dependency on `Parser`.)
 - rust-lang#134208 (coverage: Tidy up creation of covmap and covfun records)
 - rust-lang#134211 (On Neutrino QNX, reduce the need to set archiver via environment variables)

Failed merges:

 - rust-lang#133099 (forbid toggling x87 and fpregs on hard-float targets)

r? `@ghost`
`@rustbot` modify labels: rollup
  • Loading branch information
bors committed Dec 13, 2024
2 parents 3da8bfb + 0207615 commit 7dc39d6
Show file tree
Hide file tree
Showing 49 changed files with 1,035 additions and 299 deletions.
19 changes: 3 additions & 16 deletions compiler/rustc_ast/src/expand/autodiff_attrs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
use std::fmt::{self, Display, Formatter};
use std::str::FromStr;

use crate::expand::typetree::TypeTree;
use crate::expand::{Decodable, Encodable, HashStable_Generic};
use crate::ptr::P;
use crate::{Ty, TyKind};
Expand Down Expand Up @@ -79,10 +78,6 @@ pub struct AutoDiffItem {
/// The name of the function being generated
pub target: String,
pub attrs: AutoDiffAttrs,
/// Describe the memory layout of input types
pub inputs: Vec<TypeTree>,
/// Describe the memory layout of the output type
pub output: TypeTree,
}
#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
pub struct AutoDiffAttrs {
Expand Down Expand Up @@ -262,22 +257,14 @@ impl AutoDiffAttrs {
!matches!(self.mode, DiffMode::Error | DiffMode::Source)
}

pub fn into_item(
self,
source: String,
target: String,
inputs: Vec<TypeTree>,
output: TypeTree,
) -> AutoDiffItem {
AutoDiffItem { source, target, inputs, output, attrs: self }
pub fn into_item(self, source: String, target: String) -> AutoDiffItem {
AutoDiffItem { source, target, attrs: self }
}
}

impl fmt::Display for AutoDiffItem {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Differentiating {} -> {}", self.source, self.target)?;
write!(f, " with attributes: {:?}", self.attrs)?;
write!(f, " with inputs: {:?}", self.inputs)?;
write!(f, " with output: {:?}", self.output)
write!(f, " with attributes: {:?}", self.attrs)
}
}
9 changes: 9 additions & 0 deletions compiler/rustc_codegen_gcc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ use gccjit::{CType, Context, OptimizationLevel};
#[cfg(feature = "master")]
use gccjit::{TargetInfo, Version};
use rustc_ast::expand::allocator::AllocatorKind;
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
use rustc_codegen_ssa::back::write::{
CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
Expand Down Expand Up @@ -439,6 +440,14 @@ impl WriteBackendMethods for GccCodegenBackend {
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
back::write::link(cgcx, dcx, modules)
}
fn autodiff(
_cgcx: &CodegenContext<Self>,
_module: &ModuleCodegen<Self::Module>,
_diff_fncs: Vec<AutoDiffItem>,
_config: &ModuleConfig,
) -> Result<(), FatalError> {
unimplemented!()
}
}

/// This is the entrypoint for a hot plugged rustc_codegen_gccjit
Expand Down
4 changes: 4 additions & 0 deletions compiler/rustc_codegen_llvm/messages.ftl
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
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}
codegen_llvm_dynamic_linking_with_lto =
Expand Down Expand Up @@ -47,6 +49,8 @@ codegen_llvm_parse_bitcode_with_llvm_err = failed to parse bitcode for LTO modul
codegen_llvm_parse_target_machine_config =
failed to parse target machine config to target machine: {$error}
codegen_llvm_prepare_autodiff = failed to prepare autodiff: src: {$src}, target: {$target}, {$error}
codegen_llvm_prepare_autodiff_with_llvm_err = failed to prepare autodiff: {$llvm_err}, src: {$src}, target: {$target}, {$error}
codegen_llvm_prepare_thin_lto_context = failed to prepare thin LTO context
codegen_llvm_prepare_thin_lto_context_with_llvm_err = failed to prepare thin LTO context: {$llvm_err}
Expand Down
9 changes: 8 additions & 1 deletion compiler/rustc_codegen_llvm/src/back/lto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -604,7 +604,14 @@ pub(crate) fn run_pass_manager(
debug!("running the 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);
unsafe { write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage) }?;

// 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.
let first_run = true;
debug!("running llvm pm opt pipeline");
unsafe {
write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, first_run)?;
}
debug!("lto done");
Ok(())
}
Expand Down
134 changes: 128 additions & 6 deletions compiler/rustc_codegen_llvm/src/back/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use libc::{c_char, c_int, c_void, size_t};
use llvm::{
LLVMRustLLVMHasZlibCompressionForDebugSymbols, LLVMRustLLVMHasZstdCompressionForDebugSymbols,
};
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
use rustc_codegen_ssa::back::link::ensure_removed;
use rustc_codegen_ssa::back::versioned_llvm_target;
use rustc_codegen_ssa::back::write::{
Expand All @@ -28,7 +29,7 @@ use rustc_session::config::{
use rustc_span::symbol::sym;
use rustc_span::{BytePos, InnerSpan, Pos, SpanData, SyntaxContext};
use rustc_target::spec::{CodeModel, RelocModel, SanitizerSet, SplitDebuginfo, TlsModel};
use tracing::debug;
use tracing::{debug, trace};

use crate::back::lto::ThinBuffer;
use crate::back::owned_target_machine::OwnedTargetMachine;
Expand Down Expand Up @@ -530,9 +531,38 @@ pub(crate) unsafe fn llvm_optimize(
config: &ModuleConfig,
opt_level: config::OptLevel,
opt_stage: llvm::OptStage,
skip_size_increasing_opts: bool,
) -> Result<(), FatalError> {
let unroll_loops =
opt_level != config::OptLevel::Size && opt_level != config::OptLevel::SizeMin;
// Enzyme:
// The whole point of compiler based AD is to differentiate optimized IR instead of unoptimized
// source code. However, benchmarks show that optimizations increasing the code size
// tend to reduce AD performance. Therefore deactivate them before AD, then differentiate the code
// and finally re-optimize the module, now with all optimizations available.
// FIXME(ZuseZ4): In a future update we could figure out how to only optimize individual functions getting
// differentiated.

let unroll_loops;
let vectorize_slp;
let vectorize_loop;

// When we build rustc with enzyme/autodiff support, we want to postpone size-increasing
// optimizations until after differentiation. FIXME(ZuseZ4): Before shipping on nightly,
// we should make this more granular, or at least check that the user has at least one autodiff
// call in their code, to justify altering the compilation pipeline.
if skip_size_increasing_opts && cfg!(llvm_enzyme) {
unroll_loops = false;
vectorize_slp = false;
vectorize_loop = false;
} else {
unroll_loops =
opt_level != config::OptLevel::Size && opt_level != config::OptLevel::SizeMin;
vectorize_slp = config.vectorize_slp;
vectorize_loop = config.vectorize_loop;
}
trace!(
"Enzyme: Running with unroll_loops: {}, vectorize_slp: {}, vectorize_loop: {}",
unroll_loops, vectorize_slp, vectorize_loop
);
let using_thin_buffers = opt_stage == llvm::OptStage::PreLinkThinLTO || config.bitcode_needed();
let pgo_gen_path = get_pgo_gen_path(config);
let pgo_use_path = get_pgo_use_path(config);
Expand Down Expand Up @@ -596,8 +626,8 @@ pub(crate) unsafe fn llvm_optimize(
using_thin_buffers,
config.merge_functions,
unroll_loops,
config.vectorize_slp,
config.vectorize_loop,
vectorize_slp,
vectorize_loop,
config.no_builtins,
config.emit_lifetime_markers,
sanitizer_options.as_ref(),
Expand All @@ -619,6 +649,83 @@ pub(crate) unsafe fn llvm_optimize(
result.into_result().map_err(|()| llvm_err(dcx, LlvmError::RunLlvmPasses))
}

pub(crate) fn differentiate(
module: &ModuleCodegen<ModuleLlvm>,
cgcx: &CodegenContext<LlvmCodegenBackend>,
diff_items: Vec<AutoDiffItem>,
config: &ModuleConfig,
) -> Result<(), FatalError> {
for item in &diff_items {
trace!("{}", item);
}

let llmod = module.module_llvm.llmod();
let llcx = &module.module_llvm.llcx;
let diag_handler = cgcx.create_dcx();

// Before dumping the module, we want all the tt to become part of the module.
for item in diff_items.iter() {
let name = CString::new(item.source.clone()).unwrap();
let fn_def: Option<&llvm::Value> =
unsafe { llvm::LLVMGetNamedFunction(llmod, name.as_ptr()) };
let fn_def = match fn_def {
Some(x) => x,
None => {
return Err(llvm_err(diag_handler.handle(), LlvmError::PrepareAutoDiff {
src: item.source.clone(),
target: item.target.clone(),
error: "could not find source function".to_owned(),
}));
}
};
let target_name = CString::new(item.target.clone()).unwrap();
debug!("target name: {:?}", &target_name);
let fn_target: Option<&llvm::Value> =
unsafe { llvm::LLVMGetNamedFunction(llmod, target_name.as_ptr()) };
let fn_target = match fn_target {
Some(x) => x,
None => {
return Err(llvm_err(diag_handler.handle(), LlvmError::PrepareAutoDiff {
src: item.source.clone(),
target: item.target.clone(),
error: "could not find target function".to_owned(),
}));
}
};

crate::builder::generate_enzyme_call(llmod, llcx, fn_def, fn_target, item.attrs.clone());
}

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

if let Some(opt_level) = config.opt_level {
let opt_stage = match cgcx.lto {
Lto::Fat => llvm::OptStage::PreLinkFatLTO,
Lto::Thin | Lto::ThinLocal => llvm::OptStage::PreLinkThinLTO,
_ if cgcx.opts.cg.linker_plugin_lto.enabled() => llvm::OptStage::PreLinkThinLTO,
_ => llvm::OptStage::PreLinkNoLTO,
};
// This is our second opt call, so now we run all opts,
// to make sure we get the best performance.
let skip_size_increasing_opts = false;
trace!("running Module Optimization after differentiation");
unsafe {
llvm_optimize(
cgcx,
diag_handler.handle(),
module,
config,
opt_level,
opt_stage,
skip_size_increasing_opts,
)?
};
}
trace!("done with differentiate()");

Ok(())
}

// Unsafe due to LLVM calls.
pub(crate) unsafe fn optimize(
cgcx: &CodegenContext<LlvmCodegenBackend>,
Expand All @@ -641,14 +748,29 @@ pub(crate) unsafe fn optimize(
unsafe { llvm::LLVMWriteBitcodeToFile(llmod, out.as_ptr()) };
}

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

if let Some(opt_level) = config.opt_level {
let opt_stage = match cgcx.lto {
Lto::Fat => llvm::OptStage::PreLinkFatLTO,
Lto::Thin | Lto::ThinLocal => llvm::OptStage::PreLinkThinLTO,
_ if cgcx.opts.cg.linker_plugin_lto.enabled() => llvm::OptStage::PreLinkThinLTO,
_ => llvm::OptStage::PreLinkNoLTO,
};
return unsafe { llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage) };

// If we know that we will later run AD, then we disable vectorization and loop unrolling
let skip_size_increasing_opts = cfg!(llvm_enzyme);
return unsafe {
llvm_optimize(
cgcx,
dcx,
module,
config,
opt_level,
opt_stage,
skip_size_increasing_opts,
)
};
}
Ok(())
}
Expand Down
Loading

0 comments on commit 7dc39d6

Please sign in to comment.