Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle #133429

ZuseZ4 · 2024-11-25T02:19:09Z

This PR should not be merged until the rustc_codegen_llvm part is merged.
I will also alter it a little based on what get's shaved off from the cg_llvm PR,
and address some of the feedback I received in the other PR (including cleanups).

I am putting it already up to

Discuss with @jieyouxu if there is more work needed to add tests to this and
Pray that there is someone reviewing who can tell me why some of my autodiff invocations get lost.

Re 1: My test require fat-lto. I also modify the compilation pipeline. So if there are any other llvm-ir tests in the same compilation unit then I will likely break them. Luckily there are two groups who currently have the same fat-lto requirement for their GPU code which I have for my autodiff code and both groups have some plans to enable support for thin-lto. Once either that work pans out, I'll copy it over for this feature. I will also work on not changing the optimization pipeline for functions not differentiated, but that will require some thoughts and engineering, so I think it would be good to be able to run the autodiff tests isolated from the rest for now. Can you guide me here please?
For context, here are some of my tests in the samples folder: https://github.com/EnzymeAD/rustbook

Re 2: This is a pretty serious issue, since it effectively prevents publishing libraries making use of autodiff: EnzymeAD#173. For some reason my dummy code persists till the end, so the code which calls autodiff, deletes the dummy, and inserts the code to compute the derivative never gets executed. To me it looks like the rustc_autodiff attribute just get's dropped, but I don't know WHY? Any help would be super appreciated, as rustc queries look a bit voodoo to me.

Tracking:

Tracking Issue for autodiff #124509

r? @jieyouxu

ZuseZ4 · 2024-11-25T03:30:47Z

To expand on 2)
Assume you have the following code

#[autodiff(bar, Reverse, ...)]
fn foo(x: f32) -> f32 { x*x }

it will expand to

#[rustc_autodiff]
fn foo(x: f32) -> f32 {x*x}
#[rustc_autodiff(Reverse,...)]
fn bar(x: f32, scalar_factor: f32) -> (f32, f32) {
   // some_dummy_code()
}

Now I have some logic in this PR which picks up the rustc_autodiff attributes and passes them onto the backend, where for every single rustc_autodiff attribute with arguments we pick the function (thus bar here) and replace the dummy code with the right code to return the derivative. So bar would afterwards return (x*x, 2.0 * x). But as mentioned above, once you use autodiff in a library and call it in another module, the dummy code get's executed, as shown in the linked PR.
Any hints would be appreciated.

jieyouxu

I have some interim feedback for this draft PR

compiler/rustc_codegen_ssa/src/back/write.rs

jieyouxu · 2024-11-25T10:18:12Z

compiler/rustc_codegen_ssa/src/assert_module_sources.rs

        let available_cgus =
-            tcx.collect_and_partition_mono_items(()).1.iter().map(|cgu| cgu.name()).collect();
+            tcx.collect_and_partition_mono_items(()).2.iter().map(|cgu| cgu.name()).collect();


Suggestion: maybe explicitly destructure this

let (_, _, available_cgus) = tcx.collect_and_partition_mono_items(()); let available_cgus = available_cgus.iter().map(|cgu| cgu.name()).collect();

I think we should just make it a struct instead of a tuple. This can be done in a separate PR landing before this PR

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

jieyouxu · 2024-11-25T11:56:05Z

compiler/rustc_monomorphize/src/partitioning.rs

+        trace!("AUTODIFF ITEMS EXIST");
+        for item in &mut *autodiff_items {
+            trace!("{}", &item);
+        }


Remark: probably more structured tracing?

jieyouxu · 2024-11-25T11:57:09Z

compiler/rustc_session/src/options.rs

@@ -370,6 +370,7 @@ mod desc {
    pub(crate) const parse_list: &str = "a space-separated list of strings";
    pub(crate) const parse_list_with_polarity: &str =
        "a comma-separated list of strings, with elements beginning with + or -";
+    pub(crate) const parse_autodiff: &str = "various values";


Suggestion: a comma-separated list of autodiff options?

jieyouxu · 2024-11-25T11:58:00Z

compiler/rustc_session/src/options.rs

@@ -996,6 +997,35 @@ mod parse {
        }
    }

