Skip to content

Add enzyme distribution step #140244

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions src/bootstrap/src/core/build_steps/dist.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2599,3 +2599,55 @@ impl Step for Gcc {
tarball.generate()
}
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Enzyme {
pub target: TargetSelection,
}

impl Step for Enzyme {
type Output = Option<GeneratedTarball>;
const DEFAULT: bool = false;
const ONLY_HOSTS: bool = true;

fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
run.alias("enzyme")
}

fn make_run(run: RunConfig<'_>) {
run.builder.ensure(Enzyme { target: run.target });
}

fn run(self, builder: &Builder<'_>) -> Self::Output {
let mut tarball = Tarball::new(builder, "enzyme", &self.target.triple);
let enzyme_dir =
builder.ensure(super::llvm::Enzyme { target: self.target }).join("build/Enzyme");

tarball.set_overlay(OverlayKind::Enzyme);
tarball.is_preview(true);

if let Some(llvm_config) = builder.llvm_config(builder.config.build) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, this is a weird. When I have download-ci-llvm = true enabled, this LLVM config is the CI LLVM (which is version 20). But when I actually build Enzyme locally, it generates libEnzyme-18. My host system LLVM is LLVM 18, so that looks related? I don't know why is the filename containing the version of the host compiler rather than the version of the LLVM source code though.

@ZuseZ4 Does Enzyme have its own versioning separate from LLVM? I found some version 0.0.79 in the source code.

Copy link
Member

@ZuseZ4 ZuseZ4 Apr 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rg "0\.0\.79" doesn't return anything in Enzyme, where did you find it?
Enzyme has releases but doesn't follow semver (in a meaningful way only uses patch versions), so I conveniently ignore it.
https://github.com/EnzymeAD/Enzyme/releases
We also just use Enzyme's LLVM Pass which allows us to ignore most of the API, so breaking changes barely ever affect us.
They are more relevant for Julia which has a different build infra working on gh releases. We just use the source code from rust-lang/Enzyme.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just for clarification of

when I actually build Enzyme locally

I currently only try (and have it working):
A) building LLVM locally, then building Enzyme based on it.
Here we try:
B) building LLVM remote (in CI), then building Enzyme also remote based on it.
I guess you tried
C) building LLVM remote, then downloading it and building Enzyme locally against it?

I thought based on Zulip that you wanted a different solution for C, right?

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Semver is followed for user-level enzyme_autodiff and related functions. All user-level changes have been backwards compatible, so we've never bumped. Internal APIs (like LLVM) have no guarantees of stability

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I saw the version here: https://github.com/rust-lang/enzyme/blob/main/enzyme/CMakeLists.txt#L10

To clarify, I was indeed using option C), but almost inadvertedly, as download-ci-llvm is enabled in a lot of default configs and many people use it, it would be nice if bootstrap either "just worked" with it, or at least if it produced a loud error.

I don't really understand how this works though. The build step for Enzyme seems to pass LLVM_CONFIG_REAL to the llvm-config of the LLVM downloaded from CI. I would thus expect that Enzyme either attaches to that, or that it is somehow linked to the in-tree version of LLVM. But instead it seems to produce a version number tied to the host LLVM used to build Enzyme itself.

How exactly is Enzyme "tied" to a specific LLVM sysroot/source code/build/whatever?

Copy link
Member

@ZuseZ4 ZuseZ4 Apr 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Enzyme works with the LLVM it was build against, everything else could only work by coincidence. E.g. this is one of the cmake commands I use locally:
cmake .. -G Ninja -DLLVM_DIR=/home/manuel/prog/rust-middle/build/x86_64-unknown-linux-gnu/llvm/build/lib/cmake/llvm -DLLVM_EXTERNAL_LIT=/home/manuel/prog/rust-middle/src/llvm-project/llvm/utils/lit/lit.py -DCMAKE_BUILD_TYPE=Release -DCMAKE_EXPORT_COMPILE_COMMANDS=YES -DBUILD_SHARED_LIBS=On

We should build Enzyme wherever we also build LLVM. So after building LLVM, we would build Enzyme and set LLVM_DIR to the path of the newly build llvm. We should not point to the LLVM/Clang that was used to build Enzyme and LLVM.

So we should find a way to point Enzyme to a specific LLVM checkout/build.

I mean that's why I added the -DLLVM_DIR definition in rust/src/bootstrap/src/core/build_steps/llvm.rs, it just looks like we don't set it correctly in CI right now? So I guess builder.llvm_out(target) is wrong then (?)
As another puzzle piece, these folders from an apple CI run seem correct? We build using LLVM/Clang-15, but the self-build LLVM: https://github.com/rust-lang-ci/rust/actions/runs/14582626710/job/40902344579

Copy link
Contributor

@Kobzol Kobzol Apr 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, I see now, I didn't realize that LLVM_DIR is the thing that links Enzyme to LLVM, I thought it's the output directory where the Enzyme artifacts will be generated.

In that case I suppose that without download-ci-llvm it should mostly do what we want. I think that the path we give to LLVM_DIR is not exactly what it should be, as it's essentially just build/<target>/llvm; it looks like Enzyme's CMake expects it to be something a bit different (e.g. the path you use, build/<target>/llvm/build/lib/cmake/llvm. That being said, the Enzyme CMake seems to use recursive greps for finding the files that it needs, so it looks like it is able to resolve what it requires.

Ok, in that case I'm fine with just adding an error into the llvm::Enzyme step that it currently cannot be combined with download-ci-llvm.

Copy link
Member

@ZuseZ4 ZuseZ4 Apr 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, in case that the code isn't clear, @Shourya742 can you please add a comment around line 973, saying that we build enzyme against this specific LLVM (and using it with a different LLVM isn't supported)?

