Simulate memory queues in parallel for BWG

circuit_encodings/src/lib.rs

@@ -316,8 +316,6 @@ impl<
 pub struct FullWidthQueueIntermediateStates<F: SmallField, const SW: usize, const ROUNDS: usize> {
     pub head: [F; SW],
     pub tail: [F; SW],
-    pub old_head: [F; SW],
-    pub old_tail: [F; SW],
     pub num_items: u32,
     pub round_function_execution_pairs: [([F; SW], [F; SW]); ROUNDS],
 }
@@ -433,8 +431,6 @@ impl<
         let intermediate_info = FullWidthQueueIntermediateStates {
             head: self.head,
             tail: new_tail,
-            old_head: self.head,
-            old_tail,
             num_items: self.num_items,
             round_function_execution_pairs: states,
         };
@@ -473,8 +469,6 @@ impl<
         let intermediate_info = FullWidthQueueIntermediateStates {
             head: self.head,
             tail: self.tail,
-            old_head,
-            old_tail: self.tail,
             num_items: self.num_items,
             round_function_execution_pairs: states,
         };
@@ -562,6 +556,12 @@ impl<
         }
     }
 
+    pub fn replace_container(mut self, container: C) -> (Self, C) {
+        let prev_container = self.witness;
+        self.witness = container;
+        (self, prev_container)
+    }
+
     pub fn take_sponge_like_queue_state(&self) -> QueueStateWitness<F, SW> {
         let result = QueueStateWitness {
             head: self.head,
@@ -605,8 +605,6 @@ impl<
         let intermediate_info = FullWidthQueueIntermediateStates {
             head: self.head,
             tail: new_tail,
-            old_head: self.head,
-            old_tail,
             num_items: self.num_items,
             round_function_execution_pairs: states,
         };
@@ -644,7 +642,7 @@ pub struct FullWidthStackSimulator<
 > {
     pub state: [F; SW],
     pub num_items: u32,
-    pub witness: Vec<([F; N], [F; SW], I)>,
+    pub witness: Vec<([F; SW], I)>,
 }
 
 impl<
@@ -697,7 +695,7 @@ impl<
 
         let states = make_round_function_pairs(old_state, states);
 
-        self.witness.push((encoding, self.state, element));
+        self.witness.push((self.state, element));
         self.num_items += 1;
         self.state = new_state;
 
@@ -727,7 +725,7 @@ impl<
         let popped = self.witness.pop().unwrap();
         self.num_items -= 1;
 
-        let (_element_encoding, previous_state, element) = popped;
+        let (previous_state, element) = popped;
         let encoding = element.encoding_witness();
 
         let mut state = previous_state;
@@ -867,6 +865,9 @@ mod tests {
         }
         assert_eq!(queue.num_items, 10);
 
+        let old_head = queue.head;
+        let old_tail = queue.tail;
+
         // pop one element
         let (element, data) = queue.pop_and_output_intermediate_data(&round_function);
         // it should return the first one that we entered (with circuit 0).
@@ -876,8 +877,8 @@ mod tests {
         assert_eq!(data.num_items, 9);
 
         assert_eq!(data.head, tail_after_first);
-        assert_eq!(data.old_head, empty_head);
-        assert_eq!(data.old_tail, data.tail);
+        assert_eq!(old_head, empty_head);
+        assert_eq!(data.tail, old_tail);
 
         let mut parts = queue.split_by(3, &round_function);
 

circuit_encodings/src/memory_query.rs

@@ -157,3 +157,7 @@ impl<F: SmallField> CircuitEquivalentReflection<F> for MemoryQuery {
         }
     }
 }
+
+use zkevm_circuits::base_structures::memory_query::MemoryQueryWitness;
+pub type MemoryQueueStateWitnesses<F> =
+    Vec<(MemoryQueryWitness<F>, [F; FULL_SPONGE_QUEUE_STATE_WIDTH])>;

src/external_calls.rs

@@ -5,6 +5,7 @@ pub use crate::run_vms::SCHEDULER_TIMESTAMP;
 use crate::snark_wrapper::boojum::field::goldilocks::GoldilocksExt2;
 use crate::snark_wrapper::boojum::gadgets::recursion::recursive_tree_hasher::CircuitGoldilocksPoseidon2Sponge;
 use crate::toolset::GeometryConfig;
+use crate::witness::oracle::WitnessGenerationArtifact;
 use crate::witness::tree::BinarySparseStorageTree;
 use crate::witness::tree::ZkSyncStorageLeaf;
 use crate::zk_evm::abstractions::Storage;
@@ -28,15 +29,7 @@ use circuit_definitions::Field as MainField;
 ///
 /// This function will setup the environment and will run out-of-circuit and then in-circuit.
 /// GenericNoopTracer will be used as out-of-circuit tracer
-pub fn run<
-    S: Storage,
-    CB: FnMut(ZkSyncBaseLayerCircuit),
-    QSCB: FnMut(
-        u64,
-        RecursionQueueSimulator<MainField>,
-        Vec<ClosedFormInputCompactFormWitness<MainField>>,
-    ),
->(
+pub fn run<S: Storage, CB: FnMut(WitnessGenerationArtifact)>(
     caller: Address,                 // for real block must be zero
     entry_point_address: Address,    // for real block must be the bootloader
     entry_point_code: Vec<[u8; 32]>, // for read block must be a bootloader code
@@ -52,8 +45,7 @@ pub fn run<
     tree: impl BinarySparseStorageTree<256, 32, 32, 8, 32, Blake2s256, ZkSyncStorageLeaf>,
     trusted_setup_path: &str,
     eip_4844_repack_inputs: [Option<Vec<u8>>; MAX_4844_BLOBS_PER_BLOCK],
-    circuit_callback: CB,
-    queue_simulator_callback: QSCB,
+    artifacts_callback: CB,
 ) -> (
     SchedulerCircuitInstanceWitness<MainField, CircuitGoldilocksPoseidon2Sponge, GoldilocksExt2>,
     BlockAuxilaryOutputWitness<MainField>,
@@ -75,8 +67,7 @@ pub fn run<
         tree,
         trusted_setup_path,
         eip_4844_repack_inputs,
-        circuit_callback,
-        queue_simulator_callback,
+        artifacts_callback,
         &mut out_of_circuit_tracer,
     ) {
         Ok((scheduler_circuit_witness, aux_data)) => (scheduler_circuit_witness, aux_data),

src/run_vms.rs

@@ -7,6 +7,7 @@ use crate::snark_wrapper::boojum::gadgets::recursion::recursive_tree_hasher::Cir
 use crate::toolset::create_tools;
 use crate::toolset::GeometryConfig;
 use crate::witness::oracle::create_artifacts_from_tracer;
+use crate::witness::oracle::WitnessGenerationArtifact;
 use crate::witness::tracer::tracer::WitnessTracer;
 use crate::witness::tree::BinarySparseStorageTree;
 use crate::witness::tree::ZkSyncStorageLeaf;
@@ -61,15 +62,7 @@ pub type RunVMsResult = (
 /// - witness with AUX data (with information that might be useful during verification to generate the public input)
 ///
 /// This function will setup the environment and will run out-of-circuit and then in-circuit
-pub fn run_vms<
-    S: Storage,
-    CB: FnMut(ZkSyncBaseLayerCircuit),
-    QSCB: FnMut(
-        u64,
-        RecursionQueueSimulator<MainField>,
-        Vec<ClosedFormInputCompactFormWitness<MainField>>,
-    ),
->(
+pub fn run_vms<S: Storage, CB: FnMut(WitnessGenerationArtifact)>(
     caller: Address,                 // for real block must be zero
     entry_point_address: Address,    // for real block must be the bootloader
     entry_point_code: Vec<[u8; 32]>, // for read block must be a bootloader code
@@ -85,8 +78,7 @@ pub fn run_vms<
     tree: impl BinarySparseStorageTree<256, 32, 32, 8, 32, Blake2s256, ZkSyncStorageLeaf>,
     trusted_setup_path: &str,
     eip_4844_repack_inputs: [Option<Vec<u8>>; MAX_4844_BLOBS_PER_BLOCK],
-    circuit_callback: CB,
-    queue_simulator_callback: QSCB,
+    artifacts_callback: CB,
     out_of_circuit_tracer: &mut impl Tracer<SupportedMemory = SimpleMemory>,
 ) -> Result<RunVMsResult, RunVmError> {
     let round_function = ZkSyncDefaultRoundFunction::default();
@@ -254,8 +246,7 @@ pub fn run_vms<
         evm_simulator_code_hash,
         eip_4844_repack_inputs.clone(),
         trusted_setup_path,
-        circuit_callback,
-        queue_simulator_callback,
+        artifacts_callback,
     );
 
     let (scheduler_circuit_witness, aux_data) = {

src/tests/complex_tests/mod.rs

@@ -39,6 +39,7 @@ use crate::zk_evm::GenericNoopTracer;
 use crate::zkevm_circuits::eip_4844::input::*;
 use crate::zkevm_circuits::scheduler::block_header::MAX_4844_BLOBS_PER_BLOCK;
 use crate::zkevm_circuits::scheduler::input::SchedulerCircuitInstanceWitness;
+use boojum::gadgets::queue::full_state_queue::FullStateCircuitQueueRawWitness;
 use circuit_definitions::aux_definitions::witness_oracle::VmWitnessOracle;
 use circuit_definitions::circuit_definitions::aux_layer::compression::{
     self, CompressionMode1Circuit,
@@ -58,6 +59,7 @@ use circuit_definitions::{
 use circuit_definitions::{Field, RoundFunction};
 use utils::read_basic_test_artifact;
 
+use witness::oracle::WitnessGenerationArtifact;
 use zkevm_assembly::Assembly;
 
 #[ignore = "Too slow"]
@@ -255,6 +257,27 @@ pub(crate) fn generate_base_layer(
 
     let mut basic_block_circuits = vec![];
     let mut recursion_queues = vec![];
+    let mut unsorted_memory_queue_witnesses = vec![];
+    let mut sorted_memory_queue_witnesses = vec![];
+
+    let artifacts_callback = |artifact: WitnessGenerationArtifact| match artifact {
+        WitnessGenerationArtifact::BaseLayerCircuit(circuit) => basic_block_circuits.push(circuit),
+        WitnessGenerationArtifact::RecursionQueue((a, b, c)) => recursion_queues.push((
+            a,
+            b,
+            c.into_iter()
+                .map(|x| ZkSyncBaseLayerStorage::from_inner(a as u8, x))
+                .collect(),
+        )),
+        WitnessGenerationArtifact::MemoryQueueWitness((witnesses, sorted)) => {
+            if sorted {
+                sorted_memory_queue_witnesses.push(witnesses);
+            } else {
+                unsorted_memory_queue_witnesses.push(witnesses);
+            }
+        }
+    };
+
     let (scheduler_partial_input, _aux_data) = run(
         Address::zero(),
         test_artifact.entry_point_address,
@@ -271,18 +294,28 @@ pub(crate) fn generate_base_layer(
         tree,
         "kzg/src/trusted_setup.json",
         blobs,
-        |circuit| basic_block_circuits.push(circuit),
-        |a, b, c| {
-            recursion_queues.push((
-                a,
-                b,
-                c.into_iter()
-                    .map(|x| ZkSyncBaseLayerStorage::from_inner(a as u8, x))
-                    .collect(),
-            ))
-        },
+        artifacts_callback,
     );
 
+    let mut unsorted_memory_queue_witnesses_it = unsorted_memory_queue_witnesses.into_iter();
+    let mut sorted_memory_queue_witnesses = sorted_memory_queue_witnesses.into_iter();
+    for el in basic_block_circuits.iter_mut() {
+        match &el {
+            ZkSyncBaseLayerCircuit::RAMPermutation(inner) => {
+                let mut witness = inner.witness.take().unwrap();
+                witness.sorted_queue_witness = FullStateCircuitQueueRawWitness {
+                    elements: sorted_memory_queue_witnesses.next().unwrap().into(),
+                };
+                witness.unsorted_queue_witness = FullStateCircuitQueueRawWitness {
+                    elements: unsorted_memory_queue_witnesses_it.next().unwrap().into(),
+                };
+
+                inner.witness.store(Some(witness));
+            }
+            _ => {}
+        }
+    }
+
     (
         basic_block_circuits,
         recursion_queues,

src/tests/run_manually.rs

@@ -20,9 +20,11 @@ use crate::zk_evm::witness_trace::VmWitnessTracer;
 use crate::zk_evm::GenericNoopTracer;
 use crate::zkevm_circuits::base_structures::vm_state::GlobalContextWitness;
 use crate::zkevm_circuits::main_vm::main_vm_entry_point;
+use boojum::gadgets::queue::full_state_queue::FullStateCircuitQueueRawWitness;
 use circuit_definitions::aux_definitions::witness_oracle::VmWitnessOracle;
 use circuit_definitions::zk_evm::vm_state::cycle;
 use storage::{InMemoryCustomRefundStorage, StorageRefund};
+use witness::oracle::WitnessGenerationArtifact;
 use zkevm_assembly::Assembly;
 
 #[test]
@@ -260,11 +262,25 @@ pub(crate) fn run_with_options(entry_point_bytecode: Vec<[u8; 32]>, options: Opt
     save_predeployed_contracts(&mut storage_impl.storage, &mut tree, &known_contracts);
 
     let mut basic_block_circuits = vec![];
+    let mut unsorted_memory_queue_witnesses = vec![];
+    let mut sorted_memory_queue_witnesses = vec![];
 
     // we are using TestingTracer to track prints and exceptions inside out_of_circuit_vm cycles
     let mut out_of_circuit_tracer =
         TestingTracer::new(Some(storage_impl.create_refund_controller()));
 
+    let artifacts_callback = |artifact: WitnessGenerationArtifact| match artifact {
+        WitnessGenerationArtifact::BaseLayerCircuit(circuit) => basic_block_circuits.push(circuit),
+        WitnessGenerationArtifact::MemoryQueueWitness((witnesses, sorted)) => {
+            if sorted {
+                sorted_memory_queue_witnesses.push(witnesses)
+            } else {
+                unsorted_memory_queue_witnesses.push(witnesses)
+            }
+        }
+        _ => {}
+    };
+
     if let Err(err) = run_vms(
         Address::zero(),
         *BOOTLOADER_FORMAL_ADDRESS,
@@ -281,8 +297,7 @@ pub(crate) fn run_with_options(entry_point_bytecode: Vec<[u8; 32]>, options: Opt
         tree,
         "kzg/src/trusted_setup.json",
         std::array::from_fn(|_| None),
-        |circuit| basic_block_circuits.push(circuit),
-        |_, _, _| {},
+        artifacts_callback,
         &mut out_of_circuit_tracer,
     ) {
         let error_text = match err {
@@ -303,8 +318,24 @@ pub(crate) fn run_with_options(entry_point_bytecode: Vec<[u8; 32]>, options: Opt
 
     println!("Simulation and witness creation are completed");
 
+    let mut unsorted_memory_queue_witnesses_it = unsorted_memory_queue_witnesses.into_iter();
+    let mut sorted_memory_queue_witnesses = sorted_memory_queue_witnesses.into_iter();
     for el in basic_block_circuits {
         println!("Doing {} circuit", el.short_description());
+        match &el {
+            ZkSyncBaseLayerCircuit::RAMPermutation(inner) => {
+                let mut witness = inner.witness.take().unwrap();
+                witness.sorted_queue_witness = FullStateCircuitQueueRawWitness {
+                    elements: sorted_memory_queue_witnesses.next().unwrap().into(),
+                };
+                witness.unsorted_queue_witness = FullStateCircuitQueueRawWitness {
+                    elements: unsorted_memory_queue_witnesses_it.next().unwrap().into(),
+                };
+
+                inner.witness.store(Some(witness));
+            }
+            _ => {}
+        }
         base_test_circuit(el);
     }
 

src/witness/artifacts.rs

@@ -121,13 +121,6 @@ pub struct MemoryArtifacts<F: SmallField> {
         CircuitsEntryAccumulatorSparse<(u32, QueueStateWitness<F, FULL_SPONGE_QUEUE_STATE_WIDTH>)>,
 }
 
-#[derive(Derivative)]
-#[derivative(Default)]
-pub struct ImplicitMemoryArtifacts<F: SmallField> {
-    pub memory_queries: Vec<MemoryQuery>,
-    pub memory_queue_states: Vec<MemoryQueueState<F>>,
-}
-
 #[derive(Derivative)]
 #[derivative(Default)]
 pub struct LogCircuitsArtifacts<F: SmallField> {

src/witness/aux_data_structs/per_circuit_accumulator.rs

@@ -32,6 +32,10 @@ impl<T> PerCircuitAccumulatorContainer<T> {
         self.accumulated
     }
 
+    pub fn amount_of_circuits_accumulated(&self) -> usize {
+        self.circuits_data.len()
+    }
+
     pub fn last(&self) -> Option<&T> {
         let last_batch = self.circuits_data.last();
         last_batch?;
@@ -181,6 +185,10 @@ impl<T> PerCircuitAccumulator<T> {
         self.container.len()
     }
 
+    pub fn amount_of_circuits_accumulated(&self) -> usize {
+        self.container.amount_of_circuits_accumulated()
+    }
+
     pub fn push(&mut self, val: T) {
         assert!(self.container.cycles_per_circuit != 0);
         let idx = self.container.len();