adressing reviewer feedback

Author: ZuseZ4

compiler/rustc_codegen_gcc/src/lib.rs

@@ -93,6 +93,7 @@ use gccjit::{CType, Context, OptimizationLevel};
 #[cfg(feature = "master")]
 use gccjit::{TargetInfo, Version};
 use rustc_ast::expand::allocator::AllocatorKind;
+use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
 use rustc_codegen_ssa::back::write::{
     CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
@@ -439,6 +440,14 @@ impl WriteBackendMethods for GccCodegenBackend {
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
         back::write::link(cgcx, dcx, modules)
     }
+    fn autodiff(
+        _cgcx: &CodegenContext<Self>,
+        _module: &ModuleCodegen<Self::Module>,
+        _diff_fncs: Vec<AutoDiffItem>,
+        _config: &ModuleConfig,
+    ) -> Result<(), FatalError> {
+        unimplemented!()
+    }
 }
 
 /// This is the entrypoint for a hot plugged rustc_codegen_gccjit

compiler/rustc_codegen_llvm/messages.ftl

@@ -1,3 +1,5 @@
+codegen_llvm_autodiff_without_lto = using the autodiff feature requires using fat-lto
+
 codegen_llvm_copy_bitcode = failed to copy bitcode to object file: {$err}
 
 codegen_llvm_dynamic_linking_with_lto =
@@ -47,6 +49,8 @@ codegen_llvm_parse_bitcode_with_llvm_err = failed to parse bitcode for LTO modul
 codegen_llvm_parse_target_machine_config =
     failed to parse target machine config to target machine: {$error}
 
+codegen_llvm_prepare_autodiff = failed to prepare autodiff: src: {$src}, target: {$target}, {$error}
+codegen_llvm_prepare_autodiff_with_llvm_err = failed to prepare autodiff: {$llvm_err}, src: {$src}, target: {$target}, {$error}
 codegen_llvm_prepare_thin_lto_context = failed to prepare thin LTO context
 codegen_llvm_prepare_thin_lto_context_with_llvm_err = failed to prepare thin LTO context: {$llvm_err}
 
@@ -56,10 +60,6 @@ codegen_llvm_prepare_thin_lto_module_with_llvm_err = failed to prepare thin LTO
 codegen_llvm_run_passes = failed to run LLVM passes
 codegen_llvm_run_passes_with_llvm_err = failed to run LLVM passes: {$llvm_err}
 
