From 3ab4809cead66d76d318dca14a964069cdf8c6a7 Mon Sep 17 00:00:00 2001 From: siq1 Date: Tue, 8 Oct 2024 18:15:02 +0000 Subject: [PATCH 1/6] support dynamic array in rust circuit definition --- expander_compiler/src/frontend/circuit.rs | 22 ++ expander_compiler/tests/keccak_gf2_vec.rs | 287 ++++++++++++++++++++++ 2 files changed, 309 insertions(+) create mode 100644 expander_compiler/tests/keccak_gf2_vec.rs diff --git a/expander_compiler/src/frontend/circuit.rs b/expander_compiler/src/frontend/circuit.rs index 49cebc6..6f9c65d 100644 --- a/expander_compiler/src/frontend/circuit.rs +++ b/expander_compiler/src/frontend/circuit.rs @@ -12,6 +12,10 @@ macro_rules! declare_circuit_field_type { [$crate::frontend::internal::declare_circuit_field_type!(@type $elem); $n] }; + (@type [$elem:tt]) => { + Vec<$crate::frontend::internal::declare_circuit_field_type!(@type $elem)> + }; + (@type $other:ty) => { $other }; @@ -33,6 +37,12 @@ macro_rules! declare_circuit_dump_into { } }; + ($field_value:expr, @type [$elem:tt], $vars:expr, $public_vars:expr) => { + for _x in $field_value.iter() { + $crate::frontend::internal::declare_circuit_dump_into!(_x, @type $elem, $vars, $public_vars); + } + }; + ($field_value:expr, @type $other:ty, $vars:expr, $public_vars:expr) => { }; } @@ -53,6 +63,12 @@ macro_rules! declare_circuit_load_from { } }; + ($field_value:expr, @type [$elem:tt], $vars:expr, $public_vars:expr) => { + for _x in $field_value.iter_mut() { + $crate::frontend::internal::declare_circuit_load_from!(_x, @type $elem, $vars, $public_vars); + } + }; + ($field_value:expr, @type $other:ty, $vars:expr, $public_vars:expr) => { }; } @@ -71,6 +87,12 @@ macro_rules! declare_circuit_num_vars { $crate::frontend::internal::declare_circuit_num_vars!($field_value[0], @type $elem, $cnt_sec, $cnt_pub, $array_cnt * $n); }; + ($field_value:expr, @type [$elem:tt], $cnt_sec:expr, $cnt_pub:expr, $array_cnt:expr) => { + for _x in $field_value.iter() { + $crate::frontend::internal::declare_circuit_num_vars!(_x, @type $elem, $cnt_sec, $cnt_pub, $array_cnt); + } + }; + ($field_value:expr, @type $other:ty, $cnt_sec:expr, $cnt_pub:expr, $array_cnt:expr) => { }; } diff --git a/expander_compiler/tests/keccak_gf2_vec.rs b/expander_compiler/tests/keccak_gf2_vec.rs new file mode 100644 index 0000000..af8acd3 --- /dev/null +++ b/expander_compiler/tests/keccak_gf2_vec.rs @@ -0,0 +1,287 @@ +use expander_compiler::frontend::*; +use rand::{thread_rng, Rng}; +use tiny_keccak::Hasher; + +const N_HASHES: usize = 4; + +fn rc() -> Vec { + vec![ + 0x0000000000000001, + 0x0000000000008082, + 0x800000000000808A, + 0x8000000080008000, + 0x000000000000808B, + 0x0000000080000001, + 0x8000000080008081, + 0x8000000000008009, + 0x000000000000008A, + 0x0000000000000088, + 0x0000000080008009, + 0x000000008000000A, + 0x000000008000808B, + 0x800000000000008B, + 0x8000000000008089, + 0x8000000000008003, + 0x8000000000008002, + 0x8000000000000080, + 0x000000000000800A, + 0x800000008000000A, + 0x8000000080008081, + 0x8000000000008080, + 0x0000000080000001, + 0x8000000080008008, + ] +} + +fn xor_in( + api: &mut API, + mut s: Vec>, + buf: Vec>, +) -> Vec> { + for y in 0..5 { + for x in 0..5 { + if x + 5 * y < buf.len() { + s[5 * x + y] = xor(api, s[5 * x + y].clone(), buf[x + 5 * y].clone()) + } + } + } + s +} + +fn keccak_f(api: &mut API, mut a: Vec>) -> Vec> { + let mut b = vec![vec![api.constant(0); 64]; 25]; + let mut c = vec![vec![api.constant(0); 64]; 5]; + let mut d = vec![vec![api.constant(0); 64]; 5]; + let mut da = vec![vec![api.constant(0); 64]; 5]; + let rc = rc(); + + for i in 0..24 { + for j in 0..5 { + let t1 = xor(api, a[j * 5 + 1].clone(), a[j * 5 + 2].clone()); + let t2 = xor(api, a[j * 5 + 3].clone(), a[j * 5 + 4].clone()); + c[j] = xor(api, t1, t2); + } + + for j in 0..5 { + d[j] = xor( + api, + c[(j + 4) % 5].clone(), + rotate_left::(&c[(j + 1) % 5], 1), + ); + da[j] = xor( + api, + a[((j + 4) % 5) * 5].clone(), + rotate_left::(&a[((j + 1) % 5) * 5], 1), + ); + } + + for j in 0..25 { + let tmp = xor(api, da[j / 5].clone(), a[j].clone()); + a[j] = xor(api, tmp, d[j / 5].clone()); + } + + /*Rho and pi steps*/ + b[0] = a[0].clone(); + + b[8] = rotate_left::(&a[1], 36); + b[11] = rotate_left::(&a[2], 3); + b[19] = rotate_left::(&a[3], 41); + b[22] = rotate_left::(&a[4], 18); + + b[2] = rotate_left::(&a[5], 1); + b[5] = rotate_left::(&a[6], 44); + b[13] = rotate_left::(&a[7], 10); + b[16] = rotate_left::(&a[8], 45); + b[24] = rotate_left::(&a[9], 2); + + b[4] = rotate_left::(&a[10], 62); + b[7] = rotate_left::(&a[11], 6); + b[10] = rotate_left::(&a[12], 43); + b[18] = rotate_left::(&a[13], 15); + b[21] = rotate_left::(&a[14], 61); + + b[1] = rotate_left::(&a[15], 28); + b[9] = rotate_left::(&a[16], 55); + b[12] = rotate_left::(&a[17], 25); + b[15] = rotate_left::(&a[18], 21); + b[23] = rotate_left::(&a[19], 56); + + b[3] = rotate_left::(&a[20], 27); + b[6] = rotate_left::(&a[21], 20); + b[14] = rotate_left::(&a[22], 39); + b[17] = rotate_left::(&a[23], 8); + b[20] = rotate_left::(&a[24], 14); + + /*Xi state*/ + + for j in 0..25 { + let t = not(api, b[(j + 5) % 25].clone()); + let t = and(api, t, b[(j + 10) % 25].clone()); + a[j] = xor(api, b[j].clone(), t); + } + + /*Last step*/ + + for j in 0..64 { + if rc[i] >> j & 1 == 1 { + a[0][j] = api.sub(1, a[0][j]); + } + } + } + + a +} + +fn xor(api: &mut API, a: Vec, b: Vec) -> Vec { + let nbits = a.len(); + let mut bits_res = vec![api.constant(0); nbits]; + for i in 0..nbits { + bits_res[i] = api.add(a[i].clone(), b[i].clone()); + } + bits_res +} + +fn and(api: &mut API, a: Vec, b: Vec) -> Vec { + let nbits = a.len(); + let mut bits_res = vec![api.constant(0); nbits]; + for i in 0..nbits { + bits_res[i] = api.mul(a[i].clone(), b[i].clone()); + } + bits_res +} + +fn not(api: &mut API, a: Vec) -> Vec { + let mut bits_res = vec![api.constant(0); a.len()]; + for i in 0..a.len() { + bits_res[i] = api.sub(1, a[i].clone()); + } + bits_res +} + +fn rotate_left(bits: &Vec, k: usize) -> Vec { + let n = bits.len(); + let s = k & (n - 1); + let mut new_bits = bits[(n - s) as usize..].to_vec(); + new_bits.append(&mut bits[0..(n - s) as usize].to_vec()); + new_bits +} + +fn copy_out_unaligned(s: Vec>, rate: usize, output_len: usize) -> Vec { + let mut out = vec![]; + let w = 8; + let mut b = 0; + while b < output_len { + for y in 0..5 { + for x in 0..5 { + if x + 5 * y < rate / w && b < output_len { + out.append(&mut s[5 * x + y].clone()); + b += 8; + } + } + } + } + out +} + +declare_circuit!(Keccak256Circuit { + p: [[Variable]], + out: [[PublicVariable]], +}); + +fn compute_keccak(api: &mut API, p: &Vec) -> Vec { + let mut ss = vec![vec![api.constant(0); 64]; 25]; + let mut new_p = p.clone(); + let mut append_data = vec![0; 136 - 64]; + append_data[0] = 1; + append_data[135 - 64] = 0x80; + for i in 0..136 - 64 { + for j in 0..8 { + new_p.push(api.constant(((append_data[i] >> j) & 1) as u32)); + } + } + let mut p = vec![vec![api.constant(0); 64]; 17]; + for i in 0..17 { + for j in 0..64 { + p[i][j] = new_p[i * 64 + j].clone(); + } + } + ss = xor_in(api, ss, p); + ss = keccak_f(api, ss); + copy_out_unaligned(ss, 136, 32) +} + +impl Define for Keccak256Circuit { + fn define(&self, api: &mut API) { + for i in 0..N_HASHES { + let out = api.memorized_simple_call(compute_keccak, &self.p[i].to_vec()); + for j in 0..256 { + api.assert_is_equal(out[j].clone(), self.out[i][j].clone()); + } + } + } +} + +#[test] +fn keccak_gf2_vec() { + let mut circuit = Keccak256Circuit::::default(); + circuit.p = vec![vec![Variable::default(); 64 * 8]; N_HASHES]; + circuit.out = vec![vec![Variable::default(); 32 * 8]; N_HASHES]; + + let compile_result = compile(&circuit).unwrap(); + let CompileResult { + witness_solver, + layered_circuit, + } = compile_result; + + let mut assignment = Keccak256Circuit::::default(); + assignment.p = vec![vec![GF2::from(0); 64 * 8]; N_HASHES]; + assignment.out = vec![vec![GF2::from(0); 32 * 8]; N_HASHES]; + for k in 0..N_HASHES { + let mut data = vec![0u8; 64]; + for i in 0..64 { + data[i] = thread_rng().gen(); + } + let mut hash = tiny_keccak::Keccak::v256(); + hash.update(&data); + let mut output = [0u8; 32]; + hash.finalize(&mut output); + for i in 0..64 { + for j in 0..8 { + assignment.p[k][i * 8 + j] = ((data[i] >> j) as u32 & 1).into(); + } + } + for i in 0..32 { + for j in 0..8 { + assignment.out[k][i * 8 + j] = ((output[i] >> j) as u32 & 1).into(); + } + } + } + let witness = witness_solver.solve_witness(&assignment).unwrap(); + let res = layered_circuit.run(&witness); + assert_eq!(res, vec![true]); + println!("test 1 passed"); + + for k in 0..N_HASHES { + assignment.p[k][0] = assignment.p[k][0] - GF2::from(1); + } + let witness = witness_solver.solve_witness(&assignment).unwrap(); + let res = layered_circuit.run(&witness); + assert_eq!(res, vec![false]); + println!("test 2 passed"); + + let mut assignments = Vec::new(); + for _ in 0..16 { + for k in 0..N_HASHES { + assignment.p[k][0] = assignment.p[k][0] - GF2::from(1); + } + assignments.push(assignment.clone()); + } + let witness = witness_solver.solve_witnesses(&assignments).unwrap(); + let res = layered_circuit.run(&witness); + let mut expected_res = vec![false; 16]; + for i in 0..8 { + expected_res[i * 2] = true; + } + assert_eq!(res, expected_res); + println!("test 3 passed"); +} From 5141bc953a4c34448240386b33bc8dde484f3bc1 Mon Sep 17 00:00:00 2001 From: siq1 Date: Tue, 8 Oct 2024 18:21:21 +0000 Subject: [PATCH 2/6] minor changes --- expander_compiler/src/layering/mod.rs | 2 +- expander_compiler/tests/keccak_gf2_full.rs | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/expander_compiler/src/layering/mod.rs b/expander_compiler/src/layering/mod.rs index a1b39b9..c9f4cf7 100644 --- a/expander_compiler/src/layering/mod.rs +++ b/expander_compiler/src/layering/mod.rs @@ -7,7 +7,7 @@ use crate::{ mod compile; mod input; -pub mod ir_split; // TODO +pub mod ir_split; mod layer_layout; mod wire; diff --git a/expander_compiler/tests/keccak_gf2_full.rs b/expander_compiler/tests/keccak_gf2_full.rs index d05cf8a..cab1416 100644 --- a/expander_compiler/tests/keccak_gf2_full.rs +++ b/expander_compiler/tests/keccak_gf2_full.rs @@ -185,7 +185,7 @@ fn copy_out_unaligned(s: Vec>, rate: usize, output_len: usize) -> declare_circuit!(Keccak256Circuit { p: [[Variable; 64 * 8]; N_HASHES], - out: [[Variable; 256]; N_HASHES], // TODO: use public inputs + out: [[PublicVariable; 256]; N_HASHES], }); fn compute_keccak(api: &mut API, p: &Vec) -> Vec { @@ -290,6 +290,7 @@ fn keccak_gf2_full() { ); let (simd_input, simd_public_input) = witness.to_simd::(); + println!("{} {}", simd_input.len(), simd_public_input.len()); expander_circuit.layers[0].input_vals = simd_input; expander_circuit.public_input = simd_public_input.clone(); From e8b507f6131cf32d4e1ede89582293c652e6db2e Mon Sep 17 00:00:00 2001 From: siq1 Date: Tue, 8 Oct 2024 18:24:22 +0000 Subject: [PATCH 3/6] remove temporary fix for public input --- ecgo/builder/root.go | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/ecgo/builder/root.go b/ecgo/builder/root.go index b1ff6e6..8417c7d 100644 --- a/ecgo/builder/root.go +++ b/ecgo/builder/root.go @@ -44,10 +44,7 @@ func (r *Root) PublicVariable(f schema.LeafInfo) frontend.Variable { ExtraId: 2 + uint64(r.nbPublicInputs), }) r.nbPublicInputs++ - // Currently, new version of public input is not support by Expander - // So we use a hint to isolate it in witness solver - x, _ := r.NewHint(IdentityHint, 1, r.addVar()) - return x[0] + return r.addVar() } // SecretVariable creates a new secret variable for the circuit. From 159aa72f5a0c25280e610e8f5d0a749f2b4bd2bd Mon Sep 17 00:00:00 2001 From: zhenfei Date: Fri, 11 Oct 2024 12:57:27 -0400 Subject: [PATCH 4/6] update file names --- ecgo/examples/poseidon_m31/main.go | 4 ++-- expander_compiler/tests/keccak_gf2.rs | 6 ++--- expander_compiler/tests/keccak_m31_bn254.rs | 25 ++++++++++++++++----- 3 files changed, 24 insertions(+), 11 deletions(-) diff --git a/ecgo/examples/poseidon_m31/main.go b/ecgo/examples/poseidon_m31/main.go index 978c3bd..c292e6b 100644 --- a/ecgo/examples/poseidon_m31/main.go +++ b/ecgo/examples/poseidon_m31/main.go @@ -72,7 +72,7 @@ func M31CircuitBuild() { layered_circuit := circuit.GetLayeredCircuit() // circuit.GetCircuitIr().Print() - err = os.WriteFile("circuit.txt", layered_circuit.Serialize(), 0o644) + err = os.WriteFile("poseidon_120_circuit_m31.txt", layered_circuit.Serialize(), 0o644) if err != nil { panic(err) } @@ -81,7 +81,7 @@ func M31CircuitBuild() { if err != nil { panic(err) } - err = os.WriteFile("witness.txt", witness.Serialize(), 0o644) + err = os.WriteFile("poseidon_120_witness_m31.txt", witness.Serialize(), 0o644) if err != nil { panic(err) } diff --git a/expander_compiler/tests/keccak_gf2.rs b/expander_compiler/tests/keccak_gf2.rs index 76e58aa..2d3b442 100644 --- a/expander_compiler/tests/keccak_gf2.rs +++ b/expander_compiler/tests/keccak_gf2.rs @@ -288,15 +288,15 @@ fn keccak_gf2_main() { .solve_witnesses(&assignments_correct) .unwrap(); - let file = std::fs::File::create("circuit.txt").unwrap(); + let file = std::fs::File::create("circuit_gf2.txt").unwrap(); let writer = std::io::BufWriter::new(file); layered_circuit.serialize_into(writer).unwrap(); - let file = std::fs::File::create("witness.txt").unwrap(); + let file = std::fs::File::create("witness_gf2.txt").unwrap(); let writer = std::io::BufWriter::new(file); witness.serialize_into(writer).unwrap(); - let file = std::fs::File::create("witness_solver.txt").unwrap(); + let file = std::fs::File::create("witness_gf2_solver.txt").unwrap(); let writer = std::io::BufWriter::new(file); witness_solver.serialize_into(writer).unwrap(); diff --git a/expander_compiler/tests/keccak_m31_bn254.rs b/expander_compiler/tests/keccak_m31_bn254.rs index a93544f..a541074 100644 --- a/expander_compiler/tests/keccak_m31_bn254.rs +++ b/expander_compiler/tests/keccak_m31_bn254.rs @@ -284,7 +284,7 @@ impl Define for Keccak256Circuit { } } -fn keccak_big_field() { +fn keccak_big_field(field_name: &str) { let compile_result: CompileResult = compile(&Keccak256Circuit::default()).unwrap(); let CompileResult { witness_solver, @@ -355,15 +355,28 @@ fn keccak_big_field() { .solve_witnesses(&assignments_correct) .unwrap(); - let file = std::fs::File::create("circuit.txt").unwrap(); + let file = match field_name { + "m31" => std::fs::File::create("circuit_m31.txt").unwrap(), + "bn254" => std::fs::File::create("circuit_bn254.txt").unwrap(), + _ => panic!("unknown field"), + }; let writer = std::io::BufWriter::new(file); layered_circuit.serialize_into(writer).unwrap(); - let file = std::fs::File::create("witness.txt").unwrap(); + let file = match field_name { + "m31" => std::fs::File::create("witness_m31.txt").unwrap(), + "bn254" => std::fs::File::create("witness_bn254.txt").unwrap(), + _ => panic!("unknown field"), + }; + let writer = std::io::BufWriter::new(file); witness.serialize_into(writer).unwrap(); - let file = std::fs::File::create("witness_solver.txt").unwrap(); + let file = match field_name { + "m31" => std::fs::File::create("witness_m31_solver.txt").unwrap(), + "bn254" => std::fs::File::create("witness_bn254_solver.txt").unwrap(), + _ => panic!("unknown field"), + }; let writer = std::io::BufWriter::new(file); witness_solver.serialize_into(writer).unwrap(); @@ -372,10 +385,10 @@ fn keccak_big_field() { #[test] fn keccak_m31_test() { - keccak_big_field::(); + keccak_big_field::("m31"); } #[test] fn keccak_bn254_test() { - keccak_big_field::(); + keccak_big_field::("bn254"); } From e25cb5079901bf1c8e85317af41178b643ede9ff Mon Sep 17 00:00:00 2001 From: zhenfei Date: Fri, 11 Oct 2024 13:03:42 -0400 Subject: [PATCH 5/6] fix ci --- .github/workflows/ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7c52f36..d81d06c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,6 +18,8 @@ jobs: - uses: Swatinem/rust-cache@v2 with: workspaces: "expander_compiler -> expander_compiler/target" + # The prefix cache key, this can be changed to start a new cache manually. + prefix-key: "mpi-v5.0.5" # update me if brew formula changes to a new version - if: matrix.os == 'macos-latest' run: brew install openmpi - if: matrix.os == 'ubuntu-latest' From f5fc08568ad98c729f630fa68fb24e34e329296e Mon Sep 17 00:00:00 2001 From: zhenfei Date: Fri, 11 Oct 2024 13:07:52 -0400 Subject: [PATCH 6/6] fix ci --- .github/workflows/ci.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d81d06c..e0bb096 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -68,6 +68,10 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 + with: + workspaces: "expander_compiler -> expander_compiler/target" + # The prefix cache key, this can be changed to start a new cache manually. + prefix-key: "mpi-v5.0.5" # update me if brew formula changes to a new version - if: matrix.os == 'macos-latest' run: brew install openmpi - if: matrix.os == 'ubuntu-latest' @@ -81,6 +85,10 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 + with: + workspaces: "expander_compiler -> expander_compiler/target" + # The prefix cache key, this can be changed to start a new cache manually. + prefix-key: "mpi-v5.0.5" # update me if brew formula changes to a new version - run: RUSTFLAGS="-C target-cpu=native -C target-feature=+avx512f" cargo test test-go: