Apc stacking #2917

pacheco · 2025-06-18T21:31:25Z

This PR adds APC stacking to the powdr openvm pipeline.
On the openvm integration side, the main idea is adding a PowdrStackedPrecompile which may contain one or more of the existing PowdrPrecompile, referenced by opcode.
Similarly, the PowdrChip now contains a map of executors, one per precompile, by opcode.
This way, the adaptations were relatively straightforward (e.g., witgen).
To enable stacking, there's a chip_stacking_log parameter which, if set, enables chip stacking with bucketing based on the given log of the number of witness columns.

- another unknown instruction - handle interactions to the same bus id with different number of args

…king

autoprecompiles/src/lib.rs

Schaeff · 2025-06-20T16:18:49Z

autoprecompiles/src/lib.rs

    vm_config: VmConfig<T, B>,
    degree_bound: DegreeBound,
    opcode: u32,
+    strict_is_valid_guards: bool,


Can you add some documentation on this parameter?

What about this?

/// "Guarding" constraints means that the prover can choose to disable them, which is useful for /// filling up rows until the next power of two. A new boolean variable (`is_valid`) introduced /// that allows the prover to disable constraints. enum GuardingMode { /// Makes sure that setting all variables to zero satisfies the constraints. AllowZeroWitness, /// Makes sure that all constraints are satisfied as long as `is_valid` is set to 0. /// This mode always increases the constraint degree by one. Strict, }

Schaeff · 2025-06-20T16:19:12Z

autoprecompiles/src/lib.rs

 fn add_guards<T: FieldElement>(
    mut machine: SymbolicMachine<T>,
    bus_map: BusMap,
+    strict: bool,


Suggested change

strict: bool,

is_strict: bool,

Schaeff · 2025-06-20T16:20:52Z

autoprecompiles/src/lib.rs

+        machine
+            .constraints
+            .into_iter()
+            .map(|mut c| {
+                c.expr = is_valid.clone() * c.expr.clone();
+                c
+            })
+            .collect()
+    } else {


Does this disable some cheaper stuff that add_guards_constraint does in the non strict case? It seems like is_strict could be pulled into add_guards_constraint

add_guards_constraint is a recursive fn that ensures the degree doesn't increase... which is not the case here, so opted to be explicit here

Schaeff · 2025-06-20T16:24:42Z

expression/src/lib.rs

+impl<T: Clone + Ord, R: Clone + Ord> AlgebraicExpression<T, R> {
+    /// Get a canonical form by expanding the expression and reordering operands.
+    pub fn normalize(self) -> Self {
+        normalize(expand(self))
+    }
+}


Could this be done by going through quadratic expression? cc @georgwiese

Not quite. GroupedExpression (formerly QuadraticSymbolicExpression) has a canonical representation for expressions up to degree 1...

We could do this by introducing a CanonicalExpression type which is essentially a BTreeMap<BTreeSet<R>, T>. As long as it implements Add + Sub + Mul + Neg, AlgebraicExpression::to_expression can convert to it (via the ExpressionConvertible trait).

If AlgebraicExpression also implements From<CanonicalExpression>, we could implement the normalization as something like

self.to_expression( &|number| CanonicalExpression { coefficients: [(BTreeSet::new(), *number)].into_iter().collect() }, |reference| CanonicalExpression { coefficients: ([*reference].into_iter().collect(), T::one())].into_iter().collect() }).into()

Not sure if it's worth it, but seems a bit cleaner.

Schaeff · 2025-06-20T16:32:12Z

openvm/src/powdr_extension/chip.rs

