Skip to content

Commit

Permalink
perf: re-organize prover steps
Browse files Browse the repository at this point in the history 
Bypass quotient table, compute quotient codeword directly instead.
  • Loading branch information
@jan-ferdinand
jan-ferdinand committed Mar 31, 2024
2 parents fa38fa8 + d574b40 commit 86a7799
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 211 deletions.
36 changes: 16 additions & 20 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 @@ -176,41 +176,37 @@ impl Stark {

prof_start!(maybe_profiler, "Fiat-Shamir", "hash");
proof_stream.enqueue(ProofItem::MerkleRoot(ext_merkle_tree.root()));

// Get the weights with which to compress the many quotients into one.
let quotient_combination_weights =
proof_stream.sample_scalars(MasterExtTable::NUM_CONSTRAINTS);
prof_stop!(maybe_profiler, "Fiat-Shamir");
prof_stop!(maybe_profiler, "ext tables");

// low-degree extend the trace codewords so that all the quotient codewords
// can be obtained by element-wise evaluation of the AIR constraints
prof_start!(maybe_profiler, "quotient-domain codewords");
let base_quotient_domain_codewords = master_base_table.quotient_domain_table();
let ext_quotient_domain_codewords = master_ext_table.quotient_domain_table();
prof_stop!(maybe_profiler, "quotient-domain codewords");

prof_start!(maybe_profiler, "quotient codewords");
let master_quotient_table = all_quotients(
prof_start!(
maybe_profiler,
"compute and combine quotient codewords",
"CC"
);
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,
);
prof_stop!(maybe_profiler, "quotient codewords");

prof_start!(maybe_profiler, "linearly combine quotient codewords", "CC");
// Create quotient codeword. This is a part of the combination codeword. To reduce the
// amount of hashing necessary, the quotient codeword is linearly summed instead of
// hashed prior to committing to it.
let quotient_combination_weights =
proof_stream.sample_scalars(MasterExtTable::NUM_CONSTRAINTS);
let quotient_combination_weights = Array1::from(quotient_combination_weights);
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, "linearly combine quotient codewords");
prof_stop!(maybe_profiler, "compute and combine quotient codewords");

prof_start!(maybe_profiler, "commit to quotient codeword segments");
prof_start!(maybe_profiler, "LDE", "LDE");
Expand Down
129 changes: 3 additions & 126 deletions triton-vm/src/table/extension_table.rs
View file
Original file line number Diff line number Diff line change
@@ -1,23 +1,19 @@
use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt::Result as FmtResult;
use std::mem::MaybeUninit;

use itertools::Itertools;
use ndarray::parallel::prelude::*;
use ndarray::Array1;
use ndarray::ArrayView1;
use ndarray::ArrayView2;
use ndarray::ArrayViewMut2;
use ndarray::Axis;
use twenty_first::prelude::*;
use twenty_first::shared_math::mpolynomial::Degree;
use twenty_first::shared_math::traits::FiniteField;

use crate::arithmetic_domain::ArithmeticDomain;
use crate::table::challenges::Challenges;
use crate::table::master_table::MasterExtTable;

/// The implementations of these functions are automatically generated using the
/// command `cargo run --bin constraint-evaluation-generator` and live in
/// `constraints.rs`.
pub trait Evaluable<FF: FiniteField> {
/// The code for this method must be generated by running
/// `cargo run --bin constraint-evaluation-generator`
Expand Down Expand Up @@ -66,125 +62,6 @@ pub trait Quotientable: Evaluable<BFieldElement> {
+ Self::NUM_TRANSITION_CONSTRAINTS
+ Self::NUM_TERMINAL_CONSTRAINTS;

fn fill_initial_quotients(
master_base_table: ArrayView2<BFieldElement>,
master_ext_table: ArrayView2<XFieldElement>,
quot_table: &mut ArrayViewMut2<MaybeUninit<XFieldElement>>,
zerofier_inverse: ArrayView1<BFieldElement>,
challenges: &Challenges,
) {
debug_assert_eq!(zerofier_inverse.len(), master_base_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), master_ext_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), quot_table.nrows());
quot_table
.axis_iter_mut(Axis(0))
.into_par_iter()
.enumerate()
.for_each(|(row_index, quotient_table_row)| {
let base_row = master_base_table.row(row_index);
let ext_row = master_ext_table.row(row_index);
Self::evaluate_initial_constraints(base_row, ext_row, challenges)
.into_iter()
.map(|numerator| numerator * zerofier_inverse[row_index])
.map(MaybeUninit::new)
.collect::<Array1<_>>()
.move_into(quotient_table_row);
});
}

