[keccak] Refactor lookup table and rho table

[ChihChengLiang]

As mentioned in #489 (comment) we're moving toward 4 lookup gates that each have 3 table columns.

In this PR we did the following stuff:

• Create a ThreeColumnsLookup config, which defines the gate format that takes 3 advice columns and 3 table columns.
• Stackable, Fromb13, Fromb9, and Fromb2 will clone the ThreeColumnsLookup gate. They can share the same advice columns but they have to use different sets of table columns.
• During the refactor of Fromb13, we realized that there is no gate unique to Rho. We can define the whole Rho step as a function instead of config.

[CPerezz]

LGTM. Do we have any numbers in regards any performance upgrades or downgrades?

[ChihChengLiang]

hummm, it has a longer proving time. Not sure why.
(update) Ran again, and the branch was faster this time. So the performance does not change significantly.

this branch

Start:   Setup generation with degree = 20
test keccak_permutation::tests::bench_keccak_round has been running for over 60 seconds
End:     Setup generation with degree = 20 .........................................77.255s
Start:   Keccak Proof generation with 20 degree
End:     Keccak Proof generation with 20 degree ....................................14.902s
Start:   Keccak Proof verification
End:     Keccak Proof verification .................................................4.089ms
test keccak_permutation::tests::bench_keccak_round ... ok

main

Start:   Setup generation with degree = 20
test keccak_permutation::tests::bench_keccak_round has been running for over 60 seconds
End:     Setup generation with degree = 20 .........................................74.217s
Start:   Keccak Proof generation with 20 degree
End:     Keccak Proof generation with 20 degree ....................................15.185s
Start:   Keccak Proof verification
End:     Keccak Proof verification .................................................4.209ms
test keccak_permutation::tests::bench_keccak_round ... ok

[CPerezz]

Coul you edit your local halo2 version and include this modified run fn to print the costs on each magnitude we usually bench?

