use super::*;
use crate::cs::GpuSetup;
use crate::data_structures::{GenericPolynomialStorageLayout, GenericStorage};
use crate::gpu_proof_config::GpuProofConfig;
use crate::oracle::SubTree;
use crate::poly::{CosetEvaluations, LagrangeBasis, MonomialBasis};
use crate::prover::{
compute_quotient_degree, gpu_prove_from_external_witness_data_with_cache_strategy,
};
use boojum::cs::implementations::prover::ProofConfig;
use boojum::cs::implementations::transcript::Transcript;
use boojum::cs::implementations::verifier::{VerificationKey, VerificationKeyCircuitGeometry};
use boojum::cs::implementations::witness::WitnessVec;
use boojum::worker::Worker;
use era_cudart_sys::CudaError::ErrorMemoryAllocation;
use itertools::Itertools;
use std::collections::BTreeMap;
use std::hash::{DefaultHasher, Hash, Hasher};
use std::ops::Deref;
use std::rc::Rc;
pub trait ChunkLeavesSource {
fn get_chunk_leaves(&mut self, chunk_idx: usize) -> CudaResult<Rc<impl AsSingleSlice>>;
}
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub enum PolynomialsCacheStrategy {
InPlace,
CacheMonomials,
CacheMonomialsAndFirstCoset,
CacheMonomialsAndFriCosets,
CacheEvaluationsAndMonomials,
CacheEvaluationsMonomialsAndFirstCoset,
CacheEvaluationsMonomialsAndFriCosets,
CacheEvaluationsAndAllCosets,
}
use PolynomialsCacheStrategy::*;
impl PolynomialsCacheStrategy {
pub fn required_storages_count(&self, fri_lde_degree: usize, max_lde_degree: usize) -> usize {
match self {
InPlace => 1,
CacheMonomials => 1,
CacheMonomialsAndFirstCoset => 2,
CacheMonomialsAndFriCosets => 1 + fri_lde_degree,
CacheEvaluationsAndMonomials => 2,
CacheEvaluationsMonomialsAndFirstCoset => 3,
CacheEvaluationsMonomialsAndFriCosets => 2 + fri_lde_degree,
CacheEvaluationsAndAllCosets => 1 + max_lde_degree,
}
}
pub fn cache_evaluations(&self) -> bool {
match self {
InPlace | CacheMonomials | CacheMonomialsAndFirstCoset | CacheMonomialsAndFriCosets => {
false
}
CacheEvaluationsAndMonomials
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets
| CacheEvaluationsAndAllCosets => true,
}
}
pub fn cache_monomials(&self) -> bool {
match self {
InPlace | CacheEvaluationsAndAllCosets => false,
CacheEvaluationsAndMonomials
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets
| CacheMonomials
| CacheMonomialsAndFirstCoset
| CacheMonomialsAndFriCosets => true,
}
}
}
pub struct PolynomialsCache<L: GenericPolynomialStorageLayout> {
pub strategy: PolynomialsCacheStrategy,
pub layout: L,
pub domain_size: usize,
pub fri_lde_degree: usize,
pub max_lde_degree: usize,
evaluations: Option<Rc<GenericStorage<LagrangeBasis, L>>>,
monomials: Option<Rc<GenericStorage<MonomialBasis, L>>>,
coset_evaluations: BTreeMap<usize, Rc<GenericStorage<CosetEvaluations, L>>>,
uninitialized_storages: Vec<GenericStorage<Undefined, L>>,
}
impl<L: GenericPolynomialStorageLayout> PolynomialsCache<L> {
fn required_storages_count(&self) -> usize {
self.strategy
.required_storages_count(self.fri_lde_degree, self.max_lde_degree)
}
fn is_storage_borrowed(&self) -> bool {
self.uninitialized_storages.len() != self.required_storages_count()
}
fn allocate(
strategy: PolynomialsCacheStrategy,
layout: L,
domain_size: usize,
fri_lde_degree: usize,
max_lde_degree: usize,
) -> Self {
let storages_count = strategy.required_storages_count(fri_lde_degree, max_lde_degree);
let uninitialized_storages = (0..storages_count)
.map(|_| GenericStorage::allocate(layout, domain_size))
.collect();
Self {
strategy,
layout,
domain_size,
fri_lde_degree,
max_lde_degree,
evaluations: None,
monomials: None,
coset_evaluations: BTreeMap::new(),
uninitialized_storages,
}
}
fn pop_storage(&mut self) -> GenericStorage<Undefined, L> {
self.uninitialized_storages.pop().unwrap()
}
pub fn get_temp_storage(&self) -> GenericStorage<Undefined, L> {
GenericStorage::allocate(self.layout, self.domain_size)
}
pub fn borrow_storage(&mut self) -> GenericStorage<Undefined, L> {
assert_eq!(
self.uninitialized_storages.len(),
self.required_storages_count()
);
self.pop_storage()
}
pub fn num_polys(&self) -> usize {
self.layout.num_polys()
}
pub fn get_evaluations(&mut self) -> CudaResult<Rc<GenericStorage<LagrangeBasis, L>>> {
let result = match self.strategy {
InPlace => {
if let Some(evaluations) = &self.evaluations {
evaluations.clone()
} else {
let monomials = self.get_monomials()?;
drop(self.monomials.take());
let monomials = Rc::into_inner(monomials).unwrap();
let evaluations = Rc::new(monomials.into_evaluations()?);
self.evaluations = Some(evaluations.clone());
evaluations
}
}
CacheMonomials | CacheMonomialsAndFirstCoset | CacheMonomialsAndFriCosets => {
Rc::new(self.monomials.as_ref().unwrap().create_evaluations()?)
}
_ => self.evaluations.as_ref().unwrap().clone(),
};
Ok(result)
}
pub fn get_monomials(&mut self) -> CudaResult<Rc<GenericStorage<MonomialBasis, L>>> {
let result = match self.strategy {
InPlace => {
if let Some(monomials) = &self.monomials {
monomials.clone()
} else {
let monomials = if let Some(evaluations) = self.evaluations.take() {
let evaluations = Rc::into_inner(evaluations).unwrap();
evaluations.into_monomials()?
} else {
let (coset_idx, coset) = self.coset_evaluations.pop_first().unwrap();
assert!(self.coset_evaluations.is_empty());
let coset = Rc::into_inner(coset).unwrap();
coset.into_monomials(coset_idx, self.max_lde_degree)?
};
let monomials = Rc::new(monomials);
self.monomials = Some(monomials.clone());
monomials
}
}
CacheEvaluationsAndAllCosets => {
let (coset_idx, coset) = self.coset_evaluations.first_key_value().unwrap();
let monomials = coset.create_monomials(*coset_idx, self.max_lde_degree)?;
Rc::new(monomials)
}
_ => self.monomials.as_ref().unwrap().clone(),
};
Ok(result)
}
pub fn get_coset_evaluations(
&mut self,
coset_idx: usize,
) -> CudaResult<Rc<GenericStorage<CosetEvaluations, L>>> {
assert!(coset_idx < self.max_lde_degree);
let result = match self.strategy {
InPlace => {
if let Some(coset) = self.coset_evaluations.get(&coset_idx) {
coset.clone()
} else {
let monomials = self.get_monomials()?;
drop(self.monomials.take());
let monomials = Rc::into_inner(monomials).unwrap();
let coset =
Rc::new(monomials.into_coset_evaluations(coset_idx, self.max_lde_degree)?);
self.coset_evaluations.insert(coset_idx, coset.clone());
coset
}
}
CacheMonomialsAndFirstCoset | CacheEvaluationsMonomialsAndFirstCoset
if coset_idx == 0 =>
{
self.coset_evaluations.get(&coset_idx).unwrap().clone()
}
CacheMonomialsAndFriCosets | CacheEvaluationsMonomialsAndFriCosets
if coset_idx < self.fri_lde_degree =>
{
self.coset_evaluations.get(&coset_idx).unwrap().clone()
}
CacheMonomials
| CacheMonomialsAndFirstCoset
| CacheMonomialsAndFriCosets
| CacheEvaluationsAndMonomials
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets => Rc::new(
self.monomials
.as_ref()
.unwrap()
.create_coset_evaluations(coset_idx, self.max_lde_degree)?,
),
CacheEvaluationsAndAllCosets => self.coset_evaluations.get(&coset_idx).unwrap().clone(),
};
Ok(result)
}
pub fn get_coset_evaluations_subset(
&mut self,
coset_idx: usize,
subset: L::PolyType,
) -> CudaResult<Rc<GenericStorage<CosetEvaluations, L>>> {
assert!(coset_idx < self.max_lde_degree);
let result = match self.strategy {
InPlace | CacheEvaluationsAndAllCosets => self.get_coset_evaluations(coset_idx)?,
CacheMonomialsAndFirstCoset | CacheEvaluationsMonomialsAndFirstCoset
if coset_idx == 0 =>
{
self.coset_evaluations.get(&coset_idx).unwrap().clone()
}
CacheMonomialsAndFriCosets | CacheEvaluationsMonomialsAndFriCosets
if coset_idx < self.fri_lde_degree =>
{
self.coset_evaluations.get(&coset_idx).unwrap().clone()
}
CacheMonomials
| CacheMonomialsAndFirstCoset
| CacheMonomialsAndFriCosets
| CacheEvaluationsAndMonomials
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets => Rc::new(
self.monomials
.as_ref()
.unwrap()
.create_coset_evaluations_subset(coset_idx, self.max_lde_degree, subset)?,
),
};
Ok(result)
}
}
impl<L: GenericPolynomialStorageLayout> ChunkLeavesSource for PolynomialsCache<L> {
fn get_chunk_leaves(&mut self, chunk_idx: usize) -> CudaResult<Rc<impl AsSingleSlice>> {
self.get_coset_evaluations(chunk_idx)
}
}
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub enum CommitmentCacheStrategy {
CacheAllLayers,
CacheCosetCaps,
}
use CommitmentCacheStrategy::*;
impl CommitmentCacheStrategy {
pub fn required_storages_count(&self, chunks_count: usize) -> usize {
match self {
CacheAllLayers => chunks_count,
CacheCosetCaps => 0,
}
}
}
pub struct CommitmentCache<H: GpuTreeHasher> {
pub strategy: CommitmentCacheStrategy,
pub chunk_size: usize,
pub chunks_count: usize,
pub cap_size: usize,
pub chunk_cap_size: usize,
pub num_elems_per_leaf: usize,
pub subtrees_by_chunk: BTreeMap<usize, SubTree<H>>,
pub caps_by_chunk: BTreeMap<usize, Vec<H::Output>>,
pub tree: Vec<Vec<H::Output>>,
pub(crate) uninitialized_subtree_nodes: Vec<DVec<H::DigestElementType>>,
}
impl<H: GpuTreeHasher> CommitmentCache<H> {
fn get_nodes_size_per_chunk(chunk_size: usize, num_elems_per_leaf: usize) -> usize {
assert!(chunk_size.is_power_of_two());
assert!(num_elems_per_leaf.is_power_of_two());
assert!(num_elems_per_leaf <= chunk_size);
H::CAPACITY << (chunk_size.trailing_zeros() - num_elems_per_leaf.trailing_zeros() + 1)
}
pub fn allocate(
strategy: CommitmentCacheStrategy,
chunk_size: usize,
chunks_count: usize,
cap_size: usize,
num_elems_per_leaf: usize,
) -> Self {
assert!(chunks_count.is_power_of_two());
assert!(cap_size.is_power_of_two());
let chunk_cap_size = coset_cap_size(cap_size, chunks_count);
let storages_count = strategy.required_storages_count(chunks_count);
let nodes_size = Self::get_nodes_size_per_chunk(chunk_size, num_elems_per_leaf);
let uninitialized_subtree_nodes = (0..storages_count).map(|_| dvec!(nodes_size)).collect();
Self {
strategy,
chunk_size,
chunks_count,
cap_size,
chunk_cap_size,
num_elems_per_leaf,
subtrees_by_chunk: BTreeMap::new(),
caps_by_chunk: BTreeMap::new(),
tree: vec![],
uninitialized_subtree_nodes,
}
}
fn get_coset_cap_from_tree(&self, coset_index: usize) -> &[H::Output] {
let start = coset_index * self.chunk_cap_size;
let end = start + self.chunk_cap_size;
&self.tree[0][start..end]
}
pub(crate) fn initialize_for_chunk(
&mut self,
chunk_idx: usize,
chunk: &impl AsSingleSlice,
) -> CudaResult<()> {
assert!(chunk_idx < self.chunks_count);
match self.strategy {
CacheAllLayers => {
let nodes = self.uninitialized_subtree_nodes.pop().unwrap();
let (subtree, subtree_cap) =
build_subtree::<H>(chunk, self.chunk_cap_size, self.num_elems_per_leaf, nodes)?;
self.subtrees_by_chunk.insert(chunk_idx, subtree);
self.caps_by_chunk.insert(chunk_idx, subtree_cap);
}
CacheCosetCaps => {
let nodes_size =
Self::get_nodes_size_per_chunk(self.chunk_size, self.num_elems_per_leaf);
let nodes = dvec!(nodes_size);
let (_, subtree_cap) =
build_subtree::<H>(chunk, self.chunk_cap_size, self.num_elems_per_leaf, nodes)?;
self.caps_by_chunk.insert(chunk_idx, subtree_cap);
}
};
Ok(())
}
pub(crate) fn build_tree(&mut self) -> CudaResult<()> {
let is_dry_run = is_dry_run()?;
if self.tree.is_empty() {
assert_eq!(self.caps_by_chunk.len(), self.chunks_count);
self.tree.push(
(0..self.chunks_count)
.flat_map(|i| self.caps_by_chunk.get(&i).unwrap().clone())
.collect_vec(),
);
while self.tree.last().unwrap().len() != self.cap_size {
self.tree.push(
self.tree
.last()
.unwrap()
.chunks(2)
.map(|chunk| {
if is_dry_run {
H::Output::default()
} else {
H::hash_into_node(&chunk[0], &chunk[1], 0)
}
})
.collect_vec(),
);
}
} else {
for chunk_index in 0..self.chunks_count {
let tree_coset_cap = self.get_coset_cap_from_tree(chunk_index);
if let Some(map_coset_cap) = self.caps_by_chunk.get(&chunk_index) {
if !is_dry_run {
assert_eq!(tree_coset_cap, map_coset_cap);
}
} else {
self.caps_by_chunk
.insert(chunk_index, tree_coset_cap.to_vec());
}
}
}
Ok(())
}
pub fn get_tree_cap(&self) -> Vec<H::Output> {
self.tree.last().unwrap().clone()
}
pub fn batch_query_for_chunk<A: GoodAllocator>(
&self,
chunk_source: &mut impl ChunkLeavesSource,
chunk_idx: usize,
d_indexes: &DVec<u32, SmallStaticDeviceAllocator>,
h_all_leaf_elems: &mut Vec<F, A>,
h_all_proofs: &mut Vec<H::DigestElementType, A>,
) -> CudaResult<()> {
let leaf_sources = chunk_source.get_chunk_leaves(chunk_idx)?;
let subtree = if let Some(subtree) = self.subtrees_by_chunk.get(&chunk_idx) {
subtree.clone()
} else {
let coset_cap_size = coset_cap_size(self.cap_size, self.chunks_count);
let nodes_size =
Self::get_nodes_size_per_chunk(self.chunk_size, self.num_elems_per_leaf);
let nodes = dvec!(nodes_size);
let (subtree, subtree_cap) = build_subtree::<H>(
leaf_sources.deref(),
coset_cap_size,
self.num_elems_per_leaf,
nodes,
)?;
if !is_dry_run()? {
assert_eq!(&subtree_cap, self.caps_by_chunk.get(&chunk_idx).unwrap());
}
subtree
};
batch_query::<H, A>(
d_indexes,
leaf_sources.deref(),
leaf_sources.num_polys_in_base(),
&subtree,
subtree.cap_size,
self.chunk_size,
self.num_elems_per_leaf,
h_all_leaf_elems,
h_all_proofs,
)?;
if self.chunks_count > self.cap_size {
let remainder_count =
(self.chunks_count.trailing_zeros() - self.cap_size.trailing_zeros()) as usize;
assert_eq!(remainder_count + 1, self.tree.len());
let num_queries = d_indexes.len();
let mut index = chunk_idx;
for layer in self.tree.iter().take(remainder_count) {
let elements = H::DigestElements::from(layer[index ^ 1]);
index >>= 1;
elements
.iter()
.for_each(|e| h_all_proofs.extend(std::iter::repeat(e).take(num_queries)));
}
}
Ok(())
}
}
pub struct StorageCache<L: GenericPolynomialStorageLayout, H: GpuTreeHasher, T = ()> {
pub polynomials_cache: PolynomialsCache<L>,
pub commitment_cache: CommitmentCache<H>,
pub domain_size: usize,
pub fri_lde_degree: usize,
pub max_lde_degree: usize,
pub aux: T,
}
impl<L: GenericPolynomialStorageLayout, H: GpuTreeHasher, T> StorageCache<L, H, T> {
#[allow(clippy::too_many_arguments)]
pub fn allocate(
polynomials_strategy: PolynomialsCacheStrategy,
commitments_strategy: CommitmentCacheStrategy,
layout: L,
domain_size: usize,
fri_lde_degree: usize,
max_lde_degree: usize,
cap_size: usize,
num_elems_per_leaf: usize,
aux: T,
) -> Self {
let polynomials_cache = PolynomialsCache::allocate(
polynomials_strategy,
layout,
domain_size,
fri_lde_degree,
max_lde_degree,
);
let commitments_cache = CommitmentCache::allocate(
commitments_strategy,
domain_size,
fri_lde_degree,
cap_size,
num_elems_per_leaf,
);
Self {
polynomials_cache,
commitment_cache: commitments_cache,
domain_size,
fri_lde_degree,
max_lde_degree,
aux,
}
}
pub fn initialize_from_evaluations(
&mut self,
evaluations: Rc<GenericStorage<LagrangeBasis, L>>,
) -> CudaResult<()> {
let can_own_evaluations = Rc::strong_count(&evaluations) == 1;
let polynomials_strategy = self.polynomials_cache.strategy;
let cache_evaluations = polynomials_strategy.cache_evaluations();
let is_storage_borrowed = self.polynomials_cache.is_storage_borrowed();
let must_own_evaluations = !cache_evaluations && is_storage_borrowed;
assert!(can_own_evaluations || !must_own_evaluations);
let monomials = if cache_evaluations {
assert!(is_storage_borrowed);
let monomials = evaluations.fill_monomials(self.polynomials_cache.pop_storage())?;
self.polynomials_cache.evaluations = Some(evaluations);
monomials
} else if is_storage_borrowed {
Rc::into_inner(evaluations).unwrap().into_monomials()?
} else {
evaluations.fill_monomials(self.polynomials_cache.pop_storage())?
};
self.initialize_from_monomials(Rc::new(monomials))
}
pub fn initialize_from_monomials(
&mut self,
monomials: Rc<GenericStorage<MonomialBasis, L>>,
) -> CudaResult<()> {
let can_owns_monomials = Rc::strong_count(&monomials) == 1;
let polynomials_strategy = self.polynomials_cache.strategy;
let cache_monomials = polynomials_strategy.cache_monomials();
let is_storage_borrowed = self.polynomials_cache.is_storage_borrowed();
let must_own_monomials = !cache_monomials && is_storage_borrowed;
assert!(can_owns_monomials || !must_own_monomials);
let mut monomials = Some(monomials);
if polynomials_strategy.cache_evaluations() && self.polynomials_cache.evaluations.is_none()
{
let evaluations = monomials
.as_ref()
.unwrap()
.fill_evaluations(self.polynomials_cache.pop_storage())?;
self.polynomials_cache.evaluations = Some(Rc::new(evaluations));
};
let polynomials_cosets_count = match polynomials_strategy {
InPlace => 0,
CacheMonomials => 0,
CacheMonomialsAndFirstCoset => 1,
CacheMonomialsAndFriCosets => self.fri_lde_degree,
CacheEvaluationsAndMonomials => 0,
CacheEvaluationsMonomialsAndFirstCoset => 1,
CacheEvaluationsMonomialsAndFriCosets => self.fri_lde_degree,
CacheEvaluationsAndAllCosets => self.max_lde_degree,
};
let commitment_cosets_count = match self.commitment_cache.strategy {
CacheCosetCaps if !self.commitment_cache.tree.is_empty() => 0,
_ => self.fri_lde_degree,
};
let cosets_count = usize::max(polynomials_cosets_count, commitment_cosets_count);
for coset_idx in 0..cosets_count {
let coset = match polynomials_strategy {
InPlace => Rc::into_inner(monomials.take().unwrap())
.unwrap()
.into_coset_evaluations(coset_idx, self.max_lde_degree)?,
CacheMonomials | CacheEvaluationsAndMonomials => monomials
.as_ref()
.unwrap()
.create_coset_evaluations(coset_idx, self.max_lde_degree)?,
CacheMonomialsAndFirstCoset | CacheEvaluationsMonomialsAndFirstCoset
if coset_idx == 0 =>
{
monomials.as_ref().unwrap().fill_coset_evaluations(
coset_idx,
self.max_lde_degree,
self.polynomials_cache.pop_storage(),
)?
}
CacheMonomialsAndFriCosets | CacheEvaluationsMonomialsAndFriCosets
if coset_idx < self.fri_lde_degree =>
{
monomials.as_ref().unwrap().fill_coset_evaluations(
coset_idx,
self.max_lde_degree,
self.polynomials_cache.pop_storage(),
)?
}
CacheMonomialsAndFirstCoset
| CacheMonomialsAndFriCosets
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets => monomials
.as_ref()
.unwrap()
.create_coset_evaluations(coset_idx, self.max_lde_degree)?,
CacheEvaluationsAndAllCosets => {
if coset_idx + 1 == self.max_lde_degree {
Rc::into_inner(monomials.take().unwrap())
.unwrap()
.into_coset_evaluations(coset_idx, self.max_lde_degree)?
} else {
monomials.as_ref().unwrap().fill_coset_evaluations(
coset_idx,
self.max_lde_degree,
self.polynomials_cache.pop_storage(),
)?
}
}
};
if coset_idx < commitment_cosets_count {
self.commitment_cache
.initialize_for_chunk(coset_idx, &coset)?;
};
match polynomials_strategy {
InPlace => {
monomials = Some(Rc::new(
coset.into_monomials(coset_idx, self.max_lde_degree)?,
));
}
CacheMonomials | CacheEvaluationsAndMonomials => {}
CacheMonomialsAndFirstCoset | CacheEvaluationsMonomialsAndFirstCoset
if coset_idx == 0 =>
{
self.polynomials_cache
.coset_evaluations
.insert(coset_idx, Rc::new(coset));
}
CacheMonomialsAndFriCosets | CacheEvaluationsMonomialsAndFriCosets
if coset_idx < self.fri_lde_degree =>
{
self.polynomials_cache
.coset_evaluations
.insert(coset_idx, Rc::new(coset));
}
CacheMonomialsAndFirstCoset
| CacheMonomialsAndFriCosets
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets => {}
CacheEvaluationsAndAllCosets => {
self.polynomials_cache
.coset_evaluations
.insert(coset_idx, Rc::new(coset));
}
};
}
self.commitment_cache.build_tree()?;
match polynomials_strategy {
CacheEvaluationsAndAllCosets => {}
_ => {
self.polynomials_cache.monomials = Some(monomials.take().unwrap());
}
};
Ok(())
}
pub fn num_polys(&self) -> usize {
self.polynomials_cache.num_polys()
}
pub fn get_evaluations(&mut self) -> CudaResult<Rc<GenericStorage<LagrangeBasis, L>>> {
self.polynomials_cache.get_evaluations()
}
pub fn get_monomials(&mut self) -> CudaResult<Rc<GenericStorage<MonomialBasis, L>>> {
self.polynomials_cache.get_monomials()
}
pub fn get_coset_evaluations(
&mut self,
coset_idx: usize,
) -> CudaResult<Rc<GenericStorage<CosetEvaluations, L>>> {
self.polynomials_cache.get_coset_evaluations(coset_idx)
}
pub fn get_coset_evaluations_subset(
&mut self,
coset_idx: usize,
subset: L::PolyType,
) -> CudaResult<Rc<GenericStorage<CosetEvaluations, L>>> {
self.polynomials_cache
.get_coset_evaluations_subset(coset_idx, subset)
}
pub fn get_tree_cap(&mut self) -> Vec<H::Output> {
self.commitment_cache.get_tree_cap()
}
pub fn batch_query_for_coset<A: GoodAllocator>(
&mut self,
coset_idx: usize,
d_indexes: &DVec<u32, SmallStaticDeviceAllocator>,
h_all_leaf_elems: &mut Vec<F, A>,
h_all_proofs: &mut Vec<H::DigestElementType, A>,
) -> CudaResult<()> {
self.commitment_cache.batch_query_for_chunk(
&mut self.polynomials_cache,
coset_idx,
d_indexes,
h_all_leaf_elems,
h_all_proofs,
)
}
}
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct CacheStrategy {
pub setup_polynomials: PolynomialsCacheStrategy,
pub trace_polynomials: PolynomialsCacheStrategy,
pub other_polynomials: PolynomialsCacheStrategy,
pub commitment: CommitmentCacheStrategy,
}
impl CacheStrategy {
pub(crate) fn get<
TR: Transcript<F, CompatibleCap: Hash>,
H: GpuTreeHasher<Output = TR::CompatibleCap>,
POW: GPUPoWRunner,
A: GoodAllocator,
>(
config: &GpuProofConfig,
external_witness_data: &WitnessVec<F>,
proof_config: ProofConfig,
setup: &GpuSetup<H>,
vk: &VerificationKey<F, H>,
transcript_params: TR::TransciptParameters,
worker: &Worker,
) -> CudaResult<Self> {
let cap = &vk.setup_merkle_tree_cap;
let mut hasher = DefaultHasher::new();
Hash::hash_slice(cap, &mut hasher);
let cap_hash = hasher.finish();
if let Some(strategy) = _strategy_cache_get().get(&cap_hash) {
println!("reusing cache strategy");
Ok(*strategy)
} else {
let strategies =
Self::get_strategy_candidates(config, &proof_config, setup, &vk.fixed_parameters);
for (_, strategy) in strategies.iter().copied() {
_setup_cache_reset();
dry_run_start();
let result =
gpu_prove_from_external_witness_data_with_cache_strategy::<TR, H, POW, A>(
config,
external_witness_data,
proof_config.clone(),
setup,
vk,
transcript_params.clone(),
worker,
strategy,
);
_setup_cache_reset();
let result = result.and(dry_run_stop());
match result {
Ok(_) => {
println!("determined cache strategy: {:?}", strategy);
_strategy_cache_get().insert(cap_hash, strategy);
return Ok(strategy);
}
Err(ErrorMemoryAllocation) => {
continue;
}
Err(e) => return Err(e),
}
}
Err(ErrorMemoryAllocation)
}
}
pub(crate) fn get_strategy_candidates<H: GpuTreeHasher>(
config: &GpuProofConfig,
proof_config: &ProofConfig,
setup: &GpuSetup<H>,
geometry: &VerificationKeyCircuitGeometry,
) -> Vec<((usize, usize), CacheStrategy)> {
let fri_lde_degree = proof_config.fri_lde_factor;
let quotient_degree = compute_quotient_degree(config, &setup.selectors_placement);
let max_lde_degree = usize::max(quotient_degree, fri_lde_degree);
let setup_layout = setup.layout;
let domain_size = geometry.domain_size as usize;
let lookup_parameters = geometry.lookup_parameters;
let total_tables_len = geometry.total_tables_len as usize;
let num_multiplicity_cols =
lookup_parameters.num_multipicities_polys(total_tables_len, domain_size);
let trace_layout = TraceLayout {
num_variable_cols: setup.variables_hint.len(),
num_witness_cols: setup.witnesses_hint.len(),
num_multiplicity_cols,
};
let arguments_layout = ArgumentsLayout::from_trace_layout_and_lookup_params(
trace_layout,
quotient_degree,
geometry.lookup_parameters,
);
let quotient_layout = QuotientLayout::new(quotient_degree);
let setup_num_polys = setup_layout.num_polys();
let trace_num_polys = trace_layout.num_polys();
let other_num_polys = arguments_layout.num_polys() + quotient_layout.num_polys();
let commitments_cache_strategies = [CacheAllLayers, CacheCosetCaps];
let setup_polynomials_strategies = [
InPlace,
CacheMonomials,
CacheMonomialsAndFirstCoset,
CacheMonomialsAndFriCosets,
CacheEvaluationsAndMonomials,
CacheEvaluationsMonomialsAndFirstCoset,
CacheEvaluationsMonomialsAndFriCosets,
CacheEvaluationsAndAllCosets,
];
let other_polynomials_strategies = [
InPlace,
CacheMonomials,
CacheMonomialsAndFirstCoset,
CacheMonomialsAndFriCosets,
];
let mut strategies = Vec::new();
for commitment_strategy in commitments_cache_strategies.iter().copied() {
for setup_strategy in setup_polynomials_strategies.iter().copied() {
for other_strategies in (0..2)
.map(|_| other_polynomials_strategies.iter().cloned())
.multi_cartesian_product()
{
let strategy = CacheStrategy {
setup_polynomials: setup_strategy,
trace_polynomials: other_strategies[0],
other_polynomials: other_strategies[1],
commitment: commitment_strategy,
};
let setup_cost =
strategy.get_setup_cost(fri_lde_degree, max_lde_degree) * setup_num_polys;
let proof_cost_setup = strategy
.get_proof_cost_setup(fri_lde_degree, max_lde_degree)
* setup_num_polys;
let proof_cost_trace = strategy
.get_proof_cost_trace(fri_lde_degree, max_lde_degree)
* trace_num_polys;
let proof_cost_other = strategy
.get_proof_cost_other(fri_lde_degree, max_lde_degree)
* other_num_polys;
let proof_cost = proof_cost_setup + proof_cost_trace + proof_cost_other;
let costs = match commitment_strategy {
CacheAllLayers => (proof_cost, setup_cost),
CacheCosetCaps => (proof_cost << 16, setup_cost << 16),
};
strategies.push((costs, strategy));
}
}
}
strategies.sort_by_key(|x| x.0);
strategies
}
fn get_setup_cost(&self, fri_lde_degree: usize, max_lde_degree: usize) -> usize {
let f = fri_lde_degree;
let u = max_lde_degree;
match self.setup_polynomials {
InPlace => 1 + 2 * f,
CacheMonomials
| CacheMonomialsAndFirstCoset
| CacheMonomialsAndFriCosets
| CacheEvaluationsAndMonomials
| CacheEvaluationsMonomialsAndFirstCoset
| CacheEvaluationsMonomialsAndFriCosets => 1 + f,
CacheEvaluationsAndAllCosets => 1 + u,
}
}
fn get_proof_cost_setup(&self, fri_lde_degree: usize, max_lde_degree: usize) -> usize {
let f = fri_lde_degree;
let u = max_lde_degree;
match self.setup_polynomials {
InPlace => 2 + 2 * u + 2 + 2 * (f - 1) + 2 * f,
CacheMonomials => 1 + u + 1 + f + f,
CacheMonomialsAndFirstCoset => 1 + u - 1 + f - 1 + f - 1,
CacheMonomialsAndFriCosets => 1 + u - f,
CacheEvaluationsAndMonomials => u + 1 + f + f,
CacheEvaluationsMonomialsAndFirstCoset => u - 1 + f - 1 + f - 1,
CacheEvaluationsMonomialsAndFriCosets => u - f,
CacheEvaluationsAndAllCosets => 0,
}
}
fn get_proof_cost_trace(&self, fri_lde_degree: usize, max_lde_degree: usize) -> usize {
let f = fri_lde_degree;
let u = max_lde_degree;
match self.trace_polynomials {
InPlace => 1 + 2 * f + 1 + 2 * u + 2 + 2 * (f - 1) + 2 * f,
CacheMonomials => 1 + f + u + 1 + f + f,
CacheMonomialsAndFirstCoset => 1 + f + u - 1 + f - 1 + f - 1,
CacheMonomialsAndFriCosets => 1 + f + u - f,
CacheEvaluationsAndMonomials => 1 + f + u + 1 + f + f,
CacheEvaluationsMonomialsAndFirstCoset => 1 + f + u - 1 + f - 1 + f - 1,
CacheEvaluationsMonomialsAndFriCosets => 1 + f + u - f,
CacheEvaluationsAndAllCosets => 1 + u,
}
}
fn get_proof_cost_other(&self, fri_lde_degree: usize, max_lde_degree: usize) -> usize {
let f = fri_lde_degree;
let u = max_lde_degree;
match self.other_polynomials {
InPlace => 1 + 2 * f + 2 * u - 1 + 1 + 1 + 2 * (f - 1) + 2 * f,
CacheMonomials => 1 + f + u + 1 + f + f,
CacheMonomialsAndFirstCoset => 1 + f + u - 1 + f - 1 + f - 1,
CacheMonomialsAndFriCosets => 1 + f + u - f,
CacheEvaluationsAndMonomials => 1 + f + u + 1 + f + f,
CacheEvaluationsMonomialsAndFirstCoset => 1 + f + u - 1 + f - 1 + f - 1,
CacheEvaluationsMonomialsAndFriCosets => 1 + f + u - f,
CacheEvaluationsAndAllCosets => 1 + u,
}
}
}