Implement autodiff using intrinsics #142640

Sa4dUs · 2025-06-17T20:19:27Z

This PR aims to move autodiff logic to autodiff intrinsic. Allowing us to delete a great part of our frontend code and overall, simplify the compilation pipeline of autodiff functions.

fixes: #143994
fixes: #143543

workingjubilee · 2025-06-17T23:21:28Z

compiler/rustc_codegen_llvm/src/intrinsic.rs

+            _ if tcx.has_attr(def_id, sym::rustc_autodiff) => {
+                // NOTE(Sa4dUs): This is a hacky way to get the autodiff items
+                // so we can focus on the lowering of the intrinsic call
+
+                // `diff_items` is empty even when autodiff is enabled, and if we're here,
+                // it's because some function was marked as intrinsic and had the `rustc_autodiff` attr
+                let diff_items = tcx.collect_and_partition_mono_items(()).autodiff_items;
+
+                // this shouldn't happen?
+                if diff_items.is_empty() {
+                    bug!("no autodiff items found for {def_id:?}");
+                }
+
+                // TODO(Sa4dUs): generate the enzyme call itself, based on the logic in `builder.rs`
+
+                // Just gen the fallback body for now
+                return Err(ty::Instance::new_raw(def_id, instance.args));
+            }


So this is the part you got stuck on?

Kinda. I tried to use the collect_and_partition_mono_items just as a hacky way of getting autodiff_items, so i can focus on the declaration. Once that was working, I would focus on how to get that information in the best way possible.

Just for context, during yesterday's meeting, we came to the conclusion that, for the time being, it's acceptable to copy and paste the collector logic from within the provider's code, just to focus on the declaration, but still nice to have feedback since, in some moment, we'll need to get to AutoDiffItems in a decent way. collect_and_partition_mono_items(()).autodiff_items being [] in autodiff code seems weird, so it would be nice to know why it's happening, or if it is some kind of bug.

workingjubilee · 2025-06-17T23:37:58Z

@Sa4dUs Can you provide the motivation in the PR desc for context?

ZuseZ4 · 2025-07-15T23:10:33Z

tests/codegen/autodiff/batched.rs

-// CHECK-NEXT:   %19 = insertvalue [4 x float] %17, float %18, 3
-// CHECK-NEXT:   ret [4 x float] %19
-// CHECK-NEXT: }
+// CHECK-NEXT:   %4 = fadd fast float %"_2'ipl", %"_2'ipl"


