Skip to content

Share bytecode IR between CPU and GPU trace gen#3756

Draft
qwang98 wants to merge 1 commit into
powdr-labs/cpu-tracegen-optfrom
powdr-labs/bytecode-refactor
Draft

Share bytecode IR between CPU and GPU trace gen#3756
qwang98 wants to merge 1 commit into
powdr-labs/cpu-tracegen-optfrom
powdr-labs/bytecode-refactor

Conversation

@qwang98

@qwang98 qwang98 commented May 22, 2026

Copy link
Copy Markdown
Collaborator

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 mirrors openvm/cuda/src/expr_eval.cuh.

What moves

  • New module openvm/src/bytecode.rs (no feature gate) carries OpCode, ExprSpan, BusMeta, the AST-walker emit_expr, and compile_bus_to_bytecode. cuda_abi.rs re-exports the types so existing GPU code is unchanged at the type level (DevInteraction becomes a re-export alias for BusMeta).
  • cuda/mod.rs keeps compile_derived_to_gpu (derived-column-specific) but delegates bus-interaction compilation to the shared module via pub use compile_bus_to_bytecode as compile_bus_to_gpu.
  • cpu/mod.rs calls compile_bus_to_bytecode(apc_height=1) so the PushApc operand is the column index directly, then iterates BusMeta + ExprSpan per row via a new CPU eval_expr interpreter (fixed-size [F; 16] stack on the local frame, unsafe get_unchecked indexing — mirrors the GPU kernel).
  • ~200 lines of duplicate AST decomposition (decompose + mul_decompositions + Polynomial + CompiledExpr + CompiledBusInteraction) deleted from autoprecompiles/src/expression.rs.

Perf (keccak/10K/APC=30/mock, apples-to-apples on this branch)

trace_gen run 1 run 2 run 3
Polynomial (PR #3723 head) 5.05s 5.09s 2.34s
Shared bytecode (this PR) 5.29s 5.41s 2.50s
Δ +5% +6% +7%

Two intentional cost sources, both confirmed by intermediate benchmarks:

  1. match op { x if x == OpCode::X as u32 => ... } lowers to an if-chain, not a jump table. A previous prototype with Vec<BcOp<F>> enum dispatch landed at +3%.
  2. Per-PushConst F::from_canonical_unchecked(u) reconstructs the field element from the inline u32 — the GPU layout stores constants as u32, not as F inline.

Trade-offs

  • +5–6% trace_gen wall (CPU). Visible but small in absolute terms (~0.2–0.3 s per shard).
  • One source of truth for AlgebraicExpression → Vec<u32> compile. Adding/changing opcodes (e.g. InvOrZero already shared) only touches one place.
  • Runtime path becomes degree-agnostic. Polynomial had a hard assert! at degree > 2; bytecode handles arbitrary degree via more Mul opcodes.
  • Loses generic-F at the eval entry; the shared module is hardcoded to F: PrimeField32 + QuotientMap<u32> to do the u32 ↔ F round trip. The current CPU caller is BabyBear only, 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 opcode match to a jump table:

let op_tag: OpCode = unsafe { std::mem::transmute::<u32, OpCode>(op) };
match op_tag { OpCode::PushApc => ..., ... }

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

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant