feat(stark): GPU constraint/composition evaluation#798
Draft
MauroToscano wants to merge 7 commits into
Draft
Conversation
…kernel, dispatch)
Lower the captured ConstraintProgram to a flat device blob and interpret it on
the GPU (one thread per LDE row), for constraint-eval data residency — keep the
large LDE trace on-device instead of round-tripping it to evaluate constraints
on the CPU.
- constraint_ir/device.rs: DeviceProgram::lower — flat #[repr(C)] blob
(DeviceNode{op,a,b,dim} + u64/[u64;3] const tables + roots + num_base) — plus
eval_device_program, an independent CPU walk over that flat blob = the pre-GPU
parity oracle (mirrors interp.rs bit-for-bit).
- constraint_ir/gpu_interp.rs (cfg cuda): try_eval_program_gpu — TypeId gate on
Goldilocks + one device-edge reinterpret + flatten + kernel launch, Option
return (None = CPU fallback).
- math-cuda/kernels/constraint_interp.cu: the interpreter kernel — one thread
per row, per-thread value array in global memory, all-ext3 (base = {x,0,0},
embed = identity), per-constraint eval matrix output.
- math-cuda/src/constraint_interp.rs + build.rs/device.rs/lib.rs wiring: cudarc
host wrapper (all-u64 uploads, thread-capped value scratch).
- prover device parity test: lower + walk every one of the 26 production
programs and diff against the compiled prover folder.
CPU-verified: device unit test (all ops, 1000 rows) + 26-program parity green;
math-cuda and stark --features cuda compile (empty-PTX-stub path when nvcc is
absent); fmt + clippy -D warnings clean. The kernel itself awaits nvcc + a GPU
for compile and GPU<->CPU parity (the next phase).
Drive constraint_interp.cu through gpu_interp::try_eval_program_gpu over a random device-resident LDE and assert the per-constraint eval matrix is bit-for-bit identical to the CPU oracle eval_device_program (itself pinned to the compiled folder). Covers a synthetic all-ops program and all 26 production programs. cuda-gated; validated on an RTX 5090 (sm_120, CUDA 13.1).
Add constraint_composition_kernel: same node walk as the interpreter but fuses the composition-poly accumulation on-device (no per-constraint matrix, no D2H): H(row) = z_inv[row]·Σ βᵢ·Cᵢ + Σ_b z_b_inv[row]·β_b·(trace_b − value_b), the uniform-zerofier case of evaluator::evaluate. Plus the eval_composition_on_device wrapper, the try_eval_composition_gpu dispatch (gated; non-uniform → CPU fallback, which the VM never hits since it has no end-exemptions), and a GPU↔CPU H parity test against a CPU oracle that applies the exact evaluator accumulation to eval_device_program's evals.
ConstraintEvaluator::evaluate now tries try_evaluate_composition_gpu before the CPU per-row loop: when the GPU LDE handles are present, the tower is Goldilocks/degree-3, the zerofier is uniform (VM has no end-exemptions), and transition offsets are contiguous, it produces H(row) on-device (transition + boundary fused, next_step = lde_step_size) and returns it for the existing decompose+commit. Otherwise falls through to the unchanged CPU path. Additive + feature-gated; the CPU path is untouched.
Add GPU_COMPOSITION_CALLS (mirrors the other gpu_*_calls counters), incremented on a successful fused-composition dispatch, plus gpu_composition_path_fires_and_verifies: proves fib_iterative_1M with cuda, asserts the composition path fired (guards silent CPU fallback), and verifies the proof (guards a bad-H regression).
Add set_gpu_composition_disabled / LAMBDA_VM_DISABLE_GPU_COMPOSITION runtime toggle (gpu_lde.rs) checked in try_evaluate_composition_gpu, so one cuda binary can A/B the round-2 GPU composition path against the CPU accumulation with zero build differences. bench_abba_gpu.sh runs an interleaved ABBA over the ethrex workload flipping that toggle (A=GPU on, B=CPU), default 20 transfers.
…oundary shape asserts - gpu_interp: factor the shared TypeId gate + unsafe program reinterpret + lower + uniform packing into lower_and_pack() (was duplicated across both dispatch fns). - constraint_interp: assert all boundary-input shapes (b_is_aux/b_value/b_beta), not just b_z_inv, so a caller mismatch panics cleanly instead of an OOB read.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR 1 of 3 (stack: this → GPU trace residency → CUDA build robustness).
What
Evaluate the round-2 composition polynomial
Hon the GPU instead of the CPU. The constraint program is lowered once to a flat device IR (constraint_ir::device), and a fused CUDA kernel computesH(row) = z_inv·Σβᵢ·Cᵢ + boundaryreading the already-resident LDE — no per-constraint matrix, no trace D2H for constraint eval.Structure
constraint_interp.cuhas two kernels: the fused composition kernel (production, wired into round 2) and a matrix interpreter kernel that is test-only — a parity oracle proving the GPU node-walk matches the CPU folder bit-for-bit, per opcode. They share the device-IR lowering.gpu_interp.rsis the single concrete-Goldilocks dispatch seam (TypeId-gated, oneunsafereinterpret inlower_and_pack).Win
~−2.7% prove time (ethrex continuation, RTX 5090), measured via the
LAMBDA_VM_DISABLE_GPU_COMPOSITIONA/B toggle.Validation
Bit-for-bit parity: synthetic all-ops program + all 26 production AIRs vs the CPU folder; e2e prove→verify (
cuda_path_integration). Covered by a multi-agent code review (no wrong-proof findings).