autodiff rlib handling

[ZuseZ4]

As I learned recently, we now apparently support rlib builds already in some cases.
With the last hint from saethlin this seems to now cover all cases. To be sure I'll add a few more testcases before I mark it as ready.

Once this PR lands, we should to the best of my knowledge, support autodiff in almost code locations, only vtable/dyn ptr remain unsupported for now.

r? ghost

closes: #148856
closes: #137520

[bjorn3]

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.

[ZuseZ4]

@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

[bjorn3]

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.

[ZuseZ4]

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?

[Sa4dUs]

it shouldn't, basically replacing the current logic in cg_llvm::intrinsic::codegen_autodiff that starts from the DefId and generic args to obtain the value

[ZuseZ4]

We just got approval for enabling autodiff on nightly, and I care about rlib support, but we also still have some bigger open issues (e.g. re-enable macos support). Therefore, I'd prefer to have the change to fnc pointers in another PR, especially since we aren't 100% sure that it will work, and first just get this to work here. I added two fixme's, and I can open an E-easy issue for it, maybe a new contributor can chime in here. The rmake test I added at least should keep working with both approaches. On the current main branch, it panics.

cc @saethlin
r? @bjorn3 I hope that's fine for you?

[rustbot]

Some changes occurred to MIR optimizations

cc @rust-lang/wg-mir-opt

[saethlin]

Thanks for the cc. The implementation looks well factored so that it does not invade the general compiler logic.

[bjorn3]

@bors r+

[bors]

📌 Commit 2d1fc63 has been approved by bjorn3

It is now in the queue for this repository.