Rollup of 8 pull requests

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -77,6 +77,17 @@ pub struct AutoDiffAttrs {
     /// e.g. in the [JAX
     /// Documentation](https://jax.readthedocs.io/en/latest/_tutorials/advanced-autodiff.html#how-it-s-made-two-foundational-autodiff-functions).
     pub mode: DiffMode,
+    /// A user-provided, batching width. If not given, we will default to 1 (no batching).
+    /// Calling a differentiated, non-batched function through a loop 100 times is equivalent to:
+    /// - Calling the function 50 times with a batch size of 2
+    /// - Calling the function 25 times with a batch size of 4,
+    /// etc. A batched function takes more (or longer) arguments, and might be able to benefit from
+    /// cache locality, better re-usal of primal values, and other optimizations.
+    /// We will (before LLVM's vectorizer runs) just generate most LLVM-IR instructions `width`
+    /// times, so this massively increases code size. As such, values like 1024 are unlikely to
+    /// work. We should consider limiting this to u8 or u16, but will leave it at u32 for
+    /// experiments for now and focus on documenting the implications of a large width.
+    pub width: u32,
     pub ret_activity: DiffActivity,
     pub input_activity: Vec<DiffActivity>,
 }
@@ -222,13 +233,15 @@ impl AutoDiffAttrs {
     pub const fn error() -> Self {
         AutoDiffAttrs {
             mode: DiffMode::Error,
+            width: 0,
             ret_activity: DiffActivity::None,
             input_activity: Vec::new(),
         }
     }
     pub fn source() -> Self {
         AutoDiffAttrs {
             mode: DiffMode::Source,
+            width: 0,
             ret_activity: DiffActivity::None,
             input_activity: Vec::new(),
         }

compiler/rustc_builtin_macros/messages.ftl

@@ -79,6 +79,7 @@ builtin_macros_autodiff_ret_activity = invalid return activity {$act} in {$mode}
 builtin_macros_autodiff_ty_activity = {$act} can not be used for this type
 builtin_macros_autodiff_unknown_activity = did not recognize Activity: `{$act}`
 
+builtin_macros_autodiff_width = autodiff width must fit u32, but is {$width}
 builtin_macros_bad_derive_target = `derive` may only be applied to `struct`s, `enum`s and `union`s
     .label = not applicable here
     .label2 = not a `struct`, `enum` or `union`

compiler/rustc_builtin_macros/src/errors.rs

@@ -202,6 +202,14 @@ mod autodiff {
         pub(crate) mode: String,
     }
 
+    #[derive(Diagnostic)]
+    #[diag(builtin_macros_autodiff_width)]
+    pub(crate) struct AutoDiffInvalidWidth {
+        #[primary_span]
+        pub(crate) span: Span,
+        pub(crate) width: u128,
+    }
+
     #[derive(Diagnostic)]
     #[diag(builtin_macros_autodiff)]
     pub(crate) struct AutoDiffInvalidApplication {

compiler/rustc_codegen_cranelift/build_system/tests.rs

@@ -99,6 +99,34 @@ const BASE_SYSROOT_SUITE: &[TestCase] = &[
         runner.run_out_command("gen_block_iterate", &[]);
     }),
     TestCase::build_bin_and_run("aot.raw-dylib", "example/raw-dylib.rs", &[]),
+    TestCase::custom("test.sysroot", &|runner| {
+        apply_patches(
+            &runner.dirs,
+            "sysroot_tests",
+            &runner.stdlib_source.join("library"),
+            &SYSROOT_TESTS_SRC.to_path(&runner.dirs),
+        );
+
+        SYSROOT_TESTS.clean(&runner.dirs);
+
+        let mut target_compiler = runner.target_compiler.clone();
+        // coretests and alloctests produce a bunch of warnings. When running
+        // in rust's CI warnings are denied, so we have to override that here.
+        target_compiler.rustflags.push("--cap-lints=allow".to_owned());
+        // The standard library may have been compiled with -Zrandomize-layout.
+        target_compiler.rustflags.extend(["--cfg".to_owned(), "randomized_layouts".to_owned()]);
+
+        if runner.is_native {
+            let mut test_cmd = SYSROOT_TESTS.test(&target_compiler, &runner.dirs);
+            test_cmd.args(["-p", "coretests", "-p", "alloctests", "--tests", "--", "-q"]);
+            spawn_and_wait(test_cmd);
+        } else {
+            eprintln!("Cross-Compiling: Not running tests");
+            let mut build_cmd = SYSROOT_TESTS.build(&target_compiler, &runner.dirs);
+            build_cmd.args(["-p", "coretests", "-p", "alloctests", "--tests"]);
+            spawn_and_wait(build_cmd);
+        }
+    }),
 ];
 
 pub(crate) static RAND_REPO: GitRepo = GitRepo::github(
@@ -146,27 +174,6 @@ const EXTENDED_SYSROOT_SUITE: &[TestCase] = &[
             spawn_and_wait(build_cmd);
         }
     }),
