Skip to content

Commit

Permalink
refactor: Bypass quotient table
Browse files Browse the repository at this point in the history
  • Loading branch information
@aszepieniec
aszepieniec committed Mar 29, 2024
1 parent 006d795 commit ff30545
Show file tree
Hide file tree
Showing 2 changed files with 117 additions and 62 deletions.
13 changes: 4 additions & 9 deletions triton-vm/src/stark.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ use crate::proof_stream::ProofStream;
use crate::table::challenges::Challenges;
use crate::table::extension_table::Evaluable;
use crate::table::extension_table::Quotientable;
use crate::table::master_table::all_quotients;
use crate::table::master_table::all_quotients_combined;
use crate::table::master_table::interpolant_degree;
use crate::table::master_table::max_degree_with_origin;
use crate::table::master_table::MasterBaseTable;
Expand Down Expand Up @@ -180,7 +180,6 @@ impl Stark {
// Get the weights with which to compress the many quotients into one.
let quotient_combination_weights =
proof_stream.sample_scalars(MasterExtTable::NUM_CONSTRAINTS);
let quotient_combination_weights = Array1::from(quotient_combination_weights);
prof_stop!(maybe_profiler, "Fiat-Shamir");
prof_stop!(maybe_profiler, "ext tables");

Expand All @@ -196,20 +195,16 @@ impl Stark {
"compute and combine quotient codewords",
"CC"
);
let master_quotient_table = all_quotients(
let quotient_codeword = all_quotients_combined(
base_quotient_domain_codewords,
ext_quotient_domain_codewords,
master_base_table.trace_domain(),
quotient_domain,
&challenges,
&quotient_combination_weights,
maybe_profiler,
);
assert_eq!(
quotient_combination_weights.len(),
master_quotient_table.ncols()
);
let quotient_codeword =
Self::random_linear_sum(master_quotient_table.view(), quotient_combination_weights);
let quotient_codeword = Array1::from(quotient_codeword);
assert_eq!(quotient_domain.length, quotient_codeword.len());
prof_stop!(maybe_profiler, "compute and combine quotient codewords");

Expand Down
166 changes: 113 additions & 53 deletions triton-vm/src/table/master_table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::ops::MulAssign;
use std::ops::Range;

use itertools::Itertools;
use master_table::extension_table::Evaluable;
use ndarray::parallel::prelude::*;
use ndarray::prelude::*;
use ndarray::s;
Expand Down Expand Up @@ -1017,26 +1018,27 @@ pub fn terminal_quotient_zerofier_inverse(
BFieldElement::batch_inversion(zerofier_codeword).into()
}

/// Computes an array containing all quotients – the Master Quotient Table. Each column corresponds
/// to a different quotient. The quotients are ordered by category – initial, consistency,
/// Computes the quotient codeword, which the randomized linear combination of all individual
/// quotients.
///
/// About assigning weights to quotients: the quotients are ordered by category – initial, consistency,
/// transition, and then terminal. Within each category, the quotients follow the canonical order
/// of the tables. The last column holds the terminal quotient of the cross-table argument, which
/// is strictly speaking not a table.
/// The order of the quotients is not actually important. However, it must be consistent between
/// [prover] and [verifier].
///
/// The returned array is in row-major order.
///
/// [prover]: crate::stark::Stark::prove
/// [verifier]: crate::stark::Stark::verify
pub fn all_quotients(
pub fn all_quotients_combined(
quotient_domain_master_base_table: ArrayView2<BFieldElement>,
quotient_domain_master_ext_table: ArrayView2<XFieldElement>,
trace_domain: ArithmeticDomain,
quotient_domain: ArithmeticDomain,
challenges: &Challenges,
quotient_weights: &[XFieldElement],
maybe_profiler: &mut Option<TritonProfiler>,
) -> Array2<XFieldElement> {
) -> Vec<XFieldElement> {
assert_eq!(
quotient_domain.length,
quotient_domain_master_base_table.nrows(),
Expand All @@ -1046,59 +1048,117 @@ pub fn all_quotients(
quotient_domain_master_ext_table.nrows()
);

prof_start!(maybe_profiler, "malloc");
let mut quotient_table =
Array2::uninit([quotient_domain.length, MasterExtTable::NUM_CONSTRAINTS]);
prof_stop!(maybe_profiler, "malloc");

let init_section_end = MasterExtTable::NUM_INITIAL_CONSTRAINTS;
let cons_section_end = init_section_end + MasterExtTable::NUM_CONSISTENCY_CONSTRAINTS;
let tran_section_end = cons_section_end + MasterExtTable::NUM_TRANSITION_CONSTRAINTS;
let term_section_end = tran_section_end + MasterExtTable::NUM_TERMINAL_CONSTRAINTS;

prof_start!(maybe_profiler, "initial", "AIR");
MasterExtTable::fill_initial_quotients(
quotient_domain_master_base_table,
quotient_domain_master_ext_table,
&mut quotient_table.slice_mut(s![.., ..init_section_end]),
initial_quotient_zerofier_inverse(quotient_domain).view(),
challenges,
);
prof_stop!(maybe_profiler, "initial");

prof_start!(maybe_profiler, "consistency", "AIR");
MasterExtTable::fill_consistency_quotients(
quotient_domain_master_base_table,
quotient_domain_master_ext_table,
&mut quotient_table.slice_mut(s![.., init_section_end..cons_section_end]),
consistency_quotient_zerofier_inverse(trace_domain, quotient_domain).view(),
challenges,
);
prof_stop!(maybe_profiler, "consistency");

prof_start!(maybe_profiler, "transition", "AIR");
MasterExtTable::fill_transition_quotients(
quotient_domain_master_base_table,
quotient_domain_master_ext_table,
&mut quotient_table.slice_mut(s![.., cons_section_end..tran_section_end]),
transition_quotient_zerofier_inverse(trace_domain, quotient_domain).view(),
challenges,
trace_domain,
quotient_domain,

prof_start!(maybe_profiler, "zerofier inverse");
let initial_zerofier_inverse = initial_quotient_zerofier_inverse(quotient_domain);
let consistency_zerofier_inverse =
consistency_quotient_zerofier_inverse(trace_domain, quotient_domain);
let transition_zerofier_inverse =
transition_quotient_zerofier_inverse(trace_domain, quotient_domain);
let terminal_zerofier_inverse =
terminal_quotient_zerofier_inverse(trace_domain, quotient_domain);
prof_stop!(maybe_profiler, "zerofier inverse");

prof_start!(
maybe_profiler,
"evaluate AIR / divide zerofiers / combine quotients"
);
prof_stop!(maybe_profiler, "transition");

prof_start!(maybe_profiler, "terminal", "AIR");
MasterExtTable::fill_terminal_quotients(
quotient_domain_master_base_table,
quotient_domain_master_ext_table,
&mut quotient_table.slice_mut(s![.., tran_section_end..term_section_end]),
terminal_quotient_zerofier_inverse(trace_domain, quotient_domain).view(),
challenges,
let quotient_codeword = (0..quotient_domain.length)
.into_par_iter()
.map(|row| {
// get views of relevant rows
let unit_distance = quotient_domain.length / trace_domain.length;
let domain_length_bit_mask = quotient_domain.length - 1;
let next_row_index = (row + unit_distance) & domain_length_bit_mask;
let current_row_main = quotient_domain_master_base_table.slice(s![row, ..]);
let current_row_aux = quotient_domain_master_ext_table.slice(s![row, ..]);
let next_row_main = quotient_domain_master_base_table.slice(s![next_row_index, ..]);
let next_row_aux = quotient_domain_master_ext_table.slice(s![next_row_index, ..]);

// initial constraints
let initial_constraint_values = MasterExtTable::evaluate_initial_constraints(
current_row_main,
current_row_aux,
challenges,
);
let initial_inner_product = initial_constraint_values
.into_iter()
.zip_eq(
quotient_weights
.iter()
.take(MasterExtTable::NUM_INITIAL_CONSTRAINTS),
)
.map(|(v, &c)| v * c)
.sum::<XFieldElement>();
let mut quotient_value = initial_inner_product * initial_zerofier_inverse[row];

// consistency constraints
let consistency_constraint_values = MasterExtTable::evaluate_consistency_constraints(
current_row_main,
current_row_aux,
challenges,
);
let consistency_inner_product = consistency_constraint_values
.into_iter()
.zip_eq(
quotient_weights
.iter()
.skip(init_section_end)
.take(MasterExtTable::NUM_CONSISTENCY_CONSTRAINTS),
)
.map(|(v, &c)| v * c)
.sum::<XFieldElement>();
quotient_value += consistency_inner_product * consistency_zerofier_inverse[row];

// transition constraints
let transition_constraint_values = MasterExtTable::evaluate_transition_constraints(
current_row_main,
current_row_aux,
next_row_main,
next_row_aux,
challenges,
);
let transition_inner_product = transition_constraint_values
.into_iter()
.zip_eq(
quotient_weights
.iter()
.skip(cons_section_end)
.take(MasterExtTable::NUM_TRANSITION_CONSTRAINTS),
)
.map(|(v, &c)| v * c)
.sum::<XFieldElement>();
quotient_value += transition_inner_product * transition_zerofier_inverse[row];

// terminal constraints
let terminal_constraint_values = MasterExtTable::evaluate_terminal_constraints(
current_row_main,
current_row_aux,
challenges,
);
let terminal_inner_product = terminal_constraint_values
.into_iter()
.zip_eq(
quotient_weights
.iter()
.skip(tran_section_end)
.take(MasterExtTable::NUM_TERMINAL_CONSTRAINTS),
)
.map(|(v, &c)| v * c)
.sum::<XFieldElement>();
quotient_value += terminal_inner_product * terminal_zerofier_inverse[row];
quotient_value
})
.collect();
prof_stop!(
maybe_profiler,
"evaluate AIR / divide zerofiers / combine quotients"
);
prof_stop!(maybe_profiler, "terminal");

unsafe { quotient_table.assume_init() }
quotient_codeword
}

#[deprecated(
Expand Down

0 comments on commit ff30545

Please sign in to comment.