fix: run scheme compute on a dedicated large-stack rayon pool

[markosg04]

Summary

Prover/verifier compute can overflow rayon's default 2 MiB worker stacks, nondeterministically, when akita is used as a library. This PR makes the stack requirement library-owned: every AkitaCommitmentScheme entry point now runs on a dedicated, lazily-built rayon pool with 64 MiB worker stacks (run_on_compute_pool in akita_prover::compute), and adds a regression test that exercises the default-pool boundary the in-tree suite never touches.

Draft on purpose — the diagnosis is solid but there is a real design question about where you want the bound to live (see Open questions).

The failure, as observed downstream

Jolt drives this code through jolt-akita (fp128 D64, one-hot K=16/K=256 and dense flavors) and hit worker-thread stack overflows during proving:

• fatal runtime error: stack overflow, aborting / SIGABRT on an unnamed thread, killing the process mid-prove — once in GitHub Actions CI (ubuntu, during proof-fixture generation) and reproducibly on an M-series Mac at small trace shapes, where it failed ~40% of runs of an identical test (feat: Akita lattice PCS integration — jolt-prover-legacy packed prover a16z/jolt#1676, commits a2e2b21…f7172c1 carry the details and the downstream mitigation).
• The faulting frame under lldb is rayon::iter::plumbing::bridge_producer_consumer::helper on a rayon worker (guard-page hit; the frame chain doesn't unwind). RUST_MIN_STACK=134217728 makes it vanish, which fattens std-spawned rayon workers — consistent with a pure stack-depth failure.
• It is flaky because the split recursion's depth depends on work-stealing interleavings: rayon re-splits jobs when they migrate to a stealing worker, so the same input can produce very different inline recursion depths run to run, and different core counts (CI runners vs laptops) shift the distribution.

Why the in-tree suite never sees it

The e2e suite compensates globally: init_rayon_pool() installs a 256 MiB global rayon pool and run_on_large_stack() runs each test body on a 256 MiB thread (crates/akita-pcs/tests/akita_e2e.rs, tests/common/mod.rs). So the library's true stack appetite is invisible to CI here, while library callers inherit rayon's 2 MiB defaults.

Where the depth comes from

Sampling a live prove at a jolt-shaped input (one-hot K=16, 17 vars × 56 polys) shows the deepest chains are bridge_producer_consumer::helper → join_context towers bottoming out in the setup XOF matrix expansion — akita_types::proof::setup::derive_public_matrix_flat's per-entry SHAKE256 sampling (LabeledMatrixXof::entry_rng), whose closure frames carry keccak/digest state through every split level. Other kernels likely contribute; this one is the deepest chain I could observe directly.

The change

• akita_prover::compute::run_on_compute_pool<R: Send>(f: impl FnOnce() -> R + Send) -> R: a OnceLock rayon pool with stack_size(64 MiB) (lazily committed — virtual address space only) and named workers (akita-compute-N, so any future overflow identifies itself). Serial no-op fallback under not(feature = "parallel").
• All AkitaCommitmentScheme public entries funnel through it: setup_prover, setup_verifier, commit, batched_commit, commit_group, commit_final_group, batched_prove, batched_verify.
• batched_prove gains + Sync on its impl LevelProveStacks parameter (the concrete stacks already cross worker threads inside the kernels).
• New test akita-pcs/tests/default_stack_repro.rs: a jolt-shaped one-hot prove loop on the default global pool — deliberately not using init_rayon_pool/run_on_large_stack.

Honesty note on the test: on my machine it passes both before and after this change (7/7 pre-fix runs) — the overflow needs less favorable stealing dynamics than a standalone loop produces (jolt's mixed workload and CI's core counts hit it). It still guards the boundary: it is the only test in the suite that exercises library entry points on default stacks, and it turns any future regression here into a named-thread backtrace instead of a downstream mystery. A deterministic downstream repro exists: jolt @ f7172c1 with its pool wrap and trace-padding floor reverted fails ~40% of runs.

Open questions

1. Pool vs bounded kernels. An internal pool is the smallest robust fix and mirrors what both this repo's tests and jolt now do. The alternative is bounding the kernels themselves — e.g. with_min_len on the XOF expansion loop and heap-allocating its digest scratch — which would make default stacks safe without owning a pool. Happy to rework in that direction if preferred.
2. ComputeBackendSetup::prepare_setup is not wrapped (callers like jolt invoke CpuBackend.prepare_setup directly). Should it be?
3. Stack size: 64 MiB is 32× rayon's default and 4× smaller than the test suite's 256 MiB. If there are shapes that legitimately need more, the constant is one place to change.
4. If this lands, the test-suite init_rayon_pool/run_on_large_stack workarounds can probably be retired in a follow-up, which would make CI exercise the same stacks callers get.

Related (not in this PR)

The folded-only cutover (#311/#317) also removed all proof paths for small openings — measured minimum plannable num_vars at fc20716: dense ×1–2 → 13, dense ×56 → 19, one-hot ×56 → 16 ("no schedule with at least two folds"). That silently broke sub-4096-cycle traces downstream (jolt now pads to T ≥ 2^12). Happy to file it as a separate issue with the full probe data if useful.

🤖 Generated with Claude Code