Auto merge of #151063 - sgasho:aarch64-dist-enzyme, r=<try>

Add dist step for enzyme, including aarch64-apple-darwin


try-job: dist-aarch64-apple

Author: bors

.gitmodules

@@ -41,7 +41,7 @@
 	shallow = true
 [submodule "src/tools/enzyme"]
 	path = src/tools/enzyme
-	url = https://github.com/rust-lang/enzyme.git
+	url = https://github.com/sgasho/enzyme.git
 	shallow = true
 [submodule "src/gcc"]
 	path = src/gcc

compiler/rustc_codegen_ssa/src/back/linker.rs

@@ -796,6 +796,14 @@ impl<'a> Linker for GccLinker<'a> {
             }
         }
 
+        if crate_type == CrateType::Dylib
+            && self.sess.target.is_like_darwin
+            && self.sess.opts.unstable_opts.export_llvm_symbols
+            && self.sess.opts.crate_name.as_deref() == Some("rustc_driver")
+        {
+            return;
+        }
+
         // We manually create a list of exported symbols to ensure we don't expose any more.
         // The object files have far more public symbols than we actually want to export,
         // so we hide them all here.

compiler/rustc_session/src/options.rs

@@ -2367,6 +2367,8 @@ options! {
         "enable default bounds for experimental group of auto traits"),
     export_executable_symbols: bool = (false, parse_bool, [TRACKED],
         "export symbols from executables, as if they were dynamic libraries"),
+    export_llvm_symbols: bool = (false, parse_bool, [TRACKED],
+        "export LLVM symbols from rustc_driver on darwin"),
     external_clangrt: bool = (false, parse_bool, [UNTRACKED],
         "rely on user specified linker commands to find clangrt"),
     extra_const_ub_checks: bool = (false, parse_bool, [TRACKED],

src/bootstrap/src/core/build_steps/compile.rs

@@ -1261,6 +1261,10 @@ pub fn rustc_cargo(
         cargo.rustflag("-Zdefault-visibility=protected");
     }
 
+    if builder.config.llvm_enzyme && target.contains("apple") && build_compiler.stage != 0 {
+        cargo.rustflag("-Zexport-llvm-symbols");
+    }
+
     if is_lto_stage(build_compiler) {
         match builder.config.rust_lto {
             RustcLto::Thin | RustcLto::Fat => {
@@ -2292,23 +2296,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() {

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

@@ -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.

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

@@ -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 }
     }
 }
 

src/bootstrap/src/core/builder/mod.rs

@@ -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

src/bootstrap/src/utils/tarball.rs

@@ -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(),
         }
     }
 }

src/ci/docker/host-x86_64/dist-x86_64-linux/Dockerfile

@@ -91,6 +91,7 @@ ENV RUST_CONFIGURE_ARGS \
       --set llvm.thin-lto=true \
       --set llvm.ninja=false \
       --set llvm.libzstd=true \
+      --set llvm.enzyme=true \
       --set rust.jemalloc \
       --set rust.bootstrap-override-lld=true \
       --set rust.lto=thin \

src/ci/github-actions/jobs.yml

@@ -486,6 +486,7 @@ auto:
         --enable-full-tools
         --enable-sanitizers
         --enable-profiler
+        --set llvm.enzyme=true
         --set rust.jemalloc
         --set rust.lto=thin
         --set rust.codegen-units=1