Reduce memory in SumCheck::prove via move semantics#326
Conversation
Summary of ChangesHello @hczphn, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly reduces memory consumption during batched polynomial commitment opening by refactoring the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces significant memory and performance optimizations, primarily by refactoring the SumCheck::prove function to leverage move semantics for SumOfProductsPoly, thereby eliminating an unnecessary clone. Additionally, it enhances SRS (Structured Reference String) handling by implementing buffered I/O for serialization/deserialization and parallel processing for KZG SRS, which will lead to faster loading times for large SRS files. The changes also include minor performance improvements like pre-allocating vector capacity during deserialization and better code organization for PCS initialization functions.
| buffer | ||
| .par_chunks_mut(point_size) | ||
| .zip(self.powers_of_tau.par_iter()) | ||
| .for_each(|(chunk, point)| { | ||
| let uncompressed = point.to_uncompressed(); | ||
| chunk.copy_from_slice(uncompressed.as_ref()); | ||
| }); |
There was a problem hiding this comment.
Parallelizing the conversion of points to uncompressed format using par_chunks_mut and par_iter from rayon is a significant performance optimization for serializing large SRS. This will greatly reduce the time taken for SRS serialization.
| pub fn prover_init_owned(polynomials: SumOfProductsPoly<F>) -> Self { | ||
| let num_vars = polynomials.num_vars(); | ||
| let init_sum_of_vals: Vec<F> = polynomials | ||
| .f_and_g_pairs | ||
| .par_iter() | ||
| .map(|(f, g)| { | ||
| f.coeffs | ||
| .iter() | ||
| .zip(g.coeffs.iter()) | ||
| .map(|(&f, &g)| f * g) | ||
| .sum::<F>() | ||
| }) | ||
| .collect(); | ||
| let eq_prefix = vec![F::one(); polynomials.f_and_g_pairs.len()]; | ||
| Self { | ||
| challenges: Vec::with_capacity(num_vars), | ||
| round: 0, | ||
| init_num_vars: num_vars, | ||
| mle_list: polynomials, | ||
| init_sum_of_vals, | ||
| eq_prefix, | ||
| } |
There was a problem hiding this comment.
The prover_init_owned function is correctly implemented to accept polynomials by value, directly moving it into mle_list. This is the key mechanism for the memory optimization, as it prevents the SumOfProductsPoly from being cloned, which was identified as a significant memory consumer.
| let num_vars = poly_list.num_vars(); | ||
|
|
||
| let mut prover_state = IOPProverState::prover_init(poly_list); | ||
| let mut prover_state = IOPProverState::prover_init_owned(poly_list); |
There was a problem hiding this comment.
Calling prover_init_owned instead of prover_init ensures that the SumOfProductsPoly is moved into the prover state, avoiding a costly clone operation. This is a direct implementation of the PR's memory optimization objective.
| /// Takes ownership of poly_list to avoid cloning internally. | ||
| pub fn prove( | ||
| poly_list: &SumOfProductsPoly<F>, | ||
| poly_list: SumOfProductsPoly<F>, |
There was a problem hiding this comment.
Changing the poly_list parameter from a reference (&SumOfProductsPoly<F>) to a value (SumOfProductsPoly<F>) is the core change enabling move semantics. This directly addresses the memory optimization goal by transferring ownership and preventing an implicit clone within the prove function.
| sumcheck_poly.add_pair(tilde_g.clone(), tilde_eq); | ||
| } | ||
| let proof = SumCheck::<C::Scalar>::prove(&sumcheck_poly, transcript); | ||
| let proof = SumCheck::<C::Scalar>::prove(sumcheck_poly, transcript); |
There was a problem hiding this comment.
Changing SumCheck::<C::Scalar>::prove to take sumcheck_poly by value (move semantics) directly addresses the PR's goal of reducing memory consumption by avoiding an implicit clone. This is a crucial optimization for large polynomial workloads.
| /// Initialize the prover state by taking ownership of the polynomials, | ||
| /// avoiding the clone in `prover_init`. |
There was a problem hiding this comment.
The doc comment for prover_init_owned clearly explains its purpose: to initialize the prover state by taking ownership of polynomials and avoiding an internal clone. This is good documentation for a new, optimized function.
| Ok(file) => { | ||
| // file exists; deserialize SRS from file | ||
| Self::SRS::deserialize_from(&mut file).unwrap_or_else(|_| { | ||
| let mut reader = BufReader::with_capacity(SRS_IO_BUFFER_CAPACITY, file); |
There was a problem hiding this comment.
Introducing BufReader with a specified capacity is a good performance optimization for reading large SRS files. It reduces the number of system calls and can significantly speed up deserialization.
| // prover | ||
| let mut transcript = BytesHashTranscript::<Keccak256hasher>::new(); | ||
| let proof = SumCheck::<Fr>::prove(&mle_list, &mut transcript); | ||
| let proof = SumCheck::<Fr>::prove(mle_list.clone(), &mut transcript); |
There was a problem hiding this comment.
The addition of .clone() here is necessary because mle_list is used again later in the test for evaluation (mle_list.evaluate(&subclaim.point)). This correctly adapts the test to the new prove signature while preserving the original test logic.
| let claimed_sum = mle_list.sum(); | ||
|
|
||
| let proof = SumCheck::prove(&mle_list, &mut T::new()); | ||
| let proof = SumCheck::prove(mle_list.clone(), &mut T::new()); |
There was a problem hiding this comment.
Similar to the previous comment, mle_list.clone() is correctly added here because mle_list is needed for subsequent operations in the test, such as creating padded_mle_list and evaluating against subclaim.point.
| }; | ||
|
|
||
| let proof_with_padded_mle_list = SumCheck::prove(&padded_mle_list, &mut T::new()); | ||
| let proof_with_padded_mle_list = SumCheck::prove(padded_mle_list, &mut T::new()); |
There was a problem hiding this comment.
Passing padded_mle_list directly (by value) is correct here because padded_mle_list is not used after this call. This demonstrates the intended use of the new prove signature, where ownership is transferred, and no unnecessary cloning occurs.
…oning Add `prover_init_owned()` that takes `SumOfProductsPoly` by value, avoiding the expensive clone inside `prover_init()`. Change `SumCheck::prove()` to accept ownership and use the new init path. This eliminates one full copy of the polynomial list during proving, which can be significant for large batched polynomial commitments. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
niconiconi
force-pushed
the
hc/sumcheck_move_semantics
branch
from
f4c9a82 to
71777e7
Compare
2 participants
Summary
prover_init_owned()toIOPProverStatethat takesSumOfProductsPolyby value, eliminating theclone()inprover_init()SumCheck::prove()signature from&SumOfProductsPoly<F>toSumOfProductsPoly<F>(ownership)prover_merge_pointsin batching to passsumcheck_polyby moveprover_init()(borrow + clone) is preserved for callers that need itMotivation
During batched polynomial commitment opening,
SumCheck::proveinternally clones the entire polynomial list viaprover_init. For large proving workloads (e.g., ZKML with many kernels), this clone can consume significant memory. By accepting ownership, we eliminate one full copy of the polynomial data.🤖 Generated with Claude Code