Is this test currently passing on master? Also, can you compare the IR of the two rustc versions just before we run Enzyme? I wonder if we call enzyme differently, or if just llvm changed. (Ignore this if you're still working on the code and are not done yet).

yes, the test with the current fixes is passing at master. and the way we call enzyme is just the same, as i've based my work on the existing function for diff body codegen and removed unnecessary parts.

rustbot · 2025-07-27T16:46:45Z

r? @SparrowLii

rustbot has assigned @SparrowLii.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

rustbot · 2025-07-27T16:46:46Z

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

Some changes occurred in compiler/rustc_codegen_ssa/src/codegen_attrs.rs

cc @jdonszelmann

Some changes occurred to the intrinsics. Make sure the CTFE / Miri interpreter
gets adapted for the changes, if necessary.

cc @rust-lang/miri, @RalfJung, @oli-obk, @lcnr

These commits modify the Cargo.lock file. Unintentional changes to Cargo.lock can be introduced when switching branches and rebasing PRs.

If this was unintentional then you should revert the changes before this PR is merged.
Otherwise, you can ignore this comment.

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

Some changes occurred in compiler/rustc_builtin_macros/src/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs

cc @ZuseZ4

triagebot.toml has been modified, there may have been changes to the review queue.

cc @davidtwco, @wesleywiser

library/core/src/intrinsics/mod.rs

SparrowLii · 2025-07-28T03:13:11Z

r? codegen

Sa4dUs · 2025-08-14T08:20:50Z

actually, i can just remove the whole loop right?

bors · 2025-08-14T14:50:04Z

☔ The latest upstream changes (presumably #145388) made this pull request unmergeable. Please resolve the merge conflicts.

tests/codegen-llvm/autodiff/macro_hygiene.rs

ZuseZ4 · 2025-08-14T23:04:58Z

Having autodiff as an intrinsic is a good improvement which solves a few issues and this PR also unbreaks master. Thanks!
We can investigate the IR regressions which we saw in the codegen tests later, especially since they are likely not caused by this PR but changes in other parts of the compiler or enzyme.

@bors r+

bors · 2025-08-14T23:05:01Z

📌 Commit cdd4118 has been approved by ZuseZ4

It is now in the queue for this repository.

Implement autodiff using intrinsics This PR aims to move autodiff logic to `autodiff` intrinsic. Allowing us to delete a great part of our frontend code and overall, simplify the compilation pipeline of autodiff functions.

Rollup of 22 pull requests Successful merges: - #118087 (Add Ref/RefMut try_map method) - #122661 (Change the desugaring of `assert!` for better error output) - #140740 (Add `-Zindirect-branch-cs-prefix`) - #142640 (Implement autodiff using intrinsics) - #143075 (compiler: Allow `extern "interrupt" fn() -> !`) - #144865 (Fix tail calls to `#[track_caller]` functions) - #144944 (E0793: Clarify that it applies to unions as well) - #144947 (Fix description of unsigned `checked_exact_div`) - #145004 (Couple of minor cleanups) - #145005 (strip prefix of temporary file names when it exceeds filesystem name length limit) - #145012 (Tail call diagnostics to include lifetime info) - #145065 (resolve: Introduce `RibKind::Block`) - #145120 (llvm: Accept new LLVM lifetime format) - #145189 (Weekly `cargo update`) - #145235 (Minor `[const]` tweaks) - #145275 (fix(compiler/rustc_codegen_llvm): apply `target-cpu` attribute) - #145322 (Resolve the prelude import in `build_reduced_graph`) - #145331 (Make std use the edition 2024 prelude) - #145369 (Do not ICE on private type in field of unresolved struct) - #145378 (Add `FnContext` in parser for diagnostic) - #145389 ([rustdoc] Revert "rustdoc search: prefer stable items in search results") - #145392 (coverage: Remove intermediate data structures from mapping creation) r? `@ghost` `@rustbot` modify labels: rollup

Implement autodiff using intrinsics This PR aims to move autodiff logic to `autodiff` intrinsic. Allowing us to delete a great part of our frontend code and overall, simplify the compilation pipeline of autodiff functions.

Rollup of 21 pull requests Successful merges: - #118087 (Add Ref/RefMut try_map method) - #122661 (Change the desugaring of `assert!` for better error output) - #142640 (Implement autodiff using intrinsics) - #143075 (compiler: Allow `extern "interrupt" fn() -> !`) - #144865 (Fix tail calls to `#[track_caller]` functions) - #144944 (E0793: Clarify that it applies to unions as well) - #144947 (Fix description of unsigned `checked_exact_div`) - #145004 (Couple of minor cleanups) - #145005 (strip prefix of temporary file names when it exceeds filesystem name length limit) - #145012 (Tail call diagnostics to include lifetime info) - #145065 (resolve: Introduce `RibKind::Block`) - #145120 (llvm: Accept new LLVM lifetime format) - #145189 (Weekly `cargo update`) - #145235 (Minor `[const]` tweaks) - #145275 (fix(compiler/rustc_codegen_llvm): apply `target-cpu` attribute) - #145322 (Resolve the prelude import in `build_reduced_graph`) - #145331 (Make std use the edition 2024 prelude) - #145369 (Do not ICE on private type in field of unresolved struct) - #145378 (Add `FnContext` in parser for diagnostic) - #145389 ([rustdoc] Revert "rustdoc search: prefer stable items in search results") - #145392 (coverage: Remove intermediate data structures from mapping creation) r? `@ghost` `@rustbot` modify labels: rollup

Rollup merge of #142640 - Sa4dUs:ad-intrinsic, r=ZuseZ4 Implement autodiff using intrinsics This PR aims to move autodiff logic to `autodiff` intrinsic. Allowing us to delete a great part of our frontend code and overall, simplify the compilation pipeline of autodiff functions.

Rollup of 21 pull requests Successful merges: - rust-lang/rust#118087 (Add Ref/RefMut try_map method) - rust-lang/rust#122661 (Change the desugaring of `assert!` for better error output) - rust-lang/rust#142640 (Implement autodiff using intrinsics) - rust-lang/rust#143075 (compiler: Allow `extern "interrupt" fn() -> !`) - rust-lang/rust#144865 (Fix tail calls to `#[track_caller]` functions) - rust-lang/rust#144944 (E0793: Clarify that it applies to unions as well) - rust-lang/rust#144947 (Fix description of unsigned `checked_exact_div`) - rust-lang/rust#145004 (Couple of minor cleanups) - rust-lang/rust#145005 (strip prefix of temporary file names when it exceeds filesystem name length limit) - rust-lang/rust#145012 (Tail call diagnostics to include lifetime info) - rust-lang/rust#145065 (resolve: Introduce `RibKind::Block`) - rust-lang/rust#145120 (llvm: Accept new LLVM lifetime format) - rust-lang/rust#145189 (Weekly `cargo update`) - rust-lang/rust#145235 (Minor `[const]` tweaks) - rust-lang/rust#145275 (fix(compiler/rustc_codegen_llvm): apply `target-cpu` attribute) - rust-lang/rust#145322 (Resolve the prelude import in `build_reduced_graph`) - rust-lang/rust#145331 (Make std use the edition 2024 prelude) - rust-lang/rust#145369 (Do not ICE on private type in field of unresolved struct) - rust-lang/rust#145378 (Add `FnContext` in parser for diagnostic) - rust-lang/rust#145389 ([rustdoc] Revert "rustdoc search: prefer stable items in search results") - rust-lang/rust#145392 (coverage: Remove intermediate data structures from mapping creation) r? `@ghost` `@rustbot` modify labels: rollup

Implement autodiff using intrinsics This PR aims to move autodiff logic to `autodiff` intrinsic. Allowing us to delete a great part of our frontend code and overall, simplify the compilation pipeline of autodiff functions.

rustbot added A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jun 17, 2025

workingjubilee reviewed Jun 17, 2025

View reviewed changes

rustbot added the F-autodiff `#![feature(autodiff)]` label Jun 24, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch 2 times, most recently from bfb72fd to fad0b0c Compare July 1, 2025 10:53

This comment has been minimized.

Sign in to view

rustbot added the A-attributes Area: Attributes (`#[…]`, `#![…]`) label Jul 7, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from a72fd24 to 38b27a2 Compare July 7, 2025 17:30

This comment has been minimized.

Sign in to view

ZuseZ4 reviewed Jul 15, 2025

View reviewed changes

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 97f0985 to 43efce8 Compare July 25, 2025 17:38

rustbot added the A-meta Area: Issues & PRs about the rust-lang/rust repository itself label Jul 25, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 96f554c to d312894 Compare July 25, 2025 18:06

Sa4dUs marked this pull request as ready for review July 27, 2025 16:46

rustbot assigned SparrowLii Jul 27, 2025

rustbot added the S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. label Jul 27, 2025

RalfJung reviewed Jul 27, 2025

View reviewed changes

library/core/src/intrinsics/mod.rs Outdated Show resolved Hide resolved

rustbot assigned saethlin and unassigned SparrowLii Jul 28, 2025

ZuseZ4 assigned ZuseZ4 and unassigned saethlin Jul 28, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from f1c7a40 to c370d96 Compare July 28, 2025 08:52

This comment has been minimized.

Sign in to view

rust-cloud-vms bot force-pushed the ad-intrinsic branch from fe9da82 to 9115b9f Compare August 14, 2025 15:47

Sa4dUs added 4 commits August 14, 2025 16:29

Basic implementation of autodiff intrinsic

5c63104

Complete functionality and general cleanup

250d77e

Remove inlining for autodiff handling

c9c1c17

Remove lto inline logic

e1d79b9

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 9115b9f to 4f1174d Compare August 14, 2025 16:36

ZuseZ4 reviewed Aug 14, 2025

View reviewed changes

tests/codegen-llvm/autodiff/macro_hygiene.rs Outdated Show resolved Hide resolved

Update autodiff tests for the new intrinsics impl

cdd4118

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 4f1174d to cdd4118 Compare August 14, 2025 18:34

bors added S-waiting-on-bors Status: Waiting on bors to run and complete tests. Bors will change the label on completion. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Aug 14, 2025

Zalathar mentioned this pull request Aug 15, 2025

Rollup of 22 pull requests #145418

Closed

Zalathar mentioned this pull request Aug 15, 2025

Rollup of 21 pull requests #145423

Merged

bors merged commit e3a1782 into rust-lang:master Aug 15, 2025
10 checks passed

rustbot added this to the 1.91.0 milestone Aug 15, 2025

This was referenced Aug 15, 2025

autodiff does not have correct hygiene information; “cannot find value” produced if expanded by macro #143543

Closed

Autodiff with -Zbuild-std=core causes “Expected __enzyme function to have 'noinline' before adding 'alwaysinline'” #143994

Closed

github-actions bot mentioned this pull request Aug 18, 2025

Toolchain upgrade to nightly-2025-08-18 flux-rs/flux#1249

Merged

Implement autodiff using intrinsics #142640

Implement autodiff using intrinsics #142640

Uh oh!

Conversation

Sa4dUs commented Jun 17, 2025 • edited by ZuseZ4 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

workingjubilee Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

Sa4dUs Jun 18, 2025

Choose a reason for hiding this comment

Uh oh!

Sa4dUs Jun 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

workingjubilee commented Jun 17, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

ZuseZ4 Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

Sa4dUs Jul 16, 2025

Choose a reason for hiding this comment

Uh oh!

rustbot commented Jul 27, 2025

Uh oh!

rustbot commented Jul 27, 2025

Uh oh!

Uh oh!

SparrowLii commented Jul 28, 2025

Uh oh!

This comment has been minimized.

Sa4dUs commented Aug 14, 2025

Uh oh!

This comment has been minimized.

bors commented Aug 14, 2025

Uh oh!

Uh oh!

ZuseZ4 commented Aug 14, 2025

Uh oh!

bors commented Aug 14, 2025

Uh oh!

Uh oh!

Uh oh!

Sa4dUs commented Jun 17, 2025 •

edited by ZuseZ4

Loading

Sa4dUs Jun 18, 2025 •

edited

Loading