Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Single-row Optimizations on top of KeccakRound gate #1301

Closed
wants to merge 12 commits into from
Closed

Single-row Optimizations on top of KeccakRound gate #1301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
b643edc
optimizing 244 columns, and reducing by ~300 constraints the KeccakRound
querolita Nov 1, 2023
299d675
optimize 1bit rotation witness layout
querolita Nov 1, 2023
02c0127
update witness code to ~300 less columns
querolita Nov 1, 2023
af29b1a
update keccak round pattern columns
querolita Nov 1, 2023
a48ce4f
update value of keccak cols to be 2065
querolita Nov 1, 2023
e42456b
remove bound_e
querolita Nov 1, 2023
72ba91e
update order of generics
querolita Nov 1, 2023
6128078
remove lookup selectors of keccak sponge inside ocamltypes as they ar…
querolita Nov 1, 2023
ff9405b
fmt
querolita Nov 1, 2023
461f9f8
Merge branch 'zkvm/keccak/framework' into zkvm/keccak/optimized
querolita Nov 1, 2023
fde4035
Merge branch 'zkvm/keccak/framework' into zkvm/keccak/optimized
querolita Nov 2, 2023
415854d
Merge branch 'zkvm/keccak/framework' into zkvm/keccak/optimized
querolita Nov 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 21 additions & 38 deletions kimchi/src/circuits/lookup/lookups.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@ fn max_lookups_per_row(kinds: LookupPatterns) -> usize {
feature = "ocaml_types",
derive(ocaml::IntoValue, ocaml::FromValue, ocaml_gen::Struct)
)]
#[cfg_attr(feature = "wasm_types", wasm_bindgen::prelude::wasm_bindgen)]
pub struct LookupPatterns {
pub xor: bool,
pub lookup: bool,
Expand Down Expand Up @@ -148,7 +147,6 @@ impl LookupPatterns {
feature = "ocaml_types",
derive(ocaml::IntoValue, ocaml::FromValue, ocaml_gen::Struct)
)]
#[cfg_attr(feature = "wasm_types", wasm_bindgen::prelude::wasm_bindgen)]
pub struct LookupFeatures {
/// A single lookup constraint is a vector of lookup constraints to be applied at a row.
pub patterns: LookupPatterns,
Expand All @@ -174,7 +172,6 @@ impl LookupFeatures {

/// Describes the desired lookup configuration.
#[derive(Copy, Clone, Serialize, Deserialize, Debug)]
#[cfg_attr(feature = "wasm_types", wasm_bindgen::prelude::wasm_bindgen)]
pub struct LookupInfo {
/// The maximum length of an element of `kinds`. This can be computed from `kinds`.
pub max_per_row: usize,
Expand Down Expand Up @@ -415,7 +412,7 @@ impl LookupPattern {
/// Returns the maximum number of lookups per row that are used by the pattern.
pub fn max_lookups_per_row(&self) -> usize {
match self {
//LookupPattern::KeccakRound => 1760,
//LookupPattern::KeccakRound => 1620,
//LookupPattern::KeccakSponge => 800,
LookupPattern::Xor | LookupPattern::RangeCheck | LookupPattern::ForeignFieldMul => 4,
LookupPattern::Lookup => 3,
Expand Down Expand Up @@ -523,23 +520,23 @@ impl LookupPattern {
// Theta
for i in 0..20 {
// 2-column lookups
let shift0_c = curr_row(120 + i);
let dense_c = curr_row(200 + i);
let shift0_c = curr_row(100 + i);
let dense_c = curr_row(180 + i);
lookups.push(JointLookup {
table_id: LookupTableID::Constant(RESET_TABLE_ID),
entry: vec![l(dense_c), l(shift0_c)],
});
let dense_rot_c = curr_row(280 + i);
let expand_rot_c = curr_row(300 + i);
let dense_rot_c = curr_row(225 + i);
let expand_rot_c = curr_row(245 + i);
lookups.push(JointLookup {
table_id: LookupTableID::Constant(RESET_TABLE_ID),
entry: vec![l(dense_rot_c), l(expand_rot_c)],
});

// Second column lookups
let shift1_c = curr_row(140 + i);
let shift2_c = curr_row(160 + i);
let shift3_c = curr_row(180 + i);
let shift1_c = curr_row(120 + i);
let shift2_c = curr_row(140 + i);
let shift3_c = curr_row(160 + i);

lookups.push(JointLookup {
table_id: LookupTableID::Constant(SPARSE_TABLE_ID),
Expand All @@ -555,43 +552,34 @@ impl LookupPattern {
});

// First column lookups
let quotient_c = curr_row(220 + i);
let remainder_c = curr_row(240 + i);
let bound_c = curr_row(260 + i);
let remainder_c = curr_row(205 + i);

lookups.push(JointLookup {
table_id: LookupTableID::Constant(BITS16_TABLE_ID),
entry: vec![l(quotient_c)],
});
lookups.push(JointLookup {
table_id: LookupTableID::Constant(BITS16_TABLE_ID),
entry: vec![l(remainder_c)],
});
lookups.push(JointLookup {
table_id: LookupTableID::Constant(BITS16_TABLE_ID),
entry: vec![l(bound_c)],
});

}
// PiRho
for i in 0..100 {
// 2-column lookups
let shift0_e = curr_row(440 + i);
let dense_e = curr_row(840 + i);
let shift0_e = curr_row(265 + i);
let dense_e = curr_row(665 + i);
lookups.push(JointLookup {
table_id: LookupTableID::Constant(RESET_TABLE_ID),
entry: vec![l(dense_e), l(shift0_e)],
});
let dense_rot_e = curr_row(1240 + i);
let expand_rot_e = curr_row(1340 + i);
let dense_rot_e = curr_row(965 + i);
let expand_rot_e = curr_row(1065 + i);
lookups.push(JointLookup {
table_id: LookupTableID::Constant(SPARSE_TABLE_ID),
entry: vec![l(dense_rot_e), l(expand_rot_e)],
});

// Second column lookups
let shift1_e = curr_row(540 + i);
let shift2_e = curr_row(640 + i);
let shift3_e = curr_row(740 + i);
let shift1_e = curr_row(365 + i);
let shift2_e = curr_row(465 + i);
let shift3_e = curr_row(565 + i);

lookups.push(JointLookup {
table_id: LookupTableID::Constant(SPARSE_TABLE_ID),
Expand All @@ -607,9 +595,8 @@ impl LookupPattern {
});

// First column lookups
let quotient_e = curr_row(940 + i);
let remainder_e = curr_row(1040 + i);
let bound_e = curr_row(1140 + i);
let quotient_e = curr_row(765 + i);
let remainder_e = curr_row(865 + i);

lookups.push(JointLookup {
table_id: LookupTableID::Constant(BITS16_TABLE_ID),
Expand All @@ -619,16 +606,12 @@ impl LookupPattern {
table_id: LookupTableID::Constant(BITS16_TABLE_ID),
entry: vec![l(remainder_e)],
});
lookups.push(JointLookup {
table_id: LookupTableID::Constant(BITS16_TABLE_ID),
entry: vec![l(bound_e)],
});
}
// Chi
for i in 0..400 {
// Second column lookups
let shift_b: LocalPosition = curr_row(1540 + i);
let shift_sum = curr_row(1940 + i);
let shift_b: LocalPosition = curr_row(1165 + i);
let shift_sum = curr_row(1565 + i);

lookups.push(JointLookup {
table_id: LookupTableID::Constant(SPARSE_TABLE_ID),
Expand Down
129 changes: 60 additions & 69 deletions kimchi/src/circuits/polynomials/keccak/circuitgates.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@ use crate::{
auto_clone, auto_clone_array,
circuits::{
argument::{Argument, ArgumentEnv, ArgumentType},
expr::{constraints::ExprOps, Cache},
expr::{
constraints::{boolean, ExprOps},
Cache,
},
gate::GateType,
},
grid,
Expand Down Expand Up @@ -51,7 +54,7 @@ macro_rules! from_shifts {
}

//~
//~ | `KeccakRound` | [0...440) | [440...1540) | [1540...2344) |
//~ | `KeccakRound` | [0...265) | [265...1165) | [1165...1965) |
//~ | ------------- | --------- | ------------ | ------------- |
//~ | Curr | theta | pirho | chi |
//~
Expand All @@ -61,17 +64,17 @@ macro_rules! from_shifts {
//~
//~ -----------------------------------------------------------------------------------------------------------------------------------------------------------------------
//~
//~ | Columns | [0...100) | [100...120) | [120...200) | [200...220) | [220...240) | [240...260) | [260...280) | [280...300) | [300...320) | [320...340) | [340...440) |
//~ | -------- | --------- | ----------- | ----------- | ----------- | ----------- | ------------ | ----------- | ------------ | ------------ | ----------- | ----------- |
//~ | theta | state_a | state_c | shifts_c | dense_c | quotient_c | remainder_c | bound_c | dense_rot_c | expand_rot_c | state_d | state_e |
//~ | Columns | [0...100) | [100...180) | [180...200) | [200...205) | [205...225) | [225...245) | [245...265) |
//~ | -------- | --------- | ----------- | ----------- | ----------- | ------------ | ------------ | ------------ |
//~ | theta | state_a | shifts_c | dense_c | quotient_c | remainder_c | dense_rot_c | expand_rot_c |
//~
//~ | Columns | [440...840) | [840...940) | [940...1040) | [1040...1140) | [1140...1240) | [1240...1340) | [1340...1440) | [1440...1540) |
//~ | -------- | ----------- | ----------- | ------------ | ------------- | ------------- | ------------- | ------------- | ------------- |
//~ | pirho | shifts_e | dense_e | quotient_e | remainder_e | bound_e | dense_rot_e | expand_rot_e | state_b |
//~ | Columns | [265...665) | [665...765) | [765...865) | [865...965) | [965...1065) | [1065...1165) |
//~ | -------- | ----------- | ----------- | ------------ | ----------- | ------------ | ------------- |
//~ | pirho | shifts_e | dense_e | quotient_e | remainder_e | dense_rot_e | expand_rot_e |
//~
//~ | Columns | [1540...1940) | [1940...2340) | [2340...2344 |
//~ | -------- | ------------- | ------------- | ------------ |
//~ | chi | shifts_b | shifts_sum | f00 |
//~ | Columns | [1165...1565) | [1565...1965) |
//~ | -------- | ------------- | ------------- |
//~ | chi | shifts_b | shifts_sum |
//~
//~ | Columns | [0...4) | [4...100) |
//~ | -------- | ------- | --------- |
Expand All @@ -85,7 +88,7 @@ where
F: PrimeField,
{
const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::KeccakRound);
const CONSTRAINTS: u32 = 754;
const CONSTRAINTS: u32 = 389;

// Constraints for one round of the Keccak permutation function
fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
Expand All @@ -100,111 +103,99 @@ where
// LOAD STATES FROM WITNESS LAYOUT
// THETA
let state_a = grid!(100, env.witness_curr_chunk(0, 100));
let state_c = grid!(20, env.witness_curr_chunk(100, 120));
let shifts_c = grid!(80, env.witness_curr_chunk(120, 200));
let dense_c = grid!(20, env.witness_curr_chunk(200, 220));
let quotient_c = grid!(20, env.witness_curr_chunk(220, 240));
let remainder_c = grid!(20, env.witness_curr_chunk(240, 260));
let bound_c = grid!(20, env.witness_curr_chunk(260, 280));
let dense_rot_c = grid!(20, env.witness_curr_chunk(280, 300));
let expand_rot_c = grid!(20, env.witness_curr_chunk(300, 320));
let state_d = grid!(20, env.witness_curr_chunk(320, 340));
let state_e = grid!(100, env.witness_curr_chunk(340, 440));
let shifts_c = grid!(80, env.witness_curr_chunk(100, 180));
let dense_c = grid!(20, env.witness_curr_chunk(180, 200));
let quotient_c = grid!(5, env.witness_curr_chunk(200, 205));
let remainder_c = grid!(20, env.witness_curr_chunk(205, 225));
let dense_rot_c = grid!(20, env.witness_curr_chunk(225, 245));
let expand_rot_c = grid!(20, env.witness_curr_chunk(245, 265));
// PI-RHO
let shifts_e = grid!(400, env.witness_curr_chunk(440, 840));
let dense_e = grid!(100, env.witness_curr_chunk(840, 940));
let quotient_e = grid!(100, env.witness_curr_chunk(940, 1040));
let remainder_e = grid!(100, env.witness_curr_chunk(1040, 1140));
let bound_e = grid!(100, env.witness_curr_chunk(1140, 1240));
let dense_rot_e = grid!(100, env.witness_curr_chunk(1240, 1340));
let expand_rot_e = grid!(100, env.witness_curr_chunk(1340, 1440));
let state_b = grid!(100, env.witness_curr_chunk(1440, 1540));
let shifts_e = grid!(400, env.witness_curr_chunk(265, 665));
let dense_e = grid!(100, env.witness_curr_chunk(665, 765));
let quotient_e = grid!(100, env.witness_curr_chunk(765, 865));
let remainder_e = grid!(100, env.witness_curr_chunk(865, 965));
let dense_rot_e = grid!(100, env.witness_curr_chunk(965, 1065));
let expand_rot_e = grid!(100, env.witness_curr_chunk(1065, 1165));
// CHI
let shifts_b = grid!(400, env.witness_curr_chunk(1540, 1940));
let shifts_sum = grid!(400, env.witness_curr_chunk(1940, 2340));
let mut state_f: Vec<T> = env.witness_curr_chunk(2340, 2344);
let mut tail = env.witness_next_chunk(4, 100);
state_f.append(&mut tail);
let state_f = grid!(100, state_f);
let shifts_b = grid!(400, env.witness_curr_chunk(1165, 1565));
let shifts_sum = grid!(400, env.witness_curr_chunk(1565, 1965));
// IOTA
let mut state_g = env.witness_next_chunk(0, 4);
let mut tail = env.witness_next_chunk(4, 100);
state_g.append(&mut tail);
let state_g = grid!(100, state_g);
let state_g = grid!(100, env.witness_next_chunk(0, 100));

// STEP theta: 5 * ( 3 + 4 * (3 + 5 * 1) ) = 175 constraints
// Define vectors containing witness expressions which are not in the layout for efficiency
let mut state_c: Vec<Vec<T>> = vec![vec![T::zero(); QUARTERS]; DIM];
let mut state_d: Vec<Vec<T>> = vec![vec![T::zero(); QUARTERS]; DIM];
let mut state_e: Vec<Vec<Vec<T>>> = vec![vec![vec![T::zero(); QUARTERS]; DIM]; DIM];
let mut state_b: Vec<Vec<Vec<T>>> = vec![vec![vec![T::zero(); QUARTERS]; DIM]; DIM];
let mut state_f: Vec<Vec<Vec<T>>> = vec![vec![vec![T::zero(); QUARTERS]; DIM]; DIM];

// STEP theta: 5 * ( 3 + 4 * 1 ) = 35 constraints
for x in 0..DIM {
let word_c = from_quarters!(dense_c, x);
let quo_c = from_quarters!(quotient_c, x);
let rem_c = from_quarters!(remainder_c, x);
let bnd_c = from_quarters!(bound_c, x);
let rot_c = from_quarters!(dense_rot_c, x);

constraints
.push(word_c * T::two_pow(1) - (quo_c.clone() * T::two_pow(64) + rem_c.clone()));
constraints.push(rot_c - (quo_c.clone() + rem_c));
constraints.push(bnd_c - (quo_c + T::two_pow(64) - T::two_pow(1)));
.push(word_c * T::two_pow(1) - (quotient_c(x) * T::two_pow(64) + rem_c.clone()));
constraints.push(rot_c - (quotient_c(x) + rem_c));
constraints.push(boolean(&quotient_c(x)));

for q in 0..QUARTERS {
constraints.push(
state_c(x, q)
- (state_a(0, x, q)
+ state_a(1, x, q)
+ state_a(2, x, q)
+ state_a(3, x, q)
+ state_a(4, x, q)),
);
constraints.push(state_c(x, q) - from_shifts!(shifts_c, x, q));
constraints.push(
state_d(x, q)
- (shifts_c(0, (x + DIM - 1) % DIM, q) + expand_rot_c((x + 1) % DIM, q)),
);
state_c[x][q] = state_a(0, x, q)
+ state_a(1, x, q)
+ state_a(2, x, q)
+ state_a(3, x, q)
+ state_a(4, x, q);
constraints.push(state_c[x][q].clone() - from_shifts!(shifts_c, x, q));

state_d[x][q] =
shifts_c(0, (x + DIM - 1) % DIM, q) + expand_rot_c((x + 1) % DIM, q);

for y in 0..DIM {
constraints.push(state_e(y, x, q) - (state_a(y, x, q) + state_d(x, q)));
state_e[y][x][q] = state_a(y, x, q) + state_d[x][q].clone();
}
}
} // END theta

// STEP pirho: 5 * 5 * (3 + 4 * 2) = 275 constraints
// STEP pirho: 5 * 5 * (2 + 4 * 1) = 150 constraints
for (y, col) in OFF.iter().enumerate() {
for (x, off) in col.iter().enumerate() {
let word_e = from_quarters!(dense_e, y, x);
let quo_e = from_quarters!(quotient_e, y, x);
let rem_e = from_quarters!(remainder_e, y, x);
let bnd_e = from_quarters!(bound_e, y, x);
let rot_e = from_quarters!(dense_rot_e, y, x);

constraints.push(
word_e * T::two_pow(*off) - (quo_e.clone() * T::two_pow(64) + rem_e.clone()),
);
constraints.push(rot_e - (quo_e.clone() + rem_e));
constraints.push(bnd_e - (quo_e + T::two_pow(64) - T::two_pow(*off)));

for q in 0..QUARTERS {
constraints.push(state_e(y, x, q) - from_shifts!(shifts_e, y, x, q));
constraints.push(state_b((2 * x + 3 * y) % DIM, y, q) - expand_rot_e(y, x, q));
constraints.push(state_e[y][x][q].clone() - from_shifts!(shifts_e, y, x, q));
state_b[(2 * x + 3 * y) % DIM][y][q] = expand_rot_e(y, x, q);
}
}
} // END pirho

// STEP chi: 4 * 5 * 5 * 3 = 300 constraints
// STEP chi: 4 * 5 * 5 * 2 = 200 constraints
for q in 0..QUARTERS {
for x in 0..DIM {
for y in 0..DIM {
let not = T::literal(F::from(0x1111111111111111u64))
- shifts_b(0, y, (x + 1) % DIM, q);
let sum = not + shifts_b(0, y, (x + 2) % DIM, q);
let and = shifts_sum(1, y, x, q);
constraints.push(state_b(y, x, q) - from_shifts!(shifts_b, y, x, q));

constraints.push(state_b[y][x][q].clone() - from_shifts!(shifts_b, y, x, q));
constraints.push(sum - from_shifts!(shifts_sum, y, x, q));
constraints.push(state_f(y, x, q) - (shifts_b(0, y, x, q) + and));
state_f[y][x][q] = shifts_b(0, y, x, q) + and;
}
}
} // END chi

// STEP iota: 4 constraints
for (q, c) in rc.iter().enumerate() {
constraints.push(state_g(0, 0, q) - (state_f(0, 0, q) + c.clone()));
constraints.push(state_g(0, 0, q) - (state_f[0][0][q].clone() + c.clone()));
} // END iota

constraints
Expand Down
5 changes: 4 additions & 1 deletion kimchi/src/circuits/polynomials/keccak/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,16 @@ pub const QUARTERS: usize = 4;
pub const ROUNDS: usize = 24;
pub const RATE: usize = 1088 / 8;
pub const CAPACITY: usize = 512 / 8;
pub const KECCAK_COLS: usize = 2344;
pub const KECCAK_COLS: usize = 1965;

use crate::circuits::expr::constraints::ExprOps;
use ark_ff::PrimeField;

#[macro_export]
macro_rules! grid {
(5, $v:expr) => {{
|x: usize| $v[x].clone()
}};
(20, $v:expr) => {{
|x: usize, q: usize| $v[q + QUARTERS * x].clone()
}};
Expand Down
Loading
Loading