enzyme backend

Author: ZuseZ4

.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 = git@github.com:EnzymeAD/Enzyme.git
+	shallow = true

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,29 @@ 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 = match env::consts::OS {
+                "macos" => "dylib",
+                "windows" => "dll",
+                _ => "so",
+            };
+
+            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

src/bootstrap/src/core/build_steps/dist.rs

@@ -1522,6 +1522,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),
@@ -2400,6 +2401,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<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 {
+            compiler: run.builder.compiler_for(
+                run.builder.top_stage,
+                run.builder.config.build,
+                run.target,
+            ),
+            target: run.target,
+        });
+    }
+
+    fn run(self, builder: &Builder<'_>) -> Option<GeneratedTarball> {
+        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.

src/bootstrap/src/core/build_steps/llvm.rs

@@ -529,6 +529,7 @@ impl Step for Llvm {
             }
         };
 
+        // Manuel: TODO: 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,76 @@ 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 });
+
+        let out_dir = builder.enzyme_out(target);
+        let done_stamp = out_dir.join("enzyme-finished-building");
+        if done_stamp.exists() {
+            return out_dir;
+        }
+
+        builder.info(&format!("Building Enzyme for {}", target));
+        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/"));
+        // TODO: 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!(File::create(&done_stamp));
+        out_dir
+    }
+}
+
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
 pub struct Lld {
     pub target: TargetSelection,

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";

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,
@@ -1582,6 +1583,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

src/bootstrap/src/core/config/config.rs

@@ -213,6 +213,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,
@@ -876,6 +877,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",
@@ -1561,6 +1563,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;
@@ -1707,6 +1710,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);
@@ -1791,6 +1796,7 @@ impl Config {
                 release_debuginfo,
                 assertions,
                 tests,
+                enzyme,
                 plugins,
                 ccache,
                 static_libstdcpp,
@@ -1824,6 +1830,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);
@@ -2025,6 +2032,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));
 

src/bootstrap/src/lib.rs

@@ -75,6 +75,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),
@@ -138,6 +141,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,
@@ -304,6 +308,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"));
@@ -391,6 +396,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()),
@@ -849,6 +855,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.triple).join("lld")
     }

src/bootstrap/src/utils/tarball.rs

@@ -20,6 +20,7 @@ pub(crate) enum OverlayKind {
     Cargo,
     Clippy,
     Miri,
+    Enzyme,
     Rustfmt,
     Rls,
     RustAnalyzer,
@@ -61,6 +62,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",
@@ -93,6 +95,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(),
         }

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