Implement autodiff using intrinsics #142640

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

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

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

Note(Sa4dUs): `cg/generic.rs` test is passing with some tweaks

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

Sa4dUs added 29 commits August 1, 2025 06:30

Feature intrinsics in cg tests

63a0cfd

Remove sret logic

02395d7

Move logic to a dedicated enzyme_autodiff intrinsic

2fc199d

Remove attr checking from hir_analysis

9296843

FIx generics error when passing fn as param to intrinsic

d5611e4

Use Instance::new_raw instead of Instance::mono

7415234

Note(Sa4dUs): `cg/generic.rs` test is passing with some tweaks

Hacky fix for issues at trait calls

3f64b82

Fix how fns where being retrieved at intrinsic cg

1b7a540

Use Self instead of Foo placeholder

cde3c2f

Remove unused code

5b3909e

Remove primal call and collect it in mono instead

9ce6d16

Update codegen tests

4e2a005

Handle slices when extracting args from tuple

c55e7ca

Do not depend on mono anymore

70503d7

Get args from tuple using fnabi and minor fixes

09bb20e

Remove dead code

4108677

Reintroduce autodiff enable and lto errors

e347e00

Minor fixes after rebase

ba32958

Add rest of test fixes

ce2da2b

Remove cfg enzyme for ci

fde1300

FIx cg_gcc

886559b

Macro expansion cleanup

e8e0081

Remove dead code

3330692

Add enzyme_autodiff doc comment

c56d57f

Add expansion example to intrinsic docs

7f0aca5

Better error handling

5673cb8

Allow core_intrinsics when autodiff is enabled

68df065

Remove autodiff limitations subsection

b79b337

Remove inlining for autodiff handling

b7d7b73

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 577cb5d to b7d7b73 Compare August 1, 2025 09:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Implement autodiff using intrinsics #142640

Implement autodiff using intrinsics #142640

Uh oh!

Sa4dUs commented Jun 17, 2025 •

edited

Loading

Uh oh!

workingjubilee Jun 17, 2025

Uh oh!

Sa4dUs Jun 18, 2025

Uh oh!

Sa4dUs Jun 18, 2025 •

edited

Loading

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

Uh oh!

Sa4dUs Jul 16, 2025

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.

Uh oh!

Implement autodiff using intrinsics #142640

Are you sure you want to change the base?

Implement autodiff using intrinsics #142640

Uh oh!

Conversation

Sa4dUs commented Jun 17, 2025 • edited 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.

Uh oh!

Sa4dUs commented Jun 17, 2025 •

edited

Loading

Sa4dUs Jun 18, 2025 •

edited

Loading