Share bytecode IR between CPU and GPU trace gen#3756
Draft
qwang98 wants to merge 1 commit into
Draft
Conversation
Extract the stack-machine bytecode (OpCode / ExprSpan / BusMeta + the
AST-walker emit_expr / compile_bus_to_bytecode) into a new
feature-agnostic module openvm/src/bytecode.rs. cuda_abi re-exports the
types for back-compat; cuda/mod.rs keeps compile_derived_to_gpu but
delegates bus-interaction compilation to the shared module.
Add a CPU eval_expr interpreter that mirrors openvm/cuda/src/expr_eval.cuh:
fixed-size [F; 16] stack on the local frame, unsafe get_unchecked
indexing, single dispatch loop. Mirrors the GPU layout exactly:
Vec<u32> bytecode + ExprSpan{off,len} indexing into it.
Wire cpu/mod.rs to call compile_bus_to_bytecode(apc_height=1) so the
PushApc operand is the column index directly, then iterate the shared
bytecode/spans per bus interaction. Removes the per-eval Vec allocation
that the prior Vec<F> stack form had.
Drop the previous polynomial-form fast-eval (CompiledExpr / Polynomial /
decompose / mul_decompositions / CompiledBusInteraction) from
autoprecompiles/src/expression.rs - about 200 lines superseded by the
shared bytecode path.
Per keccak/10K/APC=30/mock: Polynomial 5.05/5.09/2.34s vs shared
bytecode 5.29/5.41/2.50s (+5-6%). In exchange CPU and GPU share one
compile pipeline, runtime is degree-agnostic (Polynomial was capped at
2), and ~200 LoC of duplicate AST decomposition logic is gone.
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
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.
NOT HUMAN REVIEWED YET.
Stacked on top of #3723. Refactors the CPU APC trace-gen path to reuse the GPU stack-machine bytecode pipeline introduced for
apc_apply_bus. Replaces the polynomial-form fast eval (CompiledExpr/Polynomial/decompose/mul_decompositions) with a CPU bytecode interpreter that mirrorsopenvm/cuda/src/expr_eval.cuh.What moves
openvm/src/bytecode.rs(no feature gate) carriesOpCode,ExprSpan,BusMeta, the AST-walkeremit_expr, andcompile_bus_to_bytecode.cuda_abi.rsre-exports the types so existing GPU code is unchanged at the type level (DevInteractionbecomes a re-export alias forBusMeta).cuda/mod.rskeepscompile_derived_to_gpu(derived-column-specific) but delegates bus-interaction compilation to the shared module viapub use compile_bus_to_bytecode as compile_bus_to_gpu.cpu/mod.rscallscompile_bus_to_bytecode(apc_height=1)so thePushApcoperand is the column index directly, then iteratesBusMeta+ExprSpanper row via a new CPUeval_exprinterpreter (fixed-size[F; 16]stack on the local frame,unsafe get_uncheckedindexing — mirrors the GPU kernel).decompose+mul_decompositions+Polynomial+CompiledExpr+CompiledBusInteraction) deleted fromautoprecompiles/src/expression.rs.Perf (keccak/10K/APC=30/mock, apples-to-apples on this branch)
Two intentional cost sources, both confirmed by intermediate benchmarks:
match op { x if x == OpCode::X as u32 => ... }lowers to an if-chain, not a jump table. A previous prototype withVec<BcOp<F>>enum dispatch landed at +3%.PushConstF::from_canonical_unchecked(u)reconstructs the field element from the inlineu32— the GPU layout stores constants asu32, not asFinline.Trade-offs
AlgebraicExpression → Vec<u32>compile. Adding/changing opcodes (e.g.InvOrZeroalready shared) only touches one place.Polynomialhad a hardassert!at degree > 2; bytecode handles arbitrary degree via moreMulopcodes.F: PrimeField32 + QuotientMap<u32>to do theu32 ↔ Fround trip. The current CPU caller isBabyBearonly, so no regression in callable surface.Open follow-ups
If you want the +5–6% recovered, the
transmute-to-enum dispatch trick should let the compiler lower the opcodematchto a jump table:That's ~5 lines and should land around v2's +3% based on the intermediate bench. Did not apply in this PR to keep the diff focused; happy to do it in a follow-up if the perf delta matters.
Test plan