Skip to content

Commit

Permalink
Auto merge of rust-lang#129176 - EnzymeAD:enzyme-backend, r=albertlar…
Browse files Browse the repository at this point in the history
…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:
```
[email protected]: 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
  • Loading branch information
bors committed Sep 6, 2024
2 parents a3af208 + 4f5c16d commit 59d4114
Show file tree
Hide file tree
Showing 13 changed files with 160 additions and 0 deletions.
4 changes: 4 additions & 0 deletions .gitmodules
Original file line number Diff line number Diff line change
Expand Up @@ -47,3 +47,7 @@
path = src/tools/rustc-perf
url = https://github.com/rust-lang/rustc-perf.git
shallow = true
[submodule "src/tools/enzyme"]
path = src/tools/enzyme
url = https://github.com/EnzymeAD/Enzyme.git
shallow = true
3 changes: 3 additions & 0 deletions config.example.toml
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@
# Indicates whether the LLVM plugin is enabled or not
#plugins = false

# Wheter to build Enzyme as AutoDiff backend.
#enzyme = false

# Indicates whether ccache is used when building LLVM. Set to `true` to use the first `ccache` in
# PATH, or set an absolute path to use a specific version.
#ccache = false
Expand Down
1 change: 1 addition & 0 deletions src/bootstrap/configure.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ def v(*args):
# channel, etc.
o("optimize-llvm", "llvm.optimize", "build optimized LLVM")
o("llvm-assertions", "llvm.assertions", "build LLVM with assertions")
o("llvm-enzyme", "llvm.enzyme", "build LLVM with enzyme")
o("llvm-plugins", "llvm.plugins", "build LLVM with plugin interface")
o("debug-assertions", "rust.debug-assertions", "build with debugging assertions")
o("debug-assertions-std", "rust.debug-assertions-std", "build the standard library with debugging assertions")
Expand Down
22 changes: 22 additions & 0 deletions src/bootstrap/src/core/build_steps/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1195,6 +1195,10 @@ pub fn rustc_cargo_env(
cargo.env("RUSTC_VERIFY_LLVM_IR", "1");
}

if builder.config.llvm_enzyme {
cargo.rustflag("--cfg=llvm_enzyme");
}

// Note that this is disabled if LLVM itself is disabled or we're in a check
// build. If we are in a check build we still go ahead here presuming we've
// detected that LLVM is already built and good to go which helps prevent
Expand Down Expand Up @@ -1790,6 +1794,24 @@ impl Step for Assemble {
// use that to bootstrap this compiler forward.
let mut build_compiler = builder.compiler(target_compiler.stage - 1, builder.config.build);

// Build enzyme
let enzyme_install = if builder.config.llvm_enzyme {
Some(builder.ensure(llvm::Enzyme { target: build_compiler.host }))
} else {
None
};

if let Some(enzyme_install) = enzyme_install {
let lib_ext = std::env::consts::DLL_EXTENSION;
let src_lib = enzyme_install.join("build/Enzyme/libEnzyme-19").with_extension(lib_ext);
let libdir = builder.sysroot_libdir(build_compiler, build_compiler.host);
let target_libdir = builder.sysroot_libdir(target_compiler, target_compiler.host);
let dst_lib = libdir.join("libEnzyme-19").with_extension(lib_ext);
let target_dst_lib = target_libdir.join("libEnzyme-19").with_extension(lib_ext);
builder.copy_link(&src_lib, &dst_lib);
builder.copy_link(&src_lib, &target_dst_lib);
}

// Build the libraries for this compiler to link to (i.e., the libraries
// it uses at runtime). NOTE: Crates the target compiler compiles don't
// link to these. (FIXME: Is that correct? It seems to be correct most
Expand Down
95 changes: 95 additions & 0 deletions src/bootstrap/src/core/build_steps/llvm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -529,6 +529,7 @@ impl Step for Llvm {
}
};

// FIXME(ZuseZ4): Do we need that for Enzyme too?
// When building LLVM with LLVM_LINK_LLVM_DYLIB for macOS, an unversioned
// libLLVM.dylib will be built. However, llvm-config will still look
// for a versioned path like libLLVM-14.dylib. Manually create a symbolic
Expand Down Expand Up @@ -849,6 +850,100 @@ fn get_var(var_base: &str, host: &str, target: &str) -> Option<OsString> {
.or_else(|| env::var_os(var_base))
}

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

impl Step for Enzyme {
type Output = PathBuf;
const ONLY_HOSTS: bool = true;

fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
run.path("src/tools/enzyme/enzyme")
}

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

/// Compile Enzyme for `target`.
fn run(self, builder: &Builder<'_>) -> PathBuf {
builder.require_submodule(
"src/tools/enzyme",
Some("The Enzyme sources are required for autodiff."),
);
if builder.config.dry_run() {
let out_dir = builder.enzyme_out(self.target);
return out_dir;
}
let target = self.target;

let LlvmResult { llvm_config, .. } = builder.ensure(Llvm { target: self.target });

static STAMP_HASH_MEMO: OnceLock<String> = OnceLock::new();
let smart_stamp_hash = STAMP_HASH_MEMO.get_or_init(|| {
generate_smart_stamp_hash(
builder,
&builder.config.src.join("src/tools/enzyme"),
builder.enzyme_info.sha().unwrap_or_default(),
)
});

let out_dir = builder.enzyme_out(target);
let stamp = out_dir.join("enzyme-finished-building");
let stamp = HashStamp::new(stamp, Some(smart_stamp_hash));

if stamp.is_done() {
if stamp.hash.is_none() {
builder.info(
"Could not determine the Enzyme submodule commit hash. \
Assuming that an Enzyme rebuild is not necessary.",
);
builder.info(&format!(
"To force Enzyme to rebuild, remove the file `{}`",
stamp.path.display()
));
}
return out_dir;
}

builder.info(&format!("Building Enzyme for {}", target));
t!(stamp.remove());
let _time = helpers::timeit(builder);
t!(fs::create_dir_all(&out_dir));

builder
.config
.update_submodule(Path::new("src").join("tools").join("enzyme").to_str().unwrap());
let mut cfg = cmake::Config::new(builder.src.join("src/tools/enzyme/enzyme/"));
// FIXME(ZuseZ4): Find a nicer way to use Enzyme Debug builds
//cfg.profile("Debug");
//cfg.define("CMAKE_BUILD_TYPE", "Debug");
configure_cmake(builder, target, &mut cfg, true, LdFlags::default(), &[]);

// Re-use the same flags as llvm to control the level of debug information
// generated for lld.
let profile = match (builder.config.llvm_optimize, builder.config.llvm_release_debuginfo) {
(false, _) => "Debug",
(true, false) => "Release",
(true, true) => "RelWithDebInfo",
};

cfg.out_dir(&out_dir)
.profile(profile)
.env("LLVM_CONFIG_REAL", &llvm_config)
.define("LLVM_ENABLE_ASSERTIONS", "ON")
.define("ENZYME_EXTERNAL_SHARED_LIB", "ON")
.define("LLVM_DIR", builder.llvm_out(target));

cfg.build();

t!(stamp.write());
out_dir
}
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Lld {
pub target: TargetSelection,
Expand Down
7 changes: 7 additions & 0 deletions src/bootstrap/src/core/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -798,6 +798,7 @@ impl<'a> Builder<'a> {
tool::Miri,
tool::CargoMiri,
llvm::Lld,
llvm::Enzyme,
llvm::CrtBeginEnd,
tool::RustdocGUITest,
tool::OptimizedDist,
Expand Down Expand Up @@ -1588,6 +1589,12 @@ impl<'a> Builder<'a> {
rustflags.arg(sysroot_str);
}

// https://rust-lang.zulipchat.com/#narrow/stream/182449-t-compiler.2Fhelp/topic/.E2.9C.94.20link.20new.20library.20into.20stage1.2Frustc
if self.config.llvm_enzyme {
rustflags.arg("-l");
rustflags.arg("Enzyme-19");
}

let use_new_symbol_mangling = match self.config.rust_new_symbol_mangling {
Some(setting) => {
// If an explicit setting is given, use that
Expand Down
8 changes: 8 additions & 0 deletions src/bootstrap/src/core/config/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,7 @@ pub struct Config {
// llvm codegen options
pub llvm_assertions: bool,
pub llvm_tests: bool,
pub llvm_enzyme: bool,
pub llvm_plugins: bool,
pub llvm_optimize: bool,
pub llvm_thin_lto: bool,
Expand Down Expand Up @@ -898,6 +899,7 @@ define_config! {
release_debuginfo: Option<bool> = "release-debuginfo",
assertions: Option<bool> = "assertions",
tests: Option<bool> = "tests",
enzyme: Option<bool> = "enzyme",
plugins: Option<bool> = "plugins",
ccache: Option<StringOrBool> = "ccache",
static_libstdcpp: Option<bool> = "static-libstdcpp",
Expand Down Expand Up @@ -1603,6 +1605,7 @@ impl Config {
// we'll infer default values for them later
let mut llvm_assertions = None;
let mut llvm_tests = None;
let mut llvm_enzyme = None;
let mut llvm_plugins = None;
let mut debug = None;
let mut debug_assertions = None;
Expand Down Expand Up @@ -1722,6 +1725,8 @@ impl Config {
config.llvm_tools_enabled = llvm_tools.unwrap_or(true);
config.rustc_parallel =
parallel_compiler.unwrap_or(config.channel == "dev" || config.channel == "nightly");
config.llvm_enzyme =
llvm_enzyme.unwrap_or(config.channel == "dev" || config.channel == "nightly");
config.rustc_default_linker = default_linker;
config.musl_root = musl_root.map(PathBuf::from);
config.save_toolstates = save_toolstates.map(PathBuf::from);
Expand Down Expand Up @@ -1806,6 +1811,7 @@ impl Config {
release_debuginfo,
assertions,
tests,
enzyme,
plugins,
ccache,
static_libstdcpp,
Expand Down Expand Up @@ -1839,6 +1845,7 @@ impl Config {
set(&mut config.ninja_in_file, ninja);
llvm_assertions = assertions;
llvm_tests = tests;
llvm_enzyme = enzyme;
llvm_plugins = plugins;
set(&mut config.llvm_optimize, optimize_toml);
set(&mut config.llvm_thin_lto, thin_lto);
Expand Down Expand Up @@ -2055,6 +2062,7 @@ impl Config {

config.llvm_assertions = llvm_assertions.unwrap_or(false);
config.llvm_tests = llvm_tests.unwrap_or(false);
config.llvm_enzyme = llvm_enzyme.unwrap_or(false);
config.llvm_plugins = llvm_plugins.unwrap_or(false);
config.rust_optimize = optimize.unwrap_or(RustOptimize::Bool(true));

Expand Down
10 changes: 10 additions & 0 deletions src/bootstrap/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,9 @@ const LLD_FILE_NAMES: &[&str] = &["ld.lld", "ld64.lld", "lld-link", "wasm-ld"];
#[allow(clippy::type_complexity)] // It's fine for hard-coded list and type is explained above.
const EXTRA_CHECK_CFGS: &[(Option<Mode>, &str, Option<&[&'static str]>)] = &[
(None, "bootstrap", None),
(Some(Mode::Rustc), "llvm_enzyme", None),
(Some(Mode::Codegen), "llvm_enzyme", None),
(Some(Mode::ToolRustc), "llvm_enzyme", None),
(Some(Mode::Rustc), "parallel_compiler", None),
(Some(Mode::ToolRustc), "parallel_compiler", None),
(Some(Mode::ToolRustc), "rust_analyzer", None),
Expand Down Expand Up @@ -140,6 +143,7 @@ pub struct Build {
clippy_info: GitInfo,
miri_info: GitInfo,
rustfmt_info: GitInfo,
enzyme_info: GitInfo,
in_tree_llvm_info: GitInfo,
local_rebuild: bool,
fail_fast: bool,
Expand Down Expand Up @@ -307,6 +311,7 @@ impl Build {
let clippy_info = GitInfo::new(omit_git_hash, &src.join("src/tools/clippy"));
let miri_info = GitInfo::new(omit_git_hash, &src.join("src/tools/miri"));
let rustfmt_info = GitInfo::new(omit_git_hash, &src.join("src/tools/rustfmt"));
let enzyme_info = GitInfo::new(omit_git_hash, &src.join("src/tools/enzyme"));

// we always try to use git for LLVM builds
let in_tree_llvm_info = GitInfo::new(false, &src.join("src/llvm-project"));
Expand Down Expand Up @@ -400,6 +405,7 @@ impl Build {
clippy_info,
miri_info,
rustfmt_info,
enzyme_info,
in_tree_llvm_info,
cc: RefCell::new(HashMap::new()),
cxx: RefCell::new(HashMap::new()),
Expand Down Expand Up @@ -755,6 +761,10 @@ impl Build {
}
}

fn enzyme_out(&self, target: TargetSelection) -> PathBuf {
self.out.join(&*target.triple).join("enzyme")
}

fn lld_out(&self, target: TargetSelection) -> PathBuf {
self.out.join(target).join("lld")
}
Expand Down
5 changes: 5 additions & 0 deletions src/bootstrap/src/utils/change_tracker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -245,4 +245,9 @@ pub const CONFIG_CHANGE_HISTORY: &[ChangeInfo] = &[
severity: ChangeSeverity::Warning,
summary: "Removed `rust.split-debuginfo` as it was deprecated long time ago.",
},
ChangeInfo {
change_id: 129176,
severity: ChangeSeverity::Info,
summary: "New option `llvm.enzyme` to control whether the llvm based autodiff tool (Enzyme) is built.",
},
];
1 change: 1 addition & 0 deletions src/tools/enzyme
Submodule enzyme added at 2fe516
1 change: 1 addition & 0 deletions src/tools/tidy/config/black.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,5 +12,6 @@ extend-exclude = """(\
src/llvm-project/|\
src/doc/embedded-book/|\
src/tools/rustc-perf/|\
src/tools/enzyme/|\
library/backtrace/
)"""
2 changes: 2 additions & 0 deletions src/tools/tidy/config/ruff.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ extend-exclude = [
"src/llvm-project/",
"src/doc/embedded-book/",
"library/backtrace/",
"src/tools/enzyme/",
"src/tools/rustc-perf/",
# Hack: CI runs from a subdirectory under the main checkout
"../src/doc/nomicon/",
Expand All @@ -29,6 +30,7 @@ extend-exclude = [
"../src/llvm-project/",
"../src/doc/embedded-book/",
"../library/backtrace/",
"../src/tools/enzyme/",
"../src/tools/rustc-perf/",
]

Expand Down
1 change: 1 addition & 0 deletions src/tools/tidy/src/walk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ pub fn filter_dirs(path: &Path) -> bool {
"src/tools/rust-analyzer",
"src/tools/rustc-perf",
"src/tools/rustfmt",
"src/tools/enzyme",
"src/doc/book",
"src/doc/edition-guide",
"src/doc/embedded-book",
Expand Down

0 comments on commit 59d4114

Please sign in to comment.