fix reverse_mode setting, various cleanups

Author: ZuseZ4

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -91,7 +91,7 @@ pub fn expand(
         new_decl_span,
         idents,
     );
-    let d_ident = meta_item_vec[0].meta_item().unwrap().path.segments[0].ident;
+    let d_ident = first_ident(&meta_item_vec[0]);
 
     // The first element of it is the name of the function to be generated
     let asdf = ItemKind::Fn(Box::new(ast::Fn {
@@ -102,11 +102,12 @@ pub fn expand(
     }));
     let mut rustc_ad_attr =
         P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::rustc_autodiff)));
-    let attr: ast::Attribute = ast::Attribute {
+    let mut attr: ast::Attribute = ast::Attribute {
         kind: ast::AttrKind::Normal(rustc_ad_attr.clone()),
-        id: ast::AttrId::from_u32(0),
+        //id: ast::DUMMY_TR_ID,
+        id: ast::AttrId::from_u32(12341), // TODO: fix
         style: ast::AttrStyle::Outer,
-        span: span,
+        span,
     };
     orig_item.attrs.push(attr.clone());
 
@@ -116,21 +117,15 @@ pub fn expand(
         delim: rustc_ast::token::Delimiter::Parenthesis,
         tokens: ts,
     });
-    let attr2: ast::Attribute = ast::Attribute {
-        kind: ast::AttrKind::Normal(rustc_ad_attr),
-        id: ast::AttrId::from_u32(0),
-        style: ast::AttrStyle::Outer,
-        span: span,
-    };
-    let attr_vec: rustc_ast::AttrVec = thin_vec![attr2];
-    let d_fn = ecx.item(span, d_ident, attr_vec, asdf);
+    attr.kind = ast::AttrKind::Normal(rustc_ad_attr);
+    let d_fn = ecx.item(span, d_ident, thin_vec![attr], asdf);
 
-    let orig_annotatable = Annotatable::Item(orig_item.clone());
+    let orig_annotatable = Annotatable::Item(orig_item);
     let d_annotatable = Annotatable::Item(d_fn);
     return vec![orig_annotatable, d_annotatable];
 }
 
