From 01a4251faa1aa0d5cb64879b063bb39cd72c6d7d Mon Sep 17 00:00:00 2001 From: Gali Michlevich Date: Thu, 9 Jan 2025 16:05:38 +0200 Subject: [PATCH] Add Structs using PreprocessedColumn Trait --- .../preprocessed_columns.rs | 39 +++++++++ crates/prover/src/examples/blake/mod.rs | 1 + .../examples/blake/preprocessed_columns.rs | 86 +++++++++++++++++++ crates/prover/src/examples/plonk/mod.rs | 17 ++++ .../src/examples/xor/gkr_lookups/mod.rs | 1 + .../xor/gkr_lookups/preprocessed_columns.rs | 49 +++++++++++ 6 files changed, 193 insertions(+) create mode 100644 crates/prover/src/examples/blake/preprocessed_columns.rs create mode 100644 crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs diff --git a/crates/prover/src/constraint_framework/preprocessed_columns.rs b/crates/prover/src/constraint_framework/preprocessed_columns.rs index a54ebc734..df45d7e15 100644 --- a/crates/prover/src/constraint_framework/preprocessed_columns.rs +++ b/crates/prover/src/constraint_framework/preprocessed_columns.rs @@ -3,6 +3,7 @@ use std::simd::Simd; use num_traits::{One, Zero}; use crate::core::backend::simd::m31::{PackedM31, N_LANES}; +use crate::core::backend::simd::SimdBackend; use crate::core::backend::{Backend, Col, Column}; use crate::core::fields::m31::{BaseField, M31}; use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; @@ -15,6 +16,44 @@ const SIMD_ENUMERATION_0: PackedM31 = unsafe { ])) }; +/// A column with `1` at the first position, and `0` elsewhere. +#[derive(Debug, Clone)] +pub struct IsFirst { + pub log_size: u32, +} +impl IsFirst { + pub const fn new(log_size: u32) -> Self { + Self { log_size } + } + + pub fn packed_at(&self, vec_row: usize) -> PackedM31 { + assert!(vec_row < (1 << self.log_size) / N_LANES); + if vec_row == 0 { + unsafe { + PackedM31::from_simd_unchecked(Simd::from_array(std::array::from_fn(|i| { + if i == 0 { + 1 + } else { + 0 + } + }))) + } + } else { + PackedM31::zero() + } + } + + pub fn gen_column_simd(&self) -> CircleEvaluation { + let mut col = Col::::zeros(1 << self.log_size); + col.set(0, BaseField::one()); + CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col) + } + + pub fn id(&self) -> String { + format!("preprocessed_is_first_{}", self.log_size).to_string() + } +} + // TODO(ilya): Where should this enum be placed? // TODO(Gali): Consider making it a trait, add documentation for the rest of the variants. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/crates/prover/src/examples/blake/mod.rs b/crates/prover/src/examples/blake/mod.rs index 76feb7f8b..0fd8af61d 100644 --- a/crates/prover/src/examples/blake/mod.rs +++ b/crates/prover/src/examples/blake/mod.rs @@ -16,6 +16,7 @@ use crate::core::fields::m31::BaseField; use crate::core::fields::FieldExpOps; mod air; +mod preprocessed_columns; mod round; mod scheduler; mod xor_table; diff --git a/crates/prover/src/examples/blake/preprocessed_columns.rs b/crates/prover/src/examples/blake/preprocessed_columns.rs new file mode 100644 index 000000000..07f225d66 --- /dev/null +++ b/crates/prover/src/examples/blake/preprocessed_columns.rs @@ -0,0 +1,86 @@ +use tracing::{span, Level}; + +use crate::constraint_framework::preprocessed_columns::gen_is_first; +use crate::core::backend::simd::column::BaseColumn; +use crate::core::backend::simd::SimdBackend; +use crate::core::fields::m31::BaseField; +use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; +use crate::core::poly::BitReversedOrder; +use crate::core::ColumnVec; + +// TODO(Gali): Add documentation and remove allow dead code. +#[allow(dead_code)] +#[derive(Debug)] +pub struct XorTable { + pub n_bits: u32, + pub n_expand_bits: u32, + pub index_in_table: usize, +} +impl XorTable { + // TODO(Gali): Remove allow dead code. + #[allow(dead_code)] + pub const fn new(n_bits: u32, n_expand_bits: u32, index_in_table: usize) -> Self { + Self { + n_bits, + n_expand_bits, + index_in_table, + } + } + + // TODO(Gali): Remove allow dead code. + #[allow(dead_code)] + pub fn id(&self) -> String { + format!( + "preprocessed_xor_table_{}_{}_{}", + self.n_bits, self.n_expand_bits, self.index_in_table + ) + } + + // TODO(Gali): Remove allow dead code. + #[allow(dead_code)] + pub const fn limb_bits(&self) -> u32 { + self.n_bits - self.n_expand_bits + } + + // TODO(Gali): Remove allow dead code. + #[allow(dead_code)] + pub const fn column_bits(&self) -> u32 { + 2 * self.limb_bits() + } + + /// Generates the Preprocessed trace for the xor table. + /// Returns the Preprocessed trace, the Preprocessed trace, and the claimed sum. + // TODO(Gali): Remove allow dead code. + #[allow(dead_code)] + #[allow(clippy::type_complexity)] + pub fn generate_constant_trace( + &self, + ) -> ColumnVec> { + let limb_bits = self.limb_bits(); + let _span = span!(Level::INFO, "Xor Preprocessed trace").entered(); + + // Generate the constant columns. In reality, these should be generated before the + // proof even began. + let a_col: BaseColumn = (0..(1 << self.column_bits())) + .map(|i| BaseField::from_u32_unchecked((i >> limb_bits) as u32)) + .collect(); + let b_col: BaseColumn = (0..(1 << self.column_bits())) + .map(|i| BaseField::from_u32_unchecked((i & ((1 << limb_bits) - 1)) as u32)) + .collect(); + let c_col: BaseColumn = (0..(1 << self.column_bits())) + .map(|i| { + BaseField::from_u32_unchecked( + ((i >> limb_bits) ^ (i & ((1 << limb_bits) - 1))) as u32, + ) + }) + .collect(); + + let mut constant_trace = [a_col, b_col, c_col] + .map(|x| { + CircleEvaluation::new(CanonicCoset::new(self.column_bits()).circle_domain(), x) + }) + .to_vec(); + constant_trace.push(gen_is_first(self.column_bits())); + constant_trace + } +} diff --git a/crates/prover/src/examples/plonk/mod.rs b/crates/prover/src/examples/plonk/mod.rs index a1e0362c9..a9665fe84 100644 --- a/crates/prover/src/examples/plonk/mod.rs +++ b/crates/prover/src/examples/plonk/mod.rs @@ -268,6 +268,23 @@ pub fn prove_fibonacci_plonk( (component, proof) } +/// Preprocessed columns for describing a plonk circuit. +/// Each plonk gate is described by input wires `a_wire` and `b_wire`, output wire `c_wire`, and +/// operation `op` in P. +#[derive(Debug)] +pub struct Plonk { + pub name: String, +} +impl Plonk { + pub const fn new(name: String) -> Self { + Self { name } + } + + pub fn id(&self) -> String { + format!("preprocessed_plonk_{}", self.name) + } +} + #[cfg(test)] mod tests { use std::env; diff --git a/crates/prover/src/examples/xor/gkr_lookups/mod.rs b/crates/prover/src/examples/xor/gkr_lookups/mod.rs index 6ee603eb0..7a6389a60 100644 --- a/crates/prover/src/examples/xor/gkr_lookups/mod.rs +++ b/crates/prover/src/examples/xor/gkr_lookups/mod.rs @@ -1,2 +1,3 @@ pub mod accumulation; pub mod mle_eval; +pub mod preprocessed_columns; diff --git a/crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs b/crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs new file mode 100644 index 000000000..537d57839 --- /dev/null +++ b/crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs @@ -0,0 +1,49 @@ +use num_traits::One; + +use crate::core::backend::simd::SimdBackend; +use crate::core::backend::{Col, Column}; +use crate::core::fields::m31::BaseField; +use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; +use crate::core::poly::BitReversedOrder; +use crate::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index}; + +/// A column with `1` at every `2^log_step` positions, `0` elsewhere, shifted by offset. +#[derive(Debug)] +pub struct IsStepWithOffset { + log_size: u32, + log_step: u32, + offset: usize, +} +impl IsStepWithOffset { + pub const fn new(log_size: u32, log_step: u32, offset: usize) -> Self { + Self { + log_size, + log_step, + offset, + } + } + + // TODO(andrew): Consider optimizing. Is a quotients of two coset_vanishing (use succinct rep + // for verifier). + // TODO(Gali): Remove allow dead code. + #[allow(dead_code)] + pub fn gen_column_simd(&self) -> CircleEvaluation { + let mut col = Col::::zeros(1 << self.log_size); + let size = 1 << self.log_size; + let step = 1 << self.log_step; + let step_offset = self.offset % step; + for i in (step_offset..size).step_by(step) { + let circle_domain_index = coset_index_to_circle_domain_index(i, self.log_size); + let circle_domain_index_bit_rev = bit_reverse_index(circle_domain_index, self.log_size); + col.set(circle_domain_index_bit_rev, BaseField::one()); + } + CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col) + } + + pub fn id(&self) -> String { + format!( + "preprocessed_is_step_with_offset_{}_{}_{}", + self.log_size, self.log_step, self.offset + ) + } +}