-        );
+        let height = next_power_of_two_or_zero(num_records);
+
+        let trace = if self.executors.len() == 1 {


Stupid question: why do we need an if else, aka why isn't the else block also valid for a single opcode?

it is valid, but its wasteful: when stacking we need to create the table and then copy the witness from each executor... this could be improved (TODO)

openvm/src/powdr_extension/vm.rs

georgwiese

Nice!

I wasn't able to review the code of the AIR stacking itself yet, but I did run it and skimmed the results, looks really cool!

Maybe it's also fine, because we're not sure how much we'll use this feature and we should flag it as experimental anyway at this point. I think it's more important that the feature doesn't leak too much into the rest of the code base. I have a few smaller comments that might help with that.

georgwiese · 2025-07-02T13:37:35Z

openvm/src/customize_exe/air_stacking.rs

+        let mut remapped = pcp.machine.clone();
+        remapped.pre_visit_expressions_mut(&mut |expr| {
+            if let AlgebraicExpression::Reference(r) = expr {
+                assert!(r.id <= is_valid_start);


Suggested change

assert!(r.id <= is_valid_start);

assert!(r.id < is_valid_start);

No?

georgwiese · 2025-07-02T13:39:59Z

openvm/src/customize_exe/air_stacking.rs

+                    // we assume each pcp to have a specific column named "is_valid"
+                    assert!(*r.name == "is_valid");
+                    r.id = is_valid_new_id;
+                    r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()).into();


This isn't going necessarily going to be unique, is it? Maybe idx would be better?

georgwiese · 2025-07-02T14:11:42Z

cli-openvm/src/main.rs

        input: Option<u32>,
+
+        #[arg(long)]
+        chip_stacking: Option<f32>,


Could you add a description here, with a note that the feature is experimental?

georgwiese · 2025-07-04T09:27:26Z

autoprecompiles/src/lib.rs

    vm_config: VmConfig<T, B>,
    degree_bound: DegreeBound,
    opcode: u32,
+    strict_is_valid_guards: bool,


What about this?

/// "Guarding" constraints means that the prover can choose to disable them, which is useful for /// filling up rows until the next power of two. A new boolean variable (`is_valid`) introduced /// that allows the prover to disable constraints. enum GuardingMode { /// Makes sure that setting all variables to zero satisfies the constraints. AllowZeroWitness, /// Makes sure that all constraints are satisfied as long as `is_valid` is set to 0. /// This mode always increases the constraint degree by one. Strict, }

georgwiese · 2025-07-04T09:38:32Z

expression/src/lib.rs

+            left,
+            op: AlgebraicBinaryOperator::Sub,
+            right,
+        }) => normalize(*left) - normalize(*right),


This seems sus, why would it be simpler than addition?

its simpler because there's no need to handle commutativity (no need to flatten and reorder)

georgwiese · 2025-07-04T09:56:18Z

expression/src/lib.rs

+impl<T: Clone + Ord, R: Clone + Ord> AlgebraicExpression<T, R> {
+    /// Get a canonical form by expanding the expression and reordering operands.
+    pub fn normalize(self) -> Self {
+        normalize(expand(self))
+    }
+}


Not quite. GroupedExpression (formerly QuadraticSymbolicExpression) has a canonical representation for expressions up to degree 1...

We could do this by introducing a CanonicalExpression type which is essentially a BTreeMap<BTreeSet<R>, T>. As long as it implements Add + Sub + Mul + Neg, AlgebraicExpression::to_expression can convert to it (via the ExpressionConvertible trait).

If AlgebraicExpression also implements From<CanonicalExpression>, we could implement the normalization as something like

self.to_expression( &|number| CanonicalExpression { coefficients: [(BTreeSet::new(), *number)].into_iter().collect() }, |reference| CanonicalExpression { coefficients: ([*reference].into_iter().collect(), T::one())].into_iter().collect() }).into()

Not sure if it's worth it, but seems a bit cleaner.

georgwiese · 2025-07-04T10:03:26Z

openvm/src/powdr_extension/vm.rs

