Add repo for Rust-specific soft fork of Enzyme

[traviscross]

Currently our submodule for Enzyme in rust-lang/rust points to an external repository outside of our organization. As we do for LLVM, it would be useful for us to maintain a Rust-specific soft fork of Enzyme within rust-lang/. Let's create a repository for this.

In particular, this would unblock some work being held up by CI issues on the upstream project, and it'd allow us to make some changes that might better facilitate upstreaming of this work into LLVM.

Probably we should also create a wg-enzyme or whatnot team that could maintain this. Alternatively or in addition, we could formalize wg-autodiff, for which we have a Zulip stream but no entry here. But let's handle that separately. For the moment, we just give access to the normal teams that would be likely to need to touch it, review PRs for it, manage issues on it, etc.

CC @oli-obk who seconded the compiler MCP for this work and is overseeing it on that side.

I'm the lang liaison for our approved experiment for this work.

CC @Mark-Simulacrum, @cuviper, @ehuss for thoughts on release and bootstrap and whatnot here.

CC @nikic on the LLVM aspect.

CC @ZuseZ4 who is the owner for this work (and for whom we should find a team at some point).

CC @wsmoses from upstream Enzyme with whom @ZuseZ4 talked this over.

[traviscross]

The bits here I arranged in the same way as our llvm-project repo in the expectation we'd want to manage it in a roughly similar way.

[cuviper]

I think it's a good idea to pull from a @rust-lang repo, as long as we have a goal of keeping differences minimal.

In particular, this would unblock some work being held up by CI issues on the upstream project, and it'd allow us to make some changes that might better facilitate upstreaming of this work into LLVM.

I haven't seen what issues there are, but in general I'd say that changes to fit Rust's CI should be fine. I'm more skeptical about the upstreaming part though -- why would that want to go through Rust's fork?

[traviscross]

@ZuseZ4: See @cuviper's question.

[ZuseZ4]

@cuviper My focous on this fork are these three areas. I don't intend to accept PRs for changes in other areas.

1. A working CI
2. CMake improvements
3. Documentation improvements

1) Upstream’s CI contains a large number of continuous integration tests driven by e.g. its Julia frontend, as well as its JAX-version. Due to the large number of CI tests, they tend to break quite frequently. In addition, the number of runners is limited, making the upstreaming of new (rust) benchmarks tricky. Focusing on running the CI relevant to us will make it much easier to maintain a working CI. I assume there will be some ongoing diff between our fork and upstream. However, upstream also agreed that they rarely update those, so there shouldn't be (m)any merge conflicts.

2) I wrote around 10 Enzyme PRs which improved CMake. For hot-fixes or experiments I might temporarily merge CMake improvements to our fork, but my strong preference is to have those fixes upstream first. I assume that differences will be quickly resolved by upstreaming (if valuable) or being dropped if we develop a better solution with upstream.

3) Regarding docs, upstream does presently not have enough capacity to expand its developer documentation. Some improvement efforts got stuck, others rolled back. I am in contact with developers of (partly not-yet-announced) frontends of Enzyme, who want to learn from what we did in Rust. On the other hand, there are also LLVM people who have, and are helping me to look at the Enzyme codebase to figure things out. In both cases, knowledge usually stays in DMs. I started writing some docs out-of-tree for questions which I got asked, but most of these docs should rather go into the codebase, and not on an unrelated website. Upstream docs would be better, but until they materialize I want to host a fork with a lower threshold to avoid doc improvements to starve, and accelerate the onboarding of new developers.

The documentation is what I care most about, since due to the lack of better docs, I ended up missing an enzyme optimization pass resulting in non-optimal Rust-Enzyme performance. It's clear that such things hinder us in achieving optimal performance for Rust-Enzyme, as well as onboarding new people to Rust-Enzyme. I'd argue improving on this point alone by having centralized documentation for Rust+Enzyme and the relevant interfaces between Enzyme upstream and Rust-Enzyme, would make a large difference for us in the long run, and is 100% worth it (and still shouldn't cause too many merge conflicts).

Contributing to having autodiff in upstream LLVM is more of an aspirational goal, but other things like better GPU support or an MLIR backend for rustc could also use our time. I would be fine to remove the upstreaming section to make this proposal more focused. We're not the only users, so we don't directly have to try to fix everything. Other frontends also have an interest in better Enzyme reliability and documentation, so I would suggest we run this for a while on nightly, try to keep the diff as minimal as possible, and re-evaluate this question once we have collected more data and contributors.

[wsmoses]

++ to whatever makes your guys life easier, especially if you all are working at a higher velocity than upstream.

That said, I'm a bit surprised that those are the points you need it for.

Version skew

I would've expected it more for LLVM version skew between Rust LLVM and upstream/Enzyme LLVM. For context, Enzyme currently supports a wide range of LLVM versions (which is currently the debate we're having internally re how to upstream Enzyme into LLVM itself -- if you have thoughts all are welcome at our weekly Enzyme Open Design meetings). Specifically Enzyme supports LLVM 15 - mainline (though EnzymeMLIR only supports main LLVM). Sometimes there is a little lag, but any breaking get updated in a week or two at worst. If Rust uses a LLVM-main-like but not LLVM main there might be a slight skew if an API changes, which definitely having a fork could help.

CI

Enzyme's CI usually doesn't break from integration tests of upstream changing. It goes down because our self hosted runners for integration are kind of wimpy and die if given too big a task. In particular, when @ZuseZ4 tried adding some Rust integration tests earlier, it ended up ooming the VM, crashing it and disconnecting it from Github (bringing down all our integration CI =/ making it hard to actually have nice checks). These are also shared across the org for EnzymeJaX, etc. If you have a fix for the github actions, that would be great and we'd love that upstream (maybe changing it to use docker or escape the workload). Alternatively if you have bigger runners that would be great too!

I will note that we have been doing an org-wide redo of CI in the past couple of weeks (starting with downstream repos like Reactant.jl, Enzyme-JaX, Enzyme.jl, but haven't done that for Enzyme proper yet). Contributions welcome of course!

CMake improvements

This one confused me a bit, so perhaps I'm not quite understanding what the need is here. All of your PRs in the past few weeks got a review within a day or too so I'm not sure what the issue is? (x/ref EnzymeAD/Enzyme#2249 EnzymeAD/Enzyme#2248 EnzymeAD/Enzyme#2245 EnzymeAD/Enzyme#2238 ).

Docs

Again I'm a bit confused here. @ZuseZ4 you have ownership on https://github.com/EnzymeAD/rustbook/, and for example pushed a commit without review 4 days ago (EnzymeAD/rustbook#37), which gives you ownership over enzyme.mit.edu/rust . So I'm not sure what the blocker here is for Rust?

We definitely should improve docs though, for any/all of Enzyme-Rust, Enzyme, EnzymeJaX, Enzyme.jl, etc .

[wsmoses]

oh also I'll just also throw this out there too, Enzyme has open weekly developer meetings. All are welcome!

So if at any point it's helpful for you all to chat with Enzyme devs about anything (be it design wise, something blocking, etc), please join!

[ZuseZ4]

I am happy to provide a detailed answer to the points raised above, but I think it is simpler if we focus on ++ to whatever makes your guys life easier, especially if you all are working at a higher velocity than upstream.

Both on the Rust side and the Enzyme core side there seem to be plans to improve the areas mentioned above, but having too many moving pieces can be challenging. I would suggest to give both sides a few months to implement the improvements they have in mind, and verify that the desired level of stability can be achieved with the corresponding level of features. Once that has been the case we can work on gradually reducing the scope of this soft-fork to that of our llvm fork.