From 1203575c627a2e96bc3543145526b5091a58be65 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Wed, 30 Oct 2024 18:28:03 -0400 Subject: [PATCH] add option to disable extra tt generation, fmt --- compiler/rustc_codegen_llvm/src/back/write.rs | 4 ++- compiler/rustc_codegen_llvm/src/builder.rs | 3 +- .../rustc_codegen_ssa/src/codegen_attrs.rs | 6 ++-- .../rustc_codegen_ssa/src/mir/intrinsic.rs | 30 ++++++++++++++----- .../rustc_monomorphize/src/partitioning.rs | 2 +- compiler/rustc_session/src/config.rs | 1 + compiler/rustc_session/src/options.rs | 1 + compiler/rustc_session/src/session.rs | 7 +++++ 8 files changed, 39 insertions(+), 15 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs index 3f7eec3531def..583ba426cb130 100644 --- a/compiler/rustc_codegen_llvm/src/back/write.rs +++ b/compiler/rustc_codegen_llvm/src/back/write.rs @@ -1113,7 +1113,9 @@ pub(crate) unsafe fn differentiate( let name = CString::new(item.source.clone()).unwrap(); let fn_def: &llvm::Value = unsafe { llvm::LLVMGetNamedFunction(llmod, name.as_ptr()).unwrap() }; - crate::builder::add_tt2(llmod, llcx, fn_def, tt); + if !ad.contains(&AutoDiff::NoTypeTrees) { + crate::builder::add_tt2(llmod, llcx, fn_def, tt); + } // Before dumping the module, we also might want to add dummy functions, which will // trigger the LLVMEnzyme pass to run on them, if we invoke the opt binary. diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index f640ee2ba357f..936db15be03d4 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/compiler/rustc_codegen_llvm/src/builder.rs @@ -32,12 +32,11 @@ use crate::abi::FnAbiLlvmExt; use crate::attributes; use crate::common::Funclet; use crate::context::CodegenCx; -use crate::llvm::{self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, True}; +use crate::llvm::{self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, Metadata, True}; use crate::type_::Type; use crate::type_of::LayoutLlvmExt; use crate::typetree::to_enzyme_typetree; use crate::value::Value; -use crate::llvm::Metadata; pub(crate) fn add_tt2<'ll>( llmod: &'ll llvm::Module, diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs index 79a208925891e..3edaeb9fee992 100644 --- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs +++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs @@ -1,9 +1,9 @@ use std::str::FromStr; use rustc_ast::expand::autodiff_attrs::{ - AutoDiffAttrs, DiffActivity, DiffMode, valid_ret_activity, valid_input_activity, + AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity, }; -use rustc_ast::{MetaItemInner, MetaItem, MetaItemKind, ast, attr}; +use rustc_ast::{MetaItem, MetaItemInner, MetaItemKind, ast, attr}; use rustc_attr::{InlineAttr, InstructionSetAttr, OptimizeAttr, list_contains_name}; use rustc_data_structures::fx::FxHashMap; use rustc_errors::codes::*; @@ -807,7 +807,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs { attrs.get(0).unwrap() //tcx.dcx().struct_span_err(attrs[1].span, msg_once).with_note("more than one").emit(); //return AutoDiffAttrs::error(); - }, + } }; let list = attr.meta_item_list().unwrap_or_default(); diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs index 6906bfca571a9..f6b1c88f3692e 100644 --- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs @@ -2,7 +2,7 @@ use rustc_ast::expand::typetree::{FncTree, TypeTree}; use rustc_middle::ty::layout::HasTyCtxt; use rustc_middle::ty::{self, Ty, TyCtxt, typetree_from}; use rustc_middle::{bug, span_bug}; -use rustc_session::config::OptLevel; +use rustc_session::config::{AutoDiff, OptLevel}; use rustc_span::{Span, sym}; use rustc_target::abi::WrappingRange; use rustc_target::abi::call::{FnAbi, PassMode}; @@ -26,8 +26,15 @@ fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>( ) { let tcx: TyCtxt<'_> = bx.cx().tcx(); let tt: TypeTree = typetree_from(tcx, ty); - let fnc_tree: FncTree = - FncTree { args: vec![tt.clone(), tt.clone(), TypeTree::all_ints()], ret: TypeTree::new() }; + let ad = &tcx.sess.opts.unstable_opts.autodiff; + let fnc_tree: Option = if ad.contains(&AutoDiff::NoTypeTrees) { + None + } else { + Some(FncTree { + args: vec![tt.clone(), tt.clone(), TypeTree::all_ints()], + ret: TypeTree::new(), + }) + }; let layout = bx.layout_of(ty); let size = layout.size; @@ -36,9 +43,9 @@ fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>( let flags = if volatile { MemFlags::VOLATILE } else { MemFlags::empty() }; trace!("copy: mir ty: {:?}, enzyme tt: {:?}", ty, fnc_tree); if allow_overlap { - bx.memmove(dst, align, src, align, size, flags, Some(fnc_tree)); + bx.memmove(dst, align, src, align, size, flags, fnc_tree); } else { - bx.memcpy(dst, align, src, align, size, flags, Some(fnc_tree)); + bx.memcpy(dst, align, src, align, size, flags, fnc_tree); } } @@ -52,15 +59,22 @@ fn memset_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>( ) { let tcx: TyCtxt<'_> = bx.cx().tcx(); let tt: TypeTree = typetree_from(tcx, ty); - let fnc_tree: FncTree = - FncTree { args: vec![tt.clone(), tt.clone(), TypeTree::all_ints()], ret: TypeTree::new() }; + let ad = &tcx.sess.opts.unstable_opts.autodiff; + let fnc_tree: Option = if ad.contains(&AutoDiff::NoTypeTrees) { + None + } else { + Some(FncTree { + args: vec![tt.clone(), tt.clone(), TypeTree::all_ints()], + ret: TypeTree::new(), + }) + }; let layout = bx.layout_of(ty); let size = layout.size; let align = layout.align.abi; let size = bx.mul(bx.const_usize(size.bytes()), count); let flags = if volatile { MemFlags::VOLATILE } else { MemFlags::empty() }; - bx.memset(dst, val, size, align, flags, Some(fnc_tree)); + bx.memset(dst, val, size, align, flags, fnc_tree); } impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs index 238a9b873608c..2ce6dba4f8681 100644 --- a/compiler/rustc_monomorphize/src/partitioning.rs +++ b/compiler/rustc_monomorphize/src/partitioning.rs @@ -120,8 +120,8 @@ use rustc_middle::util::Providers; use rustc_session::CodegenUnits; use rustc_session::config::{DumpMonoStatsFormat, SwitchWithOptPath}; use rustc_span::symbol::Symbol; -use rustc_target::spec::SymbolVisibility; use rustc_symbol_mangling::symbol_name_for_instance_in_crate; +use rustc_target::spec::SymbolVisibility; use tracing::{debug, trace}; use crate::collector::{self, MonoItemCollectionStrategy, UsageMap}; diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs index 3c93a7eda36f1..18cd42f6bdbcd 100644 --- a/compiler/rustc_session/src/config.rs +++ b/compiler/rustc_session/src/config.rs @@ -221,6 +221,7 @@ pub enum AutoDiff { /// TypeTree options /// TODO: Figure out how to let users construct these, /// or whether we want to leave this option in the first place. + NoTypeTrees, TTWidth(u64), TTDepth(u64), diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs index 81381e7e9855c..0eca31f97efab 100644 --- a/compiler/rustc_session/src/options.rs +++ b/compiler/rustc_session/src/options.rs @@ -1014,6 +1014,7 @@ mod parse { "PrintModAfterOpts" => AutoDiff::PrintModAfterOpts, "PrintModAfterEnzyme" => AutoDiff::PrintModAfterEnzyme, "LooseTypes" => AutoDiff::LooseTypes, + "NoTypeTree" => AutoDiff::NoTypeTrees, "OPT" => AutoDiff::OPT, "NoModOptAfter" => AutoDiff::NoModOptAfter, "EnableFncOpt" => AutoDiff::EnableFncOpt, diff --git a/compiler/rustc_session/src/session.rs b/compiler/rustc_session/src/session.rs index 1963cf4eb7c06..d907f29470aa0 100644 --- a/compiler/rustc_session/src/session.rs +++ b/compiler/rustc_session/src/session.rs @@ -160,6 +160,9 @@ pub struct Session { /// Data about code being compiled, gathered during compilation. pub code_stats: CodeStats, + /// Tracks if `-Z autodiff="NoTypeTrees"` is specified. + pub emit_type_trees: bool, + /// Tracks fuel info if `-zfuel=crate=n` is specified. optimization_fuel: Lock, @@ -1090,6 +1093,9 @@ pub fn build_session( }); let print_fuel = AtomicU64::new(0); + let emit_type_trees = + !sopts.unstable_opts.autodiff.contains(&crate::config::AutoDiff::NoTypeTrees); + let prof = SelfProfilerRef::new( self_profiler, sopts.unstable_opts.time_passes.then(|| sopts.unstable_opts.time_passes_format), @@ -1115,6 +1121,7 @@ pub fn build_session( incr_comp_session: RwLock::new(IncrCompSession::NotInitialized), prof, code_stats: Default::default(), + emit_type_trees, optimization_fuel, print_fuel, jobserver: jobserver::client(),