zf/zkcuda optim - #183
Conversation
…ilerCollection into zf/zkcuda_optim
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @Tao-Lu-X, 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!
I've focused on a significant overhaul and optimization of the zkcuda proving system. This includes a major refactoring effort to simplify generic type parameters and standardize API calls across different proving system variants. A new ExpanderNoOverSubscribe mode has been introduced to enhance performance by managing parallel execution more efficiently. Additionally, I've implemented a new configuration abstraction to streamline the setup of cryptographic components and added a comprehensive suite of integration tests to ensure the system's stability and correctness. This work also incorporates numerous dependency updates to keep the project current with the latest library versions.
Highlights
- Extensive Refactoring of ZKCuda Proving System Generics: The core
zkcudaproving system and its related data structures have undergone significant refactoring. This includes simplifying generic type parameters (e.g., removingPCSFieldfromExpanderCommitmentandExpanderProverSetup) to improve code clarity and maintainability. - Introduction of
ExpanderNoOverSubscribeProving System: A new proving system variant,ExpanderNoOverSubscribe, has been introduced. This system is designed to manage MPI parallelism more efficiently, potentially by avoiding oversubscription of CPU cores, which can lead to performance improvements. - New Configuration Abstraction for ZKCuda Components: A new
ZKCudaConfigtrait andZKCudaConfigImplstruct provide a unified and abstract way to configure different ECC and GKR engines, including options for batch Polynomial Commitment Schemes (PCS). This streamlines the setup and usage of various cryptographic backends. - New ZKCuda Integration Test Suite: A comprehensive set of new integration test binaries (
zkcuda_setup,zkcuda_prove,zkcuda_verify,zkcuda_cleanup) and an accompanyingrun.shscript have been added. This establishes a robust end-to-end testing framework for thezkcudaproving system. - Standardized
proveMethod Signature Across Proving Systems: Theprovemethod across various proving systems (e.g.,Expander,ParallelizedExpander,ExpanderPCSDefered,ExpanderNoOverSubscribe) now consistently acceptsVec<Vec<SIMDField<C>>>(an owned vector) instead of a slice reference. This API change promotes clearer ownership semantics. - Comprehensive Dependency Updates: Numerous external Rust crate dependencies (e.g.,
clap,tokio,hyper-util,rustix,serde_json) and internal crates from thePolyhedraZK/Expanderrepository have been updated to their latest versions, ensuring compatibility and incorporating upstream improvements.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 or fill out our survey to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces significant optimizations and refactoring for the zkcuda proving system. The main changes include a new ExpanderNoOverSubscribe proving system to better manage parallelism and avoid CPU oversubscription, a client-server architecture for this new system, and a major simplification of generic trait bounds throughout the codebase. The changes look solid, but I have a few suggestions for improving code quality and maintainability.
| }; | ||
| use expander_compiler::zkcuda::proving_system::{Expander, ParallelizedExpander, ProvingSystem}; | ||
| use expander_compiler::zkcuda::proving_system::expander_pcs_defered::BN254ConfigSha2UniKZG; | ||
| // use expander_compiler::zkcuda::proving_system::expander_pcs_defered::BN254ConfigSha2UniKZG; |
There was a problem hiding this comment.
This use statement is commented out and appears to be unused. It should be removed to improve code clarity.
| expander::config::ZKCudaBN254Hyrax, expander_pcs_defered::BN254ConfigSha2UniKZG, | ||
| ExpanderNoOverSubscribe, |
There was a problem hiding this comment.
The import expander_pcs_defered::BN254ConfigSha2UniKZG is unused in this file. It's good practice to remove unused imports to keep the code clean.
| expander::config::ZKCudaBN254Hyrax, expander_pcs_defered::BN254ConfigSha2UniKZG, | |
| ExpanderNoOverSubscribe, | |
| expander::config::ZKCudaBN254Hyrax, | |
| ExpanderNoOverSubscribe, |
| frontend::BN254Config, | ||
| zkcuda::proving_system::{expander_pcs_defered::BN254ConfigSha2UniKZG, ExpanderPCSDefered}, | ||
| }; | ||
| use gkr::BN254ConfigSha2Hyrax; |
There was a problem hiding this comment.
The import gkr::BN254ConfigSha2Hyrax is not used in this file and can be removed.
| pub fn prev_power_of_two(x: usize) -> usize { | ||
| if x == 0 { | ||
| return 0; | ||
| } | ||
| let mut padk: usize = 0; | ||
| while (1 << padk) <= x { | ||
| padk += 1; | ||
| } | ||
| 1 << (padk - 1) | ||
| } |
There was a problem hiding this comment.
The implementation of prev_power_of_two can be made more efficient and idiomatic by using bitwise operations instead of a loop. This avoids iteration and is generally faster.
pub fn prev_power_of_two(x: usize) -> usize {
if x == 0 {
return 0;
}
// Efficiently find the highest set bit and return the corresponding power of two.
1 << (usize::BITS - 1 - x.leading_zeros())
}| let n_local_copies = parallel_count / local_world_size; | ||
| match n_local_copies { | ||
| 1 => prove_kernel_gkr_internal::<F, F, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 2 => prove_kernel_gkr_internal::<F, BN254ConfigXN<2>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 4 => prove_kernel_gkr_internal::<F, BN254ConfigXN<4>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 8 => prove_kernel_gkr_internal::<F, BN254ConfigXN<8>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 16 => prove_kernel_gkr_internal::<F, BN254ConfigXN<16>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 32 => prove_kernel_gkr_internal::<F, BN254ConfigXN<32>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 64 => prove_kernel_gkr_internal::<F, BN254ConfigXN<64>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 128 => prove_kernel_gkr_internal::<F, BN254ConfigXN<128>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 256 => prove_kernel_gkr_internal::<F, BN254ConfigXN<256>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 512 => prove_kernel_gkr_internal::<F, BN254ConfigXN<512>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 1024 => prove_kernel_gkr_internal::<F, BN254ConfigXN<1024>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| 2048 => prove_kernel_gkr_internal::<F, BN254ConfigXN<2048>, T, ECCConfig>( | ||
| &local_mpi_config, | ||
| kernel, | ||
| commitments_values, | ||
| parallel_count, | ||
| is_broadcast, | ||
| n_bytes_profiler, | ||
| ), | ||
| _ => { | ||
| panic!("Unsupported parallel count: {parallel_count}"); | ||
| } | ||
| } |
There was a problem hiding this comment.
This large match statement contains a lot of repetitive code for handling different numbers of local copies. This pattern is a good candidate for a macro to reduce code duplication and improve maintainability. For example, you could create a macro that takes the number of copies and generates the call to prove_kernel_gkr_internal with the appropriate BN254ConfigXN type.
Also, the panic message Unsupported parallel count: {parallel_count} is a bit misleading since the match is on n_local_copies. It would be clearer if it reported the value that caused the match to fail, e.g., Unsupported number of local copies: {n_local_copies}.
let n_local_copies = parallel_count / local_world_size;
macro_rules! dispatch {
($fmulti:ty) => {
prove_kernel_gkr_internal::<F, $fmulti, T, ECCConfig>(
&local_mpi_config,
kernel,
commitments_values,
parallel_count,
is_broadcast,
n_bytes_profiler,
)
};
}
match n_local_copies {
1 => dispatch!(F),
2 => dispatch!(BN254ConfigXN<2>),
4 => dispatch!(BN254ConfigXN<4>),
8 => dispatch!(BN254ConfigXN<8>),
16 => dispatch!(BN254ConfigXN<16>),
32 => dispatch!(BN254ConfigXN<32>),
64 => dispatch!(BN254ConfigXN<64>),
128 => dispatch!(BN254ConfigXN<128>),
256 => dispatch!(BN254ConfigXN<256>),
512 => dispatch!(BN254ConfigXN<512>),
1024 => dispatch!(BN254ConfigXN<1024>),
2048 => dispatch!(BN254ConfigXN<2048>),
_ => {
panic!("Unsupported number of local copies: {n_local_copies}");
}
}|
|
||
| #[allow(clippy::zombie_processes)] | ||
| fn exec_command(cmd: &str, wait_for_completion: bool) { | ||
| println!("Executing command: {cmd}"); |
There was a problem hiding this comment.
This println! statement is useful for debugging, but it would be more idiomatic to use a logging framework like log or tracing, which are already dependencies in the project. This allows for more control over log verbosity. Consider replacing this with log::info! or log::debug!.
| .serialize_into(&mut buffer) | ||
| .expect("Failed to serialize object"); | ||
|
|
||
| println!("Object size: {}", buffer.len()); |
There was a problem hiding this comment.
This file contains several println! statements for debugging purposes (e.g., here and on lines 96, 116, 213). It's better to use a logging framework like log or tracing (which are already dependencies in the project) for this. Using log::debug! or tracing::debug! would allow these messages to be enabled or disabled based on the logging level, which is more flexible than hardcoded print statements.
4 participants
No description provided.