add option to disable extra tt generation, fmt

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.

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,

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();

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<FncTree> = 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<FncTree> = 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> {

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};

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),
 

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,

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<OptimizationFuel>,
 
@@ -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(),