Merge pull request #2139 from Kobzol/cli-match

Kobzol · web-flow · commit 9863312596d1 · 2025-05-27T11:58:35.000+02:00
Add `--exact-match` CLI argument to allow exact matching of benchmarks
diff --git a/collector/README.md b/collector/README.md
@@ -135,6 +135,9 @@ The following options alter the behaviour of the `bench_local` subcommand.
   dedicated to artifact sizes (ending with `-tiny`).
 - `--id <ID>` the identifier that will be used to identify the results in the
   database.
+- `--exact-match <BENCHMARKS>`: comma-separated list of benchmark names that should be
+  executed. The names have to match exactly. Cannot be combined with
+  `--include`/`--exclude`/`--exclude-suffix`.
 - `--include <INCLUDE>`: the inverse of `--exclude`. The argument is a
   comma-separated list of benchmark prefixes. When this option is specified, a
   benchmark is included in the run only if its name matches one of the given
diff --git a/collector/src/bin/collector.rs b/collector/src/bin/collector.rs
@@ -40,13 +40,14 @@ use collector::compile::benchmark::scenario::Scenario;
 use collector::compile::benchmark::target::Target;
 use collector::compile::benchmark::{
     compile_benchmark_dir, get_compile_benchmarks, ArtifactType, Benchmark, BenchmarkName,
+    CompileBenchmarkFilter,
 };
 use collector::compile::execute::bencher::BenchProcessor;
 use collector::compile::execute::profiler::{ProfileProcessor, Profiler};
 use collector::runtime::{
     bench_runtime, get_runtime_benchmark_groups, prepare_runtime_benchmark_suite,
-    runtime_benchmark_dir, BenchmarkFilter, BenchmarkSuite, BenchmarkSuiteCompilation,
-    CargoIsolationMode, RuntimeProfiler, DEFAULT_RUNTIME_ITERATIONS,
+    runtime_benchmark_dir, BenchmarkSuite, BenchmarkSuiteCompilation, CargoIsolationMode,
+    RuntimeBenchmarkFilter, RuntimeProfiler, DEFAULT_RUNTIME_ITERATIONS,
 };
 use collector::runtime::{profile_runtime, RuntimeCompilationOpts};
 use collector::toolchain::{
@@ -105,12 +106,12 @@ struct CompileBenchmarkConfig {
 
 struct RuntimeBenchmarkConfig {
     runtime_suite: BenchmarkSuite,
-    filter: BenchmarkFilter,
+    filter: RuntimeBenchmarkFilter,
     iterations: u32,
 }
 
 impl RuntimeBenchmarkConfig {
-    fn new(suite: BenchmarkSuite, filter: BenchmarkFilter, iterations: u32) -> Self {
+    fn new(suite: BenchmarkSuite, filter: RuntimeBenchmarkFilter, iterations: u32) -> Self {
         Self {
             runtime_suite: suite.filter(&filter),
             filter,
@@ -337,6 +338,16 @@ struct LocalOptions {
     #[arg(long, value_delimiter = ',')]
     include: Vec<String>,
 
+    /// Include only benchmarks in this comma-separated list
+    #[arg(
+        long,
+        value_delimiter = ',',
+        conflicts_with("include"),
+        conflicts_with("exclude"),
+        conflicts_with("exclude_suffix")
+    )]
+    exact_match: Vec<String>,
+
     /// Include only benchmarks belonging to the given categories.
     #[arg(long, value_parser = EnumArgParser::<Category>::default(), default_value = "Primary,Secondary")]
     category: MultiEnumValue<Category>,
@@ -688,6 +699,25 @@ enum DownloadSubcommand {
     },
 }
 
+impl<'a> From<&'a LocalOptions> for CompileBenchmarkFilter<'a> {
+    fn from(value: &'a LocalOptions) -> Self {
+        if !value.exact_match.is_empty() {
+            Self::Exact(&value.exact_match)
+        } else if !value.include.is_empty()
+            || !value.exclude.is_empty()
+            || !value.exclude_suffix.is_empty()
+        {
+            Self::Fuzzy {
+                include: &value.include,
+                exclude: &value.exclude,
+                exclude_suffix: &value.exclude_suffix,
+            }
+        } else {
+            Self::All
+        }
+    }
+}
+
 fn main_result() -> anyhow::Result<i32> {
     env_logger::init();
 
@@ -761,7 +791,7 @@ fn main_result() -> anyhow::Result<i32> {
             };
             let config = RuntimeBenchmarkConfig::new(
                 runtime_suite,
-                BenchmarkFilter::new(local.exclude, local.include),
+                RuntimeBenchmarkFilter::new(local.exclude, local.include),
                 iterations,
             );
             run_benchmarks(&mut rt, conn, shared, None, Some(config))?;
@@ -884,12 +914,7 @@ fn main_result() -> anyhow::Result<i32> {
                 target_triple,
             )?;
 
-            let mut benchmarks = get_compile_benchmarks(
-                &compile_benchmark_dir,
-                &local.include,
-                &local.exclude,
-                &local.exclude_suffix,
-            )?;
+            let mut benchmarks = get_compile_benchmarks(&compile_benchmark_dir, (&local).into())?;
             benchmarks.retain(|b| local.category.0.contains(&b.category()));
 
             let artifact_id = ArtifactId::Commit(Commit {
@@ -1006,9 +1031,11 @@ fn main_result() -> anyhow::Result<i32> {
 
                         let mut benchmarks = get_compile_benchmarks(
                             &compile_benchmark_dir,
-                            &split_args(include),
-                            &split_args(exclude),
-                            &[],
+                            CompileBenchmarkFilter::Fuzzy {
+                                include: &split_args(include),
+                                exclude: &split_args(exclude),
+                                exclude_suffix: &[],
+                            },
                         )?;
                         benchmarks.retain(|b| b.category().is_primary_or_secondary());
 
@@ -1042,7 +1069,7 @@ fn main_result() -> anyhow::Result<i32> {
 
                         let runtime_config = RuntimeBenchmarkConfig {
                             runtime_suite,
-                            filter: BenchmarkFilter::keep_all(),
+                            filter: RuntimeBenchmarkFilter::keep_all(),
                             iterations: DEFAULT_RUNTIME_ITERATIONS,
                         };
                         let shared = SharedBenchmarkConfig {
@@ -1102,12 +1129,7 @@ fn main_result() -> anyhow::Result<i32> {
             let scenarios = &opts.scenarios.0;
             let backends = &opts.codegen_backends.0;
 
-            let mut benchmarks = get_compile_benchmarks(
-                &compile_benchmark_dir,
-                &local.include,
-                &local.exclude,
-                &local.exclude_suffix,
-            )?;
+            let mut benchmarks = get_compile_benchmarks(&compile_benchmark_dir, (&local).into())?;
             benchmarks.retain(|b| local.category.0.contains(&b.category()));
 
             let mut errors = BenchmarkErrors::new();
@@ -1315,12 +1337,7 @@ fn binary_stats_compile(
         Profile::Opt => CargoProfile::Release,
         _ => return Err(anyhow::anyhow!("Only Debug and Opt profiles are supported")),
     };
-    let benchmarks = get_compile_benchmarks(
-        &compile_benchmark_dir(),
-        &local.include,
-        &local.exclude,
-        &local.exclude_suffix,
-    )?;
+    let benchmarks = get_compile_benchmarks(&compile_benchmark_dir(), (&local).into())?;
     for benchmark in benchmarks {
         println!("Stats for benchmark `{}`", benchmark.name);
         println!("{}", "-".repeat(20));
@@ -1713,7 +1730,7 @@ fn bench_published_artifact(
     };
 
     // Exclude benchmarks that don't work with a stable compiler.
-    let mut compile_benchmarks = get_compile_benchmarks(dirs.compile, &[], &[], &[])?;
+    let mut compile_benchmarks = get_compile_benchmarks(dirs.compile, CompileBenchmarkFilter::All)?;
     compile_benchmarks.retain(|b| b.category().is_stable());
 
     let runtime_suite = rt.block_on(load_runtime_benchmarks(
@@ -1745,7 +1762,7 @@ fn bench_published_artifact(
         }),
         Some(RuntimeBenchmarkConfig::new(
             runtime_suite,
-            BenchmarkFilter::keep_all(),
+            RuntimeBenchmarkFilter::keep_all(),
             DEFAULT_RUNTIME_ITERATIONS,
         )),
     )
diff --git a/collector/src/compile/benchmark/mod.rs b/collector/src/compile/benchmark/mod.rs
@@ -466,14 +466,22 @@ pub fn compile_benchmark_dir() -> PathBuf {
     PathBuf::from("collector/compile-benchmarks")
 }
 
+pub enum CompileBenchmarkFilter<'a> {
+    All,
+    /// Select benchmarks exactly matching the given benchmark names.
+    Exact(&'a [String]),
+    /// Select benchmarks matching the given prefixes/suffixes.
+    Fuzzy {
+        include: &'a [String],
+        exclude: &'a [String],
+        exclude_suffix: &'a [String],
+    },
+}
+
 pub fn get_compile_benchmarks(
     benchmark_dir: &Path,
-    include: &[String],
-    exclude: &[String],
-    exclude_suffix: &[String],
+    filter: CompileBenchmarkFilter<'_>,
 ) -> anyhow::Result<Vec<Benchmark>> {
-    let mut benchmarks = Vec::new();
-
     let mut paths = Vec::new();
     for entry in std::fs::read_dir(benchmark_dir)
         .with_context(|| format!("failed to list benchmark dir '{}'", benchmark_dir.display()))?
@@ -493,6 +501,40 @@ pub fn get_compile_benchmarks(
         paths.push((path, name));
     }
 
+    let mut benchmarks = match filter {
+        CompileBenchmarkFilter::All => paths
+            .into_iter()
+            .map(|(path, name)| Benchmark::new(name, path))
+            .collect::<anyhow::Result<Vec<Benchmark>>>()?,
+        CompileBenchmarkFilter::Exact(names) => paths
+            .into_iter()
+            .filter(|(_, name)| names.contains(name))
+            .map(|(path, name)| Benchmark::new(name, path))
+            .collect::<anyhow::Result<Vec<Benchmark>>>()?,
+        CompileBenchmarkFilter::Fuzzy {
+            include,
+            exclude,
+            exclude_suffix,
+        } => select_benchmarks_fuzzy(paths, include, exclude, exclude_suffix)?,
+    };
+
+    benchmarks.sort_by_key(|benchmark| benchmark.name.clone());
+
+    if benchmarks.is_empty() {
+        eprintln!("Warning: no benchmarks selected! Try less strict filters.");
+    }
+
+    Ok(benchmarks)
+}
+
+fn select_benchmarks_fuzzy(
+    paths: Vec<(PathBuf, String)>,
+    include: &[String],
+    exclude: &[String],
+    exclude_suffix: &[String],
+) -> anyhow::Result<Vec<Benchmark>> {
+    let mut benchmarks = Vec::new();
+
     // For each --include/--exclude entry, we count how many times it's used,
     // to enable `check_for_unused` below.
     fn to_hashmap(xyz: &[String]) -> HashMap<&str, usize> {
@@ -551,12 +593,6 @@ Expected zero or more entries or substrings from list: {:?}."#,
     check_for_unused("exclude", excludes)?;
     check_for_unused("exclude-suffix", exclude_suffixes)?;
 
-    benchmarks.sort_by_key(|benchmark| benchmark.name.clone());
-
-    if benchmarks.is_empty() {
-        eprintln!("Warning: no benchmarks selected! Try less strict filters.");
-    }
-
     Ok(benchmarks)
 }
 
@@ -578,17 +614,19 @@ fn substring_matches(
 
 #[cfg(test)]
 mod tests {
-    use crate::compile::benchmark::get_compile_benchmarks;
+    use crate::compile::benchmark::{get_compile_benchmarks, CompileBenchmarkFilter};
     use std::path::Path;
 
     #[test]
     fn check_compile_benchmarks() {
         // Check that we can deserialize all perf-config.json files in the compile benchmark
         // directory.
         let root = env!("CARGO_MANIFEST_DIR");
-        let benchmarks =
-            get_compile_benchmarks(&Path::new(root).join("compile-benchmarks"), &[], &[], &[])
-                .unwrap();
+        let benchmarks = get_compile_benchmarks(
+            &Path::new(root).join("compile-benchmarks"),
+            CompileBenchmarkFilter::All,
+        )
+        .unwrap();
         assert!(!benchmarks.is_empty());
     }
 }
diff --git a/collector/src/runtime/benchmark.rs b/collector/src/runtime/benchmark.rs
@@ -42,7 +42,7 @@ pub struct BenchmarkSuite {
 impl BenchmarkSuite {
     /// Returns a new suite containing only groups that contains at least a single benchmark
     /// that matches the filter.
-    pub fn filter(self, filter: &BenchmarkFilter) -> Self {
+    pub fn filter(self, filter: &RuntimeBenchmarkFilter) -> Self {
         let BenchmarkSuite {
             toolchain,
             groups,
@@ -64,7 +64,7 @@ impl BenchmarkSuite {
         }
     }
 
-    pub fn filtered_benchmark_count(&self, filter: &BenchmarkFilter) -> u64 {
+    pub fn filtered_benchmark_count(&self, filter: &RuntimeBenchmarkFilter) -> u64 {
         self.benchmark_names()
             .filter(|benchmark| passes_filter(benchmark, &filter.exclude, &filter.include))
             .count() as u64
@@ -86,12 +86,12 @@ impl BenchmarkSuite {
     }
 }
 
-pub struct BenchmarkFilter {
+pub struct RuntimeBenchmarkFilter {
     pub exclude: Vec<String>,
     pub include: Vec<String>,
 }
 
-impl BenchmarkFilter {
+impl RuntimeBenchmarkFilter {
     pub fn keep_all() -> Self {
         Self {
             exclude: vec![],
diff --git a/collector/src/runtime/mod.rs b/collector/src/runtime/mod.rs
@@ -9,8 +9,8 @@ use thousands::Separable;
 use benchlib::comm::messages::{BenchmarkMessage, BenchmarkResult, BenchmarkStats};
 pub use benchmark::{
     get_runtime_benchmark_groups, prepare_runtime_benchmark_suite, runtime_benchmark_dir,
-    BenchmarkFilter, BenchmarkGroup, BenchmarkGroupCrate, BenchmarkSuite,
-    BenchmarkSuiteCompilation, CargoIsolationMode,
+    BenchmarkGroup, BenchmarkGroupCrate, BenchmarkSuite, BenchmarkSuiteCompilation,
+    CargoIsolationMode, RuntimeBenchmarkFilter,
 };
 use database::{ArtifactIdNumber, CollectionId, Connection};
 
@@ -33,7 +33,7 @@ pub async fn bench_runtime(
     conn: &mut dyn Connection,
     suite: BenchmarkSuite,
     collector: &CollectorCtx,
-    filter: BenchmarkFilter,
+    filter: RuntimeBenchmarkFilter,
     iterations: u32,
 ) -> anyhow::Result<()> {
     let filtered = suite.filtered_benchmark_count(&filter);
@@ -203,7 +203,7 @@ async fn record_stats(
 /// a set of runtime benchmarks and print `BenchmarkMessage`s encoded as JSON, one per line.
 fn execute_runtime_benchmark_binary(
     binary: &Path,
-    filter: &BenchmarkFilter,
+    filter: &RuntimeBenchmarkFilter,
     iterations: u32,
 ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<BenchmarkMessage>>> {
     let mut command = prepare_command(binary);

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ pub struct BenchmarkSuite {`
`42`	`42`	`impl BenchmarkSuite {`
`43`	`43`	`/// Returns a new suite containing only groups that contains at least a single benchmark`
`44`	`44`	`/// that matches the filter.`
`45`		`- pub fn filter(self, filter: &BenchmarkFilter) -> Self {`
	`45`	`+ pub fn filter(self, filter: &RuntimeBenchmarkFilter) -> Self {`
`46`	`46`	`let BenchmarkSuite {`
`47`	`47`	`toolchain,`
`48`	`48`	`groups,`
`@@ -64,7 +64,7 @@ impl BenchmarkSuite {`
`64`	`64`	`}`
`65`	`65`	`}`
`66`	`66`
`67`		`- pub fn filtered_benchmark_count(&self, filter: &BenchmarkFilter) -> u64 {`
	`67`	`+ pub fn filtered_benchmark_count(&self, filter: &RuntimeBenchmarkFilter) -> u64 {`
`68`	`68`	`self.benchmark_names()`
`69`	`69`	`.filter(\|benchmark\| passes_filter(benchmark, &filter.exclude, &filter.include))`
`70`	`70`	`.count() as u64`
`@@ -86,12 +86,12 @@ impl BenchmarkSuite {`
`86`	`86`	`}`
`87`	`87`	`}`
`88`	`88`
`89`		`-pub struct BenchmarkFilter {`
	`89`	`+pub struct RuntimeBenchmarkFilter {`
`90`	`90`	`pub exclude: Vec<String>,`
`91`	`91`	`pub include: Vec<String>,`
`92`	`92`	`}`
`93`	`93`
`94`		`-impl BenchmarkFilter {`
	`94`	`+impl RuntimeBenchmarkFilter {`
`95`	`95`	`pub fn keep_all() -> Self {`
`96`	`96`	`Self {`
`97`	`97`	`exclude: vec![],`