-// shadow arguments must be mutable references or ptrs, because Enzyme will write into them.
+// shadow arguments in reverse mode must be mutable references or ptrs, because Enzyme will write into them.
 #[cfg(llvm_enzyme)]
 fn assure_mut_ref(ty: &ast::Ty) -> ast::Ty {
     let mut ty = ty.clone();
@@ -165,6 +160,25 @@ fn gen_enzyme_body(
 ) -> P<ast::Block> {
     let blackbox_path = ecx.std_path(&[Symbol::intern("hint"), Symbol::intern("black_box")]);
     let empty_loop_block = ecx.block(span, ThinVec::new());
+    let noop = ast::InlineAsm {
+        template: vec![ast::InlineAsmTemplatePiece::String("NOP".to_string())],
+        template_strs: Box::new([]),
+        operands: vec![],
+        clobber_abis: vec![],
+        options: ast::InlineAsmOptions::PURE & ast::InlineAsmOptions::NOMEM,
+        line_spans: vec![],
+    };
+    let noop_expr = ecx.expr_asm(span, P(noop));
+    let unsf = ast::BlockCheckMode::Unsafe(ast::UnsafeSource::CompilerGenerated);
+    let unsf_block = ast::Block {
+        stmts: thin_vec![ecx.stmt_semi(noop_expr)],
+        id: ast::DUMMY_NODE_ID,
+        tokens: None,
+        rules: unsf,
+        span,
+        could_be_bare_literal: false,
+    };
+    let unsf_expr = ecx.expr_block(P(unsf_block));
     let loop_expr = ecx.expr_loop(span, empty_loop_block);
     let blackbox_call_expr = ecx.expr_path(ecx.path(span, blackbox_path));
     let primal_call = gen_primal_call(ecx, span, primal, idents);
@@ -185,7 +199,7 @@ fn gen_enzyme_body(
     );
 
     let mut body = ecx.block(span, ThinVec::new());
-    body.stmts.push(ecx.stmt_semi(primal_call));
+    body.stmts.push(ecx.stmt_semi(unsf_expr));
     body.stmts.push(ecx.stmt_semi(black_box_primal_call));
     body.stmts.push(ecx.stmt_semi(black_box_remaining_args));
     body.stmts.push(ecx.stmt_expr(loop_expr));
@@ -234,15 +248,18 @@ fn gen_enzyme_decl(
             }
             DiffActivity::Duplicated | DiffActivity::Dual => {
                 let mut shadow_arg = arg.clone();
-                shadow_arg.ty = P(assure_mut_ref(&arg.ty));
+                // We += into the shadow in reverse mode.
+                // Otherwise copy mutability of the original argument.
+                if activity == &DiffActivity::Duplicated {
+                    shadow_arg.ty = P(assure_mut_ref(&arg.ty));
+                }
                 // adjust name depending on mode
                 let old_name = if let PatKind::Ident(_, ident, _) = arg.pat.kind {
                     ident.name
                 } else {
                     dbg!(&shadow_arg.pat);
                     panic!("not an ident?");
                 };
-                //old_names.push(old_name.to_string());
                 let name: String = match x.mode {
                     DiffMode::Reverse => format!("d{}", old_name),
                     DiffMode::Forward => format!("b{}", old_name),

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -1,5 +1,6 @@
 //! Set and unset common attributes on LLVM values.
 
+use rustc_ast::expand::autodiff_attrs::AutoDiffAttrs;
 use rustc_codegen_ssa::traits::*;
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
@@ -294,7 +295,7 @@ pub fn from_fn_attrs<'ll, 'tcx>(
     instance: ty::Instance<'tcx>,
 ) {
     let codegen_fn_attrs = cx.tcx.codegen_fn_attrs(instance.def_id());
-    let autodiff_attrs = cx.tcx.autodiff_attrs(instance.def_id());
+    let autodiff_attrs: &AutoDiffAttrs = cx.tcx.autodiff_attrs(instance.def_id());
 
     let mut to_add = SmallVec::<[_; 16]>::new();
 
@@ -313,6 +314,8 @@ pub fn from_fn_attrs<'ll, 'tcx>(
         if codegen_fn_attrs.inline == InlineAttr::None && instance.def.requires_inline(cx.tcx) {
             InlineAttr::Hint
         } else if autodiff_attrs.is_active() {
+            dbg!("autodiff_attrs.is_active()");
+            dbg!(&autodiff_attrs);
             InlineAttr::Never
         } else {
             codegen_fn_attrs.inline

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -702,6 +702,7 @@ pub(crate) unsafe fn enzyme_ad(
     diag_handler: &DiagCtxt,
     item: AutoDiffItem,
 ) -> Result<(), FatalError> {
+    dbg!("\n\n\n\n\n\n AUTO DIFF \n");
     let autodiff_mode = item.attrs.mode;
     let rust_name = item.source;
     let rust_name2 = &item.target;

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -852,7 +852,6 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
     input_tts: Vec<TypeTree>,
     output_tt: TypeTree,
 ) -> &Value {
-    let mut ret_primary_ret = false;
     let ret_activity = cdiffe_from(ret_diffactivity);
     assert!(ret_activity != CDIFFE_TYPE::DFT_OUT_DIFF);
     let mut input_activity: Vec<CDIFFE_TYPE> = vec![];
@@ -866,17 +865,12 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
         input_activity.push(act);
     }
 
-    if ret_activity == CDIFFE_TYPE::DFT_DUP_ARG {
-        if ret_primary_ret != true {
-            dbg!("overwriting ret_primary_ret!");
-        }
-        ret_primary_ret = true;
-    } else if ret_activity == CDIFFE_TYPE::DFT_DUP_NONEED {
-        if ret_primary_ret != false {
-            dbg!("overwriting ret_primary_ret!");
-        }
-        ret_primary_ret = false;
-    }
+    let ret_primary_ret = match ret_activity {
+        CDIFFE_TYPE::DFT_CONSTANT => true,
+        CDIFFE_TYPE::DFT_DUP_ARG => true,
+        CDIFFE_TYPE::DFT_DUP_NONEED => false,
+        _ => panic!("Implementation error in enzyme_rust_forward_diff."),
+    };
 
     let mut args_tree = input_tts.iter().map(|x| x.inner).collect::<Vec<_>>();
     //let mut args_tree = vec![TypeTree::new().inner; typetree.input_tt.len()];
@@ -916,7 +910,6 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
     )
 }
 
-#[allow(dead_code)]
 pub(crate) unsafe fn enzyme_rust_reverse_diff(
     logic_ref: EnzymeLogicRef,
     type_analysis: EnzymeTypeAnalysisRef,
@@ -928,7 +921,7 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
 ) -> &Value {
     let (primary_ret, ret_activity) = match ret_activity {
         DiffActivity::Const => (true, CDIFFE_TYPE::DFT_CONSTANT),
-        DiffActivity::Active => (true, CDIFFE_TYPE::DFT_DUP_ARG),
+        DiffActivity::Active => (true, CDIFFE_TYPE::DFT_OUT_DIFF),
         DiffActivity::None => (false, CDIFFE_TYPE::DFT_CONSTANT),
         _ => panic!("Invalid return activity"),
     };
@@ -962,6 +955,9 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
         KnownValues: known_values.as_mut_ptr(),
     };
 
+    dbg!(&primary_ret);
+    dbg!(&ret_activity);
+    dbg!(&input_activity);
     let res = EnzymeCreatePrimalAndGradient(
         logic_ref, // Logic
         std::ptr::null(),

compiler/rustc_codegen_ssa/src/back/lto.rs

@@ -90,7 +90,7 @@ impl<B: WriteBackendMethods> LtoModuleCodegen<B> {
             LtoModuleCodegen::Fat { ref module, .. } => {
                 B::autodiff(cgcx, &module, diff_fncs, typetrees, config)?;
             }
-            _ => {}
+            _ => panic!("autodiff called with non-fat LTO module"),
         }
 
         Ok(self)

compiler/rustc_codegen_ssa/src/back/write.rs

@@ -382,7 +382,9 @@ fn generate_lto_work<B: ExtraBackendMethods>(
     import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
 ) -> Vec<(WorkItem<B>, u64)> {
     let _prof_timer = cgcx.prof.generic_activity("codegen_generate_lto_work");
-    dbg!("Differentiating {} functions", autodiff.len());
+    //let error_msg = format!("Found {} Functions, but {} TypeTrees", autodiff.len(), typetrees.len());
+    // Don't assert yet, bc. apparently we add them later.
+    //assert!(autodiff.len() == typetrees.len(), "{}", error_msg);
 
     if !needs_fat_lto.is_empty() {
         assert!(needs_thin_lto.is_empty());

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -701,33 +701,30 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
     let attrs = tcx.get_attrs(id, sym::rustc_autodiff);
 
     let attrs = attrs
-        .into_iter()
         .filter(|attr| attr.name_or_empty() == sym::rustc_autodiff)
         .collect::<Vec<_>>();
 
-    if !attrs.is_empty() {
-        dbg!("autodiff_attrs amount = {}", attrs.len());
-    }
-
     // check for exactly one autodiff attribute on extern block
-    let msg_once = "autodiff attribute can only be applied once";
-    let attr = match &attrs[..] {
-        &[] => return AutoDiffAttrs::inactive(),
-        &[elm] => elm,
-        x => {
+    let msg_once = "cg_ssa: autodiff attribute can only be applied once";
+    let attr = match attrs.len() {
+        0 => return AutoDiffAttrs::inactive(),
+        1 => attrs.get(0).unwrap(),
+        _ => {
             tcx.sess
-                .struct_span_err(x[1].span, msg_once)
-                .span_label(x[1].span, "more than one")
+                .struct_span_err(attrs[1].span, msg_once)
+                .span_label(attrs[1].span, "more than one")
                 .emit();
-
             return AutoDiffAttrs::inactive();
         }
     };
+    dbg!("autodiff_attr = {:?}", &attr);
 
     let list = attr.meta_item_list().unwrap_or_default();
+    dbg!("autodiff_attrs list = {:?}", &list);
 
     // empty autodiff attribute macros (i.e. `#[autodiff]`) are used to mark source functions
     if list.len() == 0 {
+        dbg!("autodiff_attrs: source");
         return AutoDiffAttrs::source();
     }
 

compiler/rustc_expand/src/build.rs

@@ -412,6 +412,9 @@ impl<'a> ExtCtxt<'a> {
     pub fn expr_loop(&self, sp: Span, block: P<ast::Block>) -> P<ast::Expr> {
         self.expr(sp, ast::ExprKind::Loop(block, None, sp))
     }
+    pub fn expr_asm(&self, sp: Span, expr: P<ast::InlineAsm>) -> P<ast::Expr> {
+        self.expr(sp, ast::ExprKind::InlineAsm(expr))
+    }
 
     pub fn expr_fail(&self, span: Span, msg: Symbol) -> P<ast::Expr> {
         self.expr_call_global(

compiler/rustc_monomorphize/src/partitioning.rs

@@ -255,9 +255,11 @@ where
             export_generics,
         );
 
-        //if visibility == Visibility::Hidden && can_be_internalized {
-        let autodiff_active =
-            characteristic_def_id.map(|x| cx.tcx.autodiff_attrs(x).is_active()).unwrap_or(false);
+        // We can't differentiate something that got inlined.
+        let autodiff_active = match characteristic_def_id {
+            Some(def_id) => cx.tcx.autodiff_attrs(def_id).is_active(),
+            None => false,
+        };
 
         if !autodiff_active && visibility == Visibility::Hidden && can_be_internalized {
             internalization_candidates.insert(mono_item);
@@ -1166,6 +1168,10 @@ fn collect_and_partition_mono_items(
         .filter_map(|(item, instance)| {
             let target_id = instance.def_id();
             let target_attrs = tcx.autodiff_attrs(target_id);
+            if target_attrs.is_source() {
+                dbg!("source");
+                dbg!(&target_attrs);
+            }
             if !target_attrs.apply_autodiff() {
                 return None;
             }