Skip to content

autodiff rlib handling #149033

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
+34 −2
Draft

autodiff rlib handling #149033

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions compiler/rustc_codegen_llvm/src/intrinsic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ use rustc_target::callconv::PassMode;
use rustc_target::spec::Os;
use tracing::debug;

use rustc_session::config::CrateType;

use crate::abi::FnAbiLlvmExt;
use crate::builder::Builder;
use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
Expand Down Expand Up @@ -1136,8 +1138,17 @@ fn codegen_autodiff<'ll, 'tcx>(
if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) {
let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutEnable);
}
if tcx.sess.lto() != rustc_session::config::Lto::Fat {
let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);

let ct = tcx.crate_types();
let lto = tcx.sess.lto();
if ct.len() == 1 && ct.contains(&CrateType::Executable) {
if lto != rustc_session::config::Lto::Fat {
let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);
}
} else {
if lto != rustc_session::config::Lto::Fat && !tcx.sess.opts.cg.linker_plugin_lto.enabled(){
let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);
}
}

let fn_args = instance.args;
Expand Down
4 changes: 4 additions & 0 deletions compiler/rustc_mir_transform/src/cross_crate_inline.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,10 @@ fn cross_crate_inlinable(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
return true;
}

if tcx.has_attr(def_id, sym::autodiff_forward) || tcx.has_attr(def_id, sym::autodiff_reverse) || tcx.has_attr(def_id, sym::rustc_autodiff) {
return true;
}
Copy link
Member

@bjorn3 bjorn3 Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is rustc_monomorphize::collector::autodiff::collect_autodiff_fn even necessary in the first place. The catch_unwind intrinsic also needs the function arguments to be codegened, yet it doesn't have any special casing in the monomorphization collector. Enzyme needs the function to differentiate to be in the same module, but I would expect fat LTO to run early enough to already cause that to happen. If you remove collect_autodiff_fn then I would assume this special case in cross_crate_inlineable isn't necessary either.

Copy link
Member Author

@ZuseZ4 ZuseZ4 Nov 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@bjorn3 Unfortunately it seems to be. I've tried removing both and we are back to an old bug, could not find source function.
In the beginning we had this case when we'd differentiate f to generate g, but then only call g, not f.
During the lowering we'd realize that f is unused, delete it, and then fail on the llvm-ir level, since we now can't generate the body of g since it depended on f. Our new intrinsic in g in theory references the source function f, but it seems that the monomorphize adjustments are still needed.

error: internal compiler error: compiler/rustc_codegen_llvm/src/intrinsic.rs:1183:9: could not find source function


