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

[ZuseZ4]

std::autodiff historically didn't work with rlib builds, I spend a lot of time on it, but even with help from @oli-obk we couldn't find a solution.

I tried this code:

RUSTFLAGS="-Zautodiff=Enable" cargo +enzyme run -r

// src/lib.rs
#![feature(autodiff)]

use std::autodiff::*;

#[inline(never)]
pub fn f2(x: f64) -> f64 {
    x.sin()
}

#[autodiff_forward(df1_lib, Dual, Dual)]
pub fn _f1(x: f64) -> f64 {
    f2(x)
}

and

// src/main.rs
fn main() {
    let enzyme_y1_lib = foo::df1_lib(1.5, 1.0);
    dbg!(enzyme_y1_lib);
}

I expected to see this happen: works

Instead, this happened: crashes.

Meta

rustc --version --verbose:

main

Backtrace

➜  foo git:(main) RUST_BACKTRACE=1 RUSTFLAGS="-Zautodiff=Enable" cargo +enzyme run --bin foo -r 
   Compiling foo v0.1.0 (/home/manuel/prog/foo)

thread 'rustc' (699775) panicked at compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs:224:1:
DefId(20:7 ~ foo[f274]::_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 strip=debuginfo -Z autodiff=Enable

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
error: could not compile `foo` (bin "foo")

Caused by:
  process didn't exit successfully: `/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=178 --crate-type bin --emit=dep-info,link -C opt-level=3 -C lto=fat --check-cfg 'cfg(docsrs,test)' --check-cfg 'cfg(feature, values())' -C metadata=760ac6525d945c95 -C extra-filename=-85b868d1f750b1ac --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-6eb42b5a6b7d3396.rlib --extern num_dual=/home/manuel/prog/foo/target/release/deps/libnum_dual-1b2e4fc55de486a4.rlib -Zautodiff=Enable` (exit status: 101)

cc @oli-obk Do you have an explanation?
@Sa4dUs I'm not sure if your work might have helped here?

[ZuseZ4]

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?

[bjorn3]

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.

[ZuseZ4]

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?

[bjorn3]

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.

[saethlin]

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.

[ZuseZ4]

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:

[saethlin]

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

[ZuseZ4]

@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.