Stabilise walltime benchmarks across runs

.github/workflows/codspeed.yml

@@ -144,6 +144,10 @@ jobs:
       # crate that adds an `#[cpu_features]` benchmark is picked up without editing this workflow.
       # Building the whole workspace would do the same, but links every bench binary in it
       # — with debuginfo, from the bench profile — to measure only the tagged ones.
+      #
+      # A tagged binary must also run under `#[vortex_bench_support::main]`: ASLR is off on
+      # these runners, so without it the numbers depend on the size of the environment
+      # block, which differs between a push and a pull request. See that attribute's docs.
       - name: Select packages with tagged benchmarks
         id: select
         run: |
@@ -152,16 +156,23 @@ jobs:
 
           meta = json.loads(subprocess.check_output(
               ["cargo", "metadata", "--no-deps", "--format-version", "1"]))
-          tagged = sorted(
-              package["name"]
-              for package in meta["packages"]
-              for target in package["targets"]
-              if "bench" in target["kind"]
-              and "cpu_features]" in pathlib.Path(target["src_path"]).read_text()
-          )
+          tagged = {}
+          for package in meta["packages"]:
+              for target in package["targets"]:
+                  if "bench" not in target["kind"]:
+                      continue
+                  source = pathlib.Path(target["src_path"]).read_text()
+                  if "cpu_features]" in source:
+                      tagged[target["src_path"]] = (package["name"], "bench_support::main]" in source)
           if not tagged:
               raise SystemExit("no benchmark carries `#[cpu_features]`; this job has nothing to measure")
-          print("packages=" + " ".join(f"-p {name}" for name in dict.fromkeys(tagged)))
+          unwrapped = sorted(path for path, (_, wrapped) in tagged.items() if not wrapped)
+          if unwrapped:
+              raise SystemExit(
+                  "these carry `#[cpu_features]` but their `fn main` lacks "
+                  "`#[vortex_bench_support::main]`:\n  " + "\n  ".join(unwrapped))
+          packages = sorted({name for name, _ in tagged.values()})
+          print("packages=" + " ".join(f"-p {name}" for name in packages))
           EOF
       - name: Build benchmarks
         env:
@@ -185,15 +196,11 @@ jobs:
       # so this selects paths of the form `<bench target>::<features>::<name>`. That middle
       # component only exists because `#[cpu_features]` puts it there: an untagged benchmark has one
       # component fewer and cannot match, whatever it is called.
+      # How many samples to take and how many iterations each one runs are set on the
+      # benchmarks by `#[cpu_features]`, not here: a `DIVAN_SAMPLE_COUNT` or
+      # `DIVAN_SAMPLE_SIZE` in the environment would override every one of them.
       - name: Run benchmarks
         uses: CodSpeedHQ/action@4296e51e7041e24dadb86d1d6e8b9320d223dbe8  # v5
-        env:
-          # divan's default of 100 samples leaves these benchmarks too noisy to compare
-          # across runs: the same commit measured twice varied by up to 2.2x on the 1,024
-          # element cases and ~50-70% on the 65,536 element ones. At 1,000 samples the same
-          # experiment stays within ~6-10%, and the suite still runs in seconds, so the extra
-          # sampling is close to free next to the minute-plus spent building it.
-          DIVAN_SAMPLE_COUNT: "1000"
         with:
           run: bash scripts/bench-taskset.sh cargo codspeed run -- '.*::${{ matrix.features }}::'
           token: ${{ secrets.CODSPEED_TOKEN }}

benchmarks/bench-support/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "vortex-bench-support"
-description = "Proc macros for gating Vortex benchmarks by CPU architecture and instruction set"
+description = "Proc macros for Vortex benchmarks measured once per CPU feature set"
 authors = { workspace = true }
 categories = { workspace = true }
 edition = { workspace = true }

benchmarks/bench-support/src/lib.rs

@@ -3,15 +3,22 @@
 
 //! Support for benchmarks measured once per CPU feature set.
 //!
