Feature: Orion PCS Integration (#153)

--------- Signed-off-by: Zhiyong Fang <[email protected]> Co-authored-by: Zhiyong Fang <[email protected]>
PolyhedraZK · Jan 22, 2025 · 9a4248d · 9a4248d
1 parent cdc32b0
commit 9a4248d
Show file tree

Hide file tree

Showing 17 changed files with 772 additions and 20 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -41,7 +41,7 @@ jobs:
       matrix:
         include:
           - os: macos-latest
-          - os: ubuntu-latest
+          # - os: ubuntu-latest
 
     steps:
       - name: Checkout code

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/config/config_macros/Cargo.toml b/config/config_macros/Cargo.toml
@@ -15,6 +15,7 @@ quote = "1.0"     # For generating code
 proc-macro2 = "1.0"  # For working with tokens
 
 [dev-dependencies]
+gf2 = { path = "../../arith/gf2" }
 mersenne31 = { path = "../../arith/mersenne31/" }
 
 [lib]

diff --git a/config/config_macros/src/lib.rs b/config/config_macros/src/lib.rs
@@ -88,6 +88,7 @@ fn parse_fiat_shamir_hash_type(
 }
 
 fn parse_polynomial_commitment_type(
+    field_type: &str,
     field_config: &str,
     transcript_type: &str,
     polynomial_commitment_type: ExprPath,
@@ -99,11 +100,19 @@ fn parse_polynomial_commitment_type(
         .expect("Empty path for polynomial commitment type");
 
     let pcs_type_str = binding.ident.to_string();
-    match pcs_type_str.as_str() {
-        "Raw" => (
+    match (pcs_type_str.as_str(), field_type) {
+        ("Raw", _) => (
             "Raw".to_owned(),
             format!("RawExpanderGKR::<{field_config}, {transcript_type}>").to_owned(),
         ),
+        ("Orion", "GF2") => (
+            "Orion".to_owned(),
+            format!("OrionPCSForGKR::<{field_config}, GF2x128, {transcript_type}>").to_owned(),
+        ),
+        ("Orion", "M31") => (
+            "Orion".to_owned(),
+            format!("OrionPCSForGKR::<{field_config}, M31x16, {transcript_type}>").to_owned(),
+        ),
         _ => panic!("Unknown polynomial commitment type in config macro expansion"),
     }
 }
@@ -134,7 +143,8 @@ fn declare_gkr_config_impl(input: proc_macro::TokenStream) -> proc_macro::TokenS
     let (fiat_shamir_hash_type, transcript_type) =
         parse_fiat_shamir_hash_type(&field_type, &field_config, fiat_shamir_hash_type_expr);
     let (polynomial_commitment_enum, polynomial_commitment_type) = parse_polynomial_commitment_type(
-        field_config.as_str(),
+        &field_type,
+        &field_config,
         &transcript_type,
         polynomial_commitment_type,
     );

diff --git a/config/config_macros/tests/macro_expansion.rs b/config/config_macros/tests/macro_expansion.rs
@@ -8,9 +8,10 @@ use gkr_field_config::FieldType;
 use config::GKRConfig;
 use config_macros::declare_gkr_config;
 use field_hashers::{MiMC5FiatShamirHasher, PoseidonFiatShamirHasher};
+use gf2::GF2x128;
 use gkr_field_config::{BN254Config, GF2ExtConfig, GKRFieldConfig, M31ExtConfig};
 use mersenne31::M31x16;
-use poly_commit::raw::RawExpanderGKR;
+use poly_commit::{OrionPCSForGKR, RawExpanderGKR};
 use transcript::{BytesHashTranscript, FieldHashTranscript, Keccak256hasher, SHA256hasher};
 
 fn print_type_name<Cfg: GKRConfig>() {
@@ -26,11 +27,17 @@ fn main() {
         PolynomialCommitmentType::Raw
     );
     declare_gkr_config!(
-        M31PoseidonConfig,
+        M31PoseidonRawConfig,
         FieldType::M31,
         FiatShamirHashType::Poseidon,
         PolynomialCommitmentType::Raw
     );
+    declare_gkr_config!(
+        M31PoseidonOrionConfig,
+        FieldType::M31,
+        FiatShamirHashType::Poseidon,
+        PolynomialCommitmentType::Orion
+    );
     declare_gkr_config!(
         BN254MIMCConfig,
         FieldType::BN254,
@@ -43,8 +50,17 @@ fn main() {
         FiatShamirHashType::Keccak256,
         PolynomialCommitmentType::Raw
     );
+    declare_gkr_config!(
+        GF2Keccak256OrionConfig,
+        FieldType::GF2,
+        FiatShamirHashType::Keccak256,
+        PolynomialCommitmentType::Orion
+    );
 
     print_type_name::<M31Sha256Config>();
+    print_type_name::<M31PoseidonRawConfig>();
+    print_type_name::<M31PoseidonOrionConfig>();
     print_type_name::<BN254MIMCConfig>();
     print_type_name::<GF2Keccak256Config>();
+    print_type_name::<GF2Keccak256OrionConfig>();
 }
diff --git a/config/mpi_config/src/lib.rs b/config/mpi_config/src/lib.rs
@@ -134,16 +134,16 @@ impl MPIConfig {
 
     /// Return an u8 vector sharing THE SAME MEMORY SLOT with the input.
     #[inline]
-    unsafe fn vec_to_u8_bytes<F: Field>(vec: &Vec<F>) -> Vec<u8> {
+    unsafe fn vec_to_u8_bytes<F: Sized>(vec: &Vec<F>) -> Vec<u8> {
         Vec::<u8>::from_raw_parts(
             vec.as_ptr() as *mut u8,
-            vec.len() * F::SIZE,
-            vec.capacity() * F::SIZE,
+            vec.len() * size_of::<F>(),
+            vec.capacity() * size_of::<F>(),
         )
     }
 
     #[allow(clippy::collapsible_else_if)]
-    pub fn gather_vec<F: Field>(&self, local_vec: &Vec<F>, global_vec: &mut Vec<F>) {
+    pub fn gather_vec<F: Sized + Clone>(&self, local_vec: &Vec<F>, global_vec: &mut Vec<F>) {
         unsafe {
             if self.world_size == 1 {
                 *global_vec = local_vec.clone()

diff --git a/gkr/src/executor.rs b/gkr/src/executor.rs
@@ -12,9 +12,10 @@ use config::{
 };
 use config_macros::declare_gkr_config;
 use field_hashers::{MiMC5FiatShamirHasher, PoseidonFiatShamirHasher};
+use gf2::GF2x128;
 use gkr_field_config::{BN254Config, GF2ExtConfig, GKRFieldConfig, M31ExtConfig};
 use mersenne31::M31x16;
-use poly_commit::{expander_pcs_init_testing_only, raw::RawExpanderGKR};
+use poly_commit::{expander_pcs_init_testing_only, raw::RawExpanderGKR, OrionPCSForGKR};
 use rand::SeedableRng;
 use rand_chacha::ChaCha12Rng;
 use transcript::{BytesHashTranscript, FieldHashTranscript, SHA256hasher};
@@ -289,5 +290,5 @@ declare_gkr_config!(
     pub GF2ExtConfigSha2,
     FieldType::GF2,
     FiatShamirHashType::SHA256,
-    PolynomialCommitmentType::Raw
+    PolynomialCommitmentType::Orion
 );
diff --git a/gkr/src/main.rs b/gkr/src/main.rs
@@ -7,10 +7,11 @@ use circuit::Circuit;
 use clap::Parser;
 use config::{Config, GKRConfig, GKRScheme};
 use config_macros::declare_gkr_config;
+use gf2::GF2x128;
 use gkr_field_config::{BN254Config, GF2ExtConfig, GKRFieldConfig, M31ExtConfig};
 use mpi_config::MPIConfig;
 
-use poly_commit::{expander_pcs_init_testing_only, raw::RawExpanderGKR};
+use poly_commit::{expander_pcs_init_testing_only, raw::RawExpanderGKR, OrionPCSForGKR};
 use rand::SeedableRng;
 use rand_chacha::ChaCha12Rng;
 use transcript::{BytesHashTranscript, SHA256hasher};
@@ -72,7 +73,7 @@ fn main() {
         GF2ExtConfigSha2,
         FieldType::GF2,
         FiatShamirHashType::SHA256,
-        PolynomialCommitmentType::Raw
+        PolynomialCommitmentType::Orion
     );
 
     match args.field.as_str() {

diff --git a/gkr/src/main_mpi.rs b/gkr/src/main_mpi.rs
@@ -4,8 +4,9 @@ use config::{Config, GKRConfig, GKRScheme};
 use config_macros::declare_gkr_config;
 use mpi_config::MPIConfig;
 
+use gf2::GF2x128;
 use gkr_field_config::{BN254Config, GF2ExtConfig, GKRFieldConfig, M31ExtConfig};
-use poly_commit::{expander_pcs_init_testing_only, raw::RawExpanderGKR};
+use poly_commit::{expander_pcs_init_testing_only, raw::RawExpanderGKR, OrionPCSForGKR};
 use rand::SeedableRng;
 use rand_chacha::ChaCha12Rng;
 use transcript::{BytesHashTranscript, SHA256hasher};
@@ -63,7 +64,7 @@ fn main() {
         GF2ExtConfigSha2,
         FieldType::GF2,
         FiatShamirHashType::SHA256,
-        PolynomialCommitmentType::Raw
+        PolynomialCommitmentType::Orion
     );
 
     match args.field.as_str() {

diff --git a/gkr/src/tests/gkr_correctness.rs b/gkr/src/tests/gkr_correctness.rs
@@ -8,11 +8,11 @@ use circuit::Circuit;
 use config::{Config, FiatShamirHashType, GKRConfig, GKRScheme, PolynomialCommitmentType};
 use config_macros::declare_gkr_config;
 use field_hashers::{MiMC5FiatShamirHasher, PoseidonFiatShamirHasher};
+use gf2::GF2x128;
 use gkr_field_config::{BN254Config, FieldType, GF2ExtConfig, GKRFieldConfig, M31ExtConfig};
 use mersenne31::M31x16;
 use mpi_config::{root_println, MPIConfig};
-use poly_commit::expander_pcs_init_testing_only;
-use poly_commit::raw::RawExpanderGKR;
+use poly_commit::{expander_pcs_init_testing_only, OrionPCSForGKR, RawExpanderGKR};
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha12Rng;
 use sha2::Digest;
@@ -67,12 +67,24 @@ fn test_gkr_correctness() {
         FiatShamirHashType::MIMC5,
         PolynomialCommitmentType::Raw
     );
+    declare_gkr_config!(
+        C7,
+        FieldType::GF2,
+        FiatShamirHashType::Keccak256,
+        PolynomialCommitmentType::Orion,
+    );
     declare_gkr_config!(
         C8,
         FieldType::M31,
         FiatShamirHashType::Poseidon,
         PolynomialCommitmentType::Raw,
     );
+    declare_gkr_config!(
+        C9,
+        FieldType::M31,
+        FiatShamirHashType::Poseidon,
+        PolynomialCommitmentType::Orion,
+    );
 
     test_gkr_correctness_helper(
         &Config::<C0>::new(GKRScheme::Vanilla, mpi_config.clone()),
@@ -102,10 +114,18 @@ fn test_gkr_correctness() {
         &Config::<C6>::new(GKRScheme::Vanilla, mpi_config.clone()),
         Some("../data/gkr_proof.txt"),
     );
+    test_gkr_correctness_helper(
+        &Config::<C7>::new(GKRScheme::Vanilla, mpi_config.clone()),
+        None,
+    );
     test_gkr_correctness_helper(
         &Config::<C8>::new(GKRScheme::Vanilla, mpi_config.clone()),
         None,
     );
+    test_gkr_correctness_helper(
+        &Config::<C9>::new(GKRScheme::Vanilla, mpi_config.clone()),
+        None,
+    );
 
     MPIConfig::finalize();
 }

diff --git a/poly_commit/src/orion.rs b/poly_commit/src/orion.rs
@@ -26,6 +26,14 @@ pub use simd_field_impl::{
 #[cfg(test)]
 mod simd_field_tests;
 
+mod simd_field_agg_impl;
+
+#[cfg(test)]
+mod simd_field_agg_tests;
+
+mod pcs_for_expander_gkr;
+pub use pcs_for_expander_gkr::OrionPCSForGKR;
+
 mod pcs_trait_impl;
 pub use pcs_trait_impl::{OrionBaseFieldPCS, OrionSIMDFieldPCS};