Rollup merge of rust-lang#157181 - scottmcm:no-autodiff-alloca, r=ZuseZ4

autodiff: stop always needing an alloca

There will still be one for things that aren't just scalars or scalar pairs.

r? @ZuseZ4
cc rust-lang#153250 (the tracking issue for removing intrinsic allocas)

Author: JonathanBrouwer

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -3,6 +3,9 @@ use std::ptr;
 use rustc_ast::expand::autodiff_attrs::{DiffActivity, DiffMode};
 use rustc_ast::expand::typetree::FncTree;
 use rustc_codegen_ssa::common::TypeKind;
+use rustc_codegen_ssa::mir::IntrinsicResult;
+use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
+use rustc_codegen_ssa::mir::place::PlaceValue;
 use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods};
 use rustc_data_structures::thin_vec::ThinVec;
 use rustc_hir::attrs::RustcAutodiff;
@@ -11,7 +14,7 @@ use rustc_middle::{bug, ty};
 use rustc_target::callconv::PassMode;
 use tracing::debug;
 
-use crate::builder::{Builder, PlaceRef, UNNAMED};
+use crate::builder::{Builder, UNNAMED};
 use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
 use crate::llvm::{self, TRUE, Type, Value};
@@ -296,9 +299,10 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
     ret_ty: &'ll Type,
     fn_args: &[&'ll Value],
     attrs: &RustcAutodiff,
-    dest: PlaceRef<'tcx, &'ll Value>,
+    dest_layout: ty::layout::TyAndLayout<'tcx>,
+    dest_place: Option<PlaceValue<&'ll Value>>,
     fnc_tree: FncTree,
-) {
+) -> IntrinsicResult<'tcx, &'ll Value> {
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
     let mut ad_name: String = match attrs.mode {
         DiffMode::Forward => "__enzyme_fwddiff",
@@ -381,11 +385,18 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
     let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);
 
     let fn_ret_ty = builder.cx.val_ty(call);
-    if fn_ret_ty != builder.cx.type_void() && fn_ret_ty != builder.cx.type_struct(&[], false) {
+    if fn_ret_ty == builder.cx.type_void() || fn_ret_ty == builder.cx.type_struct(&[], false) {
         // If we return void or an empty struct, then our caller (due to how we generated it)
         // does not expect a return value. As such, we have no pointer (or place) into which
         // we could store our value, and would store into an undef, which would cause UB.
         // As such, we just ignore the return value in those cases.
-        builder.store_to_place(call, dest.val);
+        IntrinsicResult::Operand(OperandValue::ZeroSized)
+    } else if let Some(dest_place) = dest_place {
+        builder.store_to_place(call, dest_place);
+        IntrinsicResult::WroteIntoPlace
+    } else {
+        IntrinsicResult::Operand(
+            OperandRef::from_immediate_or_packed_pair(builder, call, dest_layout).val,
+        )
     }
 }

compiler/rustc_codegen_llvm/src/context.rs

@@ -978,10 +978,7 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
     }
 
     fn intrinsic_call_expects_place_always(&self, name: Symbol) -> bool {
-        matches!(
-            name,
-            sym::autodiff | sym::volatile_load | sym::unaligned_volatile_load | sym::black_box
-        )
+        matches!(name, sym::volatile_load | sym::unaligned_volatile_load | sym::black_box)
     }
 }
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -223,12 +223,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 )
             }
             sym::autodiff => {
-                let result = PlaceRef {
-                    val: result_place.unwrap(),
-                    layout: result_layout,
-                };
-                codegen_autodiff(self, tcx, instance, args, result);
-                return IntrinsicResult::WroteIntoPlace;
+                return codegen_autodiff(self, tcx, instance, args, result_layout, result_place);
             }
             sym::offload => {
                 if tcx.sess.opts.unstable_opts.offload.is_empty() {
@@ -1728,8 +1723,9 @@ fn codegen_autodiff<'ll, 'tcx>(
     tcx: TyCtxt<'tcx>,
     instance: ty::Instance<'tcx>,
     args: &[OperandRef<'tcx, &'ll Value>],
-    result: PlaceRef<'tcx, &'ll Value>,
-) {
+    result_layout: ty::layout::TyAndLayout<'tcx>,
+    result_place: Option<PlaceValue<&'ll Value>>,
+) -> IntrinsicResult<'tcx, &'ll Value> {
     if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) {
         let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutEnable);
     }
@@ -1769,9 +1765,9 @@ fn codegen_autodiff<'ll, 'tcx>(
             diff_id,
             diff_args
         ),
-        Err(_) => {
+        Err(err) => {
             // An error has already been emitted
-            return;
+            return IntrinsicResult::Err(err);
         }
     };
 
@@ -1802,9 +1798,10 @@ fn codegen_autodiff<'ll, 'tcx>(
         llret_ty,
         &val_arr,
         &diff_attrs,
-        result,
+        result_layout,
+        result_place,
         fnc_tree,
-    );
+    )
 }
 
 // Generates the LLVM code to offload a Rust function to a target device (e.g., GPU).