chore: bench blobs with eth_runner on native prover input generation

Cargo.toml

@@ -57,8 +57,8 @@ keywords = ["blockchain", "zksync", "zk", "risc-v"]
 categories = ["cryptography"]
 
 [workspace.dependencies]
-zksync_os_evm_errors = { version = "0.0.10", default-features =  false }
-zksync_os_interface = { version = "0.0.10"}
+zksync_os_evm_errors = { git = "https://github.com/matter-labs/zksync-os-interface", branch = "alocascio-prover-input-gen", default-features =  false }
+zksync_os_interface = { git = "https://github.com/matter-labs/zksync-os-interface", branch = "alocascio-prover-input-gen", version = "0.0.10"}
 
 risc_v_simulator = { git = "https://github.com/matter-labs/zksync-airbender", tag = "v0.4.3"}
 blake2s_u32 = { git = "https://github.com/matter-labs/zksync-airbender", tag = "v0.4.3"}

api/src/lib.rs

@@ -1,54 +1,5 @@
 #![feature(allocator_api)]
 
-use std::{path::PathBuf, str::FromStr};
-
-use forward_system::run::{
-    test_impl::{InMemoryPreimageSource, InMemoryTree},
-    BlockContext, StorageCommitment,
-};
-use oracle_provider::ReadWitnessSource;
-use zksync_os_interface::traits::TxListSource;
-pub mod helpers;
-
 /// Runs the batch, and returns the output (that contains gas usage, transaction status etc.).
-pub use forward_system::run::run_block;
-use zk_ee::common_structs::da_commitment_scheme::DACommitmentScheme;
-use zk_ee::common_structs::ProofData;
-
-/// Runs a block in RISC-V - using zksync_os binary - and returns the
-/// witness that can be passed to the prover subsystem.
-pub fn run_block_generate_witness(
-    block_context: BlockContext,
-    tree: InMemoryTree,
-    preimage_source: InMemoryPreimageSource,
-    tx_source: TxListSource,
-    proof_data: ProofData<StorageCommitment>,
-    da_commitment_scheme: DACommitmentScheme,
-    zksync_os_bin_path: &str,
-) -> Vec<u32> {
-    use forward_system::run::*;
-
-    let oracle = make_oracle_for_proofs_and_dumps_for_init_data(
-        block_context,
-        tree,
-        preimage_source,
-        tx_source,
-        Some(proof_data),
-        Some(da_commitment_scheme),
-        false,
-    );
-
-    // We'll wrap the source, to collect all the reads.
-    let copy_source = ReadWitnessSource::new(oracle);
-
-    let items = copy_source.get_read_items();
-    // By default - enable diagnostics is false (which makes the test run faster).
-    let path = PathBuf::from_str(zksync_os_bin_path).unwrap();
-    let output = zksync_os_runner::run(path, None, 1 << 36, copy_source);
-
-    // We return 0s in case of failure.
-    assert_ne!(output, [0u32; 8]);
-
-    let result = items.borrow().clone();
-    result
-}
+pub use forward_system::run::{generate_proof_input, run_block};
+pub mod helpers;

basic_system/src/system_functions/mod.rs

