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.

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

GKR backend for LogUp-GKR #296

Merged
merged 47 commits into from
Aug 30, 2024
Merged
Show file tree
Hide file tree
Changes from 34 commits
Commits
Show all changes
47 commits
Select commit Hold shift + click to select a range
a11700d
feat: math utilities needed for sum-check protocol
Al-Kindi-0 Aug 6, 2024
5e00b98
feat: add sum-check prover and verifier
Al-Kindi-0 Aug 6, 2024
798deb1
tests: add sanity tests for utils
Al-Kindi-0 Aug 6, 2024
c1b67d6
feat: use SmallVec
Al-Kindi-0 Aug 7, 2024
657a24b
feat: add remaining functions for sum-check verifier
Al-Kindi-0 Aug 9, 2024
1fe8cc2
chore: move prover into sub-mod
Al-Kindi-0 Aug 9, 2024
486ea9d
chore: remove utils mod
Al-Kindi-0 Aug 9, 2024
68e6097
chore: move logup evaluator trait to separate file
Al-Kindi-0 Aug 9, 2024
09a231d
feat: add multi-threading support and simplify input sum-check
Al-Kindi-0 Aug 15, 2024
d1e899d
feat: add benchmarks and address feedback
Al-Kindi-0 Aug 19, 2024
f44ba93
feat: address feedback and add benchmarks
Al-Kindi-0 Aug 19, 2024
29d5241
feat: add GKR backend for LogUp-GKR
Al-Kindi-0 Aug 9, 2024
37afc96
chore: remove old way of handling Lagrange kernel
Al-Kindi-0 Aug 12, 2024
59ce28c
chore: remove GkrVerifier
Al-Kindi-0 Aug 12, 2024
f0fdfb2
chore: correct header
Al-Kindi-0 Aug 12, 2024
3583e7c
feat: simplify sum-check for input layer
Al-Kindi-0 Aug 14, 2024
d77985d
feat: add mult-threading to sum-checks and multi-linears
Al-Kindi-0 Aug 14, 2024
0933af9
chore: remove num_rounds from sum-check
Al-Kindi-0 Aug 14, 2024
e0650c1
chore: unify serial and concurrent EQ impls
Al-Kindi-0 Aug 14, 2024
12998fb
fix: formatting
Al-Kindi-0 Aug 14, 2024
8ba5ce9
fix: concurrent feature flag
Al-Kindi-0 Aug 14, 2024
f1511e2
feat: make query a mut ref in build_query
Al-Kindi-0 Aug 19, 2024
f4f2d6c
chore: rebase on sum-check
Al-Kindi-0 Aug 20, 2024
51a33ec
chore: update var names and simplify some structs
Al-Kindi-0 Aug 20, 2024
2ac0a9e
wip: avoid extra allocation when projecting mle
Al-Kindi-0 Aug 20, 2024
8e2f7e6
chore: rebase
Al-Kindi-0 Aug 20, 2024
7de00ae
chore: fix rebase against facebook:logup-gkr
Al-Kindi-0 Aug 21, 2024
7caf36f
chore: address feedback
Al-Kindi-0 Aug 21, 2024
f974dd1
chore: remove TODO
Al-Kindi-0 Aug 21, 2024
452559e
fix: tex formatting issue
Al-Kindi-0 Aug 22, 2024
1ae173c
chore: add dummy GkrLogUp evaluator
Al-Kindi-0 Aug 23, 2024
7e3f9e4
chore: fix clippy warnings
Al-Kindi-0 Aug 23, 2024
7a40333
chore: address feedback
Al-Kindi-0 Aug 26, 2024
626a062
chore: loosen the trait bound
Al-Kindi-0 Aug 26, 2024
22a71bd
chore: address feedback 1
Al-Kindi-0 Aug 27, 2024
011480c
chore: address feedback 2
Al-Kindi-0 Aug 27, 2024
70a2c6b
feat: refactor gkr_data generation logic
Al-Kindi-0 Aug 27, 2024
78d430f
fix: doc tests and clippy
Al-Kindi-0 Aug 27, 2024
0895d3a
chore: address feedback
Al-Kindi-0 Aug 28, 2024
24e9924
chore: simplify further
Al-Kindi-0 Aug 28, 2024
798cc66
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
bbdd692
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
3595121
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
b7eef26
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
2be88cd
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
aa106d4
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
d0893bf
fix: clippy too long first paragraph
Al-Kindi-0 Aug 29, 2024
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
1 change: 1 addition & 0 deletions .cargo/katex-header.html
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
renderMathInElement(document.body, {
fleqn: false,
macros: {
"\\B": "\\mathbb{B}",
"\\F": "\\mathbb{F}",
"\\G": "\\mathbb{G}",
"\\O": "\\mathcal{O}",
Expand Down
133 changes: 70 additions & 63 deletions air/src/air/aux.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,11 @@
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use alloc::{string::ToString, vec::Vec};
use alloc::vec::Vec;

use crypto::{ElementHasher, RandomCoin, RandomCoinError};
use math::FieldElement;
use utils::Deserializable;

use super::lagrange::LagrangeKernelRandElements;
use super::{lagrange::LagrangeKernelRandElements, LogUpGkrOracle};

/// Holds the randomly generated elements necessary to build the auxiliary trace.
///
Expand All @@ -18,12 +16,12 @@ use super::lagrange::LagrangeKernelRandElements;
/// - the ones needed to build the "s" auxiliary column (when using GKR to accelerate LogUp),
/// - the ones needed to build all the other auxiliary columns
#[derive(Debug, Clone)]
pub struct AuxRandElements<E> {
pub struct AuxRandElements<E: FieldElement> {
rand_elements: Vec<E>,
gkr: Option<GkrRandElements<E>>,
gkr: Option<GkrData<E>>,
}

impl<E> AuxRandElements<E> {
impl<E: FieldElement> AuxRandElements<E> {
/// Creates a new [`AuxRandElements`], where the auxiliary trace doesn't contain a Lagrange
/// kernel column.
pub fn new(rand_elements: Vec<E>) -> Self {
Expand All @@ -32,7 +30,7 @@ impl<E> AuxRandElements<E> {

/// Creates a new [`AuxRandElements`], where the auxiliary trace contains columns needed when
/// using GKR to accelerate LogUp (i.e. a Lagrange kernel column and the "s" column).
pub fn new_with_gkr(rand_elements: Vec<E>, gkr: GkrRandElements<E>) -> Self {
pub fn new_with_gkr(rand_elements: Vec<E>, gkr: GkrData<E>) -> Self {
Self { rand_elements, gkr: Some(gkr) }
}

Expand All @@ -43,7 +41,7 @@ impl<E> AuxRandElements<E> {

/// Returns the random elements needed to build the Lagrange kernel column.
pub fn lagrange(&self) -> Option<&LagrangeKernelRandElements<E>> {
self.gkr.as_ref().map(|gkr| &gkr.lagrange)
self.gkr.as_ref().map(|gkr| &gkr.lagrange_kernel_eval_point)
}

/// Returns the random values used to linearly combine the openings returned from the GKR proof.
Expand All @@ -52,83 +50,92 @@ impl<E> AuxRandElements<E> {
pub fn gkr_openings_combining_randomness(&self) -> Option<&[E]> {
self.gkr.as_ref().map(|gkr| gkr.openings_combining_randomness.as_ref())
}

pub fn gkr_data(&self) -> Option<GkrData<E>> {
self.gkr.clone()
}
irakliyk marked this conversation as resolved.
Show resolved Hide resolved
}

/// Holds all the random elements needed when using GKR to accelerate LogUp.
/// Holds all the data needed when using LogUp-GKR in order to build the two extra auxiliary columns
/// for running the univariate IOP for multi-linear evaluations of [1].
irakliyk marked this conversation as resolved.
Show resolved Hide resolved
///
/// This consists of two sets of random values:
/// 1. The Lagrange kernel random elements (expanded on in [`LagrangeKernelRandElements`]), and
/// This consists of:
/// 1. The Lagrange kernel random elements (expanded on in [`LagrangeKernelRandElements`]). These
/// make up the evaluation point of the multi-linear extension polynomials underlying the oracles
/// in point 4 below.
/// 2. The "openings combining randomness".
/// 3. The openings of the multi-linear extension polynomials of the main trace columns involved
/// in LogUp.
/// 4. A description of the each of the oracles involved in LogUp.
///
/// After verifying the LogUp-GKR circuit, the verifier is left with unproven claims provided
/// by the prover about the evaluations of the MLEs of the main trace columns at the evaluation
/// point defining the Lagrange kernel. Those claims are (linearly) batched into one using the
/// openings combining randomness and checked against the batched oracles using univariate IOP
/// for multi-linear evaluations of [1].
///
/// After the verifying the LogUp-GKR circuit, the verifier is left with unproven claims provided
/// nondeterministically by the prover about the evaluations of the MLE of the main trace columns at
/// the Lagrange kernel random elements. Those claims are (linearly) combined into one using the
/// openings combining randomness.
/// [1]: https://eprint.iacr.org/2023/1284
#[derive(Clone, Debug)]
pub struct GkrRandElements<E> {
lagrange: LagrangeKernelRandElements<E>,
openings_combining_randomness: Vec<E>,
pub struct GkrData<E: FieldElement> {
irakliyk marked this conversation as resolved.
Show resolved Hide resolved
pub lagrange_kernel_eval_point: LagrangeKernelRandElements<E>,
pub openings_combining_randomness: Vec<E>,
pub openings: Vec<E>,
pub oracles: Vec<LogUpGkrOracle<E::BaseField>>,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it is necessary to store the oracles here

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is needed for the s-column, which is done in another PR building on this one

}

impl<E> GkrRandElements<E> {
/// Constructs a new [`GkrRandElements`] from [`LagrangeKernelRandElements`], and the openings
/// combining randomness.
impl<E: FieldElement> GkrData<E> {
/// Constructs a new [`GkrData`] from [`LagrangeKernelRandElements`], the openings combining
/// randomness and the LogUp-GKR oracles.
///
/// See [`GkrRandElements`] for a more detailed description.
/// See [`GkrData`] for a more detailed description.
pub fn new(
lagrange: LagrangeKernelRandElements<E>,
lagrange_kernel_eval_point: LagrangeKernelRandElements<E>,
openings_combining_randomness: Vec<E>,
openings: Vec<E>,
oracles: Vec<LogUpGkrOracle<E::BaseField>>,
) -> Self {
Self { lagrange, openings_combining_randomness }
Self {
lagrange_kernel_eval_point,
openings_combining_randomness,
openings,
oracles,
}
}

/// Returns the random elements needed to build the Lagrange kernel column.
pub fn lagrange_kernel_rand_elements(&self) -> &LagrangeKernelRandElements<E> {
&self.lagrange
&self.lagrange_kernel_eval_point
}

/// Returns the random values used to linearly combine the openings returned from the GKR proof.
pub fn openings_combining_randomness(&self) -> &[E] {
&self.openings_combining_randomness
}
}

/// A trait for verifying a GKR proof.
///
/// Specifically, the use case in mind is proving the constraints of a LogUp bus using GKR, as
/// described in [Improving logarithmic derivative lookups using
/// GKR](https://eprint.iacr.org/2023/1284.pdf).
pub trait GkrVerifier {
/// The GKR proof.
type GkrProof: Deserializable;
/// The error that can occur during GKR proof verification.
type Error: ToString;

/// Verifies the GKR proof, and returns the random elements that were used in building
/// the Lagrange kernel auxiliary column.
fn verify<E, Hasher>(
&self,
gkr_proof: Self::GkrProof,
public_coin: &mut impl RandomCoin<BaseField = E::BaseField, Hasher = Hasher>,
) -> Result<GkrRandElements<E>, Self::Error>
where
E: FieldElement,
Hasher: ElementHasher<BaseField = E::BaseField>;
}
pub fn openings(&self) -> &[E] {
&self.openings
}

pub fn oracles(&self) -> &[LogUpGkrOracle<E::BaseField>] {
&self.oracles
}

pub fn compute_batched_claim(&self) -> E {
self.openings[0]
+ self
.openings
.iter()
.skip(1)
.zip(self.openings_combining_randomness.iter())
.fold(E::ZERO, |acc, (a, b)| acc + *a * *b)
}

impl GkrVerifier for () {
type GkrProof = ();
type Error = RandomCoinError;

fn verify<E, Hasher>(
&self,
_gkr_proof: Self::GkrProof,
_public_coin: &mut impl RandomCoin<BaseField = E::BaseField, Hasher = Hasher>,
) -> Result<GkrRandElements<E>, Self::Error>
where
E: FieldElement,
Hasher: ElementHasher<BaseField = E::BaseField>,
{
Ok(GkrRandElements::new(LagrangeKernelRandElements::default(), Vec::new()))
pub fn compute_batched_query(&self, query: &[E::BaseField]) -> E {
E::from(query[0])
+ query
.iter()
.skip(1)
.zip(self.openings_combining_randomness.iter())
.fold(E::ZERO, |acc, (a, b)| acc + b.mul_base(*a))
}
}
24 changes: 20 additions & 4 deletions air/src/air/coefficients.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,19 @@ use math::FieldElement;
///
/// The coefficients are separated into two lists: one for transition constraints and another one
/// for boundary constraints. This separation is done for convenience only.
///
/// In addition to the above, and when LogUp-GKR is enabled, there are two extra sets of
/// constraint composition coefficients that are used, namely for:
///
/// 1. Lagrange kernel constraints, which include both transition and boundary constraints.
/// 2. S-column constraint, which is used in implementing the cohomological sum-check argument
/// of https://eprint.iacr.org/2021/930
#[derive(Debug, Clone)]
pub struct ConstraintCompositionCoefficients<E: FieldElement> {
pub transition: Vec<E>,
pub boundary: Vec<E>,
pub lagrange: Option<LagrangeConstraintsCompositionCoefficients<E>>,
pub s_col: Option<E>,
irakliyk marked this conversation as resolved.
Show resolved Hide resolved
}

/// Stores the constraint composition coefficients for the Lagrange kernel transition and boundary
Expand Down Expand Up @@ -83,8 +91,9 @@ pub struct LagrangeConstraintsCompositionCoefficients<E: FieldElement> {
/// negligible increase in soundness error. The formula for the updated error can be found in
/// Theorem 8 of https://eprint.iacr.org/2022/1216.
///
/// In the case when the trace polynomials contain a trace polynomial corresponding to a Lagrange
/// kernel column, the above expression of $Y(x)$ includes the additional term given by
/// In the case when LogUp-GKR is enabled, the trace polynomials contain an additional trace
/// polynomial corresponding to a Lagrange kernel column and the above expression of $Y(x)$
/// includes the additional term given by
///
/// $$
/// \gamma \cdot \frac{T_l(x) - p_S(x)}{Z_S(x)}
Expand All @@ -99,8 +108,13 @@ pub struct LagrangeConstraintsCompositionCoefficients<E: FieldElement> {
/// 4. $p_S(X)$ is the polynomial of minimal degree interpolating the set ${(a, T_l(a)): a \in S}$.
/// 5. $Z_S(X)$ is the polynomial of minimal degree vanishing over the set $S$.
///
/// Note that, if a Lagrange kernel trace polynomial is present, then $\rho^{+}$ from above should
/// be updated to be $\rho^{+} := \frac{\kappa + log_2(\nu) + 1}{\nu}$.
/// Note that when LogUp-GKR is enabled, we also have to take into account an additional column,
/// called s-column throughout, used in implementing the univariate IOP for multi-linear evaluation.
/// This means that we need and additional random value, in addition to $\gamma$ above, when
/// LogUp-GKR is enabled.
///
/// Note that, when LogUp-GKR is enabled, $\rho^{+}$ from above should be updated to be
/// $\rho^{+} := \frac{\kappa + log_2(\nu) + 1}{\nu}$.
#[derive(Debug, Clone)]
pub struct DeepCompositionCoefficients<E: FieldElement> {
/// Trace polynomial composition coefficients $\alpha_i$.
Expand All @@ -109,4 +123,6 @@ pub struct DeepCompositionCoefficients<E: FieldElement> {
pub constraints: Vec<E>,
/// Lagrange kernel trace polynomial composition coefficient $\gamma$.
pub lagrange: Option<E>,
/// S-column trace polynomial composition coefficient.
pub s_col: Option<E>,
}
55 changes: 34 additions & 21 deletions air/src/air/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,11 @@ pub struct AirContext<B: StarkField> {
pub(super) aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
pub(super) num_main_assertions: usize,
pub(super) num_aux_assertions: usize,
pub(super) lagrange_kernel_aux_column_idx: Option<usize>,
pub(super) ce_blowup_factor: usize,
pub(super) trace_domain_generator: B,
pub(super) lde_domain_generator: B,
pub(super) num_transition_exemptions: usize,
pub(super) logup_gkr: bool,
}

impl<B: StarkField> AirContext<B> {
Expand Down Expand Up @@ -62,7 +62,6 @@ impl<B: StarkField> AirContext<B> {
Vec::new(),
num_assertions,
0,
None,
options,
)
}
Expand Down Expand Up @@ -95,7 +94,6 @@ impl<B: StarkField> AirContext<B> {
aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
num_main_assertions: usize,
num_aux_assertions: usize,
lagrange_kernel_aux_column_idx: Option<usize>,
options: ProofOptions,
) -> Self {
assert!(
Expand All @@ -104,11 +102,11 @@ impl<B: StarkField> AirContext<B> {
);
assert!(num_main_assertions > 0, "at least one assertion must be specified");

if trace_info.is_multi_segment() {
if trace_info.is_multi_segment() && !trace_info.logup_gkr_enabled() {
assert!(
!aux_transition_constraint_degrees.is_empty(),
"at least one transition constraint degree must be specified for the auxiliary trace segment"
);
!aux_transition_constraint_degrees.is_empty(),
"at least one transition constraint degree must be specified for the auxiliary trace segment"
);
assert!(
num_aux_assertions > 0,
"at least one assertion must be specified against the auxiliary trace segment"
Expand All @@ -124,15 +122,6 @@ impl<B: StarkField> AirContext<B> {
);
}

// validate Lagrange kernel aux column, if any
if let Some(lagrange_kernel_aux_column_idx) = lagrange_kernel_aux_column_idx {
assert!(
lagrange_kernel_aux_column_idx == trace_info.get_aux_segment_width() - 1,
"Lagrange kernel column should be the last column of the auxiliary trace: index={}, but aux trace width is {}",
lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width()
);
}

// determine minimum blowup factor needed to evaluate transition constraints by taking
// the blowup factor of the highest degree constraint
let mut ce_blowup_factor = 0;
Expand Down Expand Up @@ -165,14 +154,34 @@ impl<B: StarkField> AirContext<B> {
aux_transition_constraint_degrees,
num_main_assertions,
num_aux_assertions,
lagrange_kernel_aux_column_idx,
ce_blowup_factor,
trace_domain_generator: B::get_root_of_unity(trace_length.ilog2()),
lde_domain_generator: B::get_root_of_unity(lde_domain_size.ilog2()),
num_transition_exemptions: 1,
logup_gkr: false,
}
}

pub fn with_logup_gkr(
trace_info: TraceInfo,
main_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
num_main_assertions: usize,
num_aux_assertions: usize,
options: ProofOptions,
) -> Self {
let mut air_context = Self::new_multi_segment(
trace_info,
main_transition_constraint_degrees,
aux_transition_constraint_degrees,
num_main_assertions,
num_aux_assertions,
options,
);
air_context.logup_gkr = true;
air_context
}

// PUBLIC ACCESSORS
// --------------------------------------------------------------------------------------------

Expand Down Expand Up @@ -232,12 +241,16 @@ impl<B: StarkField> AirContext<B> {

/// Returns the index of the auxiliary column which implements the Lagrange kernel, if any
pub fn lagrange_kernel_aux_column_idx(&self) -> Option<usize> {
self.lagrange_kernel_aux_column_idx
if self.logup_gkr_enabled() {
Some(self.trace_info().aux_segment_width() - 1)
} else {
None
}
}

/// Returns true if the auxiliary trace segment contains a Lagrange kernel column
pub fn has_lagrange_kernel_aux_column(&self) -> bool {
self.lagrange_kernel_aux_column_idx().is_some()
/// Returns true if LogUp-GKR is enabled.
pub fn logup_gkr_enabled(&self) -> bool {
self.logup_gkr
}

/// Returns the total number of assertions defined for a computation, excluding the Lagrange
Expand Down
Loading
Loading