Skip to content

Memory opt #3

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.

Sign up for GitHub

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Memory opt #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1a2fbbb
vec opt
xzav3r Oct 17, 2023
f4ce264
Assembly opt
xzav3r Oct 17, 2023
5870bc6
pre allocating
xzav3r Oct 19, 2023
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
66 changes: 34 additions & 32 deletions halo2_proofs/src/helpers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,7 @@ use num_derive::FromPrimitive;
use rayon::prelude::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
use std::io::Seek;
use std::marker::PhantomData;
use std::{
fs::{File, OpenOptions},
io,
ops::RangeTo,
};
use std::{fs::{File, OpenOptions}, io, mem, ops::RangeTo};

pub(crate) trait CurveRead: CurveAffine {
/// Reads a compressed element from the buffer and attempts to parse it
Expand Down Expand Up @@ -105,9 +101,9 @@ impl Serializable for (String, u32) {

impl ParaSerializable for Vec<Vec<(u32, u32)>> {
fn vec_fetch(fd: &mut File) -> io::Result<Self> {
let columns = read_u32(fd)?;
let columns = read_u32(fd)? as usize;
let mut offset = 0;
let mut offsets = vec![];
let mut offsets = Vec::with_capacity(columns);
for _ in 0..columns {
let l = read_u32(fd)?;
offsets.push((offset, l));
Expand All @@ -124,6 +120,7 @@ impl ParaSerializable for Vec<Vec<(u32, u32)>> {
.map(&fd)
.unwrap()
};
//TODO: to be optimized
let s: &[(u32, u32)] = unsafe {
std::slice::from_raw_parts(mmap.as_ptr() as *const (u32, u32), offsets[i as usize].1 as usize)
};
Expand All @@ -139,7 +136,7 @@ impl ParaSerializable for Vec<Vec<(u32, u32)>> {
let u = self.len() as u32;
u.store(fd)?;
let mut offset = 0;
let mut offsets = vec![];
let mut offsets = Vec::with_capacity(u as usize);
for i in 0..u {
let l = self[i as usize].len();
offsets.push((offset, l));
Expand All @@ -151,15 +148,15 @@ impl ParaSerializable for Vec<Vec<(u32, u32)>> {
self.into_par_iter().enumerate().for_each(|(i, s2)| {
let mut mmap = unsafe {
MmapOptions::new()
.offset(position + (offsets[i as usize].0 as u64 * 8))
.len(offsets[i as usize].1 as usize * 8)
.offset(position + (offsets[i].0 as u64 * 8))
.len(offsets[i].1 * 8)
.map_mut(&fd)
.unwrap()
};
let s: &[u8] = unsafe {
std::slice::from_raw_parts(
s2.as_ptr() as *const u8,
offsets[i as usize].1 as usize * 8,
offsets[i].1 * 8,
)
};
(&mut mmap).copy_from_slice(s);
Expand All @@ -173,10 +170,13 @@ impl<B:Clone, F: FieldExt> Serializable for Polynomial<F, B> {
let u = read_u32(reader)?;
let mut buf = vec![0u8; u as usize * 32];
reader.read_exact(&mut buf)?;
let s: &[F] = unsafe {
std::slice::from_raw_parts(buf.as_ptr() as *const F, u as usize)
let ptr = buf.as_ptr() as *mut F;
// Don't run the destructor for buf, because it used by `values`
mem::forget(buf);
let values: Vec<F> = unsafe {
Vec::from_raw_parts(ptr, u as usize, u as usize)
};
Ok(Polynomial::new(s.to_vec()))
Ok(Polynomial::new(values))
}
fn store<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
let u = self.values.len() as u32;
Expand Down Expand Up @@ -208,9 +208,10 @@ impl<C: CurveAffine> Serializable for VerifyingKey<C> {
let domain: EvaluationDomain<C::Scalar> = EvaluationDomain::new(j, k);
let cs = read_cs::<C, R>(reader)?;

let fixed_commitments: Vec<_> = (0..cs.num_fixed_columns)
.map(|_| C::read(reader))
.collect::<Result<_, _>>()?;
let mut fixed_commitments = Vec::with_capacity(cs.num_fixed_columns);
for _ in 0..cs.num_fixed_columns {
fixed_commitments.push(C::read(reader)?);
}

let permutation = permutation::VerifyingKey::read(reader, &cs.permutation)?;

Expand All @@ -233,7 +234,7 @@ impl<T: Serializable> Serializable for Vec<T> {
}
fn fetch<R: io::Read>(reader: &mut R) -> io::Result<Vec<T>> {
let len = read_u32(reader)?;
let mut v = vec![];
let mut v = Vec::with_capacity(len as usize);
for _ in 0..len {
v.push(T::fetch(reader)?);
}
Expand Down Expand Up @@ -281,8 +282,8 @@ fn write_arguments<W: std::io::Write>(
fn read_arguments<R: std::io::Read>(
reader: &mut R,
) -> std::io::Result<plonk::permutation::Argument> {
let len = read_u32(reader)?;
let mut cols = vec![];
let len = read_u32(reader)? as usize;
let mut cols = Vec::with_capacity(len);
for _ in 0..len {
cols.push(Column::<Any>::fetch(reader)?);
}
Expand Down Expand Up @@ -336,8 +337,8 @@ fn read_queries<T: ColumnType, R: std::io::Read>(
reader: &mut R,
t: T,
) -> std::io::Result<Vec<(Column<T>, Rotation)>> {
let mut queries = vec![];
let len = read_u32(reader)?;
let len = read_u32(reader)? as usize;
let mut queries = Vec::with_capacity(len);
for _ in 0..len {
let column = read_column(reader, t)?;
let rotation = read_u32(reader)?;
Expand All @@ -348,8 +349,8 @@ fn read_queries<T: ColumnType, R: std::io::Read>(
}

fn read_virtual_cells<R: std::io::Read>(reader: &mut R) -> std::io::Result<Vec<VirtualCell>> {
let mut vcells = vec![];
let len = read_u32(reader)?;
let len = read_u32(reader)? as usize;
let mut vcells = Vec::with_capacity(len);
for _ in 0..len {
let column = Column::<Any>::fetch(reader)?;
let rotation = read_u32(reader)?;
Expand Down Expand Up @@ -384,8 +385,8 @@ fn write_fixed_columns<W: std::io::Write>(
}

fn read_fixed_columns<R: std::io::Read>(reader: &mut R) -> std::io::Result<Vec<Column<Fixed>>> {
let len = read_u32(reader)?;
let mut columns = vec![];
let len = read_u32(reader)? as usize;
let mut columns = Vec::with_capacity(len);
for _ in 0..len {
columns.push(read_fixed_column(reader)?);
}
Expand Down Expand Up @@ -426,8 +427,8 @@ fn read_cs<C: CurveAffine, R: io::Read>(reader: &mut R) -> io::Result<Constraint
let num_selectors = read_u32(reader)? as usize;
let num_fixed_columns = read_u32(reader)? as usize;

let num_advice_queries_len = read_u32(reader)?;
let mut num_advice_queries = vec![];
let num_advice_queries_len = read_u32(reader)? as usize;
let mut num_advice_queries = Vec::with_capacity(num_advice_queries_len);
for _ in 0..num_advice_queries_len {
num_advice_queries.push(read_u32(reader)? as usize);
}
Expand All @@ -440,8 +441,8 @@ fn read_cs<C: CurveAffine, R: io::Read>(reader: &mut R) -> io::Result<Constraint
let fixed_queries = read_queries::<Fixed, R>(reader, Fixed)?;
let permutation = read_arguments(reader)?;

let mut lookups = vec![];
let nb_lookup = read_u32(reader)?;
let nb_lookup = read_u32(reader)? as usize;
let mut lookups = Vec::with_capacity(nb_lookup);
for _ in 0..nb_lookup {
let input_expressions = Vec::<Expression<C::Scalar>>::fetch(reader)?;
let table_expressions = Vec::<Expression<C::Scalar>>::fetch(reader)?;
Expand Down Expand Up @@ -487,8 +488,8 @@ fn write_gates<C: CurveAffine, W: std::io::Write>(
fn read_gates<C: CurveAffine, R: std::io::Read>(
reader: &mut R,
) -> std::io::Result<Vec<Gate<C::Scalar>>> {
let nb_gates = read_u32(reader)?;
let mut gates = vec![];
let nb_gates = read_u32(reader)? as usize;
let mut gates = Vec::with_capacity(nb_gates);
for _ in 0..nb_gates {
gates.push(Gate::new_with_polys_and_queries(
Vec::<Expression<C::Scalar>>::fetch(reader)?,
Expand Down Expand Up @@ -888,6 +889,7 @@ impl<'a, C: CurveAffine> AssignWitnessCollection<'a, C> {
.map(&fd)
.unwrap()
};
//TODO: to be optimized
let s: &[C::Scalar] = unsafe {
std::slice::from_raw_parts(mmap.as_ptr() as *const C::Scalar, 1 << params.k)
};
Expand Down
6 changes: 3 additions & 3 deletions halo2_proofs/src/plonk/evaluation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -716,15 +716,15 @@ impl<C: CurveAffine> Evaluator<C> {
// Polynomials required for this lookup.
// Calculated here so these only have to be kept in memory for the short time
// they are actually needed.
let product_coset = pk.vk.domain.coeff_to_extended(lookup.product_poly.clone());
let product_coset = pk.vk.domain.coeff_to_extended(&lookup.product_poly);
let permuted_input_coset = pk
.vk
.domain
.coeff_to_extended(lookup.permuted_input_poly.clone());
.coeff_to_extended(&lookup.permuted_input_poly);
let permuted_table_coset = pk
.vk
.domain
.coeff_to_extended(lookup.permuted_table_poly.clone());
.coeff_to_extended(&lookup.permuted_table_poly);

parallelize(&mut values, |values, start| {
for (i, value) in values.iter_mut().enumerate() {
Expand Down
27 changes: 19 additions & 8 deletions halo2_proofs/src/plonk/keygen.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,7 +323,7 @@ where
#[cfg(not(feature = "cuda"))]
let fixed_cosets = fixed_polys
.iter()
.map(|poly| vk.domain.coeff_to_extended(poly.clone()))
.map(|poly| vk.domain.coeff_to_extended(poly))
.collect();

let timer = start_timer!(|| "assembly build pkey");
Expand All @@ -339,7 +339,12 @@ where
l0[0] = C::Scalar::one();
let l0 = vk.domain.lagrange_to_coeff(l0);
#[cfg(not(feature = "cuda"))]
let l0 = vk.domain.coeff_to_extended(l0);
let l0 = {
let l0_ext = vk.domain.coeff_to_extended(&l0);
drop(l0);
l0_ext
};


// Compute l_blind(X) which evaluates to 1 for each blinding factor row
// and 0 otherwise over the domain.
Expand All @@ -348,14 +353,15 @@ where
*evaluation = C::Scalar::one();
}
let l_blind = vk.domain.lagrange_to_coeff(l_blind);
let l_blind_extended = vk.domain.coeff_to_extended(l_blind);
let l_blind_extended = vk.domain.coeff_to_extended(&l_blind);
drop(l_blind);

// Compute l_last(X) which evaluates to 1 on the first inactive row (just
// before the blinding factors) and 0 otherwise over the domain
let mut l_last = vk.domain.empty_lagrange();
l_last[params.n as usize - cs.blinding_factors() - 1] = C::Scalar::one();
let l_last = vk.domain.lagrange_to_coeff(l_last);
let l_last_extended = vk.domain.coeff_to_extended(l_last.clone());
let l_last_extended = vk.domain.coeff_to_extended(&l_last);

// Compute l_active_row(X)
let one = C::Scalar::one();
Expand Down Expand Up @@ -423,7 +429,7 @@ where
#[cfg(not(feature = "cuda"))]
let fixed_cosets = fixed_polys
.iter()
.map(|poly| vk.domain.coeff_to_extended(poly.clone()))
.map(|poly| vk.domain.coeff_to_extended(poly))
.collect();

let timer = start_timer!(|| "build pk time...");
Expand All @@ -437,7 +443,11 @@ where
l0[0] = C::Scalar::one();
let l0 = vk.domain.lagrange_to_coeff(l0);
#[cfg(not(feature = "cuda"))]
let l0 = vk.domain.coeff_to_extended(l0);
let l0 = {
let l0_ext = vk.domain.coeff_to_extended(&l0);
drop(l0);
l0_ext
};

// Compute l_blind(X) which evaluates to 1 for each blinding factor row
// and 0 otherwise over the domain.
Expand All @@ -446,14 +456,15 @@ where
*evaluation = C::Scalar::one();
}
let l_blind = vk.domain.lagrange_to_coeff(l_blind);
let l_blind_extended = vk.domain.coeff_to_extended(l_blind);
let l_blind_extended = vk.domain.coeff_to_extended(&l_blind);
drop(l_blind);

// Compute l_last(X) which evaluates to 1 on the first inactive row (just
// before the blinding factors) and 0 otherwise over the domain
let mut l_last = vk.domain.empty_lagrange();
l_last[params.n as usize - cs.blinding_factors() - 1] = C::Scalar::one();
let l_last = vk.domain.lagrange_to_coeff(l_last);
let l_last_extended = vk.domain.coeff_to_extended(l_last.clone());
let l_last_extended = vk.domain.coeff_to_extended(&l_last);

// Compute l_active_row(X)
let one = C::Scalar::one();
Expand Down
8 changes: 5 additions & 3 deletions halo2_proofs/src/plonk/permutation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,9 +88,11 @@ impl<C: CurveAffine> VerifyingKey<C> {
}

pub fn read<R: io::Read>(reader: &mut R, argument: &Argument) -> io::Result<Self> {
let commitments = (0..argument.columns.len())
.map(|_| C::read(reader))
.collect::<Result<Vec<_>, _>>()?;
let len = argument.columns.len();
let mut commitments = Vec::with_capacity(len);
for _ in 0..len {
commitments.push(C::read(reader)?);
}
Ok(VerifyingKey { commitments })
}
}
Expand Down
16 changes: 10 additions & 6 deletions halo2_proofs/src/plonk/permutation/keygen.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,24 @@ impl Assembly {
pub(crate) fn new(n: usize, p: &Argument) -> Self {
// Initialize the copy vector to keep track of copy constraints in all
// the permutation arguments.
let mut columns = vec![];
let mut columns = Vec::with_capacity(p.columns.len());
for i in 0..p.columns.len() {
// Computes [(i, 0), (i, 1), ..., (i, n - 1)]
columns.push((0..n).map(|j| (i as u32, j as u32)).collect());
let mut values = Vec::with_capacity(n);
for j in 0..n {
values.push((i as u32, j as u32));
}
columns.push(values);
}

// Before any equality constraints are applied, every cell in the permutation is
// in a 1-cycle; therefore mapping and aux are identical, because every cell is
// its own distinguished element.
Assembly {
columns: p.columns.clone(),
mapping: columns.clone(),
aux: columns,
sizes: vec![vec![1usize; n]; p.columns.len()],
mapping: columns,
aux: vec![], //columns,
sizes: vec![], //vec![vec![1usize; n]; p.columns.len()],
}
}

Expand Down Expand Up @@ -202,7 +206,7 @@ impl Assembly {
#[cfg(not(feature = "cuda"))]
let cosets = polys
.par_iter()
.map(|poly| domain.coeff_to_extended(poly.clone()))
.map(|poly| domain.coeff_to_extended(poly))
.collect();

ProvingKey {
Expand Down
10 changes: 5 additions & 5 deletions halo2_proofs/src/plonk/prover.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ fn create_single_instances<
#[cfg(not(feature = "cuda"))]
let instance_cosets: Vec<_> = instance_polys
.iter()
.map(|poly| domain.coeff_to_extended(poly.clone()))
.map(|poly| domain.coeff_to_extended(poly))
.collect();

Ok(InstanceSingle {
Expand Down Expand Up @@ -578,7 +578,7 @@ pub fn create_proof<

#[cfg(not(feature = "cuda"))]
let permutation_product_coset =
domain.coeff_to_extended(permutation_product_poly.clone());
domain.coeff_to_extended(&permutation_product_poly);

let permutation_product_commitment =
permutation_product_commitment_projective.to_affine();
Expand Down Expand Up @@ -643,7 +643,7 @@ pub fn create_proof<
#[cfg(not(feature = "cuda"))]
let advice_cosets: Vec<_> = advice_polys
.iter()
.map(|poly| domain.coeff_to_extended(poly.clone()))
.map(|poly| domain.coeff_to_extended(poly))
.collect();

AdviceSingle::<C> {
Expand Down Expand Up @@ -1068,7 +1068,7 @@ pub fn create_proof_from_witness<

#[cfg(not(feature = "cuda"))]
let permutation_product_coset =
domain.coeff_to_extended(permutation_product_poly.clone());
domain.coeff_to_extended(&permutation_product_poly);

let permutation_product_commitment =
permutation_product_commitment_projective.to_affine();
Expand Down Expand Up @@ -1133,7 +1133,7 @@ pub fn create_proof_from_witness<
#[cfg(not(feature = "cuda"))]
let advice_cosets: Vec<_> = advice_polys
.iter()
.map(|poly| domain.coeff_to_extended(poly.clone()))
.map(|poly| domain.coeff_to_extended(poly))
.collect();

AdviceSingle::<C> {
Expand Down
Loading