+impl<P: IntoOpenVm> From<PowdrPrecompile<P>> for PowdrStackedPrecompile<P> {
+    fn from(pcp: PowdrPrecompile<P>) -> Self {
+        Self {
+            machine: pcp.machine.clone(),


Hmm, so by default, we store the machine twice. Maybe using an enum would be better?

georgwiese · 2025-07-04T11:07:03Z

expression/src/lib.rs

    }
 }

+impl<T: Clone + Ord, R: Clone + Ord> AlgebraicExpression<T, R> {


I think the changes to this file could be nice individual PR :)

georgwiese · 2025-07-04T11:10:51Z

openvm/src/customize_exe.rs

    original_config: &OriginalVmConfig,
    pgo_config: PgoConfig,
 ) -> Vec<BlockWithApc<P>> {
+    let strict_is_valid_guards = config.chip_stacking_log.is_some();


This is passed down a few functions that could just compute it themselves from the config.

I think I'd pass the config all the way to generate_autoprecompile (currently gets separate args for degree_bound and strict_is_valid_guards) and compute it there, before passing to powdr_autoprecompiles::build.

georgwiese · 2025-07-04T11:16:18Z

openvm/src/powdr_extension/chip.rs

-        );
+        let air = PowdrAir::new(precompile.machine);
+
+        let name = if precompile.precompiles.len() == 1 {


Could be a method of PowdrStackedPrecompile (which I think should be an enum, see comment above).

Also, maybe we should store the precompile directly? Looks like e.g. the name is re-computed later again.

pacheco added 30 commits May 19, 2025 13:19

always add guards

23002ae

apc chip stacking

b05703d

lower degree bound of optimize step

7a35e47

group by powers of two

323e706

use f64::log

f09fc34

Merge remote-tracking branch 'origin/main' into apc_stacking

a2d3683

Merge remote-tracking branch 'origin/main' into apc_stacking

bdff7c9

fix for scroll

5d060c5

- another unknown instruction - handle interactions to the same bus id with different number of args

another unsupported instruction

c6443f7

handle group of 1 precompile

c950d3f

Merge remote-tracking branch 'origin/main' into apc_stacking

cb4733d

Merge remote-tracking branch 'origin/main' into apc_stacking

e28585d

Merge remote-tracking branch 'origin/main' into apc_stacking

d142918

Merge commit '2ee7358e1f054955f3824de578c89475f6249ace' into apc_stac…

df67d47

…king

Merge commit '1a11972cc7930fe97a52e9a51ef3a45bcf5b6ec9' into apc_stac…

920edba

…king

Merge remote-tracking branch 'origin/main' into apc_stacking

50a1105

Merge remote-tracking branch 'origin/main' into apc_stacking

dedee09

Merge remote-tracking branch 'origin/main' into apc_stacking

09b512a

test joining constraints

4b3c5d4

sort

a7f5ff1

seems to work!

e48bf1b

seems to work

13850f6

try handle bus interactions

669b974

stuff

9dcff42

move air_stacking to own file

c8918f5

Merge branch 'apc_stacking' into apc_stacking_test2

6ff8258

default stacking 1.1

0f5448e

Merge remote-tracking branch 'origin/main' into apc_stacking

b8b0beb

Merge branch 'apc_stacking' into apc_stacking_test2

4de95d2

print max degree after handling constraints/interactions

4dd19ef

pacheco added 17 commits June 17, 2025 11:23

fmt

1075ba2

cleanup

428550b

Merge remote-tracking branch 'origin/main' into apc_stacking_cleanup

2ab7843

normalize expr

5193df6

Merge remote-tracking branch 'origin/main' into apc_stacking_cleanup

1ddff4f

cleanup

7ea0748

witgen cleanup

4473e9a

cleanup

35c8ba0

cleanup

c139514

clenaup

44c8c58

cleanup

a5483e9

Merge remote-tracking branch 'origin/main' into apc_stacking_cleanup

b85b763

test

c01c5ca

small change

54dd766

fix

3b8dd68

fmt

3ada389

clippy

f176a0f

Schaeff reviewed Jun 20, 2025

View reviewed changes

autoprecompiles/src/lib.rs Show resolved Hide resolved

Schaeff reviewed Jun 20, 2025

View reviewed changes

openvm/src/powdr_extension/vm.rs Outdated Show resolved Hide resolved

pacheco added 4 commits June 23, 2025 10:29

comments

857997a

Merge remote-tracking branch 'origin/main' into apc_stacking_cleanup

9dc6760

Merge remote-tracking branch 'origin/main' into apc_stacking_cleanup

1905935

fmt

be0c6b6

georgwiese reviewed Jul 4, 2025

View reviewed changes

Merge remote-tracking branch 'origin/main' into apc_stacking_cleanup

ed27779

	assert!(r.id <= is_valid_start);
	assert!(r.id < is_valid_start);

Apc stacking #2917

Are you sure you want to change the base?

Apc stacking #2917

Uh oh!

Conversation

pacheco commented Jun 18, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

georgwiese left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants