Generalize multiplicity witgen #2043

georgwiese · 2024-11-05T10:25:57Z

(mostly cherry-picked from #1958)

This PR generalizes the multiplicity witness generation we had in FixedLookup in the following ways:

Instead of searching for a column of a specific name, takes the witness column from the phantom lookups
As a result, it is also able to handle arbitrarily many lookups
The witness generation should work for all machine types which can be connected via lookups: Fixed lookup, block machines, and secondary VMs.

A common use-case that is not yet covered is global range constraints, as they are not represented as machine calls in witness generation. This should be fixed in a follow-up PR.

executor/src/witgen/generator.rs

chriseth · 2024-11-07T10:20:53Z

executor/src/witgen/machines/machine_extractor.rs

+        })
+        .collect::<BTreeMap<_, _>>();
+
+    // Multiplicity columns are not connected to any machine, because we removed identities that reference


I see that the higher-stage polynomial identities involving multiplicity columns are removed, but what about the phantom lookup? Shouldn't the RHS columns include the multiplicity column and thus be connected to the rest of the columns?

So I think we should change all_row_connected_witnesses to include the multiplicity column on the RHS columns.

No, the multiplicity does not appear in the RHS! It is an extra field in the phantom lookup that is the ignored by the rest of witgen (which treats phantom lookups and native ones the same).

But shouldn't we change all_row_connected_witnesses?

I added the multiplicities to lookup_witnesses, which is the starting point of all_row_connected_witnesses. Now, they don't need to be explicitly removed anymore.

chriseth · 2024-11-07T10:25:28Z

executor/src/witgen/machines/machine_extractor.rs

@@ -186,16 +229,41 @@ pub fn split_out_machines<'a, T: FieldElement>(
    // Note that this machine comes last, because some machines do a fixed lookup
    // in their take_witness_col_values() implementation.
    // TODO: We should also split this up and have several instances instead.
+
+    // Compute sizes of fixed lookup multiplicity columns.


Is it possible to move the creation of the fixed lookup machine to a helper function?

Moved to FixedLookup::new

executor/src/witgen/machines/mod.rs

executor/src/witgen/machines/multiplicity_counter.rs

chriseth · 2024-11-07T10:29:47Z

executor/src/witgen/machines/multiplicity_counter.rs

+
+    /// A map poly_id -> (index -> count).
+    /// Has a (possibly entry) entry for each value in `identity_id_to_multiplicity_column`.
+    counts: BTreeMap<PolyID, BTreeMap<usize, usize>>,


I might be better to use Vec instead of a BTreeMap for the inner type. Did you test that?