pub fn run<ConcreteCircuit: Circuit<F>>(
        k: u32,
        circuit: &ConcreteCircuit,
        instance: Vec<Vec<F>>,
    ) -> Result<Self, Error> {
        let n = 1 << k;

        let mut cs = ConstraintSystem::default();
        let config = ConcreteCircuit::configure(&mut cs);
        let cs = cs;

        if n < cs.minimum_rows() {
            return Err(Error::not_enough_rows_available(k));
        }

        if instance.len() != cs.num_instance_columns {
            return Err(Error::InvalidInstances);
        }

        let instance = instance
            .into_iter()
            .map(|mut instance| {
                if instance.len() > n - (cs.blinding_factors() + 1) {
                    return Err(Error::InstanceTooLarge);
                }

                instance.resize(n, F::zero());
                Ok(instance)
            })
            .collect::<Result<Vec<_>, _>>()?;

        // Fixed columns contain no blinding factors.
        let fixed = vec![vec![CellValue::Unassigned; n]; cs.num_fixed_columns];
        let selectors = vec![vec![false; n]; cs.num_selectors];
        // Advice columns contain blinding factors.
        let blinding_factors = cs.blinding_factors();
        let usable_rows = n - (blinding_factors + 1);
        let advice = vec![
            {
                let mut column = vec![CellValue::Unassigned; n];
                // Poison unusable rows.
                for (i, cell) in column.iter_mut().enumerate().skip(usable_rows) {
                    *cell = CellValue::Poison(i);
                }
                column
            };
            cs.num_advice_columns
        ];

        let permutation = permutation::keygen::Assembly::new(n, &cs.permutation);
        let constants = cs.constants.clone();

        let mut prover = MockProver {
            k,
            n: n as u32,
            cs,
            regions: vec![],
            current_region: None,
            fixed,
            advice,
            instance,
            selectors,
            permutation,
            usable_rows: 0..usable_rows,
        };

        ConcreteCircuit::FloorPlanner::synthesize(&mut prover, circuit, config, constants)?;

        let (cs, selector_polys) = prover.cs.compress_selectors(prover.selectors.clone());
        prover.cs = cs.clone();
        prover.fixed.extend(selector_polys.into_iter().map(|poly| {
            let mut v = vec![CellValue::Unassigned; n];
            for (v, p) in v.iter_mut().zip(&poly[..]) {
                *v = CellValue::Assigned(*p);
            }
            v
        }));

        println!("Num advice columns: {}", cs.num_advice_columns);
        println!("Num fixed columns: {}", cs.num_fixed_columns);
        println!("Num instance columns: {}", cs.num_instance_columns);
        println!("Num lookups: {}", prover.cs.lookups.len());
        println!("Num gates: {}", prover.cs.gates.len());
        let mut num_polys = 0;
        let mut max_degree_poly = 0;
        let mut max_degree_poly_details = String::new();

        for gate in prover.cs.gates.iter() {
            num_polys += gate.polynomials().len();
            for (idx, poly) in gate.polynomials().iter().enumerate() {
                if poly.degree() > max_degree_poly {
                    max_degree_poly_details = format!("{}", gate.constraint_name(idx));
                    max_degree_poly = poly.degree();
                }
            }
        }

        let mut max_degree_lookup_input = 0;
        let mut max_degree_lookup_input_descrip = String::new();
        let mut max_degree_lookup = 0;
        let mut max_degree_lookup_descrip = String::new();

        for lookup in prover.cs.lookups.iter() {
            for input in lookup.input_expressions.iter() {
                if input.degree() > max_degree_lookup_input {
                    max_degree_lookup_input = input.degree();
                    max_degree_lookup_input_descrip = format!("lookup input '{}'", lookup.name);
                }
            }

            for table in lookup.table_expressions.iter() {
                if table.degree() > max_degree_lookup {
                    max_degree_lookup = table.degree();
                    max_degree_lookup_descrip = format!("lookup table '{}'", lookup.name);
                }
            }
        }

        println!("num polys: {}", num_polys);
        println!(
            "max_degree_poly: {} - {}",
            max_degree_poly, max_degree_poly_details
        );
        println!(
            "max_degree_lookup_input: {} - {}",
            max_degree_lookup_input, max_degree_lookup_input_descrip
        );
        println!(
            "max_degree_lookup: {} - {}",
            max_degree_lookup, max_degree_lookup_descrip
        );

        Ok(prover)
    }

Then you should just run the tests with -- --nocapture and it should print the details.
On that way you can compare them precisely

[ChihChengLiang]

Here you go

Num advice columns: 29
Num fixed columns: 16
Num instance columns: 0
Num lookups: 4
Num gates: 10
Max advice rows 41915 - Constraint out_state and out_mixing
num polys: 98
max_gate_degree_poly: 6 - 

[CPerezz]

Here you go

Num advice columns: 29
Num fixed columns: 16
Num instance columns: 0
Num lookups: 4
Num gates: 10
Max advice rows 41915 - Constraint out_state and out_mixing
num polys: 98
max_gate_degree_poly: 6 - 

These are the results of this branch? Do you have the ones for main to compare? I can extract them otherwise.

[ChihChengLiang]

The one from main is here. We didn't change columns in this PR.

Num advice columns: 29
Num fixed columns: 16
Num instance columns: 0
Num lookups: 4
Num gates: 10
Max advice rows 42515 - Constraint out_state and out_mixing
num polys: 98
max_gate_degree_poly: 6 - 

[CPerezz]

Ohh that's right @ChihChengLiang .

Have you been able to print the resulting circuit to see if there are any major changes? Or you'd say it's not necessary?

[ChihChengLiang]

It might make more sense to print the circuit layout and compare benchs in #596. This PR changed nothing important on the performance.

[han0110]

LGTM! Nice refactor.

[ChihChengLiang]

@CPerezz @han0110 Thanks for the review. @miha-stopar feel free to give a post-merge review or no review at all.