Skip to content

Add low-memory mode for batch polynomial opening#327

Merged
niconiconi merged 3 commits intomainfrom
hc/batching_low_memory
Mar 12, 2026
Merged

Add low-memory mode for batch polynomial opening#327
niconiconi merged 3 commits intomainfrom
hc/batching_low_memory

Conversation

@hczphn
Copy link
Copy Markdown
Contributor

@hczphn hczphn commented Feb 22, 2026

Summary

  • Add low-memory Cargo feature to poly_commit crate
  • When enabled, prover_merge_points moves tilde_gs into sumcheck (zero clone) and recomputes them for g_prime, reducing peak memory by eliminating one full copy of the polynomial data
  • When disabled (default), behavior is identical to before — full backward compatibility
  • Extract build_tilde_gs helper to avoid code duplication between the two paths

Design

The optimization trades compute time for memory:

  • Default (low-memory off): clone tilde_gs into sumcheck_poly, keep original for g_prime — same as current behavior
  • Low-memory (low-memory on): move tilde_gs into sumcheck_poly, recompute from polys + eq_t_i for g_prime — peak memory reduced by ~size(tilde_gs)

Usage:

# In downstream Cargo.toml
poly_commit = { path = "...", features = ["low-memory"] }

The low_memory parameter is also exposed as a bool on prover_merge_points for callers that want runtime control.

Test plan

  • cargo test -p poly_commit (default, no feature)
  • cargo test -p poly_commit --features low-memory
  • Proof output should be identical in both modes (deterministic)

🤖 Generated with Claude Code

@gemini-code-assist
Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @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 introduces a new low-memory mode for batch polynomial opening, primarily targeting the poly_commit crate. The core change involves optimizing memory usage in the prover_merge_points function by strategically moving or recomputing polynomial data, rather than cloning it, when the low-memory feature is active. This optimization trades compute time for reduced peak memory consumption, offering a configurable approach to resource management without breaking existing functionality. Additionally, it includes improvements to Structured Reference String (SRS) serialization and deserialization for better performance.

Highlights

  • New low-memory Cargo feature: A new low-memory Cargo feature has been added to the poly_commit crate. When enabled, it optimizes memory usage during batch polynomial opening.
  • Memory optimization in prover_merge_points: When low-memory is enabled, the prover_merge_points function moves tilde_gs into sumcheck (avoiding a clone) and recomputes them for g_prime construction. This significantly reduces peak memory by eliminating one full copy of the polynomial data.
  • Backward compatibility: The default behavior (when low-memory is disabled) remains identical to before, ensuring full backward compatibility.
  • Code refactoring for tilde_gs: The logic for building tilde_gs has been extracted into a helper function, build_tilde_gs, to avoid code duplication between the two memory paths.
  • Runtime control for low-memory mode: The low_memory parameter is also exposed as a boolean on prover_merge_points for callers who require runtime control over this optimization.
  • Optimized SRS serialization/deserialization: Custom ExpSerde implementations for CoefFormUniKZGSRS now use parallel, uncompressed serialization/deserialization and buffered I/O, improving performance for SRS loading.

🧠 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
  • gkr_engine/src/poly_commit/definition.rs
    • Added SRS_IO_BUFFER_CAPACITY constant for buffered I/O.
    • Integrated BufReader and BufWriter for more efficient SRS serialization and deserialization.
  • poly_commit/Cargo.toml
    • Introduced a new low-memory Cargo feature.
  • poly_commit/src/batching.rs
    • Added a low_memory: bool parameter to prover_merge_points to control memory usage.
    • Extracted the logic for building tilde_gs into a new helper function build_tilde_gs.
    • Modified the sumcheck_poly construction to either move (low_memory enabled) or clone (low_memory disabled) tilde_gs.
    • Conditionally recomputed tilde_gs for g_prime construction when low_memory mode is active.
  • poly_commit/src/hyrax/hyrax_impl.rs
    • Updated prover_merge_points call to pass the cfg!(feature = "low-memory") flag.
  • poly_commit/src/kzg/uni_kzg/hyper_kzg.rs
    • Updated prover_merge_points call to pass the cfg!(feature = "low-memory") flag.
  • poly_commit/src/kzg/uni_kzg/structs_kzg.rs
    • Removed ExpSerde derive macro from CoefFormUniKZGSRS.
    • Implemented a custom ExpSerde trait for CoefFormUniKZGSRS to support parallel and uncompressed serialization/deserialization of G1 points, improving performance and avoiding square root computations.
    • Added necessary use statements for halo2curves::group, rayon, and std::io.
  • poly_commit/src/utils.rs
    • Refactored expander_pcs_init_testing_only to call a new expander_pcs_init_with_srs_path function.
    • Introduced expander_pcs_init_with_srs_path to provide an option for caching SRS to a specified file path.
  • serdes/src/serdes.rs
    • Optimized Vec deserialization by initializing Vec with with_capacity based on the deserialized length.
  • sumcheck/src/sumcheck_generic.rs
    • Modified SumCheck::prove to take SumOfProductsPoly by value, enabling zero-copy operations in low-memory mode.
  • sumcheck/src/sumcheck_generic/prover.rs
    • Added prover_init_owned function to initialize the prover state by consuming the SumOfProductsPoly.
  • sumcheck/src/sumcheck_generic/tests.rs
    • Adjusted calls to SumCheck::prove to clone mle_list where necessary, aligning with the new ownership semantics.
