Autodiff Upstreaming - single commit #129175
Draft
+4,222
−56
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
rustbot has assigned @michaelwoerister. Use |
|
rustbot
added
S-waiting-on-review
7 tasks
|
This is cool. Over the medium to long-term, though, I think we should work to make the interface in the compiler agnostic to whatever backend/tool is used. I know that's a big ask, and I think it's valuable to land this work in the meantime. But, I don't think we want to commit ourselves to Enzyme or to e.g. supporting autodiff in bootstrap. |
michaelwoerister
removed
the
S-waiting-on-review
|
Agree, both of my Autodfiff and GPU Offloading work could be handled in a crate, by providing better access to compiler internals. Java's reflection proposal even includes Autodiff as example (though they still don't handle all necessary features), and in Julia land, Enzyme.jl is a package that just has a dependency on their compiler package GPUCompiler.jl (terrible name, independent of GPUs these days). The equivalent of my rustc-gpu work is KA.jl (KernelAbstractions) and also works as a normal package for Julia. I talked with Niko about what compiler internals I would need access too. |
jieyouxu
added
the
S-experimental
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Sep 6, 2024
…san68 Autodiff Upstreaming - enzyme backend Tracking issue: rust-lang#124509 Part of rust-lang#129175 This PR should allow building Enzyme from source on Tier 1 targets (when also building LLVM), except MSVC. It's only a small fraction (~200 lines) of the whole upstream PR, but due to bootstrapping and the number of configurations in which rustc can be build I assume that this will be the hardest to merge, so I'm starting with it. Happy to hear what changes are required to be able to upstream this code. **Content:** It contains a new configure flag `--enable-llvm-enzyme`, and will build the new Enzyme submodule when it is set. **Discussion:** Apparently Rust CI isn't able to clone repositories outside the rust-lang org? At least I'm seeing this error in CI: ``` git@github.com: Permission denied (publickey). fatal: Could not read from remote repository. Please make sure you have the correct access rights and the repository exists. ``` Does that mean we would need to mirror github.com/EnzymeAD/Enzyme in rust-lang, until LLVM upgrades Enzyme from an Incubator project to something that ships as part of the monorepo? Tracking: - rust-lang#124509
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Sep 6, 2024
…san68 Autodiff Upstreaming - enzyme backend Tracking issue: rust-lang#124509 Part of rust-lang#129175 This PR should allow building Enzyme from source on Tier 1 targets (when also building LLVM), except MSVC. It's only a small fraction (~200 lines) of the whole upstream PR, but due to bootstrapping and the number of configurations in which rustc can be build I assume that this will be the hardest to merge, so I'm starting with it. Happy to hear what changes are required to be able to upstream this code. **Content:** It contains a new configure flag `--enable-llvm-enzyme`, and will build the new Enzyme submodule when it is set. **Discussion:** Apparently Rust CI isn't able to clone repositories outside the rust-lang org? At least I'm seeing this error in CI: ``` git@github.com: Permission denied (publickey). fatal: Could not read from remote repository. Please make sure you have the correct access rights and the repository exists. ``` Does that mean we would need to mirror github.com/EnzymeAD/Enzyme in rust-lang, until LLVM upgrades Enzyme from an Incubator project to something that ships as part of the monorepo? Tracking: - rust-lang#124509
This was referenced Sep 7, 2024
lnicola
pushed a commit
to lnicola/rust-analyzer
that referenced
this pull request
Sep 25, 2024
Autodiff Upstreaming - enzyme backend Tracking issue: rust-lang/rust#124509 Part of rust-lang/rust#129175 This PR should allow building Enzyme from source on Tier 1 targets (when also building LLVM), except MSVC. It's only a small fraction (~200 lines) of the whole upstream PR, but due to bootstrapping and the number of configurations in which rustc can be build I assume that this will be the hardest to merge, so I'm starting with it. Happy to hear what changes are required to be able to upstream this code. **Content:** It contains a new configure flag `--enable-llvm-enzyme`, and will build the new Enzyme submodule when it is set. **Discussion:** Apparently Rust CI isn't able to clone repositories outside the rust-lang org? At least I'm seeing this error in CI: ``` git@github.com: Permission denied (publickey). fatal: Could not read from remote repository. Please make sure you have the correct access rights and the repository exists. ``` Does that mean we would need to mirror github.com/EnzymeAD/Enzyme in rust-lang, until LLVM upgrades Enzyme from an Incubator project to something that ships as part of the monorepo? Tracking: - rust-lang/rust#124509
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Sep 30, 2024
add has_enzyme/needs-enzyme to the test infra This unblocks merging the Enzyme / Autodiff frontend. For the full implementation, see: rust-lang#129175 We don't want to run tests that require Enzyme / Autodiff support when we build rustc without the required features. It correctly filtered out a test which started with `//@ needs-enzyme`. ``` running 80 tests i............................................................................... test result: ok. 79 passed; 0 failed; 1 ignored; 0 measured; 0 filtered out; finished in 380.41ms ``` Tracking: - rust-lang#124509 r? jieyouxu
RalfJung
pushed a commit
to RalfJung/miri
that referenced
this pull request
Oct 3, 2024
add has_enzyme/needs-enzyme to the test infra This unblocks merging the Enzyme / Autodiff frontend. For the full implementation, see: rust-lang/rust#129175 We don't want to run tests that require Enzyme / Autodiff support when we build rustc without the required features. It correctly filtered out a test which started with `//@ needs-enzyme`. ``` running 80 tests i............................................................................... test result: ok. 79 passed; 0 failed; 1 ignored; 0 measured; 0 filtered out; finished in 380.41ms ``` Tracking: - rust-lang/rust#124509 r? jieyouxu
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Oct 13, 2024
Autodiff Upstreaming - enzyme frontend This is an upstream PR for the `autodiff` rustc_builtin_macro that is part of the autodiff feature. For the full implementation, see: rust-lang#129175 **Content:** It contains a new `#[autodiff(<args>)]` rustc_builtin_macro, as well as a `#[rustc_autodiff]` builtin attribute. The autodiff macro is applied on function `f` and will expand to a second function `df` (name given by user). It will add a dummy body to `df` to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level, we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in `compiler/rustc_builtin_macros/src/autodiff.rs`, which expand the macro. Nothing except expansion is implemented for now. I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging. **Dummy function Body:** The first line is an `inline_asm` nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If `f` gets inlined too early, we can't pass it to enzyme and thus can't differentiate it. If `df` gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body. **Motivation:** The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend. This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability. As a simple example of what this will do, we can see this expansion: From: ``` #[autodiff(df, Reverse, Duplicated, Const, Active)] pub fn f1(x: &[f64], y: f64) -> f64 { unimplemented!() } ``` to ``` #[rustc_autodiff] #[inline(never)] pub fn f1(x: &[f64], y: f64) -> f64 { ::core::panicking::panic("not implemented") } #[rustc_autodiff(Reverse, Duplicated, Const, Active,)] #[inline(never)] pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 { unsafe { asm!("NOP"); }; ::core::hint::black_box(f1(x, y)); ::core::hint::black_box((dx, dret)); ::core::hint::black_box(f1(x, y)) } ``` I will add a few more tests once I figured out why rustc rebuilds every time I touch a test. Tracking: - rust-lang#124509
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Oct 14, 2024
Autodiff Upstreaming - enzyme frontend This is an upstream PR for the `autodiff` rustc_builtin_macro that is part of the autodiff feature. For the full implementation, see: rust-lang#129175 **Content:** It contains a new `#[autodiff(<args>)]` rustc_builtin_macro, as well as a `#[rustc_autodiff]` builtin attribute. The autodiff macro is applied on function `f` and will expand to a second function `df` (name given by user). It will add a dummy body to `df` to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level, we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in `compiler/rustc_builtin_macros/src/autodiff.rs`, which expand the macro. Nothing except expansion is implemented for now. I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging. **Dummy function Body:** The first line is an `inline_asm` nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If `f` gets inlined too early, we can't pass it to enzyme and thus can't differentiate it. If `df` gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body. **Motivation:** The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend. This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability. As a simple example of what this will do, we can see this expansion: From: ``` #[autodiff(df, Reverse, Duplicated, Const, Active)] pub fn f1(x: &[f64], y: f64) -> f64 { unimplemented!() } ``` to ``` #[rustc_autodiff] #[inline(never)] pub fn f1(x: &[f64], y: f64) -> f64 { ::core::panicking::panic("not implemented") } #[rustc_autodiff(Reverse, Duplicated, Const, Active,)] #[inline(never)] pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 { unsafe { asm!("NOP"); }; ::core::hint::black_box(f1(x, y)); ::core::hint::black_box((dx, dret)); ::core::hint::black_box(f1(x, y)) } ``` I will add a few more tests once I figured out why rustc rebuilds every time I touch a test. Tracking: - rust-lang#124509
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Oct 14, 2024
Autodiff Upstreaming - enzyme frontend This is an upstream PR for the `autodiff` rustc_builtin_macro that is part of the autodiff feature. For the full implementation, see: rust-lang#129175 **Content:** It contains a new `#[autodiff(<args>)]` rustc_builtin_macro, as well as a `#[rustc_autodiff]` builtin attribute. The autodiff macro is applied on function `f` and will expand to a second function `df` (name given by user). It will add a dummy body to `df` to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level, we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in `compiler/rustc_builtin_macros/src/autodiff.rs`, which expand the macro. Nothing except expansion is implemented for now. I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging. **Dummy function Body:** The first line is an `inline_asm` nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If `f` gets inlined too early, we can't pass it to enzyme and thus can't differentiate it. If `df` gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body. **Motivation:** The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend. This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability. As a simple example of what this will do, we can see this expansion: From: ``` #[autodiff(df, Reverse, Duplicated, Const, Active)] pub fn f1(x: &[f64], y: f64) -> f64 { unimplemented!() } ``` to ``` #[rustc_autodiff] #[inline(never)] pub fn f1(x: &[f64], y: f64) -> f64 { ::core::panicking::panic("not implemented") } #[rustc_autodiff(Reverse, Duplicated, Const, Active,)] #[inline(never)] pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 { unsafe { asm!("NOP"); }; ::core::hint::black_box(f1(x, y)); ::core::hint::black_box((dx, dret)); ::core::hint::black_box(f1(x, y)) } ``` I will add a few more tests once I figured out why rustc rebuilds every time I touch a test. Tracking: - rust-lang#124509
bors
added a commit
to rust-lang-ci/rust
that referenced
this pull request
Oct 15, 2024
Autodiff Upstreaming - enzyme frontend This is an upstream PR for the `autodiff` rustc_builtin_macro that is part of the autodiff feature. For the full implementation, see: rust-lang#129175 **Content:** It contains a new `#[autodiff(<args>)]` rustc_builtin_macro, as well as a `#[rustc_autodiff]` builtin attribute. The autodiff macro is applied on function `f` and will expand to a second function `df` (name given by user). It will add a dummy body to `df` to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level, we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in `compiler/rustc_builtin_macros/src/autodiff.rs`, which expand the macro. Nothing except expansion is implemented for now. I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging. **Dummy function Body:** The first line is an `inline_asm` nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If `f` gets inlined too early, we can't pass it to enzyme and thus can't differentiate it. If `df` gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body. **Motivation:** The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend. This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability. As a simple example of what this will do, we can see this expansion: From: ``` #[autodiff(df, Reverse, Duplicated, Const, Active)] pub fn f1(x: &[f64], y: f64) -> f64 { unimplemented!() } ``` to ``` #[rustc_autodiff] #[inline(never)] pub fn f1(x: &[f64], y: f64) -> f64 { ::core::panicking::panic("not implemented") } #[rustc_autodiff(Reverse, Duplicated, Const, Active,)] #[inline(never)] pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 { unsafe { asm!("NOP"); }; ::core::hint::black_box(f1(x, y)); ::core::hint::black_box((dx, dret)); ::core::hint::black_box(f1(x, y)) } ``` I will add a few more tests once I figured out why rustc rebuilds every time I touch a test. Tracking: - rust-lang#124509 try-job: dist-x86_64-msvc
github-actions bot
pushed a commit
to rust-lang/miri
that referenced
this pull request
Oct 17, 2024
Autodiff Upstreaming - enzyme frontend This is an upstream PR for the `autodiff` rustc_builtin_macro that is part of the autodiff feature. For the full implementation, see: rust-lang/rust#129175 **Content:** It contains a new `#[autodiff(<args>)]` rustc_builtin_macro, as well as a `#[rustc_autodiff]` builtin attribute. The autodiff macro is applied on function `f` and will expand to a second function `df` (name given by user). It will add a dummy body to `df` to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level, we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in `compiler/rustc_builtin_macros/src/autodiff.rs`, which expand the macro. Nothing except expansion is implemented for now. I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging. **Dummy function Body:** The first line is an `inline_asm` nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If `f` gets inlined too early, we can't pass it to enzyme and thus can't differentiate it. If `df` gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body. **Motivation:** The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend. This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability. As a simple example of what this will do, we can see this expansion: From: ``` #[autodiff(df, Reverse, Duplicated, Const, Active)] pub fn f1(x: &[f64], y: f64) -> f64 { unimplemented!() } ``` to ``` #[rustc_autodiff] #[inline(never)] pub fn f1(x: &[f64], y: f64) -> f64 { ::core::panicking::panic("not implemented") } #[rustc_autodiff(Reverse, Duplicated, Const, Active,)] #[inline(never)] pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 { unsafe { asm!("NOP"); }; ::core::hint::black_box(f1(x, y)); ::core::hint::black_box((dx, dret)); ::core::hint::black_box(f1(x, y)) } ``` I will add a few more tests once I figured out why rustc rebuilds every time I touch a test. Tracking: - rust-lang/rust#124509 try-job: dist-x86_64-msvc
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Tracking issue:
This PR/Commit is just here to show the whole picture. I'll also push individual PRs that are easier to review, starting with the backend which due to bootstrapping is the part I am the least confident about.
The documentation currently is hosted at https://enzyme.mit.edu/index.fcgi/rust/. Please see especially the Installation chapter for build instructions (For rustc devs these shouldn't be surprising).
The tests are currently hosted at https://github.com/EnzymeAD/rustbook. To simplify reviewing I would suggest we first focus on upstreaming the code, and then I'll move over the tests once that's done.
A few of the lines like were contributed by other people who helped me with my fork. Since this work touches multiple parts of rustc, it was easier to reimplement than rebase this once in a while. Once individual PRs are approved I'll look up the history and restore author information where needed.
The first individual PR for the backend (builds on it's own) is here: #129176
The frontend PR is here: #129458