Remove inlining for autodiff handling

Author: Sa4dUs

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -193,7 +193,6 @@ mod llvm_enzyme {
     /// which becomes expanded to:
     /// ```
     /// #[rustc_autodiff]
-    /// #[inline(never)]
     /// fn sin(x: &Box<f32>) -> f32 {
     ///     f32::sin(**x)
     /// }
@@ -369,7 +368,7 @@ mod llvm_enzyme {
         let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
         let inline_never = outer_normal_attr(&inline_never_attr, new_id, span);
 
-        // We're avoid duplicating the attributes `#[rustc_autodiff]` and `#[inline(never)]`.
+        // We're avoid duplicating the attribute `#[rustc_autodiff]`.
         fn same_attribute(attr: &ast::AttrKind, item: &ast::AttrKind) -> bool {
             match (attr, item) {
                 (ast::AttrKind::Normal(a), ast::AttrKind::Normal(b)) => {
@@ -382,23 +381,25 @@ mod llvm_enzyme {
             }
         }
 
+        let mut has_inline_never = false;
+
         // Don't add it multiple times:
         let orig_annotatable: Annotatable = match item {
             Annotatable::Item(ref mut iitem) => {
                 if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
                     iitem.attrs.push(attr);
                 }
-                if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
-                    iitem.attrs.push(inline_never.clone());
+                if iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
+                    has_inline_never = true;
                 }
                 Annotatable::Item(iitem.clone())
             }
             Annotatable::AssocItem(ref mut assoc_item, i @ Impl { .. }) => {
                 if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
                     assoc_item.attrs.push(attr);
                 }
-                if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
-                    assoc_item.attrs.push(inline_never.clone());
+                if assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
+                    has_inline_never = true;
                 }
                 Annotatable::AssocItem(assoc_item.clone(), i)
             }
@@ -408,9 +409,8 @@ mod llvm_enzyme {
                         if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
                             iitem.attrs.push(attr);
                         }
-                        if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind))
-                        {
-                            iitem.attrs.push(inline_never.clone());
+                        if iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
+                            has_inline_never = true;
                         }
                     }
                     _ => unreachable!("stmt kind checked previously"),
