Partially disable ABI optimizations for ad functions

Author: Sa4dUs

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -1,6 +1,6 @@
 use std::str::FromStr;
 
-use rustc_abi::ExternAbi;
+use rustc_abi::{ExternAbi, Size};
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
 use rustc_ast::{LitKind, MetaItem, MetaItemInner, attr};
 use rustc_attr_data_structures::ReprAttr::ReprAlign;
@@ -16,7 +16,7 @@ use rustc_middle::middle::codegen_fn_attrs::{
 use rustc_middle::mir::mono::Linkage;
 use rustc_middle::query::Providers;
 use rustc_middle::span_bug;
-use rustc_middle::ty::{self as ty, TyCtxt};
+use rustc_middle::ty::{self as ty, PseudoCanonicalInput, Ty, TyCtxt, TypingEnv};
 use rustc_session::parse::feature_err;
 use rustc_session::{Session, lint};
 use rustc_span::{Ident, Span, sym};
@@ -138,6 +138,28 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
             sym::rustc_allocator_zeroed => {
                 codegen_fn_attrs.flags |= CodegenFnAttrFlags::ALLOCATOR_ZEROED
             }
+            sym::rustc_autodiff => {
+                let list = attr.meta_item_list().unwrap_or_default();
+                if list.is_empty() {
+                    // Add the flag only to the primal function so LLVM can
+                    // optimize the derivative function.
+                    if let Some(sig) = fn_sig() {
+                        let sig = sig.skip_binder();
+
+                        let has_problematic_args = sig
+                            .skip_binder()
+                            .inputs()
+                            .iter()
+                            .any(|ty| is_abi_opt_sensitive(tcx, *ty));
+
+                        if has_problematic_args {
+                            codegen_fn_attrs.flags |= CodegenFnAttrFlags::RUSTC_AUTODIFF_NO_ABI_OPT;
+                        }
+                    }
+
+                    // TODO(Sa4dUs): Handle static variable passed as argument case.
+                }
+            }
             sym::naked => codegen_fn_attrs.flags |= CodegenFnAttrFlags::NAKED,
             sym::no_mangle => {
                 no_mangle_span = Some(attr.span());
@@ -899,6 +921,44 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
     Some(AutoDiffAttrs { mode, width, ret_activity, input_activity: arg_activities })
 }
 
+fn is_abi_opt_sensitive<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> bool {
+    match ty.kind() {
+        ty::Ref(_, inner, _) | ty::RawPtr(inner, _) => {
+            match inner.kind() {
+                ty::Slice(_) => {
+                    // Since we cannot guarantee that the slice length is large enough
+                    // to avoid optimization, we assume it is ABI-opt sensitive.
+                    return true;
+                }
+                ty::Array(elem_ty, len) => {
+                    let Some(len_val) = len.try_to_target_usize(tcx) else {
+                        return false;
+                    };
+
+                    let pci = PseudoCanonicalInput {
+                        typing_env: TypingEnv::fully_monomorphized(),
+                        value: *elem_ty,
+                    };
+
+                    if elem_ty.is_scalar() {
+                        let elem_size =
+                            tcx.layout_of(pci).ok().map(|layout| layout.size).unwrap_or(Size::ZERO);
+
+                        if elem_size.bytes() * len_val <= tcx.data_layout.pointer_size.bytes() * 2 {
+                            return true;
+                        }
+                    }
+                }
+                _ => {}
+            }
+
+            false
+        }
+        ty::FnPtr(_, _) => true,
+        _ => false,
+    }
+}
+
 pub(crate) fn provide(providers: &mut Providers) {
     *providers = Providers { codegen_fn_attrs, should_inherit_track_caller, ..*providers };
 }

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -136,6 +136,10 @@ bitflags::bitflags! {
         const ALLOCATOR_ZEROED          = 1 << 14;
         /// `#[no_builtins]`: indicates that disable implicit builtin knowledge of functions for the function.
         const NO_BUILTINS               = 1 << 15;
+        /// `#[rustc_autodiff_no_abi_opt]`: internal marker applied to `#[rustc_autodiff]` primal functions
+        /// whose argument layout may be sensitive to ABI-level optimizations. This marker prevents certain
+        /// optimizations that could otherwise break compatibility with Enzyme's expectations.
+        const RUSTC_AUTODIFF_NO_ABI_OPT = 1 << 16;
     }
 }
 rustc_data_structures::external_bitflags_debug! { CodegenFnAttrFlags }
@@ -175,6 +179,7 @@ impl CodegenFnAttrs {
         self.flags.contains(CodegenFnAttrFlags::NO_MANGLE)
             || self.flags.contains(CodegenFnAttrFlags::RUSTC_STD_INTERNAL_SYMBOL)
             || self.export_name.is_some()
+            || self.flags.contains(CodegenFnAttrFlags::RUSTC_AUTODIFF_NO_ABI_OPT)
             || match self.linkage {
                 // These are private, so make sure we don't try to consider
                 // them external.