std::autodiff has partial rlib support, DefId does not have a "optimized_mir" #148856

Activity

1. ZuseZ4

   added

   C-bugCategory: This is a bug.Category: This is a bug.

   on Nov 12, 2025
2. rustbot

   added

   needs-triageThis issue may need triage. Remove when done. See docs forge.rust-lang.org/release/issue-triagingThis issue may need triage. Remove when done. See docs forge.rust-lang.org/release/issue-triaging

   on Nov 12, 2025
3. ZuseZ4

   added

   F-autodiff`#![feature(autodiff)]``#![feature(autodiff)]`

   A-MIRArea: Mid-level IR (MIR) - https://blog.rust-lang.org/2016/04/19/MIR.htmlArea: Mid-level IR (MIR) - https://blog.rust-lang.org/2016/04/19/MIR.html

   and removed

   needs-triageThis issue may need triage. Remove when done. See docs forge.rust-lang.org/release/issue-triagingThis issue may need triage. Remove when done. See docs forge.rust-lang.org/release/issue-triaging

   on Nov 12, 2025
4. ZuseZ4 commented on Nov 12, 2025

   

   ZuseZ4

   on Nov 12, 2025

   Last edited by ZuseZ4

   MemberAuthor

   More actions

   As a reminder, rlib did not used to work in any case, now any slight simplification to the code above makes it compile and return the correct result. Historically, any autodiffed function in an rlib would just return 0, since autodiff didn't run, end those function ended up returning 0's as default values.

   My hope is that one of you two can spot the new problem here, now that at least some cases compile correctly? Maybe we can do with unoptimized mir in some cases?
5. bjorn3 commented on Nov 12, 2025

   

   bjorn3

   on Nov 12, 2025

   Member

   More actions

   We don't encode MIR in the crate metadata except for generic and cross-crate inlineable functions. Are you trying to codegen df1_lib and by extension _f1 in the binary crate? That isn't going to work. df1_lib has to be codegened in the rlib.
6. ZuseZ4 commented on Nov 12, 2025

   

   ZuseZ4

   on Nov 12, 2025

   Last edited by ZuseZ4

   MemberAuthor

   More actions

   Enzyme needs the llvm-ir of all functions or types which a function (to be differentiated) accesses (even indirectly/recursively).
   So if in main.rs, we want to differentiate a function a, but a calls a function b from src/lib.rs, and a function c, from a different crate, then we must have the llvm-ir of a,b,c available in a single module (hence our general fat-lto requirement).
   Embed-bc also work (in other languages), but we kept running into other issues when trying to use it in Rust.
   So in our example above, df1_lib might need llvm-ir/llvm-bc from other crates.
   Does this mean autodiff breaks a fundamental rlib requirement and won't be doable without larger changes?

   #142640 is the current intrinsic implementation, which we're using.

   Generally, if we talk about it in terms of different llvm-ir modules, I can somewhat easily describe how to get it to work, so my hope was that at least with embed-bc we can get support?

   Every autodiff intrinsic get's vaguely lowered into the following llvm-ir code (taken from https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/test/Enzyme/ReverseMode/acos.ll)

   declare <some_return_type> @__enzyme_autodiff(...)
   
   define i32 @foo(double %x) nounwind {
   entry:
     %0 = tail call double (double (double)*, ...) @__enzyme_autodiff(double (double)* nonnull @tester, double %x)
     ret double %0
   }
   

   Where foo is the new function generated by our autodiff macro, which we'll want to call in another crate/rlib/etc.

   The called function is the function to which we apply our autodiff macro, in this case tester. It's just a normal, unchanged Rust function in the same module, it just gained one more user. If tester were to call functions from other crates, we'd again need their llvm-ir accessible.

   define double @tester(double %x) {
   entry:
     %0 = tail call fast double @acos(double %x)
     ret double %0
   }
   

   So as long as we keep the llvm-ir of foo and tester around, we could run autodiff at any time. Can that be baked into the rlib handling?
