Skip to content

Commit f96f6a6

Browse files
ZuseZ4ZuseZ4
committed
Enzyme backend
1 parent 79f5c16 commit f96f6a6

File tree

18 files changed

+209
-0
lines changed

18 files changed

+209
-0
lines changed

Diff for: ‎.gitmodules

+4
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,7 @@
4747
path = src/tools/rustc-perf
4848
url = https://github.com/rust-lang/rustc-perf.git
4949
shallow = true
50+
[submodule "src/tools/enzyme"]
51+
path = src/tools/enzyme
52+
url = https://github.com/EnzymeAD/Enzyme.git
53+
shallow = true

Diff for: ‎config.example.toml

+3
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@
7878
# Indicates whether the LLVM plugin is enabled or not
7979
#plugins = false
8080

81+
# Wheter to build Enzyme as AutoDiff backend.
82+
#enzyme = false
83+
8184
# Indicates whether ccache is used when building LLVM. Set to `true` to use the first `ccache` in
8285
# PATH, or set an absolute path to use a specific version.
8386
#ccache = false

Diff for: ‎src/bootstrap/configure.py

+1
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ def v(*args):
7171
# channel, etc.
7272
o("optimize-llvm", "llvm.optimize", "build optimized LLVM")
7373
o("llvm-assertions", "llvm.assertions", "build LLVM with assertions")
74+
o("llvm-enzyme", "llvm.enzyme", "build LLVM with enzyme")
7475
o("llvm-plugins", "llvm.plugins", "build LLVM with plugin interface")
7576
o("debug-assertions", "rust.debug-assertions", "build with debugging assertions")
7677
o("debug-assertions-std", "rust.debug-assertions-std", "build the standard library with debugging assertions")

Diff for: ‎src/bootstrap/src/core/build_steps/compile.rs

+22
Original file line numberDiff line numberDiff line change
@@ -1169,6 +1169,10 @@ pub fn rustc_cargo_env(
11691169
cargo.env("RUSTC_VERIFY_LLVM_IR", "1");
11701170
}
11711171