Hmm, a Vec that would be resized on demand (because we don't know the number of rows a priori)?

I have not tested this, but:

The counts should be accessed at most once per machine call, so I wouldn't expect a BTreeMap access to be a bottleneck.

I'd expect the vector to be sparse in practice. We typically use permutations to connect machines, so in practice I'd expect LogUp to be used only for fixed lookups (typically of size $2^{16}$ to $2^{18}$) and global range constraints. Only the former is handled in this PR (range constraints will need another pass over the data in an upcoming PR).

executor/src/witgen/generator.rs

Schaeff · 2024-11-07T10:14:30Z

executor/src/witgen/machines/multiplicity_counter.rs

+
+    /// For a given identity ID, increments the count of the corresponding multiplicity column at the given index.
+    pub fn increment(&mut self, identity_id: u64, index: usize) {
+        if let Some(poly_id) = self.identity_id_to_multiplicity_column.get(&identity_id) {


Shouldn't this panic if the identity is not in there instead of currently silently continuing?

The increment function is called for any machine call, but unless the machine is connected via a phantom lookup, no multiplicity needs to be tracked and this is a no-op, because identity_id_to_multiplicity_column would not have a key for identity_id.

executor/src/witgen/machines/block_machine.rs

Schaeff · 2024-11-07T10:56:48Z

executor/src/witgen/machines/multiplicity_counter.rs

+    pub fn generate_columns_single_size<T: FieldElement>(
+        &self,
+        size: DegreeType,
+    ) -> BTreeMap<PolyID, Vec<T>> {
+        self.counts
+            .iter()
+            .map(|(poly_id, counts)| {
+                let mut column = vec![T::zero(); size as usize];
+                for (index, count) in counts {
+                    column[*index] = T::from(*count as u64);
+                }
+                (*poly_id, column)
+            })
+            .collect()
+    }
+
+    /// Materializes the multiplicity columns, using different sizes for each of them.
+    pub fn generate_columns_different_sizes<T: FieldElement>(
+        &self,
+        sizes: BTreeMap<PolyID, DegreeType>,
+    ) -> BTreeMap<PolyID, Vec<T>> {
+        self.counts
+            .iter()
+            .map(|(poly_id, counts)| {
+                let size: DegreeType = sizes[poly_id];
+                let mut column = vec![T::zero(); size as usize];
+                for (index, count) in counts {
+                    column[*index] = T::from(*count as u64);
+                }
+                (*poly_id, column)
+            })
+            .collect()


Is there any way we can reduce duplication here?

Right, I can call generate_columns_different_sizes in generate_columns_single_size.

Co-authored-by: chriseth <[email protected]>

georgwiese · 2024-11-07T12:40:54Z

Nice, adding the multiplicity to Connection instead of maintaining another map improved the code a lot I think! Now, the code to create the fixed lookup is outrageously verbose and contains many duplications, planning to fix that before the next round of review.

georgwiese · 2024-11-07T17:27:10Z

OK, I think this is ready for another look :)

executor/src/witgen/data_structures/multiplicity_counter.rs

Co-authored-by: chriseth <[email protected]>

chriseth · 2024-11-08T09:57:06Z

executor/src/witgen/machines/machine_extractor.rs

+                .insert(connection.id, *connection)
+                .is_none());
+            if let Some(multiplicity) = connection.multiplicity_column {
+                remaining_witnesses.remove(&multiplicity);


What if the multiplicity column occurs in a row-connected identity? I think we should at least check that.

In this branch, we're talking about a fixed lookup, right? So the entire RHS is fixed columns. Do you mean that there could be constraints that connect those to witness columns?

If yes, I think that would be quite esoteric (those witness columns would not depend on an input and could be derived from the fixed columns at compile time), and it was also not checked before this PR (we continued for any connection without witness columns on the RHS).

chriseth · 2024-11-08T09:57:39Z

executor/src/witgen/machines/machine_extractor.rs

+
+    for connection in &all_connections {
+        // If the RHS only consists of fixed columns, record the connection and continue.
+        if connection.is_lookup()


Maybe extract the condition into a function?

Right. I added a FixedLookup::is_responsible. In the near future, it should just have a try_new() like any other machine...

chriseth · 2024-11-08T09:58:23Z

executor/src/witgen/machines/machine_extractor.rs

+
+        // Extract all witness columns in the RHS of the lookup.
+        let lookup_witnesses = {
+            let callee_columns = refs_in_selected_expressions(connection.right)


Why not introduce a function refs_in_connection_rhs that includes the multiplicity?

Done! I think this actually fixed a bug: If there are multiple phantom lookups with distinct multiplicities, we now include all multiplicities in the machine witness.

executor/src/witgen/machines/mod.rs

Co-authored-by: chriseth <[email protected]>

georgwiese added 5 commits November 5, 2024 11:15

Remove existing multiplicity witgen

393b34d

Implement multiplicities for FixedLookup and block machines

2005166

Add multiplicity witgen for secondary VM

e372c8e

Add multiplicity counter file...

b732aa1

Fix compile errors

a1d22c9

georgwiese mentioned this pull request Nov 5, 2024

PIL LogUp implementation: Generate multiplicity column #2042

Merged

georgwiese added 3 commits November 5, 2024 12:02

Refactor

983028c

Clippy

01f0cfa

Refactor

4f2e6bc

georgwiese changed the title ~~[WIP] Generalize multiplicity witgen~~ Generalize multiplicity witgen Nov 5, 2024

georgwiese marked this pull request as ready for review November 5, 2024 18:31

georgwiese added 2 commits November 6, 2024 17:44

Merge origin/main

fbf0ff7

Fix expected panic message

158a49d

Schaeff self-requested a review November 7, 2024 08:44

chriseth reviewed Nov 7, 2024

View reviewed changes

executor/src/witgen/generator.rs Outdated Show resolved Hide resolved

chriseth reviewed Nov 7, 2024

View reviewed changes

executor/src/witgen/machines/mod.rs Outdated Show resolved Hide resolved

chriseth reviewed Nov 7, 2024

View reviewed changes

executor/src/witgen/machines/multiplicity_counter.rs Outdated Show resolved Hide resolved

chriseth reviewed Nov 7, 2024

View reviewed changes

executor/src/witgen/machines/multiplicity_counter.rs Outdated Show resolved Hide resolved

chriseth reviewed Nov 7, 2024

View reviewed changes

Schaeff reviewed Nov 7, 2024

View reviewed changes

georgwiese and others added 5 commits November 7, 2024 12:29

Apply suggestions from code review

c43a0d0

Co-authored-by: chriseth <[email protected]>

Small review feedback

34cd913

Merge origin/refactor-multiplicity-witgen

f40a47d

Move multiplicity counter

b7991a0

Add (optional) multiplicity column to Connection

f8c1b98

georgwiese added 2 commits November 7, 2024 17:13

Add multiplicities to lookup witness

c645533

Handle fixed lookup connections better

6da56a2

georgwiese added 4 commits November 7, 2024 17:33

Bug fix

4234b89

Pull sizes computation into FixedLookup::new

81af97b

Refactor

9904bd6

Undo renaming

483e038

chriseth reviewed Nov 8, 2024

View reviewed changes

executor/src/witgen/data_structures/multiplicity_counter.rs Outdated Show resolved Hide resolved

chriseth reviewed Nov 8, 2024

View reviewed changes

executor/src/witgen/data_structures/multiplicity_counter.rs Outdated Show resolved Hide resolved

chriseth reviewed Nov 8, 2024

View reviewed changes

executor/src/witgen/data_structures/multiplicity_counter.rs Show resolved Hide resolved

georgwiese commented Nov 8, 2024

View reviewed changes

executor/src/witgen/data_structures/multiplicity_counter.rs Outdated Show resolved Hide resolved

Apply suggestions from code review

0e75e4a

Co-authored-by: chriseth <[email protected]>

chriseth reviewed Nov 8, 2024

View reviewed changes

Have FixedLookup take some responsibility

7cd62f7

chriseth reviewed Nov 8, 2024

View reviewed changes

executor/src/witgen/machines/mod.rs Outdated Show resolved Hide resolved

georgwiese and others added 3 commits November 8, 2024 11:20

Update executor/src/witgen/machines/mod.rs

50fadcf

Co-authored-by: chriseth <[email protected]>

Introduce refs_in_connection_rhs()

f7e0add

Clippy

7af1c8a

chriseth approved these changes Nov 8, 2024

View reviewed changes

chriseth enabled auto-merge November 8, 2024 11:21

chriseth added this pull request to the merge queue Nov 8, 2024

Merged via the queue into main with commit 2c016d5 Nov 8, 2024
14 checks passed

chriseth deleted the refactor-multiplicity-witgen branch November 8, 2024 12:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Generalize multiplicity witgen #2043

Generalize multiplicity witgen #2043

georgwiese commented Nov 5, 2024

chriseth Nov 7, 2024

georgwiese Nov 7, 2024

chriseth Nov 7, 2024

georgwiese Nov 7, 2024

chriseth Nov 7, 2024

georgwiese Nov 7, 2024

chriseth Nov 7, 2024

georgwiese Nov 7, 2024

Schaeff Nov 7, 2024

georgwiese Nov 7, 2024 •

edited

Loading

Schaeff Nov 7, 2024

georgwiese Nov 7, 2024

georgwiese commented Nov 7, 2024

georgwiese commented Nov 7, 2024

chriseth Nov 8, 2024

georgwiese Nov 8, 2024

chriseth Nov 8, 2024

georgwiese Nov 8, 2024

chriseth Nov 8, 2024

georgwiese Nov 8, 2024

Generalize multiplicity witgen #2043

Generalize multiplicity witgen #2043

Conversation

georgwiese commented Nov 5, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

georgwiese Nov 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

georgwiese commented Nov 7, 2024

georgwiese commented Nov 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

georgwiese Nov 7, 2024 •

edited

Loading