7. bjorn3 commented on Nov 12, 2025

   

   bjorn3

   on Nov 12, 2025

   Member

   More actions

   If you compile with LTO enabled, the LLVM IR for all compiled functions will be preserved in the rlib and loaded when compiling the top level crate and in the case of fat LTO everything will be merged into the main LLVM module.
8. saethlin commented on Nov 13, 2025

   

   saethlin

   on Nov 13, 2025

   Member

   More actions

   Does this mean autodiff breaks a fundamental rlib requirement and won't be doable without larger changes?

   Not really. For example, one of the jobs of cargo-miri is to ensure that every crate in the graph is built with -Zalways-encode-mir. MIR is only encoded for some defs as an optimization.

   I bet if you make cross_crate_inlinable check for autodiff_forward, things would get a lot better. Probably still explode if your autodiffed function makes a cross-crate call though.
9. ZuseZ4 commented on Nov 14, 2025

   

   ZuseZ4

   on Nov 14, 2025

   MemberAuthor

   More actions

   Thank you two. I'll look closer into it next week, once we've some of the open autodiff/gpu PRs. For now just a quick test based on @oli-obk recommendation:

   Printing self inside of compiler/rustc_middle/src/ty/instance.rs in upstream_monomorphization
   only gives the following, it seems that it's not looking there whether _f1 can be provided from an upstream crate.

      Compiling foo v0.1.0 (/home/manuel/prog/foo)
        Running `/home/manuel/.rustup/toolchains/enzyme/bin/rustc --crate-name foo --edition=2024 src/main.rs --error-format=json --json=diagnostic-rendered-ansi,artifacts,future-incompat --diagnostic-width=318 --crate-type bin --emit=dep-info,link -C opt-level=3 -C lto=fat -C codegen-units=1 --check-cfg 'cfg(docsrs,test)' --check-cfg 'cfg(feature, values())' -C metadata=40c8f6dc887788ea -C extra-filename=-65fca01061694a9c --out-dir /home/manuel/prog/foo/target/release/deps -C strip=debuginfo -L dependency=/home/manuel/prog/foo/target/release/deps --extern foo=/home/manuel/prog/foo/target/release/deps/libfoo-5d9ace8c06482a6e.rlib -Zautodiff=Enable -Cembed-bitcode=yes`
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(2:50221 ~ core[a013]::fmt::{impl#79}::fmt),
       ),
       args: [
           '{erased},
           (f64, f64),
       ],
   }
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(2:50221 ~ core[a013]::fmt::{impl#79}::fmt),
       ),
       args: [
           '{erased},
           dyn [Binder { value: Trait(std::fmt::Debug), bound_vars: [] }] + '{erased},
       ],
   }
   
   thread 'rustc' (1012285) panicked at compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs:224:1:
   DefId(20:5 ~ foo[e563]::_f1) does not have a "optimized_mir"
   stack backtrace:
      0: __rustc::rust_begin_unwind
      1: core::panicking::panic_fmt
      2: rustc_metadata::rmeta::decoder::cstore_impl::provide_extern::optimized_mir::{closure#2}
      3: rustc_metadata::rmeta::decoder::cstore_impl::provide_extern::optimized_mir
         [... omitted 2 frames ...]
      4: <rustc_middle::ty::context::TyCtxt>::instance_mir
      5: rustc_monomorphize::collector::items_of_instance
         [... omitted 2 frames ...]
      6: rustc_monomorphize::collector::collect_items_rec
      7: rustc_monomorphize::collector::collect_items_rec
      8: rustc_monomorphize::collector::collect_items_root
      9: <rustc_data_structures::sync::parallel::ParallelGuard>::run::<(), rustc_data_structures::sync::parallel::par_for_each_in<rustc_middle::mir::mono::MonoItem, alloc::vec::Vec<rustc_middle::mir::mono::MonoItem>, rustc_monomorphize::collector::collect_crate_mono_items::{closure#1}::{closure#0}>::{closure#0}::{closure#1}::{closure#0}>
     10: rustc_data_structures::sync::parallel::par_for_each_in::<rustc_middle::mir::mono::MonoItem, alloc::vec::Vec<rustc_middle::mir::mono::MonoItem>, rustc_monomorphize::collector::collect_crate_mono_items::{closure#1}::{closure#0}>
     11: <rustc_session::session::Session>::time::<(), rustc_monomorphize::collector::collect_crate_mono_items::{closure#1}>
     12: rustc_monomorphize::collector::collect_crate_mono_items
     13: rustc_monomorphize::partitioning::collect_and_partition_mono_items
         [... omitted 2 frames ...]
     14: rustc_codegen_ssa::base::codegen_crate::<rustc_codegen_llvm::LlvmCodegenBackend>
     15: <rustc_codegen_llvm::LlvmCodegenBackend as rustc_codegen_ssa::traits::backend::CodegenBackend>::codegen_crate
     16: <rustc_session::session::Session>::time::<alloc::boxed::Box<dyn core::any::Any>, rustc_interface::passes::start_codegen::{closure#0}>
     17: rustc_interface::passes::start_codegen
     18: <rustc_interface::queries::Linker>::codegen_and_build_linker
     19: <std::thread::local::LocalKey<core::cell::Cell<*const ()>>>::with::<rustc_middle::ty::context::tls::enter_context<<rustc_middle::ty::context::GlobalCtxt>::enter<rustc_interface::passes::create_and_enter_global_ctxt<core::option::Option<rustc_interface::queries::Linker>, rustc_driver_impl::run_compiler::{closure#0}::{closure#2}>::{closure#2}::{closure#0}, core::option::Option<rustc_interface::queries::Linker>>::{closure#1}, core::option::Option<rustc_interface::queries::Linker>>::{closure#0}, core::option::Option<rustc_interface::queries::Linker>>
     20: <rustc_middle::ty::context::TyCtxt>::create_global_ctxt::<core::option::Option<rustc_interface::queries::Linker>, rustc_interface::passes::create_and_enter_global_ctxt<core::option::Option<rustc_interface::queries::Linker>, rustc_driver_impl::run_compiler::{closure#0}::{closure#2}>::{closure#2}::{closure#0}>
     21: <rustc_interface::passes::create_and_enter_global_ctxt<core::option::Option<rustc_interface::queries::Linker>, rustc_driver_impl::run_compiler::{closure#0}::{closure#2}>::{closure#2} as core::ops::function::FnOnce<(&rustc_session::session::Session, rustc_middle::ty::context::CurrentGcx, alloc::sync::Arc<rustc_data_structures::jobserver::Proxy>, &std::sync::once_lock::OnceLock<rustc_middle::ty::context::GlobalCtxt>, &rustc_data_structures::sync::worker_local::WorkerLocal<rustc_middle::arena::Arena>, &rustc_data_structures::sync::worker_local::WorkerLocal<rustc_hir::Arena>, rustc_driver_impl::run_compiler::{closure#0}::{closure#2})>>::call_once::{shim:vtable#0}
     22: rustc_interface::passes::create_and_enter_global_ctxt::<core::option::Option<rustc_interface::queries::Linker>, rustc_driver_impl::run_compiler::{closure#0}::{closure#2}>
     23: rustc_span::create_session_globals_then::<(), rustc_interface::util::run_in_thread_with_globals<rustc_interface::util::run_in_thread_pool_with_globals<rustc_interface::interface::run_compiler<(), rustc_driver_impl::run_compiler::{closure#0}>::{closure#1}, ()>::{closure#0}, ()>::{closure#0}::{closure#0}::{closure#0}>
   note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.
   
   error: the compiler unexpectedly panicked. this is a bug.
   
   note: we would appreciate a bug report: https://github.com/rust-lang/rust/issues/new?labels=C-bug%2C+I-ICE%2C+T-compiler&template=ice.md
   
   note: please make sure that you have updated to the latest nightly
   
   note: rustc 1.93.0-nightly (c8bae8c06 2025-11-11) running on x86_64-unknown-linux-gnu
   
   note: compiler flags: --crate-type bin -C opt-level=3 -C lto=fat -C codegen-units=1 -C strip=debuginfo -Z autodiff=Enable -C embed-bitcode=yes
   
   note: some of the compiler flags provided by cargo are hidden
   
   query stack during panic:
   #0 [optimized_mir] optimizing MIR for `foo::_f1`
   #1 [items_of_instance] collecting items used by `foo::_f1`
   #2 [collect_and_partition_mono_items] collect_and_partition_mono_items
   end of query stack
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(1:150 ~ std[70b1]::rt::lang_start::{closure#0}),
       ),
       args: [
           (),
           i8,
           Binder { value: extern "RustCall" fn(()) -> i32, bound_vars: [] },
           (Binder { value: fn(), bound_vars: [] },),
       ],
   }
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(1:150 ~ std[70b1]::rt::lang_start::{closure#0}),
       ),
       args: [
           (),
           i8,
           Binder { value: extern "RustCall" fn(()) -> i32, bound_vars: [] },
           (Binder { value: fn(), bound_vars: [] },),
       ],
   }
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(1:150 ~ std[70b1]::rt::lang_start::{closure#0}),
       ),
       args: [
           (),
           i8,
           Binder { value: extern "RustCall" fn(()) -> i32, bound_vars: [] },
           (Binder { value: fn(), bound_vars: [] },),
       ],
   }
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(1:8238 ~ std[70b1]::sys::backtrace::__rust_begin_short_backtrace),
       ),
       args: [
           Binder { value: fn(), bound_vars: [] },
           (),
       ],
   }
   [compiler/rustc_middle/src/ty/instance.rs:202:9] &self = Instance {
       def: Item(
           DefId(1:6237 ~ std[70b1]::process::{impl#63}::report),
       ),
       args: [],
   }
   error: could not compile `foo` (bin "foo")
   
   Caused by:
   
10. saethlin commented on Nov 14, 2025

    

    saethlin

    on Nov 14, 2025

    Member

    More actions

    Because I'm not in bed yet, I can write you the patch I was imagining:

    diff --git a/compiler/rustc_mir_transform/src/cross_crate_inline.rs b/compiler/rustc_mir_transform/src/cross_crate_inline.rs
    index 7fc9fb9cca2..fba30921fbc 100644
    --- a/compiler/rustc_mir_transform/src/cross_crate_inline.rs
    +++ b/compiler/rustc_mir_transform/src/cross_crate_inline.rs
    @@ -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) {
    +        return true;
    +    }
    +
         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

    ❤️React with ❤️1Reacted by Manuel Drehwald
11. ZuseZ4 commented on Nov 14, 2025

    

    ZuseZ4

    on Nov 14, 2025

    MemberAuthor

    More actions

    @saethlin amazing, that worked. I just had to change it to rustc_autodiff. For context, here is how the lib.rs looks like, expanded:

    #![feature(prelude_import)]
    #![feature(autodiff)]
    #[macro_use]
    extern crate std;
    #[prelude_import]
    use std::prelude::rust_2024::*;
    use std::autodiff::*;
    #[inline(never)]
    pub fn f2(x: f64) -> f64 {
        x.sin()
    }
    #[rustc_autodiff]
    pub fn _f1(x: f64) -> f64 {
        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))
    }

    I'll play with it some more to see where it breaks next, probably the cross-crate case as you mentioned. Still, happy to see we're getting close to being able to differentiate almost everywhere.
12. matthiaskrgr

    added 2 commits that reference this issue on Nov 20, 2025

    Rollup merge of rust-lang#149033 - ZuseZ4:autodiff-rlib, r=bjorn3

    ...

    c2fdcdd

    Rollup merge of rust-lang#149033 - ZuseZ4:autodiff-rlib, r=bjorn3

    ...

    710720a
13. bors

    closed this as completedin 710720aon Nov 21, 2025
14. rust-timer

    added a commit that references this issue on Nov 21, 2025

    Unrolled build for #149033

    ...

    68000b0