Performance Optimizations (#26)

* Replace is bite / bit gates / lookups with range checks

* Fix tests

* Use V1 floor planner

* Test EVM verifier with MNIST-Medium

* Fix clippy errors

* Switch back to SimpleFloorPlanner

* Update k

* Remove second constants column

* Update docstring

examples/evm-verifier.rs

@@ -29,10 +29,7 @@ use snark_verifier::{
     verifier::{self, SnarkVerifier},
 };
 use std::{env, fs, path::Path, rc::Rc};
-use zero_g::{
-    checked_in_test_data::{MNIST_TINY, TEST_IMG_PATH},
-    load_grayscale_image, load_wnn,
-};
+use zero_g::{checked_in_test_data::*, load_grayscale_image, load_wnn};
 
 type PlonkVerifier = verifier::plonk::PlonkVerifier<KzgAs<Bn256, Gwc19>>;
 
@@ -271,18 +268,23 @@ fn evm_verify(deployment_code: Vec<u8>, instances: Vec<Vec<Fr>>, proof: Vec<u8>)
 }
 
 fn validate_evm<C: Circuit<Fr> + Clone>(circuit: C, instances: Vec<Vec<Fr>>, k: u32, name: &str) {
+    println!("Generating Params...");
     let params = gen_srs(k);
+    println!("Generating PK...");
     let pk = gen_pk(&params, &circuit);
+    println!("Generating deployment code...");
     let deployment_code = gen_evm_verifier(&params, pk.get_vk(), vec![instances[0].len()], name);
 
+    println!("Generating proof...");
     let proof = gen_proof(&params, &pk, circuit.clone(), instances.clone());
+    println!("Verifying proof...");
     evm_verify(deployment_code, instances.clone(), proof);
 }
 
 fn main() {
     let args: Vec<String> = env::args().collect();
     if args[1] == "wnn" {
-        let (k, model_path) = MNIST_TINY;
+        let (k, model_path) = MNIST_MEDIUM;
 
         let wnn = load_wnn(Path::new(model_path)).unwrap();
         let image = load_grayscale_image(Path::new(TEST_IMG_PATH)).unwrap();

models/readme.md

@@ -8,7 +8,7 @@ To update the models, follow the steps in the `BTHOWeN-0g` readme, then copy the
 cp ../BTHOWeN/software_model/models/MNIST/*.hdf5 models/
 ```
 
-### "MNIST-Tiny": `model_28input_256entry_1hash_1bpi` (k = 12)
+### "MNIST-Tiny": `model_28input_256entry_1hash_1bpi` (k = 14)
 
 A very small toy model, used in the tests and the benchmark.
 Accuracy on the MNIST test set is 83.06%.

src/gadgets/encode_image.rs

@@ -2,28 +2,32 @@ use std::collections::BTreeMap;
 
 use ff::PrimeFieldBits;
 use halo2_proofs::{
-    circuit::{AssignedCell, Layouter},
-    plonk::{Advice, Column, ConstraintSystem, Error, TableColumn},
+    circuit::{AssignedCell, Layouter, Value},
+    plonk::{Advice, Column, ConstraintSystem, Error},
 };
+use itertools::Itertools;
 use ndarray::{Array2, Array3};
 
 use crate::gadgets::greater_than::GreaterThanWitnessResult;
 
-use super::greater_than::{GreaterThanChip, GreaterThanChipConfig, GreaterThanInstructions};
+use super::{
+    greater_than::{GreaterThanChip, GreaterThanChipConfig, GreaterThanInstructions},
+    range_check::RangeCheckConfig,
+};
 
 pub trait EncodeImageInstructions<F: PrimeFieldBits> {
     /// Maps an image to a bit string.
     fn encode_image(
         &self,
         layouter: impl Layouter<F>,
-        image: &Array2<u8>,
+        image: Value<Array2<u8>>,
     ) -> Result<Vec<AssignedCell<F, F>>, Error>;
 }
 
 #[derive(Clone, Debug)]
-pub struct EncodeImageChipConfig {
+pub struct EncodeImageChipConfig<F: PrimeFieldBits> {
     advice_column: Column<Advice>,
-    greater_than_chip_config: GreaterThanChipConfig,
+    greater_than_chip_config: GreaterThanChipConfig<F>,
 }
 
 /// Encodes an image into a bit string, as follows:
@@ -34,12 +38,15 @@ pub struct EncodeImageChipConfig {
 /// - Intensities belonging to the same pixel are constrained to be equal.
 pub struct EncodeImageChip<F: PrimeFieldBits> {
     greater_than_chip: GreaterThanChip<F>,
-    config: EncodeImageChipConfig,
+    config: EncodeImageChipConfig<F>,
     binarization_thresholds: Array3<u16>,
 }
 
 impl<F: PrimeFieldBits> EncodeImageChip<F> {
-    pub fn construct(config: EncodeImageChipConfig, binarization_thresholds: Array3<u16>) -> Self {
+    pub fn construct(
+        config: EncodeImageChipConfig<F>,
+        binarization_thresholds: Array3<u16>,
+    ) -> Self {
         let greater_than_chip = GreaterThanChip::construct(config.greater_than_chip_config.clone());
         Self {
             greater_than_chip,
@@ -54,10 +61,10 @@ impl<F: PrimeFieldBits> EncodeImageChip<F> {
         y: Column<Advice>,
         diff: Column<Advice>,
         is_gt: Column<Advice>,
-        byte_column: TableColumn,
-    ) -> EncodeImageChipConfig {
+        range_check_config: RangeCheckConfig<F>,
+    ) -> EncodeImageChipConfig<F> {
         let greater_than_chip_config =
-            GreaterThanChip::configure(meta, x, y, diff, is_gt, byte_column);
+            GreaterThanChip::configure(meta, x, y, diff, is_gt, range_check_config);
         EncodeImageChipConfig {
             advice_column: is_gt,
             greater_than_chip_config,
@@ -69,9 +76,15 @@ impl<F: PrimeFieldBits> EncodeImageInstructions<F> for EncodeImageChip<F> {
     fn encode_image(
         &self,
         mut layouter: impl Layouter<F>,
-        image: &Array2<u8>,
+        image: Value<Array2<u8>>,
     ) -> Result<Vec<AssignedCell<F, F>>, Error> {
-        let (width, height) = (image.shape()[0], image.shape()[1]);
+        let width = self.binarization_thresholds.shape()[0];
+        let height = self.binarization_thresholds.shape()[1];
+
+        // Turn Value<Array2<u8>> into Vec<Value<u8>>
+        let image_flat = image
+            .map(|image| image.into_iter().collect_vec())
+            .transpose_vec(width * height);
 
         let mut intensity_cells: BTreeMap<(usize, usize), AssignedCell<F, F>> = BTreeMap::new();
         let mut bit_cells = vec![];
@@ -106,12 +119,14 @@ impl<F: PrimeFieldBits> EncodeImageInstructions<F> for EncodeImageChip<F> {
 
                         match intensity_cells.get(&(i, j)) {
                             None => {
+                                let image_value =
+                                    image_flat[i * height + j].map(|x| F::from(x as u64));
                                 // For the first cell, we want to remember the intensity cell, so that we can
                                 // add a copy constraint for the other thresholds.
                                 let GreaterThanWitnessResult { x_cell, gt_cell } =
                                     self.greater_than_chip.greater_than_witness(
-                                        &mut layouter,
-                                        F::from(image[(i, j)] as u64),
+                                        layouter.namespace(|| format!("gt[{}, {}]", i, j)),
+                                        image_value,
                                         t,
                                     )?;
                                 intensity_cells.insert((i, j), x_cell);
@@ -120,7 +135,7 @@ impl<F: PrimeFieldBits> EncodeImageInstructions<F> for EncodeImageChip<F> {
                             Some(first_cell) => {
                                 // For the other cells, we want to add a copy constraint to the first cell.
                                 self.greater_than_chip.greater_than_copy(
-                                    &mut layouter,
+                                    layouter.namespace(|| format!("gt[{}, {}]", i, j)),
                                     first_cell,
                                     t,
                                 )?