-    TestCase::custom("test.sysroot", &|runner| {
-        apply_patches(
-            &runner.dirs,
-            "sysroot_tests",
-            &runner.stdlib_source.join("library"),
-            &SYSROOT_TESTS_SRC.to_path(&runner.dirs),
-        );
-
-        SYSROOT_TESTS.clean(&runner.dirs);
-
-        if runner.is_native {
-            let mut test_cmd = SYSROOT_TESTS.test(&runner.target_compiler, &runner.dirs);
-            test_cmd.args(["-p", "coretests", "-p", "alloctests", "--", "-q"]);
-            spawn_and_wait(test_cmd);
-        } else {
-            eprintln!("Cross-Compiling: Not running tests");
-            let mut build_cmd = SYSROOT_TESTS.build(&runner.target_compiler, &runner.dirs);
-            build_cmd.args(["-p", "coretests", "-p", "alloctests", "--tests"]);
-            spawn_and_wait(build_cmd);
-        }
-    }),
     TestCase::custom("test.regex", &|runner| {
         REGEX_REPO.patch(&runner.dirs);
 

compiler/rustc_codegen_cranelift/build_system/utils.rs

@@ -105,7 +105,11 @@ impl CargoProject {
             .arg(self.manifest_path(dirs))
             .arg("--target-dir")
             .arg(self.target_dir(dirs))
-            .arg("--locked");
+            .arg("--locked")
+            // bootstrap sets both RUSTC and RUSTC_WRAPPER to the same wrapper. RUSTC is already
+            // respected by the rustc-clif wrapper, but RUSTC_WRAPPER will misinterpret rustc-clif
+            // as filename, so we need to unset it.
+            .env_remove("RUSTC_WRAPPER");
 
         if dirs.frozen {
             cmd.arg("--frozen");

compiler/rustc_codegen_cranelift/config.txt

@@ -32,9 +32,9 @@ aot.issue-59326
 aot.neon
 aot.gen_block_iterate
 aot.raw-dylib
+test.sysroot
 
 testsuite.extended_sysroot
 test.rust-random/rand
-test.sysroot
 test.regex
 test.portable-simd

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -610,6 +610,8 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff], module: &mut ModuleCodegen<
             }
             // We handle this below
             config::AutoDiff::PrintModAfter => {}
+            // We handle this below
+            config::AutoDiff::PrintModFinal => {}
             // This is required and already checked
             config::AutoDiff::Enable => {}
         }
@@ -657,14 +659,20 @@ pub(crate) fn run_pass_manager(
     }
 
     if cfg!(llvm_enzyme) && enable_ad {
+        // This is the post-autodiff IR, mainly used for testing and educational purposes.
+        if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
+            unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
+        }
+
         let opt_stage = llvm::OptStage::FatLTO;
         let stage = write::AutodiffStage::PostAD;
         unsafe {
             write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
         }
 
-        // This is the final IR, so people should be able to inspect the optimized autodiff output.
-        if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
+        // This is the final IR, so people should be able to inspect the optimized autodiff output,
+        // for manual inspection.
+        if config.autodiff.contains(&config::AutoDiff::PrintModFinal) {
             unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
         }
     }