Merge pull request #964 from topos-protocol/lookup

Implement lookups with logarithmic derivatives in Plonk

plonky2/examples/bench_recursion.rs

@@ -5,12 +5,16 @@
 
 #![allow(clippy::upper_case_acronyms)]
 
+extern crate alloc;
+use alloc::sync::Arc;
 use core::num::ParseIntError;
 use core::ops::RangeInclusive;
 use core::str::FromStr;
 
 use anyhow::{anyhow, Context as _, Result};
+use itertools::Itertools;
 use log::{info, Level, LevelFilter};
+use plonky2::gadgets::lookup::TIP5_TABLE;
 use plonky2::gates::noop::NoopGate;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::witness::{PartialWitness, WitnessWrite};
@@ -59,6 +63,12 @@ struct Options {
     /// range.
     #[structopt(long, default_value="14", parse(try_from_str = parse_range_usize))]
     size: RangeInclusive<usize>,
+
+    /// Lookup type. If `lookup_type == 0` or `lookup_type > 2`, then a benchmark with NoopGates only is run.
+    /// If `lookup_type == 1`, a benchmark with one lookup is run.
+    /// If `lookup_type == 2`, a benchmark with 515 lookups is run.
+    #[structopt(long, default_value="0", parse(try_from_str = parse_hex_u64))]
+    lookup_type: u64,
 }
 
 /// Creates a dummy proof which should have `2 ** log2_size` rows.
@@ -91,6 +101,101 @@ fn dummy_proof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D
     Ok((proof, data.verifier_only, data.common))
 }
 
+fn dummy_lookup_proof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
+    config: &CircuitConfig,
+    log2_size: usize,
+) -> Result<ProofTuple<F, C, D>> {
+    let mut builder = CircuitBuilder::<F, D>::new(config.clone());
+    let tip5_table = TIP5_TABLE.to_vec();
+    let inps = 0..256;
+    let table = Arc::new(inps.zip_eq(tip5_table).collect());
+    let tip5_idx = builder.add_lookup_table_from_pairs(table);
+    let initial_a = builder.add_virtual_target();
+    builder.add_lookup_from_index(initial_a, tip5_idx);
+    builder.register_public_input(initial_a);
+
+    // 'size' is in degree, but we want the number of gates in the circuit.
+    // A non-zero amount of padding will be added and size will be rounded to the next power of two.
+    // To hit our target size, we go just under the previous power of two and hope padding is less than half the proof.
+    let targeted_num_gates = match log2_size {
+        0 => return Err(anyhow!("size must be at least 1")),
+        1 => 0,
+        2 => 1,
+        n => (1 << (n - 1)) + 1,
+    };
+    assert!(
+        targeted_num_gates >= builder.num_gates(),
+        "size is too small to support lookups"
+    );
+
+    for _ in builder.num_gates()..targeted_num_gates {
+        builder.add_gate(NoopGate, vec![]);
+    }
+    builder.print_gate_counts(0);
+
+    let data = builder.build::<C>();
+    let mut inputs = PartialWitness::<F>::new();
+    inputs.set_target(initial_a, F::ONE);
+    let mut timing = TimingTree::new("prove with one lookup", Level::Debug);
+    let proof = prove(&data.prover_only, &data.common, inputs, &mut timing)?;
+    timing.print();
+    data.verify(proof.clone())?;
+
+    Ok((proof, data.verifier_only, data.common))
+}
+
+/// Creates a dummy proof which has more than 256 lookups to one LUT
+fn dummy_many_rows_proof<
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    const D: usize,
+>(
+    config: &CircuitConfig,
+    log2_size: usize,
+) -> Result<ProofTuple<F, C, D>> {
+    let mut builder = CircuitBuilder::<F, D>::new(config.clone());
+    let tip5_table = TIP5_TABLE.to_vec();
+    let inps: Vec<u16> = (0..256).collect();
+    let tip5_idx = builder.add_lookup_table_from_table(&inps, &tip5_table);
+    let initial_a = builder.add_virtual_target();
+
+    let output = builder.add_lookup_from_index(initial_a, tip5_idx);
+    for _ in 0..514 {
+        builder.add_lookup_from_index(output, 0);
+    }
+
+    // 'size' is in degree, but we want the number of gates in the circuit.
+    // A non-zero amount of padding will be added and size will be rounded to the next power of two.
+    // To hit our target size, we go just under the previous power of two and hope padding is less than half the proof.
+    let targeted_num_gates = match log2_size {
+        0 => return Err(anyhow!("size must be at least 1")),
+        1 => 0,
+        2 => 1,
+        n => (1 << (n - 1)) + 1,
+    };
+    assert!(
+        targeted_num_gates >= builder.num_gates(),
+        "size is too small to support so many lookups"
+    );
+
+    for _ in 0..targeted_num_gates {
+        builder.add_gate(NoopGate, vec![]);
+    }
+
+    builder.register_public_input(initial_a);
+    builder.register_public_input(output);
+
+    let mut pw = PartialWitness::new();
+    pw.set_target(initial_a, F::ONE);
+    let data = builder.build::<C>();
+    let mut timing = TimingTree::new("prove with many lookups", Level::Debug);
+    let proof = prove(&data.prover_only, &data.common, pw, &mut timing)?;
+    timing.print();
+
+    data.verify(proof.clone())?;
+    Ok((proof, data.verifier_only, data.common))
+}
+
 fn recursive_proof<
     F: RichField + Extendable<D>,
     C: GenericConfig<D, F = F>,