fn fill_consistency_quotients(
master_base_table: ArrayView2<BFieldElement>,
master_ext_table: ArrayView2<XFieldElement>,
quot_table: &mut ArrayViewMut2<MaybeUninit<XFieldElement>>,
zerofier_inverse: ArrayView1<BFieldElement>,
challenges: &Challenges,
) {
debug_assert_eq!(zerofier_inverse.len(), master_base_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), master_ext_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), quot_table.nrows());
quot_table
.axis_iter_mut(Axis(0))
.into_par_iter()
.enumerate()
.for_each(|(row_index, quotient_table_row)| {
let base_row = master_base_table.row(row_index);
let ext_row = master_ext_table.row(row_index);
Self::evaluate_consistency_constraints(base_row, ext_row, challenges)
.into_iter()
.map(|numerator| numerator * zerofier_inverse[row_index])
.map(MaybeUninit::new)
.collect::<Array1<_>>()
.move_into(quotient_table_row);
});
}

fn fill_transition_quotients(
master_base_table: ArrayView2<BFieldElement>,
master_ext_table: ArrayView2<XFieldElement>,
quot_table: &mut ArrayViewMut2<MaybeUninit<XFieldElement>>,
zerofier_inverse: ArrayView1<BFieldElement>,
challenges: &Challenges,
trace_domain: ArithmeticDomain,
quotient_domain: ArithmeticDomain,
) {
debug_assert_eq!(zerofier_inverse.len(), master_base_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), master_ext_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), quot_table.nrows());

// the relation between the quotient domain and the trace domain
let unit_distance = quotient_domain.length / trace_domain.length;
let domain_length_bit_mask = quotient_domain.length - 1;

quot_table
.axis_iter_mut(Axis(0))
.into_par_iter()
.enumerate()
.for_each(|(current_row_index, quotient_table_row)| {
// bitwise logical and `domain_length_bit_mask` performs the modulo operation:
// `domain.length - 1` is a bit-mask with all 1s because `domain.length` is 2^k
// for some k.
let next_row_index = (current_row_index + unit_distance) & domain_length_bit_mask;
Self::evaluate_transition_constraints(
master_base_table.row(current_row_index),
master_ext_table.row(current_row_index),
master_base_table.row(next_row_index),
master_ext_table.row(next_row_index),
challenges,
)
.into_iter()
.map(|numerator| numerator * zerofier_inverse[current_row_index])
.map(MaybeUninit::new)
.collect::<Array1<_>>()
.move_into(quotient_table_row);
});
}

fn fill_terminal_quotients(
master_base_table: ArrayView2<BFieldElement>,
master_ext_table: ArrayView2<XFieldElement>,
quot_table: &mut ArrayViewMut2<MaybeUninit<XFieldElement>>,
zerofier_inverse: ArrayView1<BFieldElement>,
challenges: &Challenges,
) {
debug_assert_eq!(zerofier_inverse.len(), master_base_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), master_ext_table.nrows());
debug_assert_eq!(zerofier_inverse.len(), quot_table.nrows());
quot_table
.axis_iter_mut(Axis(0))
.into_par_iter()
.enumerate()
.for_each(|(row_index, quotient_table_row)| {
let base_row = master_base_table.row(row_index);
let ext_row = master_ext_table.row(row_index);
Self::evaluate_terminal_constraints(base_row, ext_row, challenges)
.into_iter()
.map(|numerator| numerator * zerofier_inverse[row_index])
.map(MaybeUninit::new)
.collect::<Array1<_>>()
.move_into(quotient_table_row);
});
}

fn initial_quotient_degree_bounds(interpolant_degree: Degree) -> Vec<Degree>;

fn consistency_quotient_degree_bounds(
Expand Down
Loading

0 comments on commit 86a7799

Please sign in to comment.