feat: send sync region

[alexander-camuto]

As discussed in #173, layouts for large circuits (100mil+) parameters can be quite time consuming.

Here we make the Region struct both Send and Sync so as to enable parallel layouts. We do so principally by constraining the inner element of Region (region: &'r mut dyn layouter::RegionLayouter<F>,) to be Send and Sync

pub trait RegionLayouter<F: Field>: fmt::Debug + Send + Sync { // add the new constraints here
    /// Enables a selector at the given offset.
...

---

Note: the copy constraint labels can be non-deterministic in the case of a parallelized layout. In particular, referring to step 2 of the algorithm in the Halo2 Book for permutation arguments:

2. Otherwise, left and right belong to different cycles. Make left the larger cycle and right the smaller one, by swapping them iff sizes(aux(left))<sizes(aux(right)).

The smaller cycle in a multi-threaded setting can be non-deterministic and as such, the constructed copy constraint labels can differ between the key generation and proof stage. A way to mitigate this is to make this label deterministic. What I've done in zkonduit@5eebade is to swap:

halo2/halo2_proofs/src/plonk/permutation/keygen.rs

Lines 83 to 85 in 17e9765

	if self.sizes[left_cycle.0][left_cycle.1] < self.sizes[right_cycle.0][right_cycle.1] {
	std::mem::swap(&mut left_cycle, &mut right_cycle);
	}

for a deterministic label propagation such as:

if right_cycle.0 >= left_cycle.0 && right_cycle.1 > left_cycle.1 {
            std::mem::swap(&mut left_cycle, &mut right_cycle);
}

The downside is that sometimes this means that the re-labelling of the right cycle can take more time (as its no longer necessarily the smaller cycle).

Given the potential impacts to the single threaded case we have remaining TODOs:

• make the cycle label propagation changes conditional on a parallel-region feature flag (such that when deactivated the cycle label propagation is as before)
• add some unit tests that leverage parallelism (rather than relying on the https://github.com/zkonduit/ezkl tests).

[CPerezz]

I'd change the feature flag name to something like: parallel-synthesis or similar.
Also, as mentioned, would add unit-tests to check that with a ton of small regions assigned in parallel we get the safety we expect.

Finally, (if possible) I'd add benchmarks and/or an example. I think with this the PR/feature addition can be accepted!
We would appreciate it so much!

[alexander-camuto]

Quick note. Had to reconfigure some of the Assembly struct used for creating permutation arguments. The current algorithm is currently non-deterministic in a multi-threaded setting. In particular though the cycles that are created have the same constituent elements, the ordering within that cycle can differ.

For an illustration of how the order in which the copy constraints are generated can create different cycle directions see:

To mitigate this we replace the order-dependent Vec used in the cycle construction algorithm with order-invariant BTreeSet

/// Struct that accumulates all the necessary data in order to construct the permutation argument.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Assembly {
    /// Columns that participate on the copy permutation argument.
    columns: Vec<Column<Any>>,
    /// Mapping of the actual copies done.
    cycles: Vec<usize, BTreeSet<(usize, usize)>>,
    /// Mapping of the actual copies done.
    aux: HashMap<(usize, usize), usize>,
    /// total length of a column
    col_len: usize,
    /// number of columns
    num_cols: usize,
}

The algorithm for constructing the cycles is then (for a copy constraint between two elements a and b)

1. Get cycle for a and b using aux (this is None if they are singletons not part of a cycle).
2. If they are part of the same cycle return.
3. Get the size of each element's cycle using the length of the cycles Vec entry (=1 if the cycle index is None)
4. Pick the smallest cycle to merge into the larger one
5. Merging cycles is as simple as:
   • Copy cycles entry for the smaller cycle into the larger cycle entry (it will do the sorting for us).
   • Remove entry for the smaller cycle.
6. Relabel aux to point to the new cycle index for the (ex)smaller cycle element.

The constructed cycles are then:

• each element of the BTreeSet points to the next element in the set
• the last element of the set points to the first element.

Note: Constant assignment, if done inside a parallel loop, will still be non-deterministic so we should add a warning somewhere against this

[alexander-camuto]

There's a chance that even in the single threaded case this is more efficient than the existing cycle creating algo -- as we don't store the entire matrix of size (num_cols * num_rows), only the elements that are part of a cycle.

[alexander-camuto]

Have added a small elementwise vector multiplication example (single threaded and multithreaded). To time the (mock) proving time for using the new mapping cycle algorithm:

cargo run --example vector-mul --features="thread-safe-region"

With the current algorithm

cargo run --example vector-mul 

To run the mock prover in parallel:

cargo run --example vector-mul-parallel-synth --features="thread-safe-region"

[alexander-camuto]

• Have added benchmarks as well

cargo bench --bench plonk-vec
cargo bench --bench plonk-vec --features thread-safe-region 
cargo bench --bench plonk-vec-parallel-synth --features thread-safe-region

• (timed) examples (see above). @CPerezz think its ready to review :)

[han0110]

Overall LGTM! For most comments left just want to save the lines of code so we could have less effort on maintaining the library, hope it make sense!

[han0110]

All LGTM now! Thanks for addressing all the comments!

[daira]

The downside is that sometimes this means that the re-labelling of the right cycle can take more time (as its no longer necessarily the smaller cycle).

Note that there are other algorithms that could be used to generate the cycles. As pointed out by @defuse, what we need to represent the sets of equality-constrained cells is a https://en.wikipedia.org/wiki/Disjoint-set_data_structure . The algorithm we used was designed for simplicity, and there are alternatives that might be more suited to parallelism.

For example:

In 1994, Richard J. Anderson and Heather Woll described a parallelized version of Union–Find that never needs to block.[8]

[daira]

Note that if you:

• use any, potentially nondeterministic, algorithm to find the disjoint sets;
• then in a separate phase —parallelised across sets— sort each set and construct the corresponding cycle

then that's almost certainly faster than maintaining sortedness during the set finding (which is what BTreeSet is doing). This is because the algorithm used to find the sets is less constrained, and because it will involve less copying — so fewer, and more local, memory accesses.

[alexander-camuto]

@daira thank you for the detailed review. I've made it such that ordering is not tracked during the cycle creation. We order once at the end in parallel across cycles.

[alexander-camuto]

heya @CPerezz lmk if there's anything else needed here :)

[CPerezz]

This looks really nice! We could file an issue (in case you consider it can be useful) regarding Daira's comment about other algorithms to create cycles.

In any case, the PR looks good to be merged!

Thank you so much for your work!!