@@ -183,16 +288,34 @@ fn test_serialization<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>,
     Ok(())
 }
 
-fn benchmark(config: &CircuitConfig, log2_inner_size: usize) -> Result<()> {
+pub fn benchmark_function(
+    config: &CircuitConfig,
+    log2_inner_size: usize,
+    lookup_type: u64,
+) -> Result<()> {
     const D: usize = 2;
     type C = PoseidonGoldilocksConfig;
     type F = <C as GenericConfig<D>>::F;
 
+    let dummy_proof_function = match lookup_type {
+        0 => dummy_proof::<F, C, D>,
+        1 => dummy_lookup_proof::<F, C, D>,
+        2 => dummy_many_rows_proof::<F, C, D>,
+        _ => dummy_proof::<F, C, D>,
+    };
+
+    let name = match lookup_type {
+        0 => "proof",
+        1 => "one lookup proof",
+        2 => "multiple lookups proof",
+        _ => "proof",
+    };
     // Start with a dummy proof of specified size
-    let inner = dummy_proof::<F, C, D>(config, log2_inner_size)?;
+    let inner = dummy_proof_function(config, log2_inner_size)?;
     let (_, _, cd) = &inner;
     info!(
-        "Initial proof degree {} = 2^{}",
+        "Initial {} degree {} = 2^{}",
+        name,
         cd.degree(),
         cd.degree_bits()
     );
@@ -201,7 +324,8 @@ fn benchmark(config: &CircuitConfig, log2_inner_size: usize) -> Result<()> {
     let middle = recursive_proof::<F, C, C, D>(&inner, config, None)?;
     let (_, _, cd) = &middle;
     info!(
-        "Single recursion proof degree {} = 2^{}",
+        "Single recursion {} degree {} = 2^{}",
+        name,
         cd.degree(),
         cd.degree_bits()
     );
@@ -210,7 +334,8 @@ fn benchmark(config: &CircuitConfig, log2_inner_size: usize) -> Result<()> {
     let outer = recursive_proof::<F, C, C, D>(&middle, config, None)?;
     let (proof, vd, cd) = &outer;
     info!(
-        "Double recursion proof degree {} = 2^{}",
+        "Double recursion {} degree {} = 2^{}",
+        name,
         cd.degree(),
         cd.degree_bits()
     );
@@ -223,7 +348,6 @@ fn benchmark(config: &CircuitConfig, log2_inner_size: usize) -> Result<()> {
 fn main() -> Result<()> {
     // Parse command line arguments, see `--help` for details.
     let options = Options::from_args_safe()?;
-
     // Initialize logging
     let mut builder = env_logger::Builder::from_default_env();
     builder.parse_filters(&options.log_filter);
@@ -246,8 +370,9 @@ fn main() -> Result<()> {
     let threads = options.threads.unwrap_or(num_cpus..=num_cpus);
 
     let config = CircuitConfig::standard_recursion_config();
+
     for log2_inner_size in options.size {
-        // Since the `size` is most likely to be and unbounded range we make that the outer iterator.
+        // Since the `size` is most likely to be an unbounded range we make that the outer iterator.
         for threads in threads.clone() {
             rayon::ThreadPoolBuilder::new()
                 .num_threads(threads)
@@ -259,8 +384,8 @@ fn main() -> Result<()> {
                         rayon::current_num_threads(),
                         num_cpus
                     );
-                    // Run the benchmark
-                    benchmark(&config, log2_inner_size)
+                    // Run the benchmark. `options.lookup_type` determines which benchmark to run.
+                    benchmark_function(&config, log2_inner_size, options.lookup_type)
                 })?;
         }
     }

plonky2/src/gadgets/lookup.rs

@@ -0,0 +1,126 @@
+use crate::field::extension::Extendable;
+use crate::gates::lookup::LookupGate;
+use crate::gates::lookup_table::{LookupTable, LookupTableGate};
+use crate::gates::noop::NoopGate;
+use crate::hash::hash_types::RichField;
+use crate::iop::target::Target;
+use crate::plonk::circuit_builder::CircuitBuilder;
+
+/// Lookup tables used in the tests and benchmarks.
+///
+/// The following table was taken from the Tip5 paper.
+pub const TIP5_TABLE: [u16; 256] = [
+    0, 7, 26, 63, 124, 215, 85, 254, 214, 228, 45, 185, 140, 173, 33, 240, 29, 177, 176, 32, 8,
+    110, 87, 202, 204, 99, 150, 106, 230, 14, 235, 128, 213, 239, 212, 138, 23, 130, 208, 6, 44,
+    71, 93, 116, 146, 189, 251, 81, 199, 97, 38, 28, 73, 179, 95, 84, 152, 48, 35, 119, 49, 88,
+    242, 3, 148, 169, 72, 120, 62, 161, 166, 83, 175, 191, 137, 19, 100, 129, 112, 55, 221, 102,
+    218, 61, 151, 237, 68, 164, 17, 147, 46, 234, 203, 216, 22, 141, 65, 57, 123, 12, 244, 54, 219,
+    231, 96, 77, 180, 154, 5, 253, 133, 165, 98, 195, 205, 134, 245, 30, 9, 188, 59, 142, 186, 197,
+    181, 144, 92, 31, 224, 163, 111, 74, 58, 69, 113, 196, 67, 246, 225, 10, 121, 50, 60, 157, 90,
+    122, 2, 250, 101, 75, 178, 159, 24, 36, 201, 11, 243, 132, 198, 190, 114, 233, 39, 52, 21, 209,
+    108, 238, 91, 187, 18, 104, 194, 37, 153, 34, 200, 143, 126, 155, 236, 118, 64, 80, 172, 89,
+    94, 193, 135, 183, 86, 107, 252, 13, 167, 206, 136, 220, 207, 103, 171, 160, 76, 182, 227, 217,
+    158, 56, 174, 4, 66, 109, 139, 162, 184, 211, 249, 47, 125, 232, 117, 43, 16, 42, 127, 20, 241,
+    25, 149, 105, 156, 51, 53, 168, 145, 247, 223, 79, 78, 226, 15, 222, 82, 115, 70, 210, 27, 41,
+    1, 170, 40, 131, 192, 229, 248, 255,
+];
+
+/// This is a table with 256 arbitrary values.
+pub const OTHER_TABLE: [u16; 256] = [
+    2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7,
+    0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10,
+    19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216,
+    247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57,
+    126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3,
+    9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35,
+    10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45,
+    216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39,
+    57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25,
+    3, 9, 7, 0, 3, 25, 35, 10, 19, 36, 45, 216, 247, 35, 39, 57, 126, 2, 6, 25, 3, 9, 7, 0, 3, 25,
+    35, 10, 19, 36, 45, 216, 247,
+];
+
+/// This is a smaller lookup table with arbitrary values.
+pub const SMALLER_TABLE: [u16; 8] = [2, 24, 56, 100, 128, 16, 20, 49];
+
+impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
+    /// Adds a lookup table to the list of stored lookup tables `self.luts` based on a table of (input, output) pairs. It returns the index of the LUT within `self.luts`.
+    pub fn add_lookup_table_from_pairs(&mut self, table: LookupTable) -> usize {
+        self.update_luts_from_pairs(table)
+    }
+
+    /// Adds a lookup table to the list of stored lookup tables `self.luts` based on a table, represented as a slice `&[u16]` of inputs and a slice `&[u16]` of outputs. It returns the index of the LUT within `self.luts`.
+    pub fn add_lookup_table_from_table(&mut self, inps: &[u16], outs: &[u16]) -> usize {
+        self.update_luts_from_table(inps, outs)
+    }
+
+    /// Adds a lookup table to the list of stored lookup tables `self.luts` based on a function. It returns the index of the LUT within `self.luts`.
+    pub fn add_lookup_table_from_fn(&mut self, f: fn(u16) -> u16, inputs: &[u16]) -> usize {
+        self.update_luts_from_fn(f, inputs)
+    }
+
+    /// Adds a lookup (input, output) pair to the stored lookups. Takes a `Target` input and returns a `Target` output.
+    pub fn add_lookup_from_index(&mut self, looking_in: Target, lut_index: usize) -> Target {
+        assert!(
+            lut_index < self.get_luts_length(),
+            "lut number {} not in luts (length = {})",
+            lut_index,
+            self.get_luts_length()
+        );
+        let looking_out = self.add_virtual_target();
+        self.update_lookups(looking_in, looking_out, lut_index);
+        looking_out
+    }
+
+    /// We call this function at the end of circuit building right before the PI gate to add all `LookupTableGate` and `LookupGate`.
+    /// It also updates `self.lookup_rows` accordingly.
+    pub fn add_all_lookups(&mut self) {
+        for lut_index in 0..self.num_luts() {
+            assert!(
+                !self.get_lut_lookups(lut_index).is_empty() || lut_index >= self.get_luts_length(),
+                "LUT number {:?} is unused",
+                lut_index
+            );
+            if !self.get_lut_lookups(lut_index).is_empty() {
+                // Create LU gates. Connect them to the stored lookups.
+                let last_lu_gate = self.num_gates();
+
+                let lut = self.get_lut(lut_index);
+
+                let lookups = self.get_lut_lookups(lut_index).to_owned();
+
+                for (looking_in, looking_out) in lookups {
+                    let gate = LookupGate::new_from_table(&self.config, lut.clone());
+                    let (gate, i) =
+                        self.find_slot(gate, &[F::from_canonical_usize(lut_index)], &[]);
+                    let gate_in = Target::wire(gate, LookupGate::wire_ith_looking_inp(i));
+                    let gate_out = Target::wire(gate, LookupGate::wire_ith_looking_out(i));
+                    self.connect(gate_in, looking_in);
+                    self.connect(gate_out, looking_out);
+                }
+
+                // Create LUT gates. Nothing is connected to them.
+                let last_lut_gate = self.num_gates();
+                let num_lut_entries = LookupTableGate::num_slots(&self.config);
+                let num_lut_rows = (self.get_luts_idx_length(lut_index) - 1) / num_lut_entries + 1;
+                let num_lut_cells = num_lut_entries * num_lut_rows;
+                for _ in 0..num_lut_cells {
+                    let gate =
+                        LookupTableGate::new_from_table(&self.config, lut.clone(), last_lut_gate);
+                    self.find_slot(gate, &[], &[]);
+                }
+
+                let first_lut_gate = self.num_gates() - 1;
+
+                // Will ensure the next row's wires will be all zeros. With this, there is no distinction between the transition constraints on the first row
+                // and on the other rows. Additionally, initial constraints become a simple zero check.
+                self.add_gate(NoopGate, vec![]);
+
+                // These elements are increasing: the gate rows are deliberately upside down.
+                // This is necessary for constraint evaluation so that you do not need values of the next
+                // row's wires, which aren't provided in the evaluation variables.
+                self.add_lookup_rows(last_lu_gate, last_lut_gate, first_lut_gate);
+            }
+        }
+    }
+}

plonky2/src/gadgets/mod.rs

@@ -2,6 +2,7 @@ pub mod arithmetic;
 pub mod arithmetic_extension;
 pub mod hash;
 pub mod interpolation;
+pub mod lookup;
 pub mod polynomial;
 pub mod random_access;
 pub mod range_check;