1172+
if builder.config.llvm_enzyme {
1173+
cargo.rustflag("--cfg=llvm_enzyme");
1174+
}
1175+
11721176
// Note that this is disabled if LLVM itself is disabled or we're in a check
11731177
// build. If we are in a check build we still go ahead here presuming we've
11741178
// detected that LLVM is already built and good to go which helps prevent
@@ -1772,6 +1776,24 @@ impl Step for Assemble {
17721776
// use that to bootstrap this compiler forward.
17731777
let mut build_compiler = builder.compiler(target_compiler.stage - 1, builder.config.build);
17741778

1779+
// Build enzyme
1780+
let enzyme_install = if builder.config.llvm_enzyme {
1781+
Some(builder.ensure(llvm::Enzyme { target: build_compiler.host }))
1782+
} else {
1783+
None
1784+
};
1785+
1786+
if let Some(enzyme_install) = enzyme_install {
1787+
let lib_ext = std::env::consts::DLL_EXTENSION;
1788+
let src_lib = enzyme_install.join("build/Enzyme/libEnzyme-19").with_extension(lib_ext);
1789+
let libdir = builder.sysroot_libdir(build_compiler, build_compiler.host);
1790+
let target_libdir = builder.sysroot_libdir(target_compiler, target_compiler.host);
1791+
let dst_lib = libdir.join("libEnzyme-19").with_extension(lib_ext);
1792+
let target_dst_lib = target_libdir.join("libEnzyme-19").with_extension(lib_ext);
1793+
builder.copy_link(&src_lib, &dst_lib);
1794+
builder.copy_link(&src_lib, &target_dst_lib);
1795+
}
1796+
17751797
// Build the libraries for this compiler to link to (i.e., the libraries
17761798
// it uses at runtime). NOTE: Crates the target compiler compiles don't
17771799
// link to these. (FIXME: Is that correct? It seems to be correct most

Diff for: ‎src/bootstrap/src/core/build_steps/dist.rs

+40
Original file line numberDiff line numberDiff line change
@@ -1514,6 +1514,7 @@ impl Step for Extended {
15141514
add_component!("llvm-components" => LlvmTools { target });
15151515
add_component!("clippy" => Clippy { compiler, target });
15161516
add_component!("miri" => Miri { compiler, target });
1517+
add_component!("enzyme" => Enzyme { compiler, target });
15171518
add_component!("analysis" => Analysis { compiler, target });
15181519
add_component!("rustc-codegen-cranelift" => CodegenBackend {
15191520
compiler: builder.compiler(stage, target),
@@ -2392,6 +2393,45 @@ impl Step for BuildManifest {
23922393
}
23932394
}
23942395

2396+
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
2397+
pub struct Enzyme {
2398+
pub compiler: Compiler,
2399+
pub target: TargetSelection,
2400+
}
2401+
2402+
impl Step for Enzyme {
2403+
type Output = Option<GeneratedTarball>;
2404+
const DEFAULT: bool = false;
2405+
const ONLY_HOSTS: bool = true;
2406+
2407+
fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
2408+
run.alias("enzyme")
2409+
}
2410+
2411+
fn make_run(run: RunConfig<'_>) {
2412+
run.builder.ensure(Enzyme {
2413+
compiler: run.builder.compiler_for(
2414+
run.builder.top_stage,
2415+
run.builder.config.build,
2416+
run.target,
2417+
),
2418+
target: run.target,
2419+
});
2420+
}
2421+
2422+
fn run(self, builder: &Builder<'_>) -> Option<GeneratedTarball> {
2423+
let compiler = self.compiler;
2424+
let target = self.target;
2425+
let enzyme = builder.ensure(tool::Enzyme { compiler, target });
2426+
let mut tarball = Tarball::new(builder, "enzyme", &target.triple);
2427+
tarball.set_overlay(OverlayKind::Enzyme);
2428+
tarball.is_preview(true);
2429+
tarball.add_file(enzyme, "bin", 0o755);
2430+
tarball.add_legal_and_readme_to("share/doc/enzyme");
2431+
Some(tarball.generate())
2432+
}
2433+
}
2434+
23952435
/// Tarball containing artifacts necessary to reproduce the build of rustc.
23962436
///
23972437
/// Currently this is the PGO profile data.

Diff for: ‎src/bootstrap/src/core/build_steps/llvm.rs

+93
Original file line numberDiff line numberDiff line change
@@ -529,6 +529,7 @@ impl Step for Llvm {
529529
}
530530
};
531531

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

853+
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
854+
pub struct Enzyme {
855+
pub target: TargetSelection,
856+
}
857+
858+
impl Step for Enzyme {
859+
type Output = PathBuf;
860+
const ONLY_HOSTS: bool = true;
861+
862+
fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
863+
run.path("src/tools/enzyme/enzyme")
864+
}
865+
866+
fn make_run(run: RunConfig<'_>) {
867+
run.builder.ensure(Enzyme { target: run.target });
868+
}
869+
870+
/// Compile Enzyme for `target`.
871+
fn run(self, builder: &Builder<'_>) -> PathBuf {
872+
builder.require_submodule(
873+
"src/tools/enzyme",
874+
Some("The Enzyme sources are required for autodiff."),
875+
);
876+
if builder.config.dry_run() {
877+
let out_dir = builder.enzyme_out(self.target);
878+
return out_dir;
879+
}
880+
let target = self.target;
881+
882+
let LlvmResult { llvm_config, .. } = builder.ensure(Llvm { target: self.target });
883+
884+
static STAMP_HASH_MEMO: OnceLock<String> = OnceLock::new();
885+
let smart_stamp_hash = STAMP_HASH_MEMO.get_or_init(|| {
886+
generate_smart_stamp_hash(
887+
builder,
888+
&builder.config.src.join("src/tools/enzyme"),
889+
builder.enzyme_info.sha().unwrap_or_default(),
890+
)
891+
});
892+
893+
let out_dir = builder.enzyme_out(target);
894+
let stamp = out_dir.join("enzyme-finished-building");
895+
let stamp = HashStamp::new(stamp, Some(smart_stamp_hash));
896+
897+
if stamp.is_done() {
898+
if stamp.hash.is_none() {
899+
builder.info(
900+
"Could not determine the Enzyme submodule commit hash. \
901+
Assuming that an Enzyme rebuild is not necessary.",
902+
);
903+
builder.info(&format!(
904+
"To force Enzyme to rebuild, remove the file `{}`",
905+
stamp.path.display()
906+
));
907+
}
908+
return out_dir;
909+
}
910+
911+
builder.info(&format!("Building Enzyme for {}", target));
912+
t!(stamp.remove());
913+
let _time = helpers::timeit(builder);
914+
t!(fs::create_dir_all(&out_dir));
915+
916+
builder.update_submodule(Path::new("src").join("tools").join("enzyme").to_str().unwrap());
917+
let mut cfg = cmake::Config::new(builder.src.join("src/tools/enzyme/enzyme/"));
918+
// FIXME(ZuseZ4): Find a nicer way to use Enzyme Debug builds
919+
//cfg.profile("Debug");
920+
//cfg.define("CMAKE_BUILD_TYPE", "Debug");
921+
configure_cmake(builder, target, &mut cfg, true, LdFlags::default(), &[]);
922+
923+
// Re-use the same flags as llvm to control the level of debug information
924+
// generated for lld.
925+
let profile = match (builder.config.llvm_optimize, builder.config.llvm_release_debuginfo) {
926+
(false, _) => "Debug",
927+
(true, false) => "Release",
928+
(true, true) => "RelWithDebInfo",
929+
};
930+
931+
cfg.out_dir(&out_dir)
932+
.profile(profile)
933+
.env("LLVM_CONFIG_REAL", &llvm_config)
934+
.define("LLVM_ENABLE_ASSERTIONS", "ON")
935+
.define("ENZYME_EXTERNAL_SHARED_LIB", "ON")
936+
.define("LLVM_DIR", builder.llvm_out(target));
937+
938+
cfg.build();
939+
940+
t!(stamp.write());
941+
out_dir
942+
}
943+
}
944+
852945
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
853946
pub struct Lld {
854947
pub target: TargetSelection,

Diff for: ‎src/bootstrap/src/core/build_steps/tool.rs

+1
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,7 @@ bootstrap_tool!(
325325
Compiletest, "src/tools/compiletest", "compiletest", is_unstable_tool = true, allow_features = "test";
326326
BuildManifest, "src/tools/build-manifest", "build-manifest";
327327
RemoteTestClient, "src/tools/remote-test-client", "remote-test-client";
328+
Enzyme, "src/tools/enzyme", "enzyme";
328329
RustInstaller, "src/tools/rust-installer", "rust-installer";
329330
RustdocTheme, "src/tools/rustdoc-themes", "rustdoc-themes";
330331
LintDocs, "src/tools/lint-docs", "lint-docs";

Diff for: ‎src/bootstrap/src/core/builder.rs

+7
Original file line numberDiff line numberDiff line change
@@ -798,6 +798,7 @@ impl<'a> Builder<'a> {
798798
tool::Miri,
799799
tool::CargoMiri,
800800
llvm::Lld,
801+
llvm::Enzyme,
801802
llvm::CrtBeginEnd,
802803
tool::RustdocGUITest,
803804
tool::OptimizedDist,
@@ -1588,6 +1589,12 @@ impl<'a> Builder<'a> {
15881589
rustflags.arg(sysroot_str);
15891590
}
15901591

1592+
// https://rust-lang.zulipchat.com/#narrow/stream/182449-t-compiler.2Fhelp/topic/.E2.9C.94.20link.20new.20library.20into.20stage1.2Frustc
1593+
if self.config.llvm_enzyme {
1594+
rustflags.arg("-l");
1595+
rustflags.arg("Enzyme-19");
1596+
}
1597+
15911598
let use_new_symbol_mangling = match self.config.rust_new_symbol_mangling {
15921599
Some(setting) => {
15931600
// If an explicit setting is given, use that

Diff for: ‎src/bootstrap/src/core/config/config.rs

+8
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@ pub struct Config {
221221
// llvm codegen options
222222
pub llvm_assertions: bool,
223223
pub llvm_tests: bool,
224+
pub llvm_enzyme: bool,
224225
pub llvm_plugins: bool,
225226
pub llvm_optimize: bool,
226227
pub llvm_thin_lto: bool,
@@ -896,6 +897,7 @@ define_config! {
896897
release_debuginfo: Option<bool> = "release-debuginfo",
897898
assertions: Option<bool> = "assertions",
898899
tests: Option<bool> = "tests",
900+
enzyme: Option<bool> = "enzyme",
899901
plugins: Option<bool> = "plugins",
900902
ccache: Option<StringOrBool> = "ccache",
901903
static_libstdcpp: Option<bool> = "static-libstdcpp",
@@ -1591,6 +1593,7 @@ impl Config {
15911593
// we'll infer default values for them later
15921594
let mut llvm_assertions = None;
15931595
let mut llvm_tests = None;
1596+
let mut llvm_enzyme = None;
15941597
let mut llvm_plugins = None;
15951598
let mut debug = None;
15961599
let mut debug_assertions = None;
@@ -1721,6 +1724,8 @@ impl Config {
17211724
config.llvm_tools_enabled = llvm_tools.unwrap_or(true);
17221725
config.rustc_parallel =
17231726
parallel_compiler.unwrap_or(config.channel == "dev" || config.channel == "nightly");
1727+
config.llvm_enzyme =
1728+
llvm_enzyme.unwrap_or(config.channel == "dev" || config.channel == "nightly");
17241729
config.rustc_default_linker = default_linker;
17251730
config.musl_root = musl_root.map(PathBuf::from);
17261731
config.save_toolstates = save_toolstates.map(PathBuf::from);
@@ -1805,6 +1810,7 @@ impl Config {
18051810
release_debuginfo,
18061811
assertions,
18071812
tests,
1813+
enzyme,
18081814
plugins,
18091815
ccache,
18101816
static_libstdcpp,
@@ -1838,6 +1844,7 @@ impl Config {
18381844
set(&mut config.ninja_in_file, ninja);
18391845
llvm_assertions = assertions;
18401846
llvm_tests = tests;
1847+
llvm_enzyme = enzyme;
18411848
llvm_plugins = plugins;
18421849
set(&mut config.llvm_optimize, optimize_toml);
18431850
set(&mut config.llvm_thin_lto, thin_lto);
@@ -2039,6 +2046,7 @@ impl Config {
20392046

20402047
config.llvm_assertions = llvm_assertions.unwrap_or(false);
20412048
config.llvm_tests = llvm_tests.unwrap_or(false);
2049+
config.llvm_enzyme = llvm_enzyme.unwrap_or(false);
20422050
config.llvm_plugins = llvm_plugins.unwrap_or(false);
20432051
config.rust_optimize = optimize.unwrap_or(RustOptimize::Bool(true));
20442052

Diff for: ‎src/bootstrap/src/lib.rs

+10
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,9 @@ const LLD_FILE_NAMES: &[&str] = &["ld.lld", "ld64.lld", "lld-link", "wasm-ld"];
7777
#[allow(clippy::type_complexity)] // It's fine for hard-coded list and type is explained above.
7878
const EXTRA_CHECK_CFGS: &[(Option<Mode>, &str, Option<&[&'static str]>)] = &[
7979
(None, "bootstrap", None),
80+
(Some(Mode::Rustc), "llvm_enzyme", None),
81+
(Some(Mode::Codegen), "llvm_enzyme", None),
82+
(Some(Mode::ToolRustc), "llvm_enzyme", None),
8083
(Some(Mode::Rustc), "parallel_compiler", None),
8184
(Some(Mode::ToolRustc), "parallel_compiler", None),
8285
(Some(Mode::ToolRustc), "rust_analyzer", None),
@@ -140,6 +143,7 @@ pub struct Build {
140143
clippy_info: GitInfo,
141144
miri_info: GitInfo,
142145
rustfmt_info: GitInfo,
146+
enzyme_info: GitInfo,
143147
in_tree_llvm_info: GitInfo,
144148
local_rebuild: bool,
145149
fail_fast: bool,
@@ -306,6 +310,7 @@ impl Build {
306310
let clippy_info = GitInfo::new(omit_git_hash, &src.join("src/tools/clippy"));
307311
let miri_info = GitInfo::new(omit_git_hash, &src.join("src/tools/miri"));
308312
let rustfmt_info = GitInfo::new(omit_git_hash, &src.join("src/tools/rustfmt"));
313+
let enzyme_info = GitInfo::new(omit_git_hash, &src.join("src/tools/enzyme"));
309314

310315
// we always try to use git for LLVM builds
311316
let in_tree_llvm_info = GitInfo::new(false, &src.join("src/llvm-project"));
@@ -393,6 +398,7 @@ impl Build {
393398
clippy_info,
394399
miri_info,
395400
rustfmt_info,
401+
enzyme_info,
396402
in_tree_llvm_info,
397403
cc: RefCell::new(HashMap::new()),
398404
cxx: RefCell::new(HashMap::new()),
@@ -848,6 +854,10 @@ impl Build {
848854
}
849855
}
850856

857+
fn enzyme_out(&self, target: TargetSelection) -> PathBuf {
858+
self.out.join(&*target.triple).join("enzyme")
859+
}
860+
851861
fn lld_out(&self, target: TargetSelection) -> PathBuf {
852862
self.out.join(target).join("lld")
853863
}

Diff for: ‎src/bootstrap/src/utils/change_tracker.rs

+5
Original file line numberDiff line numberDiff line change
@@ -230,4 +230,9 @@ pub const CONFIG_CHANGE_HISTORY: &[ChangeInfo] = &[
230230
severity: ChangeSeverity::Warning,
231231
summary: "./x test --rustc-args was renamed to --compiletest-rustc-args as it only applies there. ./x miri --rustc-args was also removed.",
232232
},
233+
ChangeInfo {
234+
change_id: 129176,
235+
severity: ChangeSeverity::Info,
236+
summary: "New option `llvm.autodiff` to control whether the llvm based autodiff tool (Enzyme) is built.",
237+
},
233238
];

Diff for: ‎src/bootstrap/src/utils/tarball.rs

+5
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ pub(crate) enum OverlayKind {
2121
Cargo,
2222
Clippy,
2323
Miri,
24+
Enzyme,
2425
Rustfmt,
2526
Rls,
2627
RustAnalyzer,
@@ -62,6 +63,7 @@ impl OverlayKind {
6263
"src/tools/rust-analyzer/LICENSE-APACHE",
6364
"src/tools/rust-analyzer/LICENSE-MIT",
6465
],
66+
OverlayKind::Enzyme => &["src/tools/enzyme/README.md", "src/tools/enzyme/LICENSE"],
6567
OverlayKind::RustcCodegenCranelift => &[
6668
"compiler/rustc_codegen_cranelift/Readme.md",
6769
"compiler/rustc_codegen_cranelift/LICENSE-APACHE",
@@ -94,6 +96,9 @@ impl OverlayKind {
9496
OverlayKind::RustAnalyzer => builder
9597
.rust_analyzer_info
9698
.version(builder, &builder.release_num("rust-analyzer/crates/rust-analyzer")),
99+
OverlayKind::Enzyme => {
100+
builder.enzyme_info.version(builder, &builder.release_num("enzyme"))
101+
}
97102
OverlayKind::RustcCodegenCranelift => builder.rust_version(),
98103
OverlayKind::LlvmBitcodeLinker => builder.rust_version(),
99104
}

Diff for: ‎src/tools/build-manifest/src/main.rs

+1
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,7 @@ impl Builder {
470470
| PkgType::Rls
471471
| PkgType::RustAnalyzer
472472
| PkgType::Rustfmt
473+
| PkgType::Enzyme
473474
| PkgType::LlvmTools
474475
| PkgType::RustAnalysis
475476
| PkgType::JsonDocs

0 commit comments

Comments
 (0)