feat: add 256-bit integer alu e2e benchmark (#487)

.github/workflows/benchmark-call.yml

@@ -11,6 +11,7 @@ on:
           - vm_verify_fibair
           - halo2_static_verify_fibair
           - tiny_e2e
+          - alu256_e2e
           - small_e2e
           - single_rw
           - single_filter
@@ -113,6 +114,14 @@ jobs:
           python3 ../sdk/scripts/bench.py $BIN_NAME $CMD_ARGS
           echo "BIN_NAME=${BIN_NAME}" >> $GITHUB_ENV
 
+      - name: Run benchmark
+        if: inputs.benchmark_name == 'alu256_e2e'
+        working-directory: recursion
+        run: |
+          BIN_NAME="alu256_e2e"
+          python3 ../sdk/scripts/bench.py $BIN_NAME $CMD_ARGS
+          echo "BIN_NAME=${BIN_NAME}" >> $GITHUB_ENV
+
       - name: Run benchmark
         if: inputs.benchmark_name == 'small_e2e'
         working-directory: recursion

.github/workflows/recursion-bench.yml

@@ -63,6 +63,20 @@ jobs:
       benchmark_name: tiny_e2e
     secrets: inherit
 
+  benchmark_alu256_e2e:
+    uses: ./.github/workflows/benchmark-call.yml
+    # run on pull request with label 'run-benchmark' or 'run-benchmark-e2e'
+    # and always run on push to main
+    if: |
+      (github.event_name == 'pull_request' &&
+      (contains(github.event.pull_request.labels.*.name, 'run-benchmark') ||
+      contains(github.event.pull_request.labels.*.name, 'run-benchmark-e2e'))) ||
+      (github.event_name == 'push' && github.ref == 'refs/heads/main')
+    with:
+      aws_instance_type: r7g.16xlarge
+      benchmark_name: alu256_e2e
+    secrets: inherit
+
   benchmark_small_e2e:
     uses: ./.github/workflows/benchmark-call.yml
     # run on non-draft pull request with label 'run-benchmark'

README.md

@@ -18,6 +18,7 @@ where `<name>` is a benchmark implemented as a rust binary (located in `src/bin`
 
 - `verify_fibair`
 - `tiny_e2e`
+- `alu256_e2e`
 - `small_e2e`
   in the `recursion` crate.
   The benchmark outputs a JSON of metrics. You can process this into markdown with:

recursion/Cargo.toml

@@ -54,6 +54,10 @@ name = "tiny_e2e"
 harness = false
 required-features = ["bench-metrics"]
 
+[[bin]]
+name = "alu256_e2e"
+required-features = ["bench-metrics"]
+
 [[bin]]
 name = "small_e2e"
 required-features = ["bench-metrics"]

recursion/src/bin/alu256_e2e.rs

@@ -0,0 +1,170 @@
+/// E2E benchmark to aggregate small program with ALU chips.
+/// Proofs:
+/// 1. Prove a program with some ALU operations.
+/// 2. Verify the proof of 1. in the inner config.
+/// 2. Verify the proof of 2. in the outer config.
+/// 3. Verify the proof of 3. using a Halo2 static verifier.
+/// 4. Wrapper Halo2 circuit to reduce the size of 4.
+use std::iter;
+
+use afs_compiler::{
+    asm::AsmBuilder,
+    conversion::CompilerOptions,
+    ir::{RVar, Var},
+};
+use afs_recursion::testing_utils::inner::build_verification_program;
+use ax_sdk::{
+    bench::run_with_metric_collection,
+    config::{
+        baby_bear_poseidon2::BabyBearPoseidon2Engine,
+        fri_params::standard_fri_params_with_100_bits_conjectured_security,
+    },
+    engine::{StarkForTest, StarkFriEngine},
+};
+use num_bigint_dig::BigUint;
+use p3_baby_bear::BabyBear;
+use p3_commit::PolynomialSpace;
+use p3_field::{extension::BinomialExtensionField, AbstractField};
+use p3_uni_stark::{Domain, StarkGenericConfig};
+use stark_vm::{program::Program, sdk::gen_vm_program_stark_for_test, vm::config::VmConfig};
+use tracing::info_span;
+
+const NUM_DIGITS: usize = 8;
+
+fn bench_program() -> Program<BabyBear> {
+    type F = BabyBear;
+    type EF = BinomialExtensionField<BabyBear, 4>;
+    let mut builder = AsmBuilder::<F, EF>::default();
+
+    let sum_digits = iter::repeat(0u32).take(NUM_DIGITS).collect::<Vec<_>>();
+    let min_digits = iter::repeat(u32::MAX).take(NUM_DIGITS).collect::<Vec<_>>();
+    let val_digits = iter::once(246)
+        .chain(iter::repeat(0u32))
+        .take(NUM_DIGITS)
+        .collect::<Vec<_>>();
+    let one_digits = iter::once(1)
+        .chain(iter::repeat(0u32))
+        .take(NUM_DIGITS)
+        .collect::<Vec<_>>();
+
+    let n: Var<_> = builder.eval(F::from_canonical_u32(32));
+    let sum = builder.eval_biguint(BigUint::new(sum_digits));
+    let min = builder.eval_biguint(BigUint::new(min_digits));
+    let val = builder.eval_biguint(BigUint::new(val_digits));
+    let one = builder.eval_biguint(BigUint::new(one_digits));
+
+    builder.range(RVar::zero(), n).for_each(|_, builder| {
+        let add = builder.add_256(&sum, &val);
+        let sub = builder.sub_256(&min, &val);
+
+        let and = builder.and_256(&add, &sub);
+        let xor = builder.xor_256(&add, &sub);
+        let or = builder.or_256(&and, &xor);
+
+        let sltu = builder.sltu_256(&add, &sub);
+        let slt = builder.slt_256(&add, &sub);
+
+        let shift_val = or.clone();
+        builder
+            .if_eq(sltu, F::from_canonical_u32(1))
+            .then(|builder| {
+                let srl = builder.srl_256(&shift_val, &one);
+                builder.assign(&shift_val, srl);
+            });
+        builder
+            .if_eq(slt, F::from_canonical_u32(0))
+            .then(|builder| {
+                let sra = builder.sra_256(&shift_val, &one);
+                builder.assign(&shift_val, sra);
+            });
+
+        let sll = builder.sll_256(&shift_val, &one);
+        let eq = builder.eq_256(&sll, &or);
+        builder.if_eq(eq, F::from_canonical_u32(0)).then(|builder| {
+            let temp = builder.add_256(&add, &one);
+            builder.assign(&add, temp);
+        });
+        builder.if_eq(eq, F::from_canonical_u32(1)).then(|builder| {
+            let temp = builder.sub_256(&sub, &one);
+            builder.assign(&sub, temp);
+        });
+
+        builder.assign(&sum, add);
+        builder.assign(&min, sub);
+    });
+
+    builder.halt();
+    builder.compile_isa_with_options(CompilerOptions {
+        word_size: 32,
+        ..Default::default()
+    })
+}
+
+fn bench_program_stark_for_test<SC: StarkGenericConfig>() -> StarkForTest<SC>
+where
+    Domain<SC>: PolynomialSpace<Val = BabyBear>,
+{
+    let program = bench_program();
+
+    let vm_config = VmConfig {
+        u256_arithmetic_enabled: true,
+        shift_256_enabled: true,
+        bigint_limb_size: 8,
+        ..Default::default()
+    };
+    gen_vm_program_stark_for_test(program, vec![], vm_config)
+}
+
+fn main() {
+    run_with_metric_collection("OUTPUT_PATH", || {
+        let vdata =
+            info_span!("Bench Program Inner", group = "bench_program_inner").in_scope(|| {
+                let program_stark = bench_program_stark_for_test();
+                program_stark
+                    .run_simple_test(&BabyBearPoseidon2Engine::new(
+                        standard_fri_params_with_100_bits_conjectured_security(4),
+                    ))
+                    .unwrap()
+            });
+
+        let compiler_options = CompilerOptions {
+            enable_cycle_tracker: true,
+            ..Default::default()
+        };
+        let vdata = info_span!("Inner Verifier", group = "inner_verifier").in_scope(|| {
+            let (program, witness_stream) = build_verification_program(vdata, compiler_options);
+            let inner_verifier_stf = gen_vm_program_stark_for_test(
+                program,
+                witness_stream,
+                VmConfig {
+                    num_public_values: 4,
+                    ..Default::default()
+                },
+            );
+            inner_verifier_stf
+                .run_simple_test(&BabyBearPoseidon2Engine::new(
+                    // log_blowup = 3 because of poseidon2 chip.
+                    standard_fri_params_with_100_bits_conjectured_security(3),
+                ))
+                .unwrap()
+        });
+
+        #[cfg(feature = "static-verifier")]
+        info_span!("Recursive Verify e2e", group = "recursive_verify_e2e").in_scope(|| {
+            let (program, witness_stream) = build_verification_program(vdata, compiler_options);
+            let outer_verifier_sft = gen_vm_program_stark_for_test(
+                program,
+                witness_stream,
+                VmConfig {
+                    num_public_values: 4,
+                    ..Default::default()
+                },
+            );
+            afs_recursion::halo2::testing_utils::run_evm_verifier_e2e_test(
+                &outer_verifier_sft,
+                // log_blowup = 3 because of poseidon2 chip.
+                Some(standard_fri_params_with_100_bits_conjectured_security(3)),
+            );
+        });
+    });
+}