Add a modexp example guest program + benchmark #1067

.cargo/config.toml

@@ -0,0 +1,3 @@
+[target.x86_64-pc-windows-msvc]
+linker = "rust-lld.exe"
+rustflags = ["-C", "link-arg=-fuse-ld=lld"]

Cargo.toml

@@ -73,9 +73,9 @@ members = [
     "examples/merkle-tree/guest",
     "examples/hash-bench",
     "examples/hash-bench/guest",
+    "examples/modexp",
+    "examples/modexp/guest",
     "zklean-extractor",
-]
-
 [features]
 host = ["jolt-sdk/host"]
 

examples/modexp/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "modexp"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
+tracing-subscriber = "0.3"
+tracing = "0.1"
+guest = { package = "modexp-guest", path = "./guest" }
+
+hex = "0.4.3"

examples/modexp/guest/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "modexp-guest"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+guest = []
+
+[dependencies]
+jolt = { package = "jolt-sdk", path = "../../../jolt-sdk", features = [] }

examples/modexp/guest/src/lib.rs

@@ -0,0 +1,37 @@
+#![cfg_attr(feature = "guest", no_std)]
+
+#[jolt::provable(memory_size = 10240, max_trace_length = 65536)]
+fn modexp_chain(base: u128, exp: u128, modulus: u128, num_iters: u32) -> u128 {
+    // multiply modulo m without overflowing by using shift-add
+    fn mul_mod(mut a: u128, mut b: u128, m: u128) -> u128 {
+        let mut res: u128 = 0;
+        a %= m;
+        while b > 0 {
+            if (b & 1) == 1 {
+                res = (res + a) % m;
+            }
+            a = (a << 1) % m;
+            b >>= 1;
+        }
+        res
+    }
+
+    fn mod_pow(mut base: u128, mut exp: u128, m: u128) -> u128 {
+        let mut result: u128 = 1 % m;
+        base %= m;
+        while exp > 0 {
+            if (exp & 1) == 1 {
+                result = mul_mod(result, base, m);
+            }
+            base = mul_mod(base, base, m);
+            exp >>= 1;
+        }
+        result
+    }
+
+    let mut out: u128 = 0;
+    for _ in 0..num_iters {
+        out = mod_pow(base, exp, modulus);
+    }
+    out
+}

examples/modexp/src/main.rs

@@ -0,0 +1,31 @@
+use std::time::Instant;
+use tracing::info;
+
+pub fn main() {
+    tracing_subscriber::fmt::init();
+
+    let target_dir = "/tmp/jolt-guest-targets";
+    let mut program = guest::compile_modexp_chain(target_dir);
+
+    let prover_preprocessing = guest::preprocess_prover_modexp_chain(&mut program);
+    let verifier_preprocessing =
+        guest::verifier_preprocessing_from_prover_modexp_chain(&prover_preprocessing);
+
+    let prove_modexp_chain = guest::build_prover_modexp_chain(program, prover_preprocessing);
+    let verify_modexp_chain = guest::build_verifier_modexp_chain(verifier_preprocessing);
+
+    let input_base = 123456789u128;
+    let input_exp = 65537u128;
+    let input_mod = 1000000007u128;
+    let iters = 100u32;
+    let native_output = guest::modexp_chain(input_base, input_exp, input_mod, iters);
+    let now = Instant::now();
+    let (output, proof, program_io) = prove_modexp_chain(input_base, input_exp, input_mod, iters);
+    info!("Prover runtime: {} s", now.elapsed().as_secs_f64());
+    let is_valid = verify_modexp_chain(input_base, input_exp, input_mod, iters, output, program_io.panic, proof);
+
+    assert_eq!(output, native_output, "output mismatch");
+    info!("output: {}", hex::encode(output.to_be_bytes()));
+    info!("native_output: {}", hex::encode(native_output.to_be_bytes()));
+    info!("valid: {is_valid}");
+}

jolt-core/benches/e2e_profiling.rs

@@ -13,6 +13,7 @@ const CYCLES_PER_SHA256: f64 = 3396.0;
 const CYCLES_PER_SHA3: f64 = 4330.0;
 const CYCLES_PER_BTREEMAP_OP: f64 = 1550.0;
 const CYCLES_PER_FIBONACCI_UNIT: f64 = 12.0;
+const CYCLES_PER_MODEXP: f64 = 50000.0;
 const SAFETY_MARGIN: f64 = 0.9; // Use 90% of max trace capacity
 
 /// Calculate number of operations to target a specific cycle count
@@ -32,6 +33,8 @@ pub enum BenchType {
     Sha2Chain,
     #[strum(serialize = "SHA3 Chain")]
     Sha3Chain,
+    #[strum(serialize = "modexp-chain")]
+    ModExp,
 }
 
 pub fn benchmarks(bench_type: BenchType) -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
@@ -41,6 +44,7 @@ pub fn benchmarks(bench_type: BenchType) -> Vec<(tracing::Span, Box<dyn FnOnce()
         BenchType::Sha3 => sha3(),
         BenchType::Sha2Chain => sha2_chain(),
         BenchType::Sha3Chain => sha3_chain(),
+        BenchType::ModExp => modexp_chain(),
         BenchType::Fibonacci => fibonacci(),
     }
 }
@@ -87,6 +91,18 @@ fn sha3_chain() -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
     prove_example("sha3-chain-guest", inputs)
 }
 
+fn modexp_chain() -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
+    // simple u128 inputs for example
+    let mut inputs = vec![];
+    // base, exp, modulus
+    inputs.append(&mut postcard::to_stdvec(&123456789u128).unwrap());
+    inputs.append(&mut postcard::to_stdvec(&65537u128).unwrap());
+    inputs.append(&mut postcard::to_stdvec(&1000000007u128).unwrap());
+    let iters = scale_to_target_ops(((1 << 24) as f64 * SAFETY_MARGIN) as usize, CYCLES_PER_MODEXP);
+    inputs.append(&mut postcard::to_stdvec(&iters).unwrap());
+    prove_example("modexp-guest", inputs)
+}
+
 pub fn master_benchmark(
     bench_type: BenchType,
     bench_scale: usize,
@@ -129,6 +145,16 @@ pub fn master_benchmark(
                 ]
                 .concat()
             }),
+            BenchType::ModExp => ("modexp-chain", |target| {
+                let iterations = scale_to_target_ops(target, CYCLES_PER_MODEXP);
+                [
+                    postcard::to_stdvec(&123456789u128).unwrap(),
+                    postcard::to_stdvec(&65537u128).unwrap(),
+                    postcard::to_stdvec(&1000000007u128).unwrap(),
+                    postcard::to_stdvec(&iterations).unwrap(),
+                ]
+                .concat()
+            }),
             BenchType::BTreeMap => ("btreemap", |target| {
                 postcard::to_stdvec(&scale_to_target_ops(target, CYCLES_PER_BTREEMAP_OP)).unwrap()
             }),

rust-toolchain.toml

@@ -1,5 +1,5 @@
 [toolchain]
-channel = "1.88"
-targets = ["riscv32imac-unknown-none-elf", "riscv64imac-unknown-none-elf"]
+channel = "nightly"
 profile = "minimal"
-components = ["cargo", "rustc", "clippy", "rustfmt"]
+components = ["rust-src", "cargo", "rustc", "clippy", "rustfmt"]
+targets = ["riscv32imac-unknown-none-elf", "riscv64imac-unknown-none-elf"]