Repository: rust-lang/rust
Pull request: Add dist step for Enzyme

<!-- homu-ignore:start -->
*[View all comments](https://triagebot.infra.rust-lang.org/gh-comments/rust-lang/rust/pull/150071)*
<!-- homu-ignore:end -->

Now that Enzyme can be distributed as a separate dylib outside libLLVM.so, we can create a dist step for it. This is the bulk of the work required to make it distributable through rustup, once CI works and we check that the built component works locally, we can add it to Rustup through the manifest.

Note that this PR does not enable Enzyme distribution in CI yet, it just adds the dist step.

r? @ZuseZ4

try-job: dist-x86_64-linux

Changed files:
filename: src/bootstrap/src/core/build_steps/compile.rs
patch: @@ -2292,23 +2292,13 @@ impl Step for Assemble {
             builder.compiler(target_compiler.stage - 1, builder.config.host_target);
 
         // Build enzyme
-        if builder.config.llvm_enzyme && !builder.config.dry_run() {
+        if builder.config.llvm_enzyme {
             debug!("`llvm_enzyme` requested");
-            let enzyme_install = builder.ensure(llvm::Enzyme { target: build_compiler.host });
-            if let Some(llvm_config) = builder.llvm_config(builder.config.host_target) {
-                let llvm_version_major = llvm::get_llvm_version_major(builder, &llvm_config);
-                let lib_ext = std::env::consts::DLL_EXTENSION;
-                let libenzyme = format!("libEnzyme-{llvm_version_major}");
-                let src_lib =
-                    enzyme_install.join("build/Enzyme").join(&libenzyme).with_extension(lib_ext);
-                let libdir = builder.sysroot_target_libdir(build_compiler, build_compiler.host);
-                let target_libdir =
-                    builder.sysroot_target_libdir(target_compiler, target_compiler.host);
-                let dst_lib = libdir.join(&libenzyme).with_extension(lib_ext);
-                let target_dst_lib = target_libdir.join(&libenzyme).with_extension(lib_ext);
-                builder.copy_link(&src_lib, &dst_lib, FileType::NativeLibrary);
-                builder.copy_link(&src_lib, &target_dst_lib, FileType::NativeLibrary);
-            }
+            let enzyme = builder.ensure(llvm::Enzyme { target: build_compiler.host });
+            let target_libdir =
+                builder.sysroot_target_libdir(target_compiler, target_compiler.host);
+            let target_dst_lib = target_libdir.join(enzyme.enzyme_filename());
+            builder.copy_link(&enzyme.enzyme_path(), &target_dst_lib, FileType::NativeLibrary);
         }
 
         if builder.config.llvm_offload && !builder.config.dry_run() {

filename: src/bootstrap/src/core/build_steps/dist.rs
patch: @@ -2717,6 +2717,55 @@ impl Step for LlvmBitcodeLinker {
     }
 }
 
+/// Distributes the `enzyme` library so that it can be used by a compiler whose host
+/// is `target`.
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+pub struct Enzyme {
+    /// Enzyme will by usable by rustc on this host.
+    pub target: TargetSelection,
+}
+
+impl Step for Enzyme {
+    type Output = Option<GeneratedTarball>;
+    const IS_HOST: bool = true;
+
+    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
+        run.alias("enzyme")
+    }
+
+    fn is_default_step(builder: &Builder<'_>) -> bool {
+        builder.config.llvm_enzyme
+    }
+
+    fn make_run(run: RunConfig<'_>) {
+        run.builder.ensure(Enzyme { target: run.target });
+    }
+
+    fn run(self, builder: &Builder<'_>) -> Option<GeneratedTarball> {
+        // This prevents Enzyme from being built for "dist"
+        // or "install" on the stable/beta channels. It is not yet stable and
+        // should not be included.
+        if !builder.build.unstable_features() {
+            return None;
+        }
+
+        let target = self.target;
+
+        let enzyme = builder.ensure(llvm::Enzyme { target });
+
+        let target_libdir = format!("lib/rustlib/{}/lib", target.triple);
+
+        // Prepare the image directory
+        let mut tarball = Tarball::new(builder, "enzyme", &target.triple);
+        tarball.set_overlay(OverlayKind::Enzyme);
+        tarball.is_preview(true);
+
+        tarball.add_file(enzyme.enzyme_path(), target_libdir, FileType::NativeLibrary);
+
+        Some(tarball.generate())
+    }
+}
+
 /// Tarball intended for internal consumption to ease rustc/std development.
 ///
 /// Should not be considered stable by end users.

filename: src/bootstrap/src/core/build_steps/llvm.rs
patch: @@ -17,6 +17,7 @@ use std::{env, fs};
 use build_helper::exit;
 use build_helper::git::PathFreshness;
 
+use crate::core::build_steps::llvm;
 use crate::core::builder::{Builder, RunConfig, ShouldRun, Step, StepMetadata};
 use crate::core::config::{Config, TargetSelection};
 use crate::utils::build_stamp::{BuildStamp, generate_smart_stamp_hash};
@@ -1077,13 +1078,28 @@ impl Step for OmpOffload {
     }
 }
 
+#[derive(Clone)]
+pub struct BuiltEnzyme {
+    /// Path to the libEnzyme dylib.
+    enzyme: PathBuf,
+}
+
+impl BuiltEnzyme {
+    pub fn enzyme_path(&self) -> PathBuf {
+        self.enzyme.clone()
+    }
+    pub fn enzyme_filename(&self) -> String {
+        self.enzyme.file_name().unwrap().to_str().unwrap().to_owned()
+    }
+}
+
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
 pub struct Enzyme {
     pub target: TargetSelection,
 }
 
 impl Step for Enzyme {
-    type Output = PathBuf;
+    type Output = BuiltEnzyme;
     const IS_HOST: bool = true;
 
     fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
@@ -1095,16 +1111,16 @@ impl Step for Enzyme {
     }
 
     /// Compile Enzyme for `target`.
-    fn run(self, builder: &Builder<'_>) -> PathBuf {
+    fn run(self, builder: &Builder<'_>) -> Self::Output {
         builder.require_submodule(
             "src/tools/enzyme",
             Some("The Enzyme sources are required for autodiff."),
         );
+        let target = self.target;
+
         if builder.config.dry_run() {
-            let out_dir = builder.enzyme_out(self.target);
-            return out_dir;
+            return BuiltEnzyme { enzyme: builder.config.tempdir().join("enzyme-dryrun") };
         }
-        let target = self.target;
 
         let LlvmResult { host_llvm_config, llvm_cmake_dir } = builder.ensure(Llvm { target });
 
@@ -1120,6 +1136,12 @@ impl Step for Enzyme {
         let out_dir = builder.enzyme_out(target);
         let stamp = BuildStamp::new(&out_dir).with_prefix("enzyme").add_stamp(smart_stamp_hash);
 
+        let llvm_version_major = llvm::get_llvm_version_major(builder, &host_llvm_config);
+        let lib_ext = std::env::consts::DLL_EXTENSION;
+        let libenzyme = format!("libEnzyme-{llvm_version_major}");
+        let build_dir = out_dir.join("lib");
+        let dylib = build_dir.join(&libenzyme).with_extension(lib_ext);
+
         trace!("checking build stamp to see if we need to rebuild enzyme artifacts");
         if stamp.is_up_to_date() {
             trace!(?out_dir, "enzyme build artifacts are up to date");
@@ -1133,7 +1155,7 @@ impl Step for Enzyme {
                     stamp.path().display()
                 ));
             }
-            return out_dir;
+            return BuiltEnzyme { enzyme: dylib };
         }
 
         if !builder.config.dry_run() && !llvm_cmake_dir.is_dir() {
@@ -1149,7 +1171,6 @@ impl Step for Enzyme {
         let _time = helpers::timeit(builder);
         t!(fs::create_dir_all(&out_dir));
 
-        builder.config.update_submodule("src/tools/enzyme");
         let mut cfg = cmake::Config::new(builder.src.join("src/tools/enzyme/enzyme/"));
         // Enzyme devs maintain upstream compatibility, but only fix deprecations when they are about
         // to turn into a hard error. As such, Enzyme generates various warnings which could make it
@@ -1178,8 +1199,18 @@ impl Step for Enzyme {
 
         cfg.build();
 
+        // At this point, `out_dir` should contain the built libEnzyme-<LLVM-version>.<dylib-ext>
+        // file.
+        if !dylib.exists() {
+            eprintln!(
+                "`{libenzyme}` not found in `{}`. Either the build has failed or Enzyme was built with a wrong version of LLVM",
+                build_dir.display()
+            );
+            exit!(1);
+        }
+
         t!(stamp.write());
-        out_dir
+        BuiltEnzyme { enzyme: dylib }
     }
 }
 

filename: src/bootstrap/src/core/builder/mod.rs
patch: @@ -980,6 +980,7 @@ impl<'a> Builder<'a> {
                 dist::LlvmTools,
                 dist::LlvmBitcodeLinker,
                 dist::RustDev,
+                dist::Enzyme,
                 dist::Bootstrap,
                 dist::Extended,
                 // It seems that PlainSourceTarball somehow changes how some of the tools

filename: src/bootstrap/src/utils/tarball.rs
patch: @@ -28,6 +28,7 @@ pub(crate) enum OverlayKind {
     RustcCodegenGcc,
     Gcc,
     LlvmBitcodeLinker,
+    Enzyme,
 }
 
 impl OverlayKind {
@@ -37,6 +38,7 @@ impl OverlayKind {
             OverlayKind::Llvm => {
                 &["src/llvm-project/llvm/LICENSE.TXT", "src/llvm-project/llvm/README.txt"]
             }
+            OverlayKind::Enzyme => &["src/tools/enzyme/LICENSE", "src/tools/enzyme/Readme.md"],
             OverlayKind::Cargo => &[
                 "src/tools/cargo/README.md",
                 "src/tools/cargo/LICENSE-MIT",
@@ -111,6 +113,7 @@ impl OverlayKind {
             OverlayKind::RustcCodegenGcc => builder.rust_version(),
             OverlayKind::LlvmBitcodeLinker => builder.rust_version(),
             OverlayKind::Gcc => builder.rust_version(),
+            OverlayKind::Enzyme => builder.rust_version(),
         }
     }
 }