@@ -29,9 +29,9 @@ fn bytereverse(input: &mut [u8]) {
 /// No std system functions implementations.
 /// All of them are following EVM specs(for precompiles and keccak opcode).
 ///
-pub struct NoStdSystemFunctions;
+pub struct NoStdSystemFunctions<const USE_ADVICE: bool>;
 
-impl<R: Resources> SystemFunctions<R> for NoStdSystemFunctions {
+impl<R: Resources, const USE_ADVICE: bool> SystemFunctions<R> for NoStdSystemFunctions<USE_ADVICE> {
     type Keccak256 = keccak256::Keccak256Impl;
     type Sha256 = sha256::Sha256Impl;
     type Secp256k1ECRecover = ecrecover::EcRecoverImpl;
@@ -47,6 +47,8 @@ impl<R: Resources> SystemFunctions<R> for NoStdSystemFunctions {
     type PointEvaluation = point_evaluation::PointEvaluationImpl;
 }
 
-impl<R: Resources> SystemFunctionsExt<R> for NoStdSystemFunctions {
-    type ModExp = modexp::ModExpImpl;
+impl<R: Resources, const USE_ADVICE: bool> SystemFunctionsExt<R>
+    for NoStdSystemFunctions<USE_ADVICE>
+{
+    type ModExp = modexp::ModExpImpl<USE_ADVICE>;
 }

basic_system/src/system_functions/modexp/delegation/bigint.rs → basic_system/src/system_functions/modexp/advice/bigint.rs

@@ -1,15 +1,15 @@
 // Representation of big integers using primitives that are friendly for our delegations
 extern crate alloc;
 
-#[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
-use super::super::{ModExpAdviceParams, MODEXP_ADVICE_QUERY_ID};
+use super::super::MODEXP_ADVICE_QUERY_ID;
 use super::u256::*;
+use crate::system_functions::modexp::ModExpAdviceParams64;
 use alloc::vec::Vec;
 use core::alloc::Allocator;
 use core::fmt::Debug;
 use core::mem::MaybeUninit;
 use crypto::{bigint_op_delegation_raw, bigint_op_delegation_with_carry_bit_raw, BigIntOps};
-#[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
+use ruint::aliases::U256;
 use zk_ee::oracle::IOOracle;
 
 // There is a small choice to make - either we do exponentiation walking as via LE or BE exponent.
@@ -812,32 +812,32 @@
     }
 }
 
-#[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
 pub(crate) struct OracleAdvisor<'a, O: IOOracle> {
     pub(crate) inner: &'a mut O,
 }
 
-#[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
 fn write_bigint(
     it: &mut impl ExactSizeIterator<Item = usize>,
     mut to_consume: usize,
     dst: &mut BigintRepr<impl Allocator + Clone>,
 ) {
-    const {
-        assert!(core::mem::size_of::<usize>() == core::mem::size_of::<u32>());
-    }
     // NOTE: even if oracle overstates the number of digits (so - iterator length), it is not important
     // as long as caller checks that number of digits is within bounds of soundness
     unsafe {
-        let num_digits = to_consume.next_multiple_of(8) / 8;
+        const BIGINT_DIGIT_USIZE_SIZE: usize = U256::BYTES / core::mem::size_of::<usize>();
+        let num_digits =
+            to_consume.next_multiple_of(BIGINT_DIGIT_USIZE_SIZE) / BIGINT_DIGIT_USIZE_SIZE;
         let dst_capacity = dst.clear_as_capacity_mut();
         for dst in dst_capacity[..num_digits].iter_mut() {
-            let dst: *mut u32 = dst.as_mut_ptr().cast::<[u32; 8]>().cast();
-            for i in 0..8 {
+            let dst: *mut usize = dst
+                .as_mut_ptr()
+                .cast::<[usize; BIGINT_DIGIT_USIZE_SIZE]>()
+                .cast();
+            for i in 0..BIGINT_DIGIT_USIZE_SIZE {
                 if to_consume > 0 {
                     to_consume -= 1;
                     let digit = it.next().unwrap();
-                    dst.add(i).write(digit as u32);
+                    dst.add(i).write(digit);
                 } else {
                     dst.add(i).write(0);
                 }
@@ -848,7 +848,6 @@
     }
 }
 
-#[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
 impl<'a, O: IOOracle> ModexpAdvisor for OracleAdvisor<'a, O> {
     fn get_reduction_op_advice<A: Allocator + Clone>(
         &mut self,
@@ -857,37 +856,79 @@
         quotient_dst: &mut BigintRepr<A>,
         remainder_dst: &mut BigintRepr<A>,
     ) {
-        let arg: ModExpAdviceParams = {
-            let a_len = a.digits;
-            let a_ptr = a.backing.as_ptr();
-
-            let modulus_len = m.digits;
-            let modulus_ptr = m.backing.as_ptr();
-
-            assert!(modulus_len > 0);
-
-            ModExpAdviceParams {
-                op: 0,
-                a_ptr: a_ptr.addr() as u32,
-                a_len: a_len as u32,
-                b_ptr: 0,
-                b_len: 0,
-                modulus_ptr: modulus_ptr.addr() as u32,
-                modulus_len: modulus_len as u32,
-            }
+        // We use different advice params depending on architecture
+        // Both are mostly the same, main difference is the width of pointers
+        #[cfg(target_arch = "riscv32")]
+        let (mut it, q_len, r_len) = {
+            use crate::system_functions::modexp::ModExpAdviceParams;
+            let arg: ModExpAdviceParams = {
+                let a_len = a.digits;
+                let a_ptr = a.backing.as_ptr();
+
+                let modulus_len = m.digits;
+                let modulus_ptr = m.backing.as_ptr();
+
+                assert!(modulus_len > 0);
+
+                ModExpAdviceParams {
+                    op: 0,
+                    a_ptr: a_ptr.addr() as u32,
+                    a_len: a_len as u32,
+                    b_ptr: 0,
+                    b_len: 0,
+                    modulus_ptr: modulus_ptr.addr() as u32,
+                    modulus_len: modulus_len as u32,
+                }
+            };
+            // We assume that oracle's response is well-formed lengths-wise, and we will check value-wise separately
+            let mut it = self
+                .inner
+                .raw_query(
+                    MODEXP_ADVICE_QUERY_ID,
+                    &((&arg as *const ModExpAdviceParams).addr() as u32),
+                )
+                .unwrap();
+            let q_len = it.next().expect("quotient length");
+            let r_len = it.next().expect("remainder length");
+            (it, q_len, r_len)
         };
 
-        // We assume that oracle's response is well-formed lengths-wise, and we will check value-wise separately
-        let mut it = self
-            .inner
-            .raw_query(
-                MODEXP_ADVICE_QUERY_ID,
-                &((&arg as *const ModExpAdviceParams).addr() as u32),
-            )
-            .unwrap();
-
-        let q_len = it.next().expect("quotient length");
-        let r_len = it.next().expect("remainder length");
+        #[cfg(not(target_arch = "riscv32"))]
+        let (mut it, q_len, r_len) = {
+            let arg: ModExpAdviceParams64 = {
+                let a_len = a.digits;
+                let a_ptr = a.backing.as_ptr();
+
+                let modulus_len = m.digits;
+                let modulus_ptr = m.backing.as_ptr();
+
+                assert!(modulus_len > 0);
+
+                ModExpAdviceParams64 {
+                    op: 0,
+                    a_ptr: a_ptr.addr() as u64,
+                    a_len: a_len as u64,
+                    b_ptr: 0,
+                    b_len: 0,
+                    modulus_ptr: modulus_ptr.addr() as u64,
+                    modulus_len: modulus_len as u64,
+                }
+            };
+            // We assume that oracle's response is well-formed lengths-wise, and we will check value-wise separately
+            let mut it = self
+                .inner
+                .raw_query(
+                    MODEXP_ADVICE_QUERY_ID,
+                    &((&arg as *const ModExpAdviceParams64).addr() as u64),
+                )
+                .unwrap();
+            // Oracle provides lengths as u32, so in this case they are
+            // packed into a single usize
+            let packed_lens = it.next().expect("packed lengths");
+            let q_len = (packed_lens & 0xFFFF_FFFF) as usize;
+            let r_len = (packed_lens >> 32) as usize;
+            (it, q_len / 2, r_len / 2)
+        };
 
         let max_quotient_digits = if a.digits < m.digits {
             0
@@ -899,10 +940,6 @@
 
         let max_remainder_digits = m.digits;
 
-        const {
-            assert!(core::mem::size_of::<usize>() == core::mem::size_of::<u32>());
-        }
-
         // check that hint is "sane" in upper bound
 
         assert!(q_len.next_multiple_of(8) / 8 <= max_quotient_digits);

basic_system/src/system_functions/modexp/delegation/mod.rs → basic_system/src/system_functions/modexp/advice/mod.rs

@@ -11,7 +11,6 @@ use zk_ee::system::logger::Logger;
 #[cfg(feature = "testing")]
 use zk_ee::system::logger::NullLogger;
 
-#[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
 pub(super) fn modexp<O: zk_ee::oracle::IOOracle, L: Logger, A: Allocator + Clone>(
     base: &[u8],
     exp: &[u8],

basic_system/src/system_functions/modexp/mod.rs

@@ -25,29 +25,32 @@ pub const MODEXP_ADVICE_QUERY_ID: u32 = ADVICE_SUBSPACE_MASK | 0x10;
 /// Used to request division advice for big integer operations during modexp
 #[repr(C)]
 #[derive(Debug, Default)]
-pub struct ModExpAdviceParams {
-    pub op: u32,          // Operation type (0 = division)
-    pub a_ptr: u32,       // Pointer to dividend
-    pub a_len: u32,       // Length of dividend in words
-    pub b_ptr: u32,       // Pointer to divisor
-    pub b_len: u32,       // Length of divisor in words
-    pub modulus_ptr: u32, // Pointer to modulus
-    pub modulus_len: u32, // Length of modulus in words
+pub struct ModExpAdviceParamsGeneric<W> {
+    pub op: W,          // Operation type (0 = division)
+    pub a_ptr: W,       // Pointer to dividend
+    pub a_len: W,       // Length of dividend in words
+    pub b_ptr: W,       // Pointer to divisor
+    pub b_len: W,       // Length of divisor in words
+    pub modulus_ptr: W, // Pointer to modulus
+    pub modulus_len: W, // Length of modulus in words
 }
 
-#[cfg(any(
-    all(target_arch = "riscv32", feature = "proving"),
-    test,
-    feature = "testing"
-))]
-pub mod delegation;
+/// Used for proving (RISC-V 32-bit)
+pub type ModExpAdviceParams = ModExpAdviceParamsGeneric<u32>;
+
+/// Used for native execution (64-bit)
+pub type ModExpAdviceParams64 = ModExpAdviceParamsGeneric<u64>;
+
+pub mod advice;
 
 ///
 /// modexp system function implementation.
 ///
-pub struct ModExpImpl;
+pub struct ModExpImpl<const USE_ADVICE: bool>;
 
-impl<R: Resources> SystemFunctionExt<R, ModExpErrors> for ModExpImpl {
+impl<R: Resources, const USE_ADVICE: bool> SystemFunctionExt<R, ModExpErrors>
+    for ModExpImpl<USE_ADVICE>
+{
     /// If the input size is less than expected - it will be padded with zeroes.
     /// If the input size is greater - redundant bytes will be ignored.
     ///
@@ -71,7 +74,9 @@ impl<R: Resources> SystemFunctionExt<R, ModExpErrors> for ModExpImpl {
         allocator: A,
     ) -> Result<(), SubsystemError<ModExpErrors>> {
         cycle_marker::wrap_with_resources!("modexp", resources, {
-            modexp_as_system_function_inner(input, output, resources, oracle, logger, allocator)
+            modexp_as_system_function_inner::<_, _, _, _, _, USE_ADVICE>(
+                input, output, resources, oracle, logger, allocator,
+            )
         })
     }
 }
@@ -106,6 +111,9 @@ fn modexp_as_system_function_inner<
     D: ?Sized + TryExtend<u8>,
     A: Allocator + Clone,
     R: Resources,
+    // Toggle delegation-based implementation for forward run, to be able
+    // to capture oracle queries.
+    const USE_ADVICE: bool,
 >(
     input: &[u8],
     dst: &mut D,
@@ -211,7 +219,7 @@ fn modexp_as_system_function_inner<
     // Call the modexp.
 
     #[cfg(any(all(target_arch = "riscv32", feature = "proving"), test))]
-    let output = self::delegation::modexp(
+    let output = self::advice::modexp(
         base.as_slice(),
         exponent.as_slice(),
         modulus.as_slice(),
@@ -221,12 +229,23 @@ fn modexp_as_system_function_inner<
     );
 
     #[cfg(not(any(all(target_arch = "riscv32", feature = "proving"), test)))]
-    let output = ::modexp::modexp(
-        base.as_slice(),
-        exponent.as_slice(),
-        modulus.as_slice(),
-        allocator,
-    );
+    let output = if USE_ADVICE {
+        self::advice::modexp(
+            base.as_slice(),
+            exponent.as_slice(),
+            modulus.as_slice(),
+            oracle,
+            logger,
+            allocator,
+        )
+    } else {
+        ::modexp::modexp(
+            base.as_slice(),
+            exponent.as_slice(),
+            modulus.as_slice(),
+            allocator,
+        )
+    };
 
     if output.len() >= mod_len {
         // truncate