@@ -431,11 +431,19 @@ mod llvm_enzyme {
 
         let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
         let d_attr = outer_normal_attr(&rustc_ad_attr, new_id, span);
+
+        // If the source function has the `#[inline(never)]` attribute, we'll also add it to the diff function
+        let mut d_attrs = thin_vec![d_attr];
+
+        if has_inline_never {
+            d_attrs.push(inline_never);
+        }
+
         let d_annotatable = match &item {
             Annotatable::AssocItem(_, _) => {
                 let assoc_item: AssocItemKind = ast::AssocItemKind::Fn(d_fn);
                 let d_fn = P(ast::AssocItem {
-                    attrs: thin_vec![d_attr],
+                    attrs: d_attrs,
                     id: ast::DUMMY_NODE_ID,
                     span,
                     vis,
@@ -445,13 +453,13 @@ mod llvm_enzyme {
                 Annotatable::AssocItem(d_fn, Impl { of_trait: false })
             }
             Annotatable::Item(_) => {
-                let mut d_fn = ecx.item(span, thin_vec![d_attr], ItemKind::Fn(d_fn));
+                let mut d_fn = ecx.item(span, d_attrs, ItemKind::Fn(d_fn));
                 d_fn.vis = vis;
 
                 Annotatable::Item(d_fn)
             }
             Annotatable::Stmt(_) => {
-                let mut d_fn = ecx.item(span, thin_vec![d_attr], ItemKind::Fn(d_fn));
+                let mut d_fn = ecx.item(span, d_attrs, ItemKind::Fn(d_fn));
                 d_fn.vis = vis;
 
                 Annotatable::Stmt(P(ast::Stmt {

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -10,10 +10,9 @@ use tracing::debug;
 use crate::builder::{Builder, PlaceRef, UNNAMED};
 use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
-use crate::llvm::AttributePlace::Function;
+use crate::llvm;
 use crate::llvm::{Metadata, True, Type};
 use crate::value::Value;
-use crate::{attributes, llvm};
 
 pub(crate) fn adjust_activity_to_abi<'tcx>(
     tcx: TyCtxt<'tcx>,
@@ -308,11 +307,6 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
         enzyme_ty,
     );
 
-    // Otherwise LLVM might inline our temporary code before the enzyme pass has a chance to
-    // do it's work.
-    let attr = llvm::AttributeKind::NoInline.create_attr(cx.llcx);
-    attributes::apply_to_llfn(ad_fn, Function, &[attr]);
-
     let num_args = llvm::LLVMCountParams(&fn_to_diff);
     let mut args = Vec::with_capacity(num_args as usize + 1);
     args.push(fn_to_diff);

tests/codegen-llvm/autodiff/batched.rs

@@ -17,6 +17,7 @@ use std::autodiff::autodiff_forward;
 #[autodiff_forward(d_square2, 4, Dual, DualOnly)]
 #[autodiff_forward(d_square1, 4, Dual, Dual)]
 #[no_mangle]
+#[inline(never)]
 fn square(x: &f32) -> f32 {
     x * x
 }

tests/codegen-llvm/autodiff/generic.rs

@@ -6,6 +6,7 @@
 use std::autodiff::autodiff_reverse;
 
 #[autodiff_reverse(d_square, Duplicated, Active)]
+#[inline(never)]
 fn square<T: std::ops::Mul<Output = T> + Copy>(x: &T) -> T {
     *x * *x
 }

tests/codegen-llvm/autodiff/identical_fnc.rs

@@ -14,11 +14,13 @@
 use std::autodiff::autodiff_reverse;
 
 #[autodiff_reverse(d_square, Duplicated, Active)]
+#[inline(never)]
 fn square(x: &f64) -> f64 {
     x * x
 }
 
 #[autodiff_reverse(d_square2, Duplicated, Active)]
+#[inline(never)]
 fn square2(x: &f64) -> f64 {
     x * x
 }
@@ -29,8 +31,10 @@ fn square2(x: &f64) -> f64 {
 // CHECK-NEXT:start:
 // CHECK-NOT:br
 // CHECK-NOT:ret
-// CHECK:call fastcc void @diffe_ZN13identical_fnc6square17h67c6eccd3051fb4cE(double %x.val, ptr %dx1)
-// CHECK-NEXT:call fastcc void @diffe_ZN13identical_fnc6square17h67c6eccd3051fb4cE(double %x.val, ptr %dx2)
+// CHECK:; call identical_fnc::d_square
+// CHECK-NEXT:call fastcc void @_ZN13identical_fnc8d_square17hcb5768e95528c35fE(double %x.val, ptr noalias noundef align 8 dereferenceable(8) %dx1)
+// CHECK:; call identical_fnc::d_square
+// CHECK-NEXT:call fastcc void @_ZN13identical_fnc8d_square17hcb5768e95528c35fE(double %x.val, ptr noalias noundef align 8 dereferenceable(8) %dx2)
 
 fn main() {
     let x = std::hint::black_box(3.0);

tests/codegen-llvm/autodiff/scalar.rs

@@ -7,6 +7,7 @@ use std::autodiff::autodiff_reverse;
 
 #[autodiff_reverse(d_square, Duplicated, Active)]
 #[no_mangle]
+#[inline(never)]
 fn square(x: &f64) -> f64 {
     x * x
 }

tests/codegen-llvm/autodiff/sret.rs

@@ -13,6 +13,7 @@ use std::autodiff::autodiff_reverse;
 
 #[no_mangle]
 #[autodiff_reverse(df, Active, Active, Active)]
+#[inline(never)]
 fn primal(x: f32, y: f32) -> f64 {
     (x * x * y) as f64
 }

tests/pretty/autodiff/autodiff_forward.pp

@@ -3,10 +3,10 @@
 //@ needs-enzyme
 
 #![feature(autodiff)]
-#[prelude_import]
-use ::std::prelude::rust_2015::*;
 #[macro_use]
 extern crate std;
+#[prelude_import]
+use ::std::prelude::rust_2015::*;
 //@ pretty-mode:expanded
 //@ pretty-compare-only
 //@ pp-exact:autodiff_forward.pp
@@ -16,7 +16,6 @@
 use std::autodiff::{autodiff_forward, autodiff_reverse};
 
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f1(x: &[f64], y: f64) -> f64 {
 
 
@@ -40,7 +39,6 @@
     ::core::intrinsics::enzyme_autodiff(f1::<>, df1::<>, (x, bx_0, y))
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f2(x: &[f64], y: f64) -> f64 {
     ::core::panicking::panic("not implemented")
 }
@@ -49,7 +47,6 @@
     ::core::intrinsics::enzyme_autodiff(f2::<>, df2::<>, (x, bx_0, y))
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f3(x: &[f64], y: f64) -> f64 {
     ::core::panicking::panic("not implemented")
 }
@@ -58,14 +55,12 @@
     ::core::intrinsics::enzyme_autodiff(f3::<>, df3::<>, (x, bx_0, y))
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f4() {}
 #[rustc_autodiff(Forward, 1, None)]
 pub fn df4() -> () {
     ::core::intrinsics::enzyme_autodiff(f4::<>, df4::<>, ())
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f5(x: &[f64], y: f64) -> f64 {
     ::core::panicking::panic("not implemented")
 }
@@ -84,7 +79,6 @@
 }
 struct DoesNotImplDefault;
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f6() -> DoesNotImplDefault {
     ::core::panicking::panic("not implemented")
 }
@@ -93,15 +87,13 @@
     ::core::intrinsics::enzyme_autodiff(f6::<>, df6::<>, ())
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f7(x: f32) -> () {}
 #[rustc_autodiff(Forward, 1, Const, None)]
 pub fn df7(x: f32) -> () {
     ::core::intrinsics::enzyme_autodiff(f7::<>, df7::<>, (x,))
 }
 #[no_mangle]
 #[rustc_autodiff]
-#[inline(never)]
 fn f8(x: &f32) -> f32 { ::core::panicking::panic("not implemented") }
 #[rustc_autodiff(Forward, 4, Dual, Dual)]
 fn f8_3(x: &f32, bx_0: &f32, bx_1: &f32, bx_2: &f32, bx_3: &f32)
@@ -121,7 +113,6 @@
 }
 pub fn f9() {
     #[rustc_autodiff]
-    #[inline(never)]
     fn inner(x: f32) -> f32 { x * x }
     #[rustc_autodiff(Forward, 1, Dual, Dual)]
     fn d_inner_2(x: f32, bx_0: f32) -> (f32, f32) {
@@ -135,7 +126,6 @@
     }
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f10<T: std::ops::Mul<Output = T> + Copy>(x: &T) -> T { *x * *x }
 #[rustc_autodiff(Reverse, 1, Duplicated, Active)]
 pub fn d_square<T: std::ops::Mul<Output = T> +

tests/pretty/autodiff/autodiff_reverse.pp

@@ -3,10 +3,10 @@
 //@ needs-enzyme
 
 #![feature(autodiff)]
-#[prelude_import]
-use ::std::prelude::rust_2015::*;
 #[macro_use]
 extern crate std;
+#[prelude_import]
+use ::std::prelude::rust_2015::*;
 //@ pretty-mode:expanded
 //@ pretty-compare-only
 //@ pp-exact:autodiff_reverse.pp
@@ -16,7 +16,6 @@
 use std::autodiff::autodiff_reverse;
 
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f1(x: &[f64], y: f64) -> f64 {
 
     // Not the most interesting derivative, but who are we to judge
@@ -33,12 +32,10 @@
     ::core::intrinsics::enzyme_autodiff(f1::<>, df1::<>, (x, dx_0, y, dret))
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f2() {}
 #[rustc_autodiff(Reverse, 1, None)]
 pub fn df2() { ::core::intrinsics::enzyme_autodiff(f2::<>, df2::<>, ()) }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f3(x: &[f64], y: f64) -> f64 {
     ::core::panicking::panic("not implemented")
 }
@@ -49,14 +46,12 @@
 enum Foo { Reverse, }
 use Foo::Reverse;
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f4(x: f32) { ::core::panicking::panic("not implemented") }
 #[rustc_autodiff(Reverse, 1, Const, None)]
 pub fn df4(x: f32) {
     ::core::intrinsics::enzyme_autodiff(f4::<>, df4::<>, (x,))
 }
 #[rustc_autodiff]
-#[inline(never)]
 pub fn f5(x: *const f32, y: &f32) {
     ::core::panicking::panic("not implemented")
 }

tests/pretty/autodiff/inherent_impl.pp

@@ -3,10 +3,10 @@
 //@ needs-enzyme
 
 #![feature(autodiff)]
-#[prelude_import]
-use ::std::prelude::rust_2015::*;
 #[macro_use]
 extern crate std;
+#[prelude_import]
+use ::std::prelude::rust_2015::*;
 //@ pretty-mode:expanded
 //@ pretty-compare-only
 //@ pp-exact:inherent_impl.pp
@@ -26,7 +26,6 @@
 
 impl MyTrait for Foo {
     #[rustc_autodiff]
-    #[inline(never)]
     fn f(&self, x: f64) -> f64 {
         self.a * 0.25 * (x * x - 1.0 - 2.0 * x.ln())
     }