+    pub(crate) fn parse_autodiff(slot: &mut Vec<AutoDiff>, v: Option<&str>) -> bool {


Remark: this has no error messages if the autodiff options failed to parse, acceptable for unstable flag but still poor UX, unacceptable when it comes to stabilization time.

It should be handled with a proper description in parse_autodiff above.

jieyouxu · 2024-11-25T11:58:44Z

compiler/rustc_session/src/config.rs

@@ -194,6 +194,39 @@ impl Default for CoverageLevel {
    }
 }

+/// The different settings that the `-Z ad` flag can have.


Problem: please spell the name out, i.e. -Z autodiff instead of using the shorthand ad. Actually I think this is indeed the case, the docs here is just outdated.

jieyouxu · 2024-11-25T12:00:16Z

compiler/rustc_session/src/options.rs

+    autodiff: Vec<crate::config::AutoDiff> = (Vec::new(), parse_autodiff, [TRACKED],
+        "a list autodiff flags to enable (comma separated)"),


Problems: see above,

// If you add a new option, please update: // - compiler/rustc_interface/src/tests.rs // - src/doc/unstable-book/src/compiler-flags

Please update compiler/rustc_interface/src/tests.rs

Please update src/doc/rustc/src/codegen-options/index.md

jieyouxu · 2024-11-25T12:12:03Z

My test require fat-lto.

For codegen/assembly/ui tests you can make it build w/ fat LTO via

//@ compile-flags: -Clto=fat

I also modify the compilation pipeline. So if there are any other llvm-ir tests in the same compilation unit then I will likely break them.

I don't quite understand the implication of this. Is there some small example I can refer to?

I will also work on not changing the optimization pipeline for functions not differentiated, but that will require some thoughts and engineering, so I think it would be good to be able to run the autodiff tests isolated from the rest for now. Can you guide me here please?

Can you elaborate on what test conditions you need? What do you mean exactly by "isolated"? Can you not run autodiff but only if there is autodiff support, or do you mean like don't run by default even if there is autodiff support?

ZuseZ4 · 2024-11-25T12:51:17Z

https://github.com/rust-lang/rust/pull/130060/files#diff-a56b374664e290a55d70fa80e456b6280913830b382b73fb70c4483d3d4cf246
adjust's the llvm opt pipeline (if autodiff is enabled at build time and used).
We have a first opt run which skips the late llvm opts (which tend to increase code size), and runs opt a second time (now with the full pipeline) once autodiff is done. I will manage to not make it optimize unrelated code in the future, but for now it means other functions are now optimized 1.5 times by llvm. And in reality llvm opts don't really run to a fixpoint, so that is highly likely to change the IR.

jieyouxu · 2024-11-27T23:54:13Z

functions are now optimized 1.5 times by llvm. And in reality llvm opts don't really run to a fixpoint, so that is highly likely to change the IR.

This just means that for now, you'll have to gate autodiff-related tests with //@ needs-autodiff or somehow allow compiletest to determine that the llvm used is built with autodiff support and the test is exercising said autodiff support.

Note that this cannot break and should not modify code that does not use autodiff at all, which is indicated by codegen tests that do not use / opt-in to autodiff support.

rustbot · 2024-12-13T00:57:29Z

This PR modifies config.example.toml.

If appropriate, please update CONFIG_CHANGE_HISTORY in src/bootstrap/src/utils/change_tracker.rs.

Some changes occurred in coverage instrumentation.

cc @Zalathar

Some changes occurred in cfg and check-cfg configuration

cc @Urgau

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.

Urgau · 2024-12-16T17:34:27Z

compiler/rustc_session/src/config/cfg.rs

