Reduce number of lookup accesses

0xPolygonZero · Jul 6, 2023 · cafd703 · cafd703
1 parent abc4850
commit cafd703
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 39 deletions.
diff --git a/plonky2/src/gates/lookup.rs b/plonky2/src/gates/lookup.rs
@@ -170,23 +170,25 @@ impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D> for Loo
         let get_wire = |wire: usize| -> F { witness.get_target(Target::wire(self.row, wire)) };
 
         let input_val = get_wire(LookupGate::wire_ith_looking_inp(self.slot_nb));
-        let output_val = if input_val
-            == F::from_canonical_u16(self.lut[input_val.to_canonical_u64() as usize].0)
-        {
-            F::from_canonical_u16(self.lut[input_val.to_canonical_u64() as usize].1)
+        let (input, output) = self.lut[input_val.to_canonical_u64() as usize];
+        if input_val == F::from_canonical_u16(input) {
+            let output_val = F::from_canonical_u16(output);
+
+            let out_wire = Target::wire(self.row, LookupGate::wire_ith_looking_out(self.slot_nb));
+            out_buffer.set_target(out_wire, output_val);
         } else {
-            let mut cur_idx = 0;
-            while input_val != F::from_canonical_u16(self.lut[cur_idx].0)
-                && cur_idx < self.lut.len()
-            {
-                cur_idx += 1;
+            for (input, output) in self.lut.iter() {
+                if input_val == F::from_canonical_u16(*input) {
+                    let output_val = F::from_canonical_u16(*output);
+
+                    let out_wire =
+                        Target::wire(self.row, LookupGate::wire_ith_looking_out(self.slot_nb));
+                    out_buffer.set_target(out_wire, output_val);
+                    return;
+                }
             }
-            assert!(cur_idx < self.lut.len(), "Incorrect input value provided");
-            F::from_canonical_u16(self.lut[cur_idx].1)
+            panic!("Incorrect input value provided");
         };
-
-        let out_wire = Target::wire(self.row, LookupGate::wire_ith_looking_out(self.slot_nb));
-        out_buffer.set_target(out_wire, output_val);
     }
 
     fn serialize(&self, dst: &mut Vec<u8>, cd: &CommonCircuitData<F, D>) -> IoResult<()> {

diff --git a/plonky2/src/gates/lookup_table.rs b/plonky2/src/gates/lookup_table.rs
@@ -193,14 +193,9 @@ impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D> for Loo
             Target::wire(self.row, LookupTableGate::wire_ith_looked_out(self.slot_nb));
 
         if slot < self.lut.len() {
-            out_buffer.set_target(
-                slot_input_target,
-                F::from_canonical_usize(self.lut[slot].0 as usize),
-            );
-            out_buffer.set_target(
-                slot_output_target,
-                F::from_canonical_usize(self.lut[slot].1.into()),
-            );
+            let (input, output) = self.lut[slot];
+            out_buffer.set_target(slot_input_target, F::from_canonical_usize(input as usize));
+            out_buffer.set_target(slot_output_target, F::from_canonical_usize(output as usize));
         } else {
             // Pad with zeros.
             out_buffer.set_target(slot_input_target, F::ZERO);

diff --git a/plonky2/src/plonk/prover.rs b/plonky2/src/plonk/prover.rs
@@ -71,15 +71,14 @@ pub fn set_lookup_wires<
         let remaining_slots = (num_entries
             - (prover_data.lut_to_lookups[lut_index].len() % num_entries))
             % num_entries;
-        let first_inp_value = F::from_canonical_u16(common_data.luts[lut_index][0].0);
-        let first_out_value = F::from_canonical_u16(common_data.luts[lut_index][0].1);
+        let (first_inp_value, first_out_value) = common_data.luts[lut_index][0];
         for slot in (num_entries - remaining_slots)..num_entries {
             let inp_target =
                 Target::wire(last_lut_gate - 1, LookupGate::wire_ith_looking_inp(slot));
             let out_target =
                 Target::wire(last_lut_gate - 1, LookupGate::wire_ith_looking_out(slot));
-            pw.set_target(inp_target, first_inp_value);
-            pw.set_target(out_target, first_out_value);
+            pw.set_target(inp_target, F::from_canonical_u16(first_inp_value));
+            pw.set_target(out_target, F::from_canonical_u16(first_out_value));
 
             multiplicities[0] += 1;
         }

diff --git a/plonky2/src/plonk/vanishing_poly.rs b/plonky2/src/plonk/vanishing_poly.rs
@@ -37,11 +37,8 @@ pub(crate) fn get_lut_poly<F: RichField + Extendable<D>, const D: usize>(
     let b = deltas[LookupChallenges::ChallengeB as usize];
     let mut coeffs = Vec::new();
     let n = common_data.luts[lut_index].len();
-    for i in 0..n {
-        coeffs.push(
-            F::from_canonical_u16(common_data.luts[lut_index][i].0)
-                + b * F::from_canonical_u16(common_data.luts[lut_index][i].1),
-        );
+    for (input, output) in common_data.luts[lut_index].iter() {
+        coeffs.push(F::from_canonical_u16(*input) + b * F::from_canonical_u16(*output));
     }
     coeffs.append(&mut vec![F::ZERO; degree - n]);
     coeffs.reverse();
@@ -767,14 +764,11 @@ pub(crate) fn get_lut_poly_circuit<F: RichField + Extendable<D>, const D: usize>
     let b = deltas[LookupChallenges::ChallengeB as usize];
     let delta = deltas[LookupChallenges::ChallengeDelta as usize];
     let n = common_data.luts[lut_index].len();
-    let mut coeffs: Vec<Target> = (0..n)
-        .map(|i| {
-            let temp =
-                builder.mul_const(F::from_canonical_u16(common_data.luts[lut_index][i].1), b);
-            builder.add_const(
-                temp,
-                F::from_canonical_u16(common_data.luts[lut_index][i].0),
-            )
+    let mut coeffs: Vec<Target> = common_data.luts[lut_index]
+        .iter()
+        .map(|(input, output)| {
+            let temp = builder.mul_const(F::from_canonical_u16(*output), b);
+            builder.add_const(temp, F::from_canonical_u16(*input))
         })
         .collect();
     for _ in n..degree {