feat: update architecture for generic `ProverBackend`s#1836

jonathanpwang · 2025-07-09T19:25:36Z

We update the overall architecture, mainly in the form of the VM extension and config traits, so that OpenVM becomes even more modular:

Update traits to better support our new design of different execution modes
Make the core system and all VM extensions modular in how trace generation is performed -- the motivation is so that trace generation may be done directly on device (e.g., GPU) instead of on CPU, depending on the ProverBackend, which is the hardware backend for stark-backend.

Some motivations were described in the description of #1809

Todo:

Update specs with new designs and also overview of new traits and structs
Fix unit testing framework (VmTestBuilder) to adapt to new design
Fix metric collection
Update extension definitions for all existing extensions
Update SDK
Handle forced trace heights
Nice to have: custom arena capacity declaration based on the RecordArena type (don't always use air width)

This is a partial PR merging into another feature branch `feat/new-exec-device` where I'll collect the overall work into generalizing chips for different tracegen and `ProverBackend`s. The integration tests in `rv32im/tests` work and `VirtualMachine, VmExecutor` were updated for the correct flow, assuming you call `set_main_widths` ahead of time (a temporary measure). The unit test framework is really not designed for this, so I only demonstrate `rand_auipc_test` works -- I'm not going to update all the tests to use `give_away_arena` since this is all going to change in the next PR. Overall summary: - chore change to `TraceStep` to remove useless `<CTX>` generic - `InstructionExecutor` is now basically `TraceStep` without the associate types, but still implemented on the old `Chip` struct that contains a trace matrix - however now the flow is that the `**Chip` is created with empty matrix, the record arena is owned by `TracegenCtx`, and then after all the records are populated, we give the record arenas to the chips via `give_me_my_arena` (function intentionally named to make sure we remove it later) The goal here is just to do the initial refactor to get the `RecordArena` into the correct place. Coming soon: blowing everything up to support different `RA, PB` generics. closes INT-4062

Previously, `VmExtension` had only a single notion of chip, where some chips could be executors, and the notions of execution, circuit, tracegen were all deeply intertwined. With the new design, we have multiple modes of execution which are more disconnected from the circuit design, and we also want tracegen to support different device / prover backends. The overall realization is that a true VM extension, meaning a modular way to extend the functionality of a working zkVM (including the prover) has three _separate_ components: - extending the execution methods for custom opcodes (this includes different execution modes) - extending the circuit itself at the level of AIRs - extending how tracegen and ultimately proving of the circuit from execution traces is actually performed In this PR, we aim to separate the notion of `VmExtension` into three separate traits `VmExecutionExtension, VmCircuitExtension, VmProverExtension` to better encapsulate this. - The `VmExecutionExtension` is in principle separate from any circuit considerations: it is just a way to specify hooks to opcodes. Although practically, there is still some connection, and we still take an enum dispatch approach (for now) for storage purposes. This is intended to be compatible with the fn pointer changes being made in parallel. - The `VmCircuitExtension` is the most security critical, and it should have **no** dependencies on the other two components. This will be made so that the verifying keys do not change even after this PR, and also this trait will not depend on any prover backend differences. - `VmProverExtension` is the most customizable, and hence unfortunately will have a bunch of generics. We will allow even the system chips to be customized according to the prover backend. Also - Deleted `NewVmChipWrapper` because the wrapper pattern is not good. Now each chip owns `mem_helper` -- it is perhaps more correct if just the adapter chip owns `mem_helper`, but since we aren't so convinced about adapter chips existing in the future, I just did the more immediate refactor. - Deleted the `Shared**Chip` wrapper structs because we don't need them anymore. I made them typedef to `Arc<..>`. ## Design Questions I am still somewhat on the fence about `Chip` being generic in `RA` versus doing the old approach where the `**Chip` struct is generic in `RA` and starts off with empty records which get swapped out with the records from the `TracegenCtx` (see before this PR). However it feels weird when the clear use case is now that the "preflight trace" is recorded elsewhere and only fed into the chip at time of tracegen.

… we are rewriting it

…mory chore: remove generic <F> from TracingMemory

added access_adapter_records and touched_memory

Toward INT-4363 implemented the new extension traits for `rv32im` and `native` extensions. `cargo check` passes for both however had to comment out bunch of stuff (mainly testing related) update: implemented keccak and sha2 as well

Segment check frequency shouldn't be larger than the height limit.

jonathanpwang · 2025-07-23T18:07:18Z

In GH comment summary, app.cells_used is main_cells_used. The difference is because the way we measure is slightly different now.
Also I've added a total_cells_used metric to include perm cols.

github-actions · 2025-07-24T01:27:27Z

group	app.proof_time_ms	app.cycles	app.cells_used	leaf.proof_time_ms	leaf.cycles	leaf.cells_used
verify_fibair	(-182 [-14.5%]) 1,070	322,677	(+401554 [+2.3%]) 17,672,132	-	-	-
fibonacci	(-346 [-14.7%]) 2,003	3,000,556	(+919286 [+1.8%]) 51,508,517	2,966	1,248,063	70,825,520
regex	(-246 [-3.5%]) 6,760	8,330,866	(-1017902 [-0.6%]) 165,431,684	8,814	3,326,702	234,204,718
ecrecover	(-359 [-26.1%]) 1,016	274,568	(+416462 [+5.1%]) 8,596,011	9,096	2,934,894	246,667,288
pairing	(-255 [-6.6%]) 3,594	3,725,930	(-9425620 [-9.2%]) 93,105,233	4,344	2,010,471	140,396,975
fib_e2e	16,792	24,000,556	410,948,443	15,001	7,462,398	435,773,140
kitchen_sink	14,829	309,192	899,579,940	22,077	7,952,497	748,994,286

Commit: e8f3eba

Benchmark Workflow

Make interface changes to coordinate with openvm-org/openvm#1836

* chore: fix benchmark for new traits Make interface changes to coordinate with openvm-org/openvm#1836 * chore: update openvm for trace_gen_time_ms

Sign in to view

+            // ATTENTION: make sure to permute everything in vm_pk that references the original AIR
+            // ID ordering:
            root_air_perm.permute(&mut vm_pk.per_air);
+            for thc in &mut vm_pk.trace_height_constraints {


* fix: update tracegen to do keygen first # Conflicts: # .cargo/config.toml # Cargo.lock # bin/client-eth/Cargo.lock # bin/client-eth/Cargo.toml * chore: update lock files * chore: update lock files * feat: impl max_{cells,trace_height} for segmentation strategy * chore: update lock files * chore: update lock files * chore: update lock files * fix: directly use vmexecutor * chore: add e2 execution to benchmark (#428) * chore: add e2 execution to benchmark * fix: typo * fix: typo * chore: update deps * chore: nit * chore: logging (#429) * chore: logging * update run.sh * update patch * tmp: fix tag until we merge main * chore: update openvm and all patches... * feat: use branch openvm-prof * update lock * ci: use ssh key * fix: wasn't patching k256 * debug bencher * update openvm-prof * Revert "debug bencher" This reverts commit c4c489a. * chore: update lock * feature flag evm-verify * update branch * update cargo dependencies (but still need patch for revm) * chore(ci): display inputs/result link in summary * fix: update default segment_max_cells * chore: show results in summary * chore: update openvm * chore: remove patching (#434) * chore: update patching to openvm:feat/new-execution * fix: update patch workflow to handle double patch * chore: update stark-backend and openvm * chore: update lock files * fix: use hyphen for metered * chore: update commit * chore: update openvm/stark-backend to fix span * chore: update deps * fix: use debug directory for dev build * fix: remove widths from metered call (#439) * chore: update lock files * chore: fix benchmark for new traits (#445) * chore: fix benchmark for new traits Make interface changes to coordinate with openvm-org/openvm#1836 * chore: update openvm for trace_gen_time_ms * feat(ci): allow passing block number to compare workflow * fix: replace --unstable-presymbolicate with --presymbolicate * feat(new-execution): upload execute profile to firefox (#448) * feat: upload execute profile to firefox; store gzipped * chore: update lock files * feat: add block number to run name * chore: update lock files * fix(new-execution): remove segmentation strategy (#450) * chore: update lock files * chore: update openvm commit * fix: use interpreter to run execute_metered (#456) * chore: update openvm traits (#462) * chore: update lock files * chore: update `openvm_init.rs` to new format (#466) * fix(ci): use cargo-openvm from target rev (#467) * chore: update openvm commit * fix: constrain storage roots (#469) We check the state root against a trusted state root, but we missed this check on storage roots. This PR fixes this bug by checking storage trie roots against account storage roots, which are stored and constrain in the state trie. * feat(new-execution): update to use new sdk interface (#470) * feat: update to use new sdk interface * fix: update params dir * fix: checkout correct openvm rev * fix: use var * fix: remove public value config override * chore: comment state root and fix handling of empty accounts (#468) * chore: comment out state root check * fix: properly handle empty accounts in the Rpc DB * fix: make guest change to handle empty accounts * fix: compute bytecode hash instead of reading from proof to sidestep different handling of empty codehash by rpc providers * chore: address comments * chore: update lock files * fix: `BenchmarkCli` max cells override (#471) * chore: update start-backend tag (#473) * chore: replace feat/new-execution with main * chore: update revm dep --------- Co-authored-by: Yi Sun <yi-sun@users.noreply.github.com> Co-authored-by: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com> Co-authored-by: Valery <valery@intrinsictech.xyz> Co-authored-by: Shayan Hosseini <shayan.hosseiny@gmail.com> Co-authored-by: HrikB <23041116+HrikB@users.noreply.github.com>

jonathanpwang and others added 30 commits June 27, 2025 20:53

Merge branch 'feat/new-execution' into feat/new-exec-device

fb168d4

wip: update ExecutionControl trait; disable InterpretedInstance since…

53eab5e

… we are rewriting it

feat: rewrite ProgramHandler with preprocessed insn lookup

79528ad

feat: fixed all ExecutionControl impls

64df2c0

feat: reworked VmSegmentExecutor

b040309

wip: still working on vm.rs

ac1c90c

chore: update rv32im for VmChipWrapper and **Filler

54472aa

chore: default derive VmConfig trait alias

eecd740

wip: comment out SingleSegmentExecutor

1a68cbd

wip: refactored VirtualMachine::preflight_execute_from_state

6e2b02e

feat: VirtualMachine holds pkey

59ff4e3

feat: remove AccessAdapterInventory but keep RecordArena in TracingMe…

aa5a0a4

…mory chore: remove generic <F> from TracingMemory

feat: stabilize SystemRecords interface

e123d7c

added access_adapter_records and touched_memory

refactor: move .finalize, .handle_touched to TracingMemory

9b1151e

feat: complete execute_preflight_from_state

feba7f1

chore: return final memory state as well

cc154f5

feat: VirtualMachine::prove

02fc909

feat: return final memory on success

81972b5

wip

c469c46

chore: fix trait bound issue in air_test_impl

2bb93e6

feat: air_test

603df86

feat: impl new ext traits (#1826)

9c4f541

Toward INT-4363 implemented the new extension traits for `rv32im` and `native` extensions. `cargo check` passes for both however had to comment out bunch of stuff (mainly testing related) update: implemented keccak and sha2 as well

chore: re-enable stark_utils

10ed1d0

chore: make AdapterTraceFiller: Send + Sync

e47bbc3

fmt

0410d48

wip: update derive macro for VmConfig

19158b1

wip: add generics attribute to derive macro

7fda251

chore: deduplicate StepExecutorE1 -> InsExecutorE1

d484814

jonathanpwang added 3 commits July 23, 2025 01:12

feat: finalize metrics

a3cdc2c

chore: update prof and metric docs

a534932

fix: sdk test (now segment check every 1000insns)

91175c6

jonathanpwang force-pushed the feat/new-exec-device branch from 5672570 to 91175c6 Compare July 23, 2025 15:40

fix: e2e benchmarks

672c9ef