@@ -176,6 +176,8 @@ pub(crate) fn default_configuration(sess: &Session) -> Cfg {
    // NOTE: These insertions should be kept in sync with
    // `CheckCfg::fill_well_known` below.

+    ins_none!(sym::autodiff_fallback);


This is making the autodiff_fallback cfg stably available.

The cfg should at least be gated behind nightly (with tests, documentation, ...), but is your goal to even let users use autodiff_fallback cfg? or is it intended to be purely an internal cfg?

ZuseZ4 · 2025-01-02T17:31:46Z

I rebased now that the other autodiff PR got merged, fixed all conflicts, and got it to compile locally.
I will work through the existing feedback over the next days.

bors · 2025-01-25T02:30:42Z

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

ZuseZ4 · 2025-01-25T21:00:05Z

@oli-obk Just to keep track, so far we have 3/4 things which should be fixed here. Potentially not all in this specific PR, but preferably before enabling it for default nightly builds.

(Most orthogonal) the macro getting lost when used in dependencies.
Performance: Have two opt runs, with AD being applied at the end of the first run. I need to look again at it, last time I wasn't 100% sure how to trigger these two runs after the refactoring.
Build setup: Don't force users to pass RUSTFLAGS="-Z llvm-plugins=/home/manuel/prog/rust-working/build/x86_64-unknown-linux-gnu/enzyme/build/Enzyme/LLVMEnzyme-19.so -C passes=enzyme. I'll probably ask on zulip/bootstrap for help with that.
(of course I'll also need to apply the other code quality recommendations above).

After thinking about it for a bit I'll probably just do 4) for now, and leave 1-3 for a follow-up PR, just for the sake of having a working version upstream. I'll need to talk to jieyouxu to see if we then add test's already here, or in the follow-up PR where I fix 3).

rust-log-analyzer · 2025-01-26T00:01:33Z

The job x86_64-gnu-llvm-18 failed! Check out the build log: (web) (plain)

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

#21 exporting to docker image format
#21 sending tarball 27.5s done
#21 DONE 33.2s
##[endgroup]
Setting extra environment values for docker:  --env ENABLE_GCC_CODEGEN=1 --env GCC_EXEC_PREFIX=/usr/lib/gcc/
[CI_JOB_NAME=x86_64-gnu-llvm-18]
debug: `DISABLE_CI_RUSTC_IF_INCOMPATIBLE` configured.
---
sccache: Starting the server...
##[group]Configure the build
configure: processing command line
configure: 
configure: build.configure-args := ['--build=x86_64-unknown-linux-gnu', '--llvm-root=/usr/lib/llvm-18', '--enable-llvm-link-shared', '--set', 'rust.randomize-layout=true', '--set', 'rust.thin-lto-import-instr-limit=10', '--enable-verbose-configure', '--enable-sccache', '--disable-manage-submodules', '--enable-locked-deps', '--enable-cargo-native-static', '--set', 'rust.codegen-units-std=1', '--set', 'dist.compression-profile=balanced', '--dist-compression-formats=xz', '--set', 'rust.lld=false', '--disable-dist-src', '--release-channel=nightly', '--enable-debug-assertions', '--enable-overflow-checks', '--enable-llvm-assertions', '--set', 'rust.verify-llvm-ir', '--set', 'rust.codegen-backends=llvm,cranelift,gcc', '--set', 'llvm.static-libstdcpp', '--enable-new-symbol-mangling']
configure: target.x86_64-unknown-linux-gnu.llvm-config := /usr/lib/llvm-18/bin/llvm-config
configure: llvm.link-shared     := True
configure: rust.randomize-layout := True
configure: rust.thin-lto-import-instr-limit := 10
---
   Compiling rustc_privacy v0.0.0 (/checkout/compiler/rustc_privacy)
   Compiling rustc_passes v0.0.0 (/checkout/compiler/rustc_passes)
   Compiling rustc_mir_build v0.0.0 (/checkout/compiler/rustc_mir_build)
   Compiling rustc_codegen_llvm v0.0.0 (/checkout/compiler/rustc_codegen_llvm)
error[E0425]: cannot find function, tuple struct or tuple variant `LLVMDumpModule` in module `llvm`
    |
    |
259 |         llvm::LLVMDumpModule(cx.llmod);

error[E0425]: cannot find function, tuple struct or tuple variant `LLVMDumpValue` in module `llvm`
   --> compiler/rustc_codegen_llvm/src/builder/autodiff.rs:261:15
    |
    |
261 |         llvm::LLVMDumpValue(last_inst);
    |               ^^^^^^^^^^^^^ not found in `llvm`

error[E0425]: cannot find function, tuple struct or tuple variant `LLVMDumpValue` in module `llvm`
   --> compiler/rustc_codegen_llvm/src/builder/autodiff.rs:266:15
    |
266 |         llvm::LLVMDumpValue(outer_fn);

error: unused import: `rustc_session::config::Lto`
 --> compiler/rustc_codegen_llvm/src/builder/autodiff.rs:7:5
  |
---

error: unused import: `llvm_optimize`
  --> compiler/rustc_codegen_llvm/src/builder/autodiff.rs:10:36
   |
10 | use crate::back::write::{llvm_err, llvm_optimize};

error: unused variable: `config`
   --> compiler/rustc_codegen_llvm/src/builder/autodiff.rs:288:5
    |

rustbot assigned jieyouxu Nov 25, 2024

traviscross mentioned this pull request Nov 4, 2024

Tracking Issue for autodiff #124509

Open

7 tasks

jieyouxu requested changes Nov 25, 2024

View reviewed changes

rustbot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Nov 25, 2024

ZuseZ4 changed the title ~~upstream rustc_codegen_ssa/rustc_middle changes for enzyme/autodiff~~ Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle Nov 26, 2024

ZuseZ4 mentioned this pull request Nov 27, 2024

Expose experimental LLVM features for automatic differentiation and GPU offloading rust-lang/rust-project-goals#109

Open

9 tasks

oli-obk self-assigned this Dec 6, 2024

ZuseZ4 marked this pull request as ready for review December 13, 2024 00:57

Urgau reviewed Dec 16, 2024

View reviewed changes

ZuseZ4 force-pushed the autodiff-middle branch from b30f1d7 to 2ad340e Compare January 2, 2025 17:17

This comment has been minimized.

Sign in to view

upstream rustc_codegen_ssa/rustc_middle changes for enzyme/autodiff

c4af0ba

ZuseZ4 force-pushed the autodiff-middle branch from 2ad340e to c4af0ba Compare January 25, 2025 03:13

This comment has been minimized.

Sign in to view

ZuseZ4 added 2 commits January 25, 2025 18:51

cleanup and temporary dbg/print statements

f1bbeb4

temporary don't link against enzyme to try -Z llvm-plugins

5858f5a

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle #133429

Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle #133429

ZuseZ4 commented Nov 25, 2024 •

edited

Loading

ZuseZ4 commented Nov 25, 2024

jieyouxu left a comment

jieyouxu Nov 25, 2024

oli-obk Dec 14, 2024

jieyouxu Nov 25, 2024

jieyouxu Nov 25, 2024

jieyouxu Nov 25, 2024

Urgau Dec 16, 2024

jieyouxu Nov 25, 2024

jieyouxu Nov 25, 2024

jieyouxu commented Nov 25, 2024 •

edited

Loading

ZuseZ4 commented Nov 25, 2024

jieyouxu commented Nov 27, 2024

rustbot commented Dec 13, 2024

Urgau Dec 16, 2024 •

edited

Loading

This comment has been minimized.

ZuseZ4 commented Jan 2, 2025

bors commented Jan 25, 2025

This comment has been minimized.

ZuseZ4 commented Jan 25, 2025 •

edited

Loading

rust-log-analyzer commented Jan 26, 2025

		autodiff: Vec<crate::config::AutoDiff> = (Vec::new(), parse_autodiff, [TRACKED],
		"a list autodiff flags to enable (comma separated)"),

Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle #133429

Are you sure you want to change the base?

Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle #133429

Conversation

ZuseZ4 commented Nov 25, 2024 • edited Loading

ZuseZ4 commented Nov 25, 2024

jieyouxu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jieyouxu commented Nov 25, 2024 • edited Loading

ZuseZ4 commented Nov 25, 2024

jieyouxu commented Nov 27, 2024

rustbot commented Dec 13, 2024

Urgau Dec 16, 2024 • edited Loading

Choose a reason for hiding this comment

This comment has been minimized.

ZuseZ4 commented Jan 2, 2025

bors commented Jan 25, 2025

This comment has been minimized.

ZuseZ4 commented Jan 25, 2025 • edited Loading

rust-log-analyzer commented Jan 26, 2025

ZuseZ4 commented Nov 25, 2024 •

edited

Loading

jieyouxu commented Nov 25, 2024 •

edited

Loading

Urgau Dec 16, 2024 •

edited

Loading

ZuseZ4 commented Jan 25, 2025 •

edited

Loading