Adjust autodiff activities for abi transformations (small tuples and structs)

Author: Sa4dUs

compiler/rustc_monomorphize/src/partitioning/autodiff.rs

@@ -1,3 +1,4 @@
+use rustc_abi::HasDataLayout;
 use rustc_ast::expand::autodiff_attrs::{AutoDiffItem, DiffActivity};
 use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_middle::bug;
@@ -16,6 +17,7 @@ fn adjust_activity_to_abi<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>, da: &mut Vec
     // 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 pointer_size = tcx.data_layout().pointer_size;
 
     let mut new_activities = vec![];
     let mut new_positions = vec![];
@@ -70,6 +72,25 @@ fn adjust_activity_to_abi<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>, da: &mut Vec
                 continue;
             }
         }
+
+        let pci = PseudoCanonicalInput { typing_env: TypingEnv::fully_monomorphized(), value: *ty };
+
+        let layout = match tcx.layout_of(pci) {
+            Ok(layout) => layout.layout,
+            Err(_) => {
+                bug!("failed to compute layout for type {:?}", ty);
+            }
+        };
+
+        let is_product = |t: Ty<'tcx>| matches!(t.kind(), ty::Tuple(_) | ty::Adt(_, _));
+
+        if layout.size() <= pointer_size * 2 && is_product(*ty) {
+            let n_scalars = count_scalar_fields(tcx, *ty);
+            for _ in 0..n_scalars.saturating_sub(1) {
+                new_activities.push(da[i].clone());
+                new_positions.push(i + 1);
+            }
+        }
     }
     // now add the extra activities coming from slices
     // Reverse order to not invalidate the indices
@@ -80,6 +101,20 @@ fn adjust_activity_to_abi<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>, da: &mut Vec
     }
 }
 
+fn count_scalar_fields<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> usize {
+    match ty.kind() {
+        ty::Float(_) | ty::Int(_) | ty::Uint(_) => 1,
+        ty::Adt(def, substs) if def.is_struct() => def
+            .non_enum_variant()
+            .fields
+            .iter()
+            .map(|f| count_scalar_fields(tcx, f.ty(tcx, substs)))
+            .sum(),
+        ty::Tuple(substs) => substs.iter().map(|t| count_scalar_fields(tcx, t)).sum(),
+        _ => 0,
+    }
+}
+
 pub(crate) fn find_autodiff_source_functions<'tcx>(
     tcx: TyCtxt<'tcx>,
     usage_map: &UsageMap<'tcx>,