-//! One attribute, [`cpu_features`]. Which feature sets exist, what each is built with, and
-//! where it runs are all in `.github/workflows/codspeed.yml`.
+//! Two attributes. [`cpu_features`] marks a benchmark as measured on every walltime leg, and
+//! [`main`] goes on the `fn main` of a binary carrying such benchmarks. Which feature sets
+//! exist, what each is built with, and where it runs are all in `.github/workflows/codspeed.yml`.
 
 use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
 use syn::ItemFn;
+use syn::ReturnType;
 use syn::parse_macro_input;
 
+/// Stack size of the thread [`main`] runs the benchmarks on: the 16 MiB the walltime runners
+/// give the main thread (`ulimit -s`), so nothing that fit before overflows now. It is only
+/// reserved, not committed, so there is no cost to matching the larger figure.
+const BENCH_THREAD_STACK_SIZE: usize = 16 << 20;
+
 /// Measure this benchmark on every walltime CPU-feature leg instead of in simulation.
 ///
 /// Takes no argument: the benchmark runs on all of them. Write it *above* `#[divan::bench]`,
@@ -23,12 +30,30 @@ use syn::parse_macro_input;
 ///
 /// ```ignore
 /// #[vortex_bench_support::cpu_features]
-/// #[divan::bench(args = INPUT_SIZE)]
+/// #[divan::bench(args = INPUT_SIZE, sample_size = 256)]
 /// fn words_gather_dispatch(bencher: Bencher, len: usize) { /* ... */ }
 /// ```
 ///
+/// `sample_size` is required: it is how many times one sample runs the benchmark, and the
+/// number reported is the sample's time divided by it. The runner brackets every sample with
+/// its own hooks and timestamps, and what that costs the sample that follows is a fixed
+/// amount that differs from one host to the next by about a microsecond. Left to itself,
+/// divan sizes samples by its first, cold run and settles on one iteration for anything over
+/// a few microseconds, so on a bad host every case in a binary read the same ~1.2 µs slower
+/// regardless of its length. Pick the size so that a sample lasts at least ~100 µs on the
+/// smallest argument: 64 for a 2 µs case, 16 for a 10 µs one, 4 for 50 µs. When the arguments
+/// span a wide range, add `sample_count` below the default of 1,000 so the largest one does
+/// not take seconds.
+///
+/// A larger sample changes what divan keeps alive: it builds a sample's `with_inputs` up
+/// front and holds its outputs until the sample ends, so with `sample_size` of each in
+/// flight, an input built per iteration or an output of any size is written to cold memory
+/// instead of a block the allocator just recycled. Build inputs once, and drop a large output
+/// inside the closure with `divan::black_box_drop` rather than returning it.
+///
 /// Spell it out in full rather than importing it: benchmark files are read a function at a
-/// time, and the path says where the behaviour comes from.
+/// time, and the path says where the behaviour comes from. The binary's `fn main` carries
+/// [`main`], which the walltime legs need for a repeatable layout.
 ///
 /// This is for code that is written once and *compiled* differently per feature set — a
 /// shipped entry point that selects its kernel through `cfg(target_feature)`, or a scalar
@@ -76,12 +101,15 @@ pub fn cpu_features(attr: TokenStream, item: TokenStream) -> TokenStream {
         }
     };
 