-codegen_llvm_prepare_autodiff = failed to prepare autodiff: src: {$src}, target: {$target}, {$error}
-codegen_llvm_prepare_autodiff_with_llvm_err = failed to prepare autodiff: {$llvm_err}, src: {$src}, target: {$target}, {$error}
-codegen_llvm_autodiff_without_lto = using the autodiff feature requires using fat-lto
-
 codegen_llvm_sanitizer_memtag_requires_mte =
     `-Zsanitizer=memtag` requires `-Ctarget-feature=+mte`
 

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -525,7 +525,7 @@ pub(crate) unsafe fn llvm_optimize(
     // source code. However, benchmarks show that optimizations increasing the code size
     // tend to reduce AD performance. Therefore deactivate them before AD, then differentiate the code
     // and finally re-optimize the module, now with all optimizations available.
-    // TODO: In a future update we could figure out how to only optimize individual functions getting
+    // FIXME(ZuseZ4): In a future update we could figure out how to only optimize individual functions getting
     // differentiated.
 
     let unroll_loops;
@@ -683,8 +683,7 @@ pub(crate) fn differentiate(
         crate::builder::generate_enzyme_call(llmod, llcx, fn_def, fn_target, item.attrs.clone());
     }
 
-    // FIXME(ZuseZ4): In the following upstream PR, we want to add code to handle SanitizeHWAddress,
-    // to prevent some illegal/unsupported optimizations.
+    // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
 
     if let Some(opt_level) = config.opt_level {
         let opt_stage = match cgcx.lto {
@@ -736,8 +735,7 @@ pub(crate) unsafe fn optimize(
         unsafe { llvm::LLVMWriteBitcodeToFile(llmod, out.as_ptr()) };
     }
 
-    // FIXME(ZuseZ4): In the following PR, we have to add code to apply the sanitize_hwaddress
-    // attribute to all functions in the module, to prevent some illegal/unsupported optimizations.
+    // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
 
     if let Some(opt_level) = config.opt_level {
         let opt_stage = match cgcx.lto {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -48,16 +48,15 @@ fn get_params(fnc: &Value) -> Vec<&Value> {
     }
 }
 
-// The lowering of one `#[autodiff]` macro happens in multiple steps.
-// First we transalte generate a new dummy function, who's llvm-ir we now have as outer_fn.
-// We kept track of the original function to which the `#[autodiff]` macro was applied to, which we
-// now have as fn_to_diff. In our current implementation, we use the enzyme pass to carry out the
-// differentiation, following naming and calling conventions documented here: <https://enzyme.mit.edu/getting_started/CallingConvention/>
-//
-// Our `outer_fn` had some dummy code inserted at higher levels, so we first remove most of the
-// existing body. We then insert an `__enzyme_<autodiff/fwddiff>_<unique_id>` call, which the pass
-// will then pick up. FIXME(ZuseZ4): We will later want to upstream safety checks to the `outer_fn`,
-// in order to cover some assumptions of enzyme/autodiff, which could lead to UB otherwise.
+/// When differentiating `fn_to_diff`, take a `outer_fn` and generate another
+/// function with expected naming and calling conventions[^1] which will be
+/// discovered by the enzyme LLVM pass and its body populated with the differentiated
+/// `fn_to_diff`. `outer_fn` is then modified to have a call to the generated
+/// function and handle the differences between the Rust calling convention and
+/// Enzyme.
+/// [^1]: <https://enzyme.mit.edu/getting_started/CallingConvention/>
+// FIXME(ZuseZ4): `outer_fn` should include upstream safety checks to
+// cover some assumptions of enzyme/autodiff, which could lead to UB otherwise.
 pub(crate) fn generate_enzyme_call<'ll>(
     llmod: &'ll llvm::Module,
     llcx: &'ll llvm::Context,
@@ -69,7 +68,7 @@ pub(crate) fn generate_enzyme_call<'ll>(
     let output = attrs.ret_activity;
 
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
-    // FIXME(ZuseZ4): The new pass based approach should not need the *First method anymore, since
+    // FIXME(ZuseZ4): The new pass based approach should not need the {Forward/Reverse}First method anymore, since
     // it will handle higher-order derivatives correctly automatically (in theory). Currently
     // higher-order derivatives fail, so we should debug that before adjusting this code.
     let mut ad_name: String = match attrs.mode {
@@ -87,16 +86,38 @@ pub(crate) fn generate_enzyme_call<'ll>(
     let outer_fn_name = std::ffi::CStr::from_bytes_with_nul(name).unwrap().to_str().unwrap();
     ad_name.push_str(outer_fn_name.to_string().as_str());
 
-    // Assuming that our fn_to_diff is the fnc square, want to generate the following llvm-ir, which
-    // would allow the enzyme pass to generate a function body for `__enzyme_autodiff_square`
+    // Let us assume the user wrote the following function square:
     //
+    // ```llvm
+    // define double @square(double %x) {
+    // entry:
+    //  %0 = fmul double %x, %x
+    //  ret double %0
+    // }
+    // ```
+    //
+    // The user now applies autodiff to the function square, in which case fn_to_diff will be `square`.
+    // Our macro generates the following placeholder code (slightly simplified):
+    //
+    // ```llvm
+    // define double @dsquare(double %x) {
+    //  ; placeholder code
+    //  return 0.0;
+    // }
+    // ```
+    //
+    // so our `outer_fn` will be `dsquare`. The unsafe code section below now removes the placeholder
+    // code and inserts an autodiff call. We also add a declaration for the __enzyme_autodiff call.
+    // Again, the arguments to all functions are slightly simplified.
+    // ```llvm
     // declare double @__enzyme_autodiff_square(...)
     //
     // define double @dsquare(double %x) {
     // entry:
     //   %0 = tail call double (...) @__enzyme_autodiff_square(double (double)* nonnull @square, double %x)
     //   ret double %0
     // }
+    // ```
     unsafe {
         // On LLVM-IR, we can luckily declare __enzyme_ functions without specifying the input
         // arguments. We do however need to declare them with their correct return type.