From 466bec90292a1b04a50a5889ab13fddd53c925ce Mon Sep 17 00:00:00 2001 From: Marcelo Domínguez Date: Mon, 14 Jul 2025 17:07:01 +0000 Subject: Adjust autodiff actitivies for ScalarPair --- compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) (limited to 'compiler/rustc_codegen_llvm/src/builder/autodiff.rs') diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index e2df3265f6f..a1cc4fdbdf6 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -80,6 +80,23 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( 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); + } + }; + + match layout.backend_repr() { + rustc_abi::BackendRepr::ScalarPair(_, _) => { + 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 -- cgit 1.4.1-3-g733a5 From 8dbd1b014ac43cc9c950e3b7e112b62d0963ec17 Mon Sep 17 00:00:00 2001 From: Marcelo Domínguez Date: Tue, 15 Jul 2025 09:56:51 +0000 Subject: doc and move single branch match to an if let --- compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'compiler/rustc_codegen_llvm/src/builder/autodiff.rs') diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index a1cc4fdbdf6..a19a0d867ac 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -90,12 +90,12 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( } }; - match layout.backend_repr() { - rustc_abi::BackendRepr::ScalarPair(_, _) => { - new_activities.push(da[i].clone()); - new_positions.push(i + 1); - } - _ => {} + // If the argument is lowered as a `ScalarPair`, we need to duplicate its activity. + // Otherwise, the number of activities won't match the number of LLVM arguments and + // this will lead to errors when verifying the Enzyme call. + if let rustc_abi::BackendRepr::ScalarPair(_, _) = layout.backend_repr() { + new_activities.push(da[i].clone()); + new_positions.push(i + 1); } } // now add the extra activities coming from slices -- cgit 1.4.1-3-g733a5 From 0bf85d35ec9ed9cbffd274def09027c5fe9a8f3c Mon Sep 17 00:00:00 2001 From: Marcelo Domínguez Date: Sun, 3 Aug 2025 11:12:34 +0000 Subject: Support ZST args --- compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 11 ++++++++++- tests/ui/autodiff/zst.rs | 17 +++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 tests/ui/autodiff/zst.rs (limited to 'compiler/rustc_codegen_llvm/src/builder/autodiff.rs') diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index a19a0d867ac..78deffa3a7a 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -29,6 +29,7 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( let mut new_activities = vec![]; let mut new_positions = vec![]; + let mut del_activities = 0; for (i, ty) in sig.inputs().iter().enumerate() { if let Some(inner_ty) = ty.builtin_deref(true) { if inner_ty.is_slice() { @@ -90,12 +91,20 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( } }; + // For ZST, just ignore and don't add its activity, as this arg won't be present + // in the LLVM passed to Enzyme. + // FIXME(Sa4dUs): Enforce ZST corresponding diff activity be `Const` + if layout.is_zst() { + del_activities += 1; + da.remove(i); + } + // If the argument is lowered as a `ScalarPair`, we need to duplicate its activity. // Otherwise, the number of activities won't match the number of LLVM arguments and // this will lead to errors when verifying the Enzyme call. if let rustc_abi::BackendRepr::ScalarPair(_, _) = layout.backend_repr() { new_activities.push(da[i].clone()); - new_positions.push(i + 1); + new_positions.push(i + 1 - del_activities); } } // now add the extra activities coming from slices diff --git a/tests/ui/autodiff/zst.rs b/tests/ui/autodiff/zst.rs new file mode 100644 index 00000000000..7b9b5f5f20b --- /dev/null +++ b/tests/ui/autodiff/zst.rs @@ -0,0 +1,17 @@ +//@ compile-flags: -Zautodiff=Enable -C opt-level=3 -Clto=fat +//@ no-prefer-dynamic +//@ needs-enzyme +//@ build-pass + +// Check that differentiating functions with ZST args does not break + +#![feature(autodiff)] + +#[core::autodiff::autodiff_forward(fd_inner, Const, Dual)] +fn f(_zst: (), _x: &mut f64) {} + +fn fd(x: &mut f64, xd: &mut f64) { + fd_inner((), x, xd); +} + +fn main() {} -- cgit 1.4.1-3-g733a5 From e04567c363e1f0417bf8bf24830c2bc536020582 Mon Sep 17 00:00:00 2001 From: Marcelo Domínguez Date: Wed, 17 Sep 2025 13:58:17 +0000 Subject: Check ZST via `PassMode` --- compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 20 +++++++++++++++++--- compiler/rustc_codegen_llvm/src/intrinsic.rs | 3 ++- 2 files changed, 19 insertions(+), 4 deletions(-) (limited to 'compiler/rustc_codegen_llvm/src/builder/autodiff.rs') diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index 78deffa3a7a..566877f4a1e 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -3,8 +3,9 @@ use std::ptr; use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode}; use rustc_codegen_ssa::common::TypeKind; use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods}; -use rustc_middle::ty::{PseudoCanonicalInput, Ty, TyCtxt, TypingEnv}; +use rustc_middle::ty::{Instance, PseudoCanonicalInput, TyCtxt, TypingEnv}; use rustc_middle::{bug, ty}; +use rustc_target::callconv::PassMode; use tracing::debug; use crate::builder::{Builder, PlaceRef, UNNAMED}; @@ -16,9 +17,12 @@ use crate::value::Value; pub(crate) fn adjust_activity_to_abi<'tcx>( tcx: TyCtxt<'tcx>, - fn_ty: Ty<'tcx>, + instance: Instance<'tcx>, + typing_env: TypingEnv<'tcx>, da: &mut Vec, ) { + let fn_ty = instance.ty(tcx, typing_env); + if !matches!(fn_ty.kind(), ty::FnDef(..)) { bug!("expected fn def for autodiff, got {:?}", fn_ty); } @@ -27,6 +31,13 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( // All we do is decide how to handle the arguments. let sig = fn_ty.fn_sig(tcx).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()))) + else { + bug!("failed to get fn_abi of instance with empty varargs"); + }; + let mut new_activities = vec![]; let mut new_positions = vec![]; let mut del_activities = 0; @@ -91,10 +102,13 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( } }; + let pass_mode = &fn_abi.args[i].mode; + // For ZST, just ignore and don't add its activity, as this arg won't be present // in the LLVM passed to Enzyme. + // Some targets pass ZST indirectly in the C ABI, in that case, handle it as a normal arg // FIXME(Sa4dUs): Enforce ZST corresponding diff activity be `Const` - if layout.is_zst() { + if *pass_mode == PassMode::Ignore { del_activities += 1; da.remove(i); } diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 06c3d8ed6bc..9e6e7606491 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1198,7 +1198,8 @@ fn codegen_autodiff<'ll, 'tcx>( adjust_activity_to_abi( tcx, - fn_source.ty(tcx, TypingEnv::fully_monomorphized()), + fn_source, + TypingEnv::fully_monomorphized(), &mut diff_attrs.input_activity, ); -- cgit 1.4.1-3-g733a5