Handle slices when extracting args from tuple

Author: Sa4dUs

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -3,7 +3,6 @@ use std::cmp::Ordering;
 
 use rustc_abi::{Align, BackendRepr, ExternAbi, Float, HasDataLayout, Primitive, Size};
 use rustc_codegen_ssa::base::{compare_simd_types, wants_msvc_seh, wants_wasm_eh};
-use rustc_codegen_ssa::codegen_attrs::autodiff_attrs;
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
 use rustc_codegen_ssa::errors::{ExpectedPointerMutability, InvalidMonomorphization};
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
@@ -1157,7 +1156,13 @@ fn codegen_enzyme_autodiff<'ll, 'tcx>(
         Instance::try_resolve(tcx, bx.cx.typing_env(), *diff_id, diff_args).unwrap().unwrap();
     let diff_symbol = symbol_name_for_instance_in_crate(tcx, fn_diff.clone(), LOCAL_CRATE);
 
-    let diff_attrs = autodiff_attrs(tcx, fn_diff.def_id());
+    // TODO(Sa4dUs): Store autodiff items in a single pass and just get them here
+    // in a O(1) step
+    let diff_attrs = tcx
+        .collect_and_partition_mono_items(())
+        .autodiff_items
+        .iter()
+        .find(|item| item.target == diff_symbol);
     let Some(diff_attrs) = diff_attrs else { bug!("could not find autodiff attrs") };
 
     // Build body
@@ -1168,7 +1173,7 @@ fn codegen_enzyme_autodiff<'ll, 'tcx>(
         &diff_symbol,
         llret_ty,
         &val_arr,
-        diff_attrs.clone(),
+        diff_attrs.attrs.clone(),
         result,
     );
 }
@@ -1185,11 +1190,22 @@ fn get_args_from_tuple<'ll, 'tcx>(
             for i in 0..tuple_place.layout.layout.0.fields.count() {
                 let field_place = tuple_place.project_field(bx, i);
                 let field_layout = tuple_place.layout.field(bx, i);
+                let field_ty = field_layout.ty;
                 let llvm_ty = field_layout.llvm_type(bx.cx);
 
                 let field_val = bx.load(llvm_ty, field_place.val.llval, field_place.val.align);
 
-                ret_arr.push(field_val)
+                match field_ty.kind() {
+                    ty::Ref(_, inner_ty, _) if matches!(inner_ty.kind(), ty::Slice(_)) => {
+                        let ptr = bx.extract_value(field_val, 0);
+                        let len = bx.extract_value(field_val, 1);
+                        ret_arr.push(ptr);
+                        ret_arr.push(len);
+                    }
+                    _ => {
+                        ret_arr.push(field_val);
+                    }
+                }
             }
 
             ret_arr

compiler/rustc_monomorphize/src/collector.rs

@@ -238,7 +238,7 @@ use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument, trace};
 
 #[cfg(llvm_enzyme)]
-use crate::collector::autodiff::collect_enzyme_autodiff_source_fn;
+use crate::collector::autodiff::collect_enzyme_autodiff_fn;
 use crate::errors::{
     self, EncounteredErrorWhileInstantiating, EncounteredErrorWhileInstantiatingGlobalAsm,
     NoOptimizedMir, RecursionLimit,
@@ -916,7 +916,7 @@ fn visit_instance_use<'tcx>(
     }
     if let Some(intrinsic) = tcx.intrinsic(instance.def_id()) {
         #[cfg(llvm_enzyme)]
-        collect_enzyme_autodiff_source_fn(tcx, instance, intrinsic, output);
+        collect_enzyme_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

compiler/rustc_monomorphize/src/collector/autodiff.rs

@@ -1,10 +1,13 @@
 use rustc_middle::bug;
-use rustc_middle::ty::{self, IntrinsicDef, TyCtxt};
-use tracing::debug;
+use rustc_middle::ty::{self, GenericArg, IntrinsicDef, TyCtxt};
 
 use crate::collector::{MonoItems, create_fn_mono_item};
 
-pub(crate) fn collect_enzyme_autodiff_source_fn<'tcx>(
+// Here, we force both primal and diff function to be collected in
+// mono so this does not interfere in `enzyme_autodiff` intrinsics
+// codegen process. If they are unused, they will be removed later and
+// won't be present at LLVM-IR.
+pub(crate) fn collect_enzyme_autodiff_fn<'tcx>(
     tcx: TyCtxt<'tcx>,
     instance: ty::Instance<'tcx>,
     intrinsic: IntrinsicDef,
@@ -14,8 +17,16 @@ pub(crate) fn collect_enzyme_autodiff_source_fn<'tcx>(
         return;
     };
 
-    debug!("enzyme_autodiff found");
-    let (primal, span) = match instance.args[0].kind() {
+    collect_autodiff_fn_from_arg(instance.args[0], tcx, output);
+    collect_autodiff_fn_from_arg(instance.args[1], tcx, output);
+}
+
+fn collect_autodiff_fn_from_arg<'tcx>(
+    arg: GenericArg<'tcx>,
+    tcx: TyCtxt<'tcx>,
+    output: &mut MonoItems<'tcx>,
+) {
+    let (instance, span) = match arg.kind() {
         rustc_middle::infer::canonical::ir::GenericArgKind::Type(ty) => match ty.kind() {
             ty::FnDef(def_id, substs) => {
                 let span = tcx.def_span(def_id);
@@ -34,5 +45,5 @@ pub(crate) fn collect_enzyme_autodiff_source_fn<'tcx>(
         _ => bug!("expected type"),
     };
 
-    output.push(create_fn_mono_item(tcx, primal, span));
+    output.push(create_fn_mono_item(tcx, instance, span));
 }