bench: Poseidon2 with user supplied number of bytes

circuits/src/generation/cpu.rs

@@ -121,6 +121,7 @@ pub fn generate_cpu_trace<F: RichField>(record: &ExecutionRecord<F>) -> Vec<CpuS
     }
 
     log::trace!("trace {:?}", trace);
+    log::info!("cpu trace len {}", trace.len().next_power_of_two());
     pad_trace_with_last(trace)
 }
 

circuits/src/generation/memory.rs

@@ -199,6 +199,10 @@ pub fn generate_memory_trace<F: RichField>(
         .collect();
 
     log::trace!("trace {:?}", merged_trace);
+    log::info!(
+        "memory trace len {:?}",
+        merged_trace.len().next_power_of_two()
+    );
     pad_mem_trace(merged_trace)
 }
 

circuits/src/poseidon2/generation.rs

@@ -154,7 +154,8 @@ pub fn generate_poseidon2_trace<F: RichField>(step_rows: &[Row<F>]) -> Vec<Posei
             .collect::<Vec<Poseidon2State<F>>>(),
         generate_poseidon2_state(&[F::ZERO; STATE_SIZE], false),
     );
-    log::trace!("Poseison2 trace {:?}", trace);
+    log::trace!("Poseidon2 trace {:?}", trace);
+    log::info!("Poseidon2 trace length {:?}", trace.len());
     trace
 }
 

cli/Cargo.toml

@@ -29,7 +29,7 @@ clio = { version = "0.3", features = ["clap-parse"] }
 env_logger = "0.11"
 itertools = "0.12"
 log = "0.4"
-mozak-examples = { path = "../examples-builder", features = ["mozak-sort"] }
+mozak-examples = { path = "../examples-builder", features = ["mozak-sort", "mozak-poseidon2"] }
 plonky2 = { workspace = true, default-features = false }
 rkyv = { version = "=0.8.0-alpha.1", default-features = false, features = ["pointer_width_32", "alloc"] }
 rkyv_derive = "=0.8.0-alpha.1"

cli/src/cli_benches/benches.rs

@@ -6,6 +6,7 @@ use clap::{Args as Args_, Subcommand};
 use super::nop::NopBench;
 use super::omni::OmniBench;
 use super::poseidon2::Poseidon2Bench;
+use super::poseidon2_elf::Poseidon2ELFBench;
 use super::sort::{SortBench, SortBenchRecursive};
 use super::xor::XorBench;
 
@@ -55,6 +56,9 @@ pub enum BenchFunction {
     SortBench {
         n: u32,
     },
+    Poseidon2ELFBench {
+        n: u32,
+    },
     SortBenchRecursive {
         n: u32,
     },
@@ -69,6 +73,7 @@ impl BenchArgs {
             BenchFunction::Poseidon2Bench { input_len } => Poseidon2Bench.bench(input_len),
             BenchFunction::SortBench { n } => SortBench.bench(n),
             BenchFunction::SortBenchRecursive { n } => SortBenchRecursive.bench(n),
+            BenchFunction::Poseidon2ELFBench { n } => Poseidon2ELFBench.bench(n),
         }
     }
 }

cli/src/cli_benches/mod.rs

@@ -2,5 +2,6 @@ pub mod benches;
 pub mod nop;
 pub mod omni;
 pub mod poseidon2;
+pub mod poseidon2_elf;
 pub mod sort;
 pub mod xor;

cli/src/cli_benches/poseidon2_elf.rs

