fix monomorphization, extra dbg output

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -712,6 +712,9 @@ pub(crate) unsafe fn enzyme_ad(
     let name2 = CString::new(rust_name2.clone()).unwrap();
     let src_fnc = llvm::LLVMGetNamedFunction(llmod, name.as_c_str().as_ptr()).unwrap();
     let target_fnc = llvm::LLVMGetNamedFunction(llmod, name2.as_ptr()).unwrap();
+    let src_num_args = llvm::LLVMCountParams(src_fnc);
+    let target_num_args = llvm::LLVMCountParams(target_fnc);
+    assert!(src_num_args <= target_num_args);
 
     // create enzyme typetrees
     let llvm_data_layout = unsafe { llvm::LLVMGetDataLayoutStr(&*llmod) };

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -914,6 +914,8 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
     assert!(ret_activity == CDIFFE_TYPE::DFT_CONSTANT || ret_activity == CDIFFE_TYPE::DFT_OUT_DIFF);
     let input_activity: Vec<CDIFFE_TYPE> = input_activity.iter().map(|&x| cdiffe_from(x)).collect();
 
+    dbg!(&fnc);
+
     if ret_activity == CDIFFE_TYPE::DFT_DUP_ARG {
         if ret_primary_ret != true {
             dbg!("overwriting ret_primary_ret!");
@@ -931,6 +933,11 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
     // We don't support volatile / extern / (global?) values.
     // Just because I didn't had time to test them, and it seems less urgent.
     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());
+    assert!(num_fnc_args == input_activity.len() as u32);
     let kv_tmp = IntList { data: std::ptr::null_mut(), size: 0 };
 
 
@@ -942,7 +949,7 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
         KnownValues: known_values.as_mut_ptr(),
     };
 
-    EnzymeCreatePrimalAndGradient(
+    let res = EnzymeCreatePrimalAndGradient(
         logic_ref, // Logic
         std::ptr::null(),
         std::ptr::null(),
@@ -963,7 +970,9 @@ pub(crate) unsafe fn enzyme_rust_reverse_diff(
         args_uncacheable.len(), // uncacheable arguments
         std::ptr::null_mut(),   // write augmented function to this
         0,
-    )
+    );
+    dbg!(&res);
+    res
 }
 pub type GetSymbolsCallback = unsafe extern "C" fn(*mut c_void, *const c_char) -> *mut c_void;
 pub type GetSymbolsErrorCallback = unsafe extern "C" fn(*const c_char) -> *mut c_void;

compiler/rustc_monomorphize/src/partitioning.rs

@@ -247,7 +247,14 @@ where
             &mut can_be_internalized,
             export_generics,
         );
-        if visibility == Visibility::Hidden && can_be_internalized {
+        //if visibility == Visibility::Hidden && can_be_internalized {
+
+        //dbg!(&characteristic_def_id);
+        let autodiff_active = characteristic_def_id
+            .map(|x| cx.tcx.autodiff_attrs(x).is_active())
+            .unwrap_or(false);
+
+        if !autodiff_active && visibility == Visibility::Hidden && can_be_internalized {
             internalization_candidates.insert(mono_item);
         }
         let size_estimate = mono_item.size_estimate(cx.tcx);