-    for reserved in ["name", "ignore"] {
-        if existing
+    let has_option = |option: &str| {
+        existing
             .clone()
             .into_iter()
-            .any(|token| matches!(&token, proc_macro2::TokenTree::Ident(i) if i == reserved))
-        {
+            .any(|token| matches!(&token, proc_macro2::TokenTree::Ident(i) if i == option))
+    };
+
+    for reserved in ["name", "ignore"] {
+        if has_option(reserved) {
             return syn::Error::new_spanned(
                 &existing,
                 format!(
@@ -93,6 +121,27 @@ pub fn cpu_features(attr: TokenStream, item: TokenStream) -> TokenStream {
         }
     }
 
+    if !has_option("sample_size") {
+        return syn::Error::new_spanned(
+            bench,
+            "`#[cpu_features]` needs `sample_size` in `#[divan::bench(..)]`: enough iterations \
+             per sample for one sample to last ~100 µs or more, so the runner's per-sample \
+             overhead, which varies by host, is amortised (see the attribute's docs)",
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    // The default used to be `DIVAN_SAMPLE_COUNT` in the workflow, but a value set there beats
+    // one set on the benchmark, and the cases that need thousands of iterations per sample
+    // need fewer samples in exchange. divan's own default of 100 samples left the same commit
+    // measured twice up to 2.2x apart on the 1,024 element cases; 1,000 holds them to ~6-10%.
+    let sample_count = if has_option("sample_count") {
+        quote!()
+    } else {
+        quote!(sample_count = 1000,)
+    };
+
     // Skipping simulation, rather than naming the legs to run on, is what keeps the leg list
     // in the workflow alone. `env!` rather than reading the environment during expansion:
     // rustc records it as a dependency of the crate, so changing legs rebuilds the benchmarks.
@@ -112,6 +161,7 @@ pub fn cpu_features(attr: TokenStream, item: TokenStream) -> TokenStream {
     bench.meta = syn::parse_quote! {
         #bench_path(
             #existing #separator
+            #sample_count
             name = concat!(
                 env!("VORTEX_BENCH_PREFIX"),
                 #name,
@@ -123,3 +173,82 @@ pub fn cpu_features(attr: TokenStream, item: TokenStream) -> TokenStream {
 
     quote!(#function).into()
 }
+
+/// Run the benchmarks on a thread of their own, so where their stack and heap land is set by
+/// the binary rather than by the process.
+///
+/// Write it on `fn main`, which must take no arguments and return `()`:
+///
+/// ```ignore
+/// #[vortex_bench_support::main]
+/// fn main() {
+///     divan::main();
+/// }
+/// ```
+///
+/// Every binary with a [`cpu_features`] benchmark needs it; the workflow refuses to measure
+/// one without it. The body runs unchanged on a spawned thread with a fixed-size stack, and a
+/// panic there is re-raised on the main thread so the exit status is what it was.
+///
+/// Some µs-scale cases have two stable timings up to 1.5× apart, and which one a run lands in
+/// depends on the layout of the process rather than on the code: `arrow_checked_add_u32` in
+/// `vortex-compute` read 13 µs or 20 µs on the neon leg from one run to the next. On the main
+/// thread, with the address space randomised, it landed in either about half the time; on a
+/// spawned thread it held one timing through every one of those runs, and through a sweep of
+/// the environment block's size (which, contrary to the obvious guess, did not move it on
+/// either thread). The spawned thread's stack is a fresh mapping and the allocators in use
+/// serve it from mappings of their own, so its layout is set by the binary alone.
+#[proc_macro_attribute]
+pub fn main(attr: TokenStream, item: TokenStream) -> TokenStream {
+    if !attr.is_empty() {
+        let attr = TokenStream2::from(attr);
+        return syn::Error::new_spanned(attr, "`#[vortex_bench_support::main]` takes no argument")
+            .to_compile_error()
+            .into();
+    }
+
+    let function = parse_macro_input!(item as ItemFn);
+    let signature = &function.sig;
+
+    if signature.ident != "main" {
+        return syn::Error::new_spanned(
+            &signature.ident,
+            "`#[vortex_bench_support::main]` goes on `fn main`",
+        )
+        .to_compile_error()
+        .into();
+    }
+    if !signature.inputs.is_empty()
+        || !matches!(signature.output, ReturnType::Default)
+        || signature.asyncness.is_some()
+    {
+        return syn::Error::new_spanned(
+            signature,
+            "`#[vortex_bench_support::main]` expects `fn main()` with no arguments and no return type",
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    let attrs = &function.attrs;
+    let vis = &function.vis;
+    let body = &function.block;
+
+    // The body becomes the closure, so a `return` in it leaves the benchmarks as before. An
+    // explicit stack size rather than the default keeps `RUST_MIN_STACK` from putting the
+    // environment back into the layout.
+    quote! {
+        #(#attrs)*
+        #vis fn main() -> ::std::io::Result<()> {
+            let benchmarks = ::std::thread::Builder::new()
+                .name("vortex-bench".to_owned())
+                .stack_size(#BENCH_THREAD_STACK_SIZE)
+                .spawn(|| #body)?;
+            if let Err(panic) = benchmarks.join() {
+                ::std::panic::resume_unwind(panic);
+            }
+            Ok(())
+        }
+    }
+    .into()
+}