Skip to content

Integrate buffer abstraction to Zook - #270

Open
zkfriendly wants to merge 12 commits into
mainfrom
zkfr/integrate-buffer-abstraction
Open

Integrate buffer abstraction to Zook#270
zkfriendly wants to merge 12 commits into
mainfrom
zkfr/integrate-buffer-abstraction

Conversation

@zkfriendly

@zkfriendly zkfriendly commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR integrates the composable buffer abstraction from #266 into the current Zook, IRS, sumcheck, and code-switch protocol flow.

Once a witness is supplied as a Buffer, witness-sized prover state remains resident through:

commit → IRS/NTT → sumcheck → code-switch → subsequent rounds → basecase

This removes the two main protocol-level host bounces:

  1. IRS downloading messages and masks to construct message || mask polynomials before the NTT.
  2. Zook converting messages and covectors between Buffer and Vec on every round.

Transcript-, query-, and security-parameter-sized values intentionally remain host-side.

What changed

Buffer-native IRS/NTT boundary

  • IRS now passes resident message and mask buffers directly to the Reed–Solomon encoder.
  • The CPU NTT accesses slices privately inside the CPU backend.
  • Protocol code no longer constructs a witness-sized host poly_buf.
  • Logical message || mask ordering and row-major (evaluation, polynomial) output are preserved.
  • Buffer and DefaultRs are selected together through one active backend module.

Resident Zook state

  • Zook committed and round state now stores messages and covectors as Buffer.
  • ProtocolConfig::commit consumes an already-resident witness buffer.
  • Sumcheck outputs move directly into code-switch.
  • Code-switch outputs move directly into the next round or basecase.
  • The basecase-only path lifts the witness through mixed_lift without host materialization.

Buffer-native code-switch arithmetic

  • Code-switch messages and claim covectors are buffers.
  • OOD message evaluation, covector updates, equality weights, row reductions, and in-domain dot products use buffer operations.
  • Source evaluation rows are collapsed through one embedding-aware resident reduction path.
  • Geometric coefficient runs are generated directly as buffers.
  • OOD answers remain host-side because they are transcript-sized; their contribution is computed from the same geometric base without uploading them again.
  • The original claim owns coefficient x⁰ = 1; OOD coefficients begin at , and the in-domain run continues from there.

Buffer capabilities

This PR adds or generalizes the operations required by the resident lane:

  • resize_zeroed
  • mixed_lift
  • mixed_mat_vec
  • accumulate_geometric(points, scalars, prefix_len)
  • backend-native geometric challenge generation
  • concatenation through BufferOps

The evaluator-object batching path and UnivariateEvaluation::accumulate_many were removed in favor of direct geometric accumulation.

Transfer regression

A scoped regression checks that the resident IRS, Zook round, and code-switch paths do not reintroduce:

  • witness-sized to_slice()
  • into_vec()
  • message/covector re-uploads
  • OOD-answer uploads
  • host poly_buf construction

CPU backend internals may still use slices.

Validation

  • cargo test — 257 passed, 30 ignored
  • cargo test --features rs_in_order — 267 passed, 30 ignored
  • cargo clippy --all-targets -- -D warnings
  • cargo clippy --all-targets --features rs_in_order -- -D warnings
  • resident transfer-budget regression
  • git diff --check

Reviewer focus

The most important invariants to review are:

  1. IRS coefficient ordering remains message_chunk || mask_row.
  2. Encoded output remains row-major (evaluation_index, polynomial_index).
  3. Geometric batching remains [1, x, x², …].
  4. ZK OOD accumulation covers the full covector while in-domain accumulation covers only source.masked_message_length().
  5. CPU slice access occurs only inside backend implementations, not the resident protocol path.

Follow-ups

  • Replace the transitional IRS-specific Messages + masks interface with one validated segmented polynomial batch. Each ordered coefficient segment should describe its resident buffers, rows per buffer, and row width. This restores the abstraction that NTT sees logical polynomials rather than message/mask semantics without physically concatenating them.

  • Replace the source-text transfer regression with an instrumented AuditBackend that counts uploads, downloads, synchronizations, and full-buffer transfers.

@codspeed-hq

codspeed-hq Bot commented Jul 23, 2026

Copy link
Copy Markdown

Merging this PR will not alter performance

✅ 10 untouched benchmarks
⏩ 22 skipped benchmarks1


Comparing zkfr/integrate-buffer-abstraction (4527ff5) with main (f9ee397)

Open in CodSpeed

Footnotes

  1. 22 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports.

@zkfriendly
zkfriendly force-pushed the zkfr/integrate-buffer-abstraction branch from 33fbecf to e0aec15 Compare July 23, 2026 08:50
# Conflicts:
#	src/algebra/ntt/mod.rs
#	src/buffer/cpu.rs
#	src/buffer/mod.rs
 Key changes include:
- Introduction of a `Messages` struct to encapsulate interleaved Reed-Solomon messages, improving clarity and usability.
- Refactoring of encoding functions to utilize the new `Messages` struct, allowing for better handling of message and mask buffers.
- Updates to the `BufferMath` trait to include new methods for geometric accumulation and mixed matrix-vector products.
- Removal of redundant code and improved type safety by leveraging the buffer abstraction in various protocols, including Zook and Sumcheck.
@zkfriendly

Copy link
Copy Markdown
Collaborator Author

@shreyas-londhe ready for review

@zkfriendly zkfriendly changed the title Change visibility and minor refactors Integrate buffer abstraction to Zook Aug 10, 2026
@zkfriendly

Copy link
Copy Markdown
Collaborator Author

@shreyas-londhe bumping this

@shreyas-londhe shreyas-londhe left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks for the PR! Overall looks correct, left some minor nits.

Comment thread src/protocols/code_switch.rs
Comment thread src/buffer/mod.rs
Comment thread src/algebra/ntt/mod.rs
Comment on lines +136 to +141
fn interleaved_encode(
&self,
messages: Messages<'_, F>,
masks: &Buffer<F>,
codeword_length: usize,
) -> Buffer<F>;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Pairing rules belong on the trait, not in irs_commit. This method now owns the message ‖ mask pairing that irs_commit::commit used to do, but the two orderings that make it well-defined aren't stated here, logical polynomial index is vector_index * interleaving_depth + chunk_index, and masks[i * mask_length + c] belongs to polynomial i. Both are currently documented only at irs_commit.rs:379-382, i.e. in one caller.

For a trait meant to be reimplemented by a GPU backend, that leaves a backend author guessing between vector-major and chunk-major flattening, with nothing asserting which they picked. Worth stating both on interleaved_encode (or on Messages).

Comment thread tests/resident_transfer_budget.rs Outdated
Comment thread src/protocols/zook/prover.rs Outdated
Comment thread src/protocols/zook/prover.rs Outdated
Comment thread src/algebra/ntt/cooley_tukey.rs Outdated
Comment thread src/algebra/ntt/cooley_tukey.rs Outdated
Comment thread src/algebra/ntt/mod.rs Outdated
@BornPsych
BornPsych self-requested a review August 12, 2026 10:30
@zkfriendly
zkfriendly force-pushed the zkfr/integrate-buffer-abstraction branch from b3187d8 to 8804e80 Compare August 18, 2026 16:00
@zkfriendly

Copy link
Copy Markdown
Collaborator Author

@BornPsych I addressed/resolved comments. The NTT interface changes are the bulk of the new diff. I originally planned the ntt interface refinement as a follow-up PR as mentioned in the description, but since we had comments here, I made the changes but we might want to iterate on the design. Next I'll update the provekit side changes

@BornPsych BornPsych left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

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.

3 participants