Repository navigation
Finish the std::offload module #109
Description
| Metadata | |
|---|---|
| Point of contact | @ZuseZ4 |
| Team(s) | compiler, lang |
| Goal document | 2025h2/finishing-gpu-offload |
Summary
This project aims to get the std::offload module finished. We currently have the ability to automatically move data to and from a GPU, but we can not launch kernels yet. We will add this ability during the next project goal period.
Tasks and status
Note: we have updated the body to match the 2025h2 goal. Your original text is preserved below.
Details
| Metadata | |
|---|---|
| Point of contact | @ZuseZ4 |
| Team(s) | compiler, lang |
| Goal document | 2025h1/GPU-Offload |
Summary
Expose experimental LLVM features for GPU offloading and allow combining it with the std::autodiff feature.
Tasks and status
Metadata
Metadata
Assignees
Type
Projects
- Status
Activity
nikomatsakis commentedon Jul 25, 2024
This issue is intended for status updates only.
For general questions or comments, please contact the owner(s) directly.
ZuseZ4 commentedon Aug 24, 2024
During the first month, I focused on automatic differentiation. I cleaned up my rustc fork and made my first two upstreaming PRs for the frontend and backend. Once they are merged, I will continue with posting PRs for the remaining middle-end. While waiting for reviews, I have been improving the docs a bit, mainly the pages about debugging Enzyme crashes. I am especially proud that due to those docs we recently got our first enzyme core issue with a full LLVM-IR reproducer from a Rust dev, even though the developer reporting that issue had no previous compiler/LLVM experience. Such detailed issues make fixing bugs for Enzyme core much easier.
On the GPU side, I mainly have to thank nikic, who reliably updates the LLVM backend of Rustc every few weeks or months. Thanks to his latest update rustc now supports a sufficiently new LLVM which ships most of the GPU/Offloading work that I want to expose on the Rust side. Once my first two Autodiff patches have settled, I'll look a bit more into setting up documentation for the GPU feature.
ZuseZ4 commentedon Sep 12, 2024
During the last three weeks, my first autodiff PR for the backend, which includes the enzyme submodule and 13 additional files, got merged! I also got a ton of feedback from reviewers, especially for my frontend PR (thanks to jieyouxu). Now that the backend is merged, I did put up my third PR, covering the changes I made to
rustc_codegen_llvm. I am currently at RustConf, so I won't be able to address much of the feedback this week, but I am happy to talk to everyone also visiting and will try to get both PRs ready to merge in the next week.Once we then have the two open PRs merged, we should have my changes to ~55/85 files upstream, so we're making good progress.
On the GPU side again not many updates due to my current autodiff focus, but thanks to another llvm submodule update we can now use some nicer APIs for our development in rustc, which recently got merged into LLVM.
ZuseZ4 commentedon Sep 16, 2024
And as another short update, my talk "When unsafe code is slow - Automatic Differentiation in Rust" got accepted as a techtalk for the LLVM dev mtg. There I'll present a lot of benchmarks and some analysis comparing Rust-Enzyme with the C++ frontend of Enzyme, and show one application which we had to port from Python/JAX to Rust/Enzyme.
The full program of the dev meeting is available here.
For that, I spent some time trying to fix the benchmark infrastructure in Enzyme core, to make sure everyone can reproduce our benchmarks.
ZuseZ4 commentedon Sep 30, 2024
Thanks to some support from the bootstrap team, dist builds with autodiff support enabled now work.
That allowed us to add Rust to our autodiff fork of the compiler explorer: https://enzyme.mit.edu/explorer/
Unfortunately, we still have some dist issues about finding
stdin the compiler explorer build, so help here would be appreciated.Other than that, this morning my PR to add Enzyme/autodiff support to the test infra got merged: rust-lang/rust#131044
This should allow to add this larger frontend PR to the merge queue later today: rust-lang/rust#129458
ZuseZ4 commentedon Oct 22, 2024
I've been travelling a lot for the last two weeks, but hope to be able to get back to work next Monday. Since the last update we got:
Help Wanted:
I would appreciate if someone could look into fixing our Rust dist build used in the Enzyme Compiler explorer. I have spend quite a few hours trying different configurations, but have been unable to get rid of the error
Any help would be appreciated, I can share more information if someone has time to investigate further.
ZuseZ4 commentedon Nov 27, 2024
The re-design of our autodiff middle/backend which I described in the last update was implemented. This led to a reduction from 2.5k to 1.1k LoC remaining to be upstreamed. I split the code up in two PRs (Autodiff Upstreaming - rustc_codegen_ssa, rustc_middle rust#133429 and Autodiff Upstreaming - rustc_codegen_llvm changes rust#130060). Both are now small enough to be reviewed, and got their first round of feedback, so they will hopefully land at the beginning of December. Afterwards, everything needed to run autodiff will be available on nightly (at least as MVP), so we can discuss to build and ship it by default.
The talks which I gave at LLVM Dev caused some interesting follow-up discussions. Most companies still use Rust "only" for classical SWE, but given that it's getting more and more common I also see more interest outside of Academia in using it for (scientific) computing, ML, HPC, etc, which I find exciting. I also got some offers from people in industry to help with the GPU work.
The preprint of the first paper making use of
std::autodiffis available on Arxive! https://arxiv.org/abs/2411.17011v1The code is also available here: https://github.com/ChemAI-Lab/molpipx/. It includes both Python/JAX and Rust implementations, because JAX Jitting times are unbearably slow here. In certain configuration it takes more than a day to JIT, but only 30 minutes to compile in Rust.
Once autodiff is upstreamed (especially including some small follow-up PRs which are needed to achieve the best performance), I will also publish some very promising runtime results that we have on a larger set of benchmarks.
Last month I asked for help with our compiler explorer, and I'm happy that we indeed got the needed support since, thank you! fix rustc installation EnzymeAD/enzyme-explorer#15 Our compiler explorer for Rust with
std::autodiffsupport is now available under https://enzyme.mit.edu/explorer/ (just select Rust).ZuseZ4 commentedon Jan 3, 2025
Happy New Year everyone! After a few more rounds of feedback, the next autodiff PR recently got merged: rust-lang/rust#130060
With that, I only have one last PR open to have a fully working autodiff MVP upstream. A few features had to be removed during upstreaming to simplify the reviewing process, but they should be easier to bring back as single PRs.
Beginning next week, I will also work on an MVP for the
batchingfeature of LLVM/Enzyme, which enables some AoS and SoA vectorization. It mostly re-uses the existing autodiff infrastructure, so I expect the PRs for it to be much smaller.On the GPU side, there has been a recent push by another developer to add a new AMD GPU target to the Rust compiler. This is something that I would have needed for the llvm offload project anyway, so I'm very happy to see movement here: rust-lang/compiler-team#823
21 remaining items
nikomatsakis commentedon Sep 16, 2025
This is a continuing project goal, and the updates below this comment will be for the new period 2025h2
nikomatsakis commentedon Sep 16, 2025
This is a continuing project goal, and the updates below this comment will be for the new period 2025h2
nikomatsakis commentedon Sep 16, 2025
This is a continuing project goal, and the updates below this comment will be for the new period 2025h2
nikomatsakis commentedon Sep 16, 2025
This is a continuing project goal, and the updates below this comment will be for the new period 2025h2
ZuseZ4 commentedon Oct 22, 2025
A longer update of the changes over the fall. We had two gsoc contributors and a lot of smaller improvements for std::autodiff. The first two improvements were already mentioned as draft PRs in the previous update, but got merged since. I also upstreamed more std::offload changes.
All-in-all I spend a lot more time on infra (dlopen, cmake, download-llvm-ci, ...) then I'd like, but on the happy side there are only so many features left that I want to support here so there is an end in sight.
I am also about to give a tech-talk at the upcoming LLVM dev meeting about safe GPU programming in Rust.
ZuseZ4 commentedon Nov 19, 2025
Automatic Differentiation
Time for the next update. By now, we've had
std::autodifffor around a year in upstream rustc, but not in nightly. In order to get some more test users, I asked the infra team to re-evaluate just shippingautodiffas-is. This means that for the moment, we will increase the binary size of rustc by ~5%, even for nightly users who don't use this feature. We still have an open issue to avoid this overhead by using dlopen, please reach out if you have time to help. Thankfully, my request was accepted, so I spent most of my time lately preparing that release.main.rsfile, but also in dependencies (eitherlib.rs, or even rely on crates that use autodiff). With the help of @saethlin I figured out how to get the remaining cases covered, hopefully the PR will land soon.lto=fat: Error if an autodiff user does not set lto=fat rust#148855I have the following elements on my TODO list as part shipping AD on nightly
GPU offload
AMD Ryzen 7 PRO 7840Uunfortunately turned out to be not supported by AMD drivers.All in all, I think we made great progress over the last month, and it's motivating that we finally have no blockers left for flipping the
llvm.enzymeconfig on our nightly builds.ZuseZ4 commentedon Dec 2, 2025
It's only been two weeks, but we got a good number of updates, so I already wanted to share them.
autodiff
dynarguments that I'd like to add, but they are currently waiting for better docs on the Enzyme side. There is also a long-term goal off replacing the fat-lto requirement with the less invasive embed-bc requirement, but this proved to be tricky in the past and only affects compile times.All in all, we have landed most of the hard work here, so that's a very comfortable position to be in before enabling it on nightly.
offload
#[cfg(target_os=X)]annotations. Under the hood, the LLVM-IR which we generate is also still a bit convoluted. In his next PRs, he'll clean up the generated IR, and introduce an offload macro that users shall call instead of the internal offload intrinsic.std::offloadin nightly. After multiple iterations and support from LLVM offload devs, we found a cmake config that does not run into bugs, should not increase Rust CI time too much, and works with both in-tree llvm/clang builds, as well as external clang's (the current case in our Rust CI).clang-linker-wrapper. Once I automated this last wrapper (and enabled offload in CI), nightly users should be able to experiment withstd::offload.ZuseZ4 commentedon Dec 26, 2025
Time for the next round of updates. Again, most of the updates were on the GPU side, but with some notable autodiff improvements too.
autodiff:
@sgasho finished his work on using dlopen to load enzyme and the pr landed. This allowed @Kobzol and me to start working on distributing Enzyme via a standalone component.
As a first step, I added a nicer error if we fail to find or dlopen our Enzyme backend. I also removed most of our autodiff fallbacks, we now unconditionally enable our macro frontend on nightly: remove llvm_enzyme and enzyme fallbacks from most places rust#150133
You may notice that
cargo expandnow works on autodiff code. This also allowed the first bug reports about ICE (internal compiler error) in our macro parser logic.Kobzol opened a PR to build Enzyme in CI. In theory, I should have been able to download that artifact, put it into my sysroot, and use the latest nightly to automatically load it. If that had worked, we could have just merged his PR, and everyone could have started using AD on nightly. Of course, things are never that easy. Even though both Enzyme, LLVM, and rustc were built in CI, the LLVM version shipped along with rustc does not seem compatible with the LLVM version Enzyme was built against. We assume some slight cmake mismatch during our CI builds, which we will have to debug.
offload:
On the gpu side, @Sa4dUs finished his cleanup PR, and along the way also fixed using multiple kernels within a single codebase. When developing the offload MVP I had taken a lot of inspiration from the LLVM-IR generated by clang - and it looks like I had gotten one of the (way too many) LLVM attributes wrong. That caused some metadata to be fused when multiple kernels are present, confusing our offload backend. We started to find more bugs when working on benchmarks, more about the fixes for those in the next update.
I finished cleaning up my offload build PR, and @oli-obk reviewed and approved it. Once the dev-guide gets synced, you should see much simpler usage instructions. Now it's just up to me to automate the last part, then you can compile offload code purely with cargo or rustc.
I also improved how we build offload, which allows us to build it both in CI and locally. CI had some very specific requirements to not increase build times, since our x86-64-dist runner is already quite slow.
Our first benchmarks directly linked against NVIDIA and AMD intrinsics on llvm-ir level. However, we already had an nvptx Rust module for a while, and since recently also an amdgpu module which nicely wraps those intrinsics. I just synced the stdarch repository into rustc a few minutes ago, so from now on, we can replace both with the corresponding Rust functions. In the near future we should get a higher level GPU module, which abstracts away naming differences between vendors.
Most of my past rustc contributions were related to LLVM projects or plugins (Offload and Enzyme), and I increasingly encountered myself asking other people for updates or backports of our LLVM submodule, since upstream LLVM has fixes which were not yet merged into our LLVM submodule. Our llvm working group is quite small and I didn't want to burden them too much with my requests, so I recently asked them to join it, which also got approved. In the future I intend to help a little with the maintenance here.
ZuseZ4 commentedon Jan 16, 2026
Key developments:
std::autodiffis moving closer to nightly, andstd::offloadis gaining various performance, feature, and hardware support improvements.autodiff
@Kobzol, @sgasho, and I continued working on enabling autodiff in nightly. We have a PR up that builds autodiff in CI, and verified that the artifacts can be installed and work on Linux. For apple however, we noticed that any autodiff usage hangs. After some investigation, it turns out that we ended up embedding two LLVM copies, one in rustc, and one in Enzyme. It should be comparably easy to get rid of the second one. Once we verified that this fixes the build, we'll merge the PR to enable autodiff on both targets in nightly.
offload
A lot of interesting updates on the performance, feature, and hardware support side.
@Sa4dUs, @kevinsala, @jdoerfert, and I started implementing the first benchmarks, since that's generally the best way to find missing features or performance issues. We were positively surprised by how good the out-of-the-box performance was. We will implement a few more benchmarks and post the results once we have verified them. We also implemented multiple PRs which implement bugfixes, cleanups, and needed features like support for scalars. We also started working on LLVM optimizations which make sure that we can achieve even better performance.
I noticed that our offload intrinsic allowed running Rust code on the GPU, but it wasn't of much help when calling gpu vendor libraries like cuBLAS. In Adding a new offload_args intrinsic, which only maps arguments rust#150683 I implemented a new helper intrinsic which allows calling those functions conveniently, without having to manually move data to or from the device. It will benefit from the same LLVM optimizations as our full offload intrinsic. It also a bit simpler to set up on the compiler and linker side, so it already works with
stdand mangled kernel names, something that we still have to improve for our main offload intrinsic.A lot of work happened on the LLVM offload side for SPIRV and Intel GPU support. At the moment, our Rust frontend is tested on NVIDIA and AMD server and consumer GPUs, as well as AMD HPC and Lapotop APUs. @karolzwolak reached out since he wants to help with with also running Rust on Intel GPUs. Offload relies on LLVM which started gaining Intel support, so hopefully we won't need much work beyond a new intel-gpu target and a new stdarch module. There is also work on a new spirv target for rustc, which we could also support if it goes through LLVM. Due to some open questions around typed pointers it does not seem clear yet whether it will, so we will have to wait.
Nikita started working on updating our submodule to LLVM 22. This hopefully does not only brings some compile and runtime performance improvements, but also greatly simplifies how we can build and use offload. Once it landed I'll refactor our bootstraping logic, and as part of that start building offload in CI.
ZuseZ4 commentedon Apr 1, 2026
Key developments:
std::autodiffis now partly in CI, andstd::offloadgot tested on a lot more benchmarks.autodiff
Work continued on enabling autodiff in nightly. Since the last update, we have enabled autodiff in some Mingw and Linux runners. Users can now download libEnzyme artifacts, place them locally in the right spot for their toolchain, and then use autodiff on their nightly compiler. Once macOS is added, we will enable a new rustup component that will handle the download for users. Before enabling autodiff on macOS, however, we want to change how we distribute LLVM on this target (from static to dynamic linking). There are a lot of workflows and users of this target, not all of which can be modelled in the Rust CI. Our last two attempts sadly broke such downstream users and local contributors, so both attempts had to be reverted. Since testing here is tricky, progress here might be on the slower side; we will see.
offload
Most of the work on the offload side lately has been invisible, since we were working on implementing more benchmarks and LLVM optimizations, as well as missing features, discovered by those benchmarks. We achieved excellent performance on those benchmarks; more details will soon be presented by @Sa4dUs at the EuroLLVM conference in two weeks!
Beyond benchmarks, there was a lot of tinkering on smaller PRs, reviewing, and housekeeping. LLVM-22 landed, so we updated our bootrstrap code to make use of new APIs, and tried to move a few smaller PRs forward, mainly around a better user experience and for making more Rust features available. Since the focus is still on benchmarks, not many of those PRs landed. They are in a mostly ready state, so it's a good time to pick them up if you're considering contributing. Please ping me on Zulip or in any PR with the offload label if you are interested!
nxsaken commentedon May 7, 2026
The 2025h2 goal period is now over, so I'm closing this issue. Thanks to everyone who contributed! The High-Level ML optimizations goal will build on this work in the 2026 period. Authors are welcome to leave a final comment about the goal's outcome.