Linear constraint remover #2425

Schaeff · 2025-02-01T16:21:24Z

what

Replace certain multilinear constraints* with intermediate polynomials in pilopt.
*with no next references, not touching public inputs

why

This comment highlights a use-case where materialising a column which is only constrained to a non-shifted multilinear polynomial is wasteful. Fixing this when we create the code is tricky, since we cannot know the degree of the constraint from pil. This is however something pilopt can do, which is what we attempt to do here

unknowns

We have similar optimizers for x = 42 and x = y which are both multilinear expressions. They do no yield intermediate polynomials, since in that case it is fine to "inline" them. This implementation currently makes sure all three optimizer apply to distinct cases, but maybe it would be better to have a single optimizer which handles all three cases, only introducing intermediate polynomials in the non trivial cases (42 and y)
- Implemented in Better optimization of intermediate polynomials #2532

ast/src/parsed/mod.rs

…aint-remover

@Schaeff

Depends on #2426. With this PR, we use the annotations added in #2412 and #2426 to find the columns being generated. This allows us to detect when the folded payload is not persisted as a witness column. The background is that @Schaeff is working on #2425, which would undo persisting the folded payloads in some cases, allowing us to spend fewer witness columns per bus interaction. With this PR, this should be fine: The column type will change from witness to intermediate, which means that the bus witgen will not output any folded columns. It can be tested by changing the `materialize_folded` bool to `false`, e.g. [here](https://github.com/powdr-labs/powdr/blob/e77d3801c1decff039fd0ec6bbeb55ed734357fb/std/protocols/bus.asm#L48) and running: ``` cargo run pil test_data/asm/block_to_block.asm \ -o output -f --linker-mode bus --prove-with mock --field gl ``` This used to fail before this PR.

…aint-remover

pilopt/tests/optimizer.rs

pilopt/src/lib.rs

qwang98 · 2025-02-10T05:29:01Z

pilopt/src/lib.rs

+    // We require `left` to be a single, non-shifted, witness column `w`
+    let w = if let AlgebraicExpression::Reference(
+        r @ AlgebraicReference {
+            poly_id:
+                PolyID {
+                    ptype: PolynomialType::Committed,
+                    ..
+                },
+            next: false,
+            ..
+        },
+    ) = left.as_ref()
+    {
+        r
+    } else {
+        return None;
+    };
+
+    // we require `y` to be a multi-linear expression over non-shifted columns
+    if right.contains_next_ref(intermediate_columns)
+        || right.degree_with_cache(intermediate_columns, &mut Default::default()) != 1
+    {
+        return None;
+    }
+
+    // we require `y` not to be a single, non-shifted, witness column, because this case is already covered and does not require an intermediate column
+    // TODO: maybe we should avoid this edge case by removing the other optimizer, see xxx
+    if matches!(
+        right.as_ref(),
+        AlgebraicExpression::Reference(AlgebraicReference {
+            poly_id: PolyID {
+                ptype: PolynomialType::Committed,
+                ..
+            },
+            next: false,
+            ..
+        })
+    ) {
+        return None;
+    }


I'm a bit confused by this chunk. It seems that we have opposite requirements for left and right being a single, non-shifted, witness column AND also no requirements for left being a multi-linear expression. In practice, shouldn't we have both left and right be single and non shifted multilinear expression OR witness column?

The current pattern is:

w = lin

where w is a non shifted witness column and lin is a multilinear expression over non-shifted columns.
Additionally, we require lin NOT to be a single non shifted column, because we already have another optimizer for this case that we don't want to interfere with here.

Are you suggesting that we could also support lin = w with the same requirements?

but both are possible cases, right? I don't see why we should consider w = lin and not do it for lin = w, could you elaborate on that?

pilopt/src/lib.rs

test_data/pil/simple_div.pil

qwang98 · 2025-02-10T07:02:42Z

backend/src/mock/mod.rs

-            .flat_map(|pil| pil.identities.iter())
-            .flat_map(|identity| identity.all_children())
+            .flat_map(|pil| {
+                pil.identities
+                    .iter()
+                    .flat_map(|identity| identity.all_children())
+                    .chain(
+                        // Note: we iterate on a `HashMap` so the ordering is not guaranteed, but this is ok since we're building another map.
+                        pil.intermediate_columns
+                            .values()
+                            .flat_map(|(_, def)| def.iter().flat_map(|d| d.all_children())),
+                    )
+                    .collect_vec()
+            })


Is this edit because we now move some identities to intermediates, so we have to use both as the preimage of challenge derivation, in order to guarantee sufficient randomness?

