perf: Add profile-generating tests for the FRI verifier

tasm-lib/src/lib.rs

@@ -53,6 +53,8 @@ use triton_vm::profiler::TritonProfiler;
 use triton_vm::table::master_table::TableId;
 pub use triton_vm::twenty_first;
 
+use crate::test_helpers::prepend_program_with_stack_setup;
+
 // The hasher type must match whatever algebraic hasher the VM is using
 pub type VmHasher = Tip5;
 pub type Digest = tip5::Digest;
@@ -299,18 +301,11 @@ pub fn prove_and_verify(
 
     // Construct the program that initializes the stack to the expected start value.
     // If this is not done, a stack underflow will occur for most programs
-    let stack_initialization_code = match init_stack {
-        Some(init_stack) => init_stack
-            .into_iter()
-            .skip(NUM_OP_STACK_REGISTERS)
-            .map(|word| triton_instr!(push word))
-            .collect(),
-        None => triton_asm!(),
+    let program = match &init_stack {
+        Some(init_stack) => prepend_program_with_stack_setup(init_stack, program),
+        None => program.to_owned(),
     };
 
-    let program =
-        Program::new(&[stack_initialization_code, program.labelled_instructions()].concat());
-
     let (aet, _) = program
         .trace_execution(PublicInput::new(std_in.to_owned()), nondeterminism.clone())
         .unwrap();

tasm-lib/src/test_helpers.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::fmt::Display;
 
 use itertools::Itertools;
+use triton_vm::op_stack::NUM_OP_STACK_REGISTERS;
 use triton_vm::prelude::*;
 
 use crate::dyn_malloc::DYN_MALLOC_ADDRESS;
@@ -493,3 +494,20 @@ pub fn test_rust_equivalence_given_execution_state<T: BasicSnippet + RustShadow>
         None,
     )
 }
+
+pub fn prepend_program_with_stack_setup(
+    init_stack: &[BFieldElement],
+    program: &Program,
+) -> Program {
+    let stack_initialization_code = init_stack
+        .iter()
+        .skip(NUM_OP_STACK_REGISTERS)
+        .map(|&word| triton_instr!(push word))
+        .collect_vec();
+
+    Program::new(&[stack_initialization_code, program.labelled_instructions()].concat())
+}
+
+pub fn prepend_program_with_sponge_init(program: &Program) -> Program {
+    Program::new(&[triton_asm!(sponge_init), program.labelled_instructions()].concat())
+}

tasm-lib/src/verifier/fri/verify.rs

@@ -1479,15 +1479,22 @@ mod test {
     }
 
     #[derive(Debug, Clone, test_strategy::Arbitrary)]
