Prevent ABI changes affect EnzymeAD #142544
Conversation
rustbot
added
A-attributes
|
The job Click to see the possible cause of the failure (guessed by this bot) |
| /// `#[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; |
There was a problem hiding this comment.
Don't say "certain optimizations", or the next person who comes along is going to make it so that these functions are treated as -O0. Identify the actual problem: LLVM will modify the ABI of functions if it can identify them as fully internalized.
There was a problem hiding this comment.
A change that bad should be caught by a reviewer, but I review a lot of PRs and I enjoy it when the codebase informs people of what is actually going on so they are more likely to have made changes that are consistent with the existing situation.
There was a problem hiding this comment.
Got it, thanks for the feedback. I’ll make sure to keep comments as specific as possible to avoid any ambiguity going forward.
| 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, |
There was a problem hiding this comment.
This should only matter when Enzyme is on, right?
Why are you ignoring ty::Array when it is not through ty::Ref? Does Enzyme not even deal in simple aggregates like that? I'm not even sure this is the correct layer to be examining things like this at, since it's well above the LLVM IR type layer. Multiple types in Rust source can wind up being lowered to the naive equivalent of this in the IR.
Anyway, please name this fn as specific to Enzyme, at least.
There was a problem hiding this comment.
Yes, this logic is only for Enzyme. I didn't add the ty::Array logic yet because I'm still not sure on how to handle it as, for example, [f32; 2] is lowered to i64. As the number of args does not change, Enzyme may not have issues with that.
There was a problem hiding this comment.
Okay, I have determined this is blocked on me to have a proper solution for it. I would like it if you opened an issue before this PR lands that points to https://github.com/rust-lang/rust/blame/86d0aef80403f095d8bbabf44d9fdecfcd45f076/compiler/rustc_target/src/callconv/mod.rs#L708 and your new code here, and says that a new variant of the adjust_for_rust_abi code that doesn't just mutate the arguments needs to exist so that this query can be answered without relying on mutable state that cannot be invoked idempotently.
| @@ -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) | |||
There was a problem hiding this comment.
This will also suppress the dead_code lint. I think the reason the symbol is still getting marked as dso_local even with this change is because it has SymbolExportLevel::Rust. Only for SymbolExportLevel::C do we tell LTO to export the symbol.
There was a problem hiding this comment.
Okay, I'll look into that and try to ensure it has the minimal number of side effects possible. Thank you :)
This PR handles ABI changes for autodiff input arguments to improve Enzyme compatibility. A couple of cases (like statics and small arrays) are still not handled.
r? @ZuseZ4