Activity
  • No human activity has been recorded on this pull request yet.
Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

The pull request introduces a low-memory mode for batch polynomial opening, which effectively trades computation for reduced peak memory usage. It also includes significant optimizations for SRS loading through buffered I/O and parallel deserialization of uncompressed points. While the core logic for the low-memory mode is sound, there are issues regarding test isolation and portability in the utility functions, and a potential security/correctness risk in the SRS loading implementation due to the use of unchecked point parsing.

Comment on lines +112 to +114
E::G1Affine::from_uncompressed_unchecked(&uncompressed)
.into_option()
.ok_or(serdes::SerdeError::DeserializeError)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-medium medium

Using from_uncompressed_unchecked skips critical validation checks, such as verifying that the points are actually on the curve. While this is faster, it introduces a risk if the SRS file is corrupted or tampered with. Since the SRS is a foundational security component, it is safer to use from_uncompressed which performs these checks, especially since the uncompressed format already avoids the most expensive part (square root computation).

Suggested change
E::G1Affine::from_uncompressed_unchecked(&uncompressed)
.into_option()
.ok_or(serdes::SerdeError::DeserializeError)
E::G1Affine::from_uncompressed(&uncompressed)
.into_option()
.ok_or(serdes::SerdeError::DeserializeError)

Comment thread poly_commit/src/utils.rs
Comment on lines +19 to +20
let srs_path = format!("/tmp/kzg_srs_{}.bin", n_input_vars);
expander_pcs_init_with_srs_path::<FieldConfig, PCS>(n_input_vars, mpi_config, Some(&srs_path))
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Hardcoding a path in /tmp with a fixed filename is problematic for several reasons:

  1. Portability: /tmp is not available on all operating systems (e.g., Windows).
  2. Test Isolation: Multiple tests running in parallel or multiple users running tests on the same machine will conflict over the same file. If a test leaves a corrupted file, subsequent tests will fail.
  3. Regression: The docstring for expander_pcs_init_testing_only states it should 'always regenerate SRS', but this change enables caching by default.

Consider passing None to expander_pcs_init_with_srs_path to maintain the 'always regenerate' behavior for this helper, or use a more robust temporary file mechanism if caching is desired for specific tests.

Suggested change
let srs_path = format!("/tmp/kzg_srs_{}.bin", n_input_vars);
expander_pcs_init_with_srs_path::<FieldConfig, PCS>(n_input_vars, mpi_config, Some(&srs_path))
expander_pcs_init_with_srs_path::<FieldConfig, PCS>(n_input_vars, mpi_config, None)

@hczphn hczphn changed the base branch from main to hc/sumcheck_move_semantics February 22, 2026 17:27
@niconiconi niconiconi force-pushed the hc/sumcheck_move_semantics branch from f4c9a82 to 71777e7 Compare March 12, 2026 02:42
hczphn and others added 3 commits March 12, 2026 10:44
When the `low-memory` feature is enabled, `prover_merge_points` moves
tilde_gs into sumcheck (instead of cloning) and recomputes them
afterwards for g_prime construction. This eliminates one full copy of
the polynomial data at the cost of redundant computation.

- Add `low-memory` feature flag to poly_commit
- Add `low_memory: bool` parameter to `prover_merge_points`
- Callers default to `cfg!(feature = "low-memory")`
- Extract `build_tilde_gs` helper to avoid code duplication
- No behavior change when feature is not enabled

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The compiler cannot statically verify that tilde_gs is not moved in
both if/else branches. Wrap in Option so take() explicitly transfers
ownership in the low_memory path, while as_ref() borrows in the
default path.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@niconiconi niconiconi force-pushed the hc/batching_low_memory branch from bf48823 to f9dcdf9 Compare March 12, 2026 02:44
@hczphn hczphn marked this pull request as draft March 12, 2026 02:55
@hczphn hczphn changed the base branch from hc/sumcheck_move_semantics to main March 12, 2026 02:56
@hczphn hczphn marked this pull request as ready for review March 12, 2026 02:56
@niconiconi niconiconi merged commit 245a75e into main Mar 12, 2026
0 of 9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants