feat(fflonk): changes for better compression

Author: saitima

crates/fflonk/src/context.rs

@@ -55,7 +55,7 @@ pub(crate) fn init_tmp_mempool() {
     unsafe {
         _TMP_MEMPOOL = Some(bc_mem_pool::new(DEFAULT_DEVICE_ID).unwrap());
     }
-    let num_tmp_bytes = 1_100_000_000;
+    let num_tmp_bytes = 3 << 29; //1.5GB
     let stream = bc_stream::new().unwrap();
     DVec::<u8, PoolAllocator>::allocate_on(num_tmp_bytes, _tmp_mempool(), stream);
 }
@@ -146,9 +146,19 @@ const POWERS_OF_COSET_OMEGA_COARSE_LOG_COUNT: u32 = 14;
 pub type DeviceContextWithSingleDevice = DeviceContext<1>;
 
 impl<const N: usize> DeviceContext<N> {
+    pub fn init_from_preloaded_crs(
+        domain_size: usize,
+        crs: Crs<CompactBn256, CrsForMonomialForm>,
+    ) -> CudaResult<Self> {
+        let context = Self::init_no_msm(domain_size)?;
+        Self::init_msm_on_static_memory(domain_size, Some(crs))?;
+
+        Ok(context)
+    }
+
     pub fn init(domain_size: usize) -> CudaResult<Self> {
         let context = Self::init_no_msm(domain_size)?;
-        Self::init_msm_on_static_memory(domain_size)?;
+        Self::init_msm_on_static_memory(domain_size, None)?;
         // Self::init_msm_on_pool(domain_size)?;
 
         Ok(context)
@@ -179,22 +189,26 @@ impl<const N: usize> DeviceContext<N> {
         Ok(DeviceContext)
     }
 
-    fn init_msm_on_static_memory(domain_size: usize) -> CudaResult<()> {
-        Self::inner_init_msm(domain_size, None, None)?;
+    fn init_msm_on_static_memory(
+        domain_size: usize,
+        crs: Option<Crs<CompactBn256, CrsForMonomialForm>>,
+    ) -> CudaResult<()> {
+        Self::inner_init_msm(domain_size, crs, None, None)?;
         Ok(())
     }
 
     // In reality we keep bases on a statically allocated buffer.
     unsafe fn init_msm_on_pool(domain_size: usize) -> CudaResult<()> {
         let pool = _msm_bases_mempool();
         let stream = bc_stream::new().unwrap();
-        Self::inner_init_msm(domain_size, Some(pool), Some(stream))?;
+        Self::inner_init_msm(domain_size, None, Some(pool), Some(stream))?;
         stream.sync().unwrap();
         Ok(())
     }
 
     fn inner_init_msm(
         domain_size: usize,
+        crs: Option<Crs<CompactBn256, CrsForMonomialForm>>,
         pool: Option<bc_mem_pool>,
         stream: Option<bc_stream>,
     ) -> CudaResult<()> {
@@ -205,7 +219,10 @@ impl<const N: usize> DeviceContext<N> {
         init_msm_result_mempool();
         // MSM impl requires bases to be located in a buffer that is
         // multiple of the domain_size
-        let crs = init_compact_crs(&bellman::worker::Worker::new(), domain_size);
+        let crs = match crs {
+            Some(preloaded_crs) => preloaded_crs,
+            None => init_compact_crs(&bellman::worker::Worker::new(), domain_size),
+        };
         let num_bases = MAX_COMBINED_DEGREE_FACTOR * domain_size;
         assert!(crs.g1_bases.len() >= num_bases);
         let bases = match (pool, stream) {

crates/fflonk/src/convenience.rs

@@ -14,8 +14,8 @@ use circuit_definitions::circuit_definitions::aux_layer::{
 };
 use fflonk::{FflonkAssembly, L1_VERIFIER_DOMAIN_SIZE_LOG};
 
-pub type FflonkSnarkVerifierCircuitDeviceSetup =
-    FflonkDeviceSetup<Bn256, FflonkSnarkVerifierCircuit>;
+pub type FflonkSnarkVerifierCircuitDeviceSetup<A> =
+    FflonkDeviceSetup<Bn256, FflonkSnarkVerifierCircuit, A>;
 
 use super::*;
 

crates/fflonk/src/lib.rs

@@ -63,7 +63,7 @@ pub use gpu_ffi;
 use gpu_ffi::{bc_event, bc_mem_pool, bc_stream};
 use std::alloc::Allocator;
 
-pub use context::{DeviceContext, DeviceContextWithSingleDevice};
+pub use context::{init_compact_crs, DeviceContext, DeviceContextWithSingleDevice};
 pub use fflonk::MAX_COMBINED_DEGREE_FACTOR;
 
 pub use convenience::FflonkSnarkVerifierCircuitDeviceSetup;
@@ -74,3 +74,5 @@ pub type FflonkSnarkVerifierCircuit = ZkSyncSnarkWrapperCircuitNoLookupCustomGat
 pub type FflonkSnarkVerifierCircuitVK = FflonkVerificationKey<Bn256, FflonkSnarkVerifierCircuit>;
 pub type FflonkSnarkVerifierCircuitProof = FflonkProof<Bn256, FflonkSnarkVerifierCircuit>;
 pub type FflonkSnarkVerifierCircuitSetup = FflonkSetup<Bn256, FflonkSnarkVerifierCircuit>;
+
+pub use allocator::{GlobalHost, HostAllocator};

crates/fflonk/src/setup.rs

@@ -307,3 +307,44 @@ impl<E: Engine, C: Circuit<E>, A: HostAllocator> FflonkDeviceSetup<E, C, A> {
         write_curve_affine(&self.g2_elems[1], writer)
     }
 }
+
+pub fn read_raw_fr_vec<F: PrimeField, R: std::io::Read, A: Allocator + Default>(
+    mut src: R,
+) -> std::io::Result<Vec<F, A>> {
+    use byteorder::{BigEndian, ReadBytesExt};
+    let num_values = src.read_u32::<BigEndian>()? as usize;
+    let mut values = Vec::with_capacity_in(num_values, A::default());
+    unsafe {
+        values.set_len(num_values);
+        let buf = std::slice::from_raw_parts_mut(
+            values.as_mut_ptr() as *mut u8,
+            num_values * std::mem::size_of::<F>(),
+        );
+        // src.read_exact(buf)?;
+        let mut dst = &mut buf[..];
+        std::io::copy(&mut src, &mut dst)?;
+    }
+
+    Ok(values)
+}
+
+pub fn write_raw_fr_slice<F: PrimeField, W: std::io::Write>(
+    src_values: &[F],
+    mut dst: W,
+) -> std::io::Result<()> {
+    use byteorder::{BigEndian, WriteBytesExt};
+    let num_values = src_values.len();
+    assert!(num_values < u32::MAX as usize);
+    dst.write_u32::<BigEndian>(num_values as u32)?;
+    unsafe {
+        let buf = std::slice::from_raw_parts_mut(
+            src_values.as_ptr() as *mut u8,
+            num_values * std::mem::size_of::<F>(),
+        );
+        let mut reader = &buf[..];
+        // dst.write_all(buf)?;
+        std::io::copy(&mut reader, &mut dst)?;
+    }
+
+    Ok(())
+}