perf: GKR-style DAG kernel prototype for apc_apply_bus#3751
Draft
qwang98 wants to merge 7 commits into
Draft
Conversation
End-to-end prototype that replaces the stack-VM bus kernel with a GKR-style
DAG evaluator reusing stark-backend's SymbolicRulesBuilder. Gated behind
`POWDR_BUS_KERNEL=dag` so the VM path stays the default. Bit-equal proof on
pairing APC=10 and APC=500.
What's new
- openvm/cuda/include/codec.cuh — vendored from stark-backend 8d36ad26
(`Rule`, `RuleHeader`, `decode_*`, `SourceInfo`, `OperationType`).
- openvm/cuda/src/apc_apply_bus_dag.cu — new kernel: per-thread walks
CSE-deduped `Rule[]`, fills `Fp inter[LOCAL_K=2048]`, dispatches to the
existing var-range / tuple2 / bitwise histograms identical to the VM.
- openvm/src/powdr_extension/trace_generator/cuda/expr_dag.rs:
- `to_symbolic` + `algebraic_to_symbolic_dag`: convert
AlgebraicExpression to SymbolicExpressionDag via SymbolicDagBuilder
(gets hash-cons + constant-fold for free).
- `compile_bus_to_gpu_dag`: build dag → run SymbolicRulesGpu (with
buffer_vars=true) → encode rules to u128 / `DevRule{low,high}` →
case-split each output rule to OutputDesc{Inter|Col|Const}.
- `dump_bus_dag_stats`: gated stats dump (rules_len, buffer_size,
node-kind histogram) so plan decisions are data-backed.
- CRITICAL: `keepalive` Vec retains every Arc<SymbolicExpression> the
converter hands to `SymbolicDagBuilder::add_expr`. Without it, the
builder's pointer-keyed `expr_to_idx` cache aliases stale ptr
addresses across freed Arcs and returns wrong dag_idxs — silently
collapsing distinct args to the same dag node.
- openvm/src/cuda_abi.rs — `DevRule`, `DevInteractionDag`, `OutputDesc`,
`OUTPUT_KIND_{INTER,COL,CONST}` + extern + safe wrapper for the new
kernel. Existing VM FFI untouched.
- openvm/src/powdr_extension/.../cuda/mod.rs — `POWDR_BUS_KERNEL=dag` env
read; both compile paths run under the same `bus_compile_h2d`
substage; both kernels run under the same `bus_kernel` substage.
Bench (pairing APC=500, nsys cuda_gpu_kern_sum, sm_120):
apc_apply_bus_kernel (VM) : 175.6 ms / 5.1% GPU
apc_apply_bus_dag_kernel (DAG) : 180.6 ms / 5.3% GPU (+2.9%)
ptxas -v on apc_apply_bus_dag.cu reports 8192-byte stack frame (the
`Fp inter[2048]` array goes to local memory, not registers, because the
slot index is runtime-decoded from rules). VM kernel: 64-byte stack frame.
Local-memory traffic cancels the DAG's per-row instruction-count win.
Companion patch: stark-backend-dag-prototype @ powdr-labs/dag-bus-prototype
(visibility flip on `rules` mod + output-slot last_use pin); Cargo.toml
points there via `[patch."https://github.com/powdr-labs/stark-backend.git"]`.
Checkpoint commit — local only, not for upstream. Next step (separate
branch) is option-2: replace `Fp inter[LOCAL_K]` with a per-launch
coalesced DeviceBuffer (`gkr_input.cu`'s `is_global=true` pattern).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The local-mode variant (Fp inter[LOCAL_K=2048]) forced an 8192-byte per-thread stack frame on chips with buffer_size > 1024 (max observed on pairing APC=500: 1762 slots). nvcc backs that as per-thread *local memory* — DRAM with per-thread random-access patterns — and L1 thrashes trying to cache a 128-thread block's 1 MB of per-thread state. Mirror gkr_input.cu's `is_global=true` path: the host allocates one DeviceBuffer<BabyBear> per launch sized at `total_threads * buffer_size`, slot-major coalesced. Each thread's `inter[]` is a base pointer `d_intermediates + thread_id`; slot `s` lives at `d_intermediates[s * total_threads + thread_id]`. A warp's 32 threads reading slot `s` hit one cache line. After change, `ptxas -v` reports: apc_apply_bus_dag_kernel: 0 byte stack frame, 0 spills, 34 regs (was 8192 byte stack frame, 0 spills, 35 regs) Bench (pairing APC=500, nsys cuda_gpu_kern_sum, sm_120): apc_apply_bus_kernel (VM) : 166.6 ms / 7.8% GPU apc_apply_bus_dag_kernel (DAG global) : 135.8 ms / 6.5% GPU (-18.5%) StdDev per launch : 271 → 188 µs (-31%) Max per launch (the 1762-slot chip) : 1.80 → 1.97 ms (+10%) Bit-equal on pairing APC=10 and APC=500 (public_values_commit matches VM byte-for-byte both times). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds a round-robin CudaStream pool (POWDR_BUS_DAG_STREAMS env, default 8) and launches each chip's bus DAG kernel on its assigned stream, with default_stream_wait at the end to preserve ordering against periphery chips that still run on cudaStreamPerThread. Result: zero speedup across pool sizes 1/4/8/16 (135.83/135.96/136.03/ 135.97 ms total). Root cause: cuda-common's memory manager routes cudaMallocAsync through cudaStreamPerThread, so the next chip's d_malloc serializes behind the prior chip's default_stream_wait. Without the sync, periphery chips race bus_kernel writes and produce NonzeroRootSum. Plan doc captures the candidate-D scope; included here so the failed attempt is reproducible from history. Reverted in the following commit to restore the option-2 winner state. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This reverts commit c30941e.
Borrowed from GKR (stark-backend logup_zerocheck/zerocheck_round0.cu:23, :214, :692): templatize the kernel on `<bool GLOBAL>`, allocate a per-thread `Fp local_buf[LOCAL_K]` for chips with buffer_size <= K, and skip the slot-major device buffer allocation entirely on those chips. LOCAL_K=32 chosen after measuring buffer_size distribution on guest-keccak APC=10 (4/29/30/50/81/97/2125 per chip — K=32 covers 6/10 chips). Dispatcher in `_apc_apply_bus_dag` picks the specialization at launch time by buffer_size. Result: bit-equal proof, but ~0% speedup. nvcc reports 128 B stack frame on the GLOBAL=false specialization (0 spill stores/loads, 34 registers) — `local_buf` ended up in .local memory rather than registers because its address escapes through the `__forceinline__`'d `evaluate_dag_entry_bus` helper. .local memory is per-thread DRAM (L1/L2 cached), the same tier as d_intermediates with slot-major access, so the cache hierarchies cancel. A/B benchmark (guest-keccak APC=100, 3 runs each, FORCE_GLOBAL env hook to flip dispatcher): SPLIT : bus_kernel = 7.72 ms (29 GLOBAL + 32 LOCAL launches) ALL-GLOBAL : bus_kernel = 7.71 ms (61 GLOBAL launches) Per-launch latency for the qualifying small chips: local-mode 24.2 us vs global-mode 24.1 us — within noise. Reviewer's R3 (pointer escape forces lmem) predicted this; data confirmed. Includes plan + critique doc and the abandoned D-2 scope doc. Reverted in the next commit to restore the option-2 winner state. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…(failed)" This reverts commit e2d4131.
Adds optional cudaProfilerStart/Stop wrapping around the launch when
POWDR_NCU_BIG=1 and n_rules > POWDR_NCU_BIG_THRESHOLD (default 1000).
Lets `ncu --profile-from-start off` capture only the keccak-class
launches regardless of launch-count/template-mangling complications.
Zero overhead when the env var is unset.
Usage:
POWDR_NCU_BIG=1 POWDR_BUS_KERNEL=dag \
ncu --profile-from-start off --set full -f -o report ./bin ...
Surfaced via this hook: the big-chip launches are memory-latency bound
on inter[slot] reads (L1TEX scoreboard stalls = 61.7% of stall time,
L1 hit 50%, L2 hit 62%, duration ~independent of num_apc_calls). See
the next round of optimization work.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
qwang98
commented
May 19, 2026
Collaborator
Author
There was a problem hiding this comment.
How this entire file differs from the stark-backend codec.cuh. Any chance we just reuse?
Collaborator
Author
There was a problem hiding this comment.
Claude answer: Byte-for-byte identical right now — diff returns no output (vendored at sha 8d36ad26).
Reuse needs a small upstream change. Powdr's openvm/build.rs already uses CudaBuilder::include_from_dep("DEP_X_INCLUDE") for cuda-common and circuit-primitives. Blocker: openvm-cuda-backend's build.rs doesn't println!("cargo:include=…"), so DEP_CUDA_BACKEND_INCLUDE resolves to nothing.
Two paths:
- Export the include from
cuda-backend/build.rs(one line in our powdr-labs fork), then drop the vendored file in favor ofinclude_from_dep("DEP_CUDA_BACKEND_INCLUDE")+#include "codec.cuh". Cleanest; cuda-backend stays source of truth. - Move
codec.cuhfromcuda-backend/cuda/include/→cuda-common/cuda/include/upstream.cuda-commonalready exports its include dir; powdr picks it up automatically. Codec is a pure data format — arguably belongs in the shared utility crate.
Vendored in this PR because the spike needed to land before the upstream coordination. #1 is ~5 LOC in our fork — happy to follow up.
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.
Summary
Replaces the stack-machine bytecode VM in
apc_apply_buswith a GKR-style DAG-evaluation kernel that reuses stark-backend'sSymbolicRulesBuilder(CSE + lifetime-analyzed slot allocation). Opt-in viaPOWDR_BUS_KERNEL=dag; default unchanged.Stacked on #3740 (
powdr-labs/apc-apply-bus-row-parallel).Measured
public_values_committo the VM path on guest-keccak APC=10.Commits in this PR
c1b3c7be4— DAG kernel prototype (correctness-only). New kernelapc_apply_bus_dag_kernelwalksRule[]fromSymbolicRulesBuilder. Codec vendored inopenvm/cuda/include/codec.cuh(synced with stark-backendcrates/cuda-backend/cuda/include/codec.cuh).e96a7baef— Coalesced global intermediates buffer. Slot-major coalesced layoutinter[slot * total_threads + tid]instead of per-thread localFp inter[K]. 5× faster than the local-mode prototype. The local-mode array (max observed K=1762) blew up the per-thread stack frame to 8 KB and trashed L1.c30941e58→38d963f2d— D-1 per-chip stream pool (failed) + revert. CudaStream pool round-robin across chips. 0 % speedup:cudaMallocAsyncserializes oncudaStreamPerThread. Kept in history.e2d413178→2ac6b0cd3— GKR-style GLOBAL/local intermediate split (failed) + revert. Template kernel on<bool GLOBAL>withFp local_buf[K=32]for small chips. 0 % speedup: nvcc fell back to per-thread.localmemory (128 B stack frame on the local-mode kernel) — same memory tier as global, just through different cache. R3 from the reviewer critique called this out before implementation. Kept in history.80ef81c91— ncu profiler gate for big-chip capture. OptionalcudaProfilerStart/Stopwrapping gated onPOWDR_NCU_BIG=1env + n_rules threshold. Zero overhead when unset.Diagnostic finding (from #5)
ncu on the big keccak chip (buffer_size=2125, ~5800 rules):
num_apc_calls→ it's a per-row latency chain of 5800 sequential rule walks, each waiting on the previous'sinter[slot]DRAM readThe big chip is memory-latency bound on slot reads, not bandwidth or compute. Layout/streams/local-array optimizations don't move it; the real levers are shared-mem
inter[](limited by 100 KB block cap to buffer_size ≲ 800), reducing the concurrent live-slot working set, or block-level cooperative DAG walks. Follow-up work tracked separately.Test plan
POWDR_BUS_KERNEL=dag cli-openvm prove guest-keccak --autoprecompiles 10 --input 10produces bit-equalpublic_values_commitPOWDR_BUS_KERNEL=dag cli-openvm prove guest-keccak --autoprecompiles 100 --input 100matches VM path🤖 Generated with Claude Code