-    struct TestCase {
+    pub(super) struct TestCase {
         #[strategy(any::<ArbitraryFriVerify>().prop_map(|x| x.fri_verify()))]
-        fri_verify: FriVerify,
+        pub(super) fri_verify: FriVerify,
 
         #[strategy(vec(arb(), #fri_verify.first_round_max_degree()))]
         polynomial_coefficients: Vec<XFieldElement>,
     }
 
     impl TestCase {
+        fn thirty_two_coefficients() -> [u64; 32] {
+            [
+                92, 42, 91, 59, 86, 64, 5, 64, 74, 53, 54, 68, 54, 23, 24, 58, 15, 44, 33, 31, 38,
+                97, 25, 69, 11, 67, 66, 33, 37, 58, 43, 14,
+            ]
+        }
+
         /// A test case with no FRI rounds.
         fn tiny_case() -> Self {
             let fri_verify = FriVerify::new(BFieldElement::one(), 2, 2, 1);
@@ -1504,10 +1511,7 @@ mod test {
             let fri_verify = FriVerify::new(BFieldElement::one(), 64, 2, 7);
             assert_eq!(2, fri_verify.num_rounds(), "test case must be meaningful");
 
-            let coefficients = [
-                92, 42, 91, 59, 86, 64, 5, 64, 74, 53, 54, 68, 54, 23, 24, 58, 15, 44, 33, 31, 38,
-                97, 25, 69, 11, 67, 66, 33, 37, 58, 43, 14,
-            ];
+            let coefficients = Self::thirty_two_coefficients();
             assert_eq!(fri_verify.first_round_max_degree() + 1, coefficients.len());
 
             Self {
@@ -1516,6 +1520,24 @@ mod test {
             }
         }
 
+        /// A test case resembling that which will be run by the STARK verifier
+        pub(super) fn recursive_case(
+            inner_padded_height: usize,
+            log2_expansion_factor: usize,
+        ) -> Self {
+            const SECURITY_LEVEL: usize = 160;
+            let stark = Stark::new(SECURITY_LEVEL, log2_expansion_factor);
+
+            let fri_verify: FriVerify = stark.derive_fri(inner_padded_height).unwrap().into();
+
+            let coefficients = Self::thirty_two_coefficients();
+
+            Self {
+                fri_verify,
+                polynomial_coefficients: coefficients.map(|n| xfe!([n; 3])).to_vec(),
+            }
+        }
+
         fn fri(&self) -> Fri<Tip5> {
             self.fri_verify.to_fri()
         }
@@ -1549,7 +1571,7 @@ mod test {
             }
         }
 
-        fn initial_state(&self) -> ProcedureInitialState {
+        pub(super) fn initial_state(&self) -> ProcedureInitialState {
             let fri_verify = self.fri_verify;
             let proof_stream = self.proof_stream();
             let snippet = FriSnippet {
@@ -1720,8 +1742,19 @@ mod test {
 
 #[cfg(test)]
 mod bench {
+    use std::fs::create_dir_all;
+    use std::fs::File;
+    use std::io::Write;
+    use std::path::Path;
+    use std::path::PathBuf;
+
+    use crate::generate_full_profile;
+    use crate::test_helpers::prepend_program_with_sponge_init;
+    use crate::test_helpers::prepend_program_with_stack_setup;
+    use crate::traits::procedure::ProcedureInitialState;
     use crate::traits::procedure::ShadowedProcedure;
     use crate::traits::rust_shadow::RustShadow;
+    use crate::verifier::fri::verify::test::TestCase;
 
     use super::*;
 
@@ -1731,7 +1764,6 @@ mod bench {
         let domain_length = 1 << 10;
         let offset = BFieldElement::new(7);
         let num_collinearity_checks = 80;
-        // tiny parameters for FRI yes, but the bench framework is awful atm
         let fri_verify = FriVerify::new(
             offset,
             domain_length,
@@ -1743,4 +1775,62 @@ mod bench {
         };
         ShadowedProcedure::new(snippet).bench();
     }
+
+    #[ignore = "Takes many minutes to run"]
+    #[test]
+    fn profile_fri() {
+        const INNER_PADDED_HEIGHTS: [usize; 4] = [1 << 8, 1 << 9, 1 << 10, 1 << 11];
+        const SECURITY_LEVEL: usize = 160;
+
+        for inner_padded_height in INNER_PADDED_HEIGHTS {
+            for log2_fri_expansion_factor in 1..=5 {
+                let stark = Stark::new(SECURITY_LEVEL, log2_fri_expansion_factor);
+                profile_and_bench_fri_run(stark, inner_padded_height);
+            }
+        }
+    }
+
+    fn profile_and_bench_fri_run(stark: Stark, inner_padded_height: usize) {
+        let test_case = TestCase::recursive_case(
+            inner_padded_height,
+            stark.fri_expansion_factor.ilog2() as usize,
+        );
+        let fri_params = stark.derive_fri(inner_padded_height).unwrap();
+        let name = format!(
+            "fri_verify_expansion_factor_inner_padded_height_{}_expansion_factor_{}",
+            inner_padded_height, fri_params.expansion_factor
+        );
+
+        let ProcedureInitialState {
+            stack: initial_stack,
+            nondeterminism,
+            public_input: stdin,
+            sponge: _,
+        } = test_case.initial_state();
+        let snippet = FriSnippet {
+            test_instance: test_case.fri_verify,
+        };
+        let code = snippet.link_for_isolated_run();
+        let program = Program::new(&code);
+        let program = prepend_program_with_stack_setup(&initial_stack, &program);
+        let program = prepend_program_with_sponge_init(&program);
+
+        let profile = generate_full_profile(&name, program, &stdin.into(), &nondeterminism);
+
+        // write profile to profile file
+        let mut path = PathBuf::new();
+        path.push("profiles");
+        create_dir_all(&path).expect("profiles directory must be created successfully");
+
+        path.push(Path::new(&name).with_extension("profile"));
+        let mut file = File::create(path).expect("open file for writing");
+        write!(file, "{profile}").unwrap();
+
+        println!("{}", profile);
+
+        // Don't do this since entrypoint label of FRI snippet is not parameterized over its
+        // variables, meaning all benchmarks are overwritten.
+        // // Also create benchmarks for quick reference
+        // ShadowedProcedure::new(snippet).bench();
+    }
 }

tasm-lib/src/verifier/stark_verify.rs

@@ -1208,10 +1208,11 @@ mod benches {
     #[ignore = "Takes a fairly long time. Intended to find optimal FRI expansion factor."]
     #[test]
     fn small_benchmark_different_fri_expansion_factors() {
-        for log2_of_fri_expansion_factor in 2..=5 {
+        for log2_of_fri_expansion_factor in 1..=5 {
             let stark = Stark::new(160, log2_of_fri_expansion_factor);
             benchmark_verifier(10, 1 << 8, stark);
             benchmark_verifier(40, 1 << 9, stark);
+            benchmark_verifier(80, 1 << 10, stark);
         }
     }