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

[rust-log-analyzer]

The job tidy failed! Check out the build log: (web) (plain enhanced) (plain)

Click to see the possible cause of the failure (guessed by this bot)

[TIMING:end] tool::ToolBuild { build_compiler: Compiler { stage: 0, host: x86_64-unknown-linux-gnu, forced_compiler: false }, target: x86_64-unknown-linux-gnu, tool: "tidy", path: "src/tools/tidy", mode: ToolBootstrap, source_type: InTree, extra_features: [], allow_features: "", cargo_args: [], artifact_kind: Binary } -- 11.350
[TIMING:end] tool::Tidy { compiler: Compiler { stage: 0, host: x86_64-unknown-linux-gnu, forced_compiler: false }, target: x86_64-unknown-linux-gnu } -- 0.000
fmt check
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:15:
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
 use rustc_middle::ty::{self, GenericArgsRef, Instance, SimdAlign, Ty, TyCtxt, TypingEnv};
 use rustc_middle::{bug, span_bug};
+use rustc_session::config::CrateType;
 use rustc_span::{Span, Symbol, sym};
 use rustc_symbol_mangling::{mangle_internal_symbol, symbol_name_for_instance_in_crate};
 use rustc_target::callconv::PassMode;
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:21:
 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};
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:1146:
             let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);
         }
     } else {
-        if lto != rustc_session::config::Lto::Fat && !tcx.sess.opts.cg.linker_plugin_lto.enabled(){
+        if lto != rustc_session::config::Lto::Fat && !tcx.sess.opts.cg.linker_plugin_lto.enabled() {
             let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);
         }
     }
Diff in /checkout/compiler/rustc_mir_transform/src/cross_crate_inline.rs:34:
         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) {
+    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;
     }
 
fmt: checked 6549 files

[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?