Copy link
Member

@ZuseZ4 ZuseZ4 Apr 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also @Kobzol now that I already have you here, do you mind looking at the apple issue here? :D https://github.com/rust-lang-ci/rust/actions/runs/14582626710/job/40902344579#step:28:28938 and checking if the llvm path looks right to you? I feel like it should be ok, 15 is used for building everything, 20 is the one we build against. But it still fails to link LLVM. Presumably Enzyme's cmake should not just link against llvm (-lLLVM), but also specify -L based on the given LLVM_DIR?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not exactly sure how linking works on Apple, but just passing -lLLVM is unlikely to work, IMO. The directory where the libLLVM.dylib (or static archive) files can be found will likely be needed.

let major = llvm::get_llvm_version_major(builder, &llvm_config);
let prefix = format!("libEnzyme-{major}");
let mut found = false;

for entry in std::fs::read_dir(&enzyme_dir)
.unwrap_or_else(|_| panic!("Failed to read {:?}", enzyme_dir))
{
let path = entry.unwrap().path();
if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
if name.starts_with(&prefix) && is_dylib(&path) {
tarball.add_file(path, "", FileType::NativeLibrary);
found = true;
}
}
}

assert!(found, "Enzyme library starting with '{}' not found", prefix);
tarball.add_legal_and_readme_to("share/doc/enzyme");
return Some(tarball.generate());
}

None
}
}
9 changes: 9 additions & 0 deletions src/bootstrap/src/core/build_steps/llvm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -903,6 +903,14 @@ impl Step for Enzyme {
),
)]
fn run(self, builder: &Builder<'_>) -> PathBuf {
// FIXME: This can be removed once we start shipping the CMake files needed by Enzyme
// to build against a standalone LLVM binary.
// For now, Enzyme must be built with the in-tree LLVM and shared LLVM (`llvm.link-shared = true`), not one downloaded from CI.
if builder.config.llvm_from_ci || !builder.config.llvm_link_shared() {
panic!(
"Enzyme must be built with an in-tree LLVM and requires `llvm.link-shared = true`.Set `llvm.download-ci-llvm = false` and `llvm.link-shared = true` to build Enzyme."
);
}
builder.require_submodule(
"src/tools/enzyme",
Some("The Enzyme sources are required for autodiff."),
Expand Down Expand Up @@ -970,6 +978,7 @@ impl Step for Enzyme {
.env("LLVM_CONFIG_REAL", &llvm_config)
.define("LLVM_ENABLE_ASSERTIONS", "ON")
.define("ENZYME_EXTERNAL_SHARED_LIB", "ON")
// Enzyme must be built against this exact LLVM build — mixing versions breaks compatibility.
.define("LLVM_DIR", builder.llvm_out(target));

cfg.build();
Expand Down
3 changes: 2 additions & 1 deletion src/bootstrap/src/core/builder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1084,7 +1084,8 @@ impl<'a> Builder<'a> {
dist::PlainSourceTarball,
dist::BuildManifest,
dist::ReproducibleArtifacts,
dist::Gcc
dist::Gcc,
dist::Enzyme
),
Kind::Install => describe!(
install::Docs,
Expand Down
3 changes: 3 additions & 0 deletions src/bootstrap/src/utils/tarball.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ pub(crate) enum OverlayKind {
RustAnalyzer,
RustcCodegenCranelift,
LlvmBitcodeLinker,
Enzyme,
}

impl OverlayKind {
Expand Down Expand Up @@ -72,6 +73,7 @@ impl OverlayKind {
"LICENSE-MIT",
"src/tools/llvm-bitcode-linker/README.md",
],
OverlayKind::Enzyme => &["src/tools/enzyme/Readme.md", "src/tools/enzyme/LICENSE"],
}
}

Expand All @@ -94,6 +96,7 @@ impl OverlayKind {
.version(builder, &builder.release_num("rust-analyzer/crates/rust-analyzer")),
OverlayKind::RustcCodegenCranelift => builder.rust_version(),
OverlayKind::LlvmBitcodeLinker => builder.rust_version(),
OverlayKind::Enzyme => builder.rust_version(),
}
}
}
Expand Down
3 changes: 2 additions & 1 deletion src/ci/docker/host-x86_64/dist-x86_64-linux/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,8 @@ ENV SCRIPT python3 ../x.py build --set rust.debug=true opt-dist && \
--include-default-paths \
build-manifest bootstrap && \
# Use GCC for building GCC, as it seems to behave badly when built with Clang
CC=/rustroot/bin/cc CXX=/rustroot/bin/c++ python3 ../x.py dist gcc
# Also build Enzyme only on the x64 Linux dist runner
CC=/rustroot/bin/cc CXX=/rustroot/bin/c++ python3 ../x.py dist gcc && ../x.py dist enzyme
ENV CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_LINKER=clang

# This is the only builder which will create source tarballs
Expand Down
2 changes: 1 addition & 1 deletion src/tools/enzyme
Submodule enzyme updated 122 files
Loading