Auto merge of #140799 - EnzymeAD:enzyme-static, r=<try>

[DO NOT MERGE] build Enzyme as archive

r? ghost

try-job: dist-x86_64-linux

Author: bors

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

@@ -1229,7 +1229,7 @@ pub fn rustc_cargo(
 
         if let Some(llvm_config) = builder.llvm_config(builder.config.host_target) {
             let llvm_version_major = llvm::get_llvm_version_major(builder, &llvm_config);
-            cargo.rustflag("-l").rustflag(&format!("Enzyme-{llvm_version_major}"));
+            cargo.rustflag("-l").rustflag(&format!("EnzymeStatic-{llvm_version_major}"));
         }
     }
 
@@ -2124,8 +2124,8 @@ impl Step for Assemble {
             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 lib_ext = "a";
+                let libenzyme = format!("libEnzymeStatic-{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);

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

@@ -541,7 +541,6 @@ 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
@@ -977,6 +976,7 @@ impl Step for Enzyme {
             .env("LLVM_CONFIG_REAL", &llvm_config)
             .define("LLVM_ENABLE_ASSERTIONS", "ON")
             .define("ENZYME_EXTERNAL_SHARED_LIB", "ON")
+            .define("ENZYME_STATIC_LIB", "ON")
             .define("LLVM_DIR", builder.llvm_out(target));
 
         cfg.build();

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

@@ -89,6 +89,7 @@ ENV RUST_CONFIGURE_ARGS \
       --set llvm.thin-lto=true \
       --set llvm.libzstd=true \
       --set llvm.ninja=false \
+      --set llvm.enzyme=true \
       --set rust.debug-assertions=false \
       --set rust.jemalloc \
       --set rust.use-lld=true \

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.use-lld=true \
       --set rust.lto=thin \

src/ci/docker/host-x86_64/x86_64-gnu/Dockerfile

@@ -29,5 +29,6 @@ ENV RUST_CONFIGURE_ARGS \
  --enable-sanitizers \
  --enable-profiler \
  --enable-compiler-docs \
- --set llvm.libzstd=true
+ --set llvm.libzstd=true \
+ --set llvm.enzyme=true
 ENV SCRIPT python3 ../x.py --stage 2 test

src/ci/github-actions/jobs.yml

@@ -437,7 +437,7 @@ auto:
   - name: dist-x86_64-apple
     env:
       SCRIPT: ./x.py dist bootstrap --include-default-paths --host=x86_64-apple-darwin --target=x86_64-apple-darwin
-      RUST_CONFIGURE_ARGS: --enable-full-tools --enable-sanitizers --enable-profiler --set rust.jemalloc --set rust.lto=thin --set rust.codegen-units=1
+      RUST_CONFIGURE_ARGS: --enable-full-tools --enable-sanitizers --enable-profiler --set rust.jemalloc --set rust.lto=thin --set llvm.enzyme=true --set rust.codegen-units=1
       # Ensure that host tooling is built to support our minimum support macOS version.
       MACOSX_DEPLOYMENT_TARGET: 10.12
       MACOSX_STD_DEPLOYMENT_TARGET: 10.12
@@ -454,7 +454,7 @@ auto:
       SCRIPT: ./x.py dist bootstrap --include-default-paths --host='' --target=aarch64-apple-ios,x86_64-apple-ios,aarch64-apple-ios-sim,aarch64-apple-ios-macabi,x86_64-apple-ios-macabi
       # Mac Catalyst cannot currently compile the sanitizer:
       # https://github.com/rust-lang/rust/issues/129069
-      RUST_CONFIGURE_ARGS: --enable-sanitizers --enable-profiler --set rust.jemalloc --set target.aarch64-apple-ios-macabi.sanitizers=false --set target.x86_64-apple-ios-macabi.sanitizers=false
+      RUST_CONFIGURE_ARGS: --enable-sanitizers --enable-profiler --set rust.jemalloc --set target.aarch64-apple-ios-macabi.sanitizers=false --set target.x86_64-apple-ios-macabi.sanitizers=false --set llvm.enzyme=true
       # Ensure that host tooling is built to support our minimum support macOS version.
       # FIXME(madsmtm): This might be redundant, as we're not building host tooling here (?)
       MACOSX_DEPLOYMENT_TARGET: 10.12
@@ -485,6 +485,7 @@ auto:
         --enable-profiler
         --set rust.jemalloc
         --set llvm.ninja=false
+        --set llvm.enzyme=true
         --set rust.lto=thin
         --set rust.codegen-units=1
       SELECT_XCODE: /Applications/Xcode_15.4.app

src/tools/enzyme

@@ -18,18 +18,13 @@
 // but each shadow argument is `width` times larger (thus 16 and 20 elements here).
 // `d_square3` instead takes `width` (4) shadow arguments, which are all the same size as the
 // original function arguments.
-//
-// FIXME(autodiff): We currently can't test `d_square1` and `d_square3` in the same file, since they
-// generate the same dummy functions which get merged by LLVM, breaking pieces of our pipeline which
-// try to rewrite the dummy functions later. We should consider to change to pure declarations both
-// in our frontend and in the llvm backend to avoid these issues.
 
 #![feature(autodiff)]
 
 use std::autodiff::autodiff;
 
 #[no_mangle]
-//#[autodiff(d_square1, Forward, Dual, Dual)]
+#[autodiff(d_square1, Forward, Dual, Dual)]
 #[autodiff(d_square2, Forward, 4, Dualv, Dualv)]
 #[autodiff(d_square3, Forward, 4, Dual, Dual)]
 fn square(x: &[f32], y: &mut [f32]) {
@@ -42,6 +37,9 @@ fn square(x: &[f32], y: &mut [f32]) {
     y[4] = 1.0 * x[0] + 2.0 * x[1] + 3.0 * x[2] + 4.0 * x[3];
 }
 
+// FIXME
+// CHECK: start:
+
 fn main() {
     let x1 = std::hint::black_box(vec![0.0, 1.0, 2.0, 3.0]);
 
@@ -78,25 +76,25 @@ fn main() {
     let mut dy3_4 = std::hint::black_box(vec![0.0; 5]);
 
     // scalar.
-    //d_square1(&x1, &z1, &mut y1, &mut dy1_1);
-    //d_square1(&x1, &z2, &mut y2, &mut dy1_2);
-    //d_square1(&x1, &z3, &mut y3, &mut dy1_3);
-    //d_square1(&x1, &z4, &mut y4, &mut dy1_4);
+    d_square1(&x1, &z1, &mut y1, &mut dy1_1);
+    d_square1(&x1, &z2, &mut y2, &mut dy1_2);
+    d_square1(&x1, &z3, &mut y3, &mut dy1_3);
+    d_square1(&x1, &z4, &mut y4, &mut dy1_4);
 
     // assert y1 == y2 == y3 == y4
-    //for i in 0..5 {
-    //    assert_eq!(y1[i], y2[i]);
-    //    assert_eq!(y1[i], y3[i]);
-    //    assert_eq!(y1[i], y4[i]);
-    //}
+    for i in 0..5 {
+        assert_eq!(y1[i], y2[i]);
+        assert_eq!(y1[i], y3[i]);
+        assert_eq!(y1[i], y4[i]);
+    }
 
     // batch mode A)
     d_square2(&x1, &z5, &mut y5, &mut dy2);
 
     // assert y1 == y2 == y3 == y4 == y5
-    //for i in 0..5 {
-    //    assert_eq!(y1[i], y5[i]);
-    //}
+    for i in 0..5 {
+        assert_eq!(y1[i], y5[i]);
+    }
 
     // batch mode B)
     d_square3(&x1, &z1, &z2, &z3, &z4, &mut y6, &mut dy3_1, &mut dy3_2, &mut dy3_3, &mut dy3_4);

tests/codegen/autodiffv2.rs renamed to tests/codegen/autodiff/autodiffv2.rs

@@ -0,0 +1,113 @@
+//@ compile-flags: -Zautodiff=Enable -C opt-level=3  -Clto=fat
+//@ no-prefer-dynamic
+//@ needs-enzyme
+//
+// In Enzyme, we test against a large range of LLVM versions (5+) and don't have overly many
+// breakages. One benefit is that we match the IR generated by Enzyme only after running it
+// through LLVM's O3 pipeline, which will remove most of the noise.
+// However, our integration test could also be affected by changes in how rustc lowers MIR into
+// LLVM-IR, which could cause additional noise and thus breakages. If that's the case, we should
+// reduce this test to only match the first lines and the ret instructions.
+//
+// The function tested here has 4 inputs and 5 outputs, so we could either call forward-mode
+// autodiff 4 times, or reverse mode 5 times. Since a forward-mode call is usually faster than
+// reverse mode, we prefer it here. This file also tests a new optimization (batch mode), which
+// allows us to call forward-mode autodiff only once, and get all 5 outputs in a single call.
+//
+// We support 2 different batch modes. `d_square2` has the same interface as scalar forward-mode,
+// but each shadow argument is `width` times larger (thus 16 and 20 elements here).
+// `d_square3` instead takes `width` (4) shadow arguments, which are all the same size as the
+// original function arguments.
+//
+// FIXME(autodiff): We currently can't test `d_square1` and `d_square3` in the same file, since they
+// generate the same dummy functions which get merged by LLVM, breaking pieces of our pipeline which
+// try to rewrite the dummy functions later. We should consider to change to pure declarations both
+// in our frontend and in the llvm backend to avoid these issues.
+
+#![feature(autodiff)]
+
+use std::autodiff::autodiff;
+
+#[no_mangle]
+#[autodiff(d_square1, Forward, Dual, Dual)]
+#[autodiff(d_square2, Forward, 4, Dualv, Dualv)]
+#[autodiff(d_square3, Forward, 4, Dual, Dual)]
+fn square(x: &[f32], y: &mut [f32]) {
+    assert!(x.len() >= 4);
+    assert!(y.len() >= 5);
+    y[0] = 4.3 * x[0] + 1.2 * x[1] + 3.4 * x[2] + 2.1 * x[3];
+    y[1] = 2.3 * x[0] + 4.5 * x[1] + 1.7 * x[2] + 6.4 * x[3];
+    y[2] = 1.1 * x[0] + 3.3 * x[1] + 2.5 * x[2] + 4.7 * x[3];
+    y[3] = 5.2 * x[0] + 1.4 * x[1] + 2.6 * x[2] + 3.8 * x[3];
+    y[4] = 1.0 * x[0] + 2.0 * x[1] + 3.0 * x[2] + 4.0 * x[3];
+}
+
+fn main() {
+    let x1 = std::hint::black_box(vec![0.0, 1.0, 2.0, 3.0]);
+
+    let dx1 = std::hint::black_box(vec![1.0; 12]);
+
+    let z1 = std::hint::black_box(vec![1.0, 0.0, 0.0, 0.0]);
+    let z2 = std::hint::black_box(vec![0.0, 1.0, 0.0, 0.0]);
+    let z3 = std::hint::black_box(vec![0.0, 0.0, 1.0, 0.0]);
+    let z4 = std::hint::black_box(vec![0.0, 0.0, 0.0, 1.0]);
+
+    let z5 = std::hint::black_box(vec![
+        1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0,
+    ]);
+
+    let mut y1 = std::hint::black_box(vec![0.0; 5]);
+    let mut y2 = std::hint::black_box(vec![0.0; 5]);
+    let mut y3 = std::hint::black_box(vec![0.0; 5]);
+    let mut y4 = std::hint::black_box(vec![0.0; 5]);
+
+    let mut y5 = std::hint::black_box(vec![0.0; 5]);
+
+    let mut y6 = std::hint::black_box(vec![0.0; 5]);
+
+    let mut dy1_1 = std::hint::black_box(vec![0.0; 5]);
+    let mut dy1_2 = std::hint::black_box(vec![0.0; 5]);
+    let mut dy1_3 = std::hint::black_box(vec![0.0; 5]);
+    let mut dy1_4 = std::hint::black_box(vec![0.0; 5]);
+
+    let mut dy2 = std::hint::black_box(vec![0.0; 20]);
+
+    let mut dy3_1 = std::hint::black_box(vec![0.0; 5]);
+    let mut dy3_2 = std::hint::black_box(vec![0.0; 5]);
+    let mut dy3_3 = std::hint::black_box(vec![0.0; 5]);
+    let mut dy3_4 = std::hint::black_box(vec![0.0; 5]);
+
+    // scalar.
+    d_square1(&x1, &z1, &mut y1, &mut dy1_1);
+    d_square1(&x1, &z2, &mut y2, &mut dy1_2);
+    d_square1(&x1, &z3, &mut y3, &mut dy1_3);
+    d_square1(&x1, &z4, &mut y4, &mut dy1_4);
+
+    // assert y1 == y2 == y3 == y4
+    for i in 0..5 {
+        assert_eq!(y1[i], y2[i]);
+        assert_eq!(y1[i], y3[i]);
+        assert_eq!(y1[i], y4[i]);
+    }
+
+    // batch mode A)
+    d_square2(&x1, &z5, &mut y5, &mut dy2);
+
+    // assert y1 == y2 == y3 == y4 == y5
+    for i in 0..5 {
+        assert_eq!(y1[i], y5[i]);
+    }
+
+    // batch mode B)
+    d_square3(&x1, &z1, &z2, &z3, &z4, &mut y6, &mut dy3_1, &mut dy3_2, &mut dy3_3, &mut dy3_4);
+    for i in 0..5 {
+        assert_eq!(y5[i], y6[i]);
+    }
+
+    for i in 0..5 {
+        assert_eq!(dy2[0..5][i], dy3_1[i]);
+        assert_eq!(dy2[5..10][i], dy3_2[i]);
+        assert_eq!(dy2[10..15][i], dy3_3[i]);
+        assert_eq!(dy2[15..20][i], dy3_4[i]);
+    }
+}