feat: programmable buffer capacity (#712)

src/circuit/modules/kzg.rs

@@ -41,7 +41,7 @@ pub struct KZGChip {
 }
 
 impl KZGChip {
-    /// Returns the number of inputs to the hash function
+    /// Commit to the message using the KZG commitment scheme
     pub fn commit(
         message: Vec<Fp>,
         degree: u32,

src/graph/mod.rs

@@ -35,6 +35,8 @@ use halo2_proofs::{
 };
 use halo2curves::bn256::{self, Bn256, Fr as Fp, G1Affine};
 use halo2curves::ff::PrimeField;
+#[cfg(not(target_arch = "wasm32"))]
+use lazy_static::lazy_static;
 use log::{debug, error, info, trace, warn};
 use maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 pub use model::*;
@@ -61,6 +63,20 @@ pub const RANGE_MULTIPLIER: i128 = 2;
 /// Max representation of a lookup table input
 pub const MAX_LOOKUP_ABS: i128 = 8 * 2_i128.pow(MAX_PUBLIC_SRS);
 
+#[cfg(not(target_arch = "wasm32"))]
+lazy_static! {
+    /// Max circuit area
+    pub static ref EZKL_MAX_CIRCUIT_AREA: Option<usize> =
+        if let Ok(max_circuit_area) = std::env::var("EZKL_MAX_CIRCUIT_AREA") {
+            Some(max_circuit_area.parse().unwrap_or(0))
+        } else {
+            None
+        };
+}
+
+#[cfg(target_arch = "wasm32")]
+const EZKL_MAX_CIRCUIT_AREA: Option<usize> = None;
+
 /// circuit related errors.
 #[derive(Debug, Error)]
 pub enum GraphError {
@@ -530,6 +546,7 @@ impl GraphSettings {
 pub struct GraphConfig {
     model_config: ModelConfig,
     module_configs: ModuleConfigs,
+    circuit_size: CircuitSize,
 }
 
 /// Defines the circuit for a computational graph / model loaded from a `.onnx` file.
@@ -1366,7 +1383,6 @@ impl GraphCircuit {
     }
 }
 
-#[cfg(not(target_arch = "wasm32"))]
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
 struct CircuitSize {
     num_instances: usize,
@@ -1374,20 +1390,22 @@ struct CircuitSize {
     num_fixed: usize,
     num_challenges: usize,
     num_selectors: usize,
+    logrows: u32,
 }
 
-#[cfg(not(target_arch = "wasm32"))]
 impl CircuitSize {
-    pub fn from_cs(cs: &ConstraintSystem<Fp>) -> Self {
+    pub fn from_cs(cs: &ConstraintSystem<Fp>, logrows: u32) -> Self {
         CircuitSize {
             num_instances: cs.num_instance_columns(),
             num_advice_columns: cs.num_advice_columns(),
             num_fixed: cs.num_fixed_columns(),
             num_challenges: cs.num_challenges(),
             num_selectors: cs.num_selectors(),
+            logrows,
         }
     }
 
+    #[cfg(not(target_arch = "wasm32"))]
     /// Export the ezkl configuration as json
     pub fn as_json(&self) -> Result<String, Box<dyn std::error::Error>> {
         let serialized = match serde_json::to_string(&self) {
@@ -1398,6 +1416,25 @@ impl CircuitSize {
         };
         Ok(serialized)
     }
+
+    /// number of columns
+    pub fn num_columns(&self) -> usize {
+        self.num_instances + self.num_advice_columns + self.num_fixed
+    }
+
+    /// area of the circuit
+    pub fn area(&self) -> usize {
+        self.num_columns() * (1 << self.logrows)
+    }
+
+    /// area less than max
+    pub fn area_less_than_max(&self) -> bool {
+        if EZKL_MAX_CIRCUIT_AREA.is_some() {
+            self.area() < EZKL_MAX_CIRCUIT_AREA.unwrap()
+        } else {
+            true
+        }
+    }
 }
 
 impl Circuit<Fp> for GraphCircuit {
@@ -1472,10 +1509,12 @@ impl Circuit<Fp> for GraphCircuit {
             (cs.degree() as f32).log2().ceil()
         );
 
+        let circuit_size = CircuitSize::from_cs(cs, params.run_args.logrows);
+
         #[cfg(not(target_arch = "wasm32"))]
         info!(
             "circuit size: \n {}",
-            CircuitSize::from_cs(cs)
+            circuit_size
                 .as_json()
                 .unwrap()
                 .to_colored_json_auto()
@@ -1485,6 +1524,7 @@ impl Circuit<Fp> for GraphCircuit {
         GraphConfig {
             model_config,
             module_configs,
+            circuit_size,
         }
     }
 
@@ -1497,6 +1537,16 @@ impl Circuit<Fp> for GraphCircuit {
         config: Self::Config,
         mut layouter: impl Layouter<Fp>,
     ) -> Result<(), PlonkError> {
+        // check if the circuit area is less than the max
+        if !config.circuit_size.area_less_than_max() {
+            error!(
+                "circuit area {} is larger than the max allowed area {}",
+                config.circuit_size.area(),
+                EZKL_MAX_CIRCUIT_AREA.unwrap()
+            );
+            return Err(PlonkError::Synthesis);
+        }
+
         trace!("Setting input in synthesize");
         let input_vis = &self.settings().run_args.input_visibility;
         let output_vis = &self.settings().run_args.output_visibility;

src/pfsys/mod.rs

@@ -40,8 +40,24 @@ use std::ops::Deref;
 use std::path::PathBuf;
 use thiserror::Error as thisError;
 
+// not wasm
+#[cfg(not(target_arch = "wasm32"))]
+use lazy_static::lazy_static;
+
 use halo2curves::bn256::{Bn256, Fr, G1Affine};
 
+#[cfg(not(target_arch = "wasm32"))]
+// Buf writer capacity
+lazy_static! {
+    static ref EZKL_BUF_CAPACITY: usize = std::env::var("EZKL_BUF_CAPACITY")
+        .unwrap_or("8000".to_string())
+        .parse()
+        .unwrap();
+}
+
+#[cfg(target_arch = "wasm32")]
+const EZKL_BUF_CAPACITY: &usize = &8000;
+
 #[allow(missing_docs)]
 #[derive(
     ValueEnum, Copy, Clone, Default, Debug, PartialEq, Eq, Deserialize, Serialize, PartialOrd,
@@ -315,7 +331,7 @@ where
     /// Saves the Proof to a specified `proof_path`.
     pub fn save(&self, proof_path: &PathBuf) -> Result<(), Box<dyn Error>> {
         let file = std::fs::File::create(proof_path)?;
-        let mut writer = BufWriter::new(file);
+        let mut writer = BufWriter::with_capacity(*EZKL_BUF_CAPACITY, file);
         serde_json::to_writer(&mut writer, &self)?;
         Ok(())
     }
@@ -678,7 +694,7 @@ where
     info!("loading verification key from {:?}", path);
     let f =
         File::open(path.clone()).map_err(|_| format!("failed to load vk at {}", path.display()))?;
-    let mut reader = BufReader::new(f);
+    let mut reader = BufReader::with_capacity(*EZKL_BUF_CAPACITY, f);
     VerifyingKey::<Scheme::Curve>::read::<_, C>(
         &mut reader,
         halo2_proofs::SerdeFormat::RawBytes,
@@ -700,7 +716,7 @@ where
     info!("loading proving key from {:?}", path);
     let f =
         File::open(path.clone()).map_err(|_| format!("failed to load pk at {}", path.display()))?;
-    let mut reader = BufReader::new(f);
+    let mut reader = BufReader::with_capacity(*EZKL_BUF_CAPACITY, f);
     ProvingKey::<Scheme::Curve>::read::<_, C>(
         &mut reader,
         halo2_proofs::SerdeFormat::RawBytes,
@@ -720,7 +736,7 @@ where
 {
     info!("saving proving key 💾");
     let f = File::create(path)?;
-    let mut writer = BufWriter::new(f);
+    let mut writer = BufWriter::with_capacity(*EZKL_BUF_CAPACITY, f);
     vk.write(&mut writer, halo2_proofs::SerdeFormat::RawBytes)?;
     writer.flush()?;
     Ok(())
@@ -737,7 +753,7 @@ where
 {
     info!("saving verification key 💾");
     let f = File::create(path)?;
-    let mut writer = BufWriter::new(f);
+    let mut writer = BufWriter::with_capacity(*EZKL_BUF_CAPACITY, f);
     vk.write(&mut writer, halo2_proofs::SerdeFormat::RawBytes)?;
     writer.flush()?;
     Ok(())
@@ -750,7 +766,7 @@ pub fn save_params<Scheme: CommitmentScheme>(
 ) -> Result<(), io::Error> {
     info!("saving parameters 💾");
     let f = File::create(path)?;
-    let mut writer = BufWriter::new(f);
+    let mut writer = BufWriter::with_capacity(*EZKL_BUF_CAPACITY, f);
     params.write(&mut writer)?;
     writer.flush()?;
     Ok(())