thread 'rustc' (815281) panicked at compiler/rustc_codegen_llvm/src/intrinsic.rs:1183:9:
Box<dyn Any>
stack backtrace:
   0:     0x72cccc49e5c2 - <std::sys::backtrace::BacktraceLock::print::DisplayBacktrace as core::fmt::Display>::fmt::h47627beda0767f4c
   1:     0x72cccc4cc9af - core::fmt::write::h8a63064c67176fd8
   2:     0x72cccc4516e3 - std::io::Write::write_fmt::h5e9b646cef91150b
   3:     0x72cccc466ae2 - std::sys::backtrace::BacktraceLock::print::h4634d27c0afbaebb
   4:     0x72cccc46ee7a - std::panicking::default_hook::{{closure}}::h23e0299d4738b766
   5:     0x72cccc46ecd4 - std::panicking::default_hook::hfeec6294e8751b62
   6:     0x72ccc8bd3a36 - std[e20d0b87aa1ec278]::panicking::update_hook::<alloc[d35db71d034d3387]::boxed::Box<rustc_driver_impl[ee6e3d57fbee4319]::install_ice_hook::{closure#1}>>::{closure#0}
   7:     0x72cccc46f52f - std::panicking::panic_with_hook::hd57cd463d793ed79
   8:     0x72cccbf6b8c3 - std[e20d0b87aa1ec278]::panicking::begin_panic::<rustc_errors[52538de81f6f8989]::ExplicitBug>::{closure#0}
   9:     0x72cccbf6a976 - std[e20d0b87aa1ec278]::sys::backtrace::__rust_end_short_backtrace::<std[e20d0b87aa1ec278]::panicking::begin_panic<rustc_errors[52538de81f6f8989]::ExplicitBug>::{closure#0}, !>
  10:     0x72cccbf63f5a - std[e20d0b87aa1ec278]::panicking::begin_panic::<rustc_errors[52538de81f6f8989]::ExplicitBug>
  11:     0x72cccbf4c652 - <rustc_errors[52538de81f6f8989]::diagnostic::BugAbort as rustc_errors[52538de81f6f8989]::diagnostic::EmissionGuarantee>::emit_producing_guarantee
  12:     0x72cccbce40c5 - rustc_middle[4d11da33ed8b74e2]::util::bug::opt_span_bug_fmt::<rustc_span[eee2686a21fb97e]::span_encoding::Span>::{closure#0}
  13:     0x72cccbce413a - rustc_middle[4d11da33ed8b74e2]::ty::context::tls::with_opt::<rustc_middle[4d11da33ed8b74e2]::util::bug::opt_span_bug_fmt<rustc_span[eee2686a21fb97e]::span_encoding::Span>::{closure#0}, !>::{closure#0}
  14:     0x72cccbce083b - rustc_middle[4d11da33ed8b74e2]::ty::context::tls::with_context_opt::<rustc_middle[4d11da33ed8b74e2]::ty::context::tls::with_opt<rustc_middle[4d11da33ed8b74e2]::util::bug::opt_span_bug_fmt<rustc_span[eee2686a21fb97e]::span_encoding::Span>::{closure#0}, !>::{closure#0}, !>
  15:     0x72cccbceabe2 - rustc_middle[4d11da33ed8b74e2]::util::bug::bug_fmt
  16:     0x72ccc8f90449 - rustc_codegen_llvm[fd0152a31030d463]::intrinsic::codegen_autodiff
  17:     0x72ccc90aa0fe - <rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx> as rustc_codegen_ssa[78843d74ba320aa6]::traits::intrinsic::IntrinsicCallBuilderMethods>::codegen_intrinsic_call
  18:     0x72ccc903e8b5 - <rustc_codegen_ssa[78843d74ba320aa6]::mir::FunctionCx<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>>>::codegen_intrinsic_call
  19:     0x72ccc9033c07 - rustc_codegen_ssa[78843d74ba320aa6]::mir::codegen_mir::<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>>
  20:     0x72ccc90cb8c1 - rustc_codegen_ssa[78843d74ba320aa6]::base::codegen_instance::<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>>
  21:     0x72ccc90bda96 - <rustc_middle[4d11da33ed8b74e2]::mir::mono::MonoItem as rustc_codegen_ssa[78843d74ba320aa6]::mono_item::MonoItemExt>::define::<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>>
  22:     0x72ccc90921c1 - rustc_codegen_llvm[fd0152a31030d463]::base::compile_codegen_unit::module_codegen
  23:     0x72ccc9082313 - rustc_codegen_llvm[fd0152a31030d463]::base::compile_codegen_unit
  24:     0x72ccc90cadec - rustc_codegen_ssa[78843d74ba320aa6]::base::codegen_crate::<rustc_codegen_llvm[fd0152a31030d463]::LlvmCodegenBackend>
  25:     0x72ccc8f630c6 - <rustc_codegen_llvm[fd0152a31030d463]::LlvmCodegenBackend as rustc_codegen_ssa[78843d74ba320aa6]::traits::backend::CodegenBackend>::codegen_crate
  26:     0x72ccc8ebf84e - <rustc_session[85ae29ebbb651820]::session::Session>::time::<alloc[d35db71d034d3387]::boxed::Box<dyn core[5b293a957d0a6e5f]::any::Any>, rustc_interface[80f9411e4dfda4b]::passes::start_codegen::{closure#0}>
  27:     0x72ccc8dee328 - rustc_interface[80f9411e4dfda4b]::passes::start_codegen
  28:     0x72ccc8ed10ba - <rustc_interface[80f9411e4dfda4b]::queries::Linker>::codegen_and_build_linker

Copy link
Member

@bjorn3 bjorn3 Nov 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right. Maybe the reason why it works for catch_unwind is that catch_unwind accepts function pointers rather than function items as arguments?

Is it possible to replace fetching #[rustc_autodiff] from the function to be differentiated with a regular argument (or const generic) passed to the autodiff intrinsic, pass a function pointer to the autodiff intrinsic and then get the symbol name from the function pointer value (under the assumption that the function pointer value is a compile-time known value, which should hold true for the #[autodiff_forward] and #[autodiff_backward] expansions) without involving Instance?

Edit: I got the source and differentiated function backwards. #[rustc_autodiff] is applied to the differentiated function which contains the autodiff call, so that one would already work. You would only need to replace the function item with a function pointer and adjust adjust_activity_to_abi to accept the function pointer type rather than a concrete Instance and then use fn_abi_of_fn_ptr instead of fn_abi_of_instance.

Copy link
Member Author

@ZuseZ4 ZuseZ4 Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That might be a reason, though it feels weird to me that a fn item does not count as use, whereas a function pointer does. And yes, they should always be compile-time known for now.
We have to add support for dyn (vtable) in the future, though, which in other languages works by differentiating all candidates, and then at runtime looking up which primal function was given and picking the responding differentiated function. That might require extra work anyway, so I'm also fine with not considering it for now.

For reference, here is the current expansion of placing #[autodiff_forward(...)] on a function _f1.

#[rustc_autodiff]
pub fn _f1(x: f64) -> f64 {
    simple_dep::f(x, x) * f2(x)
}
#[rustc_autodiff(Forward, 1, Dual, Dual)]
pub fn df1_lib(x: f64, bx_0: f64) -> (f64, f64) {
    ::core::intrinsics::autodiff(_f1, df1_lib, (x, bx_0))
}

cc @oli-obk for thoughts and @Sa4dUs who developed our new intrinsic. I think moving it shouldn't be too much work?


if tcx.has_attr(def_id, sym::rustc_intrinsic) {
// Intrinsic fallback bodies are always cross-crate inlineable.
// To ensure that the MIR inliner doesn't cluelessly try to inline fallback
Expand Down
17 changes: 17 additions & 0 deletions tests/ui/autodiff/autodiff_in_rlib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
#![feature(autodiff)]
#![crate_type = "rlib"]
//@ needs-enzyme
//@ compile-flags: -Zautodiff=Enable -C opt-level=3 -Clto=fat
//@ build-fail

// We test that we fail to compile if a user applied an autodiff_ macro in src/lib.rs,
// since autodiff doesn't work in libraries yet. In the past we used to just return zeros in the
// autodiffed functions, which is obviously confusing and wrong, so erroring is an improvement.

use std::autodiff::autodiff_reverse;
//~? ERROR: using the autodiff feature with library builds is not yet supported

#[autodiff_reverse(d_square, Duplicated, Active)]
pub fn square(x: &f64) -> f64 {
*x * *x
}
Loading