Test timings and effect of enabling avx2

.github/workflows/scripts/summarise-timings.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Aggregate pod2/plonky2 timing lines into a per-scope summary.
+
+Recognises two line shapes:
+
+    timed "MainPod::prove": 7.068943638s          <- pod2's `timed!` macro
+    [.. DEBUG plonky2::util::timing] 0.5s to fft  <- plonky2's TimingTree
+
+Usage:
+    summarise-timings.py [log]              per-scope totals (stdin if no arg)
+    summarise-timings.py base.log other.log compare two runs, scope by scope
+
+Tests run in parallel unless --test-threads=1 is passed, in which case samples
+interleave and cannot be attributed to a test. Totals across the run still tell
+you where the time goes, which is the point.
+"""
+import re
+import sys
+from collections import defaultdict
+
+POD2 = re.compile(r'timed "([^"]+)": ([0-9.]+)(ms|µs|us|ns|s)\b')
+PLONKY2 = re.compile(r"(?:\| )*([0-9.]+)s to (.+?)\s*$")
+UNIT = {"s": 1.0, "ms": 1e-3, "us": 1e-6, "µs": 1e-6, "ns": 1e-9}
+
+
+def parse(lines):
+    totals = defaultdict(list)
+    for line in lines:
+        m = POD2.search(line)
+        if m:
+            totals[m.group(1)].append(float(m.group(2)) * UNIT[m.group(3)])
+            continue
+        if "util::timing" in line or line.lstrip().startswith("| "):
+            m = PLONKY2.search(line)
+            if m:
+                totals[m.group(2)].append(float(m.group(1)))
+    return totals
+
+
+def summarise(totals):
+    rows = sorted(totals.items(), key=lambda kv: -sum(kv[1]))
+    print(f"{'scope':<44}{'n':>5}{'total':>11}{'mean':>10}{'max':>10}")
+    print("-" * 80)
+    for name, xs in rows:
+        print(
+            f"{name[:43]:<44}{len(xs):>5}{sum(xs):>10.2f}s"
+            f"{sum(xs) / len(xs):>9.2f}s{max(xs):>9.2f}s"
+        )
+
+
+def compare(base, other, base_name, other_name):
+    # Compare totals per scope. A scope missing from one side is reported as
+    # such rather than as a delta, since that means the runs were not
+    # equivalent and any percentage would be meaningless.
+    names = sorted(set(base) | set(other), key=lambda n: -sum(base.get(n, [0])))
+    print(f"{'scope':<44}{base_name:>11}{other_name:>11}{'delta':>10}")
+    print("-" * 76)
+    for name in names:
+        if name not in base or name not in other:
+            side = base_name if name in base else other_name
+            print(f"{name[:43]:<44}{'(only in ' + side + ')':>32}")
+            continue
+        a, b = sum(base[name]), sum(other[name])
+        delta = f"{(b - a) / a * 100:+.1f}%" if a > 0 else "n/a"
+        print(f"{name[:43]:<44}{a:>10.2f}s{b:>10.2f}s{delta:>10}")
+
+
+if __name__ == "__main__":
+    args = sys.argv[1:]
+    if len(args) == 2:
+        with open(args[0]) as f:
+            base = parse(f)
+        with open(args[1]) as f:
+            other = parse(f)
+        if not base or not other:
+            sys.exit("no timing lines found in one of the logs")
+        compare(base, other, "base", "avx2")
+    else:
+        src = open(args[0]) if args else sys.stdin
+        totals = parse(src)
+        if not totals:
+            sys.exit(
+                "no timing lines found "
+                "(did you pass --features time and --nocapture?)"
+            )
+        summarise(totals)

.github/workflows/test-timings.yml

@@ -0,0 +1,50 @@
+name: Test timings
+
+# Manual only: this exists to answer "what is CI actually spending its time on",
+# not to gate anything. It costs nothing until someone runs it.
+on:
+  workflow_dispatch:
+    inputs:
+      filter:
+        description: "Test name filter (empty runs the whole suite)"
+        required: false
+        default: ""
+
+jobs:
+  timings:
+    name: Test timings
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Rust
+      uses: actions-rust-lang/setup-rust-toolchain@v1
+      with:
+        # Read the tests.yml cache but never write to it: the feature set here
+        # differs, and this job should not disturb the entry PRs restore from.
+        cache-shared-key: tests-release
+        cache-save-if: false
+    - name: Run tests with timing
+      # `--nocapture` is load-bearing. The `timed!` macro and plonky2's
+      # TimingTree both print, and libtest swallows stdout without it, so the
+      # timings silently vanish. Tests still run in parallel, so samples
+      # interleave and cannot be attributed to an individual test; the totals
+      # are what this job is for.
+      #
+      # The filter goes through env rather than being interpolated into the
+      # script: a dispatch input pasted straight into `run:` would let anyone
+      # who can trigger the workflow run arbitrary commands.
+      env:
+        FILTER: ${{ inputs.filter || '' }}
+      run: |
+        set -o pipefail
+        cargo test --release --features time,db_rocksdb ${FILTER:+"$FILTER"} \
+          -- --nocapture 2>&1 | tee timings.log
+    - name: Summarise
+      if: always()
+      run: .github/workflows/scripts/summarise-timings.py < timings.log
+    - name: Upload full log
+      if: always()
+      uses: actions/upload-artifact@v4
+      with:
+        name: timings
+        path: timings.log

.github/workflows/tests.yml

@@ -28,10 +28,45 @@ jobs:
         # saving from PRs just burns into the repo's 10 GB cache budget and can
         # evict the one entry every PR actually restores from.
         cache-save-if: ${{ github.ref == 'refs/heads/main' }}
+        # Proving is field-arithmetic heavy and plonky2 picks its Goldilocks
+        # packing at compile time, so without this every run uses the scalar
+        # 1-wide path. Measured at roughly -9% on MainPod::prove, reproduced on
+        # both Intel and AMD runners.
+        #
+        # A fixed feature, never `target-cpu=native`: rust-cache restores rlibs
+        # built on a different host, and this fleet mixes hosts with and
+        # without AVX-512, so `native` would bake in instructions the next
+        # runner may not have. AVX2 alone is safe on any x86 runner.
+        #
+        # `+pclmulqdq` is not optional. Build scripts see these features via
+        # CARGO_CFG_TARGET_FEATURE, and librocksdb-sys passes `-mavx2` to its
+        # C++ compiler when it sees avx2, which switches rocksdb onto a
+        # PCLMUL-based CRC32 path. It only passes `-mpclmul` when it also sees
+        # `pclmulqdq`, so avx2 alone fails to compile crc32c.cc. Note that
+        # `-C target-cpu=x86-64-v3` does NOT imply pclmulqdq and hits the same
+        # wall. PCLMULQDQ predates AVX2 in hardware, so this adds no floor.
+        rustflags: "-D warnings -C target-feature=+avx2,+pclmulqdq"
+    - name: Check the runner has AVX2
+      # plonky2 has no runtime fallback, so a runner without AVX2 would die
+      # with SIGILL somewhere inside a proof. Fail here with a readable reason.
+      run: |
+        grep -qm1 avx2 /proc/cpuinfo || {
+          echo "::error::runner lacks AVX2; drop +avx2 from rustflags in this workflow"
+          exit 1
+        }
     - name: Run tests
       # RocksDB is disabled by default but we still want to test it.
-      run: cargo test --release --features db_rocksdb
+      #
+      # mimalloc is CI-only, via a feature rather than a default: proving
+      # allocates heavily and this measured about -5% here (more on machines
+      # with more cores). pod2 must not impose an allocator on the programs
+      # that depend on it, so the hook is `cfg(test)` gated and applications
+      # that want it set their own `#[global_allocator]`.
+      #
+      # The example steps below must use the same feature set, or pod2 gets
+      # rebuilt between the steps.
+      run: cargo test --release --features db_rocksdb,mimalloc
     - name: Run example 1
-      run: cargo run --release --features db_rocksdb --example main_pod_points -- --mock
+      run: cargo run --release --features db_rocksdb,mimalloc --example main_pod_points -- --mock
     - name: Run example 2
-      run: cargo run --release --features db_rocksdb --example signed_dict
+      run: cargo run --release --features db_rocksdb,mimalloc --example signed_dict

Cargo.toml

@@ -51,6 +51,13 @@ good_lp = { version = "1.8", default-features = false, features = [
 annotate-snippets = "0.11"
 rocksdb = { version = "0.24.0", optional = true } # keyvalue database for merkletree
 
+# Alternative global allocators, for the experiment in #398. Proving is
+# allocation-heavy, so the allocator is worth measuring. These are only ever
+# installed by pod2's own test harness and examples, never by pod2 as a
+# dependency (see the `cfg(test)` gate in src/lib.rs).
+mimalloc = { version = "0.1.52", optional = true }
+tikv-jemallocator = { version = "0.7.0", optional = true }
+
 # Uncomment for debugging with https://github.com/ed255/plonky2/ at branch `feat/debug`.  The repo directory needs to be checked out next to the pod2 repo directory.
 # [patch."https://github.com/0xPARC/plonky2"]
 # plonky2 = { path = "../plonky2/plonky2" }
@@ -84,6 +91,10 @@ db_rocksdb = ["rocksdb"]
 # dependency out of slim builds; on for the test target unconditionally via
 # `[dev-dependencies] good_lp` so the parity sweep always builds.
 milp = ["dep:good_lp"]
+# Swap the global allocator in the test harness and examples. Mutually
+# exclusive, and diagnostic only.
+mimalloc = ["dep:mimalloc"]
+jemalloc = ["dep:tikv-jemallocator"]
 
 # Uncomment in order to enable debug information in the release builds.  This allows getting panic backtraces with a performance similar to regular release.
 # [profile.release]

examples/main_pod_points.rs

@@ -10,6 +10,14 @@
 //! Run in mock mode: `cargo run --release --example main_pod_points -- --mock`
 use std::env;
 
+#[cfg(feature = "mimalloc")]
+#[global_allocator]
+static GLOBAL_ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+#[cfg(feature = "jemalloc")]
+#[global_allocator]
+static GLOBAL_ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 use pod2::{
     backends::plonky2::{
         basetypes::DEFAULT_VD_SET, mainpod::Prover, mock::mainpod::MockProver,

src/backends/plonky2/emptypod.rs

@@ -130,7 +130,10 @@ impl EmptyPod {
 
         let mut pw = PartialWitness::<F>::new();
         empty_pod_verify_target.set_targets(&mut pw, vd_set.root())?;
-        let proof = timed!("EmptyPod prove", data.prove(pw)?);
+        let proof = timed!(
+            "EmptyPod prove",
+            crate::backends::plonky2::prove_with_timing("EmptyPod prove", data, pw)?
+        );
         let common_hash = hash_common_data(&data.common).expect("hash ok");
         Ok(EmptyPod {
             params: Params::default(),

src/backends/plonky2/mod.rs

@@ -39,6 +39,51 @@ use crate::{
     timed,
 };
 
+/// Prove `circuit_data`, and under the `time` feature also report plonky2's own
+/// breakdown of the proof: witness generation, wire polynomials and their
+/// commitment, partial products, quotient polynomials and the opening proofs.
+///
+/// `CircuitData::prove` hands plonky2 a throwaway `TimingTree` and discards all
+/// of that, which leaves a proof as a single opaque number. Since proving
+/// dominates the test suite, that is the number worth splitting up.
+pub(crate) fn prove_with_timing(
+    name: &str,
+    circuit_data: &basetypes::CircuitData,
+    pw: plonky2::iop::witness::PartialWitness<F>,
+) -> anyhow::Result<basetypes::ProofWithPublicInputs> {
+    #[cfg(not(feature = "time"))]
+    {
+        let _ = name;
+        circuit_data.prove(pw)
+    }
+    #[cfg(feature = "time")]
+    {
+        use plonky2::{plonk::prover::prove, util::timing::TimingTree};
+
+        // plonky2 reports the tree through the `log` crate, so with no logger
+        // installed the timings vanish without a trace. Installing one from
+        // library code is only tolerable because this is behind `time`, a
+        // diagnostic-only feature. RUST_LOG still wins if it is set.
+        static LOGGER: std::sync::Once = std::sync::Once::new();
+        LOGGER.call_once(|| {
+            let _ = env_logger::Builder::from_env(
+                env_logger::Env::default().default_filter_or("debug"),
+            )
+            .try_init();
+        });
+
+        let mut timing = TimingTree::new(name, log::Level::Debug);
+        let proof = prove(
+            &circuit_data.prover_only,
+            &circuit_data.common,
+            pw,
+            &mut timing,
+        )?;
+        timing.print();
+        Ok(proof)
+    }
+}
+
 pub fn cache_get_standard_rec_main_pod_common_circuit_data(
 ) -> CacheEntry<CommonCircuitDataSerializer> {
     let params = Params::default();

src/backends/plonky2/recursion/circuit.rs

@@ -181,7 +181,11 @@ pub fn prove_rec_circuit<I: InnerCircuit>(
         proofs,
         verifier_datas,
     )?;
-    Ok(circuit_data.prove(pw)?)
+    Ok(crate::backends::plonky2::prove_with_timing(
+        "prove_rec_circuit",
+        circuit_data,
+        pw,
+    )?)
 }
 
 impl<I: InnerCircuit> RecursiveCircuit<I> {

src/lib.rs

@@ -4,6 +4,20 @@
 #![allow(clippy::large_enum_variant)] // TODO: Remove this in another PR
 #![feature(mapped_lock_guards)]
 
+// Alternative global allocators (#398). Gated on `cfg(test)` so this applies
+// only to pod2's own test harness: a library has no business choosing the
+// allocator for the programs that depend on it.
+#[cfg(all(feature = "mimalloc", feature = "jemalloc"))]
+compile_error!("features `mimalloc` and `jemalloc` are mutually exclusive");
+
+#[cfg(all(test, feature = "mimalloc"))]
+#[global_allocator]
+static GLOBAL_ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+#[cfg(all(test, feature = "jemalloc"))]
+#[global_allocator]
+static GLOBAL_ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 pub mod backends;
 pub mod cache;
 pub mod frontend;