Performance Optimization

Performance Optimization

Architecture Overview

Backend Hierarchy

prover/backend/
  cpu/       Reference implementation (slow, readable)
  simd/      Production implementation (fast, unsafe-heavy)

The SIMD backend is the primary production path. The CPU backend is:

• Reference for correctness verification
• Fallback for unsupported architectures
• Baseline for benchmark comparisons

SIMD Targets

PackedM31: 16 M31 elements processed in parallel regardless of platform. On platforms with < 512-bit SIMD, this is emulated with multiple operations.

Memory Optimization

Memory Pool (prover/mempool.rs):

• BaseColumnPool reuses Vec allocations across proving rounds
• Uses DashMap for thread-safe concurrent access
• Avoids repeated large allocation + deallocation during polynomial evaluation

Uninitialized Allocation:

• uninit_vec() in core/utils.rs skips zero-initialization
• All callers MUST write before read — this is enforced by code review, not the type system
• Used in FFT buffers, quotient computation, trace generation

Benchmarking

Running Benchmarks

# All benchmarks (single-threaded)
cargo bench --features prover

# With parallelism
cargo bench --features "prover,parallel"

# Specific benchmark
cargo bench --features prover --bench fft
cargo bench --features prover --bench field
cargo bench --features prover --bench fri

# With AVX512 (on supported hardware)
RUSTFLAGS="-C target-cpu=native -C target-feature=+avx512f" cargo bench --features prover

Benchmark Suites

CI Regression Detection

The CI runs run-avx512-bench and run-avx512-bench-parallel on dedicated stwo-avx runners. The benchmark-action/github-action-benchmark@v1 action:

• Compares against previous benchmark data
• Fails on alert (regression detected)
• Comments on PRs with performance changes
• Alerts @gilbens-starkware on regressions

Hot Paths

FFT (Dominant Cost)

The SIMD FFT in prover/backend/simd/fft/ is the single most performance-critical code in STWO. It uses:

• Raw pointer arithmetic for zero-overhead array access
• Platform-specific SIMD intrinsics
• In-place computation to minimize memory traffic
• Twiddle factor precomputation (prover/poly/twiddles.rs)
• Bit-reversal permutation fused with FFT layers

Files: simd/fft/rfft.rs (forward), simd/fft/ifft.rs (inverse)

Merkle Tree Hashing

Blake2s hashing for Merkle commitments is the second hot path:

• prover/backend/simd/blake2s.rs — SIMD Blake2s (16 parallel hashes)
• prover/backend/simd/blake2s_lifted.rs — Lifted Merkle variant

Quotient Computation

DEEP quotient evaluation at query points:

• prover/backend/simd/quotients.rs — SIMD quotient computation
• Uses from_simd_unchecked to skip field reduction when values are known valid

Parallelism

Feature: parallel (enables rayon)

Key parallel sections:

• Trace generation (row-parallel)
• Constraint evaluation (domain-parallel)
• Merkle tree construction (subtree-parallel)
• FRI fold (element-parallel)

Thread safety: UnsafeMut/UnsafeConst wrappers in simd/utils.rs enable passing raw pointers across rayon threads. This requires that no two threads write to the same memory location.

Profile Configuration

# Cargo.toml [profile.bench]
codegen-units = 1    # Single codegen unit for maximum optimization
lto = true           # Link-time optimization

Optimization Rules

1. Never sacrifice correctness for performance. A fast wrong answer is worse than a slow right answer.
2. Benchmark before and after. Every optimization must be measured.
3. Regression = blocking issue. Performance regressions are not tech debt.
4. SIMD must match scalar. Every SIMD optimization must produce identical results to the CPU reference.
5. Document unsafe. Every unsafe block in the hot path needs a safety comment.

Review Checklist

Before approving performance changes:

• Benchmark shows measurable improvement
• No regression on other benchmarks
• SIMD implementation matches CPU reference
• Unsafe code has documented safety invariants
• Memory allocation patterns don't leak
• Thread safety is maintained under parallel feature

Forbidden Actions

In this domain, agents must NEVER:

• Optimize by removing verification checks
• Add unsafe without a clear performance justification AND safety proof
• Change algorithmic correctness for speed (e.g., approximate reduction)
• Skip benchmarking when claiming an optimization
• Modify the bench profile without understanding CI implications