@@ -0,0 +1,49 @@
+use anyhow::Result;
+use mozak_circuits::test_utils::{prove_and_verify_mozak_stark, F};
+use mozak_examples::MOZAK_POSEIDON2_ELF;
+use mozak_runner::elf::Program;
+use mozak_runner::state::{RawTapes, State};
+use mozak_runner::vm::{step, ExecutionRecord};
+use starky::config::StarkConfig;
+
+use super::benches::Bench;
+
+pub fn poseidon2_elf_prepare(n: u32) -> Result<(Program, ExecutionRecord<F>)> {
+    let program = Program::vanilla_load_elf(MOZAK_POSEIDON2_ELF)?;
+    let raw_tapes = RawTapes {
+        public_tape: n.to_le_bytes().to_vec(),
+        ..Default::default()
+    };
+    let state = State::new(program.clone(), raw_tapes);
+    let record = step(&program, state)?;
+    Ok((program, record))
+}
+pub fn poseidon2_elf_execute(
+    result: Result<(Program, ExecutionRecord<F>)>,
+) -> Result<(), anyhow::Error> {
+    let (program, record) = result?;
+    prove_and_verify_mozak_stark(&program, &record, &StarkConfig::standard_fast_config())
+}
+
+pub(crate) struct Poseidon2ELFBench;
+
+impl Bench for Poseidon2ELFBench {
+    type Args = u32;
+    type Prepared = Result<(Program, ExecutionRecord<F>)>;
+
+    fn prepare(&self, args: &Self::Args) -> Self::Prepared { poseidon2_elf_prepare(*args) }
+
+    fn execute(&self, prepared: Self::Prepared) -> anyhow::Result<()> {
+        poseidon2_elf_execute(prepared)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{poseidon2_elf_execute, poseidon2_elf_prepare};
+    #[test]
+    fn test_poseidon2_elf_with_run() -> anyhow::Result<()> {
+        let n = 100;
+        poseidon2_elf_execute(poseidon2_elf_prepare(n))
+    }
+}

examples-builder/Cargo.toml

@@ -22,3 +22,4 @@ sha2 = []
 static-mem-access = []
 token = []
 wallet = []
+mozak-poseidon2 = []

examples-builder/build.rs

@@ -40,6 +40,7 @@ const CRATES: &[Crate] = &[
     ecrate!("static-mem-access", "STATIC_MEM_ACCESS_ELF", false),
     ecrate!("empty", "EMPTY_ELF", false),
     ecrate!("mozak-sort", "MOZAK_SORT_ELF", false),
+    ecrate!("mozak-poseidon2", "MOZAK_POSEIDON2_ELF", false),
     ecrate!("tokenbin", "TOKENBIN", false),
     ecrate!("walletbin", "WALLETBIN", false),
     ecrate!("inputtapebin", "INPUTTAPEBIN", false),

examples/Cargo.toml

@@ -11,6 +11,7 @@ members = [
   "rkyv-serialization",
   "empty",
   "mozak-sort",
+  "mozak-poseidon2",
 
   # Members that depend on more than sdk::core
   "token",

examples/mozak-poseidon2/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+edition = "2021"
+name = "mozak-poseidon2"
+version = "0.1.0"
+
+[dependencies]
+mozak-sdk = { path = "../../sdk", default-features = false }
+
+[[bin]]
+name = "mozak-poseidon2"
+path = "main.rs"

examples/mozak-poseidon2/main.rs

@@ -0,0 +1,32 @@
+//! The main objective is to benchmark the performance of the
+//! poseidon2 ecall in the context of flat hashing tapes.
+//! The tapes which we might want to flat hash are public
+//! tape, and call tape
+
+#![cfg_attr(target_os = "mozakvm", no_main)]
+#![cfg_attr(not(feature = "std"), no_std)]
+
+use core::hint::black_box;
+extern crate alloc;
+use alloc::vec::Vec;
+
+use mozak_sdk::core::ecall::{ioread_public, poseidon2};
+
+#[allow(clippy::unit_arg)]
+fn main() {
+    // number of bytes we would hash.
+    let n = {
+        let mut bytes = [0u8; 4];
+        ioread_public(bytes.as_mut_ptr(), bytes.len());
+        u32::from_le_bytes(bytes).next_multiple_of(8)
+    } as usize;
+
+    // generate a vector of size n
+    let v: Vec<u8> = black_box((0..n).map(|i| i.to_le_bytes()[0]).collect());
+
+    // flat hash v
+    let mut hash = [0u8; 32];
+    black_box(poseidon2(v.as_ptr(), n, hash.as_mut_ptr()));
+}
+
+mozak_sdk::entry!(main);

perftool/config.json

@@ -64,14 +64,20 @@
             }
         }
     },
-    "poseidon2": {
-        "description": "Benching Poseidon2 ECALL",
+    "poseidon2-elf": {
+        "description": "Benching Poseidon2 ecall inside ELF. We load `input_len` bytes into memory and hash it",
         "parameter": "input_len",
         "output": "time taken (in s)",
         "benches": {
-            "poseidon2": {
+            "poseidon2-elf": {
                 "commit": "latest",
-                "bench_function": "poseidon2-bench"
+                "bench_function": "poseidon2-elf-bench",
+                "elf": "examples/mozak-poseidon2"
+            },
+            "poseidon2-elf-load-memory-only": {
+                "commit": "a95d68f96e85bd5758f84b4b14e44772adc9714e",
+                "bench_function": "poseidon2-elf-bench",
+                "elf": "examples/mozak-poseidon2"
             }
         }
     }