No it's simpler than that:

col a = challenge(0, 0); // the challenge is declared in the intermediate definition a = a * a; // in the identity, the challenge does not appear

Here we simply look through all expressions as opposed to starting with identities and visiting the definitions of the intermediates they access, which is the same assuming that unused intermediates have been optimized away. If not, we may end up introducing challenges which are not needed.

backend/src/stwo/circuit_builder.rs

…aint-remover

chriseth · 2025-03-10T15:15:01Z

ast/src/analyzed/mod.rs

+        self.identities
+            .iter()
+            .flat_map(|identity| identity.all_children())
+            .chain(
+                // Note: we iterate on a `HashMap` so the ordering is not guaranteed, but this is ok since we're building another map.
+                self.intermediate_columns
+                    .values()
+                    .flat_map(|(_, def)| def.iter().flat_map(|d| d.all_children())),
+            )


Suggested change

self.identities

.iter()

.flat_map(|identity| identity.all_children())

.chain(

// Note: we iterate on a `HashMap` so the ordering is not guaranteed, but this is ok since we're building another map.

self.intermediate_columns

.values()

.flat_map(|(_, def)| def.iter().flat_map(|d| d.all_children())),

)

self.all_children()

chriseth · 2025-03-10T15:17:20Z

backend/src/stwo/circuit_builder.rs

+// This function creates a list of the names of the constant polynomials that have next references
+// Note that the anaylsis should also dereference next references to intermediate polynomials


Suggested change

// This function creates a list of the names of the constant polynomials that have next references

// Note that the anaylsis should also dereference next references to intermediate polynomials

/// This function creates a list of the names of the constant polynomials that have next references

/// Note that the anaylsis should also dereference next references to intermediate polynomials

chriseth · 2025-03-10T15:18:47Z

pil-analyzer/src/expressionizer.rs

@@ -58,6 +58,14 @@ pub fn try_value_to_expression<T: FieldElement>(
    Expressionizer { poly_id_to_name }.try_value_to_expression(value)
 }

+/// Tries to convert an algebraic expression to an expression.


Suggested change

/// Tries to convert an algebraic expression to an expression.

/// Tries to convert an algebraic expression to an expression that evaluates to the same value.

chriseth · 2025-03-10T15:21:29Z

pilopt/src/lib.rs

@@ -558,6 +563,178 @@ fn extract_constant_lookups<T: FieldElement>(pil_file: &mut Analyzed<T>) {
    }
 }

+/// Identifies witness columns that are constrained to a non-shifted (multi)linear expression, replaces the witness column by an intermediate polynomial


Can you extend the comment a little? What is the intermediate polynomial and why are we doing this optimization?

chriseth · 2025-03-10T15:22:08Z

pilopt/src/lib.rs

+        .collect();
+
+    // pattern match identities looking for `w * (1 - w) = 0` and `(1 - w) * w = 0` constraints
+    let boolean_constrained_witnesses = pil_file


Can this be a helper function?

chriseth · 2025-03-10T15:26:13Z

pilopt/src/lib.rs

@@ -686,6 +886,67 @@ fn constrained_to_constant<T: FieldElement>(
    None
 }

+/// Identifies witness columns that are constrained to a (multi)linear expression.


Can you specify the semantics a bit more concretely? What does it return?

chriseth · 2025-03-10T15:27:25Z

pilopt/src/lib.rs

+        return None;
+    };
+
+    // We require `left` to be a single, non-shifted, witness column `w`


Should we also try with left and right swapped?

ShuangWu121 · 2025-03-11T09:21:25Z

pilopt/src/lib.rs

@@ -891,6 +1118,65 @@ fn equal_constrained<T: FieldElement>(
    }
 }

+/// Tries to extract a boolean constrained witness column from a polynomial identity.
+/// The pattern used is `x * (1 - x) = 0` or `(1 - x) * x = 0` where `x` is a witness column.


The pattern for boolean constraints is strictly limited to x * (1-x) or (1 - x) * x, There will never be expressions like x * (x - 1), (x - 1) * x or x * x - x, This is guaranteed by the way PIL generates constraints. Is my understanding correct?

Correct! It aligns with how witgen detects boolean constraints.

I think the global_constraints.rs can detect (x-1)*x = 0 as well. From the above, only x * x - x = 0 is not detected.

…aint-remover

chriseth · 2025-03-11T14:21:14Z

backend/src/stwo/circuit_builder.rs

-                // add the name of the constant polynomial to the list
-                constant_with_next_list.insert(name);
+/// This function creates a list of the names of the constant polynomials that have next references
+/// Note that the anaylsis should also dereference next references to intermediate polynomials


