impl new modes for higher order ad (#106)

* impl new modes for higher order ad

* fix fwd void ret case

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -11,6 +11,21 @@ pub enum DiffMode {
     Source,
     Forward,
     Reverse,
+    ForwardFirst,
+    ReverseFirst,
+}
+
+pub fn is_rev(mode: DiffMode) -> bool {
+    match mode {
+        DiffMode::Reverse | DiffMode::ReverseFirst => true,
+        _ => false,
+    }
+}
+pub fn is_fwd(mode: DiffMode) -> bool {
+    match mode {
+        DiffMode::Forward | DiffMode::ForwardFirst => true,
+        _ => false,
+    }
 }
 
 impl Display for DiffMode {
@@ -20,6 +35,8 @@ impl Display for DiffMode {
             DiffMode::Source => write!(f, "Source"),
             DiffMode::Forward => write!(f, "Forward"),
             DiffMode::Reverse => write!(f, "Reverse"),
+            DiffMode::ForwardFirst => write!(f, "ForwardFirst"),
+            DiffMode::ReverseFirst => write!(f, "ReverseFirst"),
         }
     }
 }
@@ -32,12 +49,12 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
     match mode {
         DiffMode::Inactive => false,
         DiffMode::Source => false,
-        DiffMode::Forward => {
+        DiffMode::Forward | DiffMode::ForwardFirst => {
             activity == DiffActivity::Dual ||
                 activity == DiffActivity::DualOnly ||
                 activity == DiffActivity::Const
         }
-        DiffMode::Reverse => {
+        DiffMode::Reverse | DiffMode::ReverseFirst => {
             activity == DiffActivity::Const ||
                 activity == DiffActivity::Active ||
                 activity == DiffActivity::ActiveOnly
@@ -73,13 +90,13 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         return match mode {
             DiffMode::Inactive => false,
             DiffMode::Source => false,
-            DiffMode::Forward => {
+            DiffMode::Forward | DiffMode::ForwardFirst => {
                 // These are the only valid cases
                 activity == DiffActivity::Dual ||
                     activity == DiffActivity::DualOnly ||
                     activity == DiffActivity::Const
             }
-            DiffMode::Reverse => {
+            DiffMode::Reverse | DiffMode::ReverseFirst => {
                 // These are the only valid cases
                 activity == DiffActivity::Active ||
                     activity == DiffActivity::ActiveOnly ||
@@ -137,6 +154,8 @@ impl FromStr for DiffMode {
             "Source" => Ok(DiffMode::Source),
             "Forward" => Ok(DiffMode::Forward),
             "Reverse" => Ok(DiffMode::Reverse),
+            "ForwardFirst" => Ok(DiffMode::ForwardFirst),
+            "ReverseFirst" => Ok(DiffMode::ReverseFirst),
             _ => Err(()),
         }
     }

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -3,7 +3,7 @@
 //use crate::util::check_autodiff;
 
 use crate::errors;
-use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ty_for_activity};
+use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode, is_fwd, is_rev, valid_input_activity, valid_ty_for_activity};
 use rustc_ast::ptr::P;
 use rustc_ast::token::{Token, TokenKind};
 use rustc_ast::tokenstream::*;
@@ -308,7 +308,7 @@ fn gen_enzyme_body(
 
     let primal_ret = sig.decl.output.has_ret();
 
-    if primal_ret && n_active == 0 && x.mode == DiffMode::Reverse {
+    if primal_ret && n_active == 0 && is_rev(x.mode) {
         // We only have the primal ret.
         body.stmts.push(ecx.stmt_expr(black_box_primal_call.clone()));
         return body;
@@ -355,7 +355,7 @@ fn gen_enzyme_body(
             panic!("Did not expect non-tuple ret ty: {:?}", d_ret_ty);
         }
     };
-    if x.mode == DiffMode::Forward {
+    if is_fwd(x.mode) {
         if x.ret_activity == DiffActivity::Dual {
             assert!(d_ret_ty.len() == 2);
             // both should be identical, by construction
@@ -369,7 +369,7 @@ fn gen_enzyme_body(
             exprs.push(default_call_expr);
         }
     } else {
-        assert!(x.mode == DiffMode::Reverse);
+        assert!(is_rev(x.mode));
 
         if primal_ret {
             // We have extra handling above for the primal ret
@@ -508,7 +508,7 @@ fn gen_enzyme_decl(
 
     // If we return a scalar in the primal and the scalar is active,
     // then add it as last arg to the inputs.
-    if let DiffMode::Reverse = x.mode {
+    if is_rev(x.mode) {
         if let DiffActivity::Active = x.ret_activity {
             let ty = match d_decl.output {
                 FnRetTy::Ty(ref ty) => ty.clone(),
@@ -537,7 +537,7 @@ fn gen_enzyme_decl(
     }
     d_decl.inputs = d_inputs.into();
 
-    if let DiffMode::Forward = x.mode {
+    if is_fwd(x.mode) {
         if let DiffActivity::Dual = x.ret_activity {
             let ty = match d_decl.output {
                 FnRetTy::Ty(ref ty) => ty.clone(),

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -703,7 +703,7 @@ fn get_params(fnc: &Value) -> Vec<&Value> {
 
 unsafe fn create_call<'a>(tgt: &'a Value, src: &'a Value, rev_mode: bool,
     llmod: &'a llvm::Module, llcx: &llvm::Context, size_positions: &[usize]) {
-    dbg!("size_positions: {:?}", size_positions);
+
    // first, remove all calls from fnc
    let bb = LLVMGetFirstBasicBlock(tgt);
    let br = LLVMRustGetTerminator(bb);
@@ -843,7 +843,6 @@ unsafe fn create_call<'a>(tgt: &'a Value, src: &'a Value, rev_mode: bool,
 
     // Now clean up placeholder code.
     LLVMRustEraseInstBefore(bb, last_inst);
-    //dbg!(&tgt);
 
     let f_return_type = LLVMGetReturnType(LLVMGlobalGetValueType(src));
     let void_type = LLVMVoidTypeInContext(llcx);
@@ -865,6 +864,7 @@ unsafe fn create_call<'a>(tgt: &'a Value, src: &'a Value, rev_mode: bool,
     let _fnc_ok =
         LLVMVerifyFunction(tgt, llvm::LLVMVerifierFailureAction::LLVMAbortProcessAction);
 }
+
 unsafe fn get_panic_name(llmod: &llvm::Module) -> CString {
     // The names are mangled and their ending changes based on a hash, so just take whichever.
     let mut f = LLVMGetFirstFunction(llmod);
@@ -922,21 +922,7 @@ unsafe fn add_panic_msg_to_global<'a>(llmod: &'a llvm::Module, llcx: &'a llvm::C
     LLVMRustSetLinkage(global_var, Linkage::PrivateLinkage);
     LLVMSetInitializer(global_var, struct_initializer);
 
-        //let msg_global_name = "ad_safety_msg".to_string();
-        //let cmsg_global_name = CString::new(msg_global_name).unwrap();
-        //let msg = "autodiff safety check failed!";
-        //let cmsg = CString::new(msg).unwrap();
-        //let msg_len = msg.len();
-        //let i8_array_type = llvm::LLVMRustArrayType(llvm::LLVMInt8TypeInContext(llcx), msg_len as u64);
-        //let global_type  = llvm::LLVMStructTypeInContext(llcx, [i8_array_type].as_mut_ptr(), 1, 0);
-        //let string_const_val = llvm::LLVMConstStringInContext(llcx, cmsg.as_ptr() as *const c_char, msg_len as u32, 0);
-        //let initializer = llvm::LLVMConstStructInContext(llcx, [string_const_val].as_mut_ptr(), 1, 0);
-        //let global = llvm::LLVMAddGlobal(llmod, global_type, cmsg_global_name.as_ptr() as *const c_char);
-        //llvm::LLVMRustSetLinkage(global, llvm::Linkage::PrivateLinkage);
-        //llvm::LLVMSetInitializer(global, initializer);
-        //llvm::LLVMSetUnnamedAddress(global, llvm::UnnamedAddr::Global);
-
-        global_var
+    global_var
 }
 
 // As unsafe as it can be.
@@ -947,6 +933,7 @@ pub(crate) unsafe fn enzyme_ad(
     llcx: &llvm::Context,
     diag_handler: &DiagCtxt,
     item: AutoDiffItem,
+    logic_ref: EnzymeLogicRef,
 ) -> Result<(), FatalError> {
     let autodiff_mode = item.attrs.mode;
     let rust_name = item.source;
@@ -1001,13 +988,6 @@ pub(crate) unsafe fn enzyme_ad(
         item.inputs.into_iter().map(|x| to_enzyme_typetree(x, llvm_data_layout, llcx)).collect();
     let output_tt = to_enzyme_typetree(item.output, llvm_data_layout, llcx);
 
-    let mut fnc_opt = false;
-    if std::env::var("ENZYME_ENABLE_FNC_OPT").is_ok() {
-        dbg!("Disabling optimizations for Enzyme");
-        fnc_opt = true;
-    }
-
-    let logic_ref: EnzymeLogicRef = CreateEnzymeLogic(fnc_opt as u8);
     let type_analysis: EnzymeTypeAnalysisRef =
         CreateTypeAnalysis(logic_ref, std::ptr::null_mut(), std::ptr::null_mut(), 0);
 
@@ -1026,7 +1006,18 @@ pub(crate) unsafe fn enzyme_ad(
         llvm::set_print(true);
     }
 
-    let mut tmp = match item.attrs.mode {
+    let mode = match autodiff_mode {
+        DiffMode::Forward => DiffMode::Forward,
+        DiffMode::Reverse => DiffMode::Reverse,
+        DiffMode::ForwardFirst => DiffMode::Forward,
+        DiffMode::ReverseFirst => DiffMode::Reverse,
+        _ => unreachable!(),
+    };
+
+    let void_type = LLVMVoidTypeInContext(llcx);
+    let return_type = LLVMGetReturnType(LLVMGlobalGetValueType(src_fnc));
+    let void_ret = void_type == return_type;
+    let mut tmp = match mode {
         DiffMode::Forward => enzyme_rust_forward_diff(
             logic_ref,
             type_analysis,
@@ -1035,6 +1026,7 @@ pub(crate) unsafe fn enzyme_ad(
             ret_activity,
             input_tts,
             output_tt,
+            void_ret,
         ),
         DiffMode::Reverse => enzyme_rust_reverse_diff(
             logic_ref,
@@ -1052,7 +1044,6 @@ pub(crate) unsafe fn enzyme_ad(
 
     let f_return_type = LLVMGetReturnType(LLVMGlobalGetValueType(res));
 
-    let void_type = LLVMVoidTypeInContext(llcx);
     let rev_mode = item.attrs.mode == DiffMode::Reverse;
     create_call(target_fnc, res, rev_mode, llmod, llcx, &size_positions);
     // TODO: implement drop for wrapper type?
@@ -1114,8 +1105,40 @@ pub(crate) unsafe fn differentiate(
     }
 
     let differentiate = !diff_items.is_empty();
+    let mut first_order_items: Vec<AutoDiffItem> = vec![];
+    let mut higher_order_items: Vec<AutoDiffItem> = vec![];
     for item in diff_items {
-        let res = enzyme_ad(llmod, llcx, &diag_handler, item);
+        if item.attrs.mode == DiffMode::ForwardFirst || item.attrs.mode == DiffMode::ReverseFirst{
+            first_order_items.push(item);
+        } else {
+            // default
+            higher_order_items.push(item);
+        }
+    }
+
+    let mut fnc_opt = false;
+    if std::env::var("ENZYME_ENABLE_FNC_OPT").is_ok() {
+        dbg!("Enable extra optimizations for Enzyme");
+        fnc_opt = true;
+    }
+
+    // If a function is a base for some higher order ad, always optimize
+    let fnc_opt_base = true;
+    let logic_ref_opt: EnzymeLogicRef = CreateEnzymeLogic(fnc_opt_base as u8);
+
+    for item in first_order_items {
+        let res = enzyme_ad(llmod, llcx, &diag_handler, item, logic_ref_opt);
+        assert!(res.is_ok());
+    }
+
+    // For the rest, follow the user choice on debug vs release.
+    // Reuse the opt one if possible for better compile time (Enzyme internal caching).
+    let logic_ref = match fnc_opt {
+        true => logic_ref_opt,
+        false => CreateEnzymeLogic(fnc_opt as u8),
+    };
+    for item in higher_order_items {
+        let res = enzyme_ad(llmod, llcx, &diag_handler, item, logic_ref);
         assert!(res.is_ok());
     }
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -850,6 +850,7 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
     ret_diffactivity: DiffActivity,
     input_tts: Vec<TypeTree>,
     output_tt: TypeTree,
+    void_ret: bool,
 ) -> (&Value, Vec<usize>) {
     let ret_activity = cdiffe_from(ret_diffactivity);
     assert!(ret_activity != CDIFFE_TYPE::DFT_OUT_DIFF);
@@ -864,12 +865,18 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
         input_activity.push(act);
     }
 
-    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."),
+    // if we have void ret, this must be false;
+    let ret_primary_ret = if void_ret {
+        false
+    } else {
+        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."),
+        }
     };
+    trace!("ret_primary_ret: {}", &ret_primary_ret);
 
     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()];
@@ -879,8 +886,8 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
     let args_uncacheable = vec![0; input_tts.len()];
     assert!(args_uncacheable.len() == input_activity.len());
     let num_fnc_args = LLVMCountParams(fnc);
-    println!("num_fnc_args: {}", num_fnc_args);
-    println!("input_activity.len(): {}", input_activity.len());
+    trace!("num_fnc_args: {}", num_fnc_args);
+    trace!("input_activity.len(): {}", input_activity.len());
     assert!(num_fnc_args == input_activity.len() as u32);
 
     let kv_tmp = IntList { data: std::ptr::null_mut(), size: 0 };
@@ -893,6 +900,11 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
         KnownValues: known_values.as_mut_ptr(),
     };
 
+    trace!("ret_activity: {}", &ret_activity);
+    for i in &input_activity {
+        trace!("input_activity i: {}", &i);
+    }
+    trace!("before calling Enzyme");
     let res = EnzymeCreateForwardDiff(
         logic_ref, // Logic
         std::ptr::null(),
@@ -912,6 +924,7 @@ pub(crate) unsafe fn enzyme_rust_forward_diff(
         args_uncacheable.len(), // uncacheable arguments
         std::ptr::null_mut(),   // write augmented function to this
     );
+    trace!("after calling Enzyme");
     (res, vec![])
 }
 
@@ -971,11 +984,12 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
         KnownValues: known_values.as_mut_ptr(),
     };
 
-    trace!("{}", &primary_ret);
-    trace!("{}", &ret_activity);
+    trace!("primary_ret: {}", &primary_ret);
+    trace!("ret_activity: {}", &ret_activity);
     for i in &input_activity {
-        trace!("{}", &i);
+        trace!("input_activity i: {}", &i);
     }
+    trace!("before calling Enzyme");
     let res = EnzymeCreatePrimalAndGradient(
         logic_ref, // Logic
         std::ptr::null(),
@@ -998,6 +1012,7 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
         std::ptr::null_mut(),   // write augmented function to this
         0,
     );
+    trace!("after calling Enzyme");
     (res, primal_sizes)
 }
 

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -749,6 +749,8 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
     let mode = match mode.as_str() {
         "Forward" => DiffMode::Forward,
         "Reverse" => DiffMode::Reverse,
+        "ForwardFirst" => DiffMode::ForwardFirst,
+        "ReverseFirst" => DiffMode::ReverseFirst,
         _ => {
             tcx.sess
                 .struct_span_err(attr.span, msg_mode)