diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/mod.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/mod.rs index e72253fd8..d57b2ef51 100644 --- a/stwo_cairo_prover/crates/prover/src/cairo_air/mod.rs +++ b/stwo_cairo_prover/crates/prover/src/cairo_air/mod.rs @@ -2,7 +2,9 @@ pub mod air; pub mod builtins_air; mod debug_tools; pub mod opcodes_air; +pub mod poseidon; pub mod preprocessed; +pub mod preprocessed_utils; pub mod range_checks_air; use air::{lookup_sum, CairoClaimGenerator, CairoComponents, CairoInteractionElements, CairoProof}; diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/const_columns.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/const_columns.rs new file mode 100644 index 000000000..2d2d38906 --- /dev/null +++ b/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/const_columns.rs @@ -0,0 +1,110 @@ +use std::array::from_fn; + +use prover_types::cpu::{Felt252Packed27, FELT252PACKED27_N_WORDS, M31}; +use prover_types::simd::N_LANES; +use stwo_prover::constraint_framework::preprocessed_columns::PreProcessedColumnId; +use stwo_prover::core::backend::simd::column::BaseColumn; +use stwo_prover::core::backend::simd::m31::{PackedM31, LOG_N_LANES}; +use stwo_prover::core::backend::simd::SimdBackend; +use stwo_prover::core::fields::m31::BaseField; +use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation}; +use stwo_prover::core::poly::BitReversedOrder; + +use super::super::preprocessed_utils::{pack_and_pad, table_id_to_col_id}; +use super::consts::POSEIDON_ROUND_KEYS; +use crate::cairo_air::poseidon::consts::{N_ROUNDS, N_WORDS}; +use crate::cairo_air::preprocessed::PreProcessedColumn; + +pub const POSEIDON_ROUND_KEYS_TABLE: &str = "poseidon_round_keys"; +const LOG_N_ROWS: u32 = 6; +const N_PACKED_ROWS: usize = (2_u32.pow(LOG_N_ROWS)) as usize / N_LANES; + +pub fn round_keys(round: usize) -> [Felt252Packed27; 3] { + POSEIDON_ROUND_KEYS[round].map(|k| Felt252Packed27 { limbs: k }) +} + +pub fn round_keys_m31(round: usize, col: usize) -> M31 { + assert!(col < N_WORDS); + assert!(round < N_ROUNDS); + + let felt252_index = col / FELT252PACKED27_N_WORDS; + let felt_index = col % FELT252PACKED27_N_WORDS; + round_keys(round)[felt252_index].get_m31(felt_index) +} + +#[derive(Debug)] +pub struct PoseidonRoundKeysColumn { + pub packed_keys: [PackedM31; N_PACKED_ROWS], + pub col: usize, +} + +impl PoseidonRoundKeysColumn { + pub fn new(col: usize) -> Self { + let packed_keys = from_fn(|i| pack_and_pad::(round_keys_m31, i, col)); + Self { packed_keys, col } + } + + pub fn packed_at(&self, vec_row: usize) -> PackedM31 { + self.packed_keys[vec_row] + } +} + +impl PreProcessedColumn for PoseidonRoundKeysColumn { + fn log_size(&self) -> u32 { + LOG_N_ROWS + } + + fn gen_column_simd(&self) -> CircleEvaluation { + CircleEvaluation::new( + CanonicCoset::new(self.log_size()).circle_domain(), + BaseColumn::from_simd( + (0..(1 << (self.log_size() - LOG_N_LANES))) + .map(|i| self.packed_at(i)) + .collect(), + ), + ) + } + + fn id(&self) -> PreProcessedColumnId { + PreProcessedColumnId { + id: table_id_to_col_id(POSEIDON_ROUND_KEYS_TABLE, self.col), + } + } +} + +#[cfg(test)] +mod tests { + use std::array::from_fn; + + use stwo_prover::core::backend::simd::m31::N_LANES; + + use super::*; + use crate::cairo_air::poseidon::consts::N_FELT252PACKED27; + + #[test] + fn test_packed_at_round_keys() { + for vec_row in 0..N_PACKED_ROWS { + for i in 0..N_FELT252PACKED27 { + let packed: [[M31; N_LANES]; FELT252PACKED27_N_WORDS] = from_fn(|c| { + PoseidonRoundKeysColumn::new((i * FELT252PACKED27_N_WORDS) + c) + .packed_at(vec_row) + .to_array() + }); + for row_in_packed in 0..N_LANES { + let felt_limbs: [M31; FELT252PACKED27_N_WORDS] = packed + .iter() + .map(|arr| arr[row_in_packed]) + .collect::>() + .try_into() + .unwrap(); + let row = (vec_row * N_LANES) + row_in_packed; + if row < N_ROUNDS { + assert_eq!(Felt252Packed27::from_limbs(&felt_limbs), round_keys(row)[i]); + } else { + assert_eq!(Felt252Packed27::from_limbs(&felt_limbs), round_keys(0)[i]); + } + } + } + } + } +} diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/consts.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/consts.rs new file mode 100644 index 000000000..6a987fb55 --- /dev/null +++ b/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/consts.rs @@ -0,0 +1,688 @@ +use prover_types::cpu::FELT252PACKED27_N_WORDS; + +pub(super) const N_ROUNDS: usize = 35; +pub(super) const N_FELT252PACKED27: usize = 3; +pub(super) const N_WORDS: usize = FELT252PACKED27_N_WORDS * N_FELT252PACKED27; +pub(super) const POSEIDON_ROUND_KEYS: [[[u64; 4]; N_FELT252PACKED27]; N_ROUNDS] = [ + [ + [ + 9808894619969057997, + 2962375666393338310, + 17382841788414994265, + 443257643709112289, + ], + [ + 12537484503666775718, + 3256805997184644908, + 6617722259049010207, + 543112534054733059, + ], + [ + 1454046077829682943, + 14331133962181073949, + 2327346812919484995, + 379005027604567203, + ], + ], + [ + [ + 316912300309518807, + 9546737057323600779, + 4990939959663297477, + 409555158710929193, + ], + [ + 14375050875883784322, + 3258765518491372314, + 6123276414968091301, + 564945574506516589, + ], + [ + 16399159056375946558, + 12401617009203210820, + 11251954111719545950, + 433429710746163456, + ], + ], + [ + [ + 7006359817426385501, + 17203056170488800107, + 10266463410669146573, + 302632258003414824, + ], + [ + 14549101708442237159, + 5447808788302094550, + 5985460154360671870, + 377474158904626280, + ], + [ + 2016536305309843132, + 7819086366821881070, + 6549900492473011498, + 375041666190951811, + ], + ], + [ + [ + 10958522786453251491, + 9697564334799485296, + 12061515884908749995, + 382992757552468126, + ], + [ + 3123128533353263104, + 1101320594306927398, + 12277506650622088974, + 151394834303922635, + ], + [ + 10558209943415701402, + 3761550961988184469, + 3770582263098070207, + 337917216919135628, + ], + ], + [ + [ + 11059652905750625380, + 13475195141561210865, + 13294395003503408798, + 543485850395037306, + ], + [ + 2130930448912281502, + 11333634387982439184, + 8850610548639699306, + 457475955817445493, + ], + [ + 13832015370839427617, + 3536623570151876469, + 6528270901734940966, + 1727329127918258, + ], + ], + [ + [ + 8115407971155582787, + 9978560128345434391, + 5056408649803520810, + 262112615523232333, + ], + [ + 15132830981034848104, + 221062278661831986, + 642558393488344280, + 294435853161867218, + ], + [ + 17720328485765873457, + 906259302840864515, + 9886887042042513701, + 91167054851387838, + ], + ], + [ + [ + 17686069520976319126, + 357140690021361429, + 4698816318705416205, + 393981709058899502, + ], + [ + 11422699654326280778, + 16059236267229938280, + 14304086719964370791, + 408074902120160445, + ], + [ + 10098039853740591407, + 18346213706869023683, + 9856189649941491293, + 184406899276982606, + ], + ], + [ + [ + 4111625807920905640, + 13925198121954558587, + 12310073145155562618, + 235927056615592132, + ], + [ + 6384587362686501122, + 8879249956632890389, + 16548116661510213280, + 336779148206005613, + ], + [ + 16031192111936359076, + 5992351855224207619, + 12627781605286612627, + 62096344547068532, + ], + ], + [ + [ + 7132509749599922320, + 6162523224461213171, + 8812310904075333603, + 485641259025949528, + ], + [ + 9608543551775247566, + 5196481376567266799, + 4241060526105290574, + 127878632248644222, + ], + [ + 5678830138638110916, + 17150803417208083936, + 3818159621611526901, + 334934582699708306, + ], + ], + [ + [ + 2411984091808734148, + 4676885686514770974, + 16024545775274701230, + 35144855618233700, + ], + [ + 3764367625527066390, + 16203185340231970163, + 3091249709657140605, + 246768161927392025, + ], + [ + 14145295330358598665, + 8189373517343951529, + 9185965492588609992, + 84513112494072013, + ], + ], + [ + [ + 11207056279882910057, + 8343632314543039012, + 11825707279115803854, + 290028553136950168, + ], + [ + 8324430733814960957, + 1295757100659713237, + 13793768882092694703, + 170786252004278897, + ], + [ + 3470195965229650799, + 7014342387715054610, + 8068066979310954905, + 29052085287755578, + ], + ], + [ + [ + 16182857524744847037, + 13508798581767021717, + 11609421482229546503, + 516546051298708222, + ], + [ + 15033308840249833218, + 2943419432621774691, + 12721848833217608341, + 495114643959326355, + ], + [ + 15900588831390415422, + 13235287741266617180, + 199968676028566291, + 31464852864192972, + ], + ], + [ + [ + 1620930607287324279, + 5691881440575675201, + 14029655266804088527, + 85281494397439074, + ], + [ + 14358430905948328990, + 8174075507093900204, + 4259719420728355987, + 198633356278451635, + ], + [ + 9809328616759644792, + 9452995911559605327, + 14571337138054143278, + 30595350758349726, + ], + ], + [ + [ + 5052814846190951527, + 2564319269931470445, + 11667947324380057882, + 381968969372291382, + ], + [ + 10270535945835093416, + 7013539859298536233, + 12880625276280589921, + 423512085516371736, + ], + [ + 9247593602616013019, + 9800080834126373385, + 15154714092547675637, + 85949681409890944, + ], + ], + [ + [ + 16691816165224508898, + 7632563883998222450, + 7283702476287088318, + 122927053010244988, + ], + [ + 5281932702559195717, + 3912411525754476767, + 2751980448518808692, + 449246095011271218, + ], + [ + 2154533012881281233, + 12108066475824676498, + 3101185982842383519, + 23082295823839220, + ], + ], + [ + [ + 7605923575436520679, + 17775553940505278137, + 12354955681295648587, + 509503557137574051, + ], + [ + 6859516056437622059, + 15185460371714151768, + 11379739398280558941, + 467020453759984910, + ], + [ + 16906035870412618946, + 2048289172670790831, + 3913835398798558993, + 247123888097172708, + ], + ], + [ + [ + 14861675641163720274, + 13490102368184285242, + 7347027430097237399, + 12293023119986689, + ], + [ + 17015804524472763158, + 2030415039026408622, + 17809691364575575612, + 373903064080638948, + ], + [ + 8022416475434186219, + 17815483025186149958, + 17841645611508634712, + 214671237987350594, + ], + ], + [ + [ + 12836252995927977384, + 12965348847163767059, + 16404178258733598267, + 90570121994582570, + ], + [ + 3307613700375182919, + 6181136657427428089, + 13131983874186228376, + 111501226499533359, + ], + [ + 17976156798761747365, + 5762323702974965097, + 2597451573586851781, + 505236697901381580, + ], + ], + [ + [ + 10381335841935625777, + 14975760611930379479, + 14435427058050060920, + 398310795157470149, + ], + [ + 7159397558615159963, + 7188734421404393410, + 1719787959693584840, + 314760383409924924, + ], + [ + 2143638283124439842, + 13645456387251540292, + 13644498560249152495, + 40707004462247922, + ], + ], + [ + [ + 4076286098075094248, + 6047377110922245529, + 10310252161423143831, + 4203916682452144, + ], + [ + 16974446701450911254, + 13817121135003214195, + 2110477587661043379, + 408404362623939579, + ], + [ + 2551943238291520253, + 7863747909853014700, + 10038172819555036178, + 498557445795487158, + ], + ], + [ + [ + 2539457924220004741, + 6798810574668326903, + 734801439896130781, + 197318323104987578, + ], + [ + 5630865015115736926, + 6924395279250128121, + 6087898613499446423, + 97920604124542022, + ], + [ + 8238671520399847831, + 17200586436505834896, + 18050188643002787777, + 522418299476559161, + ], + ], + [ + [ + 16671769425050870867, + 9818859487908268561, + 14628982416326270968, + 105495391240150891, + ], + [ + 15264270982597349937, + 3214172504508351054, + 96620451664846254, + 107082329324411929, + ], + [ + 14217429574808978483, + 128541115086728122, + 9630653827036234478, + 337586787343095831, + ], + ], + [ + [ + 8212059728564478324, + 2347043101088709486, + 6567058597747348925, + 136303161555124818, + ], + [ + 9215571201366006957, + 18390624930749960250, + 12318590206736769157, + 94289926799047171, + ], + [ + 2681199449952507734, + 2490827916210922767, + 17337862272405306868, + 167761531076143152, + ], + ], + [ + [ + 12798085655163047736, + 13696387070792973059, + 5787352356986496426, + 499982807322000917, + ], + [ + 11741293120501172298, + 2334843635281516844, + 168280946537445205, + 83199885793358504, + ], + [ + 12631247450401696760, + 1333500347809883553, + 7960218164236031817, + 545118190714259783, + ], + ], + [ + [ + 10647924774016649377, + 7324763009135962770, + 16081867801897155361, + 428325268027940951, + ], + [ + 9916218426490841056, + 1030092679665695813, + 1263314736050787503, + 189382575143994932, + ], + [ + 13675857909985031082, + 12446306110735075056, + 5033197056310105549, + 437182141684164304, + ], + ], + [ + [ + 16017091423340680337, + 16608552316205419099, + 8224269358142049653, + 421369926480120239, + ], + [ + 9595998208446022457, + 16134937223929573489, + 14473485045201252070, + 492654444776713988, + ], + [ + 5219139941060837063, + 15853778231351784921, + 17809532253019979818, + 417855011179639500, + ], + ], + [ + [ + 10403233932030073012, + 16730279770404598818, + 5644362027553147371, + 71735369062479113, + ], + [ + 9631935352839043273, + 17054291492818491992, + 17392960852478338690, + 460451922811752815, + ], + [ + 8477235527901859940, + 9814894372396933310, + 4799990685272759189, + 400282443455003622, + ], + ], + [ + [ + 1576843956219139158, + 702430080765383770, + 17965431827539789120, + 396822748194700979, + ], + [ + 6534711024274141327, + 7130143603302325050, + 16132773823416316657, + 291127484699867608, + ], + [ + 4689783592352337689, + 13240213583589885272, + 9538724043191226625, + 53501278744087361, + ], + ], + [ + [ + 5303664022074446245, + 11845428113003974855, + 12846417839211241712, + 261206892901600457, + ], + [ + 7213274845438153874, + 3182604331988734929, + 3825798403644353861, + 498424027422114180, + ], + [ + 16374080072108357107, + 12602249521880983826, + 16789038177262469212, + 434844701481670312, + ], + ], + [ + [ + 5424818379886973048, + 12246700522389374047, + 16105275706111921973, + 171092744151981801, + ], + [ + 6597251438734065845, + 2269299153703490182, + 1681550773894047556, + 497581200702741603, + ], + [ + 448123443688089620, + 2093069703231428951, + 10690354368775868807, + 47970718237487928, + ], + ], + [ + [ + 5663590943682895948, + 1657980836047728369, + 14473564132866859515, + 63033233261729097, + ], + [ + 10253436943656002480, + 496762026162837186, + 4416882861358244138, + 410250456421992694, + ], + [ + 13092039789419326359, + 12701778245175598192, + 10053832990213334033, + 99263799527290004, + ], + ], + [ + [ + 17612850421305358241, + 8270527177451683327, + 3824004781155827930, + 23420660055616514, + ], + [ + 14792590694225191985, + 10340526750113228103, + 13907692663639317222, + 419027522602786902, + ], + [ + 6537420399978421861, + 13000095247594509242, + 4689550804574686720, + 216508944422249833, + ], + ], + [ + [ + 5284804903132760421, + 6021193890942533180, + 12919709475177128988, + 388552658857864838, + ], + [ + 14487504387541221402, + 8671715521049733970, + 11505630672145478718, + 340569955046273223, + ], + [ + 15605227016987613715, + 4467780628446399859, + 1916547247173479880, + 360408062797054196, + ], + ], + [ + [ + 2318724554222443904, + 12462857660673509117, + 1043912215626002694, + 444903370426104614, + ], + [ + 13771333397279872933, + 8504629457688506196, + 17402104297977249580, + 365482958936297625, + ], + [ + 9663191740751744091, + 17588211649412389869, + 8849849772044756264, + 441288247586026977, + ], + ], + [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], +]; diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/mod.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/mod.rs new file mode 100644 index 000000000..7353f7179 --- /dev/null +++ b/stwo_cairo_prover/crates/prover/src/cairo_air/poseidon/mod.rs @@ -0,0 +1,2 @@ +pub mod const_columns; +pub mod consts; diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs index c73ea9b6e..a62ed906b 100644 --- a/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs +++ b/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs @@ -1,6 +1,7 @@ use std::simd::{u32x16, Simd}; use itertools::Itertools; +use prover_types::cpu::FELT252PACKED27_N_WORDS; use prover_types::simd::LOG_N_LANES; use stwo_prover::constraint_framework::preprocessed_columns::PreProcessedColumnId; use stwo_prover::core::backend::simd::column::BaseColumn; @@ -11,6 +12,8 @@ use stwo_prover::core::fields::m31::{BaseField, M31}; use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation}; use stwo_prover::core::poly::BitReversedOrder; +use super::poseidon::const_columns::PoseidonRoundKeysColumn; +use super::preprocessed_utils::table_id_to_col_id; use super::LOG_MAX_ROWS; use crate::components::range_check_vector::SIMD_ENUMERATION_0; @@ -26,6 +29,7 @@ pub trait PreProcessedColumn { pub struct PreProcessedTrace { seq_columns: Vec, bitwise_xor_columns: Vec, + poseidon_round_keys_columns: Vec, } impl PreProcessedTrace { #[allow(clippy::new_without_default)] @@ -34,9 +38,14 @@ impl PreProcessedTrace { let bitwise_xor_columns = (0..3) .map(move |col_index| BitwiseXor::new(XOR_N_BITS, col_index)) .collect_vec(); + let poseidon_round_keys_columns = (0..FELT252PACKED27_N_WORDS * 3) + .map(PoseidonRoundKeysColumn::new) + .collect_vec(); + Self { seq_columns, bitwise_xor_columns, + poseidon_round_keys_columns, } } @@ -52,6 +61,11 @@ impl PreProcessedTrace { .iter() .map(|c| c as &dyn PreProcessedColumn), ); + columns.extend( + self.poseidon_round_keys_columns + .iter() + .map(|c| c as &dyn PreProcessedColumn), + ); // Sort columns by descending log size. columns @@ -157,7 +171,7 @@ impl PreProcessedColumn for BitwiseXor { fn id(&self) -> PreProcessedColumnId { PreProcessedColumnId { - id: format!("bitwise_xor_{}_{}", self.n_bits, self.col_index), + id: table_id_to_col_id(&format!("bitwise_xor_{}", self.n_bits), self.col_index), } } } diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed_utils.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed_utils.rs new file mode 100644 index 000000000..cc29ff147 --- /dev/null +++ b/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed_utils.rs @@ -0,0 +1,23 @@ +use std::array::from_fn; + +use prover_types::simd::N_LANES; +use stwo_prover::core::backend::simd::m31::PackedM31; +use stwo_prover::core::fields::m31::M31; + +pub fn table_id_to_col_id(table_id: &str, col_index: usize) -> String { + format!("{}_{}", table_id, col_index) +} + +// Packs the values of column and rows into a +// PackedM31. Pads all row below R with the first row. Uses the function to get the value +// in a given row and column. +pub fn pack_and_pad(get_m31: F, vec_row: usize, col: usize) -> PackedM31 +where + F: Fn(usize, usize) -> M31, +{ + let first_row = get_m31(0, col); + PackedM31::from_array(from_fn(|i| { + let row = (vec_row * N_LANES) + i; + (row < R).then(|| get_m31(row, col)).unwrap_or(first_row) + })) +} diff --git a/stwo_cairo_prover/crates/prover_types/src/cpu.rs b/stwo_cairo_prover/crates/prover_types/src/cpu.rs index c43b1669e..0ebd549be 100644 --- a/stwo_cairo_prover/crates/prover_types/src/cpu.rs +++ b/stwo_cairo_prover/crates/prover_types/src/cpu.rs @@ -486,9 +486,9 @@ impl Felt252 { // Convert between Felt252 and FieldElement for performing field operations. // Note that FieldElements are in Montgomery form, and for efficiency and simplicity, we skip the -// conversion in both direction. We thus have to compensate with extra factors when performing +// conversion in both directions. We thus have to compensate with extra factors when performing // multiplication and division. We also ensure the FieldElement is reduced modulo P, by adding zero -// on initial conversion. This results in subtracting P once if needed, which suffices as as the +// to the initial conversion. This results in subtracting P once if needed, which suffices as the // limbs of a Felt252 are smaller than 2**252 < 2*P. impl From for FieldElement { fn from(n: Felt252) -> FieldElement {