diff --git a/.gitmodules b/.gitmodules index b5250d493864..926807336d79 100644 --- a/.gitmodules +++ b/.gitmodules @@ -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 diff --git a/config.example.toml b/config.example.toml index 1b7de662e84a..4734378b2309 100644 --- a/config.example.toml +++ b/config.example.toml @@ -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 diff --git a/src/bootstrap/configure.py b/src/bootstrap/configure.py index 768aac912ce4..49d564642bd6 100755 --- a/src/bootstrap/configure.py +++ b/src/bootstrap/configure.py @@ -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") diff --git a/src/bootstrap/src/core/build_steps/compile.rs b/src/bootstrap/src/core/build_steps/compile.rs index 4353cfadd8d3..af3e71b005fa 100644 --- a/src/bootstrap/src/core/build_steps/compile.rs +++ b/src/bootstrap/src/core/build_steps/compile.rs @@ -1169,6 +1169,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 @@ -1772,6 +1776,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 diff --git a/src/bootstrap/src/core/build_steps/dist.rs b/src/bootstrap/src/core/build_steps/dist.rs index 530eb9b446a4..74fd83a8bd63 100644 --- a/src/bootstrap/src/core/build_steps/dist.rs +++ b/src/bootstrap/src/core/build_steps/dist.rs @@ -1514,6 +1514,7 @@ impl Step for Extended { add_component!("llvm-components" => LlvmTools { target }); add_component!("clippy" => Clippy { compiler, target }); add_component!("miri" => Miri { compiler, target }); + add_component!("enzyme" => Enzyme { compiler, target }); add_component!("analysis" => Analysis { compiler, target }); add_component!("rustc-codegen-cranelift" => CodegenBackend { compiler: builder.compiler(stage, target), @@ -2392,6 +2393,45 @@ impl Step for BuildManifest { } } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct Enzyme { + pub compiler: Compiler, + pub target: TargetSelection, +} + +impl Step for Enzyme { + type Output = Option; + 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 { + compiler: run.builder.compiler_for( + run.builder.top_stage, + run.builder.config.build, + run.target, + ), + target: run.target, + }); + } + + fn run(self, builder: &Builder<'_>) -> Option { + let compiler = self.compiler; + let target = self.target; + let enzyme = builder.ensure(tool::Enzyme { compiler, target }); + let mut tarball = Tarball::new(builder, "enzyme", &target.triple); + tarball.set_overlay(OverlayKind::Enzyme); + tarball.is_preview(true); + tarball.add_file(enzyme, "bin", 0o755); + tarball.add_legal_and_readme_to("share/doc/enzyme"); + Some(tarball.generate()) + } +} + /// Tarball containing artifacts necessary to reproduce the build of rustc. /// /// Currently this is the PGO profile data. diff --git a/src/bootstrap/src/core/build_steps/llvm.rs b/src/bootstrap/src/core/build_steps/llvm.rs index c5a1ab788016..0376b183a1c0 100644 --- a/src/bootstrap/src/core/build_steps/llvm.rs +++ b/src/bootstrap/src/core/build_steps/llvm.rs @@ -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 @@ -849,6 +850,98 @@ fn get_var(var_base: &str, host: &str, target: &str) -> Option { .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 = 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.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, diff --git a/src/bootstrap/src/core/build_steps/tool.rs b/src/bootstrap/src/core/build_steps/tool.rs index 4d573b107b58..ed9502fe1f57 100644 --- a/src/bootstrap/src/core/build_steps/tool.rs +++ b/src/bootstrap/src/core/build_steps/tool.rs @@ -325,6 +325,7 @@ bootstrap_tool!( Compiletest, "src/tools/compiletest", "compiletest", is_unstable_tool = true, allow_features = "test"; BuildManifest, "src/tools/build-manifest", "build-manifest"; RemoteTestClient, "src/tools/remote-test-client", "remote-test-client"; + Enzyme, "src/tools/enzyme", "enzyme"; RustInstaller, "src/tools/rust-installer", "rust-installer"; RustdocTheme, "src/tools/rustdoc-themes", "rustdoc-themes"; LintDocs, "src/tools/lint-docs", "lint-docs"; diff --git a/src/bootstrap/src/core/builder.rs b/src/bootstrap/src/core/builder.rs index 454cc20d1553..522c14fc1564 100644 --- a/src/bootstrap/src/core/builder.rs +++ b/src/bootstrap/src/core/builder.rs @@ -798,6 +798,7 @@ impl<'a> Builder<'a> { tool::Miri, tool::CargoMiri, llvm::Lld, + llvm::Enzyme, llvm::CrtBeginEnd, tool::RustdocGUITest, tool::OptimizedDist, @@ -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 diff --git a/src/bootstrap/src/core/config/config.rs b/src/bootstrap/src/core/config/config.rs index bdd9fd755aa2..ba2cb57a7849 100644 --- a/src/bootstrap/src/core/config/config.rs +++ b/src/bootstrap/src/core/config/config.rs @@ -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, @@ -896,6 +897,7 @@ define_config! { release_debuginfo: Option = "release-debuginfo", assertions: Option = "assertions", tests: Option = "tests", + enzyme: Option = "enzyme", plugins: Option = "plugins", ccache: Option = "ccache", static_libstdcpp: Option = "static-libstdcpp", @@ -1591,6 +1593,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; @@ -1721,6 +1724,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); @@ -1805,6 +1810,7 @@ impl Config { release_debuginfo, assertions, tests, + enzyme, plugins, ccache, static_libstdcpp, @@ -1838,6 +1844,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); @@ -2039,6 +2046,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)); diff --git a/src/bootstrap/src/lib.rs b/src/bootstrap/src/lib.rs index 784519a20a2d..5747a77b0ba8 100644 --- a/src/bootstrap/src/lib.rs +++ b/src/bootstrap/src/lib.rs @@ -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, &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), @@ -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, @@ -306,6 +310,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")); @@ -393,6 +398,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()), @@ -848,6 +854,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") } diff --git a/src/bootstrap/src/utils/change_tracker.rs b/src/bootstrap/src/utils/change_tracker.rs index c629f04c00ec..0b0791a7d62c 100644 --- a/src/bootstrap/src/utils/change_tracker.rs +++ b/src/bootstrap/src/utils/change_tracker.rs @@ -230,4 +230,9 @@ pub const CONFIG_CHANGE_HISTORY: &[ChangeInfo] = &[ severity: ChangeSeverity::Warning, summary: "./x test --rustc-args was renamed to --compiletest-rustc-args as it only applies there. ./x miri --rustc-args was also removed.", }, + ChangeInfo { + change_id: 129176, + severity: ChangeSeverity::Info, + summary: "New option `llvm.autodiff` to control whether the llvm based autodiff tool (Enzyme) is built.", + }, ]; diff --git a/src/bootstrap/src/utils/tarball.rs b/src/bootstrap/src/utils/tarball.rs index 3c6c7a7fa180..381aa0c309a8 100644 --- a/src/bootstrap/src/utils/tarball.rs +++ b/src/bootstrap/src/utils/tarball.rs @@ -21,6 +21,7 @@ pub(crate) enum OverlayKind { Cargo, Clippy, Miri, + Enzyme, Rustfmt, Rls, RustAnalyzer, @@ -62,6 +63,7 @@ impl OverlayKind { "src/tools/rust-analyzer/LICENSE-APACHE", "src/tools/rust-analyzer/LICENSE-MIT", ], + OverlayKind::Enzyme => &["src/tools/enzyme/README.md", "src/tools/enzyme/LICENSE"], OverlayKind::RustcCodegenCranelift => &[ "compiler/rustc_codegen_cranelift/Readme.md", "compiler/rustc_codegen_cranelift/LICENSE-APACHE", @@ -94,6 +96,9 @@ impl OverlayKind { OverlayKind::RustAnalyzer => builder .rust_analyzer_info .version(builder, &builder.release_num("rust-analyzer/crates/rust-analyzer")), + OverlayKind::Enzyme => { + builder.enzyme_info.version(builder, &builder.release_num("enzyme")) + } OverlayKind::RustcCodegenCranelift => builder.rust_version(), OverlayKind::LlvmBitcodeLinker => builder.rust_version(), } diff --git a/src/tools/build-manifest/src/main.rs b/src/tools/build-manifest/src/main.rs index 2b263f848e89..2bcc9ee85b91 100644 --- a/src/tools/build-manifest/src/main.rs +++ b/src/tools/build-manifest/src/main.rs @@ -470,6 +470,7 @@ impl Builder { | PkgType::Rls | PkgType::RustAnalyzer | PkgType::Rustfmt + | PkgType::Enzyme | PkgType::LlvmTools | PkgType::RustAnalysis | PkgType::JsonDocs diff --git a/src/tools/build-manifest/src/versions.rs b/src/tools/build-manifest/src/versions.rs index 495cab582eb0..c05966bb8fed 100644 --- a/src/tools/build-manifest/src/versions.rs +++ b/src/tools/build-manifest/src/versions.rs @@ -55,6 +55,7 @@ pkg_type! { RustAnalyzer = "rust-analyzer"; preview = true, Clippy = "clippy"; preview = true, Rustfmt = "rustfmt"; preview = true, + Enzyme = "enzyme"; preview = true, LlvmTools = "llvm-tools"; preview = true, Miri = "miri"; preview = true, JsonDocs = "rust-docs-json"; preview = true, @@ -83,6 +84,8 @@ impl PkgType { PkgType::Rustfmt => false, PkgType::LlvmTools => false, PkgType::Miri => false, + // Not sure? Enzyme has the same version as llvm + PkgType::Enzyme => false, PkgType::RustcCodegenCranelift => false, PkgType::Rust => true, @@ -123,6 +126,7 @@ impl PkgType { Clippy => HOSTS, Miri => HOSTS, Rustfmt => HOSTS, + Enzyme => HOSTS, RustAnalysis => TARGETS, LlvmTools => TARGETS, LlvmBitcodeLinker => HOSTS, diff --git a/src/tools/enzyme b/src/tools/enzyme new file mode 160000 index 000000000000..2fe5164a2423 --- /dev/null +++ b/src/tools/enzyme @@ -0,0 +1 @@ +Subproject commit 2fe5164a2423dd67ef25e2c4fb204fd06362494b diff --git a/src/tools/tidy/config/black.toml b/src/tools/tidy/config/black.toml index e73847a93ba4..a1e1eef987a5 100644 --- a/src/tools/tidy/config/black.toml +++ b/src/tools/tidy/config/black.toml @@ -12,5 +12,6 @@ extend-exclude = """(\ src/llvm-project/|\ src/doc/embedded-book/|\ src/tools/rustc-perf/|\ + src/tools/enzyme/|\ library/backtrace/ )""" diff --git a/src/tools/tidy/config/ruff.toml b/src/tools/tidy/config/ruff.toml index c87c17f78331..2be109010a06 100644 --- a/src/tools/tidy/config/ruff.toml +++ b/src/tools/tidy/config/ruff.toml @@ -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/", @@ -29,6 +30,7 @@ extend-exclude = [ "../src/llvm-project/", "../src/doc/embedded-book/", "../library/backtrace/", + "../src/tools/enzyme/", "../src/tools/rustc-perf/", ] diff --git a/src/tools/tidy/src/walk.rs b/src/tools/tidy/src/walk.rs index d4721b5a5fd1..4a69022ebdc5 100644 --- a/src/tools/tidy/src/walk.rs +++ b/src/tools/tidy/src/walk.rs @@ -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",