Auto merge of #153379 - TKanX:refactor/149164-simplify-autodiff-rlib, r=ZuseZ4

refactor(autodiff): Simplify Autodiff Handling of `rlib` Dependencies

### Summary:

Resolves the two FIXMEs left in #149033, per @bjorn3 guidance in [the discussion](#149033 (comment)).

Closes #149164 

r? @ZuseZ4
cc @bjorn3

Author: bors

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -20,7 +20,7 @@ mod llvm_enzyme {
     };
     use rustc_expand::base::{Annotatable, ExtCtxt};
     use rustc_hir::attrs::RustcAutodiff;
-    use rustc_span::{Ident, Span, Symbol, sym};
+    use rustc_span::{Ident, Span, Symbol, kw, sym};
     use thin_vec::{ThinVec, thin_vec};
     use tracing::{debug, trace};
 
@@ -197,7 +197,7 @@ mod llvm_enzyme {
     /// }
     /// #[rustc_autodiff(Reverse, Duplicated, Active)]
     /// fn cos_box(x: &Box<f32>, dx: &mut Box<f32>, dret: f32) -> f32 {
-    ///     std::intrinsics::autodiff(sin::<>, cos_box::<>, (x, dx, dret))
+    ///     std::intrinsics::autodiff(sin::<> as fn(..) -> .., cos_box::<>, (x, dx, dret))
     /// }
     /// ```
     /// FIXME(ZuseZ4): Once autodiff is enabled by default, make this a doc comment which is checked
@@ -326,6 +326,7 @@ mod llvm_enzyme {
                 primal,
                 first_ident(&meta_item_vec[0]),
                 span,
+                &sig,
                 &d_sig,
                 &generics,
                 is_impl,
@@ -496,18 +497,62 @@ mod llvm_enzyme {
 
     // Generate `autodiff` intrinsic call
     // ```
-    // std::intrinsics::autodiff(source, diff, (args))
+    // std::intrinsics::autodiff(source as fn(..) -> .., diff, (args))
     // ```
     fn call_autodiff(
         ecx: &ExtCtxt<'_>,
         primal: Ident,
         diff: Ident,
         span: Span,
+        p_sig: &FnSig,
         d_sig: &FnSig,
         generics: &Generics,
         is_impl: bool,
     ) -> rustc_ast::Stmt {
         let primal_path_expr = gen_turbofish_expr(ecx, primal, generics, span, is_impl);
+
+        let self_ty = || ecx.ty_path(ast::Path::from_ident(Ident::with_dummy_span(kw::SelfUpper)));
+        let fn_ptr_params: ThinVec<ast::Param> = p_sig
+            .decl
+            .inputs
+            .iter()
+            .map(|param| {
+                let ty = match &param.ty.kind {
+                    TyKind::ImplicitSelf => self_ty(),
+                    TyKind::Ref(lt, mt) if matches!(mt.ty.kind, TyKind::ImplicitSelf) => ecx.ty(
+                        span,
+                        TyKind::Ref(lt.clone(), ast::MutTy { ty: self_ty(), mutbl: mt.mutbl }),
+                    ),
+                    TyKind::Ptr(mt) if matches!(mt.ty.kind, TyKind::ImplicitSelf) => {
+                        ecx.ty(span, TyKind::Ptr(ast::MutTy { ty: self_ty(), mutbl: mt.mutbl }))
+                    }
+                    _ => param.ty.clone(),
+                };
+                ast::Param {
+                    attrs: ast::AttrVec::new(),
+                    ty,
+                    pat: Box::new(ecx.pat_wild(span)),
+                    id: ast::DUMMY_NODE_ID,
+                    span,
+                    is_placeholder: false,
+                }
+            })
+            .collect();
+        let fn_ptr_ty = ecx.ty(
+            span,
+            TyKind::FnPtr(Box::new(ast::FnPtrTy {
+                safety: p_sig.header.safety,
+                ext: p_sig.header.ext,
+                generic_params: ThinVec::new(),
+                decl: Box::new(ast::FnDecl {
+                    inputs: fn_ptr_params,
+                    output: p_sig.decl.output.clone(),
+                }),
+                decl_span: span,
+            })),
+        );
+        let primal_fn_ptr = ecx.expr(span, ast::ExprKind::Cast(primal_path_expr, fn_ptr_ty));
+
         let diff_path_expr = gen_turbofish_expr(ecx, diff, generics, span, is_impl);
 
         let tuple_expr = ecx.expr_tuple(
@@ -529,7 +574,7 @@ mod llvm_enzyme {
         let call_expr = ecx.expr_call(
             span,
             ecx.expr_path(enzyme_path),
-            vec![primal_path_expr, diff_path_expr, tuple_expr].into(),
+            vec![primal_fn_ptr, diff_path_expr, tuple_expr].into(),
         );
 
         ecx.stmt_expr(call_expr)

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -6,7 +6,7 @@ use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods};
 use rustc_data_structures::thin_vec::ThinVec;
 use rustc_hir::attrs::RustcAutodiff;
-use rustc_middle::ty::{Instance, PseudoCanonicalInput, TyCtxt, TypingEnv};
+use rustc_middle::ty::{PseudoCanonicalInput, Ty, TyCtxt, TypingEnv};
 use rustc_middle::{bug, ty};
 use rustc_target::callconv::PassMode;
 use tracing::debug;
@@ -18,25 +18,23 @@ use crate::llvm::{self, TRUE, Type, Value};
 
 pub(crate) fn adjust_activity_to_abi<'tcx>(
     tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
+    fn_ptr_ty: Ty<'tcx>,
     typing_env: TypingEnv<'tcx>,
     da: &mut ThinVec<DiffActivity>,
 ) {
-    let fn_ty = instance.ty(tcx, typing_env);
-
-    if !matches!(fn_ty.kind(), ty::FnDef(..)) {
-        bug!("expected fn def for autodiff, got {:?}", fn_ty);
+    if !matches!(fn_ptr_ty.kind(), ty::FnPtr(..)) {
+        bug!("expected fn ptr for autodiff, got {:?}", fn_ptr_ty);
     }
 
     // We don't actually pass the types back into the type system.
     // All we do is decide how to handle the arguments.
-    let sig = fn_ty.fn_sig(tcx).skip_binder();
+    let fn_sig = fn_ptr_ty.fn_sig(tcx);
+    let sig = fn_sig.skip_binder();
 
     // FIXME(Sa4dUs): pass proper varargs once we have support for differentiating variadic functions
-    let Ok(fn_abi) =
-        tcx.fn_abi_of_instance(typing_env.as_query_input((instance, ty::List::empty())))
+    let Ok(fn_abi) = tcx.fn_abi_of_fn_ptr(typing_env.as_query_input((fn_sig, ty::List::empty())))
     else {
-        bug!("failed to get fn_abi of instance with empty varargs");
+        bug!("failed to get fn_abi of fn_ptr with empty varargs");
     };
 
     let mut new_activities = vec![];

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1322,29 +1322,8 @@ fn codegen_autodiff<'ll, 'tcx>(
     let ret_ty = sig.output();
     let llret_ty = bx.layout_of(ret_ty).llvm_type(bx);
 
-    // Get source, diff, and attrs
-    let (source_id, source_args) = match fn_args.into_type_list(tcx)[0].kind() {
-        ty::FnDef(def_id, source_params) => (def_id, source_params),
-        _ => bug!("invalid autodiff intrinsic args"),
-    };
-
-    let fn_source = match Instance::try_resolve(tcx, bx.cx.typing_env(), *source_id, source_args) {
-        Ok(Some(instance)) => instance,
-        Ok(None) => bug!(
-            "could not resolve ({:?}, {:?}) to a specific autodiff instance",
-            source_id,
-            source_args
-        ),
-        Err(_) => {
-            // An error has already been emitted
-            return;
-        }
-    };
-
-    let source_symbol = symbol_name_for_instance_in_crate(tcx, fn_source.clone(), LOCAL_CRATE);
-    let Some(fn_to_diff) = bx.cx.get_function(&source_symbol) else {
-        bug!("could not find source function")
-    };
+    let source_fn_ptr_ty = fn_args.into_type_list(tcx)[0];
+    let fn_to_diff = args[0].immediate();
 
     let (diff_id, diff_args) = match fn_args.into_type_list(tcx)[1].kind() {
         ty::FnDef(def_id, diff_args) => (def_id, diff_args),
@@ -1375,13 +1354,12 @@ fn codegen_autodiff<'ll, 'tcx>(
 
     adjust_activity_to_abi(
         tcx,
-        fn_source,
+        source_fn_ptr_ty,
         TypingEnv::fully_monomorphized(),
         &mut diff_attrs.input_activity,
     );
 
-    let fnc_tree =
-        rustc_middle::ty::fnc_typetrees(tcx, fn_source.ty(tcx, TypingEnv::fully_monomorphized()));
+    let fnc_tree = rustc_middle::ty::fnc_typetrees(tcx, source_fn_ptr_ty);
 
     // Build body
     generate_enzyme_call(

compiler/rustc_mir_transform/src/cross_crate_inline.rs

@@ -35,11 +35,6 @@ fn cross_crate_inlinable(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
         return true;
     }
 
-    // FIXME(autodiff): replace this as per discussion in https://github.com/rust-lang/rust/pull/149033#discussion_r2535465880
-    if find_attr!(tcx, def_id, RustcAutodiff(..)) {
-        return true;
-    }
-
     if find_attr!(tcx, def_id, RustcIntrinsic) {
         // Intrinsic fallback bodies are always cross-crate inlineable.
         // To ensure that the MIR inliner doesn't cluelessly try to inline fallback

compiler/rustc_monomorphize/src/collector.rs

@@ -205,8 +205,6 @@
 //! this is not implemented however: a mono item will be produced
 //! regardless of whether it is actually needed or not.
 
-mod autodiff;
-
 use std::cell::OnceCell;
 use std::ops::ControlFlow;
 
@@ -240,7 +238,6 @@ use rustc_span::source_map::{Spanned, dummy_spanned, respan};
 use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument, trace};
 
-use crate::collector::autodiff::collect_autodiff_fn;
 use crate::errors::{
     self, EncounteredErrorWhileInstantiating, EncounteredErrorWhileInstantiatingGlobalAsm,
     NoOptimizedMir, RecursionLimit,
@@ -990,8 +987,6 @@ fn visit_instance_use<'tcx>(
         return;
     }
     if let Some(intrinsic) = tcx.intrinsic(instance.def_id()) {
-        collect_autodiff_fn(tcx, instance, intrinsic, output);
-
         if let Some(_requirement) = ValidityRequirement::from_intrinsic(intrinsic.name) {
             // The intrinsics assert_inhabited, assert_zero_valid, and assert_mem_uninitialized_valid will
             // be lowered in codegen to nothing or a call to panic_nounwind. So if we encounter any