Suggested change

/// Note that the anaylsis should also dereference next references to intermediate polynomials

/// Note that the analysis should also dereference next references to intermediate polynomials

chriseth · 2025-03-11T14:25:29Z

pilopt/src/lib.rs

@@ -558,6 +563,129 @@ fn extract_constant_lookups<T: FieldElement>(pil_file: &mut Analyzed<T>) {
    }
 }

+/// Identifies witness columns that are constrained to a non-shifted (multi)linear expression, replaces the witness column by an intermediate polynomial.


Suggested change

/// Identifies witness columns that are constrained to a non-shifted (multi)linear expression, replaces the witness column by an intermediate polynomial.

/// Identifies witness columns that are constrained to a non-shifted (multi)linear expression, replaces the witness column by an intermediate polynomial defined to be that expression.

(obvious, but this is what I was missing to understand it :) )

Ah and maybe also use intermediate column instead of polynomial?

chriseth · 2025-03-11T14:30:16Z

pilopt/src/lib.rs

+        .iter()
+        .filter_map(|(index, ((name, old_poly_id), expression))| {
+            // Remove the definition. If this fails, we have already replaced it by an intermediate column.
+            pil_file.definitions.remove(name).map(|(symbol, value)| {


Can you turn this into a function inside Analyzed? I think we should avoid modifying it directly, especially when we have to deal with creating new IDs and so on.

I had a look into this and it's a bit tricky:

extracting it requires going from a valid pil to a valid pil

here we're replacing a witness column by an intermediate column with a given definition

the result is a valid pil only after we replace all references to the original witness by references to the intermediate. This is because poly_id is not global but local to each polynomial type, so we cannot reuse the one from the witness :(

the function which does that is pilopt::substitute_polynomial_references, which relies on try_algebraic_expression_to_expression, which is in pil-analyzer

therefore, substitute_polynomial_references cannot be turned into a method on Analyzed without causing a circular dep

This could be simplified a lot if poly_ids were global, then we could simply replace

col witness a; a = lin

by

col a = lin

reusing the poly_id

Another option would be:

go through the columns to remove and rename them to some dummy names, do the same for their references (this could be a rename method on Analyzed)

introduce the new intermediates with the correct names, which is valid since the names are now unused

replace the references to the dummy names by references to the correct names

remove the identities involving the dummy witnesses

I think these four steps are valid, although they involve a lot of back and forth. @chriseth do you have any thoughts on these? I personally like the former more (making poly_id global) but I think it's a big change I wouldn't want this PR to be blocked by.

ShuangWu121 · 2025-03-12T08:32:12Z

pilopt/src/lib.rs

+/// - `Some(((name, poly_id), expression))` if the identity matches the pattern, where:
+///     - `name` is the name of the witness column
+///     - `poly_id` is the id of the witness column
+///     - `expression` is the (multi)linear expression that the witness column is constrained to


Suggested change

/// - `expression` is the (multi)linear expression that the witness column is constrained to

/// - `expression` is the (multi)linear expression that the witness column is constrained by

wip

ee1fe14

georgwiese mentioned this pull request Feb 3, 2025

Manual Bus Witgen: Output folded columns only if they exist #2428

Merged

chriseth reviewed Feb 3, 2025

View reviewed changes

ast/src/parsed/mod.rs Outdated Show resolved Hide resolved

georgwiese mentioned this pull request Feb 3, 2025

Add folded payload columns to PhantomBusInteraction #2426

Merged

Schaeff added 4 commits February 3, 2025 15:11

use existing algexpr->expr conversion, adjust tests

4901fa0

add pilopt

8922a92

Merge branch 'main' of github.com:powdr-labs/powdr into linear-constr…

517cd66

…aint-remover

pull main

1ef509d

Schaeff mentioned this pull request Feb 4, 2025

Batch two bus interactions to the same set of accumulator columns #2435

Merged

Schaeff added 10 commits February 5, 2025 12:29

fix wrong symbol id

f208365

Merge branch 'main' of github.com:powdr-labs/powdr into linear-constr…

58e5b23

…aint-remover

Merge branch 'main' of github.com:powdr-labs/powdr into linear-constr…

471014b

…aint-remover

visit intermediate polynomials when collecting challenges

feac232

fix more tests

9269faf

Merge branch 'main' of github.com:powdr-labs/powdr into linear-constr…

875f11e

…aint-remover

fix stwo next reference detection

862e2e6

fix external witgen tests

779a974

fix invalid witness test

4ea55f5

fix div test

48128b4