From bc2a913a9bb89b137a43cc761793ae70886dfd14 Mon Sep 17 00:00:00 2001
From: Manuel Drehwald <git@manuel.drehwald.info>
Date: Sun, 29 Sep 2024 18:27:33 -0400
Subject: [PATCH] add has_enzyme/needs-enzyme to the test infra

---
 src/bootstrap/src/core/build_steps/test.rs | 4 ++++
 src/tools/compiletest/src/command-list.rs  | 1 +
 src/tools/compiletest/src/common.rs        | 3 +++
 src/tools/compiletest/src/header.rs        | 3 +++
 src/tools/compiletest/src/header/needs.rs  | 5 +++++
 src/tools/compiletest/src/lib.rs           | 3 +++
 6 files changed, 19 insertions(+)

diff --git a/src/bootstrap/src/core/build_steps/test.rs b/src/bootstrap/src/core/build_steps/test.rs
index 870fe6a9f16..7283b0e9574 100644
--- a/src/bootstrap/src/core/build_steps/test.rs
+++ b/src/bootstrap/src/core/build_steps/test.rs
@@ -1792,6 +1792,10 @@ NOTE: if you're sure you want to do this, please open an issue as to why. In the
         cmd.arg("--host").arg(&*compiler.host.triple);
         cmd.arg("--llvm-filecheck").arg(builder.llvm_filecheck(builder.config.build));
 
+        if builder.build.config.llvm_enzyme {
+            cmd.arg("--has-enzyme");
+        }
+
         if builder.config.cmd.bless() {
             cmd.arg("--bless");
         }
diff --git a/src/tools/compiletest/src/command-list.rs b/src/tools/compiletest/src/command-list.rs
index 865aa76ddb0..a4cedbf66e2 100644
--- a/src/tools/compiletest/src/command-list.rs
+++ b/src/tools/compiletest/src/command-list.rs
@@ -139,6 +139,7 @@ const KNOWN_DIRECTIVE_NAMES: &[&str] = &[
     "needs-deterministic-layouts",
     "needs-dlltool",
     "needs-dynamic-linking",
+    "needs-enzyme",
     "needs-force-clang-based-tests",
     "needs-git-hash",
     "needs-llvm-components",
diff --git a/src/tools/compiletest/src/common.rs b/src/tools/compiletest/src/common.rs
index 2d8c0c3fa5e..adc89cad72f 100644
--- a/src/tools/compiletest/src/common.rs
+++ b/src/tools/compiletest/src/common.rs
@@ -349,6 +349,9 @@ pub struct Config {
     /// whether to run `tidy` when a rustdoc test fails
     pub has_tidy: bool,
 
+    /// whether to run `enzyme` autodiff tests
+    pub has_enzyme: bool,
+
     /// The current Rust channel
     pub channel: String,
 
diff --git a/src/tools/compiletest/src/header.rs b/src/tools/compiletest/src/header.rs
index 6a889d27793..83a10c56208 100644
--- a/src/tools/compiletest/src/header.rs
+++ b/src/tools/compiletest/src/header.rs
@@ -218,6 +218,8 @@ pub struct TestProps {
     pub filecheck_flags: Vec<String>,
     /// Don't automatically insert any `--check-cfg` args
     pub no_auto_check_cfg: bool,
+    /// Run tests which require enzyme being build
+    pub has_enzyme: bool,
 }
 
 mod directives {
@@ -322,6 +324,7 @@ impl TestProps {
             llvm_cov_flags: vec![],
             filecheck_flags: vec![],
             no_auto_check_cfg: false,
+            has_enzyme: false,
         }
     }
 
diff --git a/src/tools/compiletest/src/header/needs.rs b/src/tools/compiletest/src/header/needs.rs
index 99c0e850f1a..f5dd722ed37 100644
--- a/src/tools/compiletest/src/header/needs.rs
+++ b/src/tools/compiletest/src/header/needs.rs
@@ -79,6 +79,11 @@ pub(super) fn handle_needs(
             condition: cache.sanitizer_safestack,
             ignore_reason: "ignored on targets without SafeStack support",
         },
+        Need {
+            name: "needs-enzyme",
+            condition: config.has_enzyme,
+            ignore_reason: "ignored when LLVM Enzyme is disabled",
+        },
         Need {
             name: "needs-run-enabled",
             condition: config.run_enabled(),
diff --git a/src/tools/compiletest/src/lib.rs b/src/tools/compiletest/src/lib.rs
index cfc619f9342..a8355ee9590 100644
--- a/src/tools/compiletest/src/lib.rs
+++ b/src/tools/compiletest/src/lib.rs
@@ -83,6 +83,7 @@ pub fn parse_config(args: Vec<String>) -> Config {
         )
         .optopt("", "run", "whether to execute run-* tests", "auto | always | never")
         .optflag("", "ignored", "run tests marked as ignored")
+        .optflag("", "has-enzyme", "run tests that require enzyme")
         .optflag("", "with-debug-assertions", "whether to run tests with `ignore-debug` header")
         .optmulti(
             "",
@@ -233,6 +234,7 @@ pub fn parse_config(args: Vec<String>) -> Config {
         // Avoid spawning an external command when we know tidy won't be used.
         false
     };
+    let has_enzyme = matches.opt_present("has-enzyme");
     let filters = if mode == Mode::RunMake {
         matches
             .free
@@ -331,6 +333,7 @@ pub fn parse_config(args: Vec<String>) -> Config {
             .map(|s| s.parse().expect("invalid --compare-mode provided")),
         rustfix_coverage: matches.opt_present("rustfix-coverage"),
         has_tidy,
+        has_enzyme,
         channel: matches.opt_str("channel").unwrap(),
         git_hash: matches.opt_present("git-hash"),
         edition: matches.opt_str("edition"),