From 655783c3b8466af34587d43f333ca967604e1954 Mon Sep 17 00:00:00 2001 From: Arun Jangra Date: Mon, 25 May 2026 21:33:07 +0530 Subject: [PATCH 01/35] zkWHIR 3.0 : Updated protocols for whir round integration and added mask proximity protocol (part 2) (#251) * feat : updated sumcheck according to construction 6.3 in paper * feat : updated sumcheck and code switch for integration in main whir orchestrator * feat : added mask proximity protocol * refactor : cleanup and added few tests * tiny change * feat : added assert for interleaving depth == 1 in mask proximity protocol * resolved comments * feat : added some tests in code switch * feat : updated code switch protocol for parameter selection logic * feat : added support for interleaving depth > 1 for code switch * fix : tests in runner * nits: comment and type * feat : comments resolved, added asserts, fixed test coverage * added security comment for verifier * fix: rust nits --- .../protocols/code_switch.txt | 3 + .../protocols/mask_proximity.txt | 7 + proptest-regressions/protocols/sumcheck.txt | 7 + src/algebra/sumcheck.rs | 2 +- src/protocols/basecase.rs | 79 +- src/protocols/code_switch.rs | 751 +++++++++++------- src/protocols/irs_commit.rs | 24 +- src/protocols/mask_proximity.rs | 525 ++++++++++++ src/protocols/mod.rs | 1 + src/protocols/sumcheck.rs | 111 +-- src/protocols/whir/prover.rs | 6 +- src/protocols/whir/verifier.rs | 9 +- src/transcript/mock_sponge.rs | 6 +- 13 files changed, 1168 insertions(+), 363 deletions(-) create mode 100644 proptest-regressions/protocols/mask_proximity.txt create mode 100644 proptest-regressions/protocols/sumcheck.txt create mode 100644 src/protocols/mask_proximity.rs diff --git a/proptest-regressions/protocols/code_switch.txt b/proptest-regressions/protocols/code_switch.txt index 67bf0158..fdc02672 100644 --- a/proptest-regressions/protocols/code_switch.txt +++ b/proptest-regressions/protocols/code_switch.txt @@ -7,3 +7,6 @@ cc d1b11d6a89f5f0e74e62779bf54948e664f6780433d440ac3d977744776ed358 # shrinks to seed = 0, source_size = 3, src_lir = 1, tgt_lir = 1 cc e5c94e6ed6e6d91a8725206056b46ea1c13efd83e0c206e32f20a01929ddf654 # shrinks to seed = 0, source_size = 12, src_lir = 1, tgt_lir = 1 cc c5cdfb9b93481bf4819a4207587a6bf93c85cce8b7e2e999a9478cfc0f78cc0d # shrinks to seed = 0, source_size = 8, src_lir = 2, tgt_lir = 1 +cc ba9d37bf3442e60c7741c3593749c9413bb460cc51999b5200e54c9985794559 # shrinks to seed = 0, msg_len = 1, mask_length = 2, ood_samples = 1 +cc 582495d0662a9c5083a7bad07e789df2e433d36d9afeb27f21bf593fb2d6a9e0 # shrinks to seed = 0, config = Config { source: Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 1, vector_size: 4, mask_length: 3, codeword_length: 8, interleaving_depth: 2, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 2, leaf_hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076, merkle_tree: Config { num_leaves: 8, layers: [LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: ddd248964e320312a66775aee8e16c88c927734be59aca09b7af6deb0ad00e8c }, LayerConfig { hash_id: ddd248964e320312a66775aee8e16c88c927734be59aca09b7af6deb0ad00e8c }, LayerConfig { hash_id: aa802dcdabad8ea8a1430919893b96c30e31ff5734b385999108aa202d27dc12 }, LayerConfig { hash_id: ddd248964e320312a66775aee8e16c88c927734be59aca09b7af6deb0ad00e8c }, LayerConfig { hash_id: aa802dcdabad8ea8a1430919893b96c30e31ff5734b385999108aa202d27dc12 }] } }, johnson_slack: 0.0, in_domain_samples: 1, out_domain_samples: 7, deduplicate_in_domain: true }, target: Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 1, vector_size: 2, mask_length: 1, codeword_length: 4, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 1, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 4, layers: [LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }] } }, johnson_slack: 0.0, in_domain_samples: 0, out_domain_samples: 1, deduplicate_in_domain: true }, message_mask_length: 3, out_domain_samples: 4 } +cc a0ce555f621161cc430096035f127c5098eebb8670480a5fa39ec36a9b10d4e2 # shrinks to seed = 0, config = Config { source: Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 1, vector_size: 64, mask_length: 3, codeword_length: 32, interleaving_depth: 4, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 4, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 32, layers: [LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }] } }, johnson_slack: 0.0, in_domain_samples: 1, out_domain_samples: 10, deduplicate_in_domain: false }, target: Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 1, vector_size: 16, mask_length: 1, codeword_length: 64, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 1, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 64, layers: [LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }, LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, LayerConfig { hash_id: 018eaa247cb8957ab1e9fdac885450403c5e7bda1acaaa504e4cc8f2f76bb076 }] } }, johnson_slack: 0.0, in_domain_samples: 3, out_domain_samples: 8, deduplicate_in_domain: true }, message_mask_length: 3, out_domain_samples: 3 } diff --git a/proptest-regressions/protocols/mask_proximity.txt b/proptest-regressions/protocols/mask_proximity.txt new file mode 100644 index 00000000..71771e33 --- /dev/null +++ b/proptest-regressions/protocols/mask_proximity.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 6a96626efcf88c8458e112fceda942463479dfd2637ca89ecaa4c89e49bf8f8e # shrinks to seed = 0, config = Config { c_zk: Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 2, vector_size: 2, mask_length: 3, codeword_length: 4, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 2, leaf_hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, merkle_tree: Config { num_leaves: 4, layers: [LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }] } }, johnson_slack: 0.05590169943749475, in_domain_samples: 0, out_domain_samples: 1, deduplicate_in_domain: false }, num_masks: 1 } diff --git a/proptest-regressions/protocols/sumcheck.txt b/proptest-regressions/protocols/sumcheck.txt new file mode 100644 index 00000000..a25df3ca --- /dev/null +++ b/proptest-regressions/protocols/sumcheck.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 3a705bcc9147c4da8cf518feed3a83ca5031e1ccc09c3b20ba272427ee701cd2 # shrinks to seed = 0, config = Config { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, initial_size: 2, round_pow: Config { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, threshold: 18446744073709551615 }, num_rounds: 1, mask_length: 9, mask_config: Some(Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 1, vector_size: 9, mask_length: 0, codeword_length: 18, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 1, leaf_hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, merkle_tree: Config { num_leaves: 18, layers: [LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, LayerConfig { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }] } }, johnson_slack: 0.035355339059327376, in_domain_samples: 75, out_domain_samples: 1, deduplicate_in_domain: false }) } diff --git a/src/algebra/sumcheck.rs b/src/algebra/sumcheck.rs index aa911b3f..15d9c1df 100644 --- a/src/algebra/sumcheck.rs +++ b/src/algebra/sumcheck.rs @@ -13,7 +13,7 @@ pub fn compute_sumcheck_polynomial(a: &[F], b: &[F]) -> (F, F) { fn recurse(a0: &[F], a1: &[F], b0: &[F], b1: &[F]) -> (F, F) { debug_assert_eq!(a0.len(), b0.len()); debug_assert_eq!(a1.len(), b1.len()); - debug_assert!(a0.len() == a1.len()); + debug_assert_eq!(a0.len(), a1.len()); #[cfg(feature = "parallel")] if a0.len() * 4 > workload_size::() { diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index e142aa55..32763271 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -24,13 +24,20 @@ use crate::{ verify, }; +/// Output from the base case protocol (shared by prover and verifier). +#[must_use] +pub struct Opening { + pub evaluation_points: Vec, + pub linear_form_evaluation: F, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config { pub commit: irs_commit::Config>, pub sumcheck: sumcheck::Config, - /// Whether to mask the vectors, whichs adds HVZK. + /// Whether to mask the vectors, which adds HVZK. pub masked: bool, } @@ -46,7 +53,7 @@ impl Config { witness: &irs_commit::Witness, mut covector: Vec, mut sum: F, - ) -> (Vec, F) + ) -> Opening where H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -63,10 +70,13 @@ impl Config { assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size)); debug_assert_eq!(dot(&vector, &covector), sum); if self.size() == 0 { - return (Vec::new(), F::ZERO); + return Opening { + evaluation_points: Vec::new(), + linear_form_evaluation: F::ZERO, + }; } - // Even more trivial non-zk protocol: send f an r directly. + // Even more trivial non-zk protocol: send f and r directly. if !self.masked { prover_state.prover_messages(&vector); prover_state.prover_messages(&witness.masks); @@ -74,9 +84,12 @@ impl Config { let point = self .sumcheck .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) - .0; + .round_challenges; assert!(!vector[0].is_zero(), "Proof failed"); - return (point, covector[0]); + return Opening { + evaluation_points: point, + linear_form_evaluation: covector[0], + }; } // Create masking vector. @@ -113,7 +126,7 @@ impl Config { &mut masked_sum, &[], ) - .0; + .round_challenges; // If the MLE of `masked_vector` evaluates to zero, the verifier can not proceed. // Basically the sumcheck equation has degenerated to 0 * l(r) = 0, which provides @@ -121,8 +134,10 @@ impl Config { // This event is cryptographically unlikely as `F` is challenge sized. assert!(!masked_vector[0].is_zero(), "Proof failed"); - // Return evaluation point and value of the covector. - (point, covector[0]) + Opening { + evaluation_points: point, + linear_form_evaluation: covector[0], + } } pub fn verify( @@ -130,7 +145,7 @@ impl Config { verifier_state: &mut VerifierState, commitment: &irs_commit::Commitment, mut sum: F, - ) -> VerificationResult<(Vec, F)> + ) -> VerificationResult> where H: DuplexSpongeInterface, F: Codec<[H::U]>, @@ -144,7 +159,10 @@ impl Config { assert_eq!(self.commit.vector_size, self.sumcheck.initial_size); assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size)); if self.size() == 0 { - return Ok((Vec::new(), F::ZERO)); + return Ok(Opening { + evaluation_points: Vec::new(), + linear_form_evaluation: F::ZERO, + }); } // Unmasked protocol @@ -153,7 +171,10 @@ impl Config { let masks = verifier_state .prover_messages_vec(self.commit.mask_length * self.commit.num_messages())?; let evals = self.commit.verify(verifier_state, &[commitment])?; - let point = self.sumcheck.verify(verifier_state, &mut sum)?.0; + let point = self + .sumcheck + .verify(verifier_state, &mut sum)? + .round_challenges; for (&point, value) in zip_strict(&evals.points, evals.values(&[F::ONE])) { // We expected `f(x) + x^l · g(x)` where l = deg(f) + 1, f is the message and g the mask. @@ -165,7 +186,10 @@ impl Config { let mle = multilinear_extend(&vector, &point); verify!(!mle.is_zero()); let linear_mle = sum / mle; - return Ok((point, linear_mle)); + return Ok(Opening { + evaluation_points: point, + linear_form_evaluation: linear_mle, + }); } let mask_commitment = self.commit.receive_commitment(verifier_state)?; @@ -191,7 +215,10 @@ impl Config { // Sumcheck on masked inner product let mut masked_sum = mask_sum + mask_rlc * sum; - let point = self.sumcheck.verify(verifier_state, &mut masked_sum)?.0; + let point = self + .sumcheck + .verify(verifier_state, &mut masked_sum)? + .round_challenges; // Compute implied MLE of the linear form // f*(r) · l(r) = sum => l(r) = sum / f*(r) @@ -199,7 +226,10 @@ impl Config { verify!(!masked_mle.is_zero()); let linear_mle = masked_sum / masked_mle; - Ok((point, linear_mle)) + Ok(Opening { + evaluation_points: point, + linear_form_evaluation: linear_mle, + }) } } @@ -255,14 +285,17 @@ mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); let witness = config.commit.commit(&mut prover_state, &[&vector]); - let (point, value) = config.prove( + let prover_result = config.prove( &mut prover_state, vector.clone(), &witness, covector.clone(), sum, ); - assert_eq!(multilinear_extend(&covector, &point), value); + assert_eq!( + multilinear_extend(&covector, &prover_result.evaluation_points), + prover_result.linear_form_evaluation + ); let proof = prover_state.proof(); // Verifier @@ -271,11 +304,17 @@ mod tests { .commit .receive_commitment(&mut verifier_state) .unwrap(); - let (verifier_point, verifier_value) = config + let verifier_result = config .verify(&mut verifier_state, &commitment, sum) .unwrap(); - assert_eq!(verifier_point, point); - assert_eq!(verifier_value, value); + assert_eq!( + verifier_result.evaluation_points, + prover_result.evaluation_points + ); + assert_eq!( + verifier_result.linear_form_evaluation, + prover_result.linear_form_evaluation + ); verifier_state.check_eof().unwrap(); } diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 31ef08a3..62bbb602 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -15,192 +15,149 @@ use crate::{ algebra::{ dot, embedding::{Embedding, Identity}, - fields::FieldWithSize, - geometric_accumulate, lift, mixed_dot, random_vector, scalar_mul, univariate_evaluate, + eq_weights, geometric_accumulate, lift, mixed_dot, scalar_mul, univariate_evaluate, }, - engines::EngineId, hash::Hash, protocols::{ geometric_challenge::geometric_challenge, - irs_commit::{ - num_ood_samples, Commitment as IrsCommitment, Config as IrsConfig, - Witness as IrsWitness, - }, + irs_commit::{Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness}, }, transcript::{ Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, VerificationResult, VerifierMessage, VerifierState, }, - type_info::Typed, + verify, }; /// Code-switching IOR config with optional ZK. #[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config { - pub embedding: Typed, pub source: IrsConfig, pub target: IrsConfig>, + pub message_mask_length: usize, // l_zk pub out_domain_samples: usize, - pub mask_commit: Option>>, -} - -/// Next stage's query budgets for ZK encoding (Prop 3.19). -#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] -pub struct ZkQueryBudget { - /// Queries the next stage makes to the target oracle g. - pub target: usize, - /// Queries the next stage makes to the mask oracle s. - pub mask: usize, -} - -/// ZK mask oracle: the message and its IRS witness. -#[derive(Clone, Debug)] -pub struct MaskOracle { - pub message: Vec, - pub witness: IrsWitness, } /// Prover output from the code-switch. +#[must_use] #[derive(Clone, Debug)] pub struct Witness { pub message: Vec, pub target_witness: IrsWitness, - pub mask: Option>, } /// Verifier output from the code-switch. -#[derive(Clone, Debug)] -pub struct Commitment { - pub target: IrsCommitment, - pub mask: Option>, +pub type Commitment = IrsCommitment; + +/// Mask input for the code-switch prover. +// TODO : This may be removed after parameter selection PR +pub enum MaskInput<'a, F> { + Disabled, + Enabled(&'a [F]), } impl Config { - /// `zk`: `None` for non-ZK, `Some(budget)` for ZK with next stage's query budgets. + /// Create a code-switch config. + /// + /// The orchestrator is responsible for: + /// - Setting `target_config.mask_length` for ZK mode before passing it in. + /// - Computing `out_domain_samples` from the security budget. + /// - Setting `message_mask_length` = mask oracle message length (0 for non-ZK). pub fn new( - security_target: f64, - unique_decoding: bool, - hash_id: EngineId, source_config: IrsConfig, - target_log_inv_rate: usize, - target_interleaving_depth: usize, - zk: Option, - ) -> Self - where - M: Default, - M::Target: Default, - { - assert!(target_log_inv_rate > 0); - assert!(target_interleaving_depth > 0); - - let source_message_length = source_config.message_length(); - let target_rate = 0.5_f64.powf(target_log_inv_rate as f64); - let target_vector_size = source_message_length * target_interleaving_depth; - - let mut target_config = IrsConfig::>::new( - security_target, - unique_decoding, - hash_id, - 1, - target_vector_size, - target_interleaving_depth, - target_rate, + target_config: IrsConfig>, + out_domain_samples: usize, + message_mask_length: usize, + ) -> Self { + assert_eq!( + source_config.num_vectors, 1, + "code-switch requires a single source vector" ); - - let mask_commit = zk.map(|budget| { - assert!(budget.target > 0, "ZK requires nonzero target query budget"); - assert!(budget.mask > 0, "ZK requires nonzero mask query budget"); - // TODO : find a better way to set mask length for ZK code-switch - // leaving right now for code switching. This will change in parameter - // selection - let target_slack = - target_config.codeword_length - target_config.message_length(); + assert_eq!( + target_config.num_vectors, 1, + "code-switch requires a single target vector" + ); + // Target encodes one polynomial of length ℓ = source.message_length() + // under C' = D^{ι_t}. The IRS splits the input of length ℓ into ι_t + // parallel slices of length ℓ/ι_t, each encoded under D. + assert_eq!( + target_config.vector_size, + source_config.message_length(), + "target vector_size must equal source message_length (target encodes one polynomial of length ℓ)" + ); + assert!( + target_config.interleaving_depth.is_power_of_two(), + "target.interleaving_depth must be a power of 2" + ); + // Theorem 9.6: ℓ_zk ≥ r (mask oracle must cover source randomness). + if message_mask_length > 0 { assert!( - budget.target <= target_slack, - "budget.target ({}) exceeds target codeword slack ({codeword_len} - {msg_len} = {target_slack})", - budget.target, - codeword_len = target_config.codeword_length, - msg_len = target_config.message_length(), + message_mask_length >= source_config.mask_length, + "message_mask_length ({message_mask_length}) must be >= source randomness length ({})", + source_config.mask_length, ); - target_config.mask_length = budget.target; - target_config.recompute_security_parameters(security_target, unique_decoding); - - let source_randomness_len = source_config.mask_length * source_config.num_messages(); assert!( - source_randomness_len > 0, - "ZK code-switch requires source_config.mask_length > 0" + message_mask_length - source_config.mask_length >= out_domain_samples, + "the sampled randomness (s) length must be covering all the out of domain sample requests" ); - // TODO : move the mask config out for shared mask tree per iteration - // inputs to the new function will also change - let mut mask_config = IrsConfig::>::new( - security_target, - unique_decoding, - hash_id, - 1, - source_randomness_len, - 1, - source_config.rate(), - ); - let mask_slack = - mask_config.codeword_length - mask_config.message_length(); assert!( - budget.mask <= mask_slack, - "budget.mask ({}) exceeds mask codeword slack ({codeword_len} - {msg_len} = {mask_slack})", - budget.mask, - codeword_len = mask_config.codeword_length, - msg_len = mask_config.message_length(), + target_config.mask_length + >= target_config.in_domain_samples + target_config.out_domain_samples, + "target encoder violates: t' > r', number of queries should be covered by random mask" ); - mask_config.mask_length = budget.mask; - mask_config.recompute_security_parameters(security_target, unique_decoding); - mask_config - }); - + } assert!( - source_config.mask_length == 0 || mask_commit.is_some(), - "source with ZK randomness requires ZK code-switch (mask_commit)" - ); - - let (list_size, degree) = mask_commit.as_ref().map_or_else( - || (target_config.list_size(), source_message_length), - |mask_cfg| { - ( - target_config.list_size() * mask_cfg.list_size(), - source_message_length + mask_cfg.message_length(), - ) - }, + source_config.mask_length == 0 || message_mask_length > 0, + "source with mask_length > 0 (IRS randomness) requires ZK mode (message_mask_length > 0)" ); - let out_domain_samples = num_ood_samples( - unique_decoding, - security_target, - M::Target::field_size_bits(), - list_size, - degree, + assert!( + source_config.interleaving_depth.is_power_of_two(), + "source.interleaving_depth must be a power of 2" ); Self { - embedding: Typed::::default(), source: source_config, target: target_config, + message_mask_length, out_domain_samples, - mask_commit, } } /// Length of the covector for this code-switch. pub fn covector_length(&self) -> usize { - self.source.masked_message_length() + self.source.message_length() + self.message_mask_length } - /// Prove the code-switch + /// Prove the code-switch. + /// + /// # Soundness-critical inputs + /// + /// `folding_randomness` is the **sumcheck folding randomness `γ`** that + /// was sampled from the verifier in the preceding sumcheck protocol + /// (Construction 6.3, p.37-38). It must be the same `γ` the verifier + /// derived from the transcript — it is NOT caller-supplied randomness. + /// + /// Used by the verifier to collapse ι_s parallel codeword columns into a + /// single value of `Fold(f, γ)` via `eq_weights(γ)`. Passing different + /// randomness here breaks IOR completeness; passing locally-sampled + /// randomness breaks Fiat-Shamir soundness in the composed protocol. + /// + /// `message` is `Fold(f, γ)`, the post-sumcheck polynomial of length + /// `source.message_length()`. + /// + /// `mask_input` is `(r || s)` from the orchestrator's shared mask tree + /// (see Construction 9.7 Step 1, p.55). Must be `None` when + /// `message_mask_length == 0`. #[cfg_attr(feature = "tracing", instrument(skip_all))] pub fn prove( &self, prover_state: &mut ProverState, message: Vec, - masks: &[M::Source], witness: &IrsWitness, covector: &mut [M::Target], + folding_randomness: &[M::Target], + mask_input: &MaskInput<'_, M::Target>, ) -> Witness where H: DuplexSpongeInterface, @@ -213,52 +170,43 @@ impl Config { assert_eq!(message.len(), self.source.message_length()); assert_eq!(covector.len(), self.covector_length()); assert_eq!( - self.source.num_messages(), - 1, - "code-switch only supports num_messages() == 1 (single polynomial)" - ); - assert_eq!( - masks.len(), - self.source.mask_length * self.source.num_messages() - ); - assert!( - self.mask_commit.is_some() == (self.target.mask_length > 0), - "mask config and target mask_length must agree" + 1 << folding_randomness.len(), + self.source.interleaving_depth, + "folding_randomness must have length log2(source.interleaving_depth) ({} != log2({}))", + folding_randomness.len(), + self.source.interleaving_depth, ); + let mask_msg: Option<&[M::Target]> = match &mask_input { + MaskInput::Disabled => { + assert_eq!( + self.message_mask_length, 0, + "MaskInput::Disabled requires message_mask_length == 0" + ); + None + } + MaskInput::Enabled(mask) => { + assert_eq!( + mask.len(), + self.message_mask_length, + "mask_msg length must equal message_mask_length" + ); + Some(mask) + } + }; - // Step 1a: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 + // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 let target_witness = self.target.commit(prover_state, &[&message]); - // Step 1b: s := Enc_{C_zk}((r || padding), r'') — Construction 9.7 Step 1, p.55 - #[allow(clippy::option_if_let_else)] - let mask = if let Some(mask_config) = &self.mask_commit { - let mask_msg_len = mask_config.message_length(); - let r_embedded = lift(self.source.embedding(), masks); - let embedded_randomness_len = r_embedded.len(); - let mut mask_msg = Vec::with_capacity(mask_msg_len); - mask_msg.extend_from_slice(&r_embedded); - let random_padding: Vec = - random_vector(prover_state.rng(), mask_msg_len - embedded_randomness_len); - mask_msg.extend_from_slice(&random_padding); - let witness = mask_config.commit(prover_state, &[&mask_msg]); - Some(MaskOracle { - message: mask_msg, - witness, - }) - } else { - None - }; - // Step 2-3: OOD challenge + answers — Construction 9.7 Steps 2-3, p.55 - // TODO : check the private zero evader for code switch protocol. + // y := ze_ood(ρ) · [f; r; s] = f(α) + α^ℓ · (r,s)(α) let ood_points: Vec = prover_state.verifier_message_vec(self.out_domain_samples); let msg_len = message.len(); for &point in &ood_points { let f_eval = univariate_evaluate(&message, point); - if let Some(ref mask_oracle) = mask { - let mask_msg_eval = univariate_evaluate(&mask_oracle.message, point); + if let Some(mask) = mask_msg { + let mask_eval = univariate_evaluate(mask, point); let shift = point.pow([msg_len as u64]); - prover_state.prover_message(&(f_eval + shift * mask_msg_eval)); + prover_state.prover_message(&(f_eval + shift * mask_eval)); } else { prover_state.prover_message(&f_eval); } @@ -269,7 +217,7 @@ impl Config { // Step 4.1: batching — Construction 9.7 Step 4, p.55 let num_ood = self.out_domain_samples; - let num_in_domain = source_evaluations.matrix.len(); + let num_in_domain = source_evaluations.points.len(); let batching_coeffs = geometric_challenge::<_, M::Target>(prover_state, 1 + num_ood + num_in_domain); let (&original_sl_coeff, constraint_rlc_coeffs) = batching_coeffs.split_first().unwrap(); @@ -277,28 +225,64 @@ impl Config { // Covector update — sl' from Completeness proof (p.55-56) let eval_points = lift(self.source.embedding(), &source_evaluations.points); - let all_points: Vec<_> = ood_points.iter().chain(&eval_points).copied().collect(); - let pows: Vec<_> = ood_rlc_coeffs - .iter() - .chain(in_domain_rlc_coeffs) - .copied() - .collect(); scalar_mul(covector, original_sl_coeff); - geometric_accumulate(covector, pows, &all_points); + if self.message_mask_length == 0 { + // Non-ZK: single accumulate over all points + let all_points: Vec<_> = ood_points.iter().chain(&eval_points).copied().collect(); + let pows: Vec<_> = ood_rlc_coeffs + .iter() + .chain(in_domain_rlc_coeffs) + .copied() + .collect(); + geometric_accumulate(covector, pows, &all_points); + } else { + // ZK: OOD contributes to full [f; r; s], in-domain only to [f; r] + geometric_accumulate(covector, ood_rlc_coeffs.to_vec(), &ood_points); + geometric_accumulate( + &mut covector[..self.source.masked_message_length()], + in_domain_rlc_coeffs.to_vec(), + &eval_points, + ); + } Witness { message, target_witness, - mask, } } - /// Verify the code-switch + /// Verify the code-switch. + /// + /// `folding_randomness` is the **sumcheck folding randomness `γ`** the + /// verifier derived from the transcript during the preceding sumcheck. + /// It must match what the prover received from the same transcript — + /// not caller-supplied randomness. See `prove` doc for details. + /// + /// Returns the target commitment. In ZK mode, the caller **must** + /// additionally run `mask_proximity::verify` on the mask commitment + /// to ensure the mask oracle `(r, s)` is close to a `C_zk` codeword. + /// Without this check, soundness is not guaranteed. + /// + /// # Soundness composition note + /// + /// This verifier checks the OOD/in-domain consistency of the target + /// codeword `g` against transcript-supplied mask values `s(α_i)`. It + /// does **not** check that `s` is close to a `C_zk` codeword — that + /// is the job of mask-proximity (Construction 7.2). Without a + /// downstream mask-proximity invocation against the same `s`, a + /// prover can submit non-codeword mask values that satisfy the OOD + /// equation, breaking the soundness reduction in Theorem 9.10. + /// + /// In the orchestrated WHIR protocol, the orchestrator owns the + /// per-round mask tree containing `s` and is responsible for + /// running `mask_proximity::verify` on that same tree before + /// accepting the round. #[cfg_attr(feature = "tracing", instrument(skip_all))] pub fn verify( &self, verifier_state: &mut VerifierState, sum: &mut M::Target, + folding_randomness: &[M::Target], commitment: &IrsCommitment, ) -> VerificationResult> where @@ -308,25 +292,17 @@ impl Config { u8: Decoding<[H::U]>, Hash: ProverMessage<[H::U]>, { - assert_eq!( - self.source.num_messages(), - 1, - "code-switch only supports num_messages() == 1 (single polynomial)" - ); - assert!( - self.mask_commit.is_some() == (self.target.mask_length > 0), - "mask config and target mask_length must agree" - ); + verify!(1 << folding_randomness.len() == self.source.interleaving_depth); - // Step 1: commitments — Construction 9.7 Step 1, p.55 + let collapse_weights = eq_weights(folding_randomness); + + // Step 1: target commitment — Construction 9.7 Step 1, p.55 + // Mask oracle is committed in the shared mask tree by the orchestrator. let target_commitment = self.target.receive_commitment(verifier_state)?; - let mask_commitment = self - .mask_commit - .as_ref() - .map(|mask_cfg| mask_cfg.receive_commitment(verifier_state)) - .transpose()?; // Step 2-3: OOD — Construction 9.7 Steps 2-3, p.55 + // In ZK mode, ood_answers = f(α) + α^ℓ · (r,s)(α) where (r,s) is + // the mask oracle message committed in the shared tree. let _ood_points: Vec = verifier_state.verifier_message_vec(self.out_domain_samples); let ood_answers: Vec = @@ -334,26 +310,24 @@ impl Config { // Step 4: source opening — Construction 9.7 Step 4, p.55 let source_evaluations = self.source.verify(verifier_state, &[commitment])?; + let collapsed_values: Vec = source_evaluations + .matrix + .chunks_exact(self.source.interleaving_depth) + .map(|row| mixed_dot(self.source.embedding(), &collapse_weights, row)) + .collect(); // Step 4.1: batching + μ' — Construction 9.7 Decision phase, p.55 let num_ood = self.out_domain_samples; - let num_in_domain = source_evaluations.matrix.len(); + let num_in_domain = source_evaluations.points.len(); let coeffs = geometric_challenge(verifier_state, 1 + num_ood + num_in_domain); let (&original_sl_coeff, all_rlc_coeffs) = coeffs.split_first().unwrap(); let (ood_rlc_coeffs, in_domain_rlc_coeffs) = all_rlc_coeffs.split_at(num_ood); *sum = original_sl_coeff * *sum + dot(ood_rlc_coeffs, &ood_answers) - + mixed_dot( - self.source.embedding(), - in_domain_rlc_coeffs, - &source_evaluations.matrix, - ); + + dot(in_domain_rlc_coeffs, &collapsed_values); - Ok(Commitment { - target: target_commitment, - mask: mask_commitment, - }) + Ok(target_commitment) } } @@ -361,13 +335,12 @@ impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, - "Config(source={}, target={}, ood_samples={}", - self.source, self.target, self.out_domain_samples, - )?; - if let Some(mask) = &self.mask_commit { - write!(f, ", mask={mask}")?; - } - write!(f, ")") + "CodeSwitch(source={}, target={}, ood={}, zk={})", + self.source, + self.target, + self.out_domain_samples, + self.message_mask_length != 0, + ) } } @@ -376,7 +349,7 @@ mod tests { use ark_std::rand::{ distributions::Standard, prelude::Distribution, rngs::StdRng, Rng, SeedableRng, }; - use proptest::{bool, prelude::Strategy, proptest, sample::select, strategy::Just}; + use proptest::{bool, prelude::Strategy, prop_assume, proptest, sample::select}; use super::*; use crate::{ @@ -387,59 +360,145 @@ mod tests { impl Config { pub fn arbitrary(embedding: M) -> impl Strategy where - M: 'static, + M: Default + 'static, { - let valid_sizes = (1..=256) + // Sizes ≥ 4 to allow ι ∈ {1, 2, 4} with non-trivial message_length. + let valid_sizes = (4..=256) .filter(|&n| ntt::next_order::(n) == Some(n)) + .filter(|&n| n.is_power_of_two()) .collect::>(); - let emb1 = embedding.clone(); - let emb2 = embedding; - - (select(valid_sizes), 0_usize..=3, bool::ANY) - .prop_flat_map(move |(size, mask_length, masked)| { - let source_mask = if masked { mask_length.max(1) } else { 0 }; - let source = IrsConfig::arbitrary(emb1.clone(), 1, size, source_mask, 1); - (source, Just(masked)) - }) - .prop_flat_map(|(source, masked)| { - let msg_len = source.message_length(); - let rnd_len = source.mask_length * source.num_messages(); - let target_mask = usize::from(masked); - - let target = IrsConfig::>::arbitrary( - Identity::new(), - 1, - msg_len, - target_mask, - 1, - ); - - let mask = if masked && rnd_len > 0 { - IrsConfig::>::arbitrary( - Identity::new(), - 1, - rnd_len, - 0, - 1, - ) - .prop_map(Some) - .boxed() + let scalars = ( + select(valid_sizes), + 0_usize..=3, // src_mask_len (source IRS randomness, post-fold) + bool::ANY, // zk + 0_usize..=5, // ood (= code-switch t_ood) + 0_usize..=5, // fresh_s_len (≥ ood for assumption (c)) + select(vec![1_usize, 2, 4]), // ι_s (source interleaving) + 0_usize..=10, // target.in_domain_samples (t'_in) + 0_usize..=10, // target.out_domain_samples (t'_out) + ); + + scalars.prop_flat_map( + move |(size, src_mask_len, zk, ood, fresh_s_len, iota_s, t_in, t_out)| { + // Bound 3 assumption (c): ℓ_zk - r ≥ t_ood ⇒ fresh_s_len ≥ ood. + let fresh_s_len = if zk { + fresh_s_len.max(ood) } else { - Just(None).boxed() + fresh_s_len }; + // Bound 4 assumption (a): target.mask_length ≥ t' = t_in + t_out. + let target_mask = if zk { t_in + t_out } else { 0 }; + // ZK with source.mask_length = 0 is valid: the assert + // `source.mask_length == 0 || message_mask_length > 0` + // is trivially satisfied. Allows testing the corner + // where the mask oracle has only fresh randomness. + let source_mask = if zk { src_mask_len } else { 0 }; + + IrsConfig::arbitrary(embedding.clone(), 1, size, source_mask, iota_s) + .prop_flat_map(move |source| { + // ι_t must divide msg_len and be a power of 2. + let msg_len = source.message_length(); + let iota_t_choices: Vec = [1, 2, 4] + .into_iter() + .filter(|&i| msg_len.is_multiple_of(i)) + .collect(); + + select(iota_t_choices).prop_flat_map(move |iota_t| { + // target.vector_size = ℓ; C' = D^{ι_t} where D's + // message length = ℓ / ι_t. + let target = IrsConfig::>::arbitrary( + Identity::new(), + 1, + msg_len, + target_mask, + iota_t, + ); + let source = source.clone(); + target.prop_map(move |mut target| { + // IrsConfig::arbitrary samples query counts in + // [0,10] independently of mask_length; pin them + // to the values target_mask was sized for so + // assumption (a) holds. + if zk { + target.in_domain_samples = t_in; + target.out_domain_samples = t_out; + } + // r = post-fold randomness length (ι_s parallel + // masks fold to a single length-mask_length chunk). + let r = source.mask_length; + let message_mask_length = if zk { r + fresh_s_len } else { 0 }; + Self::new(source.clone(), target, ood, message_mask_length) + }) + }) + }) + }, + ) + } + } - (Just(source), target, mask, 0_usize..=5) - }) - .prop_map( - move |(source, target, mask_commit, out_domain_samples)| Self { - embedding: Typed::new(emb2.clone()), - source, - target, - out_domain_samples, - mask_commit, - }, - ) + /// Fold ι parallel chunks of length `chunk_len` into a single chunk via + /// eq_weights(γ). Layout: values = [chunk_0; chunk_1; ...; chunk_{ι-1}], + /// each of length `chunk_len`. Returns Σ_l eq_weights(γ)[l] · chunk_l. + fn fold_chunks(values: &[F], chunk_len: usize, folding_randomness: &[F]) -> Vec { + let iota = 1 << folding_randomness.len(); + assert_eq!(values.len(), chunk_len * iota); + if iota == 1 { + return values.to_vec(); + } + let weights = eq_weights(folding_randomness); + (0..chunk_len) + .map(|j| { + (0..iota) + .map(|l| weights[l] * values[l * chunk_len + j]) + .sum() + }) + .collect() + } + + /// Sample folding randomness of length log2(source.interleaving_depth). + fn sample_folding_randomness( + config: &Config>, + rng: &mut impl RngCore, + ) -> Vec + where + Standard: Distribution, + { + let log_iota = config.source.interleaving_depth.trailing_zeros() as usize; + random_vector(rng, log_iota) + } + + /// Simulate what the orchestrator does: build (r || fresh_s) where r is + /// the *folded* source IRS randomness. Returns empty vec in non-ZK mode. + fn build_mask_msg( + config: &Config>, + source_witness: &IrsWitness, + folding_randomness: &[F], + rng: &mut impl RngCore, + ) -> Vec + where + Standard: Distribution, + { + if config.message_mask_length == 0 { + return Vec::new(); + } + // Lift ι parallel masks (total length source.mask_length × ι) and fold + // chunks of length source.mask_length down to a single chunk. + let raw = lift(config.source.embedding(), &source_witness.masks); + let mut mask = fold_chunks(&raw, config.source.mask_length, folding_randomness); + // Append fresh padding s of length message_mask_length - source.mask_length. + mask.extend(random_vector::( + rng, + config.message_mask_length - mask.len(), + )); + mask + } + + fn mask_input(mask_msg: &[F]) -> MaskInput<'_, F> { + if mask_msg.is_empty() { + MaskInput::Disabled + } else { + MaskInput::Enabled(mask_msg) } } @@ -448,28 +507,102 @@ mod tests { Standard: Distribution, Hash: ProverMessage<[u8]>, { - crate::tests::init(); + let mut rng = StdRng::seed_from_u64(seed); + // Commit the full pre-fold vector of length source.vector_size + // (= ι · message_length), which IRS encodes as ι parallel codewords. + let f_full: Vec = random_vector(&mut rng, config.source.vector_size); + let initial_sum: F = rng.gen(); + + let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); + covector.resize(config.covector_length(), F::ZERO); let instance = U64(seed); let ds = DomainSeparator::protocol(config) .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); + let mut prover_state = ProverState::new_std(&ds); + let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + + // Sample γ for sumcheck folding (length log2(ι)). + let folding_randomness = sample_folding_randomness(config, &mut rng); + // Post-fold message Fold(f_full, γ) of length message_length. + let folded_message = + fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); + + let witness = config.prove( + &mut prover_state, + folded_message.clone(), + &source_witness, + &mut covector, + &folding_randomness, + &mask_input(&mask_msg), + ); + let proof = prover_state.proof(); + + let mut verifier_state = VerifierState::new_std(&ds, &proof); + let source_commitment = config + .source + .receive_commitment(&mut verifier_state) + .unwrap(); + let mut verifier_sum = initial_sum; + let _ = config + .verify( + &mut verifier_state, + &mut verifier_sum, + &folding_randomness, + &source_commitment, + ) + .unwrap(); + verifier_state.check_eof().unwrap(); + assert_eq!(witness.message, folded_message); + } + + fn test_ior_identity_config>(seed: u64, config: &Config>) + where + Standard: Distribution, + Hash: ProverMessage<[u8]>, + { let mut rng = StdRng::seed_from_u64(seed); - let message: Vec = random_vector(&mut rng, config.source.message_length()); - let target_mu: F = rng.gen(); + let f_full: Vec = random_vector(&mut rng, config.source.vector_size); - let msg_len = config.source.message_length(); - let mut covector: Vec = random_vector(&mut rng, msg_len); + let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let instance = U64(seed); + let ds = DomainSeparator::protocol(config) + .session(&format!("Test at {}:{}", file!(), line!())) + .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&message]); - let witness = config.prove( + let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + + let folding_randomness = sample_folding_randomness(config, &mut rng); + let folded_message = + fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); + + // h is the post-fold polynomial whose inner product with covector + // should equal the verifier sum: + // - non-ZK: h = folded_message (length message_length) + // - ZK: h = [folded_message; mask_msg] (length message_length + l_zk) + let h: Vec = if mask_msg.is_empty() { + folded_message.clone() + } else { + folded_message + .iter() + .chain(mask_msg.iter()) + .copied() + .collect() + }; + let initial_mu = dot(&h, &covector); + + let _witness = config.prove( &mut prover_state, - message.clone(), - &source_witness.masks, + folded_message, &source_witness, &mut covector, + &folding_randomness, + &mask_input(&mask_msg), ); let proof = prover_state.proof(); @@ -478,17 +611,76 @@ mod tests { .source .receive_commitment(&mut verifier_state) .unwrap(); - let mut verifier_sum = target_mu; - let commitments = config - .verify(&mut verifier_state, &mut verifier_sum, &source_commitment) + let mut verifier_sum = initial_mu; + let _ = config + .verify( + &mut verifier_state, + &mut verifier_sum, + &folding_randomness, + &source_commitment, + ) .unwrap(); + verifier_state.check_eof().unwrap(); + + assert_eq!(dot(&h, &covector), verifier_sum); + } + + fn test_tampered_ood_config>(seed: u64, config: &Config>) + where + Standard: Distribution, + Hash: ProverMessage<[u8]>, + { + let instance = U64(seed); + let ds = DomainSeparator::protocol(config) + .session(&format!("Test at {}:{}", file!(), line!())) + .instance(&instance); + let mut rng = StdRng::seed_from_u64(seed); + let f_full: Vec = random_vector(&mut rng, config.source.vector_size); + + let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); + covector.resize(config.covector_length(), F::ZERO); - // Transcript fully consumed — prover and verifier are in sync + // Commit honest f_full, fold to get the honest post-fold message. + let mut prover_state = ProverState::new_std(&ds); + let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let folding_randomness = sample_folding_randomness(config, &mut rng); + let folded_message = + fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + + // For non-ZK and source.mask_length == 0, h = folded_message and identity holds. + let initial_mu = dot(&folded_message, &covector); + + // Tamper the post-fold message before proving. + let mut tampered = folded_message.clone(); + tampered[0] += F::ONE; + let _witness = config.prove( + &mut prover_state, + tampered, + &source_witness, + &mut covector, + &folding_randomness, + &MaskInput::Disabled, + ); + let proof = prover_state.proof(); + + let mut verifier_state = VerifierState::new_std(&ds, &proof); + let source_commitment = config + .source + .receive_commitment(&mut verifier_state) + .unwrap(); + let mut verifier_sum = initial_mu; + let _ = config + .verify( + &mut verifier_state, + &mut verifier_sum, + &folding_randomness, + &source_commitment, + ) + .unwrap(); verifier_state.check_eof().unwrap(); - // Mask commitment present if ZK mode - assert_eq!(commitments.mask.is_some(), config.mask_commit.is_some()); - assert_eq!(witness.message, message); - assert_eq!(witness.mask.is_some(), config.mask_commit.is_some()); + + // Sum diverges — downstream sumcheck would reject + assert_ne!(dot(&folded_message, &covector), verifier_sum); } fn test + 'static>() @@ -496,6 +688,7 @@ mod tests { Standard: Distribution, Hash: ProverMessage<[u8]>, { + crate::tests::init(); let configs = Config::arbitrary(Identity::::new()); proptest!(|(seed: u64, config in configs)| { test_config(seed, &config); @@ -536,4 +729,30 @@ mod tests { fn test_field256() { test::(); } + + #[test] + fn test_ior_identity() { + crate::tests::init(); + let configs = Config::arbitrary(Identity::::new()); + proptest!(|(seed: u64, config in configs)| { + prop_assume!(config.source.in_domain_samples > 0); + test_ior_identity_config(seed, &config); + }); + } + + #[test] + fn test_tampered_ood() { + crate::tests::init(); + let configs = Config::arbitrary(Identity::::new()).prop_filter( + "non-ZK with ood > 0", + |config| { + config.message_mask_length == 0 + && config.source.mask_length == 0 + && config.out_domain_samples > 0 + }, + ); + proptest!(|(seed: u64, config in configs)| { + test_tampered_ood_config(seed, &config); + }); + } } diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 04c12717..b5716f3a 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -237,27 +237,6 @@ impl Config { self.out_domain_samples == 0 && self.johnson_slack == 0.0 } - /// Recompute `johnson_slack`, `in_domain_samples`, and `out_domain_samples` - /// for the current rate (which accounts for `mask_length`). - pub fn recompute_security_parameters(&mut self, security_target: f64, unique_decoding: bool) { - let rate = self.rate(); - let johnson_slack = if unique_decoding { - 0.0 - } else { - rate.sqrt() / 20. - }; - self.johnson_slack = OrderedFloat(johnson_slack); - self.in_domain_samples = num_in_domain_queries(unique_decoding, security_target, rate); - let list_size = self.list_size(); - self.out_domain_samples = num_ood_samples( - unique_decoding, - security_target, - M::Target::field_size_bits(), - list_size, - self.vector_size, - ); - } - /// Compute a list size bound. pub fn list_size(&self) -> f64 { if self.unique_decoding() { @@ -346,7 +325,8 @@ impl Config { // Handle out-of-domain points and values // TODO : Remove this logic after main whir protocol is updated - // as this is not required in the new construction. + // as this is not required in the new construction. This will be + // removed in next PR (Parameter Selection) let oods_points: Vec = prover_state.verifier_message_vec(self.out_domain_samples); let mut oods_matrix = Vec::with_capacity(self.out_domain_samples * self.num_vectors); diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs new file mode 100644 index 00000000..1d632ca0 --- /dev/null +++ b/src/protocols/mask_proximity.rs @@ -0,0 +1,525 @@ +//! Mask proximity verification via γ-combination. +//! +//! Implements Construction 7.2 (p.43-44) specialized for zero-constraint mask +//! oracles. Given a shared Merkle tree containing 2n vectors — n original masks +//! and n fresh mask-of-masks — this protocol proves that each original mask is +//! close to a C_zk codeword without revealing the mask polynomials. +//! +//! The tree layout is: +//! columns 0..n-1: original masks ξ_1, ..., ξ_n +//! columns n..2n-1: mask-of-masks s_1, ..., s_n +//! +//! Protocol: +//! 1. Verifier sends γ (combination randomness) +//! 2. Prover sends combined polynomials ξ*_i = s_i + γ·ξ_i and +//! combined IRS randomness r*_i = r'_i + γ·r_i for each mask pair +//! 3. Shared tree is opened at random positions +//! 4. Verifier checks: Enc(ξ*_i, r*_i)(y_j) = s_i(y_j) + γ·ξ_i(y_j) +//! at each opened position, using linearity of the RS encoding +//! +//! ZK safety (follows the pattern of Construction 7.2, §7.1): +//! - Only the combined ξ*_i = s_i + γ·ξ_i is revealed in full. Since s_i +//! is uniformly random, ξ*_i is uniform regardless of ξ_i. +//! - The tree is opened at ≤ t_zk positions, revealing ξ_i(y_j) and s_i(y_j) +//! at each opened position. This is within the ZK query budget per mask. +//! +//! Note: Lemma 7.3 (p.44) derives the ZK bound ζ_C + n·ζ_{C_zk} assuming +//! each mask is a separate oracle. This implementation uses a shared tree +//! for all 2n vectors, so a single opening reveals 2n values. The per-mask +//! query count is unchanged (each ξ_i is opened at t_zk positions), but the +//! shared tree means all masks are opened at the SAME positions. +//! +//! Assumed: the shared-tree variant preserves the ZK bound. The argument +//! is that the simulator can simulate all 2n values at each position +//! independently (each pair (ξ_i, s_i) is simulatable from C_zk's ZK +//! property, and the pairs are independent across i). Formal derivation +//! for the shared-tree case is pending. +//! +//! Soundness: if ξ_i is far from C_zk, the spot-check fails with high +//! probability over γ (Lemma 7.4, p.45). + +use std::fmt; + +use ark_ff::Field; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; +use serde::{Deserialize, Serialize}; + +use crate::{ + algebra::{embedding::Identity, random_vector, scalar_mul_add_new, univariate_evaluate}, + hash::Hash, + protocols::irs_commit::{ + Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness, + }, + transcript::{ + Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, VerificationResult, + VerifierMessage, VerifierState, + }, + utils::zip_strict, + verify, +}; + +/// Mask proximity configuration. +/// +/// Wraps an IRS config for the shared mask tree and the number of mask pairs. +#[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] +#[serde(bound = "")] +pub struct Config { + pub c_zk_commit: IrsConfig>, + pub num_masks: usize, +} + +/// Prover output from the commit phase. +#[must_use] +pub struct Witness { + pub(crate) mask_witness: IrsWitness, + pub(crate) fresh_msgs: Vec>, +} + +/// Verifier output from the commit phase. +pub type Commitment = IrsCommitment; + +impl Config { + pub fn new(c_zk_commit: IrsConfig>, num_masks: usize) -> Self { + assert_eq!( + c_zk_commit.num_vectors, + 2 * num_masks, + "c_zk.num_vectors must be 2 * num_masks" + ); + assert_eq!( + c_zk_commit.interleaving_depth, 1, + "mask proximity requires interleaving_depth = 1" + ); + // OOD evaluations are sent in the clear during IRS commit/receive, + // which would leak raw mask values before the γ-combination and + // break the ZK contract. The OOD path in irs_commit is slated for + // removal in the new construction; until then, enforce zero here. + assert_eq!( + c_zk_commit.out_domain_samples, 0, + "mask proximity requires out_domain_samples = 0 (OOD openings would leak raw mask evaluations)" + ); + Self { + c_zk_commit, + num_masks, + } + } + + /// Commit all masks and their mask-of-masks in a single shared tree. + /// + /// Samples n fresh mask-of-mask polynomials, combines them with the + /// provided original masks into a 2n-vector tree, and commits via IRS. + pub fn commit( + &self, + prover_state: &mut ProverState, + original_msgs: &[Vec], + ) -> Witness + where + F: Codec<[H::U]>, + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + Standard: Distribution, + Hash: ProverMessage<[H::U]>, + { + assert_eq!(original_msgs.len(), self.num_masks); + for msg in original_msgs { + assert_eq!(msg.len(), self.c_zk_commit.vector_size); + } + + // Sample fresh mask-of-masks + let fresh_msgs: Vec> = (0..self.num_masks) + .map(|_| random_vector(prover_state.rng(), self.c_zk_commit.vector_size)) + .collect(); + + // Tree layout: [originals..., freshes...] + let all_vectors: Vec<&[F]> = original_msgs + .iter() + .chain(fresh_msgs.iter()) + .map(|v| v.as_slice()) + .collect(); + + let mask_witness = self.c_zk_commit.commit(prover_state, &all_vectors); + + Witness { + mask_witness, + fresh_msgs, + } + } + + /// Receive a mask proximity commitment + pub fn receive_commitment( + &self, + verifier_state: &mut VerifierState, + ) -> VerificationResult> + where + F: Codec<[H::U]>, + H: DuplexSpongeInterface, + Hash: ProverMessage<[H::U]>, + { + self.c_zk_commit.receive_commitment(verifier_state) + } + + /// Prove that each original mask is close to a C_zk codeword. + pub fn prove( + &self, + prover_state: &mut ProverState, + witness: &Witness, + original_msgs: &[Vec], + ) where + F: Codec<[H::U]>, + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + Standard: Distribution, + u8: Decoding<[H::U]>, + Hash: ProverMessage<[H::U]>, + { + assert_eq!(original_msgs.len(), self.num_masks); + assert_eq!(witness.fresh_msgs.len(), self.num_masks); + + // Step 1: receive combination randomness γ + let gamma: F = prover_state.verifier_message(); + + // Step 2: compute and send combined polynomials + IRS randomness + let irs_masks_per_vector = + self.c_zk_commit.mask_length * self.c_zk_commit.interleaving_depth; + assert_eq!( + witness.mask_witness.masks.len(), + 2 * self.num_masks * irs_masks_per_vector + ); + for (i, (orig_msg, fresh_msg)) in original_msgs + .iter() + .zip(witness.fresh_msgs.iter()) + .enumerate() + { + // ξ*_i = s_i + γ · ξ_i + let combined_msg = scalar_mul_add_new(fresh_msg, gamma, orig_msg); + prover_state.prover_messages(&combined_msg); + + // r*_i = r'_i + γ · r_i + if irs_masks_per_vector > 0 { + let orig_r = &witness.mask_witness.masks + [i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; + let fresh_r = &witness.mask_witness.masks[(self.num_masks + i) + * irs_masks_per_vector + ..(self.num_masks + i + 1) * irs_masks_per_vector]; + let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); + prover_state.prover_messages(&combined_r); + } + } + + // Step 3: open the shared tree at random in-domain positions + self.c_zk_commit + .open(prover_state, &[&witness.mask_witness]); + } + + /// Verify that each original mask is close to a C_zk codeword. + pub fn verify( + &self, + verifier_state: &mut VerifierState, + commitment: &Commitment, + ) -> VerificationResult<()> + where + F: Codec<[H::U]>, + H: DuplexSpongeInterface, + u8: Decoding<[H::U]>, + Hash: ProverMessage<[H::U]>, + { + // Step 1: send combination randomness γ + let gamma: F = verifier_state.verifier_message(); + + // Step 2: read combined polynomials + IRS randomness + let msg_len = self.c_zk_commit.message_length(); + let irs_masks_per_vector = + self.c_zk_commit.mask_length * self.c_zk_commit.interleaving_depth; + let has_irs_masks = irs_masks_per_vector > 0; + let mut combined_msgs = Vec::with_capacity(self.num_masks); + let mut combined_rs: Option>> = + has_irs_masks.then(|| Vec::with_capacity(self.num_masks)); + for _ in 0..self.num_masks { + combined_msgs.push(verifier_state.prover_messages_vec(msg_len)?); + if let Some(ref mut rs) = combined_rs { + rs.push(verifier_state.prover_messages_vec(irs_masks_per_vector)?); + } + } + + // Step 3: verify tree openings and get codeword values at opened positions + let evaluations = self.c_zk_commit.verify(verifier_state, &[commitment])?; + + // Step 4: spot-check γ-combination at each opened position + let num_cols = self.c_zk_commit.num_cols(); + for (row, &point) in zip_strict( + evaluations.matrix.chunks_exact(num_cols), + &evaluations.points, + ) { + let shift = combined_rs.as_ref().map(|_| point.pow([msg_len as u64])); + // Direct indexing: c_zk_commit.interleaving_depth == 1 by Config::new. + for i in 0..self.num_masks { + let original_val = row[i]; + let fresh_val = row[self.num_masks + i]; + + // Enc(ξ*_i, r*_i)(point) = ξ*_i(point) + point^msg_len · r*_i(point) + let mut expected = univariate_evaluate(&combined_msgs[i], point); + if let Some((rs, shift)) = combined_rs.as_ref().zip(shift) { + expected += shift * univariate_evaluate(&rs[i], point); + } + + let actual = fresh_val + gamma * original_val; + verify!(expected == actual); + } + } + + Ok(()) + } +} + +impl fmt::Display for Config { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "masks {} tree {}", self.num_masks, self.c_zk_commit) + } +} + +#[cfg(test)] +mod tests { + use ark_std::rand::{ + distributions::{Distribution, Standard}, + rngs::StdRng, + SeedableRng, + }; + use proptest::{prelude::Strategy, prop_assume, prop_oneof, proptest, strategy::Just}; + + use super::*; + use crate::{ + algebra::{fields, ntt, random_vector}, + transcript::{codecs::U64, DomainSeparator}, + }; + + impl Config + where + Standard: Distribution, + { + pub fn arbitrary() -> impl Strategy { + let valid_sizes = (1..=256) + .filter(|&n| ntt::next_order::(n) == Some(n)) + .collect::>(); + + // mask_length=0: tests non-ZK mask tree (no IRS masking on the committed vectors). + let mask_length = prop_oneof![ + 3 => Just(0_usize), + 7 => 1_usize..=4, + ]; + ( + 1_usize..=6, + proptest::sample::select(valid_sizes), + mask_length, + ) + .prop_flat_map(|(num_masks, vector_size, mask_length)| { + let c_zk = IrsConfig::>::arbitrary( + Identity::new(), + 2 * num_masks, + vector_size, + mask_length, + 1, + ); + (Just(num_masks), c_zk) + }) + .prop_filter( + "mask proximity requires out_domain_samples = 0", + |(_, c_zk)| c_zk.out_domain_samples == 0, + ) + .prop_map(|(num_masks, c_zk)| Self::new(c_zk, num_masks)) + } + } + + fn test_config(seed: u64, config: &Config) + where + F: Field + Codec<[u8]> + 'static, + Standard: Distribution, + Hash: crate::transcript::ProverMessage<[u8]>, + { + let instance = U64(seed); + let ds = DomainSeparator::protocol(config) + .session(&format!("Test at {}:{}", file!(), line!())) + .instance(&instance); + let mut rng = StdRng::seed_from_u64(seed); + + let original_msgs: Vec> = (0..config.num_masks) + .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) + .collect(); + + let mut prover_state = ProverState::new_std(&ds); + let witness = config.commit(&mut prover_state, &original_msgs); + config.prove(&mut prover_state, &witness, &original_msgs); + let proof = prover_state.proof(); + + let mut verifier_state = VerifierState::new_std(&ds, &proof); + let commitment = config.receive_commitment(&mut verifier_state).unwrap(); + config.verify(&mut verifier_state, &commitment).unwrap(); + verifier_state.check_eof().unwrap(); + } + + fn test + 'static>() + where + Standard: Distribution, + Hash: crate::transcript::ProverMessage<[u8]>, + { + crate::tests::init(); + proptest!(|(seed: u64, config in Config::arbitrary())| { + test_config(seed, &config); + }); + } + + #[test] + fn test_field64_1() { + test::(); + } + + #[test] + #[ignore = "Somewhat expensive and redundant"] + fn test_field64_2() { + test::(); + } + + #[test] + #[ignore = "Somewhat expensive and redundant"] + fn test_field64_3() { + test::(); + } + + #[test] + #[ignore = "Somewhat expensive and redundant"] + fn test_field128() { + test::(); + } + + #[test] + #[ignore = "Somewhat expensive and redundant"] + fn test_field192() { + test::(); + } + + #[test] + #[ignore = "Somewhat expensive and redundant"] + fn test_field256() { + test::(); + } + + /// Pre-γ tamper: commit honestly, then prove with different original masks. + fn test_tampered_mask_config(seed: u64, config: &Config) + where + F: Field + Codec<[u8]> + 'static, + Standard: Distribution, + Hash: crate::transcript::ProverMessage<[u8]>, + { + let instance = U64(seed); + let ds = DomainSeparator::protocol(config) + .session(&format!("Test at {}:{}", file!(), line!())) + .instance(&instance); + let mut rng = StdRng::seed_from_u64(seed); + + let original_msgs: Vec> = (0..config.num_masks) + .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) + .collect(); + + let mut prover_state = ProverState::new_std(&ds); + let witness = config.commit(&mut prover_state, &original_msgs); + + let mut tampered_msgs = original_msgs; + tampered_msgs[0][0] += F::ONE; + config.prove(&mut prover_state, &witness, &tampered_msgs); + let proof = prover_state.proof(); + + let mut verifier_state = VerifierState::new_std(&ds, &proof); + let commitment = config.receive_commitment(&mut verifier_state).unwrap(); + assert_rejected(|| config.verify(&mut verifier_state, &commitment)); + } + + /// Post-γ tamper: commit honestly, then corrupt the combined message to + /// exercise the verifier's spot-check (step 4). + fn test_tampered_combined_msg_config(seed: u64, config: &Config) + where + F: Field + Codec<[u8]> + 'static, + Standard: Distribution, + Hash: crate::transcript::ProverMessage<[u8]>, + { + let instance = U64(seed); + let ds = DomainSeparator::protocol(config) + .session(&format!("Test at {}:{}", file!(), line!())) + .instance(&instance); + let mut rng = StdRng::seed_from_u64(seed); + + let original_msgs: Vec> = (0..config.num_masks) + .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) + .collect(); + + let mut prover_state = ProverState::new_std(&ds); + let witness = config.commit(&mut prover_state, &original_msgs); + + let gamma: F = prover_state.verifier_message(); + let irs_masks_per_vector = + config.c_zk_commit.mask_length * config.c_zk_commit.interleaving_depth; + + for (i, (orig_msg, fresh_msg)) in original_msgs + .iter() + .zip(witness.fresh_msgs.iter()) + .enumerate() + { + let mut combined_msg = scalar_mul_add_new(fresh_msg, gamma, orig_msg); + if i == 0 { + combined_msg[0] += F::ONE; + } + prover_state.prover_messages(&combined_msg); + + if irs_masks_per_vector > 0 { + let orig_r = &witness.mask_witness.masks + [i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; + let fresh_r = &witness.mask_witness.masks[(config.num_masks + i) + * irs_masks_per_vector + ..(config.num_masks + i + 1) * irs_masks_per_vector]; + let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); + prover_state.prover_messages(&combined_r); + } + } + + config + .c_zk_commit + .open(&mut prover_state, &[&witness.mask_witness]); + let proof = prover_state.proof(); + + let mut verifier_state = VerifierState::new_std(&ds, &proof); + let commitment = config.receive_commitment(&mut verifier_state).unwrap(); + assert_rejected(|| config.verify(&mut verifier_state, &commitment)); + } + + fn assert_rejected(verify: impl FnOnce() -> VerificationResult) { + #[cfg(feature = "verifier_panics")] + use std::panic::AssertUnwindSafe; + #[cfg(feature = "verifier_panics")] + { + let result = std::panic::catch_unwind(AssertUnwindSafe(verify)); + assert!( + !matches!(result, Ok(Ok(_))), + "verification unexpectedly succeeded" + ); + } + #[cfg(not(feature = "verifier_panics"))] + { + assert!(verify().is_err(), "verification unexpectedly succeeded"); + } + } + + #[test] + fn test_tampered_mask_rejected() { + crate::tests::init(); + proptest!(|(seed: u64, config in Config::::arbitrary())| { + prop_assume!(config.c_zk_commit.in_domain_samples > 0); + test_tampered_mask_config(seed, &config); + }); + } + + #[test] + fn test_tampered_combined_msg_rejected() { + crate::tests::init(); + proptest!(|(seed: u64, config in Config::::arbitrary())| { + prop_assume!(config.c_zk_commit.in_domain_samples > 0); + test_tampered_combined_msg_config(seed, &config); + }); + } +} diff --git a/src/protocols/mod.rs b/src/protocols/mod.rs index 3c232ccb..17e58040 100644 --- a/src/protocols/mod.rs +++ b/src/protocols/mod.rs @@ -13,6 +13,7 @@ pub mod challenge_indices; pub mod code_switch; pub mod geometric_challenge; pub mod irs_commit; +pub mod mask_proximity; pub mod matrix_commit; pub mod merkle_tree; pub mod proof_of_work; diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index f721edd1..5018e22d 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -23,6 +23,13 @@ use crate::{ utils::chunks_exact_or_empty, }; +/// Output from the sumcheck protocol (shared by prover and verifier). +#[must_use] +pub struct SumcheckOpening { + pub round_challenges: Vec, + pub mask_rlc: F, +} + #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config @@ -66,7 +73,7 @@ impl Config { b: &mut Vec, sum: &mut F, masks: &[F], - ) -> (Vec, F, F) + ) -> SumcheckOpening where H: DuplexSpongeInterface, R: CryptoRng + RngCore, @@ -77,6 +84,7 @@ impl Config { assert!( self.num_rounds == 0 || self.initial_size.next_power_of_two() >= 1 << self.num_rounds ); + // Mask must cover all 3 sumcheck polynomial coefficients (c0, c1, c2) assert!(self.mask_length == 0 || self.mask_length >= 3); assert_eq!(a.len(), self.initial_size); assert_eq!(b.len(), self.initial_size); @@ -84,29 +92,24 @@ impl Config { assert_eq!(masks.len(), self.num_rounds * self.mask_length); let half = F::from(2).inverse().unwrap(); - // Send mask sum and get combination randomness. let mut mask_sum = F::ZERO; let mut mask_rlc = F::ONE; - if !masks.is_empty() { + if self.mask_length > 0 && self.num_rounds > 0 { let sum_multiple = F::from(1 << self.num_rounds.saturating_sub(1)); - mask_sum = masks - .chunks_exact(self.mask_length) - .map(eval_01) // s(0) + s(1) - .sum::() - * sum_multiple; + mask_sum = masks.chunks_exact(self.mask_length).map(eval_01).sum::() * sum_multiple; prover_state.prover_message(&mask_sum); mask_rlc = prover_state.verifier_message(); } // We do a staggered Sumcheck loop so we can merge the inner fold+compute loops. let mut univariate = Vec::new(); - let mut res = Vec::with_capacity(self.num_rounds); - let mut folding_randomness = None; + let mut round_challenges = Vec::with_capacity(self.num_rounds); + let mut prev_round_challenge = None; for (round, mask) in chunks_exact_or_empty(masks, self.mask_length, self.num_rounds).enumerate() { // Fold and compute sumcheck polynomial in one pass. - let (c0, c2) = if let Some(w) = folding_randomness { + let (c0, c2) = if let Some(w) = prev_round_challenge { fold_and_compute_polynomial(a, b, w) } else { compute_sumcheck_polynomial(a, b) @@ -137,22 +140,25 @@ impl Config { // Receive the random evaluation point and update the sum self.round_pow.prove(prover_state); let r = prover_state.verifier_message::(); - res.push(r); + round_challenges.push(r); *sum = (c2 * r + c1) * r + c0; - if !masks.is_empty() { + if self.mask_length > 0 && self.num_rounds > 0 { let masked_sum = univariate_evaluate(&univariate, r); mask_sum = masked_sum - mask_rlc * *sum; } - folding_randomness = Some(r); + prev_round_challenge = Some(r); } - if let Some(w) = folding_randomness { + if let Some(w) = prev_round_challenge { // Final fold of the inputs (no polynomial computation) fold(a, w); fold(b, w); } *sum = mask_sum + mask_rlc * *sum; - (res, mask_sum, mask_rlc) + SumcheckOpening { + round_challenges, + mask_rlc, + } } #[cfg_attr(feature = "tracing", instrument(skip_all))] @@ -160,7 +166,7 @@ impl Config { &self, verifier_state: &mut VerifierState, sum: &mut F, - ) -> VerificationResult<(Vec, F)> + ) -> VerificationResult> where H: DuplexSpongeInterface, F: Codec<[H::U]>, @@ -170,6 +176,7 @@ impl Config { assert!( self.num_rounds == 0 || self.initial_size.next_power_of_two() >= 1 << self.num_rounds ); + // Mask must cover all 3 sumcheck polynomial coefficients (c0, c1, c2) assert!(self.mask_length == 0 || self.mask_length >= 3); let mut mask_rlc = F::ONE; @@ -180,7 +187,7 @@ impl Config { } let mut univariate = vec![F::ZERO; self.mask_length.max(3)]; - let mut res = Vec::with_capacity(self.num_rounds); + let mut round_challenges = Vec::with_capacity(self.num_rounds); for _ in 0..self.num_rounds { // Receive all but linear coefficient. univariate[0] = verifier_state.prover_message()?; @@ -195,13 +202,16 @@ impl Config { self.round_pow.verify(verifier_state)?; // Receive the random evaluation point - let folding_randomness = verifier_state.verifier_message::(); - res.push(folding_randomness); + let round_challenge = verifier_state.verifier_message::(); + round_challenges.push(round_challenge); // Update the sum - *sum = univariate_evaluate(&univariate, folding_randomness); + *sum = univariate_evaluate(&univariate, round_challenge); } - Ok((res, mask_rlc)) + Ok(SumcheckOpening { + round_challenges, + mask_rlc, + }) } } @@ -228,7 +238,6 @@ fn eval_01(coefficients: &[F]) -> F { #[cfg(test)] mod tests { - // TODO: Proptest based tests checking invariants and post conditions. use ark_std::rand::{ distributions::{Distribution, Standard}, rngs::StdRng, @@ -245,14 +254,16 @@ mod tests { multilinear_extend, random_vector, }, transcript::DomainSeparator, - utils::zip_strict, }; - impl Config { + impl Config + where + Standard: Distribution, + { pub fn arbitrary() -> impl Strategy { let mask_length = prop_oneof![ 3 => Just(0_usize), - 7 => 3_usize..100, + 7 => 3_usize..20, ]; (0_usize..(1 << 12), 0_usize..12, mask_length).prop_map( |(initial_size, num_rounds, mask_length)| { @@ -273,7 +284,7 @@ mod tests { #[cfg_attr(feature = "tracing", instrument)] fn test_config(seed: u64, config: &Config) where - F: Field + Codec, + F: Field + Codec<[u8]> + 'static, Standard: Distribution, { // Pseudo-random Instance @@ -292,46 +303,56 @@ mod tests { let mut covector = initial_covector.clone(); let mut sum = initial_sum; let mut prover_state = ProverState::new_std(&ds); - let (point, mask_sum, mask_rlc) = config.prove( + let SumcheckOpening { + round_challenges: point, + mask_rlc, + } = config.prove( &mut prover_state, &mut vector, &mut covector, &mut sum, &masks, ); - let expected_mask_sum = zip_strict( - chunks_exact_or_empty(&masks, config.mask_length, config.num_rounds), - &point, - ) - .map(|(m, x)| univariate_evaluate(m, *x)) - .sum::(); assert_eq!(vector.len(), config.final_size()); assert_eq!(covector.len(), config.final_size()); - assert_eq!(mask_sum, expected_mask_sum); - assert_eq!(mask_sum + mask_rlc * dot(&vector, &covector), sum); if config.final_size() == 1 { assert_eq!(multilinear_extend(&initial_vector, &point), vector[0]); assert_eq!(multilinear_extend(&initial_covector, &point), covector[0]); } else { // TODO: Check correct folding. } + + let expected_mask_sum: F = + chunks_exact_or_empty(&masks, config.mask_length, config.num_rounds) + .zip(&point) + .map(|(m, x)| univariate_evaluate(m, *x)) + .sum(); + assert_eq!(sum, expected_mask_sum + mask_rlc * dot(&vector, &covector)); + let proof = prover_state.proof(); // Verifier let mut verifier_sum = initial_sum; let mut verifier_state = VerifierState::new_std(&ds, &proof); - let (verifier_point, verifier_mask_rlc) = config + let SumcheckOpening { + round_challenges: verifier_point, + mask_rlc: verifier_mask_rlc, + } = config .verify(&mut verifier_state, &mut verifier_sum) .unwrap(); assert_eq!(verifier_point, point); assert_eq!(verifier_mask_rlc, mask_rlc); assert_eq!(verifier_sum, sum); verifier_state.check_eof().unwrap(); + + // Non-ZK path: mask_rlc defaults to ONE (no combination randomness sampled) + if config.mask_length == 0 || config.num_rounds == 0 { + assert_eq!(mask_rlc, F::ONE); + } } - fn test_sumcheck() + fn test + 'static>() where - F: Field + Codec, Standard: Distribution, { crate::tests::init(); @@ -384,36 +405,36 @@ mod tests { #[test] fn test_field64_1() { - test_sumcheck::(); + test::(); } #[test] #[ignore = "Somewhat expensive and redundant"] fn test_field64_2() { - test_sumcheck::(); + test::(); } #[test] #[ignore = "Somewhat expensive and redundant"] fn test_field64_3() { - test_sumcheck::(); + test::(); } #[test] #[ignore = "Somewhat expensive and redundant"] fn test_field128() { - test_sumcheck::(); + test::(); } #[test] #[ignore = "Somewhat expensive and redundant"] fn test_field192() { - test_sumcheck::(); + test::(); } #[test] #[ignore = "Somewhat expensive and redundant"] fn test_field256() { - test_sumcheck::(); + test::(); } } diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 87962f44..04b4473e 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -196,7 +196,7 @@ impl Config { let mut folding_randomness = if has_constraints { self.initial_sumcheck .prove(prover_state, &mut vector, &mut covector, &mut the_sum, &[]) - .0 + .round_challenges } else { // There are no constraints yet, so we can skip the sumcheck. // (If we did run it, all sumcheck vectors would be constant zero) @@ -268,7 +268,7 @@ impl Config { folding_randomness = round_config .sumcheck .prove(prover_state, &mut vector, &mut covector, &mut the_sum, &[]) - .0; + .round_challenges; evaluation_point.extend(folding_randomness.iter().copied()); debug_assert_eq!(dot(&vector, &covector), the_sum); @@ -304,7 +304,7 @@ impl Config { let final_folding_randomness = self .final_sumcheck .prove(prover_state, &mut vector, &mut covector, &mut the_sum, &[]) - .0; + .round_challenges; evaluation_point.extend(final_folding_randomness.iter().copied()); FinalClaim { diff --git a/src/protocols/whir/verifier.rs b/src/protocols/whir/verifier.rs index df9c7c71..8421ceab 100644 --- a/src/protocols/whir/verifier.rs +++ b/src/protocols/whir/verifier.rs @@ -128,7 +128,7 @@ impl Config { } else { self.initial_sumcheck .verify(verifier_state, &mut the_sum)? - .0 + .round_challenges }; round_folding_randomness.push(folding_randomness); @@ -184,7 +184,7 @@ impl Config { let folding_randomness = round_config .sumcheck .verify(verifier_state, &mut the_sum)? - .0; + .round_challenges; round_folding_randomness.push(folding_randomness); prev_commitment = RoundCommitment::Round { commitment }; @@ -226,7 +226,10 @@ impl Config { } // Final sumcheck - let final_sumcheck_randomness = self.final_sumcheck.verify(verifier_state, &mut the_sum)?.0; + let final_sumcheck_randomness = self + .final_sumcheck + .verify(verifier_state, &mut the_sum)? + .round_challenges; round_folding_randomness.push(final_sumcheck_randomness.clone()); // Compute folding randomness across all rounds diff --git a/src/transcript/mock_sponge.rs b/src/transcript/mock_sponge.rs index 822463db..df17af08 100644 --- a/src/transcript/mock_sponge.rs +++ b/src/transcript/mock_sponge.rs @@ -13,7 +13,7 @@ impl DuplexSpongeInterface for MockSponge<'_> { fn absorb(&mut self, input: &[Self::U]) -> &mut Self { if let Some(absorb) = self.absorb.as_mut() { - assert!(&absorb[..input.len()] == input); + assert_eq!(&absorb[..input.len()], input); *absorb = &absorb[input.len()..]; } self @@ -27,10 +27,10 @@ impl DuplexSpongeInterface for MockSponge<'_> { fn ratchet(&mut self) -> &mut Self { if let Some(absorb) = self.absorb.as_mut() { - assert!(&absorb[..7] == b"RATCHET"); + assert_eq!(&absorb[..7], b"RATCHET"); *absorb = &absorb[7..]; } - assert!(&self.squeeze[..7] == b"RATCHET"); + assert_eq!(&self.squeeze[..7], b"RATCHET"); self.squeeze = &self.squeeze[7..]; self } From e6f9c3eca8659439c718d31958562716a9014441 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Tue, 2 Jun 2026 13:59:01 +0200 Subject: [PATCH 02/35] wip: introduce buffer ops trait and cpu buffer --- src/algebra/buffer.rs | 460 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 460 insertions(+) create mode 100644 src/algebra/buffer.rs diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs new file mode 100644 index 00000000..89cc4807 --- /dev/null +++ b/src/algebra/buffer.rs @@ -0,0 +1,460 @@ +use std::cmp::max; + +use ark_ff::{AdditiveGroup, Field}; + +use crate::algebra::embedding::{Embedding, Identity}; +#[cfg(feature = "parallel")] +use crate::utils::workload_size; +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +pub trait BufferOps: Clone { + fn len(&self) -> usize; + fn is_empty(&self) -> bool; + // zero pad to 2**log_m + fn zero_pad(&mut self, log_m: usize); + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target; + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &impl BufferOps, + ) -> M::Target; + fn dot(&self, other: &Self) -> F; + fn as_slice(&self) -> &[F]; + fn as_ro_buffer(&self, size: usize) -> impl BufferOps; +} + +// read-only buffer ops +pub trait ROBufferOps { + fn split_at(&self, mid: usize) -> (&impl BufferOps, &impl BufferOps); +} + +#[derive(Clone)] +pub struct SliceCpuBuffer<'a, F: Field> { + data: &'a [F], +} + +#[derive(Clone)] +pub struct CpuBuffer { + data: Vec, + len: usize, +} + +impl CpuBuffer { + pub fn from_vec(source: Vec) -> Self { + let len = source.len(); + Self { data: source, len } + } + + pub fn from_slice(source: &[F]) -> Self { + Self { + data: Vec::from(source), + len: source.len(), + } + } +} + +impl BufferOps for CpuBuffer { + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } + + fn zero_pad(&mut self, log_m: usize) { + if !self.is_empty() { + self.data.resize(1 << log_m, F::ZERO); + } + } + + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + #[inline] + fn eval_exact( + embedding: &M, + evals: &[M::Source], + point: &[M::Target], + ) -> M::Target { + debug_assert_eq!(evals.len(), 1 << point.len()); + + // Helper to compute (a + (b - a) * c) efficiently with a, b in source field. + let mixed = |a, b, c| embedding.mixed_add(embedding.mixed_mul(c, b - a), a); + + match point { + [] => embedding.map(evals[0]), + [x] => mixed(evals[0], evals[1], *x), + [x0, x1] => { + let a0 = mixed(evals[0], evals[1], *x1); + let a1 = mixed(evals[2], evals[3], *x1); + a0 + (a1 - a0) * *x0 + } + [x0, x1, x2] => { + let a00 = mixed(evals[0], evals[1], *x2); + let a01 = mixed(evals[2], evals[3], *x2); + let a10 = mixed(evals[4], evals[5], *x2); + let a11 = mixed(evals[6], evals[7], *x2); + let a0 = a00 + (a01 - a00) * *x1; + let a1 = a10 + (a11 - a10) * *x1; + a0 + (a1 - a0) * *x0 + } + [x0, x1, x2, x3] => { + let a000 = mixed(evals[0], evals[1], *x3); + let a001 = mixed(evals[2], evals[3], *x3); + let a010 = mixed(evals[4], evals[5], *x3); + let a011 = mixed(evals[6], evals[7], *x3); + let a100 = mixed(evals[8], evals[9], *x3); + let a101 = mixed(evals[10], evals[11], *x3); + let a110 = mixed(evals[12], evals[13], *x3); + let a111 = mixed(evals[14], evals[15], *x3); + let a00 = a000 + (a001 - a000) * *x2; + let a01 = a010 + (a011 - a010) * *x2; + let a10 = a100 + (a101 - a100) * *x2; + let a11 = a110 + (a111 - a110) * *x2; + let a0 = a00 + (a01 - a00) * *x1; + let a1 = a10 + (a11 - a10) * *x1; + a0 + (a1 - a0) * *x0 + } + [x, tail @ ..] => { + let (f0, f1) = evals.split_at(evals.len() / 2); + #[cfg(not(feature = "parallel"))] + let (f0, f1) = ( + eval_exact(embedding, f0, tail), + eval_exact(embedding, f1, tail), + ); + + #[cfg(feature = "parallel")] + let (f0, f1) = { + use crate::utils::workload_size; + if evals.len() > workload_size::() { + rayon::join( + || eval_exact(embedding, f0, tail), + || eval_exact(embedding, f1, tail), + ) + } else { + ( + eval_exact(embedding, f0, tail), + eval_exact(embedding, f1, tail), + ) + } + }; + + f0 + (f1 - f0) * *x + } + } + } + + #[inline] + fn eval_partial( + embedding: &M, + evals: &[M::Source], + point: &[M::Target], + ) -> M::Target { + let size = 1 << point.len(); + debug_assert!(evals.len() <= size); + if evals.is_empty() { + return M::Target::ZERO; + } + if evals.len() == size { + return eval_exact(embedding, evals, point); + } + + match point { + [] => embedding.map(evals[0]), + [x, tail @ ..] => { + let half = size / 2; + + // Only low half has data; high half is all implicit zeros. + if evals.len() <= half { + let f0 = eval_partial(embedding, evals, tail); + return f0 * (M::Target::ONE - *x); + } + + // Low subtree is exact/full, high subtree is partial. + let (low, high) = evals.split_at(half); + + #[cfg(not(feature = "parallel"))] + let (f0, f1) = ( + eval_exact(embedding, low, tail), + eval_partial(embedding, high, tail), + ); + + #[cfg(feature = "parallel")] + let (f0, f1) = { + use crate::utils::workload_size; + if evals.len() > workload_size::() { + rayon::join( + || eval_exact(embedding, low, tail), + || eval_partial(embedding, high, tail), + ) + } else { + ( + eval_exact(embedding, low, tail), + eval_partial(embedding, high, tail), + ) + } + }; + + f0 + (f1 - f0) * *x + } + } + } + + eval_partial(embedding, &self.data, point) + } + + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &impl BufferOps, + ) -> M::Target { + assert_eq!(self.len(), other.len()); + + let a = other.as_slice(); + let b = self.as_slice(); + + #[cfg(feature = "parallel")] + if a.len() > workload_size::() { + return a + .par_iter() + .zip(b) + .map(|(a, b)| embedding.mixed_mul(*a, *b)) + .sum(); + } + + a.iter() + .zip(b) + .map(|(a, b)| embedding.mixed_mul(*a, *b)) + .sum() + } + + fn dot(&self, other: &Self) -> F { + self.mixed_dot(&Identity::new(), other) + } + + fn as_slice(&self) -> &[F] { + &self.data[..self.len] + } + + fn as_ro_buffer(&self, size: usize) -> impl BufferOps { + SliceCpuBuffer::from_buffer_with_size(self, size) + } +} + +impl<'a, F: Field> SliceCpuBuffer<'a, F> { + pub fn from_buffer(buffer: &'a CpuBuffer) -> Self { + Self { + data: buffer.as_slice(), + } + } + + pub fn from_buffer_with_size(buffer: &'a CpuBuffer, size: usize) -> Self { + Self { + data: &buffer.data[..max(buffer.len, size)], + } + } + + pub fn from_slice_with_size(slice: &'a &[F], size: usize) -> Self { + assert!(size <= slice.len()); + Self { + data: &slice[..size], + } + } +} + +impl<'a, F: Field> BufferOps for SliceCpuBuffer<'a, F> { + fn len(&self) -> usize { + self.data.len() + } + + fn is_empty(&self) -> bool { + self.data.is_empty() + } + + fn zero_pad(&mut self, log_m: usize) { + panic!("read only") + } + + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + #[inline] + fn eval_exact( + embedding: &M, + evals: &[M::Source], + point: &[M::Target], + ) -> M::Target { + debug_assert_eq!(evals.len(), 1 << point.len()); + + // Helper to compute (a + (b - a) * c) efficiently with a, b in source field. + let mixed = |a, b, c| embedding.mixed_add(embedding.mixed_mul(c, b - a), a); + + match point { + [] => embedding.map(evals[0]), + [x] => mixed(evals[0], evals[1], *x), + [x0, x1] => { + let a0 = mixed(evals[0], evals[1], *x1); + let a1 = mixed(evals[2], evals[3], *x1); + a0 + (a1 - a0) * *x0 + } + [x0, x1, x2] => { + let a00 = mixed(evals[0], evals[1], *x2); + let a01 = mixed(evals[2], evals[3], *x2); + let a10 = mixed(evals[4], evals[5], *x2); + let a11 = mixed(evals[6], evals[7], *x2); + let a0 = a00 + (a01 - a00) * *x1; + let a1 = a10 + (a11 - a10) * *x1; + a0 + (a1 - a0) * *x0 + } + [x0, x1, x2, x3] => { + let a000 = mixed(evals[0], evals[1], *x3); + let a001 = mixed(evals[2], evals[3], *x3); + let a010 = mixed(evals[4], evals[5], *x3); + let a011 = mixed(evals[6], evals[7], *x3); + let a100 = mixed(evals[8], evals[9], *x3); + let a101 = mixed(evals[10], evals[11], *x3); + let a110 = mixed(evals[12], evals[13], *x3); + let a111 = mixed(evals[14], evals[15], *x3); + let a00 = a000 + (a001 - a000) * *x2; + let a01 = a010 + (a011 - a010) * *x2; + let a10 = a100 + (a101 - a100) * *x2; + let a11 = a110 + (a111 - a110) * *x2; + let a0 = a00 + (a01 - a00) * *x1; + let a1 = a10 + (a11 - a10) * *x1; + a0 + (a1 - a0) * *x0 + } + [x, tail @ ..] => { + let (f0, f1) = evals.split_at(evals.len() / 2); + #[cfg(not(feature = "parallel"))] + let (f0, f1) = ( + eval_exact(embedding, f0, tail), + eval_exact(embedding, f1, tail), + ); + + #[cfg(feature = "parallel")] + let (f0, f1) = { + use crate::utils::workload_size; + if evals.len() > workload_size::() { + rayon::join( + || eval_exact(embedding, f0, tail), + || eval_exact(embedding, f1, tail), + ) + } else { + ( + eval_exact(embedding, f0, tail), + eval_exact(embedding, f1, tail), + ) + } + }; + + f0 + (f1 - f0) * *x + } + } + } + + #[inline] + fn eval_partial( + embedding: &M, + evals: &[M::Source], + point: &[M::Target], + ) -> M::Target { + let size = 1 << point.len(); + debug_assert!(evals.len() <= size); + if evals.is_empty() { + return M::Target::ZERO; + } + if evals.len() == size { + return eval_exact(embedding, evals, point); + } + + match point { + [] => embedding.map(evals[0]), + [x, tail @ ..] => { + let half = size / 2; + + // Only low half has data; high half is all implicit zeros. + if evals.len() <= half { + let f0 = eval_partial(embedding, evals, tail); + return f0 * (M::Target::ONE - *x); + } + + // Low subtree is exact/full, high subtree is partial. + let (low, high) = evals.split_at(half); + + #[cfg(not(feature = "parallel"))] + let (f0, f1) = ( + eval_exact(embedding, low, tail), + eval_partial(embedding, high, tail), + ); + + #[cfg(feature = "parallel")] + let (f0, f1) = { + use crate::utils::workload_size; + if evals.len() > workload_size::() { + rayon::join( + || eval_exact(embedding, low, tail), + || eval_partial(embedding, high, tail), + ) + } else { + ( + eval_exact(embedding, low, tail), + eval_partial(embedding, high, tail), + ) + } + }; + + f0 + (f1 - f0) * *x + } + } + } + + eval_partial(embedding, &self.data, point) + } + + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &impl BufferOps, + ) -> M::Target { + assert_eq!(self.len(), other.len()); + + let a = other.as_slice(); + let b = self.as_slice(); + + #[cfg(feature = "parallel")] + if a.len() > workload_size::() { + return a + .par_iter() + .zip(b) + .map(|(a, b)| embedding.mixed_mul(*a, *b)) + .sum(); + } + + a.iter() + .zip(b) + .map(|(a, b)| embedding.mixed_mul(*a, *b)) + .sum() + } + + fn dot(&self, other: &Self) -> F { + self.mixed_dot(&Identity::new(), other) + } + + fn as_slice(&self) -> &[F] { + &self.data + } + + fn as_ro_buffer(&self, size: usize) -> impl BufferOps { + SliceCpuBuffer::from_slice_with_size(&self.data, size) + } +} From 3d284a980909eb67d00cf188eec631cb2da06a36 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 3 Jun 2026 12:11:06 +0200 Subject: [PATCH 03/35] feat: buffer abstraction with basic interface and cpu buffer implementation --- src/algebra/buffer.rs | 607 +++++++++++++++++------------------------- 1 file changed, 241 insertions(+), 366 deletions(-) diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index 89cc4807..a38a95a4 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -1,18 +1,35 @@ -use std::cmp::max; - -use ark_ff::{AdditiveGroup, Field}; - -use crate::algebra::embedding::{Embedding, Identity}; -#[cfg(feature = "parallel")] -use crate::utils::workload_size; -#[cfg(feature = "parallel")] -use rayon::prelude::*; +use ark_ff::Field; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; +use spongefish::DuplexSpongeInterface; + +use crate::{ + algebra::{ + embedding::{Embedding, Identity}, + mixed_dot, mixed_multilinear_extend, mixed_univariate_evaluate, ntt, + }, + hash::Hash, + protocols::matrix_commit, + transcript::{ProverMessage, ProverState}, + type_info::TypeInfo, + utils::chunks_exact_or_empty, +}; pub trait BufferOps: Clone { + type Buffer: BufferOps; + type Matrix: MatrixBufferOps; + fn len(&self) -> usize; - fn is_empty(&self) -> bool; - // zero pad to 2**log_m - fn zero_pad(&mut self, log_m: usize); + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution; + + fn zero_pad(&mut self); fn mixed_extend, T: Field>( &self, embedding: &M, @@ -21,55 +38,179 @@ pub trait BufferOps: Clone { fn mixed_dot, T: Field>( &self, embedding: &M, - other: &impl BufferOps, + other: &Self::Buffer, + ) -> M::Target; + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, ) -> M::Target; + fn interleaved_rs_encode( + vectors: &[&Self], + masks: &Self, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + ) -> Self::Matrix + where + F: 'static; fn dot(&self, other: &Self) -> F; - fn as_slice(&self) -> &[F]; - fn as_ro_buffer(&self, size: usize) -> impl BufferOps; } -// read-only buffer ops -pub trait ROBufferOps { - fn split_at(&self, mid: usize) -> (&impl BufferOps, &impl BufferOps); +fn interleaved_rs_encode_slices( + vectors: &[&[F]], + masks: &[F], + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, +) -> CpuMatrix { + let messages = vectors + .iter() + .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) + .collect::>(); + CpuMatrix::from_vec( + ntt::interleaved_rs_encode(&messages, masks, codeword_length), + codeword_length, + vectors.len() * interleaving_depth, + ) } -#[derive(Clone)] -pub struct SliceCpuBuffer<'a, F: Field> { - data: &'a [F], +pub trait MatrixBufferOps { + fn len(&self) -> usize; + fn num_rows(&self) -> usize; + fn num_cols(&self) -> usize; + + fn commit_rows( + &self, + config: &matrix_commit::Config, + prover_state: &mut ProverState, + ) -> matrix_commit::Witness + where + F: TypeInfo + matrix_commit::Encodable + Send + Sync, + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + Hash: ProverMessage<[H::U]>; + + fn read_rows(&self, indices: &[usize]) -> Vec; +} + +#[derive( + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + Debug, + Default, + serde::Serialize, + serde::Deserialize, +)] +pub struct CpuMatrix { + data: Vec, + num_rows: usize, + num_cols: usize, +} + +impl CpuMatrix { + pub fn from_vec(data: Vec, num_rows: usize, num_cols: usize) -> Self { + assert_eq!(data.len(), num_rows * num_cols); + Self { + data, + num_rows, + num_cols, + } + } + + fn row(&self, row: usize) -> &[F] { + let start = row * self.num_cols; + let end = start + self.num_cols; + &self.data[start..end] + } } +impl MatrixBufferOps for CpuMatrix +where + F: Field + TypeInfo + matrix_commit::Encodable + Send + Sync, +{ + fn len(&self) -> usize { + self.data.len() + } + + fn num_rows(&self) -> usize { + self.num_rows + } + + fn num_cols(&self) -> usize { + self.num_cols + } + + fn commit_rows( + &self, + config: &matrix_commit::Config, + prover_state: &mut ProverState, + ) -> matrix_commit::Witness + where + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + Hash: ProverMessage<[H::U]>, + { + assert_eq!(config.num_rows(), self.num_rows); + assert_eq!(config.num_cols, self.num_cols); + config.commit(prover_state, &self.data) + } + + fn read_rows(&self, indices: &[usize]) -> Vec { + let mut rows = Vec::with_capacity(indices.len() * self.num_cols); + for &index in indices { + assert!(index < self.num_rows); + rows.extend_from_slice(self.row(index)); + } + rows + } +} #[derive(Clone)] pub struct CpuBuffer { data: Vec, - len: usize, } impl CpuBuffer { pub fn from_vec(source: Vec) -> Self { - let len = source.len(); - Self { data: source, len } + Self { data: source } } pub fn from_slice(source: &[F]) -> Self { Self { data: Vec::from(source), - len: source.len(), } } + + fn as_slice(&self) -> &[F] { + self.data.as_slice() + } } impl BufferOps for CpuBuffer { + type Buffer = CpuBuffer; + type Matrix = CpuMatrix; + fn len(&self) -> usize { - self.len + self.data.len() } - fn is_empty(&self) -> bool { - self.len == 0 + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + Self { + data: (0..length).map(|_| rng.gen()).collect(), + } } - fn zero_pad(&mut self, log_m: usize) { + fn zero_pad(&mut self) { if !self.is_empty() { - self.data.resize(1 << log_m, F::ZERO); + self.data.resize(self.len().next_power_of_two(), F::ZERO); } } @@ -78,209 +219,70 @@ impl BufferOps for CpuBuffer { embedding: &M, point: &[M::Target], ) -> M::Target { - #[inline] - fn eval_exact( - embedding: &M, - evals: &[M::Source], - point: &[M::Target], - ) -> M::Target { - debug_assert_eq!(evals.len(), 1 << point.len()); - - // Helper to compute (a + (b - a) * c) efficiently with a, b in source field. - let mixed = |a, b, c| embedding.mixed_add(embedding.mixed_mul(c, b - a), a); - - match point { - [] => embedding.map(evals[0]), - [x] => mixed(evals[0], evals[1], *x), - [x0, x1] => { - let a0 = mixed(evals[0], evals[1], *x1); - let a1 = mixed(evals[2], evals[3], *x1); - a0 + (a1 - a0) * *x0 - } - [x0, x1, x2] => { - let a00 = mixed(evals[0], evals[1], *x2); - let a01 = mixed(evals[2], evals[3], *x2); - let a10 = mixed(evals[4], evals[5], *x2); - let a11 = mixed(evals[6], evals[7], *x2); - let a0 = a00 + (a01 - a00) * *x1; - let a1 = a10 + (a11 - a10) * *x1; - a0 + (a1 - a0) * *x0 - } - [x0, x1, x2, x3] => { - let a000 = mixed(evals[0], evals[1], *x3); - let a001 = mixed(evals[2], evals[3], *x3); - let a010 = mixed(evals[4], evals[5], *x3); - let a011 = mixed(evals[6], evals[7], *x3); - let a100 = mixed(evals[8], evals[9], *x3); - let a101 = mixed(evals[10], evals[11], *x3); - let a110 = mixed(evals[12], evals[13], *x3); - let a111 = mixed(evals[14], evals[15], *x3); - let a00 = a000 + (a001 - a000) * *x2; - let a01 = a010 + (a011 - a010) * *x2; - let a10 = a100 + (a101 - a100) * *x2; - let a11 = a110 + (a111 - a110) * *x2; - let a0 = a00 + (a01 - a00) * *x1; - let a1 = a10 + (a11 - a10) * *x1; - a0 + (a1 - a0) * *x0 - } - [x, tail @ ..] => { - let (f0, f1) = evals.split_at(evals.len() / 2); - #[cfg(not(feature = "parallel"))] - let (f0, f1) = ( - eval_exact(embedding, f0, tail), - eval_exact(embedding, f1, tail), - ); - - #[cfg(feature = "parallel")] - let (f0, f1) = { - use crate::utils::workload_size; - if evals.len() > workload_size::() { - rayon::join( - || eval_exact(embedding, f0, tail), - || eval_exact(embedding, f1, tail), - ) - } else { - ( - eval_exact(embedding, f0, tail), - eval_exact(embedding, f1, tail), - ) - } - }; - - f0 + (f1 - f0) * *x - } - } - } - - #[inline] - fn eval_partial( - embedding: &M, - evals: &[M::Source], - point: &[M::Target], - ) -> M::Target { - let size = 1 << point.len(); - debug_assert!(evals.len() <= size); - if evals.is_empty() { - return M::Target::ZERO; - } - if evals.len() == size { - return eval_exact(embedding, evals, point); - } - - match point { - [] => embedding.map(evals[0]), - [x, tail @ ..] => { - let half = size / 2; - - // Only low half has data; high half is all implicit zeros. - if evals.len() <= half { - let f0 = eval_partial(embedding, evals, tail); - return f0 * (M::Target::ONE - *x); - } - - // Low subtree is exact/full, high subtree is partial. - let (low, high) = evals.split_at(half); - - #[cfg(not(feature = "parallel"))] - let (f0, f1) = ( - eval_exact(embedding, low, tail), - eval_partial(embedding, high, tail), - ); - - #[cfg(feature = "parallel")] - let (f0, f1) = { - use crate::utils::workload_size; - if evals.len() > workload_size::() { - rayon::join( - || eval_exact(embedding, low, tail), - || eval_partial(embedding, high, tail), - ) - } else { - ( - eval_exact(embedding, low, tail), - eval_partial(embedding, high, tail), - ) - } - }; - - f0 + (f1 - f0) * *x - } - } - } - - eval_partial(embedding, &self.data, point) + mixed_multilinear_extend(embedding, &self.data, point) } fn mixed_dot, T: Field>( &self, embedding: &M, - other: &impl BufferOps, + other: &Self::Buffer, ) -> M::Target { - assert_eq!(self.len(), other.len()); - - let a = other.as_slice(); - let b = self.as_slice(); - - #[cfg(feature = "parallel")] - if a.len() > workload_size::() { - return a - .par_iter() - .zip(b) - .map(|(a, b)| embedding.mixed_mul(*a, *b)) - .sum(); - } - - a.iter() - .zip(b) - .map(|(a, b)| embedding.mixed_mul(*a, *b)) - .sum() + mixed_dot(embedding, other.as_slice(), self.as_slice()) } fn dot(&self, other: &Self) -> F { self.mixed_dot(&Identity::new(), other) } - fn as_slice(&self) -> &[F] { - &self.data[..self.len] + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + mixed_univariate_evaluate(embedding, &self.data, point) } - fn as_ro_buffer(&self, size: usize) -> impl BufferOps { - SliceCpuBuffer::from_buffer_with_size(self, size) + fn interleaved_rs_encode( + vectors: &[&Self], + masks: &Self, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + ) -> Self::Matrix + where + F: 'static, + { + let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); + interleaved_rs_encode_slices( + &vectors, + masks.as_slice(), + message_length, + interleaving_depth, + codeword_length, + ) } } -impl<'a, F: Field> SliceCpuBuffer<'a, F> { - pub fn from_buffer(buffer: &'a CpuBuffer) -> Self { - Self { - data: buffer.as_slice(), - } - } +impl BufferOps for Vec { + type Buffer = Vec; + type Matrix = CpuMatrix; - pub fn from_buffer_with_size(buffer: &'a CpuBuffer, size: usize) -> Self { - Self { - data: &buffer.data[..max(buffer.len, size)], - } - } - - pub fn from_slice_with_size(slice: &'a &[F], size: usize) -> Self { - assert!(size <= slice.len()); - Self { - data: &slice[..size], - } - } -} - -impl<'a, F: Field> BufferOps for SliceCpuBuffer<'a, F> { fn len(&self) -> usize { - self.data.len() + self.len() } - fn is_empty(&self) -> bool { - self.data.is_empty() + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + (0..length).map(|_| rng.gen()).collect() } - fn zero_pad(&mut self, log_m: usize) { - panic!("read only") + fn zero_pad(&mut self) { + if !self.is_empty() { + self.resize(self.len().next_power_of_two(), F::ZERO); + } } fn mixed_extend, T: Field>( @@ -288,173 +290,46 @@ impl<'a, F: Field> BufferOps for SliceCpuBuffer<'a, F> { embedding: &M, point: &[M::Target], ) -> M::Target { - #[inline] - fn eval_exact( - embedding: &M, - evals: &[M::Source], - point: &[M::Target], - ) -> M::Target { - debug_assert_eq!(evals.len(), 1 << point.len()); - - // Helper to compute (a + (b - a) * c) efficiently with a, b in source field. - let mixed = |a, b, c| embedding.mixed_add(embedding.mixed_mul(c, b - a), a); - - match point { - [] => embedding.map(evals[0]), - [x] => mixed(evals[0], evals[1], *x), - [x0, x1] => { - let a0 = mixed(evals[0], evals[1], *x1); - let a1 = mixed(evals[2], evals[3], *x1); - a0 + (a1 - a0) * *x0 - } - [x0, x1, x2] => { - let a00 = mixed(evals[0], evals[1], *x2); - let a01 = mixed(evals[2], evals[3], *x2); - let a10 = mixed(evals[4], evals[5], *x2); - let a11 = mixed(evals[6], evals[7], *x2); - let a0 = a00 + (a01 - a00) * *x1; - let a1 = a10 + (a11 - a10) * *x1; - a0 + (a1 - a0) * *x0 - } - [x0, x1, x2, x3] => { - let a000 = mixed(evals[0], evals[1], *x3); - let a001 = mixed(evals[2], evals[3], *x3); - let a010 = mixed(evals[4], evals[5], *x3); - let a011 = mixed(evals[6], evals[7], *x3); - let a100 = mixed(evals[8], evals[9], *x3); - let a101 = mixed(evals[10], evals[11], *x3); - let a110 = mixed(evals[12], evals[13], *x3); - let a111 = mixed(evals[14], evals[15], *x3); - let a00 = a000 + (a001 - a000) * *x2; - let a01 = a010 + (a011 - a010) * *x2; - let a10 = a100 + (a101 - a100) * *x2; - let a11 = a110 + (a111 - a110) * *x2; - let a0 = a00 + (a01 - a00) * *x1; - let a1 = a10 + (a11 - a10) * *x1; - a0 + (a1 - a0) * *x0 - } - [x, tail @ ..] => { - let (f0, f1) = evals.split_at(evals.len() / 2); - #[cfg(not(feature = "parallel"))] - let (f0, f1) = ( - eval_exact(embedding, f0, tail), - eval_exact(embedding, f1, tail), - ); - - #[cfg(feature = "parallel")] - let (f0, f1) = { - use crate::utils::workload_size; - if evals.len() > workload_size::() { - rayon::join( - || eval_exact(embedding, f0, tail), - || eval_exact(embedding, f1, tail), - ) - } else { - ( - eval_exact(embedding, f0, tail), - eval_exact(embedding, f1, tail), - ) - } - }; - - f0 + (f1 - f0) * *x - } - } - } - - #[inline] - fn eval_partial( - embedding: &M, - evals: &[M::Source], - point: &[M::Target], - ) -> M::Target { - let size = 1 << point.len(); - debug_assert!(evals.len() <= size); - if evals.is_empty() { - return M::Target::ZERO; - } - if evals.len() == size { - return eval_exact(embedding, evals, point); - } - - match point { - [] => embedding.map(evals[0]), - [x, tail @ ..] => { - let half = size / 2; - - // Only low half has data; high half is all implicit zeros. - if evals.len() <= half { - let f0 = eval_partial(embedding, evals, tail); - return f0 * (M::Target::ONE - *x); - } - - // Low subtree is exact/full, high subtree is partial. - let (low, high) = evals.split_at(half); - - #[cfg(not(feature = "parallel"))] - let (f0, f1) = ( - eval_exact(embedding, low, tail), - eval_partial(embedding, high, tail), - ); - - #[cfg(feature = "parallel")] - let (f0, f1) = { - use crate::utils::workload_size; - if evals.len() > workload_size::() { - rayon::join( - || eval_exact(embedding, low, tail), - || eval_partial(embedding, high, tail), - ) - } else { - ( - eval_exact(embedding, low, tail), - eval_partial(embedding, high, tail), - ) - } - }; - - f0 + (f1 - f0) * *x - } - } - } - - eval_partial(embedding, &self.data, point) + mixed_multilinear_extend(embedding, self, point) } fn mixed_dot, T: Field>( &self, embedding: &M, - other: &impl BufferOps, + other: &Self::Buffer, ) -> M::Target { - assert_eq!(self.len(), other.len()); - - let a = other.as_slice(); - let b = self.as_slice(); - - #[cfg(feature = "parallel")] - if a.len() > workload_size::() { - return a - .par_iter() - .zip(b) - .map(|(a, b)| embedding.mixed_mul(*a, *b)) - .sum(); - } - - a.iter() - .zip(b) - .map(|(a, b)| embedding.mixed_mul(*a, *b)) - .sum() + mixed_dot(embedding, other, self) } - fn dot(&self, other: &Self) -> F { - self.mixed_dot(&Identity::new(), other) + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + mixed_univariate_evaluate(embedding, self, point) } - fn as_slice(&self) -> &[F] { - &self.data + fn interleaved_rs_encode( + vectors: &[&Self], + masks: &Self, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + ) -> Self::Matrix + where + F: 'static, + { + let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); + interleaved_rs_encode_slices( + &vectors, + masks, + message_length, + interleaving_depth, + codeword_length, + ) } - fn as_ro_buffer(&self, size: usize) -> impl BufferOps { - SliceCpuBuffer::from_slice_with_size(&self.data, size) + fn dot(&self, other: &Self) -> F { + self.mixed_dot(&Identity::new(), other) } } From 7f4dd420095534285a4c3a7b409b6ca695dd7eb2 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 3 Jun 2026 12:38:49 +0200 Subject: [PATCH 04/35] feat: wire irs commit path to buffer --- src/algebra/buffer.rs | 152 ++++++++++++----------------- src/algebra/mod.rs | 1 + src/bin/benchmark.rs | 7 +- src/bin/main.rs | 7 +- src/protocols/basecase.rs | 18 ++-- src/protocols/code_switch.rs | 15 ++- src/protocols/irs_commit.rs | 105 ++++++++++++-------- src/protocols/mask_proximity.rs | 38 +++++--- src/protocols/whir/mod.rs | 43 +++++--- src/protocols/whir/prover.rs | 6 +- src/protocols/whir_zk/committer.rs | 12 ++- src/protocols/whir_zk/mod.rs | 25 ++++- 12 files changed, 245 insertions(+), 184 deletions(-) diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index a38a95a4..e282707a 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -57,25 +57,9 @@ pub trait BufferOps: Clone { fn dot(&self, other: &Self) -> F; } -fn interleaved_rs_encode_slices( - vectors: &[&[F]], - masks: &[F], - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, -) -> CpuMatrix { - let messages = vectors - .iter() - .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) - .collect::>(); - CpuMatrix::from_vec( - ntt::interleaved_rs_encode(&messages, masks, codeword_length), - codeword_length, - vectors.len() * interleaving_depth, - ) -} - pub trait MatrixBufferOps { + type Witness; + fn len(&self) -> usize; fn num_rows(&self) -> usize; fn num_cols(&self) -> usize; @@ -84,13 +68,25 @@ pub trait MatrixBufferOps { &self, config: &matrix_commit::Config, prover_state: &mut ProverState, - ) -> matrix_commit::Witness + ) -> Self::Witness where F: TypeInfo + matrix_commit::Encodable + Send + Sync, H: DuplexSpongeInterface, R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>; + fn open_rows( + &self, + config: &matrix_commit::Config, + prover_state: &mut ProverState, + witness: &Self::Witness, + indices: &[usize], + ) where + F: TypeInfo + matrix_commit::Encodable + Send + Sync, + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + Hash: ProverMessage<[H::U]>; + fn read_rows(&self, indices: &[usize]) -> Vec; } @@ -133,6 +129,8 @@ impl MatrixBufferOps for CpuMatrix where F: Field + TypeInfo + matrix_commit::Encodable + Send + Sync, { + type Witness = matrix_commit::Witness; + fn len(&self) -> usize { self.data.len() } @@ -149,7 +147,7 @@ where &self, config: &matrix_commit::Config, prover_state: &mut ProverState, - ) -> matrix_commit::Witness + ) -> Self::Witness where H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -160,6 +158,20 @@ where config.commit(prover_state, &self.data) } + fn open_rows( + &self, + config: &matrix_commit::Config, + prover_state: &mut ProverState, + witness: &Self::Witness, + indices: &[usize], + ) where + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + Hash: ProverMessage<[H::U]>, + { + config.open(prover_state, witness, indices); + } + fn read_rows(&self, indices: &[usize]) -> Vec { let mut rows = Vec::with_capacity(indices.len() * self.num_cols); for &index in indices { @@ -169,7 +181,18 @@ where rows } } -#[derive(Clone)] +#[derive( + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + Debug, + Default, + serde::Serialize, + serde::Deserialize, +)] pub struct CpuBuffer { data: Vec, } @@ -185,7 +208,7 @@ impl CpuBuffer { } } - fn as_slice(&self) -> &[F] { + pub(crate) fn as_slice(&self) -> &[F] { self.data.as_slice() } } @@ -263,73 +286,20 @@ impl BufferOps for CpuBuffer { } } -impl BufferOps for Vec { - type Buffer = Vec; - type Matrix = CpuMatrix; - - fn len(&self) -> usize { - self.len() - } - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution, - { - (0..length).map(|_| rng.gen()).collect() - } - - fn zero_pad(&mut self) { - if !self.is_empty() { - self.resize(self.len().next_power_of_two(), F::ZERO); - } - } - - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - mixed_multilinear_extend(embedding, self, point) - } - - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &Self::Buffer, - ) -> M::Target { - mixed_dot(embedding, other, self) - } - - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - mixed_univariate_evaluate(embedding, self, point) - } - - fn interleaved_rs_encode( - vectors: &[&Self], - masks: &Self, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - ) -> Self::Matrix - where - F: 'static, - { - let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); - interleaved_rs_encode_slices( - &vectors, - masks, - message_length, - interleaving_depth, - codeword_length, - ) - } - - fn dot(&self, other: &Self) -> F { - self.mixed_dot(&Identity::new(), other) - } +fn interleaved_rs_encode_slices( + vectors: &[&[F]], + masks: &[F], + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, +) -> CpuMatrix { + let messages = vectors + .iter() + .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) + .collect::>(); + CpuMatrix::from_vec( + ntt::interleaved_rs_encode(&messages, masks, codeword_length), + codeword_length, + vectors.len() * interleaving_depth, + ) } diff --git a/src/algebra/mod.rs b/src/algebra/mod.rs index 2c7e4838..4d17dc3b 100644 --- a/src/algebra/mod.rs +++ b/src/algebra/mod.rs @@ -1,3 +1,4 @@ +pub mod buffer; pub mod embedding; pub mod fields; pub mod linear_form; diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index d9e79ce6..b8d4eeab 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -10,6 +10,7 @@ use clap::Parser; use serde::Serialize; use whir::{ algebra::{ + buffer::CpuBuffer, embedding::{Basefield, Embedding, Identity}, fields::{Field128, Field192, Field256, Field64, Field64_2, Field64_3}, linear_form::{Evaluate, LinearForm, MultilinearExtension}, @@ -161,7 +162,8 @@ where HASH_COUNTER.reset(); - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( &mut prover_state, @@ -238,7 +240,8 @@ where HASH_COUNTER.reset(); let whir_prover_time = Instant::now(); - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let prove_linear_forms: Vec>> = points .iter() diff --git a/src/bin/main.rs b/src/bin/main.rs index bafe16f5..a69ba1a0 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -6,6 +6,7 @@ use ark_std::rand::distributions::{Distribution, Standard}; use clap::Parser; use whir::{ algebra::{ + buffer::CpuBuffer, embedding::{Basefield, Embedding, Identity}, fields::{Field128, Field192, Field256, Field64, Field64_2, Field64_3}, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, @@ -148,9 +149,10 @@ where } let vector = (0..num_coeffs).map(M::Source::from).collect::>(); + let vector_buffer = CpuBuffer::from_slice(&vector); let whir_commit_time = Instant::now(); - let witness = params.commit(&mut prover_state, &[&vector]); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let whir_commit_time = whir_commit_time.elapsed(); // Allocate constraints @@ -314,7 +316,8 @@ where } let whir_commit_time = Instant::now(); - let witness = params.commit(&mut prover_state, &[vector.as_slice()]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let whir_commit_time = whir_commit_time.elapsed(); let whir_prove_time = Instant::now(); diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 32763271..eac97802 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -11,8 +11,8 @@ use spongefish::{Decoding, VerificationResult}; use crate::{ algebra::{ - dot, embedding::Identity, multilinear_extend, random_vector, scalar_mul_add_new, - univariate_evaluate, + buffer::CpuBuffer, dot, embedding::Identity, multilinear_extend, random_vector, + scalar_mul_add_new, univariate_evaluate, }, hash::Hash, protocols::{irs_commit, sumcheck}, @@ -79,7 +79,7 @@ impl Config { // Even more trivial non-zk protocol: send f and r directly. if !self.masked { prover_state.prover_messages(&vector); - prover_state.prover_messages(&witness.masks); + prover_state.prover_messages(witness.masks.as_slice()); let _ = self.commit.open(prover_state, &[witness]); let point = self .sumcheck @@ -94,9 +94,10 @@ impl Config { // Create masking vector. let mask = random_vector(prover_state.rng(), vector.len()); + let mask_buffer = CpuBuffer::from_slice(&mask); // Commit to the masking vector. - let mask_witness = self.commit.commit(prover_state, &[&mask]); + let mask_witness = self.commit.commit(prover_state, &[&mask_buffer]); // Compute and send linear form of mask (μ' in paper). let mask_sum = dot(&mask, &covector); @@ -109,7 +110,11 @@ impl Config { prover_state.prover_messages(&masked_vector); // Send combined IRS randomness. (r^* in paper) - let masked_masks = scalar_mul_add_new(&mask_witness.masks, mask_rlc, &witness.masks); + let masked_masks = scalar_mul_add_new( + mask_witness.masks.as_slice(), + mask_rlc, + witness.masks.as_slice(), + ); prover_state.prover_messages(&masked_masks); // Open the commitment and mask simultaneously. @@ -284,7 +289,8 @@ mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit.commit(&mut prover_state, &[&vector]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let witness = config.commit.commit(&mut prover_state, &[&vector_buffer]); let prover_result = config.prove( &mut prover_state, vector.clone(), diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 62bbb602..9da26424 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -13,6 +13,7 @@ use tracing::instrument; use crate::{ algebra::{ + buffer::CpuBuffer, dot, embedding::{Embedding, Identity}, eq_weights, geometric_accumulate, lift, mixed_dot, scalar_mul, univariate_evaluate, @@ -195,7 +196,8 @@ impl Config { }; // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 - let target_witness = self.target.commit(prover_state, &[&message]); + let message_buffer = CpuBuffer::from_slice(&message); + let target_witness = self.target.commit(prover_state, &[&message_buffer]); // Step 2-3: OOD challenge + answers — Construction 9.7 Steps 2-3, p.55 // y := ze_ood(ρ) · [f; r; s] = f(α) + α^ℓ · (r,s)(α) @@ -484,7 +486,7 @@ mod tests { } // Lift ι parallel masks (total length source.mask_length × ι) and fold // chunks of length source.mask_length down to a single chunk. - let raw = lift(config.source.embedding(), &source_witness.masks); + let raw = lift(config.source.embedding(), source_witness.masks.as_slice()); let mut mask = fold_chunks(&raw, config.source.mask_length, folding_randomness); // Append fresh padding s of length message_mask_length - source.mask_length. mask.extend(random_vector::( @@ -521,7 +523,8 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let f_full_buffer = CpuBuffer::from_slice(&f_full); + let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); // Sample γ for sumcheck folding (length log2(ι)). let folding_randomness = sample_folding_randomness(config, &mut rng); @@ -574,7 +577,8 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let f_full_buffer = CpuBuffer::from_slice(&f_full); + let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); let folded_message = @@ -642,7 +646,8 @@ mod tests { // Commit honest f_full, fold to get the honest post-fold message. let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let f_full_buffer = CpuBuffer::from_slice(&f_full); + let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); let folded_message = fold_chunks(&f_full, config.source.message_length(), &folding_randomness); diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index b5716f3a..1d9e20a8 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -20,6 +20,7 @@ use std::{ f64::{self, consts::LOG2_10}, fmt, + marker::PhantomData, ops::Neg, }; @@ -32,8 +33,13 @@ use tracing::instrument; use crate::{ algebra::{ - dot, embedding::Embedding, fields::FieldWithSize, lift, linear_form::UnivariateEvaluation, - mixed_univariate_evaluate, ntt, random_vector, + buffer::{BufferOps, CpuBuffer, CpuMatrix, MatrixBufferOps}, + dot, + embedding::Embedding, + fields::FieldWithSize, + lift, + linear_form::UnivariateEvaluation, + ntt, }, engines::EngineId, hash::Hash, @@ -43,7 +49,7 @@ use crate::{ VerifierMessage, VerifierState, }, type_info::Typed, - utils::{chunks_exact_or_empty, zip_strict}, + utils::zip_strict, verify, }; @@ -92,14 +98,17 @@ pub struct Config { #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] #[must_use] -pub struct Witness +pub struct Witness, Matrix = CpuMatrix> where G: Field, + Masks: BufferOps, + Matrix: MatrixBufferOps, { - pub masks: Vec, - pub matrix: Vec, - pub matrix_witness: matrix_commit::Witness, + pub masks: Masks, + pub matrix: Matrix, + pub matrix_witness: Matrix::Witness, pub out_of_domain: Evaluations, + source_field: PhantomData, } #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] @@ -289,12 +298,15 @@ impl Config { /// Commit to one or more vectors. #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] - pub fn commit( + pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&[M::Source]], - ) -> Witness + vectors: &[&B], + ) -> Witness where + B: BufferOps, + >::Witness: + Clone + PartialEq + Eq + PartialOrd + Ord + fmt::Debug + std::hash::Hash + Default, Standard: Distribution, H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -310,18 +322,17 @@ impl Config { assert_eq!(vectors.len(), self.num_vectors); assert!(vectors.iter().all(|p| p.len() == self.vector_size)); - // Generate random mask - let masks = random_vector(prover_state.rng(), self.mask_length * self.num_messages()); - - // Interleaved RS Encode the vectors - let messages = vectors - .iter() - .flat_map(|v| chunks_exact_or_empty(v, self.message_length(), self.interleaving_depth)) - .collect::>(); - let matrix = ntt::interleaved_rs_encode(&messages, &masks, self.codeword_length); + let masks = B::random(prover_state.rng(), self.mask_length * self.num_messages()); + let matrix = B::interleaved_rs_encode( + vectors, + &masks, + self.message_length(), + self.interleaving_depth, + self.codeword_length, + ); // Commit to the matrix - let matrix_witness = self.matrix_commit.commit(prover_state, &matrix); + let matrix_witness = matrix.commit_rows(&self.matrix_commit, prover_state); // Handle out-of-domain points and values // TODO : Remove this logic after main whir protocol is updated @@ -332,7 +343,7 @@ impl Config { let mut oods_matrix = Vec::with_capacity(self.out_domain_samples * self.num_vectors); for &point in &oods_points { for &vector in vectors { - let value = mixed_univariate_evaluate(&*self.embedding, vector, point); + let value = vector.mixed_univariate_evaluate(&*self.embedding, point); prover_state.prover_message(&value); oods_matrix.push(value); } @@ -346,6 +357,7 @@ impl Config { points: oods_points, matrix: oods_matrix, }, + source_field: PhantomData, } } @@ -382,12 +394,14 @@ impl Config { /// When there are multiple openings, the evaluation matrices will /// be horizontally concatenated. #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] - pub fn open( + pub fn open( &self, prover_state: &mut ProverState, - witnesses: &[&Witness], + witnesses: &[&Witness], ) -> Evaluations where + Masks: BufferOps, + Matrix: MatrixBufferOps, H: DuplexSpongeInterface, R: RngCore + CryptoRng, u8: Decoding<[H::U]>, @@ -409,22 +423,23 @@ impl Config { // and collect them in the evaluation matrix. let stride = witnesses.len() * self.num_cols(); let mut matrix = vec![M::Source::ZERO; indices.len() * stride]; - let mut submatrix = Vec::with_capacity(indices.len() * self.num_cols()); let mut matrix_col_offset = 0; for witness in witnesses { - submatrix.clear(); - for (point_index, &code_index) in indices.iter().enumerate() { - let row = &witness.matrix - [code_index * self.num_cols()..(code_index + 1) * self.num_cols()]; - submatrix.extend_from_slice(row); - - let matrix_row = &mut matrix[point_index * stride..(point_index + 1) * stride]; - matrix_row[matrix_col_offset..matrix_col_offset + self.num_cols()] - .copy_from_slice(row); + let submatrix = witness.matrix.read_rows(&indices); + if self.num_cols() != 0 { + for (point_index, row) in submatrix.chunks_exact(self.num_cols()).enumerate() { + let matrix_row = &mut matrix[point_index * stride..(point_index + 1) * stride]; + matrix_row[matrix_col_offset..matrix_col_offset + self.num_cols()] + .copy_from_slice(row); + } } prover_state.prover_hint_ark(&submatrix); - self.matrix_commit - .open(prover_state, &witness.matrix_witness, &indices); + witness.matrix.open_rows( + &self.matrix_commit, + prover_state, + &witness.matrix_witness, + &indices, + ); matrix_col_offset += self.num_cols(); } @@ -511,7 +526,13 @@ impl Commitment { } } -impl Witness { +impl Witness +where + F: Field, + G: Field, + Masks: BufferOps, + Matrix: MatrixBufferOps, +{ /// Returns the out-of-domain evaluations. pub const fn out_of_domain(&self) -> &Evaluations { &self.out_of_domain @@ -642,7 +663,7 @@ pub(crate) mod tests { use crate::{ algebra::{ embedding::{Basefield, Compose, Frobenius, Identity}, - fields, + fields, mixed_univariate_evaluate, ntt::NTT, random_vector, univariate_evaluate, }, @@ -728,10 +749,12 @@ pub(crate) mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit( - &mut prover_state, - &vectors.iter().map(|p| p.as_slice()).collect::>(), - ); + let vector_buffers = vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = config.commit(&mut prover_state, &vector_refs); assert_eq!( witness.out_of_domain().points.len(), config.out_domain_samples diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index 1d632ca0..deb49de5 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -45,7 +45,10 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, R use serde::{Deserialize, Serialize}; use crate::{ - algebra::{embedding::Identity, random_vector, scalar_mul_add_new, univariate_evaluate}, + algebra::{ + buffer::CpuBuffer, embedding::Identity, random_vector, scalar_mul_add_new, + univariate_evaluate, + }, hash::Hash, protocols::irs_commit::{ Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness, @@ -129,11 +132,19 @@ impl Config { .map(|_| random_vector(prover_state.rng(), self.c_zk_commit.vector_size)) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| CpuBuffer::from_slice(msg)) + .collect::>(); + let fresh_buffers = fresh_msgs + .iter() + .map(|msg| CpuBuffer::from_slice(msg)) + .collect::>(); + // Tree layout: [originals..., freshes...] - let all_vectors: Vec<&[F]> = original_msgs + let all_vectors: Vec<&CpuBuffer> = original_buffers .iter() - .chain(fresh_msgs.iter()) - .map(|v| v.as_slice()) + .chain(fresh_buffers.iter()) .collect(); let mask_witness = self.c_zk_commit.commit(prover_state, &all_vectors); @@ -180,10 +191,8 @@ impl Config { // Step 2: compute and send combined polynomials + IRS randomness let irs_masks_per_vector = self.c_zk_commit.mask_length * self.c_zk_commit.interleaving_depth; - assert_eq!( - witness.mask_witness.masks.len(), - 2 * self.num_masks * irs_masks_per_vector - ); + let irs_masks = witness.mask_witness.masks.as_slice(); + assert_eq!(irs_masks.len(), 2 * self.num_masks * irs_masks_per_vector); for (i, (orig_msg, fresh_msg)) in original_msgs .iter() .zip(witness.fresh_msgs.iter()) @@ -195,10 +204,8 @@ impl Config { // r*_i = r'_i + γ · r_i if irs_masks_per_vector > 0 { - let orig_r = &witness.mask_witness.masks - [i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; - let fresh_r = &witness.mask_witness.masks[(self.num_masks + i) - * irs_masks_per_vector + let orig_r = &irs_masks[i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; + let fresh_r = &irs_masks[(self.num_masks + i) * irs_masks_per_vector ..(self.num_masks + i + 1) * irs_masks_per_vector]; let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); prover_state.prover_messages(&combined_r); @@ -455,6 +462,7 @@ mod tests { let gamma: F = prover_state.verifier_message(); let irs_masks_per_vector = config.c_zk_commit.mask_length * config.c_zk_commit.interleaving_depth; + let irs_masks = witness.mask_witness.masks.as_slice(); for (i, (orig_msg, fresh_msg)) in original_msgs .iter() @@ -468,10 +476,8 @@ mod tests { prover_state.prover_messages(&combined_msg); if irs_masks_per_vector > 0 { - let orig_r = &witness.mask_witness.masks - [i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; - let fresh_r = &witness.mask_witness.masks[(config.num_masks + i) - * irs_masks_per_vector + let orig_r = &irs_masks[i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; + let fresh_r = &irs_masks[(config.num_masks + i) * irs_masks_per_vector ..(config.num_masks + i + 1) * irs_masks_per_vector]; let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); prover_state.prover_messages(&combined_r); diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 805a206e..9c6a31bb 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -14,6 +14,7 @@ use tracing::instrument; use crate::{ algebra::{ + buffer::CpuBuffer, embedding::{Embedding, Identity}, linear_form::LinearForm, }, @@ -79,7 +80,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&[M::Source]], + vectors: &[&CpuBuffer], ) -> Witness where Standard: Distribution, @@ -128,6 +129,7 @@ mod tests { use super::*; use crate::{ algebra::{ + buffer::CpuBuffer, embedding::Basefield, fields::{Field64, Field64_3}, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, @@ -239,7 +241,8 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit to the polynomial and generate auxiliary witness data - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let prove_linear_forms = build_prove_forms(&points, num_variables, true); @@ -380,7 +383,7 @@ mod tests { vec![F::from((i + 1) as u64); num_coeffs] }) .collect(); - let vec_refs = vectors.iter().map(|v| v.as_slice()).collect::>(); + let vec_refs = vectors.iter().collect::>(); let points: Vec<_> = (0..num_points_per_poly) .map(|_| random_vector(thread_rng(), num_variables)) @@ -401,7 +404,7 @@ mod tests { .flat_map(|linear_form| { vec_refs .iter() - .map(|vec| linear_form.evaluate(params.embedding(), vec)) + .map(|&vec| linear_form.evaluate(params.embedding(), vec)) }) .collect::>(); @@ -412,8 +415,12 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit to each polynomial and generate witnesses + let vector_buffers = vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); let mut witnesses = Vec::new(); - for &vec in &vec_refs { + for vec in &vector_buffers { let witness = params.commit(&mut prover_state, &[vec]); witnesses.push(witness); } @@ -569,8 +576,10 @@ mod tests { .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness1 = params.commit(&mut prover_state, &[&vec1]); - let witness2 = params.commit(&mut prover_state, &[&vec2]); + let vec1_buffer = CpuBuffer::from_slice(&vec1); + let vec2_buffer = CpuBuffer::from_slice(&vec2); + let witness1 = params.commit(&mut prover_state, &[&vec1_buffer]); + let witness2 = params.commit(&mut prover_state, &[&vec2_buffer]); let prove_linear_forms = build_prove_forms(&constraint_points, num_variables, false); @@ -651,7 +660,7 @@ mod tests { let all_vectors: Vec> = (0..num_witnesses * batch_size) .map(|i| vec![F::from((i + 1) as u64); num_coeffs]) .collect::>(); - let vec_refs = all_vectors.iter().map(|p| p.as_slice()).collect::>(); + let vec_refs = all_vectors.iter().collect::>(); let points: Vec<_> = (0..num_points_per_poly) .map(|_| random_vector(thread_rng(), num_variables)) @@ -672,7 +681,7 @@ mod tests { .flat_map(|linear_form| { vec_refs .iter() - .map(|vec| linear_form.evaluate(params.embedding(), vec)) + .map(|&vec| linear_form.evaluate(params.embedding(), vec)) }) .collect::>(); @@ -683,8 +692,13 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit using commit_batch (stacks batch_size polynomials per witness) + let vector_buffers = all_vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); + let buffer_refs = vector_buffers.iter().collect::>(); let mut witnesses = Vec::new(); - for witness_polys in vec_refs.chunks(batch_size) { + for witness_polys in buffer_refs.chunks(batch_size) { let witness = params.commit(&mut prover_state, witness_polys); witnesses.push(witness); } @@ -793,7 +807,7 @@ mod tests { let vectors: Vec> = (0..batch_size) .map(|_| random_vector(thread_rng(), num_coeffs)) .collect(); - let vec_refs = vectors.iter().map(|v| v.as_slice()).collect::>(); + let vec_refs = vectors.iter().collect::>(); // Generate `num_points` random points in the multilinear domain let points: Vec<_> = (0..num_points) @@ -809,7 +823,12 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Create a commitment to the polynomial and generate auxiliary witness data - let batched_witness = params.commit(&mut prover_state, &vec_refs); + let vector_buffers = vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); + let buffer_refs = vector_buffers.iter().collect::>(); + let batched_witness = params.commit(&mut prover_state, &buffer_refs); // Create a weights matrix and evaluations for each polynomial let mut linear_forms: Vec>>> = Vec::new(); diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 04b4473e..e3fc2043 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -8,6 +8,7 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ algebra::{ + buffer::CpuBuffer, dot, embedding::Embedding, eq_weights, lift, @@ -220,7 +221,10 @@ impl Config { // Execute standard WHIR rounds on the batched vectors for (round_index, round_config) in self.round_configs.iter().enumerate() { // Commit to the vector, this generates out-of-domain evaluations. - let new_witness = round_config.irs_committer.commit(prover_state, &[&vector]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let new_witness = round_config + .irs_committer + .commit(prover_state, &[&vector_buffer]); // Proof of work before in-domain challenges round_config.pow.prove(prover_state); diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index 942b87e1..01b7a863 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -5,6 +5,7 @@ use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; use crate::{ + algebra::buffer::CpuBuffer, hash::Hash, protocols::{irs_commit, whir}, transcript::{ @@ -43,7 +44,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - polynomials: &[&[F]], + polynomials: &[&CpuBuffer], ) -> Witness where Standard: Distribution, @@ -63,6 +64,7 @@ impl Config { let num_blinding_variables = self.num_blinding_variables(); let num_witness_variables = self.num_witness_variables(); for &poly in polynomials { + let poly = poly.as_slice(); let blinding = BlindingPolynomials::sample( prover_state.rng(), num_blinding_variables, @@ -86,9 +88,10 @@ impl Config { .zip(mask.iter().cycle()) .map(|(&coeff, &m)| coeff + m) .collect::>(); + let f_hat_buffer = CpuBuffer::from_slice(&f_hat_vec); let witness = self .blinded_commitment - .commit(prover_state, &[f_hat_vec.as_slice()]); + .commit(prover_state, &[&f_hat_buffer]); f_hat_vectors.push(f_hat_vec); f_hat_witnesses.push(witness); blinding_polynomials.push(blinding); @@ -110,10 +113,11 @@ impl Config { ); blinding_vectors.extend(layout); } - let blinding_vector_refs = blinding_vectors + let blinding_buffers = blinding_vectors .iter() - .map(Vec::as_slice) + .map(|v| CpuBuffer::from_slice(v)) .collect::>(); + let blinding_vector_refs = blinding_buffers.iter().collect::>(); let blinding_witness = self .blinding_commitment .commit(prover_state, &blinding_vector_refs); diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index b5e41c6e..c27ca37b 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -249,6 +249,7 @@ mod tests { use super::*; use crate::{ algebra::{ + buffer::CpuBuffer, fields::Field64, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, random_vector, @@ -350,7 +351,12 @@ mod tests { .session(&tag) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, vectors); + let vector_buffers = vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, vectors @@ -442,7 +448,12 @@ mod tests { .session(&format!("zk-stage1-negative {}:{}", file!(), line!())) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &vectors); + let vector_buffers = vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, vectors @@ -496,7 +507,12 @@ mod tests { .session(&format!("zk-stage1-tamper {}:{}", file!(), line!())) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &vectors); + let vector_buffers = vectors + .iter() + .map(|v| CpuBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, vectors @@ -557,7 +573,8 @@ mod tests { let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = CpuBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( &mut prover_state, vec![Cow::Borrowed(&vector)], From ee1b207b7e0301f9a2a03048abeef1d78db690d9 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 3 Jun 2026 17:13:47 +0200 Subject: [PATCH 05/35] wip: integrate buffer abstraction in prove and sumcheck --- src/algebra/buffer.rs | 139 +++++++++++++++++++++++++++++++++- src/bin/benchmark.rs | 8 +- src/bin/main.rs | 8 +- src/protocols/basecase.rs | 33 +++++--- src/protocols/sumcheck.rs | 61 ++++++++------- src/protocols/whir/mod.rs | 54 +++++++------ src/protocols/whir/prover.rs | 143 +++++++++++++++++------------------ 7 files changed, 300 insertions(+), 146 deletions(-) diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index e282707a..c6c8fb2b 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -1,3 +1,5 @@ +use std::{any::Any, mem}; + use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; use spongefish::DuplexSpongeInterface; @@ -5,7 +7,9 @@ use spongefish::DuplexSpongeInterface; use crate::{ algebra::{ embedding::{Embedding, Identity}, - mixed_dot, mixed_multilinear_extend, mixed_univariate_evaluate, ntt, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, + mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, ntt, + sumcheck::{compute_sumcheck_polynomial, fold, fold_and_compute_polynomial}, }, hash::Hash, protocols::matrix_commit, @@ -18,6 +22,8 @@ pub trait BufferOps: Clone { type Buffer: BufferOps; type Matrix: MatrixBufferOps; + fn from_vec(source: Vec) -> Self; + fn zeros(length: usize) -> Self; fn len(&self) -> usize; fn is_empty(&self) -> bool { @@ -29,7 +35,26 @@ pub trait BufferOps: Clone { R: RngCore + CryptoRng, Standard: Distribution; + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self; + fn zero_pad(&mut self); + fn fold(&mut self, weight: F); + fn sumcheck_polynomial(&self, other: &Self) -> (F, F); + fn fold_and_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F); + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ); + fn write_to_prover(&self, prover_state: &mut ProverState) + where + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + F: ProverMessage<[H::U]>; fn mixed_extend, T: Field>( &self, embedding: &M, @@ -45,6 +70,22 @@ pub trait BufferOps: Clone { embedding: &M, point: M::Target, ) -> M::Target; + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::Buffer; + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut Self::Buffer, + weight: M::Target, + ); + fn mixed_dot_slice>( + &self, + embedding: &M, + other: &[M::Target], + ) -> M::Target; fn interleaved_rs_encode( vectors: &[&Self], masks: &Self, @@ -217,6 +258,16 @@ impl BufferOps for CpuBuffer { type Buffer = CpuBuffer; type Matrix = CpuMatrix; + fn from_vec(source: Vec) -> Self { + Self::from_vec(source) + } + + fn zeros(length: usize) -> Self { + Self { + data: vec![F::ZERO; length], + } + } + fn len(&self) -> usize { self.data.len() } @@ -231,12 +282,64 @@ impl BufferOps for CpuBuffer { } } + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self { + assert_eq!(linear_forms.len(), rlc_coeffs.len()); + let mut covector = vec![F::ZERO; size]; + if let Some((first, linear_forms)) = linear_forms.split_first_mut() { + debug_assert_eq!(rlc_coeffs[0], F::ONE); + if let Some(covector_form) = + (first.as_mut() as &mut dyn Any).downcast_mut::>() + { + mem::swap(&mut covector, &mut covector_form.vector); + } else { + first.accumulate(&mut covector, F::ONE); + } + for (rlc_coeff, linear_form) in rlc_coeffs[1..].iter().zip(linear_forms) { + linear_form.accumulate(&mut covector, *rlc_coeff); + } + } + Self { data: covector } + } + fn zero_pad(&mut self) { if !self.is_empty() { self.data.resize(self.len().next_power_of_two(), F::ZERO); } } + fn fold(&mut self, weight: F) { + fold(&mut self.data, weight); + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + compute_sumcheck_polynomial(&self.data, other.as_slice()) + } + + fn fold_and_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + fold_and_compute_polynomial(&mut self.data, &mut other.data, weight) + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + UnivariateEvaluation::accumulate_many(evaluators, &mut self.data, scalars); + } + + fn write_to_prover(&self, prover_state: &mut ProverState) + where + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + F: ProverMessage<[H::U]>, + { + prover_state.prover_messages(&self.data); + } + fn mixed_extend, T: Field>( &self, embedding: &M, @@ -265,6 +368,40 @@ impl BufferOps for CpuBuffer { mixed_univariate_evaluate(embedding, &self.data, point) } + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::Buffer { + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return CpuBuffer::from_vec(Vec::new()); + }; + debug_assert_eq!(coeffs[0], M::Target::ONE); + let mut accumulator = crate::algebra::lift(embedding, first.as_slice()); + for (coeff, vector) in coeffs[1..].iter().zip(vectors) { + mixed_scalar_mul_add(embedding, &mut accumulator, *coeff, vector.as_slice()); + } + CpuBuffer::from_vec(accumulator) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut Self::Buffer, + weight: M::Target, + ) { + mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.as_slice()); + } + + fn mixed_dot_slice>( + &self, + embedding: &M, + other: &[M::Target], + ) -> M::Target { + mixed_dot(embedding, other, self.as_slice()) + } + fn interleaved_rs_encode( vectors: &[&Self], masks: &Self, diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index b8d4eeab..42e654b0 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -167,8 +167,8 @@ where let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vector.as_slice())], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], vec![], Cow::Owned(vec![]), ); @@ -252,8 +252,8 @@ where let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vector.as_slice())], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); diff --git a/src/bin/main.rs b/src/bin/main.rs index a69ba1a0..c63ada92 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -185,8 +185,8 @@ where let whir_prove_time = Instant::now(); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vector.as_slice())], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -323,8 +323,8 @@ where let whir_prove_time = Instant::now(); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(&vector)], - witness, + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(&evaluations), ); diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index eac97802..e6faa026 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -49,9 +49,9 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - mut vector: Vec, + vector: Vec, witness: &irs_commit::Witness, - mut covector: Vec, + covector: Vec, mut sum: F, ) -> Opening where @@ -81,14 +81,22 @@ impl Config { prover_state.prover_messages(&vector); prover_state.prover_messages(witness.masks.as_slice()); let _ = self.commit.open(prover_state, &[witness]); + let mut vector_buffer = CpuBuffer::from_vec(vector); + let mut covector_buffer = CpuBuffer::from_vec(covector); let point = self .sumcheck - .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) + .prove( + prover_state, + &mut vector_buffer, + &mut covector_buffer, + &mut sum, + &[], + ) .round_challenges; - assert!(!vector[0].is_zero(), "Proof failed"); + assert!(!vector_buffer.as_slice()[0].is_zero(), "Proof failed"); return Opening { evaluation_points: point, - linear_form_evaluation: covector[0], + linear_form_evaluation: covector_buffer.as_slice()[0], }; } @@ -106,7 +114,7 @@ impl Config { // RLC the mask with the vector let mask_rlc = prover_state.verifier_message::(); assert!(!mask_rlc.is_zero(), "Proof failed"); - let mut masked_vector = scalar_mul_add_new(&mask, mask_rlc, &vector); + let masked_vector = scalar_mul_add_new(&mask, mask_rlc, &vector); prover_state.prover_messages(&masked_vector); // Send combined IRS randomness. (r^* in paper) @@ -122,12 +130,14 @@ impl Config { // Run sumcheck to reduce linear form claim let mut masked_sum = mask_sum + mask_rlc * sum; + let mut masked_vector_buffer = CpuBuffer::from_vec(masked_vector); + let mut covector_buffer = CpuBuffer::from_vec(covector); let point = self .sumcheck .prove( prover_state, - &mut masked_vector, - &mut covector, + &mut masked_vector_buffer, + &mut covector_buffer, &mut masked_sum, &[], ) @@ -137,11 +147,14 @@ impl Config { // Basically the sumcheck equation has degenerated to 0 * l(r) = 0, which provides // no constraints on l(r) that the verifier can return. // This event is cryptographically unlikely as `F` is challenge sized. - assert!(!masked_vector[0].is_zero(), "Proof failed"); + assert!( + !masked_vector_buffer.as_slice()[0].is_zero(), + "Proof failed" + ); Opening { evaluation_points: point, - linear_form_evaluation: covector[0], + linear_form_evaluation: covector_buffer.as_slice()[0], } } diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 5018e22d..13c0c255 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -9,11 +9,7 @@ use serde::{Deserialize, Serialize}; use tracing::instrument; use crate::{ - algebra::{ - dot, - sumcheck::{compute_sumcheck_polynomial, fold, fold_and_compute_polynomial}, - univariate_evaluate, - }, + algebra::{buffer::BufferOps, univariate_evaluate}, protocols::proof_of_work, transcript::{ codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverState, VerificationResult, @@ -66,15 +62,16 @@ impl Config { /// - Applies proof-of-work grinding if required. /// - Returns the sampled folding randomness values used in each reduction step. #[cfg_attr(feature = "tracing", instrument(skip_all))] - pub fn prove( + pub fn prove( &self, prover_state: &mut ProverState, - a: &mut Vec, - b: &mut Vec, + a: &mut B, + b: &mut B, sum: &mut F, masks: &[F], ) -> SumcheckOpening where + B: BufferOps, H: DuplexSpongeInterface, R: CryptoRng + RngCore, F: Codec<[H::U]>, @@ -88,7 +85,7 @@ impl Config { assert!(self.mask_length == 0 || self.mask_length >= 3); assert_eq!(a.len(), self.initial_size); assert_eq!(b.len(), self.initial_size); - debug_assert_eq!(dot(a, b), *sum); + debug_assert_eq!(a.dot(b), *sum); assert_eq!(masks.len(), self.num_rounds * self.mask_length); let half = F::from(2).inverse().unwrap(); @@ -110,9 +107,9 @@ impl Config { { // Fold and compute sumcheck polynomial in one pass. let (c0, c2) = if let Some(w) = prev_round_challenge { - fold_and_compute_polynomial(a, b, w) + a.fold_and_sumcheck_polynomial(b, w) } else { - compute_sumcheck_polynomial(a, b) + a.sumcheck_polynomial(b) }; let c1 = *sum - c0.double() - c2; @@ -150,8 +147,8 @@ impl Config { } if let Some(w) = prev_round_challenge { // Final fold of the inputs (no polynomial computation) - fold(a, w); - fold(b, w); + a.fold(w); + b.fold(w); } *sum = mask_sum + mask_rlc * *sum; @@ -238,23 +235,24 @@ fn eval_01(coefficients: &[F]) -> F { #[cfg(test)] mod tests { - use ark_std::rand::{ - distributions::{Distribution, Standard}, - rngs::StdRng, - SeedableRng, - }; - use proptest::{prelude::Just, prop_oneof, proptest, strategy::Strategy}; - #[cfg(feature = "tracing")] - use tracing::instrument; - use super::*; use crate::{ algebra::{ + buffer::CpuBuffer, + dot, fields::{self, Field64}, multilinear_extend, random_vector, }, transcript::DomainSeparator, }; + use ark_std::rand::{ + distributions::{Distribution, Standard}, + rngs::StdRng, + SeedableRng, + }; + use proptest::{prelude::Just, prop_oneof, proptest, strategy::Strategy}; + #[cfg(feature = "tracing")] + use tracing::instrument; impl Config where @@ -299,8 +297,8 @@ mod tests { let masks = random_vector(&mut rng, config.mask_length * config.num_rounds); // Prover - let mut vector = initial_vector.clone(); - let mut covector = initial_covector.clone(); + let mut vector = CpuBuffer::from_slice(&initial_vector); + let mut covector = CpuBuffer::from_slice(&initial_covector); let mut sum = initial_sum; let mut prover_state = ProverState::new_std(&ds); let SumcheckOpening { @@ -316,8 +314,14 @@ mod tests { assert_eq!(vector.len(), config.final_size()); assert_eq!(covector.len(), config.final_size()); if config.final_size() == 1 { - assert_eq!(multilinear_extend(&initial_vector, &point), vector[0]); - assert_eq!(multilinear_extend(&initial_covector, &point), covector[0]); + assert_eq!( + multilinear_extend(&initial_vector, &point), + vector.as_slice()[0] + ); + assert_eq!( + multilinear_extend(&initial_covector, &point), + covector.as_slice()[0] + ); } else { // TODO: Check correct folding. } @@ -327,7 +331,10 @@ mod tests { .zip(&point) .map(|(m, x)| univariate_evaluate(m, *x)) .sum(); - assert_eq!(sum, expected_mask_sum + mask_rlc * dot(&vector, &covector)); + assert_eq!( + sum, + expected_mask_sum + mask_rlc * dot(vector.as_slice(), covector.as_slice()) + ); let proof = prover_state.proof(); diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 9c6a31bb..0f20a14b 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -14,7 +14,7 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::CpuBuffer, + buffer::{BufferOps, CpuBuffer}, embedding::{Embedding, Identity}, linear_form::LinearForm, }, @@ -46,7 +46,13 @@ pub struct RoundConfig { pub pow: proof_of_work::Config, } -pub type Witness> = irs_commit::Witness; +pub type Witness, B = CpuBuffer<::Source>> = + irs_commit::Witness< + ::Source, + F, + B, + ::Source>>::Matrix, + >; pub type Commitment = irs_commit::Commitment; #[must_use = "The final claim must be checked if there where any linear forms."] @@ -76,13 +82,19 @@ impl FinalClaim { impl Config { /// Commit to one or more vectors. - #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().unwrap().len())))] - pub fn commit( + #[cfg_attr( + feature = "tracing", + instrument(skip_all, fields(size = vectors.first().unwrap().len())) + )] + pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&CpuBuffer], - ) -> Witness + vectors: &[&B], + ) -> Witness where + B: BufferOps, + >::Witness: + Clone + PartialEq + Eq + PartialOrd + Ord + Debug + std::hash::Hash + Default, Standard: Distribution, H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -249,8 +261,8 @@ mod tests { // Generate a proof for the given statement and witness let _ = params.prove( &mut prover_state, - vec![Cow::from(vector)], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -430,11 +442,8 @@ mod tests { // Batch prove all polynomials together let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|v| Cow::Borrowed(v.as_slice())) - .collect(), - witnesses.into_iter().map(Cow::Owned).collect(), + &vector_buffers.iter().collect::>(), + witnesses.iter().collect(), prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -584,10 +593,11 @@ mod tests { let prove_linear_forms = build_prove_forms(&constraint_points, num_variables, false); // Generate proof with mismatched polynomials + let vec_wrong_buffer = CpuBuffer::from_vec(vec_wrong); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vec1.as_slice()), Cow::from(vec_wrong)], - vec![Cow::Owned(witness1), Cow::Owned(witness2)], + &[&vec1_buffer, &vec_wrong_buffer], + vec![&witness1, &witness2], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -708,11 +718,8 @@ mod tests { // Batch prove all witnesses together let _ = params.prove( &mut prover_state, - all_vectors - .iter() - .map(|v| Cow::Borrowed(v.as_slice())) - .collect(), - witnesses.into_iter().map(Cow::Owned).collect(), + &vector_buffers.iter().collect::>(), + witnesses.iter().collect(), prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -858,11 +865,8 @@ mod tests { .collect::>(); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|v| Cow::Borrowed(v.as_slice())) - .collect(), - vec![Cow::Owned(batched_witness)], + &buffer_refs, + vec![&batched_witness], prove_linear_forms, Cow::Borrowed(values.as_slice()), ); diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index e3fc2043..01c87f3c 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -1,4 +1,4 @@ -use std::{any::Any, borrow::Cow, mem}; +use std::{borrow::Cow, fmt::Debug}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; @@ -8,13 +8,11 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ algebra::{ - buffer::CpuBuffer, + buffer::{BufferOps, MatrixBufferOps}, dot, embedding::Embedding, - eq_weights, lift, - linear_form::{Covector, Evaluate, LinearForm, UnivariateEvaluation}, - mixed_scalar_mul_add, - sumcheck::fold, + eq_weights, + linear_form::LinearForm, tensor_product, }, hash::Hash, @@ -26,9 +24,15 @@ use crate::{ utils::zip_strict, }; -enum RoundWitness<'a, F: Field, M: Embedding> { - Initial(Vec>>), - Round(irs_commit::Witness), +enum RoundWitness<'a, F, M, B> +where + F: Field, + M: Embedding, + B: BufferOps, + B::Buffer: BufferOps, +{ + Initial(Vec<&'a irs_commit::Witness>), + Round(irs_commit::Witness, as BufferOps>::Matrix>), } impl Config { @@ -48,15 +52,29 @@ impl Config { /// #[cfg_attr(feature = "tracing", instrument(skip_all))] #[allow(clippy::too_many_lines, clippy::cognitive_complexity)] - pub fn prove<'a, H, R>( + pub fn prove<'a, H, R, B>( &self, prover_state: &mut ProverState, - vectors: Vec>, - witnesses: Vec>>, + vectors: &[&B], + witnesses: Vec<&'a Witness>, linear_forms: Vec>>, evaluations: Cow<'a, [M::Target]>, ) -> FinalClaim where + B: BufferOps, + B::Buffer: BufferOps, + >::Witness: + Clone + PartialEq + Eq + PartialOrd + Ord + Debug + std::hash::Hash + Default, + < as BufferOps>::Matrix as MatrixBufferOps< + M::Target, + >>::Witness: Clone + + PartialEq + + Eq + + PartialOrd + + Ord + + Debug + + std::hash::Hash + + Default, Standard: Distribution + Distribution, H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -74,7 +92,7 @@ impl Config { witnesses.len() * self.initial_committer.num_vectors ); assert_eq!(evaluations.len(), num_vectors * linear_forms.len()); - for vector in &vectors { + for vector in vectors { assert_eq!(vector.len(), self.initial_size()); } for linear_form in &linear_forms { @@ -86,8 +104,11 @@ impl Config { { use crate::algebra::linear_form::Covector; let covector = Covector::from(&**linear_form); - for (vector, evaluation) in zip_strict(&vectors, evaluations) { - debug_assert_eq!(covector.evaluate(self.embedding(), vector), *evaluation); + for (vector, evaluation) in zip_strict(vectors, evaluations) { + debug_assert_eq!( + vector.mixed_dot_slice(self.embedding(), &covector.vector), + *evaluation + ); } } if vectors.is_empty() { @@ -111,12 +132,13 @@ impl Config { if j >= vector_offset && j < oods_row.len() + vector_offset { debug_assert_eq!( oods_row[j - vector_offset], - oods_eval.evaluate(self.embedding(), vector) + vector.mixed_univariate_evaluate(self.embedding(), oods_eval.point) ); oods_matrix.push(oods_row[j - vector_offset]); } else { - let eval = oods_eval.evaluate(self.embedding(), vector); + let eval = + vector.mixed_univariate_evaluate(self.embedding(), oods_eval.point); prover_state.prover_message(&eval); oods_matrix.push(eval); } @@ -131,18 +153,9 @@ impl Config { // Random linear combination of the vectors. let mut vector_rlc_coeffs: Vec = geometric_challenge(prover_state, num_vectors); assert_eq!(vector_rlc_coeffs[0], M::Target::ONE); - // Recycle the first input as the accumulator (its coefficient is always ONE). - let mut vectors = vectors.into_iter(); - let first = vectors.next().expect("non-empty"); - let mut vector = match first { - Cow::Borrowed(slice) => lift(self.embedding(), slice), - Cow::Owned(vec) => self.embedding().map_vec(vec), - }; - for (rlc_coeff, input_vector) in zip_strict(&vector_rlc_coeffs[1..], vectors) { - mixed_scalar_mul_add(self.embedding(), &mut vector, *rlc_coeff, &input_vector); - } + let mut vector = B::mixed_linear_combination(self.embedding(), vectors, &vector_rlc_coeffs); - let mut prev_witness: RoundWitness<'a, M::Target, M> = RoundWitness::Initial(witnesses); + let mut prev_witness: RoundWitness<'a, M::Target, M, B> = RoundWitness::Initial(witnesses); // Random linear combination of the constraints. let constraint_rlc_coeffs: Vec = @@ -150,26 +163,16 @@ impl Config { let has_constraints = !constraint_rlc_coeffs.is_empty(); let (initial_forms_rlc_coeffs, oods_rlc_coeffs) = constraint_rlc_coeffs.split_at(linear_forms.len()); - // Try to recycle the first linear form as Covector. - let mut covector = vec![]; let mut linear_forms = linear_forms; - if let Some((first, linear_forms)) = linear_forms.split_first_mut() { - debug_assert_eq!(initial_forms_rlc_coeffs[0], M::Target::ONE); - if let Some(covector_form) = - (first.as_mut() as &mut dyn Any).downcast_mut::>() - { - mem::swap(&mut covector, &mut covector_form.vector); - } else { - covector.resize(self.initial_size(), M::Target::ZERO); - first.accumulate(&mut covector, M::Target::ONE); - } - for (rlc_coeff, linear_form) in zip_strict(&initial_forms_rlc_coeffs[1..], linear_forms) - { - linear_form.accumulate(&mut covector, *rlc_coeff); - } - } else if has_constraints { - covector.resize(self.initial_size(), M::Target::ZERO); - } + let mut covector = if has_constraints { + as BufferOps>::linear_forms_rlc( + self.initial_size(), + &mut linear_forms, + initial_forms_rlc_coeffs, + ) + } else { + as BufferOps>::zeros(0) + }; drop(linear_forms); // Compute "The Sum" @@ -181,17 +184,17 @@ impl Config { .sum(); drop(evaluations); - debug_assert!(!has_constraints || dot(&vector, &covector) == the_sum); + debug_assert!(!has_constraints || vector.dot(&covector) == the_sum); // Add OODS constraints - UnivariateEvaluation::accumulate_many(&oods_evals, &mut covector, oods_rlc_coeffs); + covector.accumulate_univariate_evaluations(&oods_evals, oods_rlc_coeffs); the_sum += zip_strict(oods_rlc_coeffs, oods_matrix.chunks_exact(num_vectors)) .map(|(poly_coeff, row)| *poly_coeff * dot(&vector_rlc_coeffs, row)) .sum::(); drop(oods_evals); drop(oods_matrix); - debug_assert!(!has_constraints || dot(&vector, &covector) == the_sum); + debug_assert!(!has_constraints || vector.dot(&covector) == the_sum); // Run initial sumcheck on batched vectors with combined statement let mut folding_randomness = if has_constraints { @@ -208,35 +211,32 @@ impl Config { self.initial_skip_pow.prove(prover_state); // Fold vector for &f in &folding_randomness { - fold(&mut vector, f); + vector.fold(f); } // Covector must be all zeros. - covector = vec![M::Target::ZERO; self.initial_sumcheck.final_size()]; + covector = as BufferOps>::zeros( + self.initial_sumcheck.final_size(), + ); folding_randomness }; let mut evaluation_point = folding_randomness.clone(); - debug_assert_eq!(dot(&vector, &covector), the_sum); + debug_assert_eq!(vector.dot(&covector), the_sum); // Execute standard WHIR rounds on the batched vectors for (round_index, round_config) in self.round_configs.iter().enumerate() { // Commit to the vector, this generates out-of-domain evaluations. - let vector_buffer = CpuBuffer::from_slice(&vector); - let new_witness = round_config - .irs_committer - .commit(prover_state, &[&vector_buffer]); + let new_witness = round_config.irs_committer.commit(prover_state, &[&vector]); // Proof of work before in-domain challenges round_config.pow.prove(prover_state); // Open the previous round's witness. let in_domain = match prev_witness { - RoundWitness::Initial(init_witnesses) => { - let witness_refs: Vec<&_> = init_witnesses.iter().map(|c| &**c).collect(); - self.initial_committer - .open(prover_state, &witness_refs) - .lift(self.embedding()) - } + RoundWitness::Initial(init_witnesses) => self + .initial_committer + .open(prover_state, &init_witnesses) + .lift(self.embedding()), RoundWitness::Round(old_witness) => { let prev_round_config = &self.round_configs[round_index - 1]; prev_round_config @@ -260,13 +260,9 @@ impl Config { ))) .collect::>(); let stir_rlc_coeffs = geometric_challenge(prover_state, stir_challenges.len()); - UnivariateEvaluation::accumulate_many( - &stir_challenges, - &mut covector, - &stir_rlc_coeffs, - ); + covector.accumulate_univariate_evaluations(&stir_challenges, &stir_rlc_coeffs); the_sum += dot(&stir_rlc_coeffs, &stir_evaluations); - debug_assert_eq!(dot(&vector, &covector), the_sum); + debug_assert_eq!(vector.dot(&covector), the_sum); // Run sumcheck for this round folding_randomness = round_config @@ -275,7 +271,7 @@ impl Config { .round_challenges; evaluation_point.extend(folding_randomness.iter().copied()); - debug_assert_eq!(dot(&vector, &covector), the_sum); + debug_assert_eq!(vector.dot(&covector), the_sum); prev_witness = RoundWitness::Round(new_witness); vector_rlc_coeffs = vec![M::Target::ONE]; @@ -283,9 +279,7 @@ impl Config { // Directly send the vector to the verifier. assert_eq!(vector.len(), self.final_sumcheck.initial_size); - for coeff in &vector { - prover_state.prover_message(coeff); - } + vector.write_to_prover(prover_state); // PoW self.final_pow.prove(prover_state); @@ -293,8 +287,7 @@ impl Config { // Open and consume the final previous witness. match prev_witness { RoundWitness::Initial(init_witnesses) => { - let witness_refs: Vec<&_> = init_witnesses.iter().map(|c| &**c).collect(); - let _in_domain = self.initial_committer.open(prover_state, &witness_refs); + let _in_domain = self.initial_committer.open(prover_state, &init_witnesses); } RoundWitness::Round(old_witness) => { let prev_config = self.round_configs.last().unwrap(); From 595b19d4be6928fc6ad2e229b1b2977a78bc347e Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 3 Jun 2026 17:30:35 +0200 Subject: [PATCH 06/35] chore: add temp ignore notes --- src/protocols/basecase.rs | 2 ++ src/protocols/code_switch.rs | 2 ++ src/protocols/mask_proximity.rs | 2 ++ src/protocols/whir_zk/committer.rs | 2 ++ src/protocols/whir_zk/mod.rs | 2 ++ 5 files changed, 10 insertions(+) diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index e6faa026..d1e7e73b 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + //! Base Case Linear Opening Protocol //! //! It support honest verifier zero-knowledge (HVZK), but is not succinct. diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 9da26424..aa2a6370 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + //! Code-switching IOR: R_{C, C_zk, sl} → R_{C', C_zk, sl'} //! //! Reduces a proximity claim about oracle f (source code C) to a proximity diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index deb49de5..7413158c 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + //! Mask proximity verification via γ-combination. //! //! Implements Construction 7.2 (p.43-44) specialized for zero-constraint mask diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index 01b7a863..dbc875ec 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution}; #[cfg(feature = "tracing")] diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index c27ca37b..493ea32e 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + #![cfg(feature = "rs_in_order")] // TODO: Support permuted. mod committer; mod prover; From c09d54c89c9c172ff59d8a534ea43b4eca1186ad Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 10 Jun 2026 14:01:18 +0200 Subject: [PATCH 07/35] refactor: remove transcript interactions from buffer design --- docs/buffer-prove-design.md | 284 ++++++++++++++++++++++++++++++++ src/algebra/buffer.rs | 287 ++++++++------------------------- src/protocols/irs_commit.rs | 52 +++--- src/protocols/matrix_commit.rs | 51 +++--- src/protocols/merkle_tree.rs | 2 +- src/protocols/sumcheck.rs | 17 +- src/protocols/whir/mod.rs | 18 +-- src/protocols/whir/prover.rs | 53 +++--- 8 files changed, 437 insertions(+), 327 deletions(-) create mode 100644 docs/buffer-prove-design.md diff --git a/docs/buffer-prove-design.md b/docs/buffer-prove-design.md new file mode 100644 index 00000000..3d08229f --- /dev/null +++ b/docs/buffer-prove-design.md @@ -0,0 +1,284 @@ +# Buffer Abstraction Design + +## Purpose + +The prover manipulates large vectors, encoded Reed-Solomon matrices, and Merkle +node arrays. On CPU, those objects can be ordinary `Vec`s. On an accelerator, +asking for a slice of the same object can force a synchronization and a full +readback. + +The abstraction exists to make ownership explicit: + +```text +Large prover data lives in an active backend object. +Protocol code only reads proof-sized data. +``` + +Proof-sized data means transcript scalars, Merkle roots, selected rows, selected +authentication nodes, and final folded vectors that are sent directly into the +proof. Full vectors, full encoded matrices, and full Merkle trees should stay in +`ActiveBuffer` or `ActiveMatrix`. + +## Design Rules + +The buffer layer is a data and math boundary. It should not own protocol logic. + +```text +Protocol layer: + transcript order + challenge sampling + proof layout + round structure + hash choice + verification semantics + +Buffer layer: + backend-owned vectors + backend-owned matrices + full-buffer math loops + selected reads + explicit full reads +``` + +That split rules out protocol-specific buffer APIs. WHIR should keep one prove +path, IRS should keep one commit path, and buffers should not write to the +transcript. Each protocol should place backend-owned data only at the boundary +where that protocol actually consumes large resident data. + +## Active Types + +Protocol code refers to active aliases instead of hard-coding a backend: + +```rust +pub type ActiveBuffer = selected_backend::Buffer; +pub type ActiveMatrix = selected_backend::Matrix; +``` + +The current selected backend is CPU, but that is a local implementation detail +inside the buffer module. The aliases are intentionally central. WHIR, IRS, +sumcheck, matrix commitment, and Merkle opening operate on `ActiveBuffer` and +`ActiveMatrix`, so a later backend can be selected at this boundary without +changing protocol callsites. + +The active types also expose the explicit construction and read boundaries used +by protocol code: + +```rust +// Setup or protocol-owned construction. +ActiveBuffer::from_slice(...); +ActiveBuffer::from_vec(...); +ActiveBuffer::zeros(...); +ActiveBuffer::random(...); +ActiveMatrix::from_vec(...); + +// Proof-boundary materialization. +buffer.read(); +buffer.read_index(...); +buffer.read_indices(...); +matrix.read_rows(...); +``` + +These are intentionally not part of `BufferOps`. They are ownership and +materialization boundaries, while `BufferOps` is only the full-buffer math +surface. + +## Core Trait + +`BufferOps` covers operations that are naturally full-buffer operations. Loops +over protocol rounds remain outside the trait. + +```rust +pub trait BufferOps: Clone { + // Same backend family for another field. + type Buffer: BufferOps; + + // Metadata. + fn len(&self) -> usize; + fn is_empty(&self) -> bool; + + // In-place full-buffer transforms. + fn fold(&mut self, weight: F); + + // Full-buffer reductions used by sumcheck and WHIR. + fn sumcheck_polynomial(&self, other: &Self) -> (F, F); + fn fold_and_sumcheck_polynomial( + &mut self, + other: &mut Self, + weight: F, + ) -> (F, F); + + // Mixed-field full-buffer reductions used by protocol checks. + fn mixed_dot>( + &self, + embedding: &M, + other: &Self::Buffer, + ) -> M::Target; + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target; + + // Writes into backend-owned state instead of returning a large vector. + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut Self::Buffer, + weight: M::Target, + ); + + fn geometric_accumulate(&mut self, scalars: &[F], points: &[F]); +} +``` + +Construction and host materialization are explicit active-type boundaries, not +part of this math trait. Protocol code may upload data at commitment setup and +may read proof-sized data at transcript/opening boundaries, but those actions +should stay visible at the call site. + +Scalar-returning methods are still synchronization points. That is acceptable +when the scalar is transcript data. If a later backend kernel should consume the +result directly, the right design is a fused operation or a backend-owned scalar, +not a hidden readback followed by an upload. + +## Matrix Shape + +`ActiveMatrix` is the backend-owned row matrix produced by interleaved RS +encoding and consumed by matrix commitment. + +The protocol-visible matrix surface is intentionally small: shape metadata and +selected logical row readback. + +The caller owns the protocol metadata: row count, column count, query indices, +and how selected rows are written into the proof. The backend owns physical +layout. For example, an accelerated matrix may store rows in a different order +or column-major internally, but `read_rows(indices)` returns logical rows. + +## IRS Boundary + +IRS remains protocol code. It samples masks, calls the RS encoder, commits the +encoded rows, samples OOD points, and records OOD evaluations. + +The shared RS primitive is: + +```rust +pub fn interleaved_rs_encode( + vectors: &[&ActiveBuffer], + masks: &ActiveBuffer, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, +) -> ActiveMatrix; +``` + +This replaces the old call shape that required `&[&[F]]` and a mask slice. IRS +does not need to expose host slices just to encode. + +IRS witness names describe what is actually stored: + +```rust +pub struct Witness +where + G: Field, +{ + pub masks: ActiveBuffer, + pub encoded_matrix: ActiveMatrix, + pub encoded_matrix_witness: matrix_commit::Witness, + pub out_of_domain: Evaluations, +} +``` + +The masks and encoded matrix remain backend-owned. OOD evaluations are still +host vectors because they are proof/transcript data in the current protocol. + +## Matrix Commitment Boundary + +Matrix commitment owns the hash choice and transcript write. The buffer layer +does not commit rows and does not select a hash function. + +The commit path is: + +```text +ActiveMatrix + -> matrix_commit::Config::build_nodes(...) + -> ActiveBuffer + -> merkle_tree::Config::commit_nodes(...) + -> root written to transcript +``` + +The important ownership change is that the Merkle witness stores the node array +as an active hash buffer: + +```rust +pub struct Witness { + nodes: ActiveBuffer, +} +``` + +Opening a Merkle tree reads only the required authentication nodes. That keeps +the full tree resident and turns openings into proof-sized reads. + +## WHIR Prove Boundary + +WHIR `prove` accepts active buffers for the committed vectors and borrowed +witness references: + +```rust +pub fn prove<'a, H, R>( + &self, + prover_state: &mut ProverState, + vectors: &[&ActiveBuffer], + witnesses: &'a [&'a Witness], + linear_forms: Vec>>, + evaluations: Cow<'a, [M::Target]>, +) -> FinalClaim +``` + +The WHIR prover still owns the protocol decisions: + +```text +sample vector RLC coefficients +build a backend-owned linear combination +sample constraint RLC coefficients +build the covector from linear forms +add OOD and STIR constraints +run sumcheck rounds +commit and open IRS rounds +send the final folded vector +``` + +The buffer layer only performs the large vector work inside those steps: + +```rust +vector.mixed_scalar_mul_add_to(...); +covector.geometric_accumulate(...); +sumcheck.prove(..., &mut vector, &mut covector, ...); +``` + +## Readback Policy + +Expected readbacks are proof-sized: Merkle roots, authentication nodes, selected +IRS rows, transcript scalars, and the final folded WHIR vector. Suspicious +readbacks are full committed vectors, full encoded matrices, full Merkle trees, +or a readback immediately followed by upload into another backend buffer. + +The current CPU implementation may borrow slices inside CPU-only primitives, +but that access stays crate-private. IRS, WHIR, and matrix commitment callers +should see active buffers and explicit proof-boundary reads. + +## Open Work + +Scalar-returning reductions currently return host scalars. That is fine for +transcript values, but later backend work may need fused operations or +backend-owned scalar handles when the scalar feeds another backend computation. + +The ZK wrapper remains a host-side wrapper for now. Its outer API accepts host +slices/owned vectors, builds masked and blinding vectors on the host, and only +constructs active buffers at the point where it calls the inner WHIR +commit/prove paths. That keeps this refactor focused on the regular WHIR +residency boundary without pushing buffer requirements through unfinished ZK +code. + +Future backend selection should happen behind `ActiveBuffer` and `ActiveMatrix`. +Protocol APIs should not grow backend-specific entrypoints. diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index c6c8fb2b..10f65f7e 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -1,60 +1,58 @@ +use core::num; use std::{any::Any, mem}; use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; -use spongefish::DuplexSpongeInterface; use crate::{ algebra::{ embedding::{Embedding, Identity}, linear_form::{Covector, LinearForm, UnivariateEvaluation}, mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, ntt, - sumcheck::{compute_sumcheck_polynomial, fold, fold_and_compute_polynomial}, + sumcheck::{compute_sumcheck_polynomial, fold}, }, - hash::Hash, - protocols::matrix_commit, - transcript::{ProverMessage, ProverState}, - type_info::TypeInfo, utils::chunks_exact_or_empty, }; -pub trait BufferOps: Clone { - type Buffer: BufferOps; - type Matrix: MatrixBufferOps; - - fn from_vec(source: Vec) -> Self; - fn zeros(length: usize) -> Self; +pub trait BufferOps { + fn from_vec(source: Vec) -> Self; + fn from_slice(source: &[T]) -> Self; + fn as_slice(&self) -> &[T]; + fn num_rows(&self, num_cols: usize) -> usize { + self.len() / num_cols + } + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; fn len(&self) -> usize; - fn is_empty(&self) -> bool { self.len() == 0 } +} + +pub trait FieldOps: Clone { + type TargetBuffer: FieldOps; + + fn zeros(length: usize) -> Self; + // consider removing this to avoid dependency on ark for Buffer design fn random(rng: &mut R, length: usize) -> Self where R: RngCore + CryptoRng, Standard: Distribution; - - fn linear_forms_rlc( - size: usize, - linear_forms: &mut [Box>], - rlc_coeffs: &[F], - ) -> Self; - fn zero_pad(&mut self); + fn dot(&self, other: &Self) -> F; fn fold(&mut self, weight: F); fn sumcheck_polynomial(&self, other: &Self) -> (F, F); - fn fold_and_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F); fn accumulate_univariate_evaluations( &mut self, evaluators: &[UnivariateEvaluation], scalars: &[F], ); - fn write_to_prover(&self, prover_state: &mut ProverState) - where - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - F: ProverMessage<[H::U]>; + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self; + fn mixed_extend, T: Field>( &self, embedding: &M, @@ -63,7 +61,7 @@ pub trait BufferOps: Clone { fn mixed_dot, T: Field>( &self, embedding: &M, - other: &Self::Buffer, + other: &Self::TargetBuffer, ) -> M::Target; fn mixed_univariate_evaluate>( &self, @@ -74,61 +72,20 @@ pub trait BufferOps: Clone { embedding: &M, vectors: &[&Self], coeffs: &[M::Target], - ) -> Self::Buffer; + ) -> Self::TargetBuffer; fn mixed_scalar_mul_add_to>( &self, embedding: &M, - accumulator: &mut Self::Buffer, + accumulator: &mut Self::TargetBuffer, weight: M::Target, ); - fn mixed_dot_slice>( - &self, - embedding: &M, - other: &[M::Target], - ) -> M::Target; fn interleaved_rs_encode( vectors: &[&Self], masks: &Self, message_length: usize, interleaving_depth: usize, codeword_length: usize, - ) -> Self::Matrix - where - F: 'static; - fn dot(&self, other: &Self) -> F; -} - -pub trait MatrixBufferOps { - type Witness; - - fn len(&self) -> usize; - fn num_rows(&self) -> usize; - fn num_cols(&self) -> usize; - - fn commit_rows( - &self, - config: &matrix_commit::Config, - prover_state: &mut ProverState, - ) -> Self::Witness - where - F: TypeInfo + matrix_commit::Encodable + Send + Sync, - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - Hash: ProverMessage<[H::U]>; - - fn open_rows( - &self, - config: &matrix_commit::Config, - prover_state: &mut ProverState, - witness: &Self::Witness, - indices: &[usize], - ) where - F: TypeInfo + matrix_commit::Encodable + Send + Sync, - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - Hash: ProverMessage<[H::U]>; - - fn read_rows(&self, indices: &[usize]) -> Vec; + ) -> Self; } #[derive( @@ -143,124 +100,54 @@ pub trait MatrixBufferOps { serde::Serialize, serde::Deserialize, )] -pub struct CpuMatrix { - data: Vec, - num_rows: usize, - num_cols: usize, +pub struct CpuBuffer { + data: Vec, } -impl CpuMatrix { - pub fn from_vec(data: Vec, num_rows: usize, num_cols: usize) -> Self { - assert_eq!(data.len(), num_rows * num_cols); +impl CpuBuffer { + pub fn from_vec(source: Vec) -> Self { + Self { data: source } + } + + pub fn from_slice(source: &[T]) -> Self { Self { - data, - num_rows, - num_cols, + data: Vec::from(source), } } - fn row(&self, row: usize) -> &[F] { - let start = row * self.num_cols; - let end = start + self.num_cols; - &self.data[start..end] + pub(crate) fn as_slice(&self) -> &[T] { + self.data.as_slice() } } -impl MatrixBufferOps for CpuMatrix -where - F: Field + TypeInfo + matrix_commit::Encodable + Send + Sync, -{ - type Witness = matrix_commit::Witness; +impl BufferOps for CpuBuffer { + fn as_slice(&self) -> &[T] { + &self.data + } fn len(&self) -> usize { self.data.len() } - fn num_rows(&self) -> usize { - self.num_rows - } - - fn num_cols(&self) -> usize { - self.num_cols - } - - fn commit_rows( - &self, - config: &matrix_commit::Config, - prover_state: &mut ProverState, - ) -> Self::Witness - where - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - Hash: ProverMessage<[H::U]>, - { - assert_eq!(config.num_rows(), self.num_rows); - assert_eq!(config.num_cols, self.num_cols); - config.commit(prover_state, &self.data) - } - - fn open_rows( - &self, - config: &matrix_commit::Config, - prover_state: &mut ProverState, - witness: &Self::Witness, - indices: &[usize], - ) where - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - Hash: ProverMessage<[H::U]>, - { - config.open(prover_state, witness, indices); - } - - fn read_rows(&self, indices: &[usize]) -> Vec { - let mut rows = Vec::with_capacity(indices.len() * self.num_cols); - for &index in indices { - assert!(index < self.num_rows); - rows.extend_from_slice(self.row(index)); + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { + let mut result = Vec::with_capacity(indices.len() * num_cols); + for i in indices { + result.extend_from_slice(&self.data[i * num_cols..(i + 1) * num_cols]); } - rows - } -} -#[derive( - Clone, - PartialEq, - Eq, - PartialOrd, - Ord, - Hash, - Debug, - Default, - serde::Serialize, - serde::Deserialize, -)] -pub struct CpuBuffer { - data: Vec, -} - -impl CpuBuffer { - pub fn from_vec(source: Vec) -> Self { - Self { data: source } + result } - pub fn from_slice(source: &[F]) -> Self { - Self { - data: Vec::from(source), - } + fn from_vec(source: Vec) -> Self { + Self::from_vec(source) } - pub(crate) fn as_slice(&self) -> &[F] { - self.data.as_slice() + fn from_slice(source: &[T]) -> Self { + Self::from_slice(source) } } -impl BufferOps for CpuBuffer { - type Buffer = CpuBuffer; - type Matrix = CpuMatrix; - - fn from_vec(source: Vec) -> Self { - Self::from_vec(source) - } +impl FieldOps for CpuBuffer { + type TargetBuffer = CpuBuffer; fn zeros(length: usize) -> Self { Self { @@ -268,10 +155,6 @@ impl BufferOps for CpuBuffer { } } - fn len(&self) -> usize { - self.data.len() - } - fn random(rng: &mut R, length: usize) -> Self where R: RngCore + CryptoRng, @@ -319,10 +202,6 @@ impl BufferOps for CpuBuffer { compute_sumcheck_polynomial(&self.data, other.as_slice()) } - fn fold_and_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - fold_and_compute_polynomial(&mut self.data, &mut other.data, weight) - } - fn accumulate_univariate_evaluations( &mut self, evaluators: &[UnivariateEvaluation], @@ -331,15 +210,6 @@ impl BufferOps for CpuBuffer { UnivariateEvaluation::accumulate_many(evaluators, &mut self.data, scalars); } - fn write_to_prover(&self, prover_state: &mut ProverState) - where - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - F: ProverMessage<[H::U]>, - { - prover_state.prover_messages(&self.data); - } - fn mixed_extend, T: Field>( &self, embedding: &M, @@ -351,7 +221,7 @@ impl BufferOps for CpuBuffer { fn mixed_dot, T: Field>( &self, embedding: &M, - other: &Self::Buffer, + other: &Self::TargetBuffer, ) -> M::Target { mixed_dot(embedding, other.as_slice(), self.as_slice()) } @@ -372,7 +242,7 @@ impl BufferOps for CpuBuffer { embedding: &M, vectors: &[&Self], coeffs: &[M::Target], - ) -> Self::Buffer { + ) -> Self::TargetBuffer { assert_eq!(vectors.len(), coeffs.len()); let Some((first, vectors)) = vectors.split_first() else { return CpuBuffer::from_vec(Vec::new()); @@ -388,55 +258,28 @@ impl BufferOps for CpuBuffer { fn mixed_scalar_mul_add_to>( &self, embedding: &M, - accumulator: &mut Self::Buffer, + accumulator: &mut Self::TargetBuffer, weight: M::Target, ) { mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.as_slice()); } - fn mixed_dot_slice>( - &self, - embedding: &M, - other: &[M::Target], - ) -> M::Target { - mixed_dot(embedding, other, self.as_slice()) - } - fn interleaved_rs_encode( vectors: &[&Self], masks: &Self, message_length: usize, interleaving_depth: usize, codeword_length: usize, - ) -> Self::Matrix - where - F: 'static, - { + ) -> Self { let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); - interleaved_rs_encode_slices( - &vectors, + let messages = vectors + .iter() + .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) + .collect::>(); + Self::from_vec(ntt::interleaved_rs_encode( + &messages, masks.as_slice(), - message_length, - interleaving_depth, codeword_length, - ) + )) } } - -fn interleaved_rs_encode_slices( - vectors: &[&[F]], - masks: &[F], - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, -) -> CpuMatrix { - let messages = vectors - .iter() - .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) - .collect::>(); - CpuMatrix::from_vec( - ntt::interleaved_rs_encode(&messages, masks, codeword_length), - codeword_length, - vectors.len() * interleaving_depth, - ) -} diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 1d9e20a8..8bb6facf 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -24,16 +24,18 @@ use std::{ ops::Neg, }; +use crate::{algebra::buffer::FieldOps, protocols::merkle_tree}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use ordered_float::OrderedFloat; use serde::{Deserialize, Serialize}; + #[cfg(feature = "tracing")] use tracing::instrument; use crate::{ algebra::{ - buffer::{BufferOps, CpuBuffer, CpuMatrix, MatrixBufferOps}, + buffer::{BufferOps, CpuBuffer}, dot, embedding::Embedding, fields::FieldWithSize, @@ -98,15 +100,13 @@ pub struct Config { #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] #[must_use] -pub struct Witness, Matrix = CpuMatrix> +pub struct Witness where G: Field, - Masks: BufferOps, - Matrix: MatrixBufferOps, { - pub masks: Masks, - pub matrix: Matrix, - pub matrix_witness: Matrix::Witness, + pub masks: CpuBuffer, + pub matrix: CpuBuffer, + pub matrix_witness: matrix_commit::Witness, pub out_of_domain: Evaluations, source_field: PhantomData, } @@ -298,15 +298,12 @@ impl Config { /// Commit to one or more vectors. #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] - pub fn commit( + pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&B], - ) -> Witness + vectors: &[&CpuBuffer], + ) -> Witness where - B: BufferOps, - >::Witness: - Clone + PartialEq + Eq + PartialOrd + Ord + fmt::Debug + std::hash::Hash + Default, Standard: Distribution, H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -322,8 +319,11 @@ impl Config { assert_eq!(vectors.len(), self.num_vectors); assert!(vectors.iter().all(|p| p.len() == self.vector_size)); - let masks = B::random(prover_state.rng(), self.mask_length * self.num_messages()); - let matrix = B::interleaved_rs_encode( + let masks = CpuBuffer::::random( + prover_state.rng(), + self.mask_length * self.num_messages(), + ); + let matrix = CpuBuffer::interleaved_rs_encode( vectors, &masks, self.message_length(), @@ -332,7 +332,7 @@ impl Config { ); // Commit to the matrix - let matrix_witness = matrix.commit_rows(&self.matrix_commit, prover_state); + let matrix_witness = self.matrix_commit.commit(prover_state, &matrix); // Handle out-of-domain points and values // TODO : Remove this logic after main whir protocol is updated @@ -394,14 +394,12 @@ impl Config { /// When there are multiple openings, the evaluation matrices will /// be horizontally concatenated. #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] - pub fn open( + pub fn open( &self, prover_state: &mut ProverState, - witnesses: &[&Witness], + witnesses: &[&Witness], ) -> Evaluations where - Masks: BufferOps, - Matrix: MatrixBufferOps, H: DuplexSpongeInterface, R: RngCore + CryptoRng, u8: Decoding<[H::U]>, @@ -425,7 +423,7 @@ impl Config { let mut matrix = vec![M::Source::ZERO; indices.len() * stride]; let mut matrix_col_offset = 0; for witness in witnesses { - let submatrix = witness.matrix.read_rows(&indices); + let submatrix = witness.matrix.read_rows(self.num_cols(), &indices); if self.num_cols() != 0 { for (point_index, row) in submatrix.chunks_exact(self.num_cols()).enumerate() { let matrix_row = &mut matrix[point_index * stride..(point_index + 1) * stride]; @@ -434,12 +432,8 @@ impl Config { } } prover_state.prover_hint_ark(&submatrix); - witness.matrix.open_rows( - &self.matrix_commit, - prover_state, - &witness.matrix_witness, - &indices, - ); + self.matrix_commit + .open(prover_state, &witness.matrix_witness, &indices); matrix_col_offset += self.num_cols(); } @@ -526,12 +520,10 @@ impl Commitment { } } -impl Witness +impl Witness where F: Field, G: Field, - Masks: BufferOps, - Matrix: MatrixBufferOps, { /// Returns the out-of-domain evaluations. pub const fn out_of_domain(&self) -> &Evaluations { diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index 1e823162..33ffac52 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -12,6 +12,7 @@ use tracing::instrument; use zerocopy::{Immutable, IntoBytes}; use crate::{ + algebra::buffer::{BufferOps, CpuBuffer}, engines::EngineId, hash::{self, Hash}, protocols::merkle_tree, @@ -111,7 +112,12 @@ where pub merkle_tree: merkle_tree::Config, } -pub type Witness = merkle_tree::Witness; +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] +pub struct BufferWitness { + pub nodes: CpuBuffer, +} + +pub type Witness = BufferWitness; pub type Commitment = merkle_tree::Commitment; @@ -168,7 +174,7 @@ impl Encoder for ZeroCopyEncoder { } } -impl Config { +impl Config { /// Create a new matrix commit configuration with the recommended hash function. pub fn new(num_rows: usize, num_cols: usize) -> Self { // Select a leaf hash function. @@ -208,7 +214,11 @@ impl Config { /// Commit the matrix (in row-major order). #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self, size = matrix.len(), engine)))] - pub fn commit(&self, prover_state: &mut ProverState, matrix: &[T]) -> Witness + pub fn commit( + &self, + prover_state: &mut ProverState, + matrix: &CpuBuffer, + ) -> Witness where H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -225,10 +235,12 @@ impl Config { // Compute leaf hashes let mut leaves = Vec::with_capacity(self.merkle_tree.num_nodes()); leaves.resize(self.merkle_tree.num_leaves, Hash::default()); - hash_rows(&*engine, matrix, &mut leaves[..self.num_rows()]); + hash_rows(&*engine, matrix.as_slice(), &mut leaves[..self.num_rows()]); // Commit the leaf hashes - self.merkle_tree.commit(prover_state, leaves) + BufferWitness { + nodes: CpuBuffer::::from_vec(self.merkle_tree.commit(prover_state, leaves).nodes), + } } #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] @@ -258,7 +270,13 @@ impl Config { R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>, { - self.merkle_tree.open(prover_state, witness, indices); + self.merkle_tree.open( + prover_state, + &merkle_tree::Witness { + nodes: Vec::from(witness.nodes.as_slice()), + }, + indices, + ); } /// Verifies the commitment at the provided row indices. @@ -292,7 +310,7 @@ impl Config { } } -impl fmt::Display for Config { +impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "MatrixCommit({} x {})", self.num_rows(), self.num_cols) } @@ -428,22 +446,9 @@ pub(crate) mod tests { .instance(&Empty); // Instance - let matrix: Vec = (0..config.size()).map(|_| rng.gen()).collect(); - let submatrix: Vec = if num_cols > 0 { - indices - .iter() - .flat_map(|&index| { - matrix - .chunks_exact(num_cols) - .nth(index) - .unwrap() - .iter() - .cloned() - }) - .collect::>() - } else { - Vec::new() - }; + let matrix: CpuBuffer = + CpuBuffer::from_vec((0..config.size()).map(|_| rng.gen()).collect()); + let submatrix: Vec = matrix.read_rows(num_cols, indices); // Prover let mut prover_state = ProverState::new_std(&ds); diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index fcf44ba0..3c1a7ea7 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -58,7 +58,7 @@ pub struct Commitment { #[must_use] pub struct Witness { /// The nodes in the Merkle tree, starting with the leaf hash layer. - nodes: Vec, + pub nodes: Vec, } impl Config { diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 13c0c255..e1354e61 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -2,6 +2,7 @@ use std::fmt; +use crate::algebra::buffer::FieldOps; use ark_ff::Field; use ark_std::rand::{CryptoRng, RngCore}; use serde::{Deserialize, Serialize}; @@ -9,7 +10,10 @@ use serde::{Deserialize, Serialize}; use tracing::instrument; use crate::{ - algebra::{buffer::BufferOps, univariate_evaluate}, + algebra::{ + buffer::{BufferOps, CpuBuffer}, + univariate_evaluate, + }, protocols::proof_of_work, transcript::{ codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverState, VerificationResult, @@ -62,16 +66,15 @@ impl Config { /// - Applies proof-of-work grinding if required. /// - Returns the sampled folding randomness values used in each reduction step. #[cfg_attr(feature = "tracing", instrument(skip_all))] - pub fn prove( + pub fn prove( &self, prover_state: &mut ProverState, - a: &mut B, - b: &mut B, + a: &mut CpuBuffer, + b: &mut CpuBuffer, sum: &mut F, masks: &[F], ) -> SumcheckOpening where - B: BufferOps, H: DuplexSpongeInterface, R: CryptoRng + RngCore, F: Codec<[H::U]>, @@ -107,7 +110,9 @@ impl Config { { // Fold and compute sumcheck polynomial in one pass. let (c0, c2) = if let Some(w) = prev_round_challenge { - a.fold_and_sumcheck_polynomial(b, w) + a.fold(w); + b.fold(w); + a.sumcheck_polynomial(b) } else { a.sumcheck_polynomial(b) }; diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 0f20a14b..0430b28e 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -46,13 +46,8 @@ pub struct RoundConfig { pub pow: proof_of_work::Config, } -pub type Witness, B = CpuBuffer<::Source>> = - irs_commit::Witness< - ::Source, - F, - B, - ::Source>>::Matrix, - >; +pub type Witness> = + irs_commit::Witness<::Source, F>; pub type Commitment = irs_commit::Commitment; #[must_use = "The final claim must be checked if there where any linear forms."] @@ -86,15 +81,12 @@ impl Config { feature = "tracing", instrument(skip_all, fields(size = vectors.first().unwrap().len())) )] - pub fn commit( + pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&B], - ) -> Witness + vectors: &[&CpuBuffer], + ) -> Witness where - B: BufferOps, - >::Witness: - Clone + PartialEq + Eq + PartialOrd + Ord + Debug + std::hash::Hash + Default, Standard: Distribution, H: DuplexSpongeInterface, R: RngCore + CryptoRng, diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 01c87f3c..4fcc266a 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -1,5 +1,6 @@ use std::{borrow::Cow, fmt::Debug}; +use crate::algebra::buffer::FieldOps; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; #[cfg(feature = "tracing")] @@ -8,7 +9,7 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ algebra::{ - buffer::{BufferOps, MatrixBufferOps}, + buffer::{BufferOps, CpuBuffer}, dot, embedding::Embedding, eq_weights, @@ -24,15 +25,13 @@ use crate::{ utils::zip_strict, }; -enum RoundWitness<'a, F, M, B> +enum RoundWitness<'a, F, M> where F: Field, M: Embedding, - B: BufferOps, - B::Buffer: BufferOps, { - Initial(Vec<&'a irs_commit::Witness>), - Round(irs_commit::Witness, as BufferOps>::Matrix>), + Initial(Vec<&'a irs_commit::Witness>), + Round(irs_commit::Witness), } impl Config { @@ -52,29 +51,15 @@ impl Config { /// #[cfg_attr(feature = "tracing", instrument(skip_all))] #[allow(clippy::too_many_lines, clippy::cognitive_complexity)] - pub fn prove<'a, H, R, B>( + pub fn prove<'a, H, R>( &self, prover_state: &mut ProverState, - vectors: &[&B], - witnesses: Vec<&'a Witness>, + vectors: &[&CpuBuffer], + witnesses: Vec<&'a Witness>, linear_forms: Vec>>, evaluations: Cow<'a, [M::Target]>, ) -> FinalClaim where - B: BufferOps, - B::Buffer: BufferOps, - >::Witness: - Clone + PartialEq + Eq + PartialOrd + Ord + Debug + std::hash::Hash + Default, - < as BufferOps>::Matrix as MatrixBufferOps< - M::Target, - >>::Witness: Clone - + PartialEq - + Eq - + PartialOrd - + Ord - + Debug - + std::hash::Hash - + Default, Standard: Distribution + Distribution, H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -106,7 +91,10 @@ impl Config { let covector = Covector::from(&**linear_form); for (vector, evaluation) in zip_strict(vectors, evaluations) { debug_assert_eq!( - vector.mixed_dot_slice(self.embedding(), &covector.vector), + vector.mixed_dot( + self.embedding(), + &CpuBuffer::from_slice(covector.vector.as_slice()) + ), *evaluation ); } @@ -153,9 +141,10 @@ impl Config { // Random linear combination of the vectors. let mut vector_rlc_coeffs: Vec = geometric_challenge(prover_state, num_vectors); assert_eq!(vector_rlc_coeffs[0], M::Target::ONE); - let mut vector = B::mixed_linear_combination(self.embedding(), vectors, &vector_rlc_coeffs); + let mut vector = + FieldOps::mixed_linear_combination(self.embedding(), vectors, &vector_rlc_coeffs); - let mut prev_witness: RoundWitness<'a, M::Target, M, B> = RoundWitness::Initial(witnesses); + let mut prev_witness: RoundWitness<'a, M::Target, M> = RoundWitness::Initial(witnesses); // Random linear combination of the constraints. let constraint_rlc_coeffs: Vec = @@ -165,13 +154,13 @@ impl Config { constraint_rlc_coeffs.split_at(linear_forms.len()); let mut linear_forms = linear_forms; let mut covector = if has_constraints { - as BufferOps>::linear_forms_rlc( + FieldOps::linear_forms_rlc( self.initial_size(), &mut linear_forms, initial_forms_rlc_coeffs, ) } else { - as BufferOps>::zeros(0) + FieldOps::zeros(0) }; drop(linear_forms); @@ -214,9 +203,7 @@ impl Config { vector.fold(f); } // Covector must be all zeros. - covector = as BufferOps>::zeros( - self.initial_sumcheck.final_size(), - ); + covector = FieldOps::zeros(self.initial_sumcheck.final_size()); folding_randomness }; let mut evaluation_point = folding_randomness.clone(); @@ -279,7 +266,9 @@ impl Config { // Directly send the vector to the verifier. assert_eq!(vector.len(), self.final_sumcheck.initial_size); - vector.write_to_prover(prover_state); + for coeff in vector.as_slice() { + prover_state.prover_message(coeff); + } // PoW self.final_pow.prove(prover_state); From d54e9b12a0f3a45ba4a64830c162fe172d156c65 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 10 Jun 2026 17:28:51 +0200 Subject: [PATCH 08/35] refactor: use active buffer instead hardcoded CpuBuffer --- Cargo.lock | 94 ++++++++++++++++++++++++++++++ Cargo.toml | 8 +++ src/algebra/buffer.rs | 10 +++- src/bin/benchmark.rs | 6 +- src/bin/main.rs | 6 +- src/protocols/basecase.rs | 14 ++--- src/protocols/code_switch.rs | 10 ++-- src/protocols/irs_commit.rs | 16 ++--- src/protocols/mask_proximity.rs | 8 +-- src/protocols/matrix_commit.rs | 14 +++-- src/protocols/sumcheck.rs | 12 ++-- src/protocols/whir/mod.rs | 20 +++---- src/protocols/whir/prover.rs | 8 +-- src/protocols/whir_zk/committer.rs | 8 +-- src/protocols/whir_zk/mod.rs | 10 ++-- 15 files changed, 178 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index acc6ca37..e1bd056b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -195,6 +195,12 @@ dependencies = [ "digest", ] +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + [[package]] name = "block-buffer" version = "0.10.4" @@ -388,6 +394,33 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "core-graphics-types" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb" +dependencies = [ + "bitflags", + "core-foundation", + "libc", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -544,6 +577,33 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + [[package]] name = "generic-array" version = "0.14.7" @@ -700,6 +760,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + [[package]] name = "matchers" version = "0.2.0" @@ -715,6 +784,21 @@ version = "2.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" +[[package]] +name = "metal" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7047791b5bc903b8cd963014b355f71dc9864a9a0b727057676c1dcae5cbc15" +dependencies = [ + "bitflags", + "block", + "core-graphics-types", + "foreign-types", + "log", + "objc", + "paste", +] + [[package]] name = "nix" version = "0.30.1" @@ -764,6 +848,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -1547,6 +1640,7 @@ dependencies = [ "digest", "hex", "hex-literal", + "metal", "ordered-float", "proptest", "rayon", diff --git a/Cargo.toml b/Cargo.toml index 515ee2b8..8ab53a45 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -56,6 +56,9 @@ arrayvec = "0.7.6" derive-where = { version = "1.6.0", features = ["safe"] } ordered-float = { version = "5.1.0", features = ["serde"] } +[target.'cfg(target_os = "macos")'.dependencies] +metal = { version = "0.33.0", optional = true } + [dev-dependencies] proptest = "1.0" serde_json = "1.0" @@ -74,6 +77,7 @@ parallel = ["dep:rayon"] rayon = ["dep:rayon"] asm = ["ark-ff/asm"] tracing = ["dep:tracing"] +metal = ["dep:metal"] # Do not permute evaluation pointsin RS # This flag is to be removed after whir_zk supports it. rs_in_order = [] @@ -86,6 +90,10 @@ harness = false name = "sumcheck" harness = false +[[bench]] +name = "metal_buffer" +harness = false + # Disable untill fixed. # [[bench]] # name = "zk_whir" diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index 10f65f7e..e6a22abc 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -1,4 +1,3 @@ -use core::num; use std::{any::Any, mem}; use ark_ff::Field; @@ -14,6 +13,15 @@ use crate::{ utils::chunks_exact_or_empty, }; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub use super::metal_buffer::MetalBuffer; + +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type ActiveBuffer = MetalBuffer; + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type ActiveBuffer = CpuBuffer; + pub trait BufferOps { fn from_vec(source: Vec) -> Self; fn from_slice(source: &[T]) -> Self; diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index 42e654b0..f8204c2b 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -10,7 +10,7 @@ use clap::Parser; use serde::Serialize; use whir::{ algebra::{ - buffer::CpuBuffer, + buffer::ActiveBuffer, embedding::{Basefield, Embedding, Identity}, fields::{Field128, Field192, Field256, Field64, Field64_2, Field64_3}, linear_form::{Evaluate, LinearForm, MultilinearExtension}, @@ -162,7 +162,7 @@ where HASH_COUNTER.reset(); - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( @@ -240,7 +240,7 @@ where HASH_COUNTER.reset(); let whir_prover_time = Instant::now(); - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let witness = params.commit(&mut prover_state, &[&vector_buffer]); let prove_linear_forms: Vec>> = points diff --git a/src/bin/main.rs b/src/bin/main.rs index c63ada92..1c19922b 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -6,7 +6,7 @@ use ark_std::rand::distributions::{Distribution, Standard}; use clap::Parser; use whir::{ algebra::{ - buffer::CpuBuffer, + buffer::ActiveBuffer, embedding::{Basefield, Embedding, Identity}, fields::{Field128, Field192, Field256, Field64, Field64_2, Field64_3}, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, @@ -149,7 +149,7 @@ where } let vector = (0..num_coeffs).map(M::Source::from).collect::>(); - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let whir_commit_time = Instant::now(); let witness = params.commit(&mut prover_state, &[&vector_buffer]); @@ -316,7 +316,7 @@ where } let whir_commit_time = Instant::now(); - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let witness = params.commit(&mut prover_state, &[&vector_buffer]); let whir_commit_time = whir_commit_time.elapsed(); diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index d1e7e73b..77eeeabb 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -13,7 +13,7 @@ use spongefish::{Decoding, VerificationResult}; use crate::{ algebra::{ - buffer::CpuBuffer, dot, embedding::Identity, multilinear_extend, random_vector, + buffer::ActiveBuffer, dot, embedding::Identity, multilinear_extend, random_vector, scalar_mul_add_new, univariate_evaluate, }, hash::Hash, @@ -83,8 +83,8 @@ impl Config { prover_state.prover_messages(&vector); prover_state.prover_messages(witness.masks.as_slice()); let _ = self.commit.open(prover_state, &[witness]); - let mut vector_buffer = CpuBuffer::from_vec(vector); - let mut covector_buffer = CpuBuffer::from_vec(covector); + let mut vector_buffer = ActiveBuffer::from_vec(vector); + let mut covector_buffer = ActiveBuffer::from_vec(covector); let point = self .sumcheck .prove( @@ -104,7 +104,7 @@ impl Config { // Create masking vector. let mask = random_vector(prover_state.rng(), vector.len()); - let mask_buffer = CpuBuffer::from_slice(&mask); + let mask_buffer = ActiveBuffer::from_slice(&mask); // Commit to the masking vector. let mask_witness = self.commit.commit(prover_state, &[&mask_buffer]); @@ -132,8 +132,8 @@ impl Config { // Run sumcheck to reduce linear form claim let mut masked_sum = mask_sum + mask_rlc * sum; - let mut masked_vector_buffer = CpuBuffer::from_vec(masked_vector); - let mut covector_buffer = CpuBuffer::from_vec(covector); + let mut masked_vector_buffer = ActiveBuffer::from_vec(masked_vector); + let mut covector_buffer = ActiveBuffer::from_vec(covector); let point = self .sumcheck .prove( @@ -304,7 +304,7 @@ mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let witness = config.commit.commit(&mut prover_state, &[&vector_buffer]); let prover_result = config.prove( &mut prover_state, diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index aa2a6370..ec457cb1 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -15,7 +15,7 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::CpuBuffer, + buffer::ActiveBuffer, dot, embedding::{Embedding, Identity}, eq_weights, geometric_accumulate, lift, mixed_dot, scalar_mul, univariate_evaluate, @@ -198,7 +198,7 @@ impl Config { }; // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 - let message_buffer = CpuBuffer::from_slice(&message); + let message_buffer = ActiveBuffer::from_slice(&message); let target_witness = self.target.commit(prover_state, &[&message_buffer]); // Step 2-3: OOD challenge + answers — Construction 9.7 Steps 2-3, p.55 @@ -525,7 +525,7 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let f_full_buffer = CpuBuffer::from_slice(&f_full); + let f_full_buffer = ActiveBuffer::from_slice(&f_full); let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); // Sample γ for sumcheck folding (length log2(ι)). @@ -579,7 +579,7 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let f_full_buffer = CpuBuffer::from_slice(&f_full); + let f_full_buffer = ActiveBuffer::from_slice(&f_full); let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); @@ -648,7 +648,7 @@ mod tests { // Commit honest f_full, fold to get the honest post-fold message. let mut prover_state = ProverState::new_std(&ds); - let f_full_buffer = CpuBuffer::from_slice(&f_full); + let f_full_buffer = ActiveBuffer::from_slice(&f_full); let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); let folded_message = diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 8bb6facf..8bc722f9 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -24,7 +24,7 @@ use std::{ ops::Neg, }; -use crate::{algebra::buffer::FieldOps, protocols::merkle_tree}; +use crate::algebra::buffer::FieldOps; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use ordered_float::OrderedFloat; @@ -35,7 +35,7 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::{BufferOps, CpuBuffer}, + buffer::{ActiveBuffer, BufferOps}, dot, embedding::Embedding, fields::FieldWithSize, @@ -104,8 +104,8 @@ pub struct Witness where G: Field, { - pub masks: CpuBuffer, - pub matrix: CpuBuffer, + pub masks: ActiveBuffer, + pub matrix: ActiveBuffer, pub matrix_witness: matrix_commit::Witness, pub out_of_domain: Evaluations, source_field: PhantomData, @@ -301,7 +301,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&CpuBuffer], + vectors: &[&ActiveBuffer], ) -> Witness where Standard: Distribution, @@ -319,11 +319,11 @@ impl Config { assert_eq!(vectors.len(), self.num_vectors); assert!(vectors.iter().all(|p| p.len() == self.vector_size)); - let masks = CpuBuffer::::random( + let masks = ActiveBuffer::::random( prover_state.rng(), self.mask_length * self.num_messages(), ); - let matrix = CpuBuffer::interleaved_rs_encode( + let matrix = ActiveBuffer::interleaved_rs_encode( vectors, &masks, self.message_length(), @@ -743,7 +743,7 @@ pub(crate) mod tests { let mut prover_state = ProverState::new_std(&ds); let vector_buffers = vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let vector_refs = vector_buffers.iter().collect::>(); let witness = config.commit(&mut prover_state, &vector_refs); diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index 7413158c..ae88e428 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -48,7 +48,7 @@ use serde::{Deserialize, Serialize}; use crate::{ algebra::{ - buffer::CpuBuffer, embedding::Identity, random_vector, scalar_mul_add_new, + buffer::ActiveBuffer, embedding::Identity, random_vector, scalar_mul_add_new, univariate_evaluate, }, hash::Hash, @@ -136,15 +136,15 @@ impl Config { let original_buffers = original_msgs .iter() - .map(|msg| CpuBuffer::from_slice(msg)) + .map(|msg| ActiveBuffer::from_slice(msg)) .collect::>(); let fresh_buffers = fresh_msgs .iter() - .map(|msg| CpuBuffer::from_slice(msg)) + .map(|msg| ActiveBuffer::from_slice(msg)) .collect::>(); // Tree layout: [originals..., freshes...] - let all_vectors: Vec<&CpuBuffer> = original_buffers + let all_vectors: Vec<&ActiveBuffer> = original_buffers .iter() .chain(fresh_buffers.iter()) .collect(); diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index 33ffac52..d851a407 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -12,7 +12,7 @@ use tracing::instrument; use zerocopy::{Immutable, IntoBytes}; use crate::{ - algebra::buffer::{BufferOps, CpuBuffer}, + algebra::buffer::{ActiveBuffer, BufferOps}, engines::EngineId, hash::{self, Hash}, protocols::merkle_tree, @@ -114,7 +114,7 @@ where #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] pub struct BufferWitness { - pub nodes: CpuBuffer, + pub nodes: ActiveBuffer, } pub type Witness = BufferWitness; @@ -217,7 +217,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - matrix: &CpuBuffer, + matrix: &ActiveBuffer, ) -> Witness where H: DuplexSpongeInterface, @@ -239,7 +239,9 @@ impl Config { // Commit the leaf hashes BufferWitness { - nodes: CpuBuffer::::from_vec(self.merkle_tree.commit(prover_state, leaves).nodes), + nodes: ActiveBuffer::::from_vec( + self.merkle_tree.commit(prover_state, leaves).nodes, + ), } } @@ -446,8 +448,8 @@ pub(crate) mod tests { .instance(&Empty); // Instance - let matrix: CpuBuffer = - CpuBuffer::from_vec((0..config.size()).map(|_| rng.gen()).collect()); + let matrix: ActiveBuffer = + ActiveBuffer::from_vec((0..config.size()).map(|_| rng.gen()).collect()); let submatrix: Vec = matrix.read_rows(num_cols, indices); // Prover diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index e1354e61..31f30b5e 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -11,7 +11,7 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::{BufferOps, CpuBuffer}, + buffer::{ActiveBuffer, BufferOps}, univariate_evaluate, }, protocols::proof_of_work, @@ -69,8 +69,8 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - a: &mut CpuBuffer, - b: &mut CpuBuffer, + a: &mut ActiveBuffer, + b: &mut ActiveBuffer, sum: &mut F, masks: &[F], ) -> SumcheckOpening @@ -243,7 +243,7 @@ mod tests { use super::*; use crate::{ algebra::{ - buffer::CpuBuffer, + buffer::ActiveBuffer, dot, fields::{self, Field64}, multilinear_extend, random_vector, @@ -302,8 +302,8 @@ mod tests { let masks = random_vector(&mut rng, config.mask_length * config.num_rounds); // Prover - let mut vector = CpuBuffer::from_slice(&initial_vector); - let mut covector = CpuBuffer::from_slice(&initial_covector); + let mut vector = ActiveBuffer::from_slice(&initial_vector); + let mut covector = ActiveBuffer::from_slice(&initial_covector); let mut sum = initial_sum; let mut prover_state = ProverState::new_std(&ds); let SumcheckOpening { diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 0430b28e..2a8506f7 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -14,7 +14,7 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::{BufferOps, CpuBuffer}, + buffer::ActiveBuffer, embedding::{Embedding, Identity}, linear_form::LinearForm, }, @@ -84,7 +84,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&CpuBuffer], + vectors: &[&ActiveBuffer], ) -> Witness where Standard: Distribution, @@ -133,7 +133,7 @@ mod tests { use super::*; use crate::{ algebra::{ - buffer::CpuBuffer, + buffer::ActiveBuffer, embedding::Basefield, fields::{Field64, Field64_3}, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, @@ -245,7 +245,7 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit to the polynomial and generate auxiliary witness data - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let witness = params.commit(&mut prover_state, &[&vector_buffer]); let prove_linear_forms = build_prove_forms(&points, num_variables, true); @@ -421,7 +421,7 @@ mod tests { // Commit to each polynomial and generate witnesses let vector_buffers = vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let mut witnesses = Vec::new(); for vec in &vector_buffers { @@ -577,15 +577,15 @@ mod tests { .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let vec1_buffer = CpuBuffer::from_slice(&vec1); - let vec2_buffer = CpuBuffer::from_slice(&vec2); + let vec1_buffer = ActiveBuffer::from_slice(&vec1); + let vec2_buffer = ActiveBuffer::from_slice(&vec2); let witness1 = params.commit(&mut prover_state, &[&vec1_buffer]); let witness2 = params.commit(&mut prover_state, &[&vec2_buffer]); let prove_linear_forms = build_prove_forms(&constraint_points, num_variables, false); // Generate proof with mismatched polynomials - let vec_wrong_buffer = CpuBuffer::from_vec(vec_wrong); + let vec_wrong_buffer = ActiveBuffer::from_vec(vec_wrong); let _ = params.prove( &mut prover_state, &[&vec1_buffer, &vec_wrong_buffer], @@ -696,7 +696,7 @@ mod tests { // Commit using commit_batch (stacks batch_size polynomials per witness) let vector_buffers = all_vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let buffer_refs = vector_buffers.iter().collect::>(); let mut witnesses = Vec::new(); @@ -824,7 +824,7 @@ mod tests { // Create a commitment to the polynomial and generate auxiliary witness data let vector_buffers = vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let buffer_refs = vector_buffers.iter().collect::>(); let batched_witness = params.commit(&mut prover_state, &buffer_refs); diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 4fcc266a..55649e62 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -1,4 +1,4 @@ -use std::{borrow::Cow, fmt::Debug}; +use std::borrow::Cow; use crate::algebra::buffer::FieldOps; use ark_ff::{AdditiveGroup, Field}; @@ -9,7 +9,7 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ algebra::{ - buffer::{BufferOps, CpuBuffer}, + buffer::{ActiveBuffer, BufferOps}, dot, embedding::Embedding, eq_weights, @@ -54,7 +54,7 @@ impl Config { pub fn prove<'a, H, R>( &self, prover_state: &mut ProverState, - vectors: &[&CpuBuffer], + vectors: &[&ActiveBuffer], witnesses: Vec<&'a Witness>, linear_forms: Vec>>, evaluations: Cow<'a, [M::Target]>, @@ -93,7 +93,7 @@ impl Config { debug_assert_eq!( vector.mixed_dot( self.embedding(), - &CpuBuffer::from_slice(covector.vector.as_slice()) + &ActiveBuffer::from_slice(covector.vector.as_slice()) ), *evaluation ); diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index dbc875ec..c8c73c07 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -7,7 +7,7 @@ use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; use crate::{ - algebra::buffer::CpuBuffer, + algebra::buffer::ActiveBuffer, hash::Hash, protocols::{irs_commit, whir}, transcript::{ @@ -46,7 +46,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - polynomials: &[&CpuBuffer], + polynomials: &[&ActiveBuffer], ) -> Witness where Standard: Distribution, @@ -90,7 +90,7 @@ impl Config { .zip(mask.iter().cycle()) .map(|(&coeff, &m)| coeff + m) .collect::>(); - let f_hat_buffer = CpuBuffer::from_slice(&f_hat_vec); + let f_hat_buffer = ActiveBuffer::from_slice(&f_hat_vec); let witness = self .blinded_commitment .commit(prover_state, &[&f_hat_buffer]); @@ -117,7 +117,7 @@ impl Config { } let blinding_buffers = blinding_vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let blinding_vector_refs = blinding_buffers.iter().collect::>(); let blinding_witness = self diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index 493ea32e..0f2ed6f2 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -251,7 +251,7 @@ mod tests { use super::*; use crate::{ algebra::{ - buffer::CpuBuffer, + buffer::ActiveBuffer, fields::Field64, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, random_vector, @@ -355,7 +355,7 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); let vector_buffers = vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let vector_refs = vector_buffers.iter().collect::>(); let witness = params.commit(&mut prover_state, &vector_refs); @@ -452,7 +452,7 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); let vector_buffers = vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let vector_refs = vector_buffers.iter().collect::>(); let witness = params.commit(&mut prover_state, &vector_refs); @@ -511,7 +511,7 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); let vector_buffers = vectors .iter() - .map(|v| CpuBuffer::from_slice(v)) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); let vector_refs = vector_buffers.iter().collect::>(); let witness = params.commit(&mut prover_state, &vector_refs); @@ -575,7 +575,7 @@ mod tests { let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let mut prover_state = ProverState::new_std(&ds); - let vector_buffer = CpuBuffer::from_slice(&vector); + let vector_buffer = ActiveBuffer::from_slice(&vector); let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( &mut prover_state, From 23a408a26ac722cae5bcac68223bc2b12a72d455 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 10 Jun 2026 20:01:35 +0200 Subject: [PATCH 09/35] wip: metal buffer backend --- benches/metal_buffer.rs | 86 ++ scripts/phase_benchmark_cpu_gpu.sh | 118 ++ scripts/plot_phase_benchmark.py | 631 ++++++++ src/algebra/metal_buffer.rs | 2296 ++++++++++++++++++++++++++++ src/algebra/mod.rs | 2 + src/bin/gpu_proof_cpu_verify.rs | 189 +++ src/bin/phase_benchmark.rs | 504 ++++++ src/bin/phase_benchmark_report.rs | 127 ++ src/hash/metal_profile.rs | 123 ++ src/hash/metal_sha2_engine.rs | 715 +++++++++ src/hash/mod.rs | 10 + src/protocols/matrix_commit.rs | 100 +- src/protocols/proof_of_work.rs | 10 + 13 files changed, 4910 insertions(+), 1 deletion(-) create mode 100644 benches/metal_buffer.rs create mode 100644 scripts/phase_benchmark_cpu_gpu.sh create mode 100644 scripts/plot_phase_benchmark.py create mode 100644 src/algebra/metal_buffer.rs create mode 100644 src/bin/gpu_proof_cpu_verify.rs create mode 100644 src/bin/phase_benchmark.rs create mode 100644 src/bin/phase_benchmark_report.rs create mode 100644 src/hash/metal_profile.rs create mode 100644 src/hash/metal_sha2_engine.rs diff --git a/benches/metal_buffer.rs b/benches/metal_buffer.rs new file mode 100644 index 00000000..d91b8f42 --- /dev/null +++ b/benches/metal_buffer.rs @@ -0,0 +1,86 @@ +#[cfg(all(feature = "metal", target_os = "macos"))] +mod bench { + use divan::{black_box, AllocProfiler, Bencher}; + use whir::{ + algebra::{ + buffer::{CpuBuffer, FieldOps, MetalBuffer}, + fields::Field256 as F, + }, + hash::{Hash, HashEngine, MetalSha2, Sha2}, + }; + + #[global_allocator] + static ALLOC: AllocProfiler = AllocProfiler::system(); + + const FIELD_SIZES: &[usize] = &[1 << 8, 1 << 10]; + const HASH_ROWS: &[usize] = &[1 << 10, 1 << 12]; + const HASH_ROW_SIZE: usize = 128; + + #[divan::bench(args = FIELD_SIZES)] + fn cpu_bn254_dot(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| { + let a = CpuBuffer::from_vec((0..size).map(|i| F::from(i as u64)).collect()); + let b = CpuBuffer::from_vec((0..size).map(|i| F::from((i + 7) as u64)).collect()); + (a, b) + }) + .bench_values(|(a, b)| black_box(a.dot(&b))); + } + + #[divan::bench(args = FIELD_SIZES)] + fn metal_bn254_dot(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| { + let a = MetalBuffer::from_vec((0..size).map(|i| F::from(i as u64)).collect()); + let b = MetalBuffer::from_vec((0..size).map(|i| F::from((i + 7) as u64)).collect()); + (a, b) + }) + .bench_values(|(a, b)| black_box(a.dot(&b))); + } + + #[divan::bench(args = HASH_ROWS)] + fn cpu_sha256_hash_many(bencher: Bencher, rows: usize) { + bencher + .with_inputs(|| { + let input = (0..rows * HASH_ROW_SIZE) + .map(|i| (i & 0xff) as u8) + .collect::>(); + let output = vec![Hash::default(); rows]; + (input, output) + }) + .bench_values(|(input, mut output)| { + Sha2::new().hash_many(HASH_ROW_SIZE, &input, &mut output); + black_box(output) + }); + } + + #[divan::bench(args = HASH_ROWS)] + fn metal_sha256_hash_many(bencher: Bencher, rows: usize) { + bencher + .with_inputs(|| { + let input = (0..rows * HASH_ROW_SIZE) + .map(|i| (i & 0xff) as u8) + .collect::>(); + let output = vec![Hash::default(); rows]; + (input, output) + }) + .bench_values(|(input, mut output)| { + MetalSha2::new().hash_many(HASH_ROW_SIZE, &input, &mut output); + black_box(output) + }); + } + + pub fn main() { + divan::main(); + } +} + +#[cfg(all(feature = "metal", target_os = "macos"))] +fn main() { + bench::main(); +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +fn main() { + eprintln!("metal_buffer benchmark requires macOS and --features metal"); +} diff --git a/scripts/phase_benchmark_cpu_gpu.sh b/scripts/phase_benchmark_cpu_gpu.sh new file mode 100644 index 00000000..028998eb --- /dev/null +++ b/scripts/phase_benchmark_cpu_gpu.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +set -euo pipefail + +out="${WHIR_PHASE_BENCH_OUTPUT:-outputs/phase_benchmark_cpu_gpu.jsonl}" +min_log_size=16 +max_log_size=28 +folds="1,2,3,4,6" +rates="1,2,3" +phases="commit,sumcheck,e2e" +article_grid=false +common_args=() +while [[ $# -gt 0 ]]; do + case "$1" in + --output) + out="$2" + shift 2 + ;; + --output=*) + out="${1#--output=}" + shift + ;; + --min-log-size) + min_log_size="$2" + shift 2 + ;; + --min-log-size=*) + min_log_size="${1#--min-log-size=}" + shift + ;; + --max-log-size) + max_log_size="$2" + shift 2 + ;; + --max-log-size=*) + max_log_size="${1#--max-log-size=}" + shift + ;; + --folds) + folds="$2" + shift 2 + ;; + --folds=*) + folds="${1#--folds=}" + shift + ;; + --rates) + rates="$2" + shift 2 + ;; + --rates=*) + rates="${1#--rates=}" + shift + ;; + --phases) + phases="$2" + shift 2 + ;; + --phases=*) + phases="${1#--phases=}" + shift + ;; + --article-grid) + article_grid=true + shift + ;; + *) + common_args+=("$1") + shift + ;; + esac +done + +mkdir -p "$(dirname "$out")" +: > "$out" + +cargo build --release --bin phase_benchmark +CARGO_TARGET_DIR=target/metal cargo build --release --features metal --bin phase_benchmark +cargo build --release --bin phase_benchmark_report + +IFS=, read -r -a fold_values <<< "$folds" +IFS=, read -r -a rate_values <<< "$rates" +IFS=, read -r -a phase_values <<< "$phases" + +cpu_bin="target/release/phase_benchmark" +gpu_bin="target/metal/release/phase_benchmark" + +for log_size in $(seq "$min_log_size" "$max_log_size"); do + for fold in "${fold_values[@]}"; do + if [[ "$fold" -gt "$log_size" ]]; then + printf 'skip n=%s fold=%s: fold must be <= n\n' "$log_size" "$fold" >&2 + continue + fi + for rate in "${rate_values[@]}"; do + if [[ "$article_grid" == true && "$log_size" -ge 24 && "$rate" -gt 1 ]]; then + continue + fi + + for phase in "${phase_values[@]}"; do + case_args=( + --output "$out" + --min-log-size "$log_size" + --max-log-size "$log_size" + --folds "$fold" + --rates "$rate" + --phases "$phase" + ) + + "$cpu_bin" "${case_args[@]}" "${common_args[@]}" + "$gpu_bin" "${case_args[@]}" "${common_args[@]}" + done + done + done +done + +target/release/phase_benchmark_report "$out" > "${out%.jsonl}.csv" + +printf 'wrote %s\n' "$out" +printf 'wrote %s\n' "${out%.jsonl}.csv" diff --git a/scripts/plot_phase_benchmark.py b/scripts/plot_phase_benchmark.py new file mode 100644 index 00000000..4fea4422 --- /dev/null +++ b/scripts/plot_phase_benchmark.py @@ -0,0 +1,631 @@ +#!/usr/bin/env python3 +import argparse +import csv +import html +import json +import math +from collections import defaultdict +from pathlib import Path + + +PHASES = ["commit", "sumcheck", "e2e_prove"] +PHASE_LABEL = { + "commit": "commit", + "sumcheck": "sumcheck", + "e2e_prove": "e2e prove", +} +BACKEND_LABEL = { + "cpu": "CPU", + "gpu-metal": "GPU Metal", +} +BACKEND_COLOR = { + "cpu": "#2f4b7c", + "gpu-metal": "#d95f02", +} +FOLD_COLOR = { + 1: "#2f4b7c", + 2: "#d95f02", + 3: "#1b9e77", + 4: "#7570b3", + 6: "#e7298a", +} +PROFILE_COLOR = { + "total": "#1f2937", + "command wait": "#7570b3", + "upload": "#d95f02", + "readback": "#1b9e77", + "blit": "#e7298a", +} + + +def load_rows(path): + rows = [] + with open(path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + return rows + + +def row_index(rows): + out = {} + for r in rows: + key = (r["backend"], r["phase"], r["log_size"], r["fold"], r["rate"]) + out[key] = r + return out + + +def paired_rows(rows): + idx = row_index(rows) + keys = sorted( + { + (r["phase"], r["log_size"], r["fold"], r["rate"]) + for r in rows + }, + key=lambda k: (k[1], k[2], k[3], k[0]), + ) + pairs = [] + for phase, log_size, fold, rate in keys: + cpu = idx.get(("cpu", phase, log_size, fold, rate)) + gpu = idx.get(("gpu-metal", phase, log_size, fold, rate)) + if cpu or gpu: + pairs.append((phase, log_size, fold, rate, cpu, gpu)) + return pairs + + +def gib(x): + return x / 1024.0 / 1024.0 / 1024.0 + + +def mib(x): + return x / 1024.0 / 1024.0 + + +def fmt_ms(ms): + if ms is None: + return "OOM" + if ms < 1000: + return f"{ms:.1f} ms" + return f"{ms / 1000:.2f} s" + + +def fmt_speedup(x): + if x is None: + return "OOM" + return f"{x:.2f}x" + + +def fmt_gib(x): + if x is None: + return "OOM" + return f"{gib(x):.1f}" + + +def nice_ticks(vmin, vmax, log_y=False, count=5): + if not math.isfinite(vmin) or not math.isfinite(vmax) or vmax <= 0: + return [] + if log_y: + lo = math.floor(math.log10(max(vmin, 1e-9))) + hi = math.ceil(math.log10(vmax)) + ticks = [] + for p in range(lo, hi + 1): + for m in (1, 2, 5): + value = m * (10 ** p) + if vmin <= value <= vmax: + ticks.append(value) + if len(ticks) > 7: + stride = max(1, math.ceil(len(ticks) / 7)) + ticks = ticks[::stride] + return ticks + if vmax == vmin: + return [vmin] + raw = (vmax - vmin) / max(1, count - 1) + mag = 10 ** math.floor(math.log10(raw)) + step = min((1, 2, 5, 10), key=lambda s: abs(raw - s * mag)) * mag + start = math.floor(vmin / step) * step + ticks = [] + value = start + while value <= vmax + step * 0.5: + if value >= vmin - step * 0.5: + ticks.append(value) + value += step + return ticks[:8] + + +def svg_text(x, y, text, size=12, anchor="middle", weight="400", fill="#111827"): + text = html.escape(str(text)) + return ( + f'{text}' + ) + + +def svg_poly(points, color, width=2.2, dash=False): + if len(points) < 2: + return "" + pts = " ".join(f"{x:.1f},{y:.1f}" for x, y in points) + dash_attr = ' stroke-dasharray="5 4"' if dash else "" + return ( + f'' + ) + + +def render_line_chart( + path, + title, + panels, + y_label, + x_label="log2(size)", + log_y=False, + shared_y=True, + width=1280, + panel_height=260, +): + panel_count = len(panels) + height = 96 + panel_count * panel_height + 28 + margin_l, margin_r, margin_t, margin_b = 76, 24, 74, 46 + plot_w = width - margin_l - margin_r + all_x = [ + x + for panel in panels + for series in panel["series"].values() + for x, y in series + if y is not None and y > 0 + ] + xmin, xmax = min(all_x), max(all_x) + + def y_range(panel): + values = [ + y + for series in panel["series"].values() + for x, y in series + if y is not None and y > 0 and math.isfinite(y) + ] + if not values: + return 0.1, 1.0 + lo, hi = min(values), max(values) + if log_y: + return lo / 1.35, hi * 1.35 + pad = (hi - lo) * 0.12 if hi > lo else max(1.0, hi * 0.1) + return max(0.0, lo - pad), hi + pad + + shared_range = None + if shared_y: + vals = [] + for panel in panels: + a, b = y_range(panel) + vals.extend([a, b]) + shared_range = min(vals), max(vals) + + parts = [ + f'', + '', + svg_text(width / 2, 32, title, 21, weight="700"), + svg_text(width / 2, height - 8, x_label, 12), + svg_text(18, height / 2, y_label, 12, anchor="middle"), + ] + parts.append(f'') + + legend_items = {} + for panel in panels: + for name, color in panel.get("colors", {}).items(): + legend_items[name] = color + lx = margin_l + ly = 54 + for name, color in legend_items.items(): + parts.append(f'') + parts.append(svg_text(lx + 34, ly + 4, name, 12, anchor="start")) + lx += 150 + + for i, panel in enumerate(panels): + top = margin_t + i * panel_height + bottom = top + panel_height - margin_b + y0 = top + 26 + plot_h = bottom - y0 + vmin, vmax = shared_range if shared_y else y_range(panel) + if log_y: + vmin = max(vmin, 1e-9) + lmin, lmax = math.log10(vmin), math.log10(vmax) + + def sy(v): + return bottom - (math.log10(max(v, 1e-9)) - lmin) / (lmax - lmin) * plot_h + + else: + + def sy(v): + return bottom - (v - vmin) / (vmax - vmin) * plot_h if vmax > vmin else bottom + + def sx(x): + return margin_l + (x - xmin) / (xmax - xmin) * plot_w if xmax > xmin else margin_l + plot_w / 2 + + parts.append(svg_text(margin_l, top + 14, panel["title"], 15, anchor="start", weight="700")) + parts.append(f'') + parts.append(f'') + + x_span = int(xmax) - int(xmin) + x_step = 1 if x_span <= 18 else 2 if x_span <= 36 else 4 + for x in range(int(xmin), int(xmax) + 1, x_step): + px = sx(x) + parts.append(f'') + parts.append(svg_text(px, bottom + 20, str(x), 11)) + for tick in nice_ticks(vmin, vmax, log_y=log_y): + py = sy(tick) + parts.append(f'') + label = f"{tick:g}" if tick < 1000 else f"{tick / 1000:g}k" + parts.append(svg_text(margin_l - 9, py + 4, label, 10, anchor="end", fill="#374151")) + + for name, series in panel["series"].items(): + color = panel.get("colors", {}).get(name, "#111827") + points = [(sx(x), sy(y)) for x, y in series if y is not None and y > 0] + parts.append(svg_poly(points, color)) + for px, py in points: + parts.append(f'') + + for note in panel.get("annotations", []): + x, y, text = note + parts.append(svg_text(sx(x), sy(y), text, 11, fill="#b91c1c", weight="700")) + + parts.append("") + Path(path).write_text("\n".join(parts), encoding="utf-8") + + +def write_csv(path, rows): + pairs = paired_rows(rows) + with open(path, "w", newline="", encoding="utf-8") as f: + writer = csv.writer(f) + writer.writerow( + [ + "log_size", + "size", + "fold", + "rate", + "phase", + "cpu_ms", + "gpu_ms", + "speedup", + "cpu_peak_gib", + "gpu_peak_gib", + "gpu_upload_gib", + "gpu_upload_ms", + "gpu_readback_mib", + "gpu_readback_ms", + "gpu_command_wait_ms", + "gpu_blit_gib", + "gpu_blit_wait_ms", + ] + ) + for phase, log_size, fold, rate, cpu, gpu in pairs: + cpu_ms = cpu and cpu.get("duration_ms") + gpu_ms = gpu and gpu.get("duration_ms") + speedup = cpu_ms / gpu_ms if cpu_ms and gpu_ms else None + writer.writerow( + [ + log_size, + 1 << log_size, + fold, + rate, + phase, + f"{cpu_ms:.6f}" if cpu_ms else "", + f"{gpu_ms:.6f}" if gpu_ms else "", + f"{speedup:.6f}" if speedup else "", + f"{gib(cpu['peak_allocated_bytes']):.6f}" if cpu else "", + f"{gib(gpu['peak_allocated_bytes']):.6f}" if gpu else "", + f"{gib(gpu.get('metal_upload_bytes', 0)):.6f}" if gpu else "", + f"{gpu.get('metal_upload_ms', 0):.6f}" if gpu else "", + f"{mib(gpu.get('metal_readback_bytes', 0)):.6f}" if gpu else "", + f"{gpu.get('metal_readback_ms', 0):.6f}" if gpu else "", + f"{gpu.get('metal_command_wait_ms', 0):.6f}" if gpu else "", + f"{gib(gpu.get('metal_blit_bytes', 0)):.6f}" if gpu else "", + f"{gpu.get('metal_blit_wait_ms', 0):.6f}" if gpu else "", + ] + ) + + +def markdown_table(headers, rows): + out = ["| " + " | ".join(headers) + " |"] + out.append("| " + " | ".join(["---"] * len(headers)) + " |") + for row in rows: + out.append("| " + " | ".join(str(x) for x in row) + " |") + return "\n".join(out) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("jsonl") + parser.add_argument("--plot-dir", default="outputs/plots") + parser.add_argument("--report", default="outputs/article_phase_benchmark_16_27_hybrid_report.md") + parser.add_argument("--paired-csv", default="outputs/article_phase_benchmark_16_27_hybrid_paired.csv") + parser.add_argument("--requested-min-log-size", type=int) + parser.add_argument("--requested-max-log-size", type=int) + args = parser.parse_args() + + rows = load_rows(args.jsonl) + idx = row_index(rows) + plot_dir = Path(args.plot_dir) + plot_dir.mkdir(parents=True, exist_ok=True) + write_csv(args.paired_csv, rows) + + logs = sorted({r["log_size"] for r in rows}) + folds = sorted({r["fold"] for r in rows}) + rates = sorted({r["rate"] for r in rows}) + + def get(backend, phase, log_size, fold, rate): + return idx.get((backend, phase, log_size, fold, rate)) + + def series_backend(backend, phase, fold, rate, metric): + points = [] + for log_size in logs: + r = get(backend, phase, log_size, fold, rate) + points.append((log_size, metric(r) if r else None)) + return points + + e2e_panels = [] + for fold in folds: + annotations = [] + if not get("cpu", "e2e_prove", 27, fold, 1) or not get("gpu-metal", "e2e_prove", 27, fold, 1): + max_y = max( + [ + r["duration_ms"] + for b in ("cpu", "gpu-metal") + for log_size in logs + for r in [get(b, "e2e_prove", log_size, fold, 1)] + if r + ] + or [1] + ) + annotations.append((27, max_y, "missing/killed")) + e2e_panels.append( + { + "title": f"fold={fold}, rate=1", + "series": { + "CPU": series_backend("cpu", "e2e_prove", fold, 1, lambda r: r["duration_ms"] if r else None), + "GPU Metal": series_backend("gpu-metal", "e2e_prove", fold, 1, lambda r: r["duration_ms"] if r else None), + }, + "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}, + "annotations": annotations, + } + ) + render_line_chart( + plot_dir / "e2e_runtime_rate1.svg", + "E2E prove runtime, rate=1", + e2e_panels, + "milliseconds, log scale", + log_y=True, + shared_y=True, + ) + + speed_series = {} + for fold in folds: + pts = [] + for log_size in logs: + cpu = get("cpu", "e2e_prove", log_size, fold, 1) + gpu = get("gpu-metal", "e2e_prove", log_size, fold, 1) + pts.append((log_size, cpu["duration_ms"] / gpu["duration_ms"] if cpu and gpu else None)) + speed_series[f"fold={fold}"] = pts + render_line_chart( + plot_dir / "e2e_speedup_rate1.svg", + "E2E prove speedup, rate=1", + [{"title": "CPU time / GPU Metal time", "series": speed_series, "colors": {f"fold={f}": FOLD_COLOR[f] for f in folds}}], + "speedup", + log_y=False, + shared_y=False, + panel_height=360, + ) + + phase_panels = [] + for phase in PHASES: + phase_panels.append( + { + "title": f"{PHASE_LABEL[phase]}, fold=1, rate=1", + "series": { + "CPU": series_backend("cpu", phase, 1, 1, lambda r: r["duration_ms"] if r else None), + "GPU Metal": series_backend("gpu-metal", phase, 1, 1, lambda r: r["duration_ms"] if r else None), + }, + "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}, + } + ) + render_line_chart( + plot_dir / "phase_runtime_fold1_rate1.svg", + "Phase runtime, fold=1, rate=1", + phase_panels, + "milliseconds, log scale", + log_y=True, + shared_y=False, + ) + + mem_panels = [] + for fold in folds: + mem_panels.append( + { + "title": f"fold={fold}, rate=1", + "series": { + "CPU": series_backend("cpu", "e2e_prove", fold, 1, lambda r: gib(r["peak_allocated_bytes"]) if r else None), + "GPU Metal": series_backend("gpu-metal", "e2e_prove", fold, 1, lambda r: gib(r["peak_allocated_bytes"]) if r else None), + }, + "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}, + } + ) + render_line_chart( + plot_dir / "peak_memory_e2e_rate1.svg", + "Peak allocated memory during E2E prove, rate=1", + mem_panels, + "GiB", + log_y=False, + shared_y=True, + ) + + max_mem = {"CPU": [], "GPU Metal": []} + for log_size in logs: + cpu_vals = [ + r["peak_allocated_bytes"] + for fold in folds + for r in [get("cpu", "e2e_prove", log_size, fold, 1)] + if r + ] + gpu_vals = [ + r["peak_allocated_bytes"] + for fold in folds + for r in [get("gpu-metal", "e2e_prove", log_size, fold, 1)] + if r + ] + max_mem["CPU"].append((log_size, gib(max(cpu_vals)) if cpu_vals else None)) + max_mem["GPU Metal"].append((log_size, gib(max(gpu_vals)) if gpu_vals else None)) + render_line_chart( + plot_dir / "peak_memory_max_e2e_rate1.svg", + "Max peak allocated memory by size, E2E prove rate=1", + [{"title": "max across successful folds", "series": max_mem, "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}}], + "GiB", + log_y=False, + shared_y=False, + panel_height=360, + ) + + profile = defaultdict(list) + for log_size in logs: + r = get("gpu-metal", "e2e_prove", log_size, 4, 1) + profile["total"].append((log_size, r["duration_ms"] if r else None)) + profile["command wait"].append((log_size, r.get("metal_command_wait_ms", 0) if r else None)) + profile["upload"].append((log_size, r.get("metal_upload_ms", 0) if r else None)) + profile["readback"].append((log_size, r.get("metal_readback_ms", 0) if r else None)) + profile["blit"].append((log_size, r.get("metal_blit_wait_ms", 0) if r else None)) + render_line_chart( + plot_dir / "gpu_profile_e2e_fold4_rate1.svg", + "GPU E2E profile counters, fold=4, rate=1", + [{"title": "wall time components", "series": dict(profile), "colors": PROFILE_COLOR}], + "milliseconds, log scale", + log_y=True, + shared_y=False, + panel_height=360, + ) + + speedups = [] + for fold in folds: + for log_size in logs: + cpu = get("cpu", "e2e_prove", log_size, fold, 1) + gpu = get("gpu-metal", "e2e_prove", log_size, fold, 1) + if cpu and gpu: + speedups.append(cpu["duration_ms"] / gpu["duration_ms"]) + speedups_sorted = sorted(speedups) + median_speedup = speedups_sorted[len(speedups_sorted) // 2] if speedups_sorted else None + + e2e_speed_table = [] + for log_size in logs: + row = [f"2^{log_size}"] + for fold in folds: + if fold > log_size: + row.append("n/a") + continue + cpu = get("cpu", "e2e_prove", log_size, fold, 1) + gpu = get("gpu-metal", "e2e_prove", log_size, fold, 1) + row.append(fmt_speedup(cpu["duration_ms"] / gpu["duration_ms"]) if cpu and gpu else "missing") + e2e_speed_table.append(row) + + mem_table = [] + for log_size in logs: + cpu_vals = [ + r["peak_allocated_bytes"] + for fold in folds + for r in [get("cpu", "e2e_prove", log_size, fold, 1)] + if r + ] + gpu_vals = [ + r["peak_allocated_bytes"] + for fold in folds + for r in [get("gpu-metal", "e2e_prove", log_size, fold, 1)] + if r + ] + c = max(cpu_vals) if cpu_vals else None + g = max(gpu_vals) if gpu_vals else None + mem_table.append([f"2^{log_size}", fmt_gib(c), fmt_gib(g), f"{c / g:.2f}x" if c and g else "OOM"]) + + phase_table = [] + phase_sample_logs = [x for x in [8, 16, 20, 24, 27] if x in logs] + for log_size in phase_sample_logs: + for phase in PHASES: + cpu = get("cpu", phase, log_size, 4, 1) + gpu = get("gpu-metal", phase, log_size, 4, 1) + phase_table.append( + [ + f"2^{log_size}", + PHASE_LABEL[phase], + fmt_ms(cpu["duration_ms"] if cpu else None), + fmt_ms(gpu["duration_ms"] if gpu else None), + fmt_speedup(cpu["duration_ms"] / gpu["duration_ms"] if cpu and gpu else None), + ] + ) + + profile_table = [] + for log_size in phase_sample_logs: + r = get("gpu-metal", "e2e_prove", log_size, 4, 1) + profile_table.append( + [ + f"2^{log_size}", + fmt_ms(r["duration_ms"] if r else None), + fmt_ms(r.get("metal_command_wait_ms") if r else None), + f"{gib(r.get('metal_upload_bytes', 0)):.2f} GiB / {r.get('metal_upload_ms', 0):.1f} ms" if r else "OOM", + f"{mib(r.get('metal_readback_bytes', 0)):.2f} MiB / {r.get('metal_readback_ms', 0):.1f} ms" if r else "OOM", + f"{gib(r.get('metal_blit_bytes', 0)):.2f} GiB / {r.get('metal_blit_wait_ms', 0):.1f} ms" if r else "OOM", + ] + ) + + expected = { + (r["log_size"], r["fold"], r["rate"]) + for r in rows + } + missing = [] + for log_size, fold, rate in sorted(expected): + for phase in PHASES: + if not get("cpu", phase, log_size, fold, rate): + missing.append(f"CPU {PHASE_LABEL[phase]} 2^{log_size} fold={fold} rate={rate}") + if not get("gpu-metal", phase, log_size, fold, rate): + missing.append(f"GPU {PHASE_LABEL[phase]} 2^{log_size} fold={fold} rate={rate}") + + report = [] + report.append("# WHIR CPU vs Metal GPU phase benchmark") + report.append("") + requested_min = args.requested_min_log_size if args.requested_min_log_size is not None else min(logs) + requested_max = args.requested_max_log_size if args.requested_max_log_size is not None else max(logs) + report.append("Hardware: Apple M4 Max, 40-core GPU, 48 GiB unified memory.") + report.append(f"Parameters: requested `log_size={requested_min}..{requested_max}`, actual rows `log_size={min(logs)}..{max(logs)}`, folds `1,2,3,4,6`, rates `1,2,3` where the article grid fits memory and `rate=1` for the largest sizes, phases `commit,sumcheck,e2e`, `pow_bits=20`, `security_level=128`.") + if requested_min < min(logs): + report.append(f"`2^{requested_min}` has no rows because the benchmark rejects the configured folds there (`fold` must be `<= n`).") + report.append("") + report.append(f"Raw rows: `{len(rows)}`. Paired CSV: `{args.paired_csv}`.") + if missing: + report.append("Missing/killed rows: " + "; ".join(missing) + ".") + else: + report.append("Missing/killed rows: none.") + if speedups: + report.append(f"E2E rate=1 paired speedup: min `{min(speedups):.2f}x`, median `{median_speedup:.2f}x`, max `{max(speedups):.2f}x`.") + report.append("") + report.append("## Charts") + for name in [ + "e2e_runtime_rate1.svg", + "e2e_speedup_rate1.svg", + "phase_runtime_fold1_rate1.svg", + "peak_memory_e2e_rate1.svg", + "peak_memory_max_e2e_rate1.svg", + "gpu_profile_e2e_fold4_rate1.svg", + ]: + report.append(f"- `{plot_dir / name}`") + report.append("") + report.append("## E2E speedup, rate=1") + report.append(markdown_table(["size"] + [f"fold={f}" for f in folds], e2e_speed_table)) + report.append("") + report.append("## Max peak allocated memory, E2E rate=1") + report.append(markdown_table(["size", "CPU GiB", "GPU GiB", "CPU/GPU"], mem_table)) + report.append("") + report.append("## Phase speedup, fold=4 rate=1") + report.append(markdown_table(["size", "phase", "CPU", "GPU", "CPU/GPU"], phase_table)) + report.append("") + report.append("## GPU transfer/profile counters, E2E fold=4 rate=1") + report.append(markdown_table(["size", "total", "command wait", "upload", "readback", "blit"], profile_table)) + report.append("") + report.append("Notes: standalone `sumcheck` measures the phase in isolation, so large rows include host-to-GPU upload cost. The E2E path keeps more data resident and is the better signal for real prover throughput. Peak memory is allocator peak from the benchmark process; Metal counters are the explicit upload/readback/blit instrumentation emitted by the benchmark.") + Path(args.report).write_text("\n".join(report) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs new file mode 100644 index 00000000..615dac2d --- /dev/null +++ b/src/algebra/metal_buffer.rs @@ -0,0 +1,2296 @@ +use std::{ + any::type_name, + cell::OnceCell, + cmp::Ordering, + collections::HashMap, + hash::{Hash, Hasher}, + marker::PhantomData, + os::raw::c_void, + sync::{Mutex, OnceLock}, + time::Instant, +}; + +use ark_ff::{AdditiveGroup, BigInt, FftField, Field, Fp, MontBackend}; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; +use metal::{ + Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, MTLResourceOptions, MTLSize, +}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use zerocopy::IntoBytes; + +use crate::{ + algebra::{ + buffer::{BufferOps, FieldOps}, + embedding::{Embedding, Identity}, + fields::{BN254Config, Field256}, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, + }, + hash::{metal_profile, Hash as Digest, MetalSha2}, +}; + +const METAL_SOURCE: &str = r#" +#include +using namespace metal; + +struct F { ulong v[4]; }; + +constant ulong MODULUS[4] = { + 0x43e1f593f0000001UL, + 0x2833e84879b97091UL, + 0xb85045b68181585dUL, + 0x30644e72e131a029UL, +}; + +constant ulong BN254_N0PRIME = 0xc2e1f593efffffffUL; +constant F CANONICAL_ONE = {{1UL, 0UL, 0UL, 0UL}}; +constant F MONT_ONE = {{ + 0xac96341c4ffffffbUL, + 0x36fc76959f60cd29UL, + 0x666ea36f7879462eUL, + 0xe0a77c19a07df2fUL, +}}; + +struct StageConfig { + uint row_len; + uint half_m; + uint twiddle_offset; + uint _pad0; +}; + +struct BitReverseParams { + uint row_len; + uint log_n; + uint total_elements; + uint _pad0; +}; + +struct TransposeParams { + uint rows; + uint cols; + uint total_elements; +}; + +struct ReplicateCosetsParams { + uint row_len; + uint coset_size; + uint trailing_elements; +}; + +struct PackSingleVectorParams { + uint row_count; + uint message_length; + uint codeword_length; + uint coset_size; + uint total_elements; +}; + +struct FieldBytesParams { + uint rows; + uint cols; +}; + +struct ApplyCosetTwiddlesParams { + uint row_count; + uint num_cosets; + uint coset_size; + uint codeword_length; + uint total_elements; +}; + +inline F zero_f() { + F r; + r.v[0] = 0; r.v[1] = 0; r.v[2] = 0; r.v[3] = 0; + return r; +} + +inline F one_f() { + return MONT_ONE; +} + +inline F load_f(device const ulong *data, uint idx) { + F r; + uint base = idx * 4; + r.v[0] = data[base + 0]; + r.v[1] = data[base + 1]; + r.v[2] = data[base + 2]; + r.v[3] = data[base + 3]; + return r; +} + +inline void store_f(device ulong *data, uint idx, F x) { + uint base = idx * 4; + data[base + 0] = x.v[0]; + data[base + 1] = x.v[1]; + data[base + 2] = x.v[2]; + data[base + 3] = x.v[3]; +} + +inline bool ge_mod(F a) { + for (int i = 3; i >= 0; i--) { + if (a.v[i] > MODULUS[i]) return true; + if (a.v[i] < MODULUS[i]) return false; + } + return true; +} + +inline bool ge_f(F a, F b) { + for (int i = 3; i >= 0; i--) { + if (a.v[i] > b.v[i]) return true; + if (a.v[i] < b.v[i]) return false; + } + return true; +} + +inline F sub_raw(F a, thread const ulong *b, thread bool &borrow) { + F r; + borrow = false; + for (uint i = 0; i < 4; i++) { + ulong bi = b[i] + (borrow ? 1UL : 0UL); + bool bcarry = borrow && bi == 0; + r.v[i] = a.v[i] - bi; + borrow = bcarry || a.v[i] < bi; + } + return r; +} + +inline F sub_modulus(F a) { + F r; + bool borrow = false; + for (uint i = 0; i < 4; i++) { + ulong bi = MODULUS[i] + (borrow ? 1UL : 0UL); + bool bcarry = borrow && bi == 0; + r.v[i] = a.v[i] - bi; + borrow = bcarry || a.v[i] < bi; + } + return r; +} + +inline F add_modulus(F a) { + F r; + bool carry = false; + for (uint i = 0; i < 4; i++) { + ulong mi = MODULUS[i]; + ulong sum = a.v[i] + mi; + bool c0 = sum < a.v[i]; + ulong sum2 = sum + (carry ? 1UL : 0UL); + bool c1 = carry && sum2 == 0; + r.v[i] = sum2; + carry = c0 || c1; + } + return r; +} + +inline F add_f(F a, F b) { + F r; + bool carry = false; + for (uint i = 0; i < 4; i++) { + ulong sum = a.v[i] + b.v[i]; + bool c0 = sum < a.v[i]; + ulong sum2 = sum + (carry ? 1UL : 0UL); + bool c1 = carry && sum2 == 0; + r.v[i] = sum2; + carry = c0 || c1; + } + if (carry || ge_mod(r)) { + r = sub_modulus(r); + } + return r; +} + +inline F sub_f(F a, F b) { + ulong bv[4] = { b.v[0], b.v[1], b.v[2], b.v[3] }; + bool borrow = false; + F r = sub_raw(a, bv, borrow); + if (borrow) { + r = add_modulus(r); + } + return r; +} + +inline F double_f(F a) { + return add_f(a, a); +} + +inline ulong add_with_carry(ulong a, ulong b, thread ulong &carry) { + ulong sum = a + b; + ulong c1 = sum < a ? 1UL : 0UL; + ulong sum_with_carry = sum + carry; + ulong c2 = sum_with_carry < sum ? 1UL : 0UL; + carry = c1 + c2; + return sum_with_carry; +} + +inline void add_scaled_step(thread ulong &dst, ulong s, ulong a, thread ulong &carry) { + ulong product_lo = s * a; + ulong product_hi = mulhi(s, a); + + ulong sum = dst + product_lo; + ulong carry0 = sum < dst ? 1UL : 0UL; + ulong sum_with_carry = sum + carry; + ulong carry1 = sum_with_carry < sum ? 1UL : 0UL; + + dst = sum_with_carry; + carry = product_hi + carry0 + carry1; +} + +inline void add_scaled(thread ulong *dst, ulong s, ulong a0, ulong a1, ulong a2, ulong a3) { + ulong carry = 0; + add_scaled_step(dst[0], s, a0, carry); + add_scaled_step(dst[1], s, a1, carry); + add_scaled_step(dst[2], s, a2, carry); + add_scaled_step(dst[3], s, a3, carry); + dst[4] += carry; +} + +inline F mont_mul(F lhs, F rhs) { + ulong buf[9] = {0}; + uint off = 0; + +#pragma clang loop unroll(enable) + for (uint i = 0; i < 4; i++) { + add_scaled(&buf[off], lhs.v[i], rhs.v[0], rhs.v[1], rhs.v[2], rhs.v[3]); + + ulong m = buf[off] * BN254_N0PRIME; + add_scaled( + &buf[off], + m, + MODULUS[0], + MODULUS[1], + MODULUS[2], + MODULUS[3] + ); + + off += 1; + buf[off + 4] = 0; + } + + F result; + result.v[0] = buf[off + 0]; + result.v[1] = buf[off + 1]; + result.v[2] = buf[off + 2]; + result.v[3] = buf[off + 3]; + if (ge_mod(result)) { + result = sub_modulus(result); + } + return result; +} + +inline F from_mont(F value) { + F result = mont_mul(value, CANONICAL_ONE); + if (ge_mod(result)) { + result = sub_modulus(result); + } + return result; +} + +inline F mul_f(F a, F b) { + return mont_mul(a, b); +} + +inline F pow_f(F base, uint exp) { + F acc = one_f(); + F x = base; + while (exp != 0) { + if ((exp & 1) != 0) { + acc = mul_f(acc, x); + } + exp >>= 1; + if (exp != 0) { + x = mul_f(x, x); + } + } + return acc; +} + +kernel void bn254_fold( + device ulong *values [[buffer(0)]], + device const ulong *weight_buf [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &fold_half [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= fold_half) return; + F low = gid < len ? load_f(values, gid) : zero_f(); + F high = gid + fold_half < len ? load_f(values, gid + fold_half) : zero_f(); + F weight = load_f(weight_buf, 0); + store_f(values, gid, add_f(low, mul_f(sub_f(high, low), weight))); +} + +kernel void bn254_scalar_mul_add( + device ulong *acc [[buffer(0)]], + device const ulong *vector [[buffer(1)]], + device const ulong *weight_buf [[buffer(2)]], + constant uint &len [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F weight = load_f(weight_buf, 0); + store_f(acc, gid, add_f(load_f(acc, gid), mul_f(weight, load_f(vector, gid)))); +} + +kernel void bn254_dot( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + F acc = zero_f(); + for (uint i = 0; i < len; i++) { + acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); + } + store_f(out, 0, acc); +} + +kernel void bn254_sumcheck( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &fold_half [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = 0; i < fold_half; i++) { + F a0 = i < len ? load_f(a, i) : zero_f(); + F b0 = i < len ? load_f(b, i) : zero_f(); + F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); + F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); + c0 = add_f(c0, mul_f(a0, b0)); + c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); + } + store_f(out, 0, c0); + store_f(out, 1, c2); +} + +kernel void bn254_dot_chunks( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &chunk_size [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + F acc = zero_f(); + for (uint i = start; i < end; i++) { + acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); + } + store_f(out, gid, acc); +} + +kernel void bn254_sum_chunks( + device const ulong *input [[buffer(0)]], + device ulong *out [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &offset [[buffer(3)]], + constant uint &chunk_size [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + F acc = zero_f(); + for (uint i = start; i < end; i++) { + acc = add_f(acc, load_f(input, offset + i)); + } + store_f(out, gid, acc); +} + +kernel void bn254_sumcheck_chunks( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &fold_half [[buffer(4)]], + constant uint &chunk_size [[buffer(5)]], + constant uint &partial_count [[buffer(6)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= fold_half) return; + uint end = min(start + chunk_size, fold_half); + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = start; i < end; i++) { + F a0 = i < len ? load_f(a, i) : zero_f(); + F b0 = i < len ? load_f(b, i) : zero_f(); + F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); + F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); + c0 = add_f(c0, mul_f(a0, b0)); + c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); + } + store_f(out, gid, c0); + store_f(out, partial_count + gid, c2); +} + +kernel void bn254_geometric_accumulate( + device ulong *acc [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *scalars [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &num_points [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F value = load_f(acc, gid); + for (uint j = 0; j < num_points; j++) { + value = add_f(value, mul_f(load_f(scalars, j), pow_f(load_f(points, j), gid))); + } + store_f(acc, gid, value); +} + +kernel void bn254_geometric_accumulate_chunks( + device ulong *acc [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *scalars [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &num_points [[buffer(4)]], + constant uint &chunk_size [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + for (uint j = 0; j < num_points; j++) { + F point = load_f(points, j); + F scalar = load_f(scalars, j); + F power = pow_f(point, start); + for (uint i = start; i < end; i++) { + F value = load_f(acc, i); + value = add_f(value, mul_f(scalar, power)); + store_f(acc, i, value); + power = mul_f(power, point); + } + } +} + +kernel void bn254_geometric_accumulate_chunks_strided( + device ulong *acc [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *point_steps [[buffer(2)]], + device const ulong *scalars [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &num_points [[buffer(5)]], + constant uint &chunk_size [[buffer(6)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + for (uint j = 0; j < num_points; j++) { + F point = load_f(points, j); + F scalar = load_f(scalars, j); + F power = pow_f(load_f(point_steps, j), gid); + for (uint i = start; i < end; i++) { + F value = load_f(acc, i); + value = add_f(value, mul_f(scalar, power)); + store_f(acc, i, value); + power = mul_f(power, point); + } + } +} + +kernel void bn254_geometric_accumulate_point_blocks( + device ulong *partials [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *point_steps [[buffer(2)]], + device const ulong *scalars [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &num_points [[buffer(5)]], + constant uint &chunk_size [[buffer(6)]], + constant uint &point_block_size [[buffer(7)]], + constant uint &point_blocks [[buffer(8)]], + uint gid [[thread_position_in_grid]] +) { + uint chunk = gid / point_blocks; + uint point_block = gid - chunk * point_blocks; + uint start = chunk * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + uint point_start = point_block * point_block_size; + if (point_start >= num_points) return; + uint point_end = min(point_start + point_block_size, num_points); + + F sums[16]; + for (uint k = 0; k < chunk_size; k++) { + sums[k] = zero_f(); + } + + for (uint j = point_start; j < point_end; j++) { + F point = load_f(points, j); + F scalar = load_f(scalars, j); + F power = pow_f(load_f(point_steps, j), chunk); + for (uint i = start, k = 0; i < end; i++, k++) { + sums[k] = add_f(sums[k], mul_f(scalar, power)); + power = mul_f(power, point); + } + } + + uint partial_offset = point_block * len; + for (uint i = start, k = 0; i < end; i++, k++) { + store_f(partials, partial_offset + i, sums[k]); + } +} + +kernel void bn254_geometric_accumulate_reduce_point_blocks( + device ulong *acc [[buffer(0)]], + device const ulong *partials [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &point_blocks [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F value = load_f(acc, gid); + for (uint block = 0; block < point_blocks; block++) { + value = add_f(value, load_f(partials, block * len + gid)); + } + store_f(acc, gid, value); +} + +kernel void bn254_univariate_evaluate( + device const ulong *coeffs [[buffer(0)]], + device const ulong *point_buf [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + if (len == 0) { + store_f(out, 0, zero_f()); + return; + } + F point = load_f(point_buf, 0); + F acc = load_f(coeffs, len - 1); + for (uint i = len - 1; i > 0; i--) { + acc = add_f(mul_f(acc, point), load_f(coeffs, i - 1)); + } + store_f(out, 0, acc); +} + +kernel void bn254_univariate_eval_chunks( + device const ulong *coeffs [[buffer(0)]], + device const ulong *point_buf [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &chunk_size [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + F point = load_f(point_buf, 0); + F power = pow_f(point, start); + F acc = zero_f(); + for (uint i = start; i < end; i++) { + acc = add_f(acc, mul_f(load_f(coeffs, i), power)); + power = mul_f(power, point); + } + store_f(out, gid, acc); +} + +kernel void bn254_interleaved_rs_encode( + device const ulong *coeffs [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &num_messages [[buffer(3)]], + constant uint &coeff_len [[buffer(4)]], + constant uint &total [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= total) return; + uint row = gid / num_messages; + uint message = gid - row * num_messages; + uint base = message * coeff_len; + F point = load_f(points, row); + F acc = load_f(coeffs, base + coeff_len - 1); + for (uint i = coeff_len - 1; i > 0; i--) { + acc = add_f(mul_f(acc, point), load_f(coeffs, base + i - 1)); + } + store_f(out, gid, acc); +} + +kernel void bn254_interleaved_rs_encode_single_vector( + device const ulong *vector [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &interleaving_depth [[buffer(3)]], + constant uint &message_length [[buffer(4)]], + constant uint &total [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= total) return; + uint row = gid / interleaving_depth; + uint message = gid - row * interleaving_depth; + uint base = message * message_length; + F point = load_f(points, row); + F acc = load_f(vector, base + message_length - 1); + for (uint i = message_length - 1; i > 0; i--) { + acc = add_f(mul_f(acc, point), load_f(vector, base + i - 1)); + } + store_f(out, gid, acc); +} + +inline uint reverse_bits_width(uint value, uint width) { + return reverse_bits(value) >> (32u - width); +} + +kernel void bn254_pack_single_vector_cosets( + device const ulong *vector [[buffer(0)]], + device ulong *out [[buffer(1)]], + constant PackSingleVectorParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint row = gid / params.codeword_length; + uint col = gid - row * params.codeword_length; + if (col < params.message_length) { + store_f(out, gid, load_f(vector, row * params.message_length + col)); + } else { + store_f(out, gid, zero_f()); + } +} + +kernel void bn254_replicate_first_coset( + device ulong *buffer [[buffer(0)]], + constant ReplicateCosetsParams ¶ms [[buffer(1)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.trailing_elements) return; + uint repeats_per_row = params.row_len - params.coset_size; + uint row = gid / repeats_per_row; + uint within = gid - row * repeats_per_row; + uint dst = row * params.row_len + params.coset_size + within; + uint src = row * params.row_len + (within % params.coset_size); + store_f(buffer, dst, load_f(buffer, src)); +} + +kernel void bn254_bit_reverse_permute_rows_in_place( + device ulong *values [[buffer(0)]], + constant BitReverseParams &config [[buffer(1)]], + uint index [[thread_position_in_grid]] +) { + if (index >= config.total_elements || config.row_len <= 1u) return; + uint row = index / config.row_len; + uint within = index - row * config.row_len; + uint reversed = reverse_bits_width(within, config.log_n); + if (reversed <= within) return; + + uint row_base = row * config.row_len; + uint mate = row_base + reversed; + uint current = row_base + within; + F tmp = load_f(values, current); + store_f(values, current, load_f(values, mate)); + store_f(values, mate, tmp); +} + +kernel void bn254_radix2_ntt_stage_rows_in_place( + device ulong *values [[buffer(0)]], + device const ulong *twiddles [[buffer(1)]], + constant StageConfig &config [[buffer(2)]], + uint index [[thread_position_in_grid]] +) { + uint butterflies_per_row = config.row_len >> 1u; + uint row = index / butterflies_per_row; + uint local = index - row * butterflies_per_row; + uint half_m = config.half_m; + uint pair_in_group = local % half_m; + uint group = local / half_m; + uint row_base = row * config.row_len; + uint base = row_base + group * (half_m << 1u) + pair_in_group; + uint mate = base + half_m; + + F even = load_f(values, base); + F odd = load_f(values, mate); + F twiddle = load_f(twiddles, config.twiddle_offset + pair_in_group); + F t = mul_f(twiddle, odd); + + store_f(values, base, add_f(even, t)); + store_f(values, mate, sub_f(even, t)); +} + +kernel void bn254_transpose_matrix_reverse_rows( + device const ulong *input [[buffer(0)]], + device ulong *output [[buffer(1)]], + constant TransposeParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint row = gid / params.cols; + uint col = gid - row * params.cols; + uint row_bits = 31u - clz(params.cols); + uint dst_row = reverse_bits_width(col, row_bits); + uint dst = dst_row * params.rows + row; + store_f(output, dst, load_f(input, gid)); +} + +kernel void bn254_transpose_matrix( + device const ulong *input [[buffer(0)]], + device ulong *output [[buffer(1)]], + constant TransposeParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint row = gid / params.cols; + uint col = gid - row * params.cols; + uint dst = col * params.rows + row; + store_f(output, dst, load_f(input, gid)); +} + +kernel void bn254_apply_coset_twiddles( + device ulong *values [[buffer(0)]], + device const ulong *root_powers [[buffer(1)]], + constant ApplyCosetTwiddlesParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint within_codeword = gid % params.codeword_length; + uint coset = within_codeword / params.coset_size; + uint col = within_codeword - coset * params.coset_size; + if (coset == 0 || col == 0) return; + uint root_index = (coset * col) % params.codeword_length; + store_f(values, gid, mul_f(load_f(values, gid), load_f(root_powers, root_index))); +} + +kernel void bn254_encode_field_rows_le( + device const ulong *input [[buffer(0)]], + device uchar *output [[buffer(1)]], + constant FieldBytesParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + uint total_elements = params.rows * params.cols; + if (gid >= total_elements) return; + + F canonical = from_mont(load_f(input, gid)); + uint byte_offset = gid * 32u; + for (uint limb = 0; limb < 4; ++limb) { + ulong value = canonical.v[limb]; + for (uint byte = 0; byte < 8; ++byte) { + output[byte_offset + limb * 8u + byte] = uchar((value >> (byte * 8u)) & 0xffUL); + } + } +} + +kernel void bn254_read_rows( + device const ulong *input [[buffer(0)]], + device const uint *indices [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &num_cols [[buffer(3)]], + constant uint &total [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= total) return; + uint row = gid / num_cols; + uint col = gid - row * num_cols; + uint src = indices[row] * num_cols + col; + store_f(out, gid, load_f(input, src)); +} + +kernel void bn254_multilinear_extend( + device const ulong *values [[buffer(0)]], + device const ulong *point [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &num_vars [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + F acc = zero_f(); + F one = one_f(); + for (uint i = 0; i < len; i++) { + F weight = one; + for (uint j = 0; j < num_vars; j++) { + F r = load_f(point, num_vars - 1 - j); + if (((i >> j) & 1) != 0) { + weight = mul_f(weight, r); + } else { + weight = mul_f(weight, sub_f(one, r)); + } + } + acc = add_f(acc, mul_f(load_f(values, i), weight)); + } + store_f(out, 0, acc); +} +"#; + +const REDUCTION_CHUNK_SIZE: usize = 64; +const SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 8; +const LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 16; +const LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD: usize = 1 << 18; +const GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE: usize = 16; +const GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD: usize = 1 << 17; +const GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS: usize = 64; + +#[derive(Clone, Debug)] +struct MetalFieldBuffer { + limbs: Buffer, +} + +#[derive(Clone, Debug)] +struct MetalHashBuffer { + bytes: Buffer, +} + +#[derive(Clone, Debug)] +pub struct MetalBuffer { + len: usize, + host_cache: OnceCell>, + field: Option, + hash: Option, + _marker: PhantomData, +} + +impl PartialEq for MetalBuffer +where + T: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for MetalBuffer {} + +impl PartialOrd for MetalBuffer { + fn partial_cmp(&self, other: &Self) -> Option { + self.as_slice().partial_cmp(other.as_slice()) + } +} + +impl Ord for MetalBuffer { + fn cmp(&self, other: &Self) -> Ordering { + self.as_slice().cmp(other.as_slice()) + } +} + +impl Hash for MetalBuffer { + fn hash(&self, state: &mut H) { + self.as_slice().hash(state); + } +} + +impl Default for MetalBuffer { + fn default() -> Self { + Self { + len: 0, + host_cache: OnceCell::new(), + field: None, + hash: None, + _marker: PhantomData, + } + } +} + +impl Serialize for MetalBuffer { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.as_slice().serialize(serializer) + } +} + +impl<'de, T> Deserialize<'de> for MetalBuffer +where + T: Clone + Deserialize<'de>, +{ + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let data = Vec::::deserialize(deserializer)?; + Ok(Self::from_vec(data)) + } +} + +impl MetalBuffer { + pub fn warmup() { + let _ = runtime(); + } + + pub fn from_vec(source: Vec) -> Self { + let len = source.len(); + let field = maybe_upload_bn254(&source); + let host_cache = if field.is_some() { + OnceCell::new() + } else { + OnceCell::from(source) + }; + Self { + len, + host_cache, + field, + hash: None, + _marker: PhantomData, + } + } + + pub fn from_slice(source: &[T]) -> Self { + Self::from_vec(Vec::from(source)) + } + + pub(crate) fn as_slice(&self) -> &[T] { + self.host_cache + .get_or_init(|| self.download_host_cache()) + .as_slice() + } + + pub(crate) fn hash_bn254_rows_sha2(&self, num_cols: usize, out: &mut [Digest]) -> bool { + if type_name::() != type_name::() || self.field.is_none() { + return false; + } + assert_eq!(self.len(), num_cols * out.len()); + let message_size = num_cols * size_of::(); + let encoded = encode_field_rows_le( + &self + .field + .as_ref() + .expect("missing Metal field buffer") + .limbs, + out.len(), + num_cols, + ); + MetalSha2::new().hash_many_buffer(message_size, &encoded, out.len(), out); + true + } + + pub(crate) fn commit_bn254_rows_sha2_merkle( + &self, + num_cols: usize, + num_rows: usize, + layers: usize, + ) -> Option> { + if type_name::() != type_name::() || self.field.is_none() { + return None; + } + if num_rows != (1usize << layers) { + return None; + } + assert_eq!(self.len(), num_cols * num_rows); + let message_size = num_cols * size_of::(); + let encoded = encode_field_rows_le( + &self + .field + .as_ref() + .expect("missing Metal field buffer") + .limbs, + num_rows, + num_cols, + ); + let sha = MetalSha2::new(); + let nodes = sha.build_merkle_tree_buffer_from_messages_buffer( + message_size, + &encoded, + num_rows, + layers, + ); + Some(MetalBuffer::::from_digest_buffer( + nodes, + (1usize << (layers + 1)) - 1, + )) + } +} + +impl MetalBuffer { + pub(crate) fn from_digest_buffer(bytes: Buffer, len: usize) -> Self { + Self { + len, + host_cache: OnceCell::new(), + field: None, + hash: Some(MetalHashBuffer { bytes }), + _marker: PhantomData, + } + } + + pub(crate) fn read_hash_at(&self, index: usize) -> Option { + self.read_hash_indices(&[index]) + .map(|mut values| values.pop().expect("missing hash")) + } + + pub(crate) fn read_hash_indices(&self, indices: &[usize]) -> Option> { + let buffer = self.hash.as_ref()?; + Some(download_hash_indices(&buffer.bytes, self.len, indices)) + } +} + +impl BufferOps for MetalBuffer { + fn as_slice(&self) -> &[T] { + self.as_slice() + } + + fn len(&self) -> usize { + self.len + } + + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { + if type_name::() == type_name::() && self.field.is_some() { + return read_bn254_rows( + self.field.as_ref().expect("missing Metal field buffer"), + num_cols, + indices, + ) + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + let data = self.as_slice(); + let mut result = Vec::with_capacity(indices.len() * num_cols); + for i in indices { + result.extend_from_slice(&data[i * num_cols..(i + 1) * num_cols]); + } + result + } + + fn from_vec(source: Vec) -> Self { + Self::from_vec(source) + } + + fn from_slice(source: &[T]) -> Self { + Self::from_slice(source) + } +} + +impl FieldOps for MetalBuffer { + type TargetBuffer = MetalBuffer; + + fn zeros(length: usize) -> Self { + assert_bn254::(); + Self::from_vec(vec![F::ZERO; length]) + } + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + assert_bn254::(); + Self::from_vec((0..length).map(|_| rng.gen()).collect()) + } + + fn zero_pad(&mut self) { + assert_bn254::(); + if !self.is_empty() { + let mut data = self.as_slice().to_vec(); + data.resize(self.len().next_power_of_two(), F::ZERO); + *self = Self::from_vec(data); + } + } + + fn dot(&self, other: &Self) -> F { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + let this = self.bn254_buffer(); + let other = other.bn254_buffer(); + field256_to_f::(parallel_dot(&this, &other, self.len())) + } + + fn fold(&mut self, weight: F) { + assert_bn254::(); + if self.len() <= 1 { + return; + } + let len = self.len(); + let fold_half = len.next_power_of_two() >> 1; + let weight = upload_field(&[f_to_field256(weight)]); + let field = self.bn254_buffer(); + run_in_place( + "bn254_fold", + &[&field.limbs, &weight.limbs], + &[len as u32, fold_half as u32], + fold_half, + ); + self.len = fold_half; + self.invalidate_host_cache(); + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + assert_bn254::(); + let len = self.len().min(other.len()); + if len == 0 { + return (F::ZERO, F::ZERO); + } + let fold_half = len.next_power_of_two() >> 1; + let this = self.bn254_buffer(); + let other = other.bn254_buffer(); + let (c0, c2) = parallel_sumcheck(&this, &other, len, fold_half); + (field256_to_f::(c0), field256_to_f::(c2)) + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + assert_bn254::(); + assert_eq!(evaluators.len(), scalars.len()); + let Some(size) = evaluators.first().map(|e| e.size) else { + return; + }; + assert_eq!(self.len(), size); + for evaluator in evaluators { + assert_eq!(evaluator.size, size); + } + let points = evaluators + .iter() + .map(|e| f_to_field256(e.point)) + .collect::>(); + let scalars = scalars + .iter() + .copied() + .map(f_to_field256) + .collect::>(); + let points = upload_field(&points); + let scalars = upload_field(&scalars); + let field = self.bn254_buffer(); + let chunk_size = geometric_accumulate_chunk_size(self.len()); + if std::env::var_os("WHIR_METAL_TRACE").is_some() { + eprintln!( + "metal geometric shape len={} points={} chunk={} chunks={}", + self.len(), + evaluators.len(), + chunk_size, + self.len().div_ceil(chunk_size) + ); + } + if chunk_size <= 1 { + run_in_place( + "bn254_geometric_accumulate", + &[&field.limbs, &points.limbs, &scalars.limbs], + &[self.len() as u32, evaluators.len() as u32], + self.len(), + ); + } else { + let point_steps = evaluators + .iter() + .map(|e| f_to_field256(e.point.pow([chunk_size as u64]))) + .collect::>(); + let point_steps = upload_field(&point_steps); + if should_use_geometric_point_blocks(self.len(), evaluators.len(), chunk_size) { + parallel_geometric_accumulate_point_blocks( + &field, + &points, + &point_steps, + &scalars, + self.len(), + evaluators.len(), + chunk_size, + ); + } else { + run_in_place( + "bn254_geometric_accumulate_chunks_strided", + &[ + &field.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[ + self.len() as u32, + evaluators.len() as u32, + chunk_size as u32, + ], + self.len().div_ceil(chunk_size), + ); + } + } + self.invalidate_host_cache(); + } + + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self { + assert_bn254::(); + assert_eq!(linear_forms.len(), rlc_coeffs.len()); + let Some((first, rest)) = linear_forms.split_first_mut() else { + return Self::zeros(size); + }; + let first = (first.as_mut() as &mut dyn std::any::Any) + .downcast_mut::>() + .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); + let mut accumulator = Self::from_slice(&first.vector); + for (coeff, linear_form) in rlc_coeffs[1..].iter().zip(rest) { + let covector = (linear_form.as_mut() as &mut dyn std::any::Any) + .downcast_mut::>() + .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); + let vector = Self::from_slice(&covector.vector); + vector.mixed_scalar_mul_add_to(&Identity::new(), &mut accumulator, *coeff); + } + accumulator + } + + fn mixed_extend, T: Field>( + &self, + _embedding: &M, + point: &[M::Target], + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let num_vars = point.len(); + let point = point + .iter() + .copied() + .map(target_to_field256) + .collect::>(); + let point = upload_field(&point); + let this = self.bn254_buffer(); + let out = run_reduce( + "bn254_multilinear_extend", + &[&this.limbs, &point.limbs], + &[self.len() as u32, num_vars as u32], + 1, + ); + field256_to_target::(out[0]) + } + + fn mixed_dot, T: Field>( + &self, + _embedding: &M, + other: &Self::TargetBuffer, + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let this = self.bn254_buffer(); + let other = other.bn254_buffer_target(); + let value = field256_to_f::(parallel_dot(&this, &other, self.len())); + field256_to_target::(f_to_field256(value)) + } + + fn mixed_univariate_evaluate>( + &self, + _embedding: &M, + point: M::Target, + ) -> M::Target { + assert_bn254::(); + let point = target_to_field256(point); + let point = upload_field(&[point]); + let this = self.bn254_buffer(); + field256_to_target::(parallel_univariate_evaluate(&this, &point, self.len())) + } + + fn mixed_linear_combination>( + _embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer { + assert_bn254::(); + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return MetalBuffer::from_vec(Vec::new()); + }; + let mut accumulator = MetalBuffer:: { + len: first.len(), + host_cache: OnceCell::new(), + field: Some(copy_field_buffer(&first.bn254_buffer(), first.len())), + hash: None, + _marker: PhantomData, + }; + for (coeff, vector) in coeffs[1..].iter().copied().zip(vectors) { + vector.mixed_scalar_mul_add_to(_embedding, &mut accumulator, coeff); + } + accumulator + } + + fn mixed_scalar_mul_add_to>( + &self, + _embedding: &M, + accumulator: &mut Self::TargetBuffer, + weight: M::Target, + ) { + assert_bn254::(); + let weight = upload_field(&[target_to_field256(weight)]); + let vector = self.bn254_buffer(); + let acc = accumulator.bn254_buffer_target(); + run_in_place( + "bn254_scalar_mul_add", + &[&acc.limbs, &vector.limbs, &weight.limbs], + &[self.len() as u32], + self.len(), + ); + accumulator.invalidate_host_cache(); + } + + fn interleaved_rs_encode( + vectors: &[&Self], + masks: &Self, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + ) -> Self { + assert_bn254::(); + let num_messages = vectors.len() * interleaving_depth; + if num_messages == 0 { + return Self::from_vec(Vec::new()); + } + assert!(masks.len().is_multiple_of(num_messages)); + let mask_length = masks.len() / num_messages; + if vectors.len() == 1 && mask_length == 0 { + return encode_single_vector_coset_ntt( + vectors[0], + message_length, + interleaving_depth, + codeword_length, + ); + } + + panic!("MetalBuffer BN254 RS encoding supports only one unmasked vector") + } +} + +impl MetalBuffer { + fn bn254_buffer(&self) -> MetalFieldBuffer { + self.field + .clone() + .unwrap_or_else(|| upload_field(&to_field256_slice(self.as_slice()))) + } + + fn invalidate_host_cache(&mut self) { + let _ = self.host_cache.take(); + } + + fn download_host_cache(&self) -> Vec { + if self.field.is_some() && type_name::() == type_name::() { + return download_field( + &self + .field + .as_ref() + .expect("missing Metal field buffer") + .limbs, + self.len, + ) + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + if self.hash.is_some() && type_name::() == type_name::() { + return download_hash_indices( + &self.hash.as_ref().expect("missing Metal hash buffer").bytes, + self.len, + &(0..self.len).collect::>(), + ) + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + panic!( + "MetalBuffer<{}> has no host cache and cannot be materialized", + type_name::() + ); + } +} + +impl MetalBuffer { + fn bn254_buffer_target(&self) -> MetalFieldBuffer { + self.field + .clone() + .unwrap_or_else(|| upload_field(&to_field256_slice(self.as_slice()))) + } +} + +struct MetalRuntime { + device: Device, + queue: CommandQueue, + fold: ComputePipelineState, + scalar_mul_add: ComputePipelineState, + dot: ComputePipelineState, + sumcheck: ComputePipelineState, + dot_chunks: ComputePipelineState, + sum_chunks: ComputePipelineState, + sumcheck_chunks: ComputePipelineState, + geometric_accumulate: ComputePipelineState, + geometric_accumulate_chunks: ComputePipelineState, + geometric_accumulate_chunks_strided: ComputePipelineState, + geometric_accumulate_point_blocks: ComputePipelineState, + geometric_accumulate_reduce_point_blocks: ComputePipelineState, + univariate_evaluate: ComputePipelineState, + univariate_eval_chunks: ComputePipelineState, + interleaved_rs_encode: ComputePipelineState, + interleaved_rs_encode_single_vector: ComputePipelineState, + multilinear_extend: ComputePipelineState, + pack_single_vector_cosets: ComputePipelineState, + apply_coset_twiddles: ComputePipelineState, + replicate_first_coset: ComputePipelineState, + bit_reverse_rows: ComputePipelineState, + ntt_stage_rows: ComputePipelineState, + transpose: ComputePipelineState, + transpose_reverse_rows: ComputePipelineState, + encode_field_rows_le: ComputePipelineState, + read_rows: ComputePipelineState, + ntt_roots: Mutex>, + root_powers: Mutex>, +} + +fn runtime() -> &'static MetalRuntime { + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + let device = Device::system_default().expect("Metal device is not available"); + let library = device + .new_library_with_source(METAL_SOURCE, &CompileOptions::new()) + .expect("failed to compile Metal BN254 kernels"); + let pipeline = |name: &str| { + let function = library + .get_function(name, None) + .unwrap_or_else(|_| panic!("missing Metal kernel {name}")); + device + .new_compute_pipeline_state_with_function(&function) + .unwrap_or_else(|err| panic!("failed to compile Metal kernel {name}: {err}")) + }; + MetalRuntime { + queue: device.new_command_queue(), + fold: pipeline("bn254_fold"), + scalar_mul_add: pipeline("bn254_scalar_mul_add"), + dot: pipeline("bn254_dot"), + sumcheck: pipeline("bn254_sumcheck"), + dot_chunks: pipeline("bn254_dot_chunks"), + sum_chunks: pipeline("bn254_sum_chunks"), + sumcheck_chunks: pipeline("bn254_sumcheck_chunks"), + geometric_accumulate: pipeline("bn254_geometric_accumulate"), + geometric_accumulate_chunks: pipeline("bn254_geometric_accumulate_chunks"), + geometric_accumulate_chunks_strided: pipeline( + "bn254_geometric_accumulate_chunks_strided", + ), + geometric_accumulate_point_blocks: pipeline("bn254_geometric_accumulate_point_blocks"), + geometric_accumulate_reduce_point_blocks: pipeline( + "bn254_geometric_accumulate_reduce_point_blocks", + ), + univariate_evaluate: pipeline("bn254_univariate_evaluate"), + univariate_eval_chunks: pipeline("bn254_univariate_eval_chunks"), + interleaved_rs_encode: pipeline("bn254_interleaved_rs_encode"), + interleaved_rs_encode_single_vector: pipeline( + "bn254_interleaved_rs_encode_single_vector", + ), + multilinear_extend: pipeline("bn254_multilinear_extend"), + pack_single_vector_cosets: pipeline("bn254_pack_single_vector_cosets"), + apply_coset_twiddles: pipeline("bn254_apply_coset_twiddles"), + replicate_first_coset: pipeline("bn254_replicate_first_coset"), + bit_reverse_rows: pipeline("bn254_bit_reverse_permute_rows_in_place"), + ntt_stage_rows: pipeline("bn254_radix2_ntt_stage_rows_in_place"), + transpose: pipeline("bn254_transpose_matrix"), + transpose_reverse_rows: pipeline("bn254_transpose_matrix_reverse_rows"), + encode_field_rows_le: pipeline("bn254_encode_field_rows_le"), + read_rows: pipeline("bn254_read_rows"), + ntt_roots: Mutex::new(HashMap::new()), + root_powers: Mutex::new(HashMap::new()), + device, + } + }) +} + +fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { + match name { + "bn254_fold" => &rt.fold, + "bn254_scalar_mul_add" => &rt.scalar_mul_add, + "bn254_dot" => &rt.dot, + "bn254_sumcheck" => &rt.sumcheck, + "bn254_dot_chunks" => &rt.dot_chunks, + "bn254_sum_chunks" => &rt.sum_chunks, + "bn254_sumcheck_chunks" => &rt.sumcheck_chunks, + "bn254_geometric_accumulate" => &rt.geometric_accumulate, + "bn254_geometric_accumulate_chunks" => &rt.geometric_accumulate_chunks, + "bn254_geometric_accumulate_chunks_strided" => &rt.geometric_accumulate_chunks_strided, + "bn254_geometric_accumulate_point_blocks" => &rt.geometric_accumulate_point_blocks, + "bn254_geometric_accumulate_reduce_point_blocks" => { + &rt.geometric_accumulate_reduce_point_blocks + } + "bn254_univariate_evaluate" => &rt.univariate_evaluate, + "bn254_univariate_eval_chunks" => &rt.univariate_eval_chunks, + "bn254_interleaved_rs_encode" => &rt.interleaved_rs_encode, + "bn254_interleaved_rs_encode_single_vector" => &rt.interleaved_rs_encode_single_vector, + "bn254_multilinear_extend" => &rt.multilinear_extend, + "bn254_pack_single_vector_cosets" => &rt.pack_single_vector_cosets, + "bn254_apply_coset_twiddles" => &rt.apply_coset_twiddles, + "bn254_replicate_first_coset" => &rt.replicate_first_coset, + "bn254_bit_reverse_permute_rows_in_place" => &rt.bit_reverse_rows, + "bn254_radix2_ntt_stage_rows_in_place" => &rt.ntt_stage_rows, + "bn254_transpose_matrix" => &rt.transpose, + "bn254_transpose_matrix_reverse_rows" => &rt.transpose_reverse_rows, + "bn254_encode_field_rows_le" => &rt.encode_field_rows_le, + "bn254_read_rows" => &rt.read_rows, + _ => panic!("unknown Metal kernel {name}"), + } +} + +fn new_shared_buffer(rt: &MetalRuntime, bytes: u64) -> Buffer { + metal_profile::record_alloc(bytes); + rt.device + .new_buffer(bytes, MTLResourceOptions::StorageModeShared) +} + +fn new_shared_buffer_with_data(rt: &MetalRuntime, data: *const c_void, bytes: u64) -> Buffer { + let start = Instant::now(); + let buffer = rt + .device + .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); + metal_profile::record_alloc(bytes); + metal_profile::record_upload(bytes, start.elapsed()); + buffer +} + +fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { + command.commit(); + let start = Instant::now(); + command.wait_until_completed(); + let elapsed = start.elapsed(); + if std::env::var_os("WHIR_METAL_TRACE").is_some() { + eprintln!( + "metal command {label} {:.3} ms", + elapsed.as_secs_f64() * 1_000.0 + ); + } + metal_profile::record_command_wait(elapsed); +} + +fn wait_for_blit(command: &metal::CommandBufferRef, bytes: u64) { + command.commit(); + let start = Instant::now(); + command.wait_until_completed(); + metal_profile::record_blit(bytes, start.elapsed()); +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct BitReverseParams { + row_len: u32, + log_n: u32, + total_elements: u32, + _pad0: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct NttStageParams { + row_len: u32, + half_m: u32, + twiddle_offset: u32, + _pad0: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct TransposeParams { + rows: u32, + cols: u32, + total_elements: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct ReplicateCosetsParams { + row_len: u32, + coset_size: u32, + trailing_elements: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct PackSingleVectorParams { + row_count: u32, + message_length: u32, + codeword_length: u32, + coset_size: u32, + total_elements: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct FieldBytesParams { + rows: u32, + cols: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct ApplyCosetTwiddlesParams { + row_count: u32, + num_cosets: u32, + coset_size: u32, + codeword_length: u32, + total_elements: u32, +} + +fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field256 { + if len == 0 { + return Field256::ZERO; + } + let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), + }; + run_in_place( + "bn254_dot_chunks", + &[&a.limbs, &b.limbs, &partials.limbs], + &[len as u32, REDUCTION_CHUNK_SIZE as u32], + partial_count, + ); + reduce_field_buffer(&partials, partial_count, 0) +} + +fn parallel_sumcheck( + a: &MetalFieldBuffer, + b: &MetalFieldBuffer, + len: usize, + fold_half: usize, +) -> (Field256, Field256) { + let partial_count = fold_half.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), + }; + run_in_place( + "bn254_sumcheck_chunks", + &[&a.limbs, &b.limbs, &partials.limbs], + &[ + len as u32, + fold_half as u32, + REDUCTION_CHUNK_SIZE as u32, + partial_count as u32, + ], + partial_count, + ); + ( + reduce_field_buffer(&partials, partial_count, 0), + reduce_field_buffer(&partials, partial_count, partial_count), + ) +} + +fn parallel_univariate_evaluate( + coeffs: &MetalFieldBuffer, + point: &MetalFieldBuffer, + len: usize, +) -> Field256 { + if len == 0 { + return Field256::ZERO; + } + let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), + }; + run_in_place( + "bn254_univariate_eval_chunks", + &[&coeffs.limbs, &point.limbs, &partials.limbs], + &[len as u32, REDUCTION_CHUNK_SIZE as u32], + partial_count, + ); + reduce_field_buffer(&partials, partial_count, 0) +} + +fn parallel_geometric_accumulate_point_blocks( + acc: &MetalFieldBuffer, + points: &MetalFieldBuffer, + point_steps: &MetalFieldBuffer, + scalars: &MetalFieldBuffer, + len: usize, + num_points: usize, + chunk_size: usize, +) { + assert!(chunk_size <= 16); + let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); + let chunks = len.div_ceil(chunk_size); + let partial_len = point_blocks + .checked_mul(len) + .expect("Metal geometric partial size overflow"); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), + }; + run_in_place( + "bn254_geometric_accumulate_point_blocks", + &[ + &partials.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[ + len as u32, + num_points as u32, + chunk_size as u32, + GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, + point_blocks as u32, + ], + chunks * point_blocks, + ); + run_in_place( + "bn254_geometric_accumulate_reduce_point_blocks", + &[&acc.limbs, &partials.limbs], + &[len as u32, point_blocks as u32], + len, + ); +} + +fn geometric_accumulate_chunk_size(len: usize) -> usize { + std::env::var("WHIR_METAL_GEOM_CHUNK") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&value| value > 0) + .unwrap_or(if len >= LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD { + LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + } else { + SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + }) +} + +fn should_use_geometric_point_blocks(len: usize, num_points: usize, chunk_size: usize) -> bool { + chunk_size <= 16 + && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS + && len <= GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD +} + +fn reduce_field_buffer(input: &MetalFieldBuffer, len: usize, offset: usize) -> Field256 { + if len == 0 { + return Field256::ZERO; + } + let mut current = input.clone(); + let mut current_len = len; + let mut current_offset = offset; + while current_len > 1 { + let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let next = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (next_len * size_of::()) as u64), + }; + run_in_place( + "bn254_sum_chunks", + &[¤t.limbs, &next.limbs], + &[ + current_len as u32, + current_offset as u32, + REDUCTION_CHUNK_SIZE as u32, + ], + next_len, + ); + current = next; + current_len = next_len; + current_offset = 0; + } + download_field_at(¤t.limbs, current_offset) +} + +fn encode_single_vector_coset_ntt( + vector: &MetalBuffer, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, +) -> MetalBuffer { + assert!(codeword_length.is_power_of_two()); + assert!(Field256::get_root_of_unity(codeword_length as u64).is_some()); + assert_eq!(vector.len(), message_length * interleaving_depth); + + let coset_size = message_length.next_power_of_two(); + assert!(codeword_length.is_multiple_of(coset_size)); + let num_cosets = codeword_length / coset_size; + let total_elements = interleaving_depth + .checked_mul(codeword_length) + .expect("Metal RS encode size overflow"); + assert!(total_elements <= u32::MAX as usize); + + let rt = runtime(); + let source = vector.bn254_buffer(); + let current = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); + let transposed = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); + let codeword_root_powers = root_powers_buffer(codeword_length); + let coset_roots = roots_buffer(coset_size); + + let command = rt.queue.new_command_buffer(); + + encode_kernel( + &command, + &rt.pack_single_vector_cosets, + &[&source.limbs, ¤t], + &PackSingleVectorParams { + row_count: interleaving_depth as u32, + message_length: message_length as u32, + codeword_length: codeword_length as u32, + coset_size: coset_size as u32, + total_elements: total_elements as u32, + }, + total_elements, + ); + + let trailing_elements = interleaving_depth.saturating_mul(codeword_length - coset_size); + if trailing_elements != 0 { + encode_kernel( + &command, + &rt.replicate_first_coset, + &[¤t], + &ReplicateCosetsParams { + row_len: codeword_length as u32, + coset_size: coset_size as u32, + trailing_elements: trailing_elements as u32, + }, + trailing_elements, + ); + } + + encode_kernel( + &command, + &rt.apply_coset_twiddles, + &[¤t, &codeword_root_powers.limbs], + &ApplyCosetTwiddlesParams { + row_count: interleaving_depth as u32, + num_cosets: num_cosets as u32, + coset_size: coset_size as u32, + codeword_length: codeword_length as u32, + total_elements: total_elements as u32, + }, + total_elements, + ); + + let stage_count = coset_size.trailing_zeros() as usize; + encode_kernel( + &command, + &rt.bit_reverse_rows, + &[¤t], + &BitReverseParams { + row_len: coset_size as u32, + log_n: stage_count as u32, + total_elements: total_elements as u32, + _pad0: 0, + }, + total_elements, + ); + + let total_butterflies = total_elements / 2; + let mut twiddle_offset = 0usize; + for stage in 0..stage_count { + let half_m = 1usize << stage; + encode_kernel( + &command, + &rt.ntt_stage_rows, + &[¤t, &coset_roots.limbs], + &NttStageParams { + row_len: coset_size as u32, + half_m: half_m as u32, + twiddle_offset: twiddle_offset as u32, + _pad0: 0, + }, + total_butterflies, + ); + twiddle_offset += 1usize << stage; + } + + encode_kernel( + &command, + &rt.transpose, + &[¤t, &transposed], + &TransposeParams { + rows: interleaving_depth as u32, + cols: codeword_length as u32, + total_elements: total_elements as u32, + }, + total_elements, + ); + + wait_for_command_named(&command, "bn254_rs_encode"); + + MetalBuffer { + len: total_elements, + host_cache: OnceCell::new(), + field: Some(MetalFieldBuffer { limbs: transposed }), + hash: None, + _marker: PhantomData, + } +} + +fn encode_kernel( + command: &metal::CommandBufferRef, + pipeline: &ComputePipelineState, + buffers: &[&Buffer], + params: &P, + threads: usize, +) { + let encoder = command.new_compute_command_encoder(); + encoder.set_compute_pipeline_state(pipeline); + for (index, buffer) in buffers.iter().enumerate() { + encoder.set_buffer(index as u64, Some(buffer), 0); + } + encoder.set_bytes( + buffers.len() as u64, + size_of::

() as u64, + (params as *const P).cast::(), + ); + dispatch(&encoder, pipeline, threads.max(1)); + encoder.end_encoding(); +} + +fn roots_buffer(codeword_length: usize) -> MetalFieldBuffer { + let rt = runtime(); + if let Some(buffer) = rt.ntt_roots.lock().unwrap().get(&codeword_length).cloned() { + return buffer; + } + let root = Field256::get_root_of_unity(codeword_length as u64) + .expect("BN254 root of unity unavailable for Metal NTT"); + let stage_count = codeword_length.trailing_zeros() as usize; + let mut roots = Vec::with_capacity(codeword_length.saturating_sub(1)); + for stage in 0..stage_count { + let stage_size = 1usize << (stage + 1); + let half_stage = stage_size >> 1; + let stage_root = root.pow([(codeword_length / stage_size) as u64]); + let mut current = Field256::ONE; + for _ in 0..half_stage { + roots.push(current); + current *= stage_root; + } + } + let buffer = upload_field(&roots); + rt.ntt_roots + .lock() + .unwrap() + .insert(codeword_length, buffer.clone()); + buffer +} + +fn root_powers_buffer(codeword_length: usize) -> MetalFieldBuffer { + let rt = runtime(); + if let Some(buffer) = rt + .root_powers + .lock() + .unwrap() + .get(&codeword_length) + .cloned() + { + return buffer; + } + let root = Field256::get_root_of_unity(codeword_length as u64) + .expect("BN254 root of unity unavailable for Metal RS twiddles"); + let mut powers = Vec::with_capacity(codeword_length); + let mut current = Field256::ONE; + for _ in 0..codeword_length { + powers.push(current); + current *= root; + } + let buffer = upload_field(&powers); + rt.root_powers + .lock() + .unwrap() + .insert(codeword_length, buffer.clone()); + buffer +} + +fn encode_field_rows_le(input: &Buffer, rows: usize, cols: usize) -> Buffer { + assert!(rows <= u32::MAX as usize); + assert!(cols <= u32::MAX as usize); + let rt = runtime(); + let output = new_shared_buffer(rt, (rows * cols * size_of::()) as u64); + let command = rt.queue.new_command_buffer(); + encode_kernel( + &command, + &rt.encode_field_rows_le, + &[input, &output], + &FieldBytesParams { + rows: rows as u32, + cols: cols as u32, + }, + rows * cols, + ); + wait_for_command_named(&command, "bn254_encode_field_rows_le"); + output +} + +fn read_bn254_rows(source: &MetalFieldBuffer, num_cols: usize, indices: &[usize]) -> Vec { + if indices.is_empty() || num_cols == 0 { + return Vec::new(); + } + assert!(num_cols <= u32::MAX as usize); + assert!(indices.iter().all(|&index| index <= u32::MAX as usize)); + let total = indices + .len() + .checked_mul(num_cols) + .expect("Metal read_rows size overflow"); + assert!(total <= u32::MAX as usize); + + let rt = runtime(); + let indices = indices + .iter() + .copied() + .map(|index| index as u32) + .collect::>(); + let indices_buffer = new_shared_buffer_with_data( + rt, + indices.as_ptr().cast(), + (indices.len() * size_of::()) as u64, + ); + let out = new_shared_buffer(rt, (total * size_of::()) as u64); + run_in_place( + "bn254_read_rows", + &[&source.limbs, &indices_buffer, &out], + &[num_cols as u32, total as u32], + total, + ); + download_field(&out, total) +} + +fn upload_field(values: &[Field256]) -> MetalFieldBuffer { + let limbs = values + .iter() + .flat_map(|value| value.0 .0) + .collect::>(); + let rt = runtime(); + let buffer = if limbs.is_empty() { + new_shared_buffer(rt, 0) + } else { + new_shared_buffer_with_data( + rt, + limbs.as_ptr().cast(), + (limbs.len() * size_of::()) as u64, + ) + }; + MetalFieldBuffer { limbs: buffer } +} + +fn maybe_upload_bn254(values: &[T]) -> Option { + (type_name::() == type_name::()).then(|| upload_field(&to_field256_slice(values))) +} + +fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer { + let rt = runtime(); + let byte_len = (len * 4 * size_of::()) as u64; + let target = new_shared_buffer(rt, byte_len); + let command = rt.queue.new_command_buffer(); + let blit = command.new_blit_command_encoder(); + blit.copy_from_buffer(&source.limbs, 0, &target, 0, byte_len); + blit.end_encoding(); + wait_for_blit(&command, byte_len); + MetalFieldBuffer { limbs: target } +} + +fn upload_u32(value: u32) -> Buffer { + let rt = runtime(); + new_shared_buffer_with_data(rt, (&value as *const u32).cast(), size_of::() as u64) +} + +fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + let encoder = command.new_compute_command_encoder(); + let pipeline = pipeline(rt, name); + encoder.set_compute_pipeline_state(pipeline); + let mut index = 0; + for buffer in buffers { + encoder.set_buffer(index, Some(buffer), 0); + index += 1; + } + let constant_buffers = constants + .iter() + .copied() + .map(upload_u32) + .collect::>(); + for buffer in &constant_buffers { + encoder.set_buffer(index, Some(buffer), 0); + index += 1; + } + dispatch(&encoder, pipeline, threads.max(1)); + encoder.end_encoding(); + wait_for_command_named(&command, name); +} + +fn run_reduce(name: &str, buffers: &[&Buffer], constants: &[u32], out_len: usize) -> Vec { + let rt = runtime(); + let out = new_shared_buffer(rt, (out_len * 4 * size_of::()) as u64); + let mut all = buffers.to_vec(); + all.push(&out); + run_in_place(name, &all, constants, 1); + download_field(&out, out_len) +} + +fn dispatch( + encoder: &metal::ComputeCommandEncoderRef, + pipeline: &ComputePipelineState, + threads: usize, +) { + let width = pipeline.thread_execution_width().max(1); + let group_width = width.min(threads as u64).max(1); + encoder.dispatch_threads( + MTLSize { + width: threads as u64, + height: 1, + depth: 1, + }, + MTLSize { + width: group_width, + height: 1, + depth: 1, + }, + ); +} + +fn download_field(buffer: &Buffer, len: usize) -> Vec { + if len == 0 { + return Vec::new(); + } + let start = Instant::now(); + let limbs = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 4) }; + let result = limbs + .chunks_exact(4) + .map(|chunk| { + Fp::, 4>( + BigInt([chunk[0], chunk[1], chunk[2], chunk[3]]), + PhantomData, + ) + }) + .collect(); + metal_profile::record_readback((len * size_of::()) as u64, start.elapsed()); + result +} + +fn download_field_at(buffer: &Buffer, index: usize) -> Field256 { + let start = Instant::now(); + let limbs = + unsafe { std::slice::from_raw_parts(buffer.contents().cast::().add(index * 4), 4) }; + let result = Fp::, 4>( + BigInt([limbs[0], limbs[1], limbs[2], limbs[3]]), + PhantomData, + ); + metal_profile::record_readback(size_of::() as u64, start.elapsed()); + result +} + +fn download_hash_indices(buffer: &Buffer, len: usize, indices: &[usize]) -> Vec { + let start = Instant::now(); + let bytes = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 32) }; + let mut result = Vec::with_capacity(indices.len()); + for &index in indices { + assert!(index < len, "Metal hash index out of bounds"); + let mut hash = Digest::default(); + hash.as_mut_bytes() + .copy_from_slice(&bytes[index * 32..(index + 1) * 32]); + result.push(hash); + } + metal_profile::record_readback( + (indices.len() * size_of::()) as u64, + start.elapsed(), + ); + result +} + +fn assert_bn254() { + assert_eq!( + type_name::(), + type_name::(), + "MetalBuffer only supports BN254 Field256 field operations" + ); +} + +fn f_to_field256(value: F) -> Field256 { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +fn field256_to_f(value: Field256) -> F { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +fn target_to_field256(value: T) -> Field256 { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +fn field256_to_target(value: Field256) -> T { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +fn to_field256_slice(values: &[T]) -> Vec { + assert_eq!( + type_name::(), + type_name::(), + "MetalBuffer only supports BN254 Field256 buffers" + ); + values + .iter() + .map(|value| unsafe { std::mem::transmute_copy(value) }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::algebra::buffer::CpuBuffer; + + fn values(len: usize, offset: u64) -> Vec { + (0..len) + .map(|i| Field256::from(i as u64 + offset)) + .collect() + } + + #[test] + fn metal_bn254_dot_matches_cpu() { + let a = values(33, 1); + let b = values(33, 9); + let cpu_a = CpuBuffer::from_slice(&a); + let cpu_b = CpuBuffer::from_slice(&b); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_b = MetalBuffer::from_slice(&b); + assert_eq!(gpu_a.dot(&gpu_b), cpu_a.dot(&cpu_b)); + } + + #[test] + fn metal_bn254_fold_matches_cpu() { + let mut cpu = CpuBuffer::from_vec(values(31, 2)); + let mut gpu = MetalBuffer::from_vec(values(31, 2)); + let weight = Field256::from(42); + cpu.fold(weight); + gpu.fold(weight); + assert_eq!(gpu.as_slice(), cpu.as_slice()); + } + + #[test] + fn metal_bn254_sumcheck_matches_cpu() { + let a = values(27, 3); + let b = values(27, 11); + let cpu_a = CpuBuffer::from_slice(&a); + let cpu_b = CpuBuffer::from_slice(&b); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_b = MetalBuffer::from_slice(&b); + assert_eq!( + gpu_a.sumcheck_polynomial(&gpu_b), + cpu_a.sumcheck_polynomial(&cpu_b) + ); + } + + #[test] + fn metal_bn254_scalar_mul_add_matches_cpu() { + let mut cpu = CpuBuffer::from_vec(values(19, 1)); + let mut gpu = MetalBuffer::from_vec(values(19, 1)); + let vector = values(19, 5); + let cpu_vector = CpuBuffer::from_slice(&vector); + let gpu_vector = MetalBuffer::from_slice(&vector); + let weight = Field256::from(7); + cpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut cpu, weight); + gpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut gpu, weight); + assert_eq!(gpu.as_slice(), cpu.as_slice()); + } + + #[test] + fn metal_bn254_interleaved_rs_encode_matches_cpu() { + let a = values(8, 1); + let cpu_a = CpuBuffer::from_slice(&a); + let cpu_masks = CpuBuffer::from_slice(&[]); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_masks = MetalBuffer::from_slice(&[]); + let cpu = CpuBuffer::interleaved_rs_encode(&[&cpu_a], &cpu_masks, 4, 2, 8); + let gpu = MetalBuffer::interleaved_rs_encode(&[&gpu_a], &gpu_masks, 4, 2, 8); + assert_eq!(gpu.as_slice(), cpu.as_slice()); + } + + #[test] + fn metal_bn254_mixed_extend_matches_cpu() { + let values = values(8, 3); + let point = vec![Field256::from(2), Field256::from(5), Field256::from(9)]; + let cpu = CpuBuffer::from_slice(&values); + let gpu = MetalBuffer::from_slice(&values); + assert_eq!( + gpu.mixed_extend(&Identity::new(), &point), + cpu.mixed_extend(&Identity::new(), &point) + ); + } +} diff --git a/src/algebra/mod.rs b/src/algebra/mod.rs index 4d17dc3b..592a9109 100644 --- a/src/algebra/mod.rs +++ b/src/algebra/mod.rs @@ -2,6 +2,8 @@ pub mod buffer; pub mod embedding; pub mod fields; pub mod linear_form; +#[cfg(all(feature = "metal", target_os = "macos"))] +mod metal_buffer; mod multilinear; pub mod ntt; pub mod sumcheck; diff --git a/src/bin/gpu_proof_cpu_verify.rs b/src/bin/gpu_proof_cpu_verify.rs new file mode 100644 index 00000000..3a72cf66 --- /dev/null +++ b/src/bin/gpu_proof_cpu_verify.rs @@ -0,0 +1,189 @@ +use std::{borrow::Cow, fs, path::PathBuf}; + +use clap::{Parser, Subcommand}; +use serde::{Deserialize, Serialize}; +use whir::{ + algebra::{ + buffer::ActiveBuffer, embedding::Identity, fields::Field256, linear_form::LinearForm, + }, + hash, + parameters::ProtocolParameters, + protocols::whir::Config as WhirConfig, + transcript::{codecs::Empty, DomainSeparator, Proof, ProverState, VerifierState}, +}; + +type F = Field256; +type M = Identity; + +#[derive(Parser, Debug)] +#[command( + author, + version, + about = "Produce a WHIR proof and verify it from a serialized artifact" +)] +struct Args { + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand, Debug)] +enum Command { + Prove(RoundtripArgs), + Verify { + #[arg(long)] + input: PathBuf, + }, +} + +#[derive(Parser, Debug)] +struct RoundtripArgs { + #[arg(long)] + output: PathBuf, + + #[arg(long, default_value_t = 16)] + log_size: usize, + + #[arg(long, default_value_t = 4)] + fold: usize, + + #[arg(long, default_value_t = 1)] + rate: usize, + + #[arg(long, default_value_t = 20)] + pow_bits: usize, + + #[arg(long, default_value_t = 128)] + security_level: usize, +} + +#[derive(Serialize, Deserialize)] +struct Artifact { + log_size: usize, + fold: usize, + rate: usize, + pow_bits: usize, + security_level: usize, + proof: Proof, + proof_bytes: usize, +} + +fn main() { + let args = Args::parse(); + match args.command { + Command::Prove(args) => prove(args), + Command::Verify { input } => verify(input), + } +} + +fn prove(args: RoundtripArgs) { + assert!(args.fold <= args.log_size, "fold must be <= log_size"); + let size = 1usize << args.log_size; + let whir_params = protocol_parameters(args.security_level, args.pow_bits, args.fold, args.rate); + let params = WhirConfig::::new(size, &whir_params); + let ds = DomainSeparator::protocol(¶ms) + .session(&"gpu proof cpu verify") + .instance(&Empty); + + let vector = input_vector(size); + let vector_buffer = ActiveBuffer::from_slice(&vector); + + let mut prover_state = ProverState::new_std(&ds); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); + let _ = params.prove( + &mut prover_state, + &[&vector_buffer], + vec![&witness], + vec![], + Cow::Owned(vec![]), + ); + let proof = prover_state.proof(); + let proof_bytes = proof.narg_string.len() + proof.hints.len(); + let artifact = Artifact { + log_size: args.log_size, + fold: args.fold, + rate: args.rate, + pow_bits: args.pow_bits, + security_level: args.security_level, + proof, + proof_bytes, + }; + let encoded = serde_json::to_vec_pretty(&artifact).expect("serialize proof artifact"); + fs::write(&args.output, encoded).expect("write proof artifact"); + println!( + "wrote proof artifact backend={} log_size={} fold={} rate={} pow_bits={} proof_bytes={}", + backend_name(), + args.log_size, + args.fold, + args.rate, + args.pow_bits, + proof_bytes + ); +} + +fn verify(input: PathBuf) { + let bytes = fs::read(&input).expect("read proof artifact"); + let artifact: Artifact = serde_json::from_slice(&bytes).expect("decode proof artifact"); + let size = 1usize << artifact.log_size; + let whir_params = protocol_parameters( + artifact.security_level, + artifact.pow_bits, + artifact.fold, + artifact.rate, + ); + let params = WhirConfig::::new(size, &whir_params); + let ds = DomainSeparator::protocol(¶ms) + .session(&"gpu proof cpu verify") + .instance(&Empty); + + let mut verifier_state = VerifierState::new_std(&ds, &artifact.proof); + let commitment = params + .receive_commitment(&mut verifier_state) + .expect("receive commitment"); + let final_claim = params + .verify(&mut verifier_state, &[&commitment], &[]) + .expect("verify WHIR proof"); + let no_forms: Vec<&dyn LinearForm> = Vec::new(); + final_claim.verify(no_forms).expect("verify final claim"); + verifier_state.check_eof().expect("proof EOF"); + println!( + "verified proof artifact backend={} log_size={} fold={} rate={} pow_bits={} proof_bytes={}", + backend_name(), + artifact.log_size, + artifact.fold, + artifact.rate, + artifact.pow_bits, + artifact.proof_bytes + ); +} + +fn protocol_parameters( + security_level: usize, + pow_bits: usize, + fold: usize, + rate: usize, +) -> ProtocolParameters { + ProtocolParameters { + security_level, + pow_bits, + initial_folding_factor: fold, + folding_factor: fold, + unique_decoding: false, + starting_log_inv_rate: rate, + batch_size: 1, + hash_id: hash::SHA2, + } +} + +fn input_vector(size: usize) -> Vec { + (0..size).map(|i| F::from(i as u64)).collect() +} + +#[cfg(all(feature = "metal", target_os = "macos"))] +fn backend_name() -> &'static str { + "gpu-metal" +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +fn backend_name() -> &'static str { + "cpu" +} diff --git a/src/bin/phase_benchmark.rs b/src/bin/phase_benchmark.rs new file mode 100644 index 00000000..38f7f22e --- /dev/null +++ b/src/bin/phase_benchmark.rs @@ -0,0 +1,504 @@ +use std::{ + alloc::{GlobalAlloc, Layout, System}, + borrow::Cow, + fs::OpenOptions, + hint::black_box, + io::Write, + sync::atomic::{AtomicUsize, Ordering}, + time::{Duration, Instant}, +}; + +use clap::Parser; +use serde::Serialize; +use whir::{ + algebra::{ + buffer::{ActiveBuffer, FieldOps}, + embedding::Identity, + fields::Field256, + }, + hash::{self, HASH_COUNTER}, + parameters::ProtocolParameters, + protocols::whir::Config as WhirConfig, + transcript::{codecs::Empty, DomainSeparator, ProverState}, +}; + +#[global_allocator] +static ALLOC: TrackingAllocator = TrackingAllocator; + +static CURRENT_ALLOCATED: AtomicUsize = AtomicUsize::new(0); +static PEAK_ALLOCATED: AtomicUsize = AtomicUsize::new(0); + +struct TrackingAllocator; + +unsafe impl GlobalAlloc for TrackingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + add_allocated(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + CURRENT_ALLOCATED.fetch_sub(layout.size(), Ordering::Relaxed); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + match new_size.cmp(&layout.size()) { + std::cmp::Ordering::Greater => add_allocated(new_size - layout.size()), + std::cmp::Ordering::Less => { + CURRENT_ALLOCATED.fetch_sub(layout.size() - new_size, Ordering::Relaxed); + } + std::cmp::Ordering::Equal => {} + } + } + new_ptr + } +} + +fn add_allocated(bytes: usize) { + let current = CURRENT_ALLOCATED.fetch_add(bytes, Ordering::Relaxed) + bytes; + let mut peak = PEAK_ALLOCATED.load(Ordering::Relaxed); + while current > peak { + match PEAK_ALLOCATED.compare_exchange_weak( + peak, + current, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(next) => peak = next, + } + } +} + +#[derive(Parser, Debug)] +#[command(author, version, about = "Phase benchmark for CPU/GPU WHIR proving")] +struct Args { + #[arg(long, default_value_t = 16)] + min_log_size: usize, + + #[arg(long, default_value_t = 28)] + max_log_size: usize, + + #[arg(long, default_value = "1,2,3,4,6")] + folds: String, + + #[arg(long, default_value = "1,2,3")] + rates: String, + + /// Apply the article-style cap where rates above 1 are skipped for n >= 24. + #[arg(long)] + article_grid: bool, + + #[arg(long, default_value = "commit,sumcheck,e2e")] + phases: String, + + #[arg(long, default_value_t = 128)] + security_level: usize, + + #[arg(long, default_value_t = 20)] + pow_bits: usize, + + #[arg(long, default_value_t = false)] + unique_decoding: bool, + + #[arg(long, default_value = "outputs/phase_benchmark.jsonl")] + output: String, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Phase { + Commit, + Sumcheck, + E2e, +} + +#[derive(Serialize)] +struct PhaseRow { + backend: &'static str, + phase: &'static str, + log_size: usize, + size: usize, + fold: usize, + rate: usize, + duration_ms: f64, + current_allocated_bytes: usize, + peak_allocated_bytes: usize, + peak_phase_delta_bytes: usize, + hashes: usize, + proof_bytes: Option, + status: &'static str, + metal_upload_count: u64, + metal_upload_bytes: u64, + metal_upload_ms: f64, + metal_readback_count: u64, + metal_readback_bytes: u64, + metal_readback_ms: f64, + metal_alloc_count: u64, + metal_alloc_bytes: u64, + metal_command_count: u64, + metal_command_wait_ms: f64, + metal_blit_count: u64, + metal_blit_bytes: u64, + metal_blit_wait_ms: f64, +} + +struct Measurement { + value: T, + duration: Duration, + current_allocated_bytes: usize, + peak_allocated_bytes: usize, + peak_phase_delta_bytes: usize, + metal: MetalPhaseProfile, +} + +#[derive(Clone, Copy, Debug, Default)] +struct MetalPhaseProfile { + upload_count: u64, + upload_bytes: u64, + upload_ms: f64, + readback_count: u64, + readback_bytes: u64, + readback_ms: f64, + alloc_count: u64, + alloc_bytes: u64, + command_count: u64, + command_wait_ms: f64, + blit_count: u64, + blit_bytes: u64, + blit_wait_ms: f64, +} + +#[cfg(all(feature = "metal", target_os = "macos"))] +const BACKEND: &str = "gpu-metal"; + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +const BACKEND: &str = "cpu"; + +type F = Field256; +type M = Identity; + +fn main() { + let args = Args::parse(); + assert!(args.min_log_size <= args.max_log_size); + warm_backend(); + + let folds = parse_usize_list(&args.folds); + let rates = parse_usize_list(&args.rates); + let phases = parse_phases(&args.phases); + + std::fs::create_dir_all("outputs").unwrap(); + let mut output = OpenOptions::new() + .append(true) + .create(true) + .open(&args.output) + .unwrap(); + + for log_size in args.min_log_size..=args.max_log_size { + for &fold in &folds { + if fold > log_size { + eprintln!("skip n={log_size} fold={fold}: fold must be <= n"); + continue; + } + for &rate in &rates { + if args.article_grid && log_size >= 24 && rate > 1 { + continue; + } + let size = 1usize << log_size; + let params = whir_params(&args, fold, rate); + for &phase in &phases { + let row = match phase { + Phase::Commit => bench_commit(log_size, size, fold, rate, ¶ms), + Phase::Sumcheck => bench_sumcheck(log_size, size, fold, rate, ¶ms), + Phase::E2e => bench_e2e(log_size, size, fold, rate, ¶ms), + }; + writeln!(output, "{}", serde_json::to_string(&row).unwrap()).unwrap(); + output.flush().unwrap(); + eprintln!( + "{} n={log_size} fold={fold} rate={rate} phase={} {:.3} ms peak={} MiB upload={} KiB/{:.3} ms readback={} KiB/{:.3} ms cmd_wait={:.3} ms", + BACKEND, + row.phase, + row.duration_ms, + row.peak_allocated_bytes / (1024 * 1024), + row.metal_upload_bytes / 1024, + row.metal_upload_ms, + row.metal_readback_bytes / 1024, + row.metal_readback_ms, + row.metal_command_wait_ms, + ); + } + } + } + } +} + +#[cfg(all(feature = "metal", target_os = "macos"))] +fn warm_backend() { + whir::algebra::buffer::MetalBuffer::::warmup(); + whir::hash::MetalSha2::warmup(); +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +fn warm_backend() {} + +fn whir_params(args: &Args, fold: usize, rate: usize) -> ProtocolParameters { + ProtocolParameters { + security_level: args.security_level, + pow_bits: args.pow_bits, + initial_folding_factor: fold, + folding_factor: fold, + unique_decoding: args.unique_decoding, + starting_log_inv_rate: rate, + batch_size: 1, + hash_id: hash::SHA2, + } +} + +fn bench_commit( + log_size: usize, + size: usize, + fold: usize, + rate: usize, + whir_params: &ProtocolParameters, +) -> PhaseRow { + let vector = input_vector(size); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let params = WhirConfig::::new(size, whir_params); + let ds = DomainSeparator::protocol(¶ms) + .session(&"phase benchmark commit") + .instance(&Empty); + + HASH_COUNTER.reset(); + let measured = measure_phase(|| { + let mut prover_state = ProverState::new_std(&ds); + let _ = black_box(params.commit(&mut prover_state, &[&vector_buffer])); + }); + row( + Phase::Commit, + log_size, + size, + fold, + rate, + measured, + HASH_COUNTER.get(), + None, + ) +} + +fn bench_sumcheck( + log_size: usize, + size: usize, + fold: usize, + rate: usize, + whir_params: &ProtocolParameters, +) -> PhaseRow { + let params = WhirConfig::::new(size, whir_params); + let config = params.initial_sumcheck.clone(); + let ds = DomainSeparator::protocol(&config) + .session(&"phase benchmark sumcheck") + .instance(&Empty); + + HASH_COUNTER.reset(); + let measured = measure_phase(|| { + let mut a = ActiveBuffer::from_vec(input_vector(size)); + let mut b = ActiveBuffer::from_vec((0..size).map(|i| F::from(i as u64 + 17)).collect()); + let mut sum = a.dot(&b); + let mut prover_state = ProverState::new_std(&ds); + let _ = black_box(config.prove(&mut prover_state, &mut a, &mut b, &mut sum, &[])); + }); + row( + Phase::Sumcheck, + log_size, + size, + fold, + rate, + measured, + HASH_COUNTER.get(), + None, + ) +} + +fn bench_e2e( + log_size: usize, + size: usize, + fold: usize, + rate: usize, + whir_params: &ProtocolParameters, +) -> PhaseRow { + let vector = input_vector(size); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let params = WhirConfig::::new(size, whir_params); + let ds = DomainSeparator::protocol(¶ms) + .session(&"phase benchmark e2e") + .instance(&Empty); + + HASH_COUNTER.reset(); + let measured = measure_phase(|| { + let mut prover_state = ProverState::new_std(&ds); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); + let _ = params.prove( + &mut prover_state, + &[&vector_buffer], + vec![&witness], + vec![], + Cow::Owned(vec![]), + ); + let proof = prover_state.proof(); + proof.narg_string.len() + proof.hints.len() + }); + let proof_bytes = measured.value; + row( + Phase::E2e, + log_size, + size, + fold, + rate, + measured, + HASH_COUNTER.get(), + Some(proof_bytes), + ) +} + +fn input_vector(size: usize) -> Vec { + (0..size).map(|i| F::from(i as u64)).collect() +} + +fn measure_phase(f: impl FnOnce() -> T) -> Measurement { + let before = CURRENT_ALLOCATED.load(Ordering::Relaxed); + PEAK_ALLOCATED.store(before, Ordering::Relaxed); + let metal_before = metal_profile_snapshot(); + let start = Instant::now(); + let value = f(); + let duration = start.elapsed(); + let metal_after = metal_profile_snapshot(); + let current = CURRENT_ALLOCATED.load(Ordering::Relaxed); + let peak = PEAK_ALLOCATED.load(Ordering::Relaxed); + Measurement { + value, + duration, + current_allocated_bytes: current, + peak_allocated_bytes: peak, + peak_phase_delta_bytes: peak.saturating_sub(before), + metal: metal_profile_delta(metal_before, metal_after), + } +} + +fn row( + phase: Phase, + log_size: usize, + size: usize, + fold: usize, + rate: usize, + measurement: Measurement, + hashes: usize, + proof_bytes: Option, +) -> PhaseRow { + PhaseRow { + backend: BACKEND, + phase: phase.name(), + log_size, + size, + fold, + rate, + duration_ms: measurement.duration.as_secs_f64() * 1_000.0, + current_allocated_bytes: measurement.current_allocated_bytes, + peak_allocated_bytes: measurement.peak_allocated_bytes, + peak_phase_delta_bytes: measurement.peak_phase_delta_bytes, + hashes, + proof_bytes, + status: "ok", + metal_upload_count: measurement.metal.upload_count, + metal_upload_bytes: measurement.metal.upload_bytes, + metal_upload_ms: measurement.metal.upload_ms, + metal_readback_count: measurement.metal.readback_count, + metal_readback_bytes: measurement.metal.readback_bytes, + metal_readback_ms: measurement.metal.readback_ms, + metal_alloc_count: measurement.metal.alloc_count, + metal_alloc_bytes: measurement.metal.alloc_bytes, + metal_command_count: measurement.metal.command_count, + metal_command_wait_ms: measurement.metal.command_wait_ms, + metal_blit_count: measurement.metal.blit_count, + metal_blit_bytes: measurement.metal.blit_bytes, + metal_blit_wait_ms: measurement.metal.blit_wait_ms, + } +} + +#[cfg(all(feature = "metal", target_os = "macos"))] +fn metal_profile_snapshot() -> whir::hash::MetalProfileSnapshot { + whir::hash::metal_profile_snapshot() +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +fn metal_profile_snapshot() -> MetalProfileSnapshotCompat { + MetalProfileSnapshotCompat::default() +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +#[derive(Clone, Copy, Debug, Default)] +struct MetalProfileSnapshotCompat; + +#[cfg(all(feature = "metal", target_os = "macos"))] +fn metal_profile_delta( + before: whir::hash::MetalProfileSnapshot, + after: whir::hash::MetalProfileSnapshot, +) -> MetalPhaseProfile { + let delta = after.delta_since(before); + MetalPhaseProfile { + upload_count: delta.upload_count, + upload_bytes: delta.upload_bytes, + upload_ms: delta.upload_ms(), + readback_count: delta.readback_count, + readback_bytes: delta.readback_bytes, + readback_ms: delta.readback_ms(), + alloc_count: delta.alloc_count, + alloc_bytes: delta.alloc_bytes, + command_count: delta.command_count, + command_wait_ms: delta.command_wait_ms(), + blit_count: delta.blit_count, + blit_bytes: delta.blit_bytes, + blit_wait_ms: delta.blit_wait_ms(), + } +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +fn metal_profile_delta( + _before: MetalProfileSnapshotCompat, + _after: MetalProfileSnapshotCompat, +) -> MetalPhaseProfile { + MetalPhaseProfile::default() +} + +impl Phase { + const fn name(self) -> &'static str { + match self { + Self::Commit => "commit", + Self::Sumcheck => "sumcheck", + Self::E2e => "e2e_prove", + } + } +} + +fn parse_usize_list(source: &str) -> Vec { + source + .split(',') + .filter(|part| !part.is_empty()) + .map(|part| part.parse().unwrap()) + .collect() +} + +fn parse_phases(source: &str) -> Vec { + source + .split(',') + .filter(|part| !part.is_empty()) + .map(|part| match part { + "commit" => Phase::Commit, + "sumcheck" => Phase::Sumcheck, + "e2e" | "e2e_prove" => Phase::E2e, + other => panic!("unknown phase {other}; use commit,sumcheck,e2e"), + }) + .collect() +} diff --git a/src/bin/phase_benchmark_report.rs b/src/bin/phase_benchmark_report.rs new file mode 100644 index 00000000..146462f8 --- /dev/null +++ b/src/bin/phase_benchmark_report.rs @@ -0,0 +1,127 @@ +use std::{ + collections::BTreeMap, + fs::File, + io::{BufRead, BufReader}, +}; + +use serde::Deserialize; + +#[derive(Clone, Debug, Deserialize)] +struct Row { + backend: String, + phase: String, + log_size: usize, + size: usize, + fold: usize, + rate: usize, + duration_ms: f64, + peak_allocated_bytes: usize, + peak_phase_delta_bytes: usize, + hashes: usize, + proof_bytes: Option, + status: Option, + #[serde(default)] + metal_upload_bytes: u64, + #[serde(default)] + metal_upload_ms: f64, + #[serde(default)] + metal_readback_bytes: u64, + #[serde(default)] + metal_readback_ms: f64, + #[serde(default)] + metal_alloc_bytes: u64, + #[serde(default)] + metal_command_wait_ms: f64, + #[serde(default)] + metal_blit_bytes: u64, + #[serde(default)] + metal_blit_wait_ms: f64, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct Key { + phase: String, + log_size: usize, + size: usize, + fold: usize, + rate: usize, +} + +#[derive(Default)] +struct Pair { + cpu: Option, + gpu: Option, +} + +fn main() { + let path = std::env::args() + .nth(1) + .unwrap_or_else(|| "outputs/phase_benchmark_cpu_gpu.jsonl".to_string()); + let file = File::open(&path).unwrap_or_else(|err| panic!("failed to open {path}: {err}")); + let mut rows = BTreeMap::::new(); + for line in BufReader::new(file).lines() { + let line = line.unwrap(); + if line.trim().is_empty() { + continue; + } + let row: Row = serde_json::from_str(&line).unwrap(); + let key = Key { + phase: row.phase.clone(), + log_size: row.log_size, + size: row.size, + fold: row.fold, + rate: row.rate, + }; + let pair = rows.entry(key).or_default(); + match row.backend.as_str() { + "cpu" => pair.cpu = Some(row), + "gpu-metal" => pair.gpu = Some(row), + _ => {} + } + } + + println!( + "log_size,size,fold,rate,phase,cpu_status,gpu_status,cpu_ms,gpu_ms,speedup,cpu_peak_bytes,gpu_peak_bytes,cpu_peak_delta_bytes,gpu_peak_delta_bytes,cpu_hashes,gpu_hashes,proof_bytes,gpu_upload_bytes,gpu_upload_ms,gpu_readback_bytes,gpu_readback_ms,gpu_alloc_bytes,gpu_command_wait_ms,gpu_blit_bytes,gpu_blit_wait_ms" + ); + for (key, pair) in rows { + let (Some(cpu), Some(gpu)) = (pair.cpu, pair.gpu) else { + continue; + }; + let speedup = if gpu.duration_ms > 0.0 { + cpu.duration_ms / gpu.duration_ms + } else { + 0.0 + }; + let proof_bytes = cpu.proof_bytes.or(gpu.proof_bytes).unwrap_or(0); + let cpu_status = cpu.status.as_deref().unwrap_or("ok"); + let gpu_status = gpu.status.as_deref().unwrap_or("ok"); + println!( + "{},{},{},{},{},{},{},{:.6},{:.6},{:.4},{},{},{},{},{},{},{},{},{:.6},{},{:.6},{},{:.6},{},{:.6}", + key.log_size, + key.size, + key.fold, + key.rate, + key.phase, + cpu_status, + gpu_status, + cpu.duration_ms, + gpu.duration_ms, + speedup, + cpu.peak_allocated_bytes, + gpu.peak_allocated_bytes, + cpu.peak_phase_delta_bytes, + gpu.peak_phase_delta_bytes, + cpu.hashes, + gpu.hashes, + proof_bytes, + gpu.metal_upload_bytes, + gpu.metal_upload_ms, + gpu.metal_readback_bytes, + gpu.metal_readback_ms, + gpu.metal_alloc_bytes, + gpu.metal_command_wait_ms, + gpu.metal_blit_bytes, + gpu.metal_blit_wait_ms, + ); + } +} diff --git a/src/hash/metal_profile.rs b/src/hash/metal_profile.rs new file mode 100644 index 00000000..17134c13 --- /dev/null +++ b/src/hash/metal_profile.rs @@ -0,0 +1,123 @@ +use std::{ + sync::atomic::{AtomicU64, Ordering}, + time::Duration, +}; + +#[derive(Clone, Copy, Debug, Default, serde::Serialize, serde::Deserialize)] +pub struct MetalProfileSnapshot { + pub upload_count: u64, + pub upload_bytes: u64, + pub upload_nanos: u64, + pub readback_count: u64, + pub readback_bytes: u64, + pub readback_nanos: u64, + pub alloc_count: u64, + pub alloc_bytes: u64, + pub command_count: u64, + pub command_wait_nanos: u64, + pub blit_count: u64, + pub blit_bytes: u64, + pub blit_wait_nanos: u64, +} + +impl MetalProfileSnapshot { + pub fn delta_since(self, before: Self) -> Self { + Self { + upload_count: self.upload_count.saturating_sub(before.upload_count), + upload_bytes: self.upload_bytes.saturating_sub(before.upload_bytes), + upload_nanos: self.upload_nanos.saturating_sub(before.upload_nanos), + readback_count: self.readback_count.saturating_sub(before.readback_count), + readback_bytes: self.readback_bytes.saturating_sub(before.readback_bytes), + readback_nanos: self.readback_nanos.saturating_sub(before.readback_nanos), + alloc_count: self.alloc_count.saturating_sub(before.alloc_count), + alloc_bytes: self.alloc_bytes.saturating_sub(before.alloc_bytes), + command_count: self.command_count.saturating_sub(before.command_count), + command_wait_nanos: self + .command_wait_nanos + .saturating_sub(before.command_wait_nanos), + blit_count: self.blit_count.saturating_sub(before.blit_count), + blit_bytes: self.blit_bytes.saturating_sub(before.blit_bytes), + blit_wait_nanos: self.blit_wait_nanos.saturating_sub(before.blit_wait_nanos), + } + } + + pub fn upload_ms(self) -> f64 { + nanos_to_ms(self.upload_nanos) + } + + pub fn readback_ms(self) -> f64 { + nanos_to_ms(self.readback_nanos) + } + + pub fn command_wait_ms(self) -> f64 { + nanos_to_ms(self.command_wait_nanos) + } + + pub fn blit_wait_ms(self) -> f64 { + nanos_to_ms(self.blit_wait_nanos) + } +} + +static UPLOAD_COUNT: AtomicU64 = AtomicU64::new(0); +static UPLOAD_BYTES: AtomicU64 = AtomicU64::new(0); +static UPLOAD_NANOS: AtomicU64 = AtomicU64::new(0); +static READBACK_COUNT: AtomicU64 = AtomicU64::new(0); +static READBACK_BYTES: AtomicU64 = AtomicU64::new(0); +static READBACK_NANOS: AtomicU64 = AtomicU64::new(0); +static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0); +static ALLOC_BYTES: AtomicU64 = AtomicU64::new(0); +static COMMAND_COUNT: AtomicU64 = AtomicU64::new(0); +static COMMAND_WAIT_NANOS: AtomicU64 = AtomicU64::new(0); +static BLIT_COUNT: AtomicU64 = AtomicU64::new(0); +static BLIT_BYTES: AtomicU64 = AtomicU64::new(0); +static BLIT_WAIT_NANOS: AtomicU64 = AtomicU64::new(0); + +pub fn snapshot() -> MetalProfileSnapshot { + MetalProfileSnapshot { + upload_count: UPLOAD_COUNT.load(Ordering::Relaxed), + upload_bytes: UPLOAD_BYTES.load(Ordering::Relaxed), + upload_nanos: UPLOAD_NANOS.load(Ordering::Relaxed), + readback_count: READBACK_COUNT.load(Ordering::Relaxed), + readback_bytes: READBACK_BYTES.load(Ordering::Relaxed), + readback_nanos: READBACK_NANOS.load(Ordering::Relaxed), + alloc_count: ALLOC_COUNT.load(Ordering::Relaxed), + alloc_bytes: ALLOC_BYTES.load(Ordering::Relaxed), + command_count: COMMAND_COUNT.load(Ordering::Relaxed), + command_wait_nanos: COMMAND_WAIT_NANOS.load(Ordering::Relaxed), + blit_count: BLIT_COUNT.load(Ordering::Relaxed), + blit_bytes: BLIT_BYTES.load(Ordering::Relaxed), + blit_wait_nanos: BLIT_WAIT_NANOS.load(Ordering::Relaxed), + } +} + +pub fn record_upload(bytes: u64, duration: Duration) { + UPLOAD_COUNT.fetch_add(1, Ordering::Relaxed); + UPLOAD_BYTES.fetch_add(bytes, Ordering::Relaxed); + UPLOAD_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); +} + +pub fn record_readback(bytes: u64, duration: Duration) { + READBACK_COUNT.fetch_add(1, Ordering::Relaxed); + READBACK_BYTES.fetch_add(bytes, Ordering::Relaxed); + READBACK_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); +} + +pub fn record_alloc(bytes: u64) { + ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); + ALLOC_BYTES.fetch_add(bytes, Ordering::Relaxed); +} + +pub fn record_command_wait(duration: Duration) { + COMMAND_COUNT.fetch_add(1, Ordering::Relaxed); + COMMAND_WAIT_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); +} + +pub fn record_blit(bytes: u64, duration: Duration) { + BLIT_COUNT.fetch_add(1, Ordering::Relaxed); + BLIT_BYTES.fetch_add(bytes, Ordering::Relaxed); + BLIT_WAIT_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); +} + +fn nanos_to_ms(nanos: u64) -> f64 { + nanos as f64 / 1_000_000.0 +} diff --git a/src/hash/metal_sha2_engine.rs b/src/hash/metal_sha2_engine.rs new file mode 100644 index 00000000..8879d77b --- /dev/null +++ b/src/hash/metal_sha2_engine.rs @@ -0,0 +1,715 @@ +use std::{borrow::Cow, sync::OnceLock, time::Instant}; + +use const_oid::ObjectIdentifier; +use digest::const_oid::AssociatedOid; +use metal::{ + Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, MTLResourceOptions, MTLSize, +}; +use sha2::Sha256; +use zerocopy::IntoBytes; + +use super::{metal_profile, Hash, HashEngine, HASH_COUNTER}; + +const USE_SHA256_64_SPECIALIZATION: bool = false; +const POW_MAX_WINDOW: u32 = 1 << 20; +const POW_MIN_WINDOW: u32 = 1 << 12; +const POW_THREADS_PER_GROUP: u64 = 256; + +const SHA256_METAL: &str = r#" +#include +using namespace metal; + +constant uint K[64] = { + 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, + 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, + 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, + 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, + 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, + 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, + 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, + 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, + 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, + 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, + 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, + 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, + 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, + 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, + 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, + 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2, +}; + +inline uint rotr(uint x, uint n) { + return (x >> n) | (x << (32 - n)); +} + +inline uchar padded_byte( + device const uchar *input, + uint message, + uint message_size, + uint block_offset, + uint padded_size +) { + ulong bit_len = ((ulong)message_size) * 8UL; + if (block_offset < message_size) { + return input[message * message_size + block_offset]; + } + if (block_offset == message_size) { + return 0x80; + } + uint len_start = padded_size - 8; + if (block_offset >= len_start) { + uint shift = (7 - (block_offset - len_start)) * 8; + return (uchar)((bit_len >> shift) & 0xff); + } + return 0; +} + +kernel void sha256_many( + device const uchar *input [[buffer(0)]], + device uchar *output [[buffer(1)]], + constant uint &message_size [[buffer(2)]], + constant uint &messages [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= messages) return; + + uint padded_size = ((message_size + 9 + 63) / 64) * 64; + uint h0 = 0x6a09e667; + uint h1 = 0xbb67ae85; + uint h2 = 0x3c6ef372; + uint h3 = 0xa54ff53a; + uint h4 = 0x510e527f; + uint h5 = 0x9b05688c; + uint h6 = 0x1f83d9ab; + uint h7 = 0x5be0cd19; + + for (uint block = 0; block < padded_size; block += 64) { + uint w[64]; + for (uint i = 0; i < 16; i++) { + uint off = block + i * 4; + w[i] = + ((uint)padded_byte(input, gid, message_size, off + 0, padded_size) << 24) | + ((uint)padded_byte(input, gid, message_size, off + 1, padded_size) << 16) | + ((uint)padded_byte(input, gid, message_size, off + 2, padded_size) << 8) | + ((uint)padded_byte(input, gid, message_size, off + 3, padded_size)); + } + for (uint i = 16; i < 64; i++) { + uint s0 = rotr(w[i - 15], 7) ^ rotr(w[i - 15], 18) ^ (w[i - 15] >> 3); + uint s1 = rotr(w[i - 2], 17) ^ rotr(w[i - 2], 19) ^ (w[i - 2] >> 10); + w[i] = w[i - 16] + s0 + w[i - 7] + s1; + } + + uint a = h0, b = h1, c = h2, d = h3; + uint e = h4, f = h5, g = h6, h = h7; + for (uint i = 0; i < 64; i++) { + uint s1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25); + uint ch = (e & f) ^ ((~e) & g); + uint temp1 = h + s1 + ch + K[i] + w[i]; + uint s0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22); + uint maj = (a & b) ^ (a & c) ^ (b & c); + uint temp2 = s0 + maj; + h = g; + g = f; + f = e; + e = d + temp1; + d = c; + c = b; + b = a; + a = temp1 + temp2; + } + + h0 += a; h1 += b; h2 += c; h3 += d; + h4 += e; h5 += f; h6 += g; h7 += h; + } + + uint hs[8] = { h0, h1, h2, h3, h4, h5, h6, h7 }; + uint out = gid * 32; + for (uint i = 0; i < 8; i++) { + output[out + i * 4 + 0] = (uchar)((hs[i] >> 24) & 0xff); + output[out + i * 4 + 1] = (uchar)((hs[i] >> 16) & 0xff); + output[out + i * 4 + 2] = (uchar)((hs[i] >> 8) & 0xff); + output[out + i * 4 + 3] = (uchar)(hs[i] & 0xff); + } +} + +inline uint load_be32(device const uchar *input, uint off) { + return ((uint)input[off + 0] << 24) | + ((uint)input[off + 1] << 16) | + ((uint)input[off + 2] << 8) | + ((uint)input[off + 3]); +} + +inline void sha256_compress_16(thread uint h[8], thread uint w[16]) { + uint a = h[0], b = h[1], c = h[2], d = h[3]; + uint e = h[4], f = h[5], g = h[6], hh = h[7]; + + for (uint i = 0; i < 64; i++) { + uint wi; + if (i < 16) { + wi = w[i]; + } else { + uint s0 = rotr(w[(i + 1) & 15], 7) ^ rotr(w[(i + 1) & 15], 18) ^ (w[(i + 1) & 15] >> 3); + uint s1 = rotr(w[(i + 14) & 15], 17) ^ rotr(w[(i + 14) & 15], 19) ^ (w[(i + 14) & 15] >> 10); + wi = w[i & 15] + s0 + w[(i + 9) & 15] + s1; + w[i & 15] = wi; + } + uint S1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25); + uint ch = (e & f) ^ ((~e) & g); + uint temp1 = hh + S1 + ch + K[i] + wi; + uint S0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22); + uint maj = (a & b) ^ (a & c) ^ (b & c); + uint temp2 = S0 + maj; + hh = g; + g = f; + f = e; + e = d + temp1; + d = c; + c = b; + b = a; + a = temp1 + temp2; + } + + h[0] += a; h[1] += b; h[2] += c; h[3] += d; + h[4] += e; h[5] += f; h[6] += g; h[7] += hh; +} + +kernel void sha256_many_64( + device const uchar *input [[buffer(0)]], + device uchar *output [[buffer(1)]], + constant uint &messages [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= messages) return; + + uint h[8] = { + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, + 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19 + }; + + uint w[16]; + uint base = gid * 64; + for (uint i = 0; i < 16; i++) { + w[i] = load_be32(input, base + i * 4); + } + sha256_compress_16(h, w); + + w[0] = 0x80000000; + for (uint i = 1; i < 15; i++) { + w[i] = 0; + } + w[15] = 512; + sha256_compress_16(h, w); + + uint out = gid * 32; + for (uint i = 0; i < 8; i++) { + output[out + i * 4 + 0] = (uchar)((h[i] >> 24) & 0xff); + output[out + i * 4 + 1] = (uchar)((h[i] >> 16) & 0xff); + output[out + i * 4 + 2] = (uchar)((h[i] >> 8) & 0xff); + output[out + i * 4 + 3] = (uchar)(h[i] & 0xff); + } +} + +inline ulong digest_threshold_value(uint h0, uint h1) { + ulong b0 = (ulong)((h0 >> 24) & 0xff); + ulong b1 = (ulong)((h0 >> 16) & 0xff); + ulong b2 = (ulong)((h0 >> 8) & 0xff); + ulong b3 = (ulong)(h0 & 0xff); + ulong b4 = (ulong)((h1 >> 24) & 0xff); + ulong b5 = (ulong)((h1 >> 16) & 0xff); + ulong b6 = (ulong)((h1 >> 8) & 0xff); + ulong b7 = (ulong)(h1 & 0xff); + return b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) | + (b4 << 32) | (b5 << 40) | (b6 << 48) | (b7 << 56); +} + +kernel void sha256_pow_64( + device const uchar *challenge [[buffer(0)]], + device ulong *candidates [[buffer(1)]], + constant ulong &start_nonce [[buffer(2)]], + constant ulong &threshold [[buffer(3)]], + constant uint &count [[buffer(4)]], + uint gid [[thread_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint tgid [[threadgroup_position_in_grid]] +) { + threadgroup ulong local_best[256]; + ulong candidate = 0xffffffffffffffffUL; + + if (gid < count) { + ulong nonce = start_nonce + (ulong)gid; + + uint h[8] = { + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, + 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19 + }; + + uint w[16]; + for (uint i = 0; i < 8; i++) { + w[i] = load_be32(challenge, i * 4); + } + w[8] = ((uint)(nonce & 0xffUL) << 24) | + ((uint)((nonce >> 8) & 0xffUL) << 16) | + ((uint)((nonce >> 16) & 0xffUL) << 8) | + ((uint)((nonce >> 24) & 0xffUL)); + w[9] = ((uint)((nonce >> 32) & 0xffUL) << 24) | + ((uint)((nonce >> 40) & 0xffUL) << 16) | + ((uint)((nonce >> 48) & 0xffUL) << 8) | + ((uint)((nonce >> 56) & 0xffUL)); + for (uint i = 10; i < 16; i++) { + w[i] = 0; + } + sha256_compress_16(h, w); + + w[0] = 0x80000000; + for (uint i = 1; i < 15; i++) { + w[i] = 0; + } + w[15] = 512; + sha256_compress_16(h, w); + + ulong value = digest_threshold_value(h[0], h[1]); + if (value <= threshold) { + candidate = nonce; + } + } + + local_best[tid] = candidate; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 128; stride > 0; stride >>= 1) { + if (tid < stride && local_best[tid + stride] < local_best[tid]) { + local_best[tid] = local_best[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0) { + candidates[tgid] = local_best[0]; + } +} +"#; + +#[derive(Clone, Copy, Debug, Default)] +pub struct MetalSha2; + +struct MetalSha2Runtime { + device: Device, + queue: CommandQueue, + pipeline: ComputePipelineState, + pipeline_64: ComputePipelineState, + pow_pipeline: ComputePipelineState, +} + +fn new_shared_buffer(rt: &MetalSha2Runtime, bytes: u64) -> Buffer { + metal_profile::record_alloc(bytes); + rt.device + .new_buffer(bytes, MTLResourceOptions::StorageModeShared) +} + +fn new_shared_buffer_with_data( + rt: &MetalSha2Runtime, + data: *const std::ffi::c_void, + bytes: u64, +) -> Buffer { + let start = Instant::now(); + let buffer = rt + .device + .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); + metal_profile::record_alloc(bytes); + metal_profile::record_upload(bytes, start.elapsed()); + buffer +} + +fn upload_u32(rt: &MetalSha2Runtime, value: u32) -> Buffer { + new_shared_buffer_with_data(rt, (&value as *const u32).cast(), size_of::() as u64) +} + +fn upload_u64(rt: &MetalSha2Runtime, value: u64) -> Buffer { + new_shared_buffer_with_data(rt, (&value as *const u64).cast(), size_of::() as u64) +} + +fn wait_for_command(command: &metal::CommandBufferRef) { + wait_for_command_named(command, "sha256"); +} + +fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { + command.commit(); + let start = Instant::now(); + command.wait_until_completed(); + let elapsed = start.elapsed(); + if std::env::var_os("WHIR_METAL_TRACE").is_some() { + eprintln!( + "metal command {label} {:.3} ms", + elapsed.as_secs_f64() * 1_000.0 + ); + } + metal_profile::record_command_wait(elapsed); +} + +fn runtime() -> &'static MetalSha2Runtime { + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + let device = Device::system_default().expect("Metal device is not available"); + let library = device + .new_library_with_source(SHA256_METAL, &CompileOptions::new()) + .expect("failed to compile Metal SHA-256 kernel"); + let function = library + .get_function("sha256_many", None) + .expect("missing Metal SHA-256 kernel"); + let function_64 = library + .get_function("sha256_many_64", None) + .expect("missing Metal SHA-256 64-byte kernel"); + let pow_function = library + .get_function("sha256_pow_64", None) + .expect("missing Metal SHA-256 PoW kernel"); + let pipeline = device + .new_compute_pipeline_state_with_function(&function) + .expect("failed to create Metal SHA-256 pipeline"); + let pipeline_64 = device + .new_compute_pipeline_state_with_function(&function_64) + .expect("failed to create Metal SHA-256 64-byte pipeline"); + let pow_pipeline = device + .new_compute_pipeline_state_with_function(&pow_function) + .expect("failed to create Metal SHA-256 PoW pipeline"); + let queue = device.new_command_queue(); + MetalSha2Runtime { + device, + queue, + pipeline, + pipeline_64, + pow_pipeline, + } + }) +} + +impl MetalSha2 { + pub const fn new() -> Self { + Self + } + + pub fn warmup() { + let _ = runtime(); + } + + pub fn prove_pow_64(challenge: &[u8; 32], threshold: u64) -> u64 { + let rt = runtime(); + let window = pow_window_for_threshold(threshold); + let challenge_buffer = + new_shared_buffer_with_data(rt, challenge.as_ptr().cast(), challenge.len() as u64); + let candidate_groups = u64::from(window).div_ceil(POW_THREADS_PER_GROUP); + let candidate_bytes = candidate_groups * size_of::() as u64; + let candidates = new_shared_buffer(rt, candidate_bytes); + let threshold = upload_u64(rt, threshold); + let count = upload_u32(rt, window); + + let mut start_nonce = 0u64; + loop { + let start = upload_u64(rt, start_nonce); + let command = rt.queue.new_command_buffer(); + let encoder = command.new_compute_command_encoder(); + encoder.set_compute_pipeline_state(&rt.pow_pipeline); + encoder.set_buffer(0, Some(&challenge_buffer), 0); + encoder.set_buffer(1, Some(&candidates), 0); + encoder.set_buffer(2, Some(&start), 0); + encoder.set_buffer(3, Some(&threshold), 0); + encoder.set_buffer(4, Some(&count), 0); + encoder.dispatch_thread_groups( + MTLSize { + width: candidate_groups, + height: 1, + depth: 1, + }, + MTLSize { + width: POW_THREADS_PER_GROUP, + height: 1, + depth: 1, + }, + ); + encoder.end_encoding(); + wait_for_command_named(&command, "sha256_pow"); + + let read_start = Instant::now(); + let candidates = unsafe { + std::slice::from_raw_parts( + candidates.contents().cast::(), + candidate_groups as usize, + ) + }; + let nonce = candidates.iter().copied().min().unwrap_or(u64::MAX); + metal_profile::record_readback(candidate_bytes, read_start.elapsed()); + if nonce != u64::MAX { + HASH_COUNTER.add((nonce - start_nonce + 1) as usize); + return nonce; + } + HASH_COUNTER.add(window as usize); + start_nonce = start_nonce + .checked_add(window as u64) + .expect("PoW nonce range exhausted"); + } + } + + pub(crate) fn build_merkle_tree_buffer_from_messages_buffer( + &self, + message_size: usize, + messages: &Buffer, + num_leaves: usize, + layers: usize, + ) -> Buffer { + assert_eq!(num_leaves, 1usize << layers); + let rt = runtime(); + let num_nodes = (1usize << (layers + 1)) - 1; + let nodes_bytes = (num_nodes * size_of::()) as u64; + let nodes = new_shared_buffer(rt, nodes_bytes); + + let command = rt.queue.new_command_buffer(); + let mut constants = Vec::with_capacity((layers + 1) * 2); + self.encode_hash_many_into_command( + &command, + &mut constants, + message_size, + messages, + 0, + num_leaves, + &nodes, + 0, + ); + + let mut previous_offset = 0usize; + let mut previous_len = num_leaves; + let mut next_offset = num_leaves; + for _ in 0..layers { + let current_len = previous_len / 2; + self.encode_hash_many_into_command( + &command, + &mut constants, + 64, + &nodes, + (previous_offset * size_of::()) as u64, + current_len, + &nodes, + (next_offset * size_of::()) as u64, + ); + previous_offset = next_offset; + previous_len = current_len; + next_offset += current_len; + } + wait_for_command_named(&command, "sha256_merkle_fused"); + drop(constants); + nodes + } + + fn hash_many_buffer_into_buffer( + &self, + size: usize, + input_buffer: &Buffer, + input_offset: u64, + output_len: usize, + output_buffer: &Buffer, + output_offset: u64, + ) { + if output_len == 0 { + return; + } + + let command = runtime().queue.new_command_buffer(); + let mut constants = Vec::with_capacity(2); + self.encode_hash_many_into_command( + &command, + &mut constants, + size, + input_buffer, + input_offset, + output_len, + output_buffer, + output_offset, + ); + wait_for_command(&command); + drop(constants); + } + + #[allow(clippy::too_many_arguments)] + fn encode_hash_many_into_command( + &self, + command: &metal::CommandBufferRef, + constants: &mut Vec, + size: usize, + input_buffer: &Buffer, + input_offset: u64, + output_len: usize, + output_buffer: &Buffer, + output_offset: u64, + ) { + if output_len == 0 { + return; + } + let rt = runtime(); + let messages_buffer = upload_u32(rt, output_len as u32); + let encoder = command.new_compute_command_encoder(); + let pipeline = if size == 64 && USE_SHA256_64_SPECIALIZATION { + encoder.set_compute_pipeline_state(&rt.pipeline_64); + encoder.set_buffer(0, Some(input_buffer), input_offset); + encoder.set_buffer(1, Some(output_buffer), output_offset); + encoder.set_buffer(2, Some(&messages_buffer), 0); + &rt.pipeline_64 + } else { + let size_buffer = upload_u32(rt, size as u32); + encoder.set_compute_pipeline_state(&rt.pipeline); + encoder.set_buffer(0, Some(input_buffer), input_offset); + encoder.set_buffer(1, Some(output_buffer), output_offset); + encoder.set_buffer(2, Some(&size_buffer), 0); + encoder.set_buffer(3, Some(&messages_buffer), 0); + constants.push(size_buffer); + &rt.pipeline + }; + constants.push(messages_buffer); + let group_width = pipeline + .thread_execution_width() + .min(output_len as u64) + .max(1); + encoder.dispatch_threads( + MTLSize { + width: output_len as u64, + height: 1, + depth: 1, + }, + MTLSize { + width: group_width, + height: 1, + depth: 1, + }, + ); + encoder.end_encoding(); + HASH_COUNTER.add(output_len); + } + + pub(crate) fn hash_many_buffer( + &self, + size: usize, + input_buffer: &Buffer, + output_len: usize, + output: &mut [Hash], + ) { + assert_eq!(output_len, output.len()); + if output.is_empty() { + return; + } + + let rt = runtime(); + + let output_bytes = (output.len() * size_of::()) as u64; + let output_buffer = new_shared_buffer(rt, output_bytes); + self.hash_many_buffer_into_buffer(size, input_buffer, 0, output.len(), &output_buffer, 0); + + let start = Instant::now(); + let bytes = unsafe { + std::slice::from_raw_parts(output_buffer.contents().cast::(), output.len() * 32) + }; + for (out, bytes) in output.iter_mut().zip(bytes.chunks_exact(32)) { + out.as_mut_bytes().copy_from_slice(bytes); + } + metal_profile::record_readback(output_bytes, start.elapsed()); + } +} + +fn pow_window_for_threshold(threshold: u64) -> u32 { + if threshold == 0 { + return POW_MAX_WINDOW; + } + + let expected_attempts = u64::MAX + .checked_div(threshold) + .unwrap_or(u64::MAX) + .saturating_add(1); + let target = expected_attempts.saturating_mul(4); + let window = target + .checked_next_power_of_two() + .unwrap_or(u64::from(POW_MAX_WINDOW)) + .clamp(u64::from(POW_MIN_WINDOW), u64::from(POW_MAX_WINDOW)); + window as u32 +} + +impl HashEngine for MetalSha2 { + fn name(&self) -> Cow<'_, str> { + "sha2".into() + } + + fn oid(&self) -> Option { + Some(Sha256::OID) + } + + fn supports_size(&self, _size: usize) -> bool { + Device::system_default().is_some() + } + + fn preferred_batch_size(&self) -> usize { + 256 + } + + fn hash_many(&self, size: usize, input: &[u8], output: &mut [Hash]) { + assert_eq!( + input.len(), + size * output.len(), + "Input length ({}) should be size * output.len() = {size} * {}", + input.len(), + output.len() + ); + if output.is_empty() { + return; + } + + let input_buffer = if input.is_empty() { + new_shared_buffer(runtime(), 0) + } else { + new_shared_buffer_with_data(runtime(), input.as_ptr().cast(), input.len() as u64) + }; + self.hash_many_buffer(size, &input_buffer, output.len(), output); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hash::Sha2; + + #[test] + fn metal_sha2_matches_cpu() { + let rows = 17; + let size = 73; + let input = (0..rows * size) + .map(|i: usize| (i.wrapping_mul(31) & 0xff) as u8) + .collect::>(); + let mut cpu = vec![Hash::default(); rows]; + let mut gpu = vec![Hash::default(); rows]; + Sha2::new().hash_many(size, &input, &mut cpu); + MetalSha2::new().hash_many(size, &input, &mut gpu); + assert_eq!(gpu, cpu); + } + + #[test] + fn metal_sha2_64_matches_cpu() { + let rows = 257; + let size = 64; + let input = (0..rows * size) + .map(|i: usize| (i.wrapping_mul(17).wrapping_add(3) & 0xff) as u8) + .collect::>(); + let mut cpu = vec![Hash::default(); rows]; + let mut gpu = vec![Hash::default(); rows]; + Sha2::new().hash_many(size, &input, &mut cpu); + MetalSha2::new().hash_many(size, &input, &mut gpu); + assert_eq!(gpu, cpu); + } + + #[test] + fn metal_sha2_pow_matches_threshold() { + let challenge = [7u8; 32]; + let threshold = u64::MAX >> 8; + let nonce = MetalSha2::prove_pow_64(&challenge, threshold); + let mut input = [0u8; 64]; + input[..32].copy_from_slice(&challenge); + input[32..40].copy_from_slice(&nonce.to_le_bytes()); + let mut hash = Hash::default(); + Sha2::new().hash_many(64, &input, std::slice::from_mut(&mut hash)); + let value = u64::from_le_bytes(hash.0[..8].try_into().unwrap()); + assert!(value <= threshold); + } +} diff --git a/src/hash/mod.rs b/src/hash/mod.rs index 53a7e84b..1d53c405 100644 --- a/src/hash/mod.rs +++ b/src/hash/mod.rs @@ -2,6 +2,10 @@ mod blake3_engine; mod copy_engine; mod digest_engine; mod hash_counter; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub(crate) mod metal_profile; +#[cfg(all(feature = "metal", target_os = "macos"))] +mod metal_sha2_engine; use core::fmt; use std::{ @@ -14,6 +18,10 @@ use serde::{Deserialize, Serialize}; use static_assertions::{assert_impl_all, assert_obj_safe}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned}; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub use self::metal_profile::{snapshot as metal_profile_snapshot, MetalProfileSnapshot}; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub use self::metal_sha2_engine::MetalSha2; pub use self::{ blake3_engine::{Blake3, BLAKE3}, copy_engine::{Copy, COPY}, @@ -32,6 +40,8 @@ pub static ENGINES: LazyLock> = LazyLock::new(|| { engines.register(Arc::new(Keccak::new())); engines.register(Arc::new(Sha3::new())); engines.register(Arc::new(Blake3::detect())); + #[cfg(all(feature = "metal", target_os = "macos"))] + engines.register(Arc::new(MetalSha2::new())); engines }); diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index d851a407..ef13bc7b 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -232,10 +232,38 @@ impl Config { #[cfg(feature = "tracing")] tracing::Span::current().record("engine", engine.name().as_ref()); + #[cfg(all(feature = "metal", target_os = "macos"))] + if self.leaf_hash_id == hash::SHA2 + && self + .merkle_tree + .layers + .iter() + .all(|layer| layer.hash_id == hash::SHA2) + { + if let Some(nodes) = matrix.commit_bn254_rows_sha2_merkle( + self.num_cols, + self.num_rows(), + self.merkle_tree.layers.len(), + ) { + let root = nodes + .read_hash_at(self.merkle_tree.num_nodes() - 1) + .expect("missing Metal Merkle root"); + prover_state.prover_message(&root); + return BufferWitness { nodes }; + } + } + // Compute leaf hashes let mut leaves = Vec::with_capacity(self.merkle_tree.num_nodes()); leaves.resize(self.merkle_tree.num_leaves, Hash::default()); - hash_rows(&*engine, matrix.as_slice(), &mut leaves[..self.num_rows()]); + #[cfg(all(feature = "metal", target_os = "macos"))] + let hashed_on_metal = self.leaf_hash_id == hash::SHA2 + && matrix.hash_bn254_rows_sha2(self.num_cols, &mut leaves[..self.num_rows()]); + #[cfg(not(all(feature = "metal", target_os = "macos")))] + let hashed_on_metal = false; + if !hashed_on_metal { + hash_rows(&*engine, matrix.as_slice(), &mut leaves[..self.num_rows()]); + } // Commit the leaf hashes BufferWitness { @@ -272,6 +300,19 @@ impl Config { R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>, { + #[cfg(all(feature = "metal", target_os = "macos"))] + if let Some(hints) = metal_merkle_opening_hints( + &witness.nodes, + self.merkle_tree.num_leaves, + self.merkle_tree.layers.len(), + indices, + ) { + for hint in hints { + prover_state.prover_hint(&hint); + } + return; + } + self.merkle_tree.open( prover_state, &merkle_tree::Witness { @@ -312,6 +353,49 @@ impl Config { } } +#[cfg(all(feature = "metal", target_os = "macos"))] +fn metal_merkle_opening_hints( + nodes: &ActiveBuffer, + num_leaves: usize, + layers: usize, + indices: &[usize], +) -> Option> { + if num_leaves != (1usize << layers) { + return None; + } + assert!(indices.iter().all(|&i| i < num_leaves)); + + let mut indices = indices.to_vec(); + indices.sort_unstable(); + indices.dedup(); + + let mut node_indices = Vec::new(); + let mut layer_offset = 0usize; + let mut layer_len = 1usize << layers; + while layer_len > 1 { + let mut next_indices = Vec::with_capacity(indices.len()); + let mut iter = indices.iter().copied().peekable(); + loop { + match (iter.next(), iter.peek()) { + (Some(a), Some(&b)) if b == a ^ 1 => { + next_indices.push(a >> 1); + iter.next(); + } + (Some(a), _) => { + node_indices.push(layer_offset + (a ^ 1)); + next_indices.push(a >> 1); + } + (None, _) => break, + } + } + indices = next_indices; + layer_offset += layer_len; + layer_len /= 2; + } + + nodes.read_hash_indices(&node_indices) +} + impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "MatrixCommit({} x {})", self.num_rows(), self.num_cols) @@ -522,4 +606,18 @@ pub(crate) mod tests { fn test_field256() { proptest::(); } + + #[test] + #[cfg(all(feature = "metal", target_os = "macos"))] + fn test_field256_sha2_metal_commit_open() { + test::( + StdRng::seed_from_u64(7), + hash::SHA2, + hash::SHA2, + 4, + 16, + 2, + &[0, 1, 3, 8, 8, 15], + ); + } } diff --git a/src/protocols/proof_of_work.rs b/src/protocols/proof_of_work.rs index a2a8aabc..ac6f10d7 100644 --- a/src/protocols/proof_of_work.rs +++ b/src/protocols/proof_of_work.rs @@ -20,6 +20,9 @@ use crate::{ verify, }; +#[cfg(all(feature = "metal", target_os = "macos"))] +use crate::hash::{MetalSha2, SHA2}; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct Config { pub hash_id: EngineId, @@ -87,6 +90,13 @@ impl Config { let challenge: [u8; 32] = prover_state.verifier_message(); + #[cfg(all(feature = "metal", target_os = "macos"))] + if self.hash_id == SHA2 { + let nonce = MetalSha2::prove_pow_64(&challenge, self.threshold); + prover_state.prover_message(&U64(nonce)); + return; + } + #[cfg(not(feature = "parallel"))] let nonce = (0_u64..) .step_by(batch_size) From f020e81f939e4c646fa9620255affe80eca3b033 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 11 Jun 2026 07:52:28 +0200 Subject: [PATCH 10/35] wip: metal buffer refactor and minor changes --- proptest-regressions/protocols/irs_commit.txt | 8 + src/algebra/buffer.rs | 8 + src/algebra/metal_buffer.rs | 563 +++++++++++++++--- src/bin/phase_benchmark.rs | 6 +- src/protocols/sumcheck.rs | 7 +- 5 files changed, 506 insertions(+), 86 deletions(-) create mode 100644 proptest-regressions/protocols/irs_commit.txt diff --git a/proptest-regressions/protocols/irs_commit.txt b/proptest-regressions/protocols/irs_commit.txt new file mode 100644 index 00000000..c02b8454 --- /dev/null +++ b/proptest-regressions/protocols/irs_commit.txt @@ -0,0 +1,8 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 33aef60251379abf1e94f5821e0fbd4f69d0b742576f047a92921eeb9438c0f8 # shrinks to seed = 0, config = Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 0, vector_size: 1, mask_length: 0, codeword_length: 1, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 0, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 1, layers: [] } }, johnson_slack: 0.0, in_domain_samples: 0, out_domain_samples: 0, deduplicate_in_domain: false } +cc 2c8600bcbca6f088a51752815a2baef31e96a3c8ae7e516c28b3a7c407904c3f # shrinks to seed = 0, config = Config { embedding: Basefield { extension: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 2 } }, num_vectors: 0, vector_size: 1, mask_length: 0, codeword_length: 1, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 0, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 1, layers: [] } }, johnson_slack: 0.0, in_domain_samples: 0, out_domain_samples: 0, deduplicate_in_domain: false } diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index e6a22abc..7430e66b 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -49,7 +49,15 @@ pub trait FieldOps: Clone { fn zero_pad(&mut self); fn dot(&self, other: &Self) -> F; fn fold(&mut self, weight: F); + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); + } fn sumcheck_polynomial(&self, other: &Self) -> (F, F); + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) + } fn accumulate_univariate_evaluations( &mut self, evaluators: &[UnivariateEvaluation], diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs index 615dac2d..77a36b2a 100644 --- a/src/algebra/metal_buffer.rs +++ b/src/algebra/metal_buffer.rs @@ -316,6 +316,26 @@ kernel void bn254_fold( store_f(values, gid, add_f(low, mul_f(sub_f(high, low), weight))); } +inline F fold_value_at(device const ulong *values, uint len, uint fold_half, uint idx, F weight) { + F low = idx < len ? load_f(values, idx) : zero_f(); + F high = idx + fold_half < len ? load_f(values, idx + fold_half) : zero_f(); + return add_f(low, mul_f(sub_f(high, low), weight)); +} + +kernel void bn254_fold_pair( + device ulong *a [[buffer(0)]], + device ulong *b [[buffer(1)]], + device const ulong *weight_buf [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &fold_half [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= fold_half) return; + F weight = load_f(weight_buf, 0); + store_f(a, gid, fold_value_at(a, len, fold_half, gid, weight)); + store_f(b, gid, fold_value_at(b, len, fold_half, gid, weight)); +} + kernel void bn254_scalar_mul_add( device ulong *acc [[buffer(0)]], device const ulong *vector [[buffer(1)]], @@ -402,6 +422,27 @@ kernel void bn254_sum_chunks( store_f(out, gid, acc); } +kernel void bn254_sumcheck_reduce_chunks( + device const ulong *input [[buffer(0)]], + device ulong *out [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &chunk_size [[buffer(3)]], + constant uint &out_len [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= out_len) return; + uint start = gid * chunk_size; + uint end = min(start + chunk_size, len); + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = start; i < end; i++) { + c0 = add_f(c0, load_f(input, i)); + c2 = add_f(c2, load_f(input, len + i)); + } + store_f(out, gid, c0); + store_f(out, out_len + gid, c2); +} + kernel void bn254_sumcheck_chunks( device const ulong *a [[buffer(0)]], device const ulong *b [[buffer(1)]], @@ -429,6 +470,43 @@ kernel void bn254_sumcheck_chunks( store_f(out, partial_count + gid, c2); } +kernel void bn254_fold_pair_sumcheck_chunks( + device ulong *a [[buffer(0)]], + device ulong *b [[buffer(1)]], + device const ulong *weight_buf [[buffer(2)]], + device ulong *out [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &fold_half [[buffer(5)]], + constant uint &sum_half [[buffer(6)]], + constant uint &chunk_size [[buffer(7)]], + constant uint &partial_count [[buffer(8)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= sum_half) return; + uint end = min(start + chunk_size, sum_half); + F weight = load_f(weight_buf, 0); + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = start; i < end; i++) { + uint right = i + sum_half; + F a0 = fold_value_at(a, len, fold_half, i, weight); + F b0 = fold_value_at(b, len, fold_half, i, weight); + F a1 = right < fold_half ? fold_value_at(a, len, fold_half, right, weight) : zero_f(); + F b1 = right < fold_half ? fold_value_at(b, len, fold_half, right, weight) : zero_f(); + store_f(a, i, a0); + store_f(b, i, b0); + if (right < fold_half) { + store_f(a, right, a1); + store_f(b, right, b1); + } + c0 = add_f(c0, mul_f(a0, b0)); + c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); + } + store_f(out, gid, c0); + store_f(out, partial_count + gid, c2); +} + kernel void bn254_geometric_accumulate( device ulong *acc [[buffer(0)]], device const ulong *points [[buffer(1)]], @@ -483,17 +561,27 @@ kernel void bn254_geometric_accumulate_chunks_strided( uint start = gid * chunk_size; if (start >= len) return; uint end = min(start + chunk_size, len); + uint count = end - start; + + // Accumulate all points into registers so `acc` is read and written + // once per element instead of once per (element, point). + F sums[32]; + for (uint k = 0; k < count; k++) { + sums[k] = zero_f(); + } for (uint j = 0; j < num_points; j++) { F point = load_f(points, j); - F scalar = load_f(scalars, j); - F power = pow_f(load_f(point_steps, j), gid); - for (uint i = start; i < end; i++) { - F value = load_f(acc, i); - value = add_f(value, mul_f(scalar, power)); - store_f(acc, i, value); + // Fold the scalar into the running power so the inner loop needs a + // single multiplication per element. + F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), gid)); + for (uint k = 0; k < count; k++) { + sums[k] = add_f(sums[k], power); power = mul_f(power, point); } } + for (uint i = start, k = 0; i < end; i++, k++) { + store_f(acc, i, add_f(load_f(acc, i), sums[k])); + } } kernel void bn254_geometric_accumulate_point_blocks( @@ -517,17 +605,16 @@ kernel void bn254_geometric_accumulate_point_blocks( if (point_start >= num_points) return; uint point_end = min(point_start + point_block_size, num_points); - F sums[16]; + F sums[32]; for (uint k = 0; k < chunk_size; k++) { sums[k] = zero_f(); } for (uint j = point_start; j < point_end; j++) { F point = load_f(points, j); - F scalar = load_f(scalars, j); - F power = pow_f(load_f(point_steps, j), chunk); + F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); for (uint i = start, k = 0; i < end; i++, k++) { - sums[k] = add_f(sums[k], mul_f(scalar, power)); + sums[k] = add_f(sums[k], power); power = mul_f(power, point); } } @@ -538,6 +625,49 @@ kernel void bn254_geometric_accumulate_point_blocks( } } +kernel void bn254_geometric_accumulate_point_blocks_range( + device ulong *partials [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *point_steps [[buffer(2)]], + device const ulong *scalars [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &num_points [[buffer(5)]], + constant uint &chunk_size [[buffer(6)]], + constant uint &point_block_size [[buffer(7)]], + constant uint &point_block_offset [[buffer(8)]], + constant uint &batch_point_blocks [[buffer(9)]], + uint gid [[thread_position_in_grid]] +) { + uint chunk = gid / batch_point_blocks; + uint local_point_block = gid - chunk * batch_point_blocks; + uint global_point_block = point_block_offset + local_point_block; + uint start = chunk * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + uint point_start = global_point_block * point_block_size; + if (point_start >= num_points) return; + uint point_end = min(point_start + point_block_size, num_points); + + F sums[32]; + for (uint k = 0; k < chunk_size; k++) { + sums[k] = zero_f(); + } + + for (uint j = point_start; j < point_end; j++) { + F point = load_f(points, j); + F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); + for (uint i = start, k = 0; i < end; i++, k++) { + sums[k] = add_f(sums[k], power); + power = mul_f(power, point); + } + } + + uint partial_offset = local_point_block * len; + for (uint i = start, k = 0; i < end; i++, k++) { + store_f(partials, partial_offset + i, sums[k]); + } +} + kernel void bn254_geometric_accumulate_reduce_point_blocks( device ulong *acc [[buffer(0)]], device const ulong *partials [[buffer(1)]], @@ -820,11 +950,14 @@ kernel void bn254_multilinear_extend( const REDUCTION_CHUNK_SIZE: usize = 64; const SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 8; -const LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 16; +const LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; +/// Must match the `sums` register array size in the strided kernel. +const MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; const LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD: usize = 1 << 18; const GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE: usize = 16; const GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD: usize = 1 << 17; const GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS: usize = 64; +const GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES: usize = 256 << 20; #[derive(Clone, Debug)] struct MetalFieldBuffer { @@ -1060,7 +1193,14 @@ impl FieldOps for MetalBuffer { fn zeros(length: usize) -> Self { assert_bn254::(); - Self::from_vec(vec![F::ZERO; length]) + // Montgomery zero is all-zero bytes, so a device-side fill suffices. + Self { + len: length, + host_cache: OnceCell::new(), + field: Some(zeroed_field_buffer(length)), + hash: None, + _marker: PhantomData, + } } fn random(rng: &mut R, length: usize) -> Self @@ -1108,12 +1248,38 @@ impl FieldOps for MetalBuffer { self.invalidate_host_cache(); } + fn fold_pair(&mut self, other: &mut Self, weight: F) { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + if self.len() <= 1 { + return; + } + let len = self.len(); + let fold_half = len.next_power_of_two() >> 1; + let weight = upload_field(&[f_to_field256(weight)]); + let this = self.bn254_buffer(); + let other_buffer = other.bn254_buffer(); + run_in_place( + "bn254_fold_pair", + &[&this.limbs, &other_buffer.limbs, &weight.limbs], + &[len as u32, fold_half as u32], + fold_half, + ); + self.len = fold_half; + other.len = fold_half; + self.invalidate_host_cache(); + other.invalidate_host_cache(); + } + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { assert_bn254::(); let len = self.len().min(other.len()); if len == 0 { return (F::ZERO, F::ZERO); } + if len == 1 { + return (self.as_slice()[0] * other.as_slice()[0], F::ZERO); + } let fold_half = len.next_power_of_two() >> 1; let this = self.bn254_buffer(); let other = other.bn254_buffer(); @@ -1121,6 +1287,29 @@ impl FieldOps for MetalBuffer { (field256_to_f::(c0), field256_to_f::(c2)) } + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + if self.len() <= 1 { + return self.sumcheck_polynomial(other); + } + let len = self.len(); + let fold_half = len.next_power_of_two() >> 1; + if fold_half == 1 { + self.fold_pair(other, weight); + return self.sumcheck_polynomial(other); + } + let weight = upload_field(&[f_to_field256(weight)]); + let this = self.bn254_buffer(); + let other_buffer = other.bn254_buffer(); + let (c0, c2) = parallel_fold_pair_sumcheck(&this, &other_buffer, &weight, len, fold_half); + self.len = fold_half; + other.len = fold_half; + self.invalidate_host_cache(); + other.invalidate_host_cache(); + (field256_to_f::(c0), field256_to_f::(c2)) + } + fn accumulate_univariate_evaluations( &mut self, evaluators: &[UnivariateEvaluation], @@ -1180,6 +1369,20 @@ impl FieldOps for MetalBuffer { evaluators.len(), chunk_size, ); + } else if should_use_geometric_point_blocks_batched( + self.len(), + evaluators.len(), + chunk_size, + ) { + parallel_geometric_accumulate_point_blocks_batched( + &field, + &points, + &point_steps, + &scalars, + self.len(), + evaluators.len(), + chunk_size, + ); } else { run_in_place( "bn254_geometric_accumulate_chunks_strided", @@ -1347,7 +1550,7 @@ impl MetalBuffer { fn bn254_buffer(&self) -> MetalFieldBuffer { self.field .clone() - .unwrap_or_else(|| upload_field(&to_field256_slice(self.as_slice()))) + .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) } fn invalidate_host_cache(&mut self) { @@ -1389,7 +1592,7 @@ impl MetalBuffer { fn bn254_buffer_target(&self) -> MetalFieldBuffer { self.field .clone() - .unwrap_or_else(|| upload_field(&to_field256_slice(self.as_slice()))) + .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) } } @@ -1397,16 +1600,20 @@ struct MetalRuntime { device: Device, queue: CommandQueue, fold: ComputePipelineState, + fold_pair: ComputePipelineState, scalar_mul_add: ComputePipelineState, dot: ComputePipelineState, sumcheck: ComputePipelineState, dot_chunks: ComputePipelineState, sum_chunks: ComputePipelineState, + sumcheck_reduce_chunks: ComputePipelineState, sumcheck_chunks: ComputePipelineState, + fold_pair_sumcheck_chunks: ComputePipelineState, geometric_accumulate: ComputePipelineState, geometric_accumulate_chunks: ComputePipelineState, geometric_accumulate_chunks_strided: ComputePipelineState, geometric_accumulate_point_blocks: ComputePipelineState, + geometric_accumulate_point_blocks_range: ComputePipelineState, geometric_accumulate_reduce_point_blocks: ComputePipelineState, univariate_evaluate: ComputePipelineState, univariate_eval_chunks: ComputePipelineState, @@ -1444,18 +1651,24 @@ fn runtime() -> &'static MetalRuntime { MetalRuntime { queue: device.new_command_queue(), fold: pipeline("bn254_fold"), + fold_pair: pipeline("bn254_fold_pair"), scalar_mul_add: pipeline("bn254_scalar_mul_add"), dot: pipeline("bn254_dot"), sumcheck: pipeline("bn254_sumcheck"), dot_chunks: pipeline("bn254_dot_chunks"), sum_chunks: pipeline("bn254_sum_chunks"), + sumcheck_reduce_chunks: pipeline("bn254_sumcheck_reduce_chunks"), sumcheck_chunks: pipeline("bn254_sumcheck_chunks"), + fold_pair_sumcheck_chunks: pipeline("bn254_fold_pair_sumcheck_chunks"), geometric_accumulate: pipeline("bn254_geometric_accumulate"), geometric_accumulate_chunks: pipeline("bn254_geometric_accumulate_chunks"), geometric_accumulate_chunks_strided: pipeline( "bn254_geometric_accumulate_chunks_strided", ), geometric_accumulate_point_blocks: pipeline("bn254_geometric_accumulate_point_blocks"), + geometric_accumulate_point_blocks_range: pipeline( + "bn254_geometric_accumulate_point_blocks_range", + ), geometric_accumulate_reduce_point_blocks: pipeline( "bn254_geometric_accumulate_reduce_point_blocks", ), @@ -1485,16 +1698,22 @@ fn runtime() -> &'static MetalRuntime { fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { match name { "bn254_fold" => &rt.fold, + "bn254_fold_pair" => &rt.fold_pair, "bn254_scalar_mul_add" => &rt.scalar_mul_add, "bn254_dot" => &rt.dot, "bn254_sumcheck" => &rt.sumcheck, "bn254_dot_chunks" => &rt.dot_chunks, "bn254_sum_chunks" => &rt.sum_chunks, + "bn254_sumcheck_reduce_chunks" => &rt.sumcheck_reduce_chunks, "bn254_sumcheck_chunks" => &rt.sumcheck_chunks, + "bn254_fold_pair_sumcheck_chunks" => &rt.fold_pair_sumcheck_chunks, "bn254_geometric_accumulate" => &rt.geometric_accumulate, "bn254_geometric_accumulate_chunks" => &rt.geometric_accumulate_chunks, "bn254_geometric_accumulate_chunks_strided" => &rt.geometric_accumulate_chunks_strided, "bn254_geometric_accumulate_point_blocks" => &rt.geometric_accumulate_point_blocks, + "bn254_geometric_accumulate_point_blocks_range" => { + &rt.geometric_accumulate_point_blocks_range + } "bn254_geometric_accumulate_reduce_point_blocks" => { &rt.geometric_accumulate_reduce_point_blocks } @@ -1623,13 +1842,17 @@ fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field let partials = MetalFieldBuffer { limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), }; - run_in_place( - "bn254_dot_chunks", + let command = rt.queue.new_command_buffer(); + encode_u32_kernel( + command, + pipeline(rt, "bn254_dot_chunks"), &[&a.limbs, &b.limbs, &partials.limbs], &[len as u32, REDUCTION_CHUNK_SIZE as u32], partial_count, ); - reduce_field_buffer(&partials, partial_count, 0) + let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); + wait_for_command_named(command, "bn254_dot_chunks"); + download_field_at(&result.limbs, offset) } fn parallel_sumcheck( @@ -1643,8 +1866,10 @@ fn parallel_sumcheck( let partials = MetalFieldBuffer { limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), }; - run_in_place( - "bn254_sumcheck_chunks", + let command = rt.queue.new_command_buffer(); + encode_u32_kernel( + command, + pipeline(rt, "bn254_sumcheck_chunks"), &[&a.limbs, &b.limbs, &partials.limbs], &[ len as u32, @@ -1654,10 +1879,42 @@ fn parallel_sumcheck( ], partial_count, ); - ( - reduce_field_buffer(&partials, partial_count, 0), - reduce_field_buffer(&partials, partial_count, partial_count), - ) + let result = encode_sumcheck_reduction(command, partials, partial_count); + wait_for_command_named(command, "bn254_sumcheck_chunks"); + download_sumcheck_pair(&result) +} + +fn parallel_fold_pair_sumcheck( + a: &MetalFieldBuffer, + b: &MetalFieldBuffer, + weight: &MetalFieldBuffer, + len: usize, + fold_half: usize, +) -> (Field256, Field256) { + let sum_half = fold_half.next_power_of_two() >> 1; + debug_assert!(sum_half > 0); + let partial_count = sum_half.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), + }; + let command = rt.queue.new_command_buffer(); + encode_u32_kernel( + command, + pipeline(rt, "bn254_fold_pair_sumcheck_chunks"), + &[&a.limbs, &b.limbs, &weight.limbs, &partials.limbs], + &[ + len as u32, + fold_half as u32, + sum_half as u32, + REDUCTION_CHUNK_SIZE as u32, + partial_count as u32, + ], + partial_count, + ); + let result = encode_sumcheck_reduction(command, partials, partial_count); + wait_for_command_named(command, "bn254_fold_pair_sumcheck_chunks"); + download_sumcheck_pair(&result) } fn parallel_univariate_evaluate( @@ -1673,13 +1930,17 @@ fn parallel_univariate_evaluate( let partials = MetalFieldBuffer { limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), }; - run_in_place( - "bn254_univariate_eval_chunks", + let command = rt.queue.new_command_buffer(); + encode_u32_kernel( + command, + pipeline(rt, "bn254_univariate_eval_chunks"), &[&coeffs.limbs, &point.limbs, &partials.limbs], &[len as u32, REDUCTION_CHUNK_SIZE as u32], partial_count, ); - reduce_field_buffer(&partials, partial_count, 0) + let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); + wait_for_command_named(command, "bn254_univariate_eval_chunks"); + download_field_at(&result.limbs, offset) } fn parallel_geometric_accumulate_point_blocks( @@ -1691,7 +1952,7 @@ fn parallel_geometric_accumulate_point_blocks( num_points: usize, chunk_size: usize, ) { - assert!(chunk_size <= 16); + assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); let chunks = len.div_ceil(chunk_size); let partial_len = point_blocks @@ -1726,6 +1987,65 @@ fn parallel_geometric_accumulate_point_blocks( ); } +fn parallel_geometric_accumulate_point_blocks_batched( + acc: &MetalFieldBuffer, + points: &MetalFieldBuffer, + point_steps: &MetalFieldBuffer, + scalars: &MetalFieldBuffer, + len: usize, + num_points: usize, + chunk_size: usize, +) { + assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); + let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); + let chunks = len.div_ceil(chunk_size); + let bytes_per_point_block = len + .checked_mul(size_of::()) + .expect("Metal geometric partial size overflow"); + let default_batch_blocks = + (GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES / bytes_per_point_block).max(1); + let batch_blocks = std::env::var("WHIR_METAL_GEOM_POINT_BLOCK_BATCH") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&value| value > 0) + .unwrap_or(default_batch_blocks) + .min(point_blocks); + let rt = runtime(); + for point_block_offset in (0..point_blocks).step_by(batch_blocks) { + let current_batch = batch_blocks.min(point_blocks - point_block_offset); + let partial_len = current_batch + .checked_mul(len) + .expect("Metal geometric partial size overflow"); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), + }; + run_in_place( + "bn254_geometric_accumulate_point_blocks_range", + &[ + &partials.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[ + len as u32, + num_points as u32, + chunk_size as u32, + GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, + point_block_offset as u32, + current_batch as u32, + ], + chunks * current_batch, + ); + run_in_place( + "bn254_geometric_accumulate_reduce_point_blocks", + &[&acc.limbs, &partials.limbs], + &[len as u32, current_batch as u32], + len, + ); + } +} + fn geometric_accumulate_chunk_size(len: usize) -> usize { std::env::var("WHIR_METAL_GEOM_CHUNK") .ok() @@ -1736,29 +2056,45 @@ fn geometric_accumulate_chunk_size(len: usize) -> usize { } else { SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE }) + .min(MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE) } fn should_use_geometric_point_blocks(len: usize, num_points: usize, chunk_size: usize) -> bool { - chunk_size <= 16 + chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS && len <= GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD } -fn reduce_field_buffer(input: &MetalFieldBuffer, len: usize, offset: usize) -> Field256 { - if len == 0 { - return Field256::ZERO; - } - let mut current = input.clone(); +fn should_use_geometric_point_blocks_batched( + len: usize, + num_points: usize, + chunk_size: usize, +) -> bool { + chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS + && len > GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD +} + +/// Encodes all tree-reduction levels into `command` and returns the buffer +/// and offset holding the final scalar. Runs with a single command wait. +fn encode_field_reduction( + command: &metal::CommandBufferRef, + input: MetalFieldBuffer, + len: usize, + offset: usize, +) -> (MetalFieldBuffer, usize) { + let rt = runtime(); + let mut current = input; let mut current_len = len; let mut current_offset = offset; while current_len > 1 { let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); let next = MetalFieldBuffer { limbs: new_shared_buffer(rt, (next_len * size_of::()) as u64), }; - run_in_place( - "bn254_sum_chunks", + encode_u32_kernel( + command, + pipeline(rt, "bn254_sum_chunks"), &[¤t.limbs, &next.limbs], &[ current_len as u32, @@ -1771,7 +2107,44 @@ fn reduce_field_buffer(input: &MetalFieldBuffer, len: usize, offset: usize) -> F current_len = next_len; current_offset = 0; } - download_field_at(¤t.limbs, current_offset) + (current, current_offset) +} + +/// Encodes all (c0, c2) tree-reduction levels into `command` and returns the +/// buffer holding the final pair. Runs with a single command wait. +fn encode_sumcheck_reduction( + command: &metal::CommandBufferRef, + input: MetalFieldBuffer, + len: usize, +) -> MetalFieldBuffer { + let rt = runtime(); + let mut current = input; + let mut current_len = len; + while current_len > 1 { + let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); + let next = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (next_len * 2 * size_of::()) as u64), + }; + encode_u32_kernel( + command, + pipeline(rt, "bn254_sumcheck_reduce_chunks"), + &[¤t.limbs, &next.limbs], + &[ + current_len as u32, + REDUCTION_CHUNK_SIZE as u32, + next_len as u32, + ], + next_len, + ); + current = next; + current_len = next_len; + } + current +} + +fn download_sumcheck_pair(buffer: &MetalFieldBuffer) -> (Field256, Field256) { + let values = download_field(&buffer.limbs, 2); + (values[0], values[1]) } fn encode_single_vector_coset_ntt( @@ -2029,25 +2402,37 @@ fn read_bn254_rows(source: &MetalFieldBuffer, num_cols: usize, indices: &[usize] } fn upload_field(values: &[Field256]) -> MetalFieldBuffer { - let limbs = values - .iter() - .flat_map(|value| value.0 .0) - .collect::>(); let rt = runtime(); - let buffer = if limbs.is_empty() { + let buffer = if values.is_empty() { new_shared_buffer(rt, 0) } else { + // Field256 is 4 contiguous u64 limbs; upload directly without + // flattening into an intermediate Vec. new_shared_buffer_with_data( rt, - limbs.as_ptr().cast(), - (limbs.len() * size_of::()) as u64, + values.as_ptr().cast(), + std::mem::size_of_val(values) as u64, ) }; MetalFieldBuffer { limbs: buffer } } +fn zeroed_field_buffer(len: usize) -> MetalFieldBuffer { + let rt = runtime(); + let bytes = (len * size_of::()) as u64; + let buffer = new_shared_buffer(rt, bytes); + if bytes > 0 { + let command = rt.queue.new_command_buffer(); + let blit = command.new_blit_command_encoder(); + blit.fill_buffer(&buffer, metal::NSRange::new(0, bytes), 0); + blit.end_encoding(); + wait_for_blit(command, bytes); + } + MetalFieldBuffer { limbs: buffer } +} + fn maybe_upload_bn254(values: &[T]) -> Option { - (type_name::() == type_name::()).then(|| upload_field(&to_field256_slice(values))) + (type_name::() == type_name::()).then(|| upload_field(as_field256_slice(values))) } fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer { @@ -2062,34 +2447,39 @@ fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer MetalFieldBuffer { limbs: target } } -fn upload_u32(value: u32) -> Buffer { - let rt = runtime(); - new_shared_buffer_with_data(rt, (&value as *const u32).cast(), size_of::() as u64) -} - -fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { - let rt = runtime(); - let command = rt.queue.new_command_buffer(); +/// Encodes a kernel dispatch with `u32` constants bound as inline bytes +/// (no per-constant buffer allocations). +fn encode_u32_kernel( + command: &metal::CommandBufferRef, + pipeline: &ComputePipelineState, + buffers: &[&Buffer], + constants: &[u32], + threads: usize, +) { let encoder = command.new_compute_command_encoder(); - let pipeline = pipeline(rt, name); encoder.set_compute_pipeline_state(pipeline); let mut index = 0; for buffer in buffers { encoder.set_buffer(index, Some(buffer), 0); index += 1; } - let constant_buffers = constants - .iter() - .copied() - .map(upload_u32) - .collect::>(); - for buffer in &constant_buffers { - encoder.set_buffer(index, Some(buffer), 0); + for constant in constants { + encoder.set_bytes( + index, + size_of::() as u64, + (constant as *const u32).cast(), + ); index += 1; } dispatch(&encoder, pipeline, threads.max(1)); encoder.end_encoding(); - wait_for_command_named(&command, name); +} + +fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + encode_u32_kernel(command, pipeline(rt, name), buffers, constants, threads); + wait_for_command_named(command, name); } fn run_reduce(name: &str, buffers: &[&Buffer], constants: &[u32], out_len: usize) -> Vec { @@ -2106,7 +2496,9 @@ fn dispatch( pipeline: &ComputePipelineState, threads: usize, ) { - let width = pipeline.thread_execution_width().max(1); + // Use full threadgroups (capped at 256) instead of a single execution + // width; the pipeline limit already accounts for register pressure. + let width = pipeline.max_total_threads_per_threadgroup().clamp(1, 256); let group_width = width.min(threads as u64).max(1); encoder.dispatch_threads( MTLSize { @@ -2199,16 +2591,13 @@ fn field256_to_target(value: Field256) -> T { unsafe { std::mem::transmute_copy(&value) } } -fn to_field256_slice(values: &[T]) -> Vec { +fn as_field256_slice(values: &[T]) -> &[Field256] { assert_eq!( type_name::(), type_name::(), "MetalBuffer only supports BN254 Field256 buffers" ); - values - .iter() - .map(|value| unsafe { std::mem::transmute_copy(value) }) - .collect() + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } } #[cfg(test)] @@ -2245,16 +2634,34 @@ mod tests { #[test] fn metal_bn254_sumcheck_matches_cpu() { - let a = values(27, 3); - let b = values(27, 11); - let cpu_a = CpuBuffer::from_slice(&a); - let cpu_b = CpuBuffer::from_slice(&b); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_b = MetalBuffer::from_slice(&b); - assert_eq!( - gpu_a.sumcheck_polynomial(&gpu_b), - cpu_a.sumcheck_polynomial(&cpu_b) - ); + for len in [1, 2, 27, 64, 65] { + let a = values(len, 3); + let b = values(len, 11); + let cpu_a = CpuBuffer::from_slice(&a); + let cpu_b = CpuBuffer::from_slice(&b); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_b = MetalBuffer::from_slice(&b); + assert_eq!( + gpu_a.sumcheck_polynomial(&gpu_b), + cpu_a.sumcheck_polynomial(&cpu_b) + ); + } + } + + #[test] + fn metal_bn254_fold_pair_sumcheck_matches_cpu() { + for len in [2, 3, 27, 64, 65] { + let mut cpu_a = CpuBuffer::from_vec(values(len, 3)); + let mut cpu_b = CpuBuffer::from_vec(values(len, 11)); + let mut gpu_a = MetalBuffer::from_vec(values(len, 3)); + let mut gpu_b = MetalBuffer::from_vec(values(len, 11)); + let weight = Field256::from(42); + let cpu_result = cpu_a.fold_pair_sumcheck_polynomial(&mut cpu_b, weight); + let gpu_result = gpu_a.fold_pair_sumcheck_polynomial(&mut gpu_b, weight); + assert_eq!(gpu_result, cpu_result); + assert_eq!(gpu_a.as_slice(), cpu_a.as_slice()); + assert_eq!(gpu_b.as_slice(), cpu_b.as_slice()); + } } #[test] diff --git a/src/bin/phase_benchmark.rs b/src/bin/phase_benchmark.rs index 38f7f22e..782cb53d 100644 --- a/src/bin/phase_benchmark.rs +++ b/src/bin/phase_benchmark.rs @@ -301,12 +301,12 @@ fn bench_sumcheck( let ds = DomainSeparator::protocol(&config) .session(&"phase benchmark sumcheck") .instance(&Empty); + let mut a = ActiveBuffer::from_vec(input_vector(size)); + let mut b = ActiveBuffer::from_vec((0..size).map(|i| F::from(i as u64 + 17)).collect()); + let mut sum = a.dot(&b); HASH_COUNTER.reset(); let measured = measure_phase(|| { - let mut a = ActiveBuffer::from_vec(input_vector(size)); - let mut b = ActiveBuffer::from_vec((0..size).map(|i| F::from(i as u64 + 17)).collect()); - let mut sum = a.dot(&b); let mut prover_state = ProverState::new_std(&ds); let _ = black_box(config.prove(&mut prover_state, &mut a, &mut b, &mut sum, &[])); }); diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 31f30b5e..4bef3bbc 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -110,9 +110,7 @@ impl Config { { // Fold and compute sumcheck polynomial in one pass. let (c0, c2) = if let Some(w) = prev_round_challenge { - a.fold(w); - b.fold(w); - a.sumcheck_polynomial(b) + a.fold_pair_sumcheck_polynomial(b, w) } else { a.sumcheck_polynomial(b) }; @@ -152,8 +150,7 @@ impl Config { } if let Some(w) = prev_round_challenge { // Final fold of the inputs (no polynomial computation) - a.fold(w); - b.fold(w); + a.fold_pair(b, w); } *sum = mask_sum + mask_rlc * *sum; From dfddb0e064d9e6b25bdf6aaaf78f0a16aa6099ad Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Mon, 15 Jun 2026 12:17:59 +0200 Subject: [PATCH 11/35] feat: add experimental Metal backend --- docs/buffer-prove-design.md | 284 ------------- scripts/phase_benchmark_cpu_gpu.sh | 118 ------ scripts/plot_phase_benchmark.py | 631 ----------------------------- src/algebra/metal_buffer.rs | 8 +- src/bin/phase_benchmark.rs | 504 ----------------------- src/bin/phase_benchmark_report.rs | 127 ------ src/hash/metal_profile.rs | 25 ++ src/hash/metal_sha2_engine.rs | 8 +- src/hash/mod.rs | 5 +- 9 files changed, 41 insertions(+), 1669 deletions(-) delete mode 100644 docs/buffer-prove-design.md delete mode 100644 scripts/phase_benchmark_cpu_gpu.sh delete mode 100644 scripts/plot_phase_benchmark.py delete mode 100644 src/bin/phase_benchmark.rs delete mode 100644 src/bin/phase_benchmark_report.rs diff --git a/docs/buffer-prove-design.md b/docs/buffer-prove-design.md deleted file mode 100644 index 3d08229f..00000000 --- a/docs/buffer-prove-design.md +++ /dev/null @@ -1,284 +0,0 @@ -# Buffer Abstraction Design - -## Purpose - -The prover manipulates large vectors, encoded Reed-Solomon matrices, and Merkle -node arrays. On CPU, those objects can be ordinary `Vec`s. On an accelerator, -asking for a slice of the same object can force a synchronization and a full -readback. - -The abstraction exists to make ownership explicit: - -```text -Large prover data lives in an active backend object. -Protocol code only reads proof-sized data. -``` - -Proof-sized data means transcript scalars, Merkle roots, selected rows, selected -authentication nodes, and final folded vectors that are sent directly into the -proof. Full vectors, full encoded matrices, and full Merkle trees should stay in -`ActiveBuffer` or `ActiveMatrix`. - -## Design Rules - -The buffer layer is a data and math boundary. It should not own protocol logic. - -```text -Protocol layer: - transcript order - challenge sampling - proof layout - round structure - hash choice - verification semantics - -Buffer layer: - backend-owned vectors - backend-owned matrices - full-buffer math loops - selected reads - explicit full reads -``` - -That split rules out protocol-specific buffer APIs. WHIR should keep one prove -path, IRS should keep one commit path, and buffers should not write to the -transcript. Each protocol should place backend-owned data only at the boundary -where that protocol actually consumes large resident data. - -## Active Types - -Protocol code refers to active aliases instead of hard-coding a backend: - -```rust -pub type ActiveBuffer = selected_backend::Buffer; -pub type ActiveMatrix = selected_backend::Matrix; -``` - -The current selected backend is CPU, but that is a local implementation detail -inside the buffer module. The aliases are intentionally central. WHIR, IRS, -sumcheck, matrix commitment, and Merkle opening operate on `ActiveBuffer` and -`ActiveMatrix`, so a later backend can be selected at this boundary without -changing protocol callsites. - -The active types also expose the explicit construction and read boundaries used -by protocol code: - -```rust -// Setup or protocol-owned construction. -ActiveBuffer::from_slice(...); -ActiveBuffer::from_vec(...); -ActiveBuffer::zeros(...); -ActiveBuffer::random(...); -ActiveMatrix::from_vec(...); - -// Proof-boundary materialization. -buffer.read(); -buffer.read_index(...); -buffer.read_indices(...); -matrix.read_rows(...); -``` - -These are intentionally not part of `BufferOps`. They are ownership and -materialization boundaries, while `BufferOps` is only the full-buffer math -surface. - -## Core Trait - -`BufferOps` covers operations that are naturally full-buffer operations. Loops -over protocol rounds remain outside the trait. - -```rust -pub trait BufferOps: Clone { - // Same backend family for another field. - type Buffer: BufferOps; - - // Metadata. - fn len(&self) -> usize; - fn is_empty(&self) -> bool; - - // In-place full-buffer transforms. - fn fold(&mut self, weight: F); - - // Full-buffer reductions used by sumcheck and WHIR. - fn sumcheck_polynomial(&self, other: &Self) -> (F, F); - fn fold_and_sumcheck_polynomial( - &mut self, - other: &mut Self, - weight: F, - ) -> (F, F); - - // Mixed-field full-buffer reductions used by protocol checks. - fn mixed_dot>( - &self, - embedding: &M, - other: &Self::Buffer, - ) -> M::Target; - - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target; - - // Writes into backend-owned state instead of returning a large vector. - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut Self::Buffer, - weight: M::Target, - ); - - fn geometric_accumulate(&mut self, scalars: &[F], points: &[F]); -} -``` - -Construction and host materialization are explicit active-type boundaries, not -part of this math trait. Protocol code may upload data at commitment setup and -may read proof-sized data at transcript/opening boundaries, but those actions -should stay visible at the call site. - -Scalar-returning methods are still synchronization points. That is acceptable -when the scalar is transcript data. If a later backend kernel should consume the -result directly, the right design is a fused operation or a backend-owned scalar, -not a hidden readback followed by an upload. - -## Matrix Shape - -`ActiveMatrix` is the backend-owned row matrix produced by interleaved RS -encoding and consumed by matrix commitment. - -The protocol-visible matrix surface is intentionally small: shape metadata and -selected logical row readback. - -The caller owns the protocol metadata: row count, column count, query indices, -and how selected rows are written into the proof. The backend owns physical -layout. For example, an accelerated matrix may store rows in a different order -or column-major internally, but `read_rows(indices)` returns logical rows. - -## IRS Boundary - -IRS remains protocol code. It samples masks, calls the RS encoder, commits the -encoded rows, samples OOD points, and records OOD evaluations. - -The shared RS primitive is: - -```rust -pub fn interleaved_rs_encode( - vectors: &[&ActiveBuffer], - masks: &ActiveBuffer, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, -) -> ActiveMatrix; -``` - -This replaces the old call shape that required `&[&[F]]` and a mask slice. IRS -does not need to expose host slices just to encode. - -IRS witness names describe what is actually stored: - -```rust -pub struct Witness -where - G: Field, -{ - pub masks: ActiveBuffer, - pub encoded_matrix: ActiveMatrix, - pub encoded_matrix_witness: matrix_commit::Witness, - pub out_of_domain: Evaluations, -} -``` - -The masks and encoded matrix remain backend-owned. OOD evaluations are still -host vectors because they are proof/transcript data in the current protocol. - -## Matrix Commitment Boundary - -Matrix commitment owns the hash choice and transcript write. The buffer layer -does not commit rows and does not select a hash function. - -The commit path is: - -```text -ActiveMatrix - -> matrix_commit::Config::build_nodes(...) - -> ActiveBuffer - -> merkle_tree::Config::commit_nodes(...) - -> root written to transcript -``` - -The important ownership change is that the Merkle witness stores the node array -as an active hash buffer: - -```rust -pub struct Witness { - nodes: ActiveBuffer, -} -``` - -Opening a Merkle tree reads only the required authentication nodes. That keeps -the full tree resident and turns openings into proof-sized reads. - -## WHIR Prove Boundary - -WHIR `prove` accepts active buffers for the committed vectors and borrowed -witness references: - -```rust -pub fn prove<'a, H, R>( - &self, - prover_state: &mut ProverState, - vectors: &[&ActiveBuffer], - witnesses: &'a [&'a Witness], - linear_forms: Vec>>, - evaluations: Cow<'a, [M::Target]>, -) -> FinalClaim -``` - -The WHIR prover still owns the protocol decisions: - -```text -sample vector RLC coefficients -build a backend-owned linear combination -sample constraint RLC coefficients -build the covector from linear forms -add OOD and STIR constraints -run sumcheck rounds -commit and open IRS rounds -send the final folded vector -``` - -The buffer layer only performs the large vector work inside those steps: - -```rust -vector.mixed_scalar_mul_add_to(...); -covector.geometric_accumulate(...); -sumcheck.prove(..., &mut vector, &mut covector, ...); -``` - -## Readback Policy - -Expected readbacks are proof-sized: Merkle roots, authentication nodes, selected -IRS rows, transcript scalars, and the final folded WHIR vector. Suspicious -readbacks are full committed vectors, full encoded matrices, full Merkle trees, -or a readback immediately followed by upload into another backend buffer. - -The current CPU implementation may borrow slices inside CPU-only primitives, -but that access stays crate-private. IRS, WHIR, and matrix commitment callers -should see active buffers and explicit proof-boundary reads. - -## Open Work - -Scalar-returning reductions currently return host scalars. That is fine for -transcript values, but later backend work may need fused operations or -backend-owned scalar handles when the scalar feeds another backend computation. - -The ZK wrapper remains a host-side wrapper for now. Its outer API accepts host -slices/owned vectors, builds masked and blinding vectors on the host, and only -constructs active buffers at the point where it calls the inner WHIR -commit/prove paths. That keeps this refactor focused on the regular WHIR -residency boundary without pushing buffer requirements through unfinished ZK -code. - -Future backend selection should happen behind `ActiveBuffer` and `ActiveMatrix`. -Protocol APIs should not grow backend-specific entrypoints. diff --git a/scripts/phase_benchmark_cpu_gpu.sh b/scripts/phase_benchmark_cpu_gpu.sh deleted file mode 100644 index 028998eb..00000000 --- a/scripts/phase_benchmark_cpu_gpu.sh +++ /dev/null @@ -1,118 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -out="${WHIR_PHASE_BENCH_OUTPUT:-outputs/phase_benchmark_cpu_gpu.jsonl}" -min_log_size=16 -max_log_size=28 -folds="1,2,3,4,6" -rates="1,2,3" -phases="commit,sumcheck,e2e" -article_grid=false -common_args=() -while [[ $# -gt 0 ]]; do - case "$1" in - --output) - out="$2" - shift 2 - ;; - --output=*) - out="${1#--output=}" - shift - ;; - --min-log-size) - min_log_size="$2" - shift 2 - ;; - --min-log-size=*) - min_log_size="${1#--min-log-size=}" - shift - ;; - --max-log-size) - max_log_size="$2" - shift 2 - ;; - --max-log-size=*) - max_log_size="${1#--max-log-size=}" - shift - ;; - --folds) - folds="$2" - shift 2 - ;; - --folds=*) - folds="${1#--folds=}" - shift - ;; - --rates) - rates="$2" - shift 2 - ;; - --rates=*) - rates="${1#--rates=}" - shift - ;; - --phases) - phases="$2" - shift 2 - ;; - --phases=*) - phases="${1#--phases=}" - shift - ;; - --article-grid) - article_grid=true - shift - ;; - *) - common_args+=("$1") - shift - ;; - esac -done - -mkdir -p "$(dirname "$out")" -: > "$out" - -cargo build --release --bin phase_benchmark -CARGO_TARGET_DIR=target/metal cargo build --release --features metal --bin phase_benchmark -cargo build --release --bin phase_benchmark_report - -IFS=, read -r -a fold_values <<< "$folds" -IFS=, read -r -a rate_values <<< "$rates" -IFS=, read -r -a phase_values <<< "$phases" - -cpu_bin="target/release/phase_benchmark" -gpu_bin="target/metal/release/phase_benchmark" - -for log_size in $(seq "$min_log_size" "$max_log_size"); do - for fold in "${fold_values[@]}"; do - if [[ "$fold" -gt "$log_size" ]]; then - printf 'skip n=%s fold=%s: fold must be <= n\n' "$log_size" "$fold" >&2 - continue - fi - for rate in "${rate_values[@]}"; do - if [[ "$article_grid" == true && "$log_size" -ge 24 && "$rate" -gt 1 ]]; then - continue - fi - - for phase in "${phase_values[@]}"; do - case_args=( - --output "$out" - --min-log-size "$log_size" - --max-log-size "$log_size" - --folds "$fold" - --rates "$rate" - --phases "$phase" - ) - - "$cpu_bin" "${case_args[@]}" "${common_args[@]}" - "$gpu_bin" "${case_args[@]}" "${common_args[@]}" - done - done - done -done - -target/release/phase_benchmark_report "$out" > "${out%.jsonl}.csv" - -printf 'wrote %s\n' "$out" -printf 'wrote %s\n' "${out%.jsonl}.csv" diff --git a/scripts/plot_phase_benchmark.py b/scripts/plot_phase_benchmark.py deleted file mode 100644 index 4fea4422..00000000 --- a/scripts/plot_phase_benchmark.py +++ /dev/null @@ -1,631 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import csv -import html -import json -import math -from collections import defaultdict -from pathlib import Path - - -PHASES = ["commit", "sumcheck", "e2e_prove"] -PHASE_LABEL = { - "commit": "commit", - "sumcheck": "sumcheck", - "e2e_prove": "e2e prove", -} -BACKEND_LABEL = { - "cpu": "CPU", - "gpu-metal": "GPU Metal", -} -BACKEND_COLOR = { - "cpu": "#2f4b7c", - "gpu-metal": "#d95f02", -} -FOLD_COLOR = { - 1: "#2f4b7c", - 2: "#d95f02", - 3: "#1b9e77", - 4: "#7570b3", - 6: "#e7298a", -} -PROFILE_COLOR = { - "total": "#1f2937", - "command wait": "#7570b3", - "upload": "#d95f02", - "readback": "#1b9e77", - "blit": "#e7298a", -} - - -def load_rows(path): - rows = [] - with open(path, "r", encoding="utf-8") as f: - for line in f: - line = line.strip() - if line: - rows.append(json.loads(line)) - return rows - - -def row_index(rows): - out = {} - for r in rows: - key = (r["backend"], r["phase"], r["log_size"], r["fold"], r["rate"]) - out[key] = r - return out - - -def paired_rows(rows): - idx = row_index(rows) - keys = sorted( - { - (r["phase"], r["log_size"], r["fold"], r["rate"]) - for r in rows - }, - key=lambda k: (k[1], k[2], k[3], k[0]), - ) - pairs = [] - for phase, log_size, fold, rate in keys: - cpu = idx.get(("cpu", phase, log_size, fold, rate)) - gpu = idx.get(("gpu-metal", phase, log_size, fold, rate)) - if cpu or gpu: - pairs.append((phase, log_size, fold, rate, cpu, gpu)) - return pairs - - -def gib(x): - return x / 1024.0 / 1024.0 / 1024.0 - - -def mib(x): - return x / 1024.0 / 1024.0 - - -def fmt_ms(ms): - if ms is None: - return "OOM" - if ms < 1000: - return f"{ms:.1f} ms" - return f"{ms / 1000:.2f} s" - - -def fmt_speedup(x): - if x is None: - return "OOM" - return f"{x:.2f}x" - - -def fmt_gib(x): - if x is None: - return "OOM" - return f"{gib(x):.1f}" - - -def nice_ticks(vmin, vmax, log_y=False, count=5): - if not math.isfinite(vmin) or not math.isfinite(vmax) or vmax <= 0: - return [] - if log_y: - lo = math.floor(math.log10(max(vmin, 1e-9))) - hi = math.ceil(math.log10(vmax)) - ticks = [] - for p in range(lo, hi + 1): - for m in (1, 2, 5): - value = m * (10 ** p) - if vmin <= value <= vmax: - ticks.append(value) - if len(ticks) > 7: - stride = max(1, math.ceil(len(ticks) / 7)) - ticks = ticks[::stride] - return ticks - if vmax == vmin: - return [vmin] - raw = (vmax - vmin) / max(1, count - 1) - mag = 10 ** math.floor(math.log10(raw)) - step = min((1, 2, 5, 10), key=lambda s: abs(raw - s * mag)) * mag - start = math.floor(vmin / step) * step - ticks = [] - value = start - while value <= vmax + step * 0.5: - if value >= vmin - step * 0.5: - ticks.append(value) - value += step - return ticks[:8] - - -def svg_text(x, y, text, size=12, anchor="middle", weight="400", fill="#111827"): - text = html.escape(str(text)) - return ( - f'{text}' - ) - - -def svg_poly(points, color, width=2.2, dash=False): - if len(points) < 2: - return "" - pts = " ".join(f"{x:.1f},{y:.1f}" for x, y in points) - dash_attr = ' stroke-dasharray="5 4"' if dash else "" - return ( - f'' - ) - - -def render_line_chart( - path, - title, - panels, - y_label, - x_label="log2(size)", - log_y=False, - shared_y=True, - width=1280, - panel_height=260, -): - panel_count = len(panels) - height = 96 + panel_count * panel_height + 28 - margin_l, margin_r, margin_t, margin_b = 76, 24, 74, 46 - plot_w = width - margin_l - margin_r - all_x = [ - x - for panel in panels - for series in panel["series"].values() - for x, y in series - if y is not None and y > 0 - ] - xmin, xmax = min(all_x), max(all_x) - - def y_range(panel): - values = [ - y - for series in panel["series"].values() - for x, y in series - if y is not None and y > 0 and math.isfinite(y) - ] - if not values: - return 0.1, 1.0 - lo, hi = min(values), max(values) - if log_y: - return lo / 1.35, hi * 1.35 - pad = (hi - lo) * 0.12 if hi > lo else max(1.0, hi * 0.1) - return max(0.0, lo - pad), hi + pad - - shared_range = None - if shared_y: - vals = [] - for panel in panels: - a, b = y_range(panel) - vals.extend([a, b]) - shared_range = min(vals), max(vals) - - parts = [ - f'', - '', - svg_text(width / 2, 32, title, 21, weight="700"), - svg_text(width / 2, height - 8, x_label, 12), - svg_text(18, height / 2, y_label, 12, anchor="middle"), - ] - parts.append(f'') - - legend_items = {} - for panel in panels: - for name, color in panel.get("colors", {}).items(): - legend_items[name] = color - lx = margin_l - ly = 54 - for name, color in legend_items.items(): - parts.append(f'') - parts.append(svg_text(lx + 34, ly + 4, name, 12, anchor="start")) - lx += 150 - - for i, panel in enumerate(panels): - top = margin_t + i * panel_height - bottom = top + panel_height - margin_b - y0 = top + 26 - plot_h = bottom - y0 - vmin, vmax = shared_range if shared_y else y_range(panel) - if log_y: - vmin = max(vmin, 1e-9) - lmin, lmax = math.log10(vmin), math.log10(vmax) - - def sy(v): - return bottom - (math.log10(max(v, 1e-9)) - lmin) / (lmax - lmin) * plot_h - - else: - - def sy(v): - return bottom - (v - vmin) / (vmax - vmin) * plot_h if vmax > vmin else bottom - - def sx(x): - return margin_l + (x - xmin) / (xmax - xmin) * plot_w if xmax > xmin else margin_l + plot_w / 2 - - parts.append(svg_text(margin_l, top + 14, panel["title"], 15, anchor="start", weight="700")) - parts.append(f'') - parts.append(f'') - - x_span = int(xmax) - int(xmin) - x_step = 1 if x_span <= 18 else 2 if x_span <= 36 else 4 - for x in range(int(xmin), int(xmax) + 1, x_step): - px = sx(x) - parts.append(f'') - parts.append(svg_text(px, bottom + 20, str(x), 11)) - for tick in nice_ticks(vmin, vmax, log_y=log_y): - py = sy(tick) - parts.append(f'') - label = f"{tick:g}" if tick < 1000 else f"{tick / 1000:g}k" - parts.append(svg_text(margin_l - 9, py + 4, label, 10, anchor="end", fill="#374151")) - - for name, series in panel["series"].items(): - color = panel.get("colors", {}).get(name, "#111827") - points = [(sx(x), sy(y)) for x, y in series if y is not None and y > 0] - parts.append(svg_poly(points, color)) - for px, py in points: - parts.append(f'') - - for note in panel.get("annotations", []): - x, y, text = note - parts.append(svg_text(sx(x), sy(y), text, 11, fill="#b91c1c", weight="700")) - - parts.append("") - Path(path).write_text("\n".join(parts), encoding="utf-8") - - -def write_csv(path, rows): - pairs = paired_rows(rows) - with open(path, "w", newline="", encoding="utf-8") as f: - writer = csv.writer(f) - writer.writerow( - [ - "log_size", - "size", - "fold", - "rate", - "phase", - "cpu_ms", - "gpu_ms", - "speedup", - "cpu_peak_gib", - "gpu_peak_gib", - "gpu_upload_gib", - "gpu_upload_ms", - "gpu_readback_mib", - "gpu_readback_ms", - "gpu_command_wait_ms", - "gpu_blit_gib", - "gpu_blit_wait_ms", - ] - ) - for phase, log_size, fold, rate, cpu, gpu in pairs: - cpu_ms = cpu and cpu.get("duration_ms") - gpu_ms = gpu and gpu.get("duration_ms") - speedup = cpu_ms / gpu_ms if cpu_ms and gpu_ms else None - writer.writerow( - [ - log_size, - 1 << log_size, - fold, - rate, - phase, - f"{cpu_ms:.6f}" if cpu_ms else "", - f"{gpu_ms:.6f}" if gpu_ms else "", - f"{speedup:.6f}" if speedup else "", - f"{gib(cpu['peak_allocated_bytes']):.6f}" if cpu else "", - f"{gib(gpu['peak_allocated_bytes']):.6f}" if gpu else "", - f"{gib(gpu.get('metal_upload_bytes', 0)):.6f}" if gpu else "", - f"{gpu.get('metal_upload_ms', 0):.6f}" if gpu else "", - f"{mib(gpu.get('metal_readback_bytes', 0)):.6f}" if gpu else "", - f"{gpu.get('metal_readback_ms', 0):.6f}" if gpu else "", - f"{gpu.get('metal_command_wait_ms', 0):.6f}" if gpu else "", - f"{gib(gpu.get('metal_blit_bytes', 0)):.6f}" if gpu else "", - f"{gpu.get('metal_blit_wait_ms', 0):.6f}" if gpu else "", - ] - ) - - -def markdown_table(headers, rows): - out = ["| " + " | ".join(headers) + " |"] - out.append("| " + " | ".join(["---"] * len(headers)) + " |") - for row in rows: - out.append("| " + " | ".join(str(x) for x in row) + " |") - return "\n".join(out) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("jsonl") - parser.add_argument("--plot-dir", default="outputs/plots") - parser.add_argument("--report", default="outputs/article_phase_benchmark_16_27_hybrid_report.md") - parser.add_argument("--paired-csv", default="outputs/article_phase_benchmark_16_27_hybrid_paired.csv") - parser.add_argument("--requested-min-log-size", type=int) - parser.add_argument("--requested-max-log-size", type=int) - args = parser.parse_args() - - rows = load_rows(args.jsonl) - idx = row_index(rows) - plot_dir = Path(args.plot_dir) - plot_dir.mkdir(parents=True, exist_ok=True) - write_csv(args.paired_csv, rows) - - logs = sorted({r["log_size"] for r in rows}) - folds = sorted({r["fold"] for r in rows}) - rates = sorted({r["rate"] for r in rows}) - - def get(backend, phase, log_size, fold, rate): - return idx.get((backend, phase, log_size, fold, rate)) - - def series_backend(backend, phase, fold, rate, metric): - points = [] - for log_size in logs: - r = get(backend, phase, log_size, fold, rate) - points.append((log_size, metric(r) if r else None)) - return points - - e2e_panels = [] - for fold in folds: - annotations = [] - if not get("cpu", "e2e_prove", 27, fold, 1) or not get("gpu-metal", "e2e_prove", 27, fold, 1): - max_y = max( - [ - r["duration_ms"] - for b in ("cpu", "gpu-metal") - for log_size in logs - for r in [get(b, "e2e_prove", log_size, fold, 1)] - if r - ] - or [1] - ) - annotations.append((27, max_y, "missing/killed")) - e2e_panels.append( - { - "title": f"fold={fold}, rate=1", - "series": { - "CPU": series_backend("cpu", "e2e_prove", fold, 1, lambda r: r["duration_ms"] if r else None), - "GPU Metal": series_backend("gpu-metal", "e2e_prove", fold, 1, lambda r: r["duration_ms"] if r else None), - }, - "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}, - "annotations": annotations, - } - ) - render_line_chart( - plot_dir / "e2e_runtime_rate1.svg", - "E2E prove runtime, rate=1", - e2e_panels, - "milliseconds, log scale", - log_y=True, - shared_y=True, - ) - - speed_series = {} - for fold in folds: - pts = [] - for log_size in logs: - cpu = get("cpu", "e2e_prove", log_size, fold, 1) - gpu = get("gpu-metal", "e2e_prove", log_size, fold, 1) - pts.append((log_size, cpu["duration_ms"] / gpu["duration_ms"] if cpu and gpu else None)) - speed_series[f"fold={fold}"] = pts - render_line_chart( - plot_dir / "e2e_speedup_rate1.svg", - "E2E prove speedup, rate=1", - [{"title": "CPU time / GPU Metal time", "series": speed_series, "colors": {f"fold={f}": FOLD_COLOR[f] for f in folds}}], - "speedup", - log_y=False, - shared_y=False, - panel_height=360, - ) - - phase_panels = [] - for phase in PHASES: - phase_panels.append( - { - "title": f"{PHASE_LABEL[phase]}, fold=1, rate=1", - "series": { - "CPU": series_backend("cpu", phase, 1, 1, lambda r: r["duration_ms"] if r else None), - "GPU Metal": series_backend("gpu-metal", phase, 1, 1, lambda r: r["duration_ms"] if r else None), - }, - "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}, - } - ) - render_line_chart( - plot_dir / "phase_runtime_fold1_rate1.svg", - "Phase runtime, fold=1, rate=1", - phase_panels, - "milliseconds, log scale", - log_y=True, - shared_y=False, - ) - - mem_panels = [] - for fold in folds: - mem_panels.append( - { - "title": f"fold={fold}, rate=1", - "series": { - "CPU": series_backend("cpu", "e2e_prove", fold, 1, lambda r: gib(r["peak_allocated_bytes"]) if r else None), - "GPU Metal": series_backend("gpu-metal", "e2e_prove", fold, 1, lambda r: gib(r["peak_allocated_bytes"]) if r else None), - }, - "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}, - } - ) - render_line_chart( - plot_dir / "peak_memory_e2e_rate1.svg", - "Peak allocated memory during E2E prove, rate=1", - mem_panels, - "GiB", - log_y=False, - shared_y=True, - ) - - max_mem = {"CPU": [], "GPU Metal": []} - for log_size in logs: - cpu_vals = [ - r["peak_allocated_bytes"] - for fold in folds - for r in [get("cpu", "e2e_prove", log_size, fold, 1)] - if r - ] - gpu_vals = [ - r["peak_allocated_bytes"] - for fold in folds - for r in [get("gpu-metal", "e2e_prove", log_size, fold, 1)] - if r - ] - max_mem["CPU"].append((log_size, gib(max(cpu_vals)) if cpu_vals else None)) - max_mem["GPU Metal"].append((log_size, gib(max(gpu_vals)) if gpu_vals else None)) - render_line_chart( - plot_dir / "peak_memory_max_e2e_rate1.svg", - "Max peak allocated memory by size, E2E prove rate=1", - [{"title": "max across successful folds", "series": max_mem, "colors": {"CPU": BACKEND_COLOR["cpu"], "GPU Metal": BACKEND_COLOR["gpu-metal"]}}], - "GiB", - log_y=False, - shared_y=False, - panel_height=360, - ) - - profile = defaultdict(list) - for log_size in logs: - r = get("gpu-metal", "e2e_prove", log_size, 4, 1) - profile["total"].append((log_size, r["duration_ms"] if r else None)) - profile["command wait"].append((log_size, r.get("metal_command_wait_ms", 0) if r else None)) - profile["upload"].append((log_size, r.get("metal_upload_ms", 0) if r else None)) - profile["readback"].append((log_size, r.get("metal_readback_ms", 0) if r else None)) - profile["blit"].append((log_size, r.get("metal_blit_wait_ms", 0) if r else None)) - render_line_chart( - plot_dir / "gpu_profile_e2e_fold4_rate1.svg", - "GPU E2E profile counters, fold=4, rate=1", - [{"title": "wall time components", "series": dict(profile), "colors": PROFILE_COLOR}], - "milliseconds, log scale", - log_y=True, - shared_y=False, - panel_height=360, - ) - - speedups = [] - for fold in folds: - for log_size in logs: - cpu = get("cpu", "e2e_prove", log_size, fold, 1) - gpu = get("gpu-metal", "e2e_prove", log_size, fold, 1) - if cpu and gpu: - speedups.append(cpu["duration_ms"] / gpu["duration_ms"]) - speedups_sorted = sorted(speedups) - median_speedup = speedups_sorted[len(speedups_sorted) // 2] if speedups_sorted else None - - e2e_speed_table = [] - for log_size in logs: - row = [f"2^{log_size}"] - for fold in folds: - if fold > log_size: - row.append("n/a") - continue - cpu = get("cpu", "e2e_prove", log_size, fold, 1) - gpu = get("gpu-metal", "e2e_prove", log_size, fold, 1) - row.append(fmt_speedup(cpu["duration_ms"] / gpu["duration_ms"]) if cpu and gpu else "missing") - e2e_speed_table.append(row) - - mem_table = [] - for log_size in logs: - cpu_vals = [ - r["peak_allocated_bytes"] - for fold in folds - for r in [get("cpu", "e2e_prove", log_size, fold, 1)] - if r - ] - gpu_vals = [ - r["peak_allocated_bytes"] - for fold in folds - for r in [get("gpu-metal", "e2e_prove", log_size, fold, 1)] - if r - ] - c = max(cpu_vals) if cpu_vals else None - g = max(gpu_vals) if gpu_vals else None - mem_table.append([f"2^{log_size}", fmt_gib(c), fmt_gib(g), f"{c / g:.2f}x" if c and g else "OOM"]) - - phase_table = [] - phase_sample_logs = [x for x in [8, 16, 20, 24, 27] if x in logs] - for log_size in phase_sample_logs: - for phase in PHASES: - cpu = get("cpu", phase, log_size, 4, 1) - gpu = get("gpu-metal", phase, log_size, 4, 1) - phase_table.append( - [ - f"2^{log_size}", - PHASE_LABEL[phase], - fmt_ms(cpu["duration_ms"] if cpu else None), - fmt_ms(gpu["duration_ms"] if gpu else None), - fmt_speedup(cpu["duration_ms"] / gpu["duration_ms"] if cpu and gpu else None), - ] - ) - - profile_table = [] - for log_size in phase_sample_logs: - r = get("gpu-metal", "e2e_prove", log_size, 4, 1) - profile_table.append( - [ - f"2^{log_size}", - fmt_ms(r["duration_ms"] if r else None), - fmt_ms(r.get("metal_command_wait_ms") if r else None), - f"{gib(r.get('metal_upload_bytes', 0)):.2f} GiB / {r.get('metal_upload_ms', 0):.1f} ms" if r else "OOM", - f"{mib(r.get('metal_readback_bytes', 0)):.2f} MiB / {r.get('metal_readback_ms', 0):.1f} ms" if r else "OOM", - f"{gib(r.get('metal_blit_bytes', 0)):.2f} GiB / {r.get('metal_blit_wait_ms', 0):.1f} ms" if r else "OOM", - ] - ) - - expected = { - (r["log_size"], r["fold"], r["rate"]) - for r in rows - } - missing = [] - for log_size, fold, rate in sorted(expected): - for phase in PHASES: - if not get("cpu", phase, log_size, fold, rate): - missing.append(f"CPU {PHASE_LABEL[phase]} 2^{log_size} fold={fold} rate={rate}") - if not get("gpu-metal", phase, log_size, fold, rate): - missing.append(f"GPU {PHASE_LABEL[phase]} 2^{log_size} fold={fold} rate={rate}") - - report = [] - report.append("# WHIR CPU vs Metal GPU phase benchmark") - report.append("") - requested_min = args.requested_min_log_size if args.requested_min_log_size is not None else min(logs) - requested_max = args.requested_max_log_size if args.requested_max_log_size is not None else max(logs) - report.append("Hardware: Apple M4 Max, 40-core GPU, 48 GiB unified memory.") - report.append(f"Parameters: requested `log_size={requested_min}..{requested_max}`, actual rows `log_size={min(logs)}..{max(logs)}`, folds `1,2,3,4,6`, rates `1,2,3` where the article grid fits memory and `rate=1` for the largest sizes, phases `commit,sumcheck,e2e`, `pow_bits=20`, `security_level=128`.") - if requested_min < min(logs): - report.append(f"`2^{requested_min}` has no rows because the benchmark rejects the configured folds there (`fold` must be `<= n`).") - report.append("") - report.append(f"Raw rows: `{len(rows)}`. Paired CSV: `{args.paired_csv}`.") - if missing: - report.append("Missing/killed rows: " + "; ".join(missing) + ".") - else: - report.append("Missing/killed rows: none.") - if speedups: - report.append(f"E2E rate=1 paired speedup: min `{min(speedups):.2f}x`, median `{median_speedup:.2f}x`, max `{max(speedups):.2f}x`.") - report.append("") - report.append("## Charts") - for name in [ - "e2e_runtime_rate1.svg", - "e2e_speedup_rate1.svg", - "phase_runtime_fold1_rate1.svg", - "peak_memory_e2e_rate1.svg", - "peak_memory_max_e2e_rate1.svg", - "gpu_profile_e2e_fold4_rate1.svg", - ]: - report.append(f"- `{plot_dir / name}`") - report.append("") - report.append("## E2E speedup, rate=1") - report.append(markdown_table(["size"] + [f"fold={f}" for f in folds], e2e_speed_table)) - report.append("") - report.append("## Max peak allocated memory, E2E rate=1") - report.append(markdown_table(["size", "CPU GiB", "GPU GiB", "CPU/GPU"], mem_table)) - report.append("") - report.append("## Phase speedup, fold=4 rate=1") - report.append(markdown_table(["size", "phase", "CPU", "GPU", "CPU/GPU"], phase_table)) - report.append("") - report.append("## GPU transfer/profile counters, E2E fold=4 rate=1") - report.append(markdown_table(["size", "total", "command wait", "upload", "readback", "blit"], profile_table)) - report.append("") - report.append("Notes: standalone `sumcheck` measures the phase in isolation, so large rows include host-to-GPU upload cost. The E2E path keeps more data resident and is the better signal for real prover throughput. Peak memory is allocator peak from the benchmark process; Metal counters are the explicit upload/readback/blit instrumentation emitted by the benchmark.") - Path(args.report).write_text("\n".join(report) + "\n", encoding="utf-8") - - -if __name__ == "__main__": - main() diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs index 77a36b2a..2e5c225b 100644 --- a/src/algebra/metal_buffer.rs +++ b/src/algebra/metal_buffer.rs @@ -1737,8 +1737,11 @@ fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { fn new_shared_buffer(rt: &MetalRuntime, bytes: u64) -> Buffer { metal_profile::record_alloc(bytes); - rt.device - .new_buffer(bytes, MTLResourceOptions::StorageModeShared) + let buffer = rt + .device + .new_buffer(bytes, MTLResourceOptions::StorageModeShared); + metal_profile::record_device_allocated(rt.device.current_allocated_size()); + buffer } fn new_shared_buffer_with_data(rt: &MetalRuntime, data: *const c_void, bytes: u64) -> Buffer { @@ -1748,6 +1751,7 @@ fn new_shared_buffer_with_data(rt: &MetalRuntime, data: *const c_void, bytes: u6 .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); metal_profile::record_alloc(bytes); metal_profile::record_upload(bytes, start.elapsed()); + metal_profile::record_device_allocated(rt.device.current_allocated_size()); buffer } diff --git a/src/bin/phase_benchmark.rs b/src/bin/phase_benchmark.rs deleted file mode 100644 index 782cb53d..00000000 --- a/src/bin/phase_benchmark.rs +++ /dev/null @@ -1,504 +0,0 @@ -use std::{ - alloc::{GlobalAlloc, Layout, System}, - borrow::Cow, - fs::OpenOptions, - hint::black_box, - io::Write, - sync::atomic::{AtomicUsize, Ordering}, - time::{Duration, Instant}, -}; - -use clap::Parser; -use serde::Serialize; -use whir::{ - algebra::{ - buffer::{ActiveBuffer, FieldOps}, - embedding::Identity, - fields::Field256, - }, - hash::{self, HASH_COUNTER}, - parameters::ProtocolParameters, - protocols::whir::Config as WhirConfig, - transcript::{codecs::Empty, DomainSeparator, ProverState}, -}; - -#[global_allocator] -static ALLOC: TrackingAllocator = TrackingAllocator; - -static CURRENT_ALLOCATED: AtomicUsize = AtomicUsize::new(0); -static PEAK_ALLOCATED: AtomicUsize = AtomicUsize::new(0); - -struct TrackingAllocator; - -unsafe impl GlobalAlloc for TrackingAllocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - let ptr = unsafe { System.alloc(layout) }; - if !ptr.is_null() { - add_allocated(layout.size()); - } - ptr - } - - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - unsafe { System.dealloc(ptr, layout) }; - CURRENT_ALLOCATED.fetch_sub(layout.size(), Ordering::Relaxed); - } - - unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; - if !new_ptr.is_null() { - match new_size.cmp(&layout.size()) { - std::cmp::Ordering::Greater => add_allocated(new_size - layout.size()), - std::cmp::Ordering::Less => { - CURRENT_ALLOCATED.fetch_sub(layout.size() - new_size, Ordering::Relaxed); - } - std::cmp::Ordering::Equal => {} - } - } - new_ptr - } -} - -fn add_allocated(bytes: usize) { - let current = CURRENT_ALLOCATED.fetch_add(bytes, Ordering::Relaxed) + bytes; - let mut peak = PEAK_ALLOCATED.load(Ordering::Relaxed); - while current > peak { - match PEAK_ALLOCATED.compare_exchange_weak( - peak, - current, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => break, - Err(next) => peak = next, - } - } -} - -#[derive(Parser, Debug)] -#[command(author, version, about = "Phase benchmark for CPU/GPU WHIR proving")] -struct Args { - #[arg(long, default_value_t = 16)] - min_log_size: usize, - - #[arg(long, default_value_t = 28)] - max_log_size: usize, - - #[arg(long, default_value = "1,2,3,4,6")] - folds: String, - - #[arg(long, default_value = "1,2,3")] - rates: String, - - /// Apply the article-style cap where rates above 1 are skipped for n >= 24. - #[arg(long)] - article_grid: bool, - - #[arg(long, default_value = "commit,sumcheck,e2e")] - phases: String, - - #[arg(long, default_value_t = 128)] - security_level: usize, - - #[arg(long, default_value_t = 20)] - pow_bits: usize, - - #[arg(long, default_value_t = false)] - unique_decoding: bool, - - #[arg(long, default_value = "outputs/phase_benchmark.jsonl")] - output: String, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Phase { - Commit, - Sumcheck, - E2e, -} - -#[derive(Serialize)] -struct PhaseRow { - backend: &'static str, - phase: &'static str, - log_size: usize, - size: usize, - fold: usize, - rate: usize, - duration_ms: f64, - current_allocated_bytes: usize, - peak_allocated_bytes: usize, - peak_phase_delta_bytes: usize, - hashes: usize, - proof_bytes: Option, - status: &'static str, - metal_upload_count: u64, - metal_upload_bytes: u64, - metal_upload_ms: f64, - metal_readback_count: u64, - metal_readback_bytes: u64, - metal_readback_ms: f64, - metal_alloc_count: u64, - metal_alloc_bytes: u64, - metal_command_count: u64, - metal_command_wait_ms: f64, - metal_blit_count: u64, - metal_blit_bytes: u64, - metal_blit_wait_ms: f64, -} - -struct Measurement { - value: T, - duration: Duration, - current_allocated_bytes: usize, - peak_allocated_bytes: usize, - peak_phase_delta_bytes: usize, - metal: MetalPhaseProfile, -} - -#[derive(Clone, Copy, Debug, Default)] -struct MetalPhaseProfile { - upload_count: u64, - upload_bytes: u64, - upload_ms: f64, - readback_count: u64, - readback_bytes: u64, - readback_ms: f64, - alloc_count: u64, - alloc_bytes: u64, - command_count: u64, - command_wait_ms: f64, - blit_count: u64, - blit_bytes: u64, - blit_wait_ms: f64, -} - -#[cfg(all(feature = "metal", target_os = "macos"))] -const BACKEND: &str = "gpu-metal"; - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -const BACKEND: &str = "cpu"; - -type F = Field256; -type M = Identity; - -fn main() { - let args = Args::parse(); - assert!(args.min_log_size <= args.max_log_size); - warm_backend(); - - let folds = parse_usize_list(&args.folds); - let rates = parse_usize_list(&args.rates); - let phases = parse_phases(&args.phases); - - std::fs::create_dir_all("outputs").unwrap(); - let mut output = OpenOptions::new() - .append(true) - .create(true) - .open(&args.output) - .unwrap(); - - for log_size in args.min_log_size..=args.max_log_size { - for &fold in &folds { - if fold > log_size { - eprintln!("skip n={log_size} fold={fold}: fold must be <= n"); - continue; - } - for &rate in &rates { - if args.article_grid && log_size >= 24 && rate > 1 { - continue; - } - let size = 1usize << log_size; - let params = whir_params(&args, fold, rate); - for &phase in &phases { - let row = match phase { - Phase::Commit => bench_commit(log_size, size, fold, rate, ¶ms), - Phase::Sumcheck => bench_sumcheck(log_size, size, fold, rate, ¶ms), - Phase::E2e => bench_e2e(log_size, size, fold, rate, ¶ms), - }; - writeln!(output, "{}", serde_json::to_string(&row).unwrap()).unwrap(); - output.flush().unwrap(); - eprintln!( - "{} n={log_size} fold={fold} rate={rate} phase={} {:.3} ms peak={} MiB upload={} KiB/{:.3} ms readback={} KiB/{:.3} ms cmd_wait={:.3} ms", - BACKEND, - row.phase, - row.duration_ms, - row.peak_allocated_bytes / (1024 * 1024), - row.metal_upload_bytes / 1024, - row.metal_upload_ms, - row.metal_readback_bytes / 1024, - row.metal_readback_ms, - row.metal_command_wait_ms, - ); - } - } - } - } -} - -#[cfg(all(feature = "metal", target_os = "macos"))] -fn warm_backend() { - whir::algebra::buffer::MetalBuffer::::warmup(); - whir::hash::MetalSha2::warmup(); -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -fn warm_backend() {} - -fn whir_params(args: &Args, fold: usize, rate: usize) -> ProtocolParameters { - ProtocolParameters { - security_level: args.security_level, - pow_bits: args.pow_bits, - initial_folding_factor: fold, - folding_factor: fold, - unique_decoding: args.unique_decoding, - starting_log_inv_rate: rate, - batch_size: 1, - hash_id: hash::SHA2, - } -} - -fn bench_commit( - log_size: usize, - size: usize, - fold: usize, - rate: usize, - whir_params: &ProtocolParameters, -) -> PhaseRow { - let vector = input_vector(size); - let vector_buffer = ActiveBuffer::from_slice(&vector); - let params = WhirConfig::::new(size, whir_params); - let ds = DomainSeparator::protocol(¶ms) - .session(&"phase benchmark commit") - .instance(&Empty); - - HASH_COUNTER.reset(); - let measured = measure_phase(|| { - let mut prover_state = ProverState::new_std(&ds); - let _ = black_box(params.commit(&mut prover_state, &[&vector_buffer])); - }); - row( - Phase::Commit, - log_size, - size, - fold, - rate, - measured, - HASH_COUNTER.get(), - None, - ) -} - -fn bench_sumcheck( - log_size: usize, - size: usize, - fold: usize, - rate: usize, - whir_params: &ProtocolParameters, -) -> PhaseRow { - let params = WhirConfig::::new(size, whir_params); - let config = params.initial_sumcheck.clone(); - let ds = DomainSeparator::protocol(&config) - .session(&"phase benchmark sumcheck") - .instance(&Empty); - let mut a = ActiveBuffer::from_vec(input_vector(size)); - let mut b = ActiveBuffer::from_vec((0..size).map(|i| F::from(i as u64 + 17)).collect()); - let mut sum = a.dot(&b); - - HASH_COUNTER.reset(); - let measured = measure_phase(|| { - let mut prover_state = ProverState::new_std(&ds); - let _ = black_box(config.prove(&mut prover_state, &mut a, &mut b, &mut sum, &[])); - }); - row( - Phase::Sumcheck, - log_size, - size, - fold, - rate, - measured, - HASH_COUNTER.get(), - None, - ) -} - -fn bench_e2e( - log_size: usize, - size: usize, - fold: usize, - rate: usize, - whir_params: &ProtocolParameters, -) -> PhaseRow { - let vector = input_vector(size); - let vector_buffer = ActiveBuffer::from_slice(&vector); - let params = WhirConfig::::new(size, whir_params); - let ds = DomainSeparator::protocol(¶ms) - .session(&"phase benchmark e2e") - .instance(&Empty); - - HASH_COUNTER.reset(); - let measured = measure_phase(|| { - let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &[&vector_buffer]); - let _ = params.prove( - &mut prover_state, - &[&vector_buffer], - vec![&witness], - vec![], - Cow::Owned(vec![]), - ); - let proof = prover_state.proof(); - proof.narg_string.len() + proof.hints.len() - }); - let proof_bytes = measured.value; - row( - Phase::E2e, - log_size, - size, - fold, - rate, - measured, - HASH_COUNTER.get(), - Some(proof_bytes), - ) -} - -fn input_vector(size: usize) -> Vec { - (0..size).map(|i| F::from(i as u64)).collect() -} - -fn measure_phase(f: impl FnOnce() -> T) -> Measurement { - let before = CURRENT_ALLOCATED.load(Ordering::Relaxed); - PEAK_ALLOCATED.store(before, Ordering::Relaxed); - let metal_before = metal_profile_snapshot(); - let start = Instant::now(); - let value = f(); - let duration = start.elapsed(); - let metal_after = metal_profile_snapshot(); - let current = CURRENT_ALLOCATED.load(Ordering::Relaxed); - let peak = PEAK_ALLOCATED.load(Ordering::Relaxed); - Measurement { - value, - duration, - current_allocated_bytes: current, - peak_allocated_bytes: peak, - peak_phase_delta_bytes: peak.saturating_sub(before), - metal: metal_profile_delta(metal_before, metal_after), - } -} - -fn row( - phase: Phase, - log_size: usize, - size: usize, - fold: usize, - rate: usize, - measurement: Measurement, - hashes: usize, - proof_bytes: Option, -) -> PhaseRow { - PhaseRow { - backend: BACKEND, - phase: phase.name(), - log_size, - size, - fold, - rate, - duration_ms: measurement.duration.as_secs_f64() * 1_000.0, - current_allocated_bytes: measurement.current_allocated_bytes, - peak_allocated_bytes: measurement.peak_allocated_bytes, - peak_phase_delta_bytes: measurement.peak_phase_delta_bytes, - hashes, - proof_bytes, - status: "ok", - metal_upload_count: measurement.metal.upload_count, - metal_upload_bytes: measurement.metal.upload_bytes, - metal_upload_ms: measurement.metal.upload_ms, - metal_readback_count: measurement.metal.readback_count, - metal_readback_bytes: measurement.metal.readback_bytes, - metal_readback_ms: measurement.metal.readback_ms, - metal_alloc_count: measurement.metal.alloc_count, - metal_alloc_bytes: measurement.metal.alloc_bytes, - metal_command_count: measurement.metal.command_count, - metal_command_wait_ms: measurement.metal.command_wait_ms, - metal_blit_count: measurement.metal.blit_count, - metal_blit_bytes: measurement.metal.blit_bytes, - metal_blit_wait_ms: measurement.metal.blit_wait_ms, - } -} - -#[cfg(all(feature = "metal", target_os = "macos"))] -fn metal_profile_snapshot() -> whir::hash::MetalProfileSnapshot { - whir::hash::metal_profile_snapshot() -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -fn metal_profile_snapshot() -> MetalProfileSnapshotCompat { - MetalProfileSnapshotCompat::default() -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -#[derive(Clone, Copy, Debug, Default)] -struct MetalProfileSnapshotCompat; - -#[cfg(all(feature = "metal", target_os = "macos"))] -fn metal_profile_delta( - before: whir::hash::MetalProfileSnapshot, - after: whir::hash::MetalProfileSnapshot, -) -> MetalPhaseProfile { - let delta = after.delta_since(before); - MetalPhaseProfile { - upload_count: delta.upload_count, - upload_bytes: delta.upload_bytes, - upload_ms: delta.upload_ms(), - readback_count: delta.readback_count, - readback_bytes: delta.readback_bytes, - readback_ms: delta.readback_ms(), - alloc_count: delta.alloc_count, - alloc_bytes: delta.alloc_bytes, - command_count: delta.command_count, - command_wait_ms: delta.command_wait_ms(), - blit_count: delta.blit_count, - blit_bytes: delta.blit_bytes, - blit_wait_ms: delta.blit_wait_ms(), - } -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -fn metal_profile_delta( - _before: MetalProfileSnapshotCompat, - _after: MetalProfileSnapshotCompat, -) -> MetalPhaseProfile { - MetalPhaseProfile::default() -} - -impl Phase { - const fn name(self) -> &'static str { - match self { - Self::Commit => "commit", - Self::Sumcheck => "sumcheck", - Self::E2e => "e2e_prove", - } - } -} - -fn parse_usize_list(source: &str) -> Vec { - source - .split(',') - .filter(|part| !part.is_empty()) - .map(|part| part.parse().unwrap()) - .collect() -} - -fn parse_phases(source: &str) -> Vec { - source - .split(',') - .filter(|part| !part.is_empty()) - .map(|part| match part { - "commit" => Phase::Commit, - "sumcheck" => Phase::Sumcheck, - "e2e" | "e2e_prove" => Phase::E2e, - other => panic!("unknown phase {other}; use commit,sumcheck,e2e"), - }) - .collect() -} diff --git a/src/bin/phase_benchmark_report.rs b/src/bin/phase_benchmark_report.rs deleted file mode 100644 index 146462f8..00000000 --- a/src/bin/phase_benchmark_report.rs +++ /dev/null @@ -1,127 +0,0 @@ -use std::{ - collections::BTreeMap, - fs::File, - io::{BufRead, BufReader}, -}; - -use serde::Deserialize; - -#[derive(Clone, Debug, Deserialize)] -struct Row { - backend: String, - phase: String, - log_size: usize, - size: usize, - fold: usize, - rate: usize, - duration_ms: f64, - peak_allocated_bytes: usize, - peak_phase_delta_bytes: usize, - hashes: usize, - proof_bytes: Option, - status: Option, - #[serde(default)] - metal_upload_bytes: u64, - #[serde(default)] - metal_upload_ms: f64, - #[serde(default)] - metal_readback_bytes: u64, - #[serde(default)] - metal_readback_ms: f64, - #[serde(default)] - metal_alloc_bytes: u64, - #[serde(default)] - metal_command_wait_ms: f64, - #[serde(default)] - metal_blit_bytes: u64, - #[serde(default)] - metal_blit_wait_ms: f64, -} - -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] -struct Key { - phase: String, - log_size: usize, - size: usize, - fold: usize, - rate: usize, -} - -#[derive(Default)] -struct Pair { - cpu: Option, - gpu: Option, -} - -fn main() { - let path = std::env::args() - .nth(1) - .unwrap_or_else(|| "outputs/phase_benchmark_cpu_gpu.jsonl".to_string()); - let file = File::open(&path).unwrap_or_else(|err| panic!("failed to open {path}: {err}")); - let mut rows = BTreeMap::::new(); - for line in BufReader::new(file).lines() { - let line = line.unwrap(); - if line.trim().is_empty() { - continue; - } - let row: Row = serde_json::from_str(&line).unwrap(); - let key = Key { - phase: row.phase.clone(), - log_size: row.log_size, - size: row.size, - fold: row.fold, - rate: row.rate, - }; - let pair = rows.entry(key).or_default(); - match row.backend.as_str() { - "cpu" => pair.cpu = Some(row), - "gpu-metal" => pair.gpu = Some(row), - _ => {} - } - } - - println!( - "log_size,size,fold,rate,phase,cpu_status,gpu_status,cpu_ms,gpu_ms,speedup,cpu_peak_bytes,gpu_peak_bytes,cpu_peak_delta_bytes,gpu_peak_delta_bytes,cpu_hashes,gpu_hashes,proof_bytes,gpu_upload_bytes,gpu_upload_ms,gpu_readback_bytes,gpu_readback_ms,gpu_alloc_bytes,gpu_command_wait_ms,gpu_blit_bytes,gpu_blit_wait_ms" - ); - for (key, pair) in rows { - let (Some(cpu), Some(gpu)) = (pair.cpu, pair.gpu) else { - continue; - }; - let speedup = if gpu.duration_ms > 0.0 { - cpu.duration_ms / gpu.duration_ms - } else { - 0.0 - }; - let proof_bytes = cpu.proof_bytes.or(gpu.proof_bytes).unwrap_or(0); - let cpu_status = cpu.status.as_deref().unwrap_or("ok"); - let gpu_status = gpu.status.as_deref().unwrap_or("ok"); - println!( - "{},{},{},{},{},{},{},{:.6},{:.6},{:.4},{},{},{},{},{},{},{},{},{:.6},{},{:.6},{},{:.6},{},{:.6}", - key.log_size, - key.size, - key.fold, - key.rate, - key.phase, - cpu_status, - gpu_status, - cpu.duration_ms, - gpu.duration_ms, - speedup, - cpu.peak_allocated_bytes, - gpu.peak_allocated_bytes, - cpu.peak_phase_delta_bytes, - gpu.peak_phase_delta_bytes, - cpu.hashes, - gpu.hashes, - proof_bytes, - gpu.metal_upload_bytes, - gpu.metal_upload_ms, - gpu.metal_readback_bytes, - gpu.metal_readback_ms, - gpu.metal_alloc_bytes, - gpu.metal_command_wait_ms, - gpu.metal_blit_bytes, - gpu.metal_blit_wait_ms, - ); - } -} diff --git a/src/hash/metal_profile.rs b/src/hash/metal_profile.rs index 17134c13..afdf4cb6 100644 --- a/src/hash/metal_profile.rs +++ b/src/hash/metal_profile.rs @@ -18,6 +18,10 @@ pub struct MetalProfileSnapshot { pub blit_count: u64, pub blit_bytes: u64, pub blit_wait_nanos: u64, + /// Device-wide allocated GPU memory at snapshot time. + pub device_current_bytes: u64, + /// Peak device-wide allocated GPU memory since the last `reset_device_peak`. + pub device_peak_bytes: u64, } impl MetalProfileSnapshot { @@ -38,6 +42,9 @@ impl MetalProfileSnapshot { blit_count: self.blit_count.saturating_sub(before.blit_count), blit_bytes: self.blit_bytes.saturating_sub(before.blit_bytes), blit_wait_nanos: self.blit_wait_nanos.saturating_sub(before.blit_wait_nanos), + // Gauges, not counters: keep the latest values. + device_current_bytes: self.device_current_bytes, + device_peak_bytes: self.device_peak_bytes, } } @@ -71,6 +78,8 @@ static COMMAND_WAIT_NANOS: AtomicU64 = AtomicU64::new(0); static BLIT_COUNT: AtomicU64 = AtomicU64::new(0); static BLIT_BYTES: AtomicU64 = AtomicU64::new(0); static BLIT_WAIT_NANOS: AtomicU64 = AtomicU64::new(0); +static DEVICE_CURRENT_BYTES: AtomicU64 = AtomicU64::new(0); +static DEVICE_PEAK_BYTES: AtomicU64 = AtomicU64::new(0); pub fn snapshot() -> MetalProfileSnapshot { MetalProfileSnapshot { @@ -87,9 +96,25 @@ pub fn snapshot() -> MetalProfileSnapshot { blit_count: BLIT_COUNT.load(Ordering::Relaxed), blit_bytes: BLIT_BYTES.load(Ordering::Relaxed), blit_wait_nanos: BLIT_WAIT_NANOS.load(Ordering::Relaxed), + device_current_bytes: DEVICE_CURRENT_BYTES.load(Ordering::Relaxed), + device_peak_bytes: DEVICE_PEAK_BYTES.load(Ordering::Relaxed), } } +/// Record the device-wide allocated size (`MTLDevice.currentAllocatedSize`), +/// sampled after each buffer allocation. Peak live memory always coincides +/// with an allocation, so per-alloc sampling captures the true peak. +pub fn record_device_allocated(bytes: u64) { + DEVICE_CURRENT_BYTES.store(bytes, Ordering::Relaxed); + DEVICE_PEAK_BYTES.fetch_max(bytes, Ordering::Relaxed); +} + +/// Reset the peak gauge to the current value (e.g. at the start of a +/// measured phase). +pub fn reset_device_peak() { + DEVICE_PEAK_BYTES.store(DEVICE_CURRENT_BYTES.load(Ordering::Relaxed), Ordering::Relaxed); +} + pub fn record_upload(bytes: u64, duration: Duration) { UPLOAD_COUNT.fetch_add(1, Ordering::Relaxed); UPLOAD_BYTES.fetch_add(bytes, Ordering::Relaxed); diff --git a/src/hash/metal_sha2_engine.rs b/src/hash/metal_sha2_engine.rs index 8879d77b..1ad6cb65 100644 --- a/src/hash/metal_sha2_engine.rs +++ b/src/hash/metal_sha2_engine.rs @@ -302,8 +302,11 @@ struct MetalSha2Runtime { fn new_shared_buffer(rt: &MetalSha2Runtime, bytes: u64) -> Buffer { metal_profile::record_alloc(bytes); - rt.device - .new_buffer(bytes, MTLResourceOptions::StorageModeShared) + let buffer = rt + .device + .new_buffer(bytes, MTLResourceOptions::StorageModeShared); + metal_profile::record_device_allocated(rt.device.current_allocated_size()); + buffer } fn new_shared_buffer_with_data( @@ -317,6 +320,7 @@ fn new_shared_buffer_with_data( .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); metal_profile::record_alloc(bytes); metal_profile::record_upload(bytes, start.elapsed()); + metal_profile::record_device_allocated(rt.device.current_allocated_size()); buffer } diff --git a/src/hash/mod.rs b/src/hash/mod.rs index 1d53c405..94c6570f 100644 --- a/src/hash/mod.rs +++ b/src/hash/mod.rs @@ -19,7 +19,10 @@ use static_assertions::{assert_impl_all, assert_obj_safe}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned}; #[cfg(all(feature = "metal", target_os = "macos"))] -pub use self::metal_profile::{snapshot as metal_profile_snapshot, MetalProfileSnapshot}; +pub use self::metal_profile::{ + reset_device_peak as metal_reset_device_peak, snapshot as metal_profile_snapshot, + MetalProfileSnapshot, +}; #[cfg(all(feature = "metal", target_os = "macos"))] pub use self::metal_sha2_engine::MetalSha2; pub use self::{ From 1cfda8c9886d3a681d95df896eca554cf32be0f7 Mon Sep 17 00:00:00 2001 From: ocdbytes Date: Mon, 15 Jun 2026 19:47:08 +0530 Subject: [PATCH 12/35] rebase and merge commit --- Cargo.lock | 21 + Cargo.toml | 1 + .../protocols/params/basecase.txt | 10 + .../protocols/params/code_switch.txt | 13 + .../protocols/params/derive.txt | 7 + .../protocols/params/irs_commit.txt | 8 + .../protocols/params/sumcheck.txt | 12 + src/bin/benchmark.rs | 13 +- src/bin/main.rs | 13 +- src/bits.rs | 28 + src/parameters.rs | 14 +- src/protocols/basecase.rs | 338 +++++---- src/protocols/code_switch.rs | 386 +++++++---- src/protocols/irs_commit.rs | 555 +++++++-------- src/protocols/mask_proximity.rs | 101 ++- src/protocols/mod.rs | 1 + src/protocols/params/basecase.rs | 198 ++++++ src/protocols/params/bounds.rs | 40 ++ src/protocols/params/branch.rs | 65 ++ src/protocols/params/build_round.rs | 238 +++++++ src/protocols/params/code_switch.rs | 436 ++++++++++++ src/protocols/params/derive.rs | 653 ++++++++++++++++++ src/protocols/params/error.rs | 196 ++++++ src/protocols/params/irs_commit.rs | 203 ++++++ src/protocols/params/layout.rs | 229 ++++++ src/protocols/params/mask_proximity.rs | 167 +++++ src/protocols/params/mod.rs | 37 + src/protocols/params/protocol_config.rs | 522 ++++++++++++++ src/protocols/params/regime.rs | 358 ++++++++++ src/protocols/params/solved.rs | 57 ++ src/protocols/params/spec.rs | 415 +++++++++++ src/protocols/params/sumcheck.rs | 285 ++++++++ src/protocols/params/test_utils.rs | 171 +++++ src/protocols/proof_of_work.rs | 42 +- src/protocols/sumcheck.rs | 312 ++++++--- src/protocols/whir/config.rs | 298 ++++---- src/protocols/whir/mod.rs | 103 ++- src/protocols/whir/prover.rs | 36 +- src/protocols/whir/verifier.rs | 40 +- src/protocols/whir_zk/committer.rs | 12 +- src/protocols/whir_zk/mod.rs | 27 +- src/protocols/whir_zk/prover.rs | 7 +- src/protocols/whir_zk/verifier.rs | 6 +- 43 files changed, 5747 insertions(+), 927 deletions(-) create mode 100644 proptest-regressions/protocols/params/basecase.txt create mode 100644 proptest-regressions/protocols/params/code_switch.txt create mode 100644 proptest-regressions/protocols/params/derive.txt create mode 100644 proptest-regressions/protocols/params/irs_commit.txt create mode 100644 proptest-regressions/protocols/params/sumcheck.txt create mode 100644 src/protocols/params/basecase.rs create mode 100644 src/protocols/params/bounds.rs create mode 100644 src/protocols/params/branch.rs create mode 100644 src/protocols/params/build_round.rs create mode 100644 src/protocols/params/code_switch.rs create mode 100644 src/protocols/params/derive.rs create mode 100644 src/protocols/params/error.rs create mode 100644 src/protocols/params/irs_commit.rs create mode 100644 src/protocols/params/layout.rs create mode 100644 src/protocols/params/mask_proximity.rs create mode 100644 src/protocols/params/mod.rs create mode 100644 src/protocols/params/protocol_config.rs create mode 100644 src/protocols/params/regime.rs create mode 100644 src/protocols/params/solved.rs create mode 100644 src/protocols/params/spec.rs create mode 100644 src/protocols/params/sumcheck.rs create mode 100644 src/protocols/params/test_utils.rs diff --git a/Cargo.lock b/Cargo.lock index acc6ca37..62e5a0d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1360,6 +1360,26 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "thread_local" version = "1.1.9" @@ -1556,6 +1576,7 @@ dependencies = [ "sha3", "spongefish", "static_assertions", + "thiserror", "tracing", "tracing-subscriber", "zerocopy", diff --git a/Cargo.toml b/Cargo.toml index 515ee2b8..52fc40e5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,7 @@ const-oid = "0.9.6" arrayvec = "0.7.6" derive-where = { version = "1.6.0", features = ["safe"] } ordered-float = { version = "5.1.0", features = ["serde"] } +thiserror = "2.0" [dev-dependencies] proptest = "1.0" diff --git a/proptest-regressions/protocols/params/basecase.txt b/proptest-regressions/protocols/params/basecase.txt new file mode 100644 index 00000000..c3b26074 --- /dev/null +++ b/proptest-regressions/protocols/params/basecase.txt @@ -0,0 +1,10 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc b6e8ae0b3e6a9769901e0e0e489da34965bf0a8df7dd049aef66e0541bf10baf # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, decoding_regime: Johnson, target_security_bits: 30, pow_budget: Forbidden, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_size, log_inv_rate) = (1, 1) +cc a2f771fc5031440200810b95ea2d347da895f8eb2e1a87f53fd69ad224287e84 # shrinks to spec = SecuritySpec { mode: Standard, decoding_regime: Johnson, target_security_bits: 30, pow_budget: Forbidden, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_size, log_inv_rate) = (2, 2) +cc f66c89bc700c79bca5f4b7234f1345129962c78e5a2036a6430564f615f19b30 # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, decoding_regime: Johnson, target_security_bits: 30, pow_budget: Forbidden, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_size, log_inv_rate) = (1, 1) +cc 312fec8a96f6f55f5d3c0346bdb85690f23150aadf88888d453041e99b05d414 # shrinks to spec = SecuritySpec { mode: Standard, decoding_regime: Johnson, target_security_bits: 39, pow_budget: Forbidden, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_size, log_inv_rate) = (1, 1) diff --git a/proptest-regressions/protocols/params/code_switch.txt b/proptest-regressions/protocols/params/code_switch.txt new file mode 100644 index 00000000..de6c40cf --- /dev/null +++ b/proptest-regressions/protocols/params/code_switch.txt @@ -0,0 +1,13 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 7a7df094ea650db7a295d162b75dd9da9b52d1fc36947d2b07df8150cd9d906f # shrinks to spec = SecuritySpec { mode: Standard { unique_decoding: false }, target_security_bits: 80, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 4, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, log_inv_rate = 1, folding_factor = 3, num_vars = 4 +cc b42c982074a04c7110df07cf00f45156607be547e176b1ddd5f9d994ad491ddb # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 80, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 4, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, log_inv_rate = 1, folding_factor = 3, num_vars = 4 +cc eaf09a2b6bdffa86026264679f008326498ca800260dd2f17d4370df9fb3f801 # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 80, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 4, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, log_inv_rate = 1, folding_factor = 3, num_vars = 4 +cc 3887a5fa698c99109e8262e843dbd24ea94b9c9d420791e4520b5c9211a3eca0 # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 100, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 4, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, (log_inv_rate, folding_factor, num_vars) = (3, 2, 7) +cc b3e128084f721e6f43e263e05acf2e2de6fcd05dccf3811f063eeb0b63d78f8e # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 47, max_pow_bits: Some(15), hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_inv_rate, folding_factor, num_vars) = (3, 2, 4) +cc b71da9002ceac9e4a74af097a7b087557a5b916fe8da47e39c4682375d749f88 # shrinks to spec = SecuritySpec { mode: Standard, decoding_regime: Johnson, target_security_bits: 50, pow_budget: Forbidden, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_inv_rate, folding_factor, num_vars) = (1, 2, 4) +cc 1981509d857e56772dd4a79f8692619e968891aa3d84576ea1857f6d9a484a2d # shrinks to spec = SecuritySpec { mode: Standard, decoding_regime: Johnson, target_security_bits: 50, pow_budget: Forbidden, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, (log_inv_rate, folding_factor, num_vars) = (1, 2, 4) diff --git a/proptest-regressions/protocols/params/derive.txt b/proptest-regressions/protocols/params/derive.txt new file mode 100644 index 00000000..ca83c1b3 --- /dev/null +++ b/proptest-regressions/protocols/params/derive.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 104921a4117ed8255308c1ea5d3e12c72356ef72ef0d93fc0f24ed29f93fdd3a # shrinks to tuning = TuningSpec { vector_size: 32, starting_log_inv_rate: 3, folding_factor: Constant(1) } diff --git a/proptest-regressions/protocols/params/irs_commit.txt b/proptest-regressions/protocols/params/irs_commit.txt new file mode 100644 index 00000000..4d35f3df --- /dev/null +++ b/proptest-regressions/protocols/params/irs_commit.txt @@ -0,0 +1,8 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 0b6dd03179c9a4e38b29b34b241b88fba69348a2c8938af7253314b7035bea82 # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 80, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 4, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 16, log_inv_rate: 1, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 }, out_domain = 0, seed = 0 +cc 7e49f7a2d53f55cfa2f09114d17ab4123678b45ddf69e0cfbc646b246de2f042 # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 80, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c }, ctx = RoundContext { vector_size: 128, log_inv_rate: 2, folding_factor: 2 }, out_domain = 11 diff --git a/proptest-regressions/protocols/params/sumcheck.txt b/proptest-regressions/protocols/params/sumcheck.txt new file mode 100644 index 00000000..d6f6e6ed --- /dev/null +++ b/proptest-regressions/protocols/params/sumcheck.txt @@ -0,0 +1,12 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 0ffdc71948ed0315f4cf55fb8f2dd25bf71f7e41f53cd4fe35ee9da6fb125a20 # shrinks to spec = SecuritySpec { mode: Standard { unique_decoding: false }, target_security_bits: 86, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 3, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 32, log_inv_rate: 2, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 } +cc e8ab6549772cf6bf4c3af116ebcba3dbf295ffbe2aee4a94be7df4b9f45d61ec # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 91, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 3, max_pow_bits: Some(10), hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 16, log_inv_rate: 1, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 } +cc 8c4300cc375640956f81e9da5aef9ea11ef476ddc4dd253dc560afa07609262d # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 98, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 3, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 16, log_inv_rate: 1, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 } +cc 8ea40f13c63b4c0021386369ce698a5d9289381a39dc85db43d2d69b9b4877bb # shrinks to spec = SecuritySpec { mode: Standard { unique_decoding: false }, target_security_bits: 88, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 3, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 16, log_inv_rate: 3, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 } +cc f1dca600886474c74d857c547baea0c2b4faf45b2946036f21a008106396eb1c # shrinks to spec = SecuritySpec { mode: Standard { unique_decoding: false }, target_security_bits: 80, vector_size: 256, starting_log_inv_rate: 1, initial_folding_factor: 4, folding_factor: 3, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 256, log_inv_rate: 3, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 } +cc 36d0f5929e8099fa8644b0511229cf11634e5a7a66d99c06099c304f5f7a8c6e # shrinks to spec = SecuritySpec { mode: ZeroKnowledge, target_security_bits: 47, max_pow_bits: None, hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, _embedding: PhantomData, 1>>> }, ctx = RoundContext { round_index: 0, vector_size: 128, log_inv_rate: 4, folding_factor: 1, prev_round_in_domain_samples: 0, prev_round_query_error: 0.0 } diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index d9e79ce6..77574c42 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -18,6 +18,7 @@ use whir::{ cmdline_utils::{AvailableFields, AvailableHash}, hash::HASH_COUNTER, parameters::ProtocolParameters, + protocols::params::DecodingRegime, transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; @@ -48,8 +49,8 @@ struct Args { #[arg(short = 'k', long = "fold", default_value = "4")] folding_factor: usize, - #[arg(long = "unique-decoding", default_value_t = false)] - unique_decoding: bool, + #[arg(long = "decoding-regime", default_value = "Johnson")] + decoding_regime: DecodingRegime, #[arg(short = 'f', long = "field", default_value = "Goldilocks3")] field: AvailableFields, @@ -67,7 +68,7 @@ struct BenchmarkOutput { repetitions: usize, initial_folding_factor: usize, folding_factor: usize, - unique_decoding: bool, + decoding_regime: DecodingRegime, field: AvailableFields, hash: AvailableHash, @@ -117,7 +118,7 @@ where let reps = args.verifier_repetitions; let folding_factor = args.folding_factor; let first_round_folding_factor = args.first_round_folding_factor; - let unique_decoding = args.unique_decoding; + let decoding_regime = args.decoding_regime; std::fs::create_dir_all("outputs").unwrap(); @@ -128,7 +129,7 @@ where pow_bits, initial_folding_factor: first_round_folding_factor, folding_factor, - unique_decoding, + decoding_regime, starting_log_inv_rate: starting_rate, batch_size: 1, hash_id: args.hash.hash_id(), @@ -298,7 +299,7 @@ where repetitions: reps, initial_folding_factor: first_round_folding_factor, folding_factor, - unique_decoding, + decoding_regime, field: args.field, hash: args.hash, diff --git a/src/bin/main.rs b/src/bin/main.rs index bafe16f5..0faf4b36 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -14,6 +14,7 @@ use whir::{ cmdline_utils::{AvailableFields, AvailableHash}, hash::HASH_COUNTER, parameters::ProtocolParameters, + protocols::params::DecodingRegime, transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; @@ -48,9 +49,9 @@ struct Args { #[arg(short = 'k', long = "fold", default_value = "4")] folding_factor: usize, - /// Restrict PCS to the Unique Decoding regime. LDT is always UD. - #[arg(long = "unique-decoding", default_value_t = false)] - unique_decoding: bool, + /// Reed–Solomon decoding regime: Unique or Johnson (list-decoding). + #[arg(long = "decoding-regime", default_value = "Johnson")] + decoding_regime: DecodingRegime, #[arg(short = 'f', long = "field", default_value = "Goldilocks3")] field: AvailableFields, @@ -109,7 +110,7 @@ where let reps = args.verifier_repetitions; let first_round_folding_factor = args.first_round_folding_factor; let folding_factor = args.folding_factor; - let unique_decoding = args.unique_decoding; + let decoding_regime = args.decoding_regime; let num_evaluations = args.num_evaluations; let num_linear_constraints = args.num_linear_constraints; let hash_id = args.hash.hash_id(); @@ -125,7 +126,7 @@ where pow_bits, initial_folding_factor: first_round_folding_factor, folding_factor, - unique_decoding, + decoding_regime, starting_log_inv_rate: starting_rate, batch_size: 1, hash_id, @@ -254,7 +255,7 @@ where let num_coeffs = 1 << num_variables; let whir_params = ProtocolParameters { - unique_decoding: args.unique_decoding, + decoding_regime: args.decoding_regime, security_level, pow_bits, initial_folding_factor: first_round_folding_factor, diff --git a/src/bits.rs b/src/bits.rs index 537c8ab2..74f51276 100644 --- a/src/bits.rs +++ b/src/bits.rs @@ -2,6 +2,8 @@ use std::{ cmp::Ordering, fmt::{self, Display}, hash::Hash, + iter::Sum, + ops::{Add, Sub}, }; use serde::{Deserialize, Serialize}; @@ -27,6 +29,32 @@ impl Bits { pub fn is_zero(&self) -> bool { self.0 == 0.0 } + + /// Absolute difference, e.g. for tolerance comparisons. + #[must_use] + pub fn abs_diff(self, other: Self) -> Self { + Self::new((self.0 - other.0).abs()) + } +} + +impl Add for Bits { + type Output = Self; + fn add(self, rhs: Self) -> Self { + Self::new(self.0 + rhs.0) + } +} + +impl Sub for Bits { + type Output = Self; + fn sub(self, rhs: Self) -> Self { + Self::new(self.0 - rhs.0) + } +} + +impl Sum for Bits { + fn sum>(iter: I) -> Self { + iter.fold(Self(0.0), Add::add) + } } impl From for Bits { diff --git a/src/parameters.rs b/src/parameters.rs index 242b9749..8e257c79 100644 --- a/src/parameters.rs +++ b/src/parameters.rs @@ -2,13 +2,13 @@ use std::fmt::{Debug, Display}; use serde::{Deserialize, Serialize}; -use crate::engines::EngineId; +use crate::{engines::EngineId, protocols::params::DecodingRegime}; /// Configuration parameters for WHIR proofs. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct ProtocolParameters { - /// Whether to require unique decoding. - pub unique_decoding: bool, + /// Reed–Solomon decoding regime: `Unique` or `Johnson` (list-decoding). + pub decoding_regime: DecodingRegime, /// The logarithmic inverse rate for sampling. pub starting_log_inv_rate: usize, /// Folding factor for the initial round. @@ -30,13 +30,7 @@ impl Display for ProtocolParameters { writeln!( f, "Targeting {}-bits of security with {}-bits of PoW using {} decoding", - self.security_level, - self.pow_bits, - if self.unique_decoding { - "unique" - } else { - "list" - } + self.security_level, self.pow_bits, self.decoding_regime, )?; writeln!( f, diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 32763271..ec22a628 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -1,8 +1,4 @@ -//! Base Case Linear Opening Protocol -//! -//! It support honest verifier zero-knowledge (HVZK), but is not succinct. -//! -//! § 7. +//! Non-succinct linear opening (Construction 7.2, p.43). HVZK in ZK mode. use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; @@ -15,7 +11,7 @@ use crate::{ univariate_evaluate, }, hash::Hash, - protocols::{irs_commit, sumcheck}, + protocols::{irs_commit, proof_of_work, sumcheck}, transcript::{ codecs::U64, Codec, DuplexSpongeInterface, ProverMessage, ProverState, VerifierMessage, VerifierState, @@ -24,26 +20,80 @@ use crate::{ verify, }; -/// Output from the base case protocol (shared by prover and verifier). #[must_use] pub struct Opening { pub evaluation_points: Vec, pub linear_form_evaluation: F, } +/// Standard / ZeroKnowledge selector for basecase. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum BasecaseMode { + Standard, + ZeroKnowledge, +} + +#[must_use] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config { - pub commit: irs_commit::Config>, - pub sumcheck: sumcheck::Config, - - /// Whether to mask the vectors, which adds HVZK. - pub masked: bool, + commit: irs_commit::Config>, + sumcheck: sumcheck::Config, + mode: BasecaseMode, + pow: proof_of_work::Config, } impl Config { + /// Standard basecase has no γ challenge — PoW must be `none()`. ZK + /// basecase has a γ-combination slot (Lemma 7.4) and may or may not need + /// PoW depending on whether the analytic floor already clears the target + /// (under unique decoding it often does). + pub fn new( + commit: irs_commit::Config>, + sumcheck: sumcheck::Config, + mode: BasecaseMode, + pow: proof_of_work::Config, + ) -> Self { + let has_pow = pow != proof_of_work::Config::none(); + debug_assert!( + !matches!(mode, BasecaseMode::Standard) || !has_pow, + "Standard basecase has no γ challenge — pow must be none()", + ); + Self { + commit, + sumcheck, + mode, + pow, + } + } + + pub const fn commit(&self) -> &irs_commit::Config> { + &self.commit + } + + pub const fn sumcheck(&self) -> &sumcheck::Config { + &self.sumcheck + } + + pub const fn mode(&self) -> BasecaseMode { + self.mode + } + + pub const fn pow(&self) -> proof_of_work::Config { + self.pow + } + + #[cfg(test)] + pub(crate) const fn set_pow_for_test(&mut self, pow: proof_of_work::Config) { + self.pow = pow; + } + pub const fn size(&self) -> usize { - self.sumcheck.initial_size + self.sumcheck.initial_size() + } + + pub const fn is_zk(&self) -> bool { + matches!(self.mode, BasecaseMode::ZeroKnowledge) } pub fn prove( @@ -64,10 +114,10 @@ impl Config { Hash: ProverMessage<[H::U]>, Standard: Distribution, { - assert_eq!(self.commit.interleaving_depth, 1); - assert_eq!(self.commit.num_vectors, 1); - assert_eq!(self.commit.vector_size, self.sumcheck.initial_size); - assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size)); + assert_eq!(self.commit.interleaving_depth(), 1); + assert_eq!(self.commit.num_vectors(), 1); + assert_eq!(self.commit.vector_size(), self.sumcheck.initial_size()); + assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size())); debug_assert_eq!(dot(&vector, &covector), sum); if self.size() == 0 { return Opening { @@ -76,63 +126,22 @@ impl Config { }; } - // Even more trivial non-zk protocol: send f and r directly. - if !self.masked { - prover_state.prover_messages(&vector); - prover_state.prover_messages(&witness.masks); - let _ = self.commit.open(prover_state, &[witness]); - let point = self - .sumcheck - .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) - .round_challenges; - assert!(!vector[0].is_zero(), "Proof failed"); - return Opening { - evaluation_points: point, - linear_form_evaluation: covector[0], - }; - } - - // Create masking vector. - let mask = random_vector(prover_state.rng(), vector.len()); - - // Commit to the masking vector. - let mask_witness = self.commit.commit(prover_state, &[&mask]); + let blinding_witness = + self.maybe_blind_prove(prover_state, &mut vector, witness, &covector, &mut sum); - // Compute and send linear form of mask (μ' in paper). - let mask_sum = dot(&mask, &covector); - prover_state.prover_message(&mask_sum); + let witnesses: Vec<&irs_commit::Witness> = blinding_witness + .as_ref() + .map_or_else(|| vec![witness], |b| vec![b, witness]); + let _ = self.commit.open(prover_state, &witnesses); - // RLC the mask with the vector - let mask_rlc = prover_state.verifier_message::(); - assert!(!mask_rlc.is_zero(), "Proof failed"); - let mut masked_vector = scalar_mul_add_new(&mask, mask_rlc, &vector); - prover_state.prover_messages(&masked_vector); - - // Send combined IRS randomness. (r^* in paper) - let masked_masks = scalar_mul_add_new(&mask_witness.masks, mask_rlc, &witness.masks); - prover_state.prover_messages(&masked_masks); - - // Open the commitment and mask simultaneously. - let _ = self.commit.open(prover_state, &[&mask_witness, witness]); - - // Run sumcheck to reduce linear form claim - let mut masked_sum = mask_sum + mask_rlc * sum; let point = self .sumcheck - .prove( - prover_state, - &mut masked_vector, - &mut covector, - &mut masked_sum, - &[], - ) + .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) .round_challenges; - // If the MLE of `masked_vector` evaluates to zero, the verifier can not proceed. - // Basically the sumcheck equation has degenerated to 0 * l(r) = 0, which provides - // no constraints on l(r) that the verifier can return. - // This event is cryptographically unlikely as `F` is challenge sized. - assert!(!masked_vector[0].is_zero(), "Proof failed"); + // Negligible event over a challenge-sized field; without it the verifier + // cannot derive `l(r) = sum / vector_mle(r)`. + assert!(!vector[0].is_zero(), "Proof failed"); Opening { evaluation_points: point, @@ -140,10 +149,64 @@ impl Config { } } + /// ZK: commits a blinding codeword, runs the RLC, mutates `vector`/`sum` to + /// the combined values, sends them cleartext. Standard: sends `vector` and + /// `witness.masks` cleartext (no ZK). + fn maybe_blind_prove( + &self, + prover_state: &mut ProverState, + vector: &mut Vec, + witness: &irs_commit::Witness, + covector: &[F], + sum: &mut F, + ) -> Option> + where + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + F: Codec<[H::U]>, + [u8; 32]: Decoding<[H::U]>, + U64: Codec<[H::U]>, + Hash: ProverMessage<[H::U]>, + Standard: Distribution, + { + match self.mode { + BasecaseMode::Standard => { + prover_state.prover_messages(vector); + prover_state.prover_messages(&witness.masks); + None + } + BasecaseMode::ZeroKnowledge => { + let blinding_vector = random_vector(prover_state.rng(), vector.len()); + let blinding_witness = self.commit.commit(prover_state, &[&blinding_vector]); + let blinding_inner_product = dot(&blinding_vector, covector); + prover_state.prover_message(&blinding_inner_product); + + // Grind the Theorem 7.1 γ-combination gap before γ is sampled. + self.pow.prove(prover_state); + + let combination_randomness = prover_state.verifier_message::(); + assert!(!combination_randomness.is_zero(), "Proof failed"); + + *vector = scalar_mul_add_new(&blinding_vector, combination_randomness, vector); + prover_state.prover_messages(vector); + + let combined_irs_randomness = scalar_mul_add_new( + &blinding_witness.masks, + combination_randomness, + &witness.masks, + ); + prover_state.prover_messages(&combined_irs_randomness); + + *sum = blinding_inner_product + combination_randomness * *sum; + Some(blinding_witness) + } + } + } + pub fn verify( &self, verifier_state: &mut VerifierState, - commitment: &irs_commit::Commitment, + commitment: &irs_commit::Commitment, mut sum: F, ) -> VerificationResult> where @@ -154,10 +217,10 @@ impl Config { U64: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { - assert_eq!(self.commit.interleaving_depth, 1); - assert_eq!(self.commit.num_vectors, 1); - assert_eq!(self.commit.vector_size, self.sumcheck.initial_size); - assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size)); + assert_eq!(self.commit.interleaving_depth(), 1); + assert_eq!(self.commit.num_vectors(), 1); + assert_eq!(self.commit.vector_size(), self.sumcheck.initial_size()); + assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size())); if self.size() == 0 { return Ok(Opening { evaluation_points: Vec::new(), @@ -165,72 +228,71 @@ impl Config { }); } - // Unmasked protocol - if !self.masked { - let vector = verifier_state.prover_messages_vec(self.commit.vector_size)?; - let masks = verifier_state - .prover_messages_vec(self.commit.mask_length * self.commit.num_messages())?; - let evals = self.commit.verify(verifier_state, &[commitment])?; - let point = self - .sumcheck - .verify(verifier_state, &mut sum)? - .round_challenges; - - for (&point, value) in zip_strict(&evals.points, evals.values(&[F::ONE])) { - // We expected `f(x) + x^l · g(x)` where l = deg(f) + 1, f is the message and g the mask. - let expected = univariate_evaluate(&vector, point) - + point.pow([self.commit.message_length() as u64]) - * univariate_evaluate(&masks, point); - verify!(value == expected); - } - let mle = multilinear_extend(&vector, &point); - verify!(!mle.is_zero()); - let linear_mle = sum / mle; - return Ok(Opening { - evaluation_points: point, - linear_form_evaluation: linear_mle, - }); - } + let blind = self.maybe_receive_blind(verifier_state, &mut sum)?; - let mask_commitment = self.commit.receive_commitment(verifier_state)?; - let mask_sum: F = verifier_state.prover_message()?; - let mask_rlc: F = verifier_state.verifier_message(); - verify!(!mask_rlc.is_zero()); - let masked_vector: Vec = verifier_state.prover_messages_vec(self.commit.vector_size)?; - let masked_masks: Vec = verifier_state.prover_messages_vec(self.commit.mask_length)?; + let vector = verifier_state.prover_messages_vec(self.commit.vector_size())?; + let irs_randomness = verifier_state + .prover_messages_vec(self.commit.mask_length() * self.commit.num_messages())?; - // Open the commitment and mask simultaneously. - let evals = self - .commit - .verify(verifier_state, &[&mask_commitment, commitment])?; + let (commitments, weights): (Vec<&irs_commit::Commitment>, Vec) = match &blind { + Some((b, gamma)) => (vec![b, commitment], vec![F::ONE, *gamma]), + None => (vec![commitment], vec![F::ONE]), + }; + let evals = self.commit.verify(verifier_state, &commitments)?; - // Spot check evaluations. - for (&point, value) in zip_strict(&evals.points, evals.values(&[F::ONE, mask_rlc])) { - // We expected `f(x) + x^l · g(x)` where l = deg(f) + 1, f is the message and g the mask. - let expected = univariate_evaluate(&masked_vector, point) + // Spot-check: Enc_C(vector, irs_randomness)(x) = Σ weights · opened_row(x). + for (&point, value) in zip_strict(&evals.points, evals.values(&weights)) { + let expected = univariate_evaluate(&vector, point) + point.pow([self.commit.message_length() as u64]) - * univariate_evaluate(&masked_masks, point); + * univariate_evaluate(&irs_randomness, point); verify!(value == expected); } - // Sumcheck on masked inner product - let mut masked_sum = mask_sum + mask_rlc * sum; let point = self .sumcheck - .verify(verifier_state, &mut masked_sum)? + .verify(verifier_state, &mut sum)? .round_challenges; - // Compute implied MLE of the linear form - // f*(r) · l(r) = sum => l(r) = sum / f*(r) - let masked_mle = multilinear_extend(&masked_vector, &point); - verify!(!masked_mle.is_zero()); - let linear_mle = masked_sum / masked_mle; + // l(r) = sum / vector_mle(r), where l is the implicit linear form. + let mle = multilinear_extend(&vector, &point); + verify!(!mle.is_zero()); + let linear_mle = sum / mle; Ok(Opening { evaluation_points: point, linear_form_evaluation: linear_mle, }) } + + /// ZK: reads the blinding commitment + μ' + γ, mutates `sum` to the + /// combined value, returns `(commitment, γ)`. Standard: no-op. + fn maybe_receive_blind( + &self, + verifier_state: &mut VerifierState, + sum: &mut F, + ) -> VerificationResult> + where + H: DuplexSpongeInterface, + F: Codec<[H::U]>, + u8: Decoding<[H::U]>, + [u8; 32]: Decoding<[H::U]>, + U64: Codec<[H::U]>, + Hash: ProverMessage<[H::U]>, + { + match self.mode { + BasecaseMode::Standard => Ok(None), + BasecaseMode::ZeroKnowledge => { + let blinding_commitment = self.commit.receive_commitment(verifier_state)?; + let blinding_inner_product: F = verifier_state.prover_message()?; + // Grind the Theorem 7.1 γ-combination gap before γ is sampled. + self.pow.verify(verifier_state)?; + let combination_randomness: F = verifier_state.verifier_message(); + verify!(!combination_randomness.is_zero()); + *sum = blinding_inner_product + combination_randomness * *sum; + Ok(Some((blinding_commitment, combination_randomness))) + } + } + } } #[cfg(test)] @@ -241,27 +303,26 @@ mod tests { use tracing::instrument; use super::*; - use crate::{ - algebra::fields, protocols::proof_of_work, transcript::DomainSeparator, type_info::Type, - }; + use crate::{algebra::fields, protocols::proof_of_work, transcript::DomainSeparator}; impl Config { pub fn arbitrary(size: usize, mask_length: usize) -> impl Strategy { let commit = irs_commit::Config::arbitrary(Identity::::new(), 1, size, mask_length, 1); - (commit, bool::weighted(0.8)).prop_map(move |(commit, masked)| Self { - commit: irs_commit::Config { - out_domain_samples: 0, - ..commit - }, - sumcheck: sumcheck::Config { - field: Type::new(), - initial_size: size, - round_pow: proof_of_work::Config::none(), - num_rounds: size.next_power_of_two().trailing_zeros() as usize, - mask_length: 0, + (commit, bool::weighted(0.8)).prop_map(move |(commit, is_zk)| Self { + commit, + sumcheck: sumcheck::Config::new( + size, + proof_of_work::Config::none(), + size.next_power_of_two().trailing_zeros() as usize, + sumcheck::SumcheckMode::Standard, + ), + mode: if is_zk { + BasecaseMode::ZeroKnowledge + } else { + BasecaseMode::Standard }, - masked, + pow: proof_of_work::Config::none(), }) } } @@ -272,7 +333,6 @@ mod tests { F: Field + Codec, Standard: Distribution, { - // Pseudo-random Instance let instance = U64(seed); let ds = DomainSeparator::protocol(config) .session(&format!("Test at {}:{}", file!(), line!())) @@ -282,7 +342,6 @@ mod tests { let covector = random_vector(&mut rng, config.size()); let sum = dot(&vector, &covector); - // Prover let mut prover_state = ProverState::new_std(&ds); let witness = config.commit.commit(&mut prover_state, &[&vector]); let prover_result = config.prove( @@ -298,7 +357,6 @@ mod tests { ); let proof = prover_state.proof(); - // Verifier let mut verifier_state = VerifierState::new_std(&ds, &proof); let commitment = config .commit diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 62bbb602..cfaf67b1 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -3,7 +3,7 @@ //! Reduces a proximity claim about oracle f (source code C) to a proximity //! claim about oracle g (target code C'). Supports optional ZK via mask oracle. -use std::fmt; +use std::{fmt, num::NonZeroUsize}; use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; @@ -21,22 +21,32 @@ use crate::{ protocols::{ geometric_challenge::geometric_challenge, irs_commit::{Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness}, + proof_of_work, }, transcript::{ - Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, VerificationResult, - VerifierMessage, VerifierState, + codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, + VerificationResult, VerifierMessage, VerifierState, }, verify, }; +/// Standard / ZeroKnowledge selector for code-switch. +#[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] +pub enum CodeSwitchMode { + Standard, + ZeroKnowledge { message_mask_length: NonZeroUsize }, +} + /// Code-switching IOR config with optional ZK. +#[must_use] #[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config { - pub source: IrsConfig, - pub target: IrsConfig>, - pub message_mask_length: usize, // l_zk - pub out_domain_samples: usize, + source: IrsConfig, + target: IrsConfig>, + mode: CodeSwitchMode, + out_domain_samples: usize, + pow: proof_of_work::Config, } /// Prover output from the code-switch. @@ -48,85 +58,132 @@ pub struct Witness { } /// Verifier output from the code-switch. -pub type Commitment = IrsCommitment; - -/// Mask input for the code-switch prover. -// TODO : This may be removed after parameter selection PR -pub enum MaskInput<'a, F> { - Disabled, - Enabled(&'a [F]), -} +pub type Commitment = IrsCommitment; impl Config { /// Create a code-switch config. - /// - /// The orchestrator is responsible for: - /// - Setting `target_config.mask_length` for ZK mode before passing it in. - /// - Computing `out_domain_samples` from the security budget. - /// - Setting `message_mask_length` = mask oracle message length (0 for non-ZK). pub fn new( source_config: IrsConfig, target_config: IrsConfig>, out_domain_samples: usize, - message_mask_length: usize, + mode: CodeSwitchMode, + pow: proof_of_work::Config, ) -> Self { assert_eq!( - source_config.num_vectors, 1, + source_config.num_vectors(), + 1, "code-switch requires a single source vector" ); assert_eq!( - target_config.num_vectors, 1, + target_config.num_vectors(), + 1, "code-switch requires a single target vector" ); + // Construction 9.7 needs at least one OOD challenge; unique-decoding + // Standard mode (`t_ood = 0`) is incompatible with code-switch. + assert!( + out_domain_samples > 0, + "code-switch requires t_ood ≥ 1 (Construction 9.7)", + ); // Target encodes one polynomial of length ℓ = source.message_length() // under C' = D^{ι_t}. The IRS splits the input of length ℓ into ι_t // parallel slices of length ℓ/ι_t, each encoded under D. assert_eq!( - target_config.vector_size, + target_config.vector_size(), source_config.message_length(), "target vector_size must equal source message_length (target encodes one polynomial of length ℓ)" ); assert!( - target_config.interleaving_depth.is_power_of_two(), + target_config.interleaving_depth().is_power_of_two(), "target.interleaving_depth must be a power of 2" ); - // Theorem 9.6: ℓ_zk ≥ r (mask oracle must cover source randomness). - if message_mask_length > 0 { + assert!( + source_config.interleaving_depth().is_power_of_two(), + "source.interleaving_depth must be a power of 2" + ); + if let CodeSwitchMode::ZeroKnowledge { + message_mask_length, + } = &mode + { + let l_zk = message_mask_length.get(); + // Theorem 9.6: ℓ_zk ≥ r (mask oracle must cover source randomness). assert!( - message_mask_length >= source_config.mask_length, - "message_mask_length ({message_mask_length}) must be >= source randomness length ({})", - source_config.mask_length, + l_zk >= source_config.mask_length(), + "message_mask_length ({l_zk}) must be >= source randomness length ({})", + source_config.mask_length(), ); assert!( - message_mask_length - source_config.mask_length >= out_domain_samples, - "the sampled randomness (s) length must be covering all the out of domain sample requests" + l_zk - source_config.mask_length() >= out_domain_samples, + "sampled randomness (s) length must cover all out-of-domain sample requests" ); + // t' = target in-domain queries + OOD queries (Construction 9.7 step 4). + // Definition 3.16: a t'-query ZK encoding requires r' ≥ t'; here + // r' = target.mask_length. assert!( - target_config.mask_length - >= target_config.in_domain_samples + target_config.out_domain_samples, - "target encoder violates: t' > r', number of queries should be covered by random mask" + target_config.mask_length() + >= target_config.in_domain_samples() + out_domain_samples, + "target encoder violates t' ≤ r': queries must be covered by target mask" + ); + } else { + assert_eq!( + source_config.mask_length(), + 0, + "source with IRS randomness requires ZK mode", ); } - assert!( - source_config.mask_length == 0 || message_mask_length > 0, - "source with mask_length > 0 (IRS randomness) requires ZK mode (message_mask_length > 0)" - ); - assert!( - source_config.interleaving_depth.is_power_of_two(), - "source.interleaving_depth must be a power of 2" - ); Self { source: source_config, target: target_config, - message_mask_length, + mode, out_domain_samples, + pow, } } + pub const fn source(&self) -> &IrsConfig { + &self.source + } + + pub const fn target(&self) -> &IrsConfig> { + &self.target + } + + pub const fn mode(&self) -> &CodeSwitchMode { + &self.mode + } + + pub const fn out_domain_samples(&self) -> usize { + self.out_domain_samples + } + + pub const fn pow(&self) -> proof_of_work::Config { + self.pow + } + + #[cfg(test)] + pub(crate) const fn target_mut_for_test(&mut self) -> &mut IrsConfig> { + &mut self.target + } + + /// Mask oracle length `ℓ_zk`. Returns 0 in Standard mode. + pub const fn message_mask_length(&self) -> usize { + match &self.mode { + CodeSwitchMode::Standard => 0, + CodeSwitchMode::ZeroKnowledge { + message_mask_length, + } => message_mask_length.get(), + } + } + + /// `true` iff the protocol is configured for ZK. + pub const fn is_zk(&self) -> bool { + matches!(&self.mode, CodeSwitchMode::ZeroKnowledge { .. }) + } + /// Length of the covector for this code-switch. pub fn covector_length(&self) -> usize { - self.source.message_length() + self.message_mask_length + self.source.message_length() + self.message_mask_length() } /// Prove the code-switch. @@ -146,18 +203,18 @@ impl Config { /// `message` is `Fold(f, γ)`, the post-sumcheck polynomial of length /// `source.message_length()`. /// - /// `mask_input` is `(r || s)` from the orchestrator's shared mask tree - /// (see Construction 9.7 Step 1, p.55). Must be `None` when - /// `message_mask_length == 0`. + /// `mask` is `(r || s)` from the orchestrator's shared mask tree + /// (see Construction 9.7 Step 1, p.55). Length must equal + /// `self.message_mask_length()` — pass an empty slice in Standard mode. #[cfg_attr(feature = "tracing", instrument(skip_all))] pub fn prove( &self, prover_state: &mut ProverState, message: Vec, - witness: &IrsWitness, + witness: &IrsWitness, covector: &mut [M::Target], folding_randomness: &[M::Target], - mask_input: &MaskInput<'_, M::Target>, + mask: &[M::Target], ) -> Witness where H: DuplexSpongeInterface, @@ -165,52 +222,30 @@ impl Config { Standard: Distribution, M::Target: Codec<[H::U]>, u8: Decoding<[H::U]>, + [u8; 32]: Decoding<[H::U]>, + U64: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { assert_eq!(message.len(), self.source.message_length()); assert_eq!(covector.len(), self.covector_length()); + assert_eq!(mask.len(), self.message_mask_length()); assert_eq!( 1 << folding_randomness.len(), - self.source.interleaving_depth, + self.source.interleaving_depth(), "folding_randomness must have length log2(source.interleaving_depth) ({} != log2({}))", folding_randomness.len(), - self.source.interleaving_depth, + self.source.interleaving_depth(), ); - let mask_msg: Option<&[M::Target]> = match &mask_input { - MaskInput::Disabled => { - assert_eq!( - self.message_mask_length, 0, - "MaskInput::Disabled requires message_mask_length == 0" - ); - None - } - MaskInput::Enabled(mask) => { - assert_eq!( - mask.len(), - self.message_mask_length, - "mask_msg length must equal message_mask_length" - ); - Some(mask) - } - }; // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 let target_witness = self.target.commit(prover_state, &[&message]); + // Grind Lemma 9.9 OOD gap before α is sampled. + self.pow.prove(prover_state); + // Step 2-3: OOD challenge + answers — Construction 9.7 Steps 2-3, p.55 - // y := ze_ood(ρ) · [f; r; s] = f(α) + α^ℓ · (r,s)(α) let ood_points: Vec = prover_state.verifier_message_vec(self.out_domain_samples); - let msg_len = message.len(); - for &point in &ood_points { - let f_eval = univariate_evaluate(&message, point); - if let Some(mask) = mask_msg { - let mask_eval = univariate_evaluate(mask, point); - let shift = point.pow([msg_len as u64]); - prover_state.prover_message(&(f_eval + shift * mask_eval)); - } else { - prover_state.prover_message(&f_eval); - } - } + self.maybe_send_ood_answers(prover_state, &message, mask, &ood_points); // Step 4: in-domain queries — Construction 9.7 Step 4, p.55 let source_evaluations = self.source.open(prover_state, &[witness]); @@ -226,24 +261,13 @@ impl Config { // Covector update — sl' from Completeness proof (p.55-56) let eval_points = lift(self.source.embedding(), &source_evaluations.points); scalar_mul(covector, original_sl_coeff); - if self.message_mask_length == 0 { - // Non-ZK: single accumulate over all points - let all_points: Vec<_> = ood_points.iter().chain(&eval_points).copied().collect(); - let pows: Vec<_> = ood_rlc_coeffs - .iter() - .chain(in_domain_rlc_coeffs) - .copied() - .collect(); - geometric_accumulate(covector, pows, &all_points); - } else { - // ZK: OOD contributes to full [f; r; s], in-domain only to [f; r] - geometric_accumulate(covector, ood_rlc_coeffs.to_vec(), &ood_points); - geometric_accumulate( - &mut covector[..self.source.masked_message_length()], - in_domain_rlc_coeffs.to_vec(), - &eval_points, - ); - } + self.update_covector( + covector, + ood_rlc_coeffs, + &ood_points, + in_domain_rlc_coeffs, + &eval_points, + ); Witness { message, @@ -251,6 +275,67 @@ impl Config { } } + /// Send OOD answers `y_i = f(α_i) [+ α_i^ℓ · (r ‖ s)(α_i)]`. + /// In Standard mode the bracketed term is omitted. + fn maybe_send_ood_answers( + &self, + prover_state: &mut ProverState, + message: &[M::Target], + mask: &[M::Target], + ood_points: &[M::Target], + ) where + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + M::Target: Codec<[H::U]>, + { + let msg_len = message.len(); + for &point in ood_points { + let f_eval = univariate_evaluate(message, point); + let answer = match &self.mode { + CodeSwitchMode::Standard => f_eval, + CodeSwitchMode::ZeroKnowledge { .. } => { + let mask_eval = univariate_evaluate(mask, point); + let shift = point.pow([msg_len as u64]); + f_eval + shift * mask_eval + } + }; + prover_state.prover_message(&answer); + } + } + + /// Accumulate OOD and in-domain weights into the covector. + /// Standard mode treats all points uniformly; ZK mode applies OOD over + /// the full `[f; r; s]` and in-domain over the `[f; r]` prefix only. + fn update_covector( + &self, + covector: &mut [M::Target], + ood_rlc_coeffs: &[M::Target], + ood_points: &[M::Target], + in_domain_rlc_coeffs: &[M::Target], + in_domain_points: &[M::Target], + ) { + match &self.mode { + CodeSwitchMode::Standard => { + let all_points: Vec<_> = + ood_points.iter().chain(in_domain_points).copied().collect(); + let pows: Vec<_> = ood_rlc_coeffs + .iter() + .chain(in_domain_rlc_coeffs) + .copied() + .collect(); + geometric_accumulate(covector, pows, &all_points); + } + CodeSwitchMode::ZeroKnowledge { .. } => { + geometric_accumulate(covector, ood_rlc_coeffs.to_vec(), ood_points); + geometric_accumulate( + &mut covector[..self.source.masked_message_length()], + in_domain_rlc_coeffs.to_vec(), + in_domain_points, + ); + } + } + } + /// Verify the code-switch. /// /// `folding_randomness` is the **sumcheck folding randomness `γ`** the @@ -283,16 +368,18 @@ impl Config { verifier_state: &mut VerifierState, sum: &mut M::Target, folding_randomness: &[M::Target], - commitment: &IrsCommitment, - ) -> VerificationResult> + commitment: &IrsCommitment, + ) -> VerificationResult where H: DuplexSpongeInterface, Standard: Distribution, M::Target: Codec<[H::U]>, u8: Decoding<[H::U]>, + [u8; 32]: Decoding<[H::U]>, + U64: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { - verify!(1 << folding_randomness.len() == self.source.interleaving_depth); + verify!(1 << folding_randomness.len() == self.source.interleaving_depth()); let collapse_weights = eq_weights(folding_randomness); @@ -300,6 +387,9 @@ impl Config { // Mask oracle is committed in the shared mask tree by the orchestrator. let target_commitment = self.target.receive_commitment(verifier_state)?; + // Grind Lemma 9.9 OOD gap before α is sampled. + self.pow.verify(verifier_state)?; + // Step 2-3: OOD — Construction 9.7 Steps 2-3, p.55 // In ZK mode, ood_answers = f(α) + α^ℓ · (r,s)(α) where (r,s) is // the mask oracle message committed in the shared tree. @@ -312,7 +402,7 @@ impl Config { let source_evaluations = self.source.verify(verifier_state, &[commitment])?; let collapsed_values: Vec = source_evaluations .matrix - .chunks_exact(self.source.interleaving_depth) + .chunks_exact(self.source.interleaving_depth()) .map(|row| mixed_dot(self.source.embedding(), &collapse_weights, row)) .collect(); @@ -339,7 +429,7 @@ impl fmt::Display for Config { self.source, self.target, self.out_domain_samples, - self.message_mask_length != 0, + self.is_zk(), ) } } @@ -372,27 +462,25 @@ mod tests { select(valid_sizes), 0_usize..=3, // src_mask_len (source IRS randomness, post-fold) bool::ANY, // zk - 0_usize..=5, // ood (= code-switch t_ood) + 1_usize..=5, // ood (= code-switch t_ood; ≥ 1 per Construction 9.7) 0_usize..=5, // fresh_s_len (≥ ood for assumption (c)) select(vec![1_usize, 2, 4]), // ι_s (source interleaving) 0_usize..=10, // target.in_domain_samples (t'_in) - 0_usize..=10, // target.out_domain_samples (t'_out) ); scalars.prop_flat_map( - move |(size, src_mask_len, zk, ood, fresh_s_len, iota_s, t_in, t_out)| { + move |(size, src_mask_len, zk, ood, fresh_s_len, iota_s, t_in)| { // Bound 3 assumption (c): ℓ_zk - r ≥ t_ood ⇒ fresh_s_len ≥ ood. + // Also enforce `ℓ_zk = r + fresh_s_len > 0` so NonZeroUsize + // construction below is total in ZK mode. let fresh_s_len = if zk { - fresh_s_len.max(ood) + let min_fresh = usize::from(src_mask_len == 0); + fresh_s_len.max(ood).max(min_fresh) } else { fresh_s_len }; - // Bound 4 assumption (a): target.mask_length ≥ t' = t_in + t_out. - let target_mask = if zk { t_in + t_out } else { 0 }; - // ZK with source.mask_length = 0 is valid: the assert - // `source.mask_length == 0 || message_mask_length > 0` - // is trivially satisfied. Allows testing the corner - // where the mask oracle has only fresh randomness. + // Bound 4 assumption (a): target.mask_length ≥ t' = t_in + ood. + let target_mask = if zk { t_in + ood } else { 0 }; let source_mask = if zk { src_mask_len } else { 0 }; IrsConfig::arbitrary(embedding.clone(), 1, size, source_mask, iota_s) @@ -416,19 +504,31 @@ mod tests { ); let source = source.clone(); target.prop_map(move |mut target| { - // IrsConfig::arbitrary samples query counts in - // [0,10] independently of mask_length; pin them - // to the values target_mask was sized for so + // IrsConfig::arbitrary samples in_domain_samples + // in [0,10] independently of mask_length; pin it + // to the value target_mask was sized for so // assumption (a) holds. if zk { - target.in_domain_samples = t_in; - target.out_domain_samples = t_out; + target.set_in_domain_samples_for_test(t_in); } // r = post-fold randomness length (ι_s parallel // masks fold to a single length-mask_length chunk). - let r = source.mask_length; - let message_mask_length = if zk { r + fresh_s_len } else { 0 }; - Self::new(source.clone(), target, ood, message_mask_length) + let r = source.mask_length(); + let mode = if zk { + CodeSwitchMode::ZeroKnowledge { + message_mask_length: NonZeroUsize::new(r + fresh_s_len) + .expect("ZK ⇒ r + fresh_s_len > 0"), + } + } else { + CodeSwitchMode::Standard + }; + Self::new( + source.clone(), + target, + ood, + mode, + proof_of_work::Config::none(), + ) }) }) }) @@ -464,7 +564,7 @@ mod tests { where Standard: Distribution, { - let log_iota = config.source.interleaving_depth.trailing_zeros() as usize; + let log_iota = config.source.interleaving_depth().trailing_zeros() as usize; random_vector(rng, log_iota) } @@ -479,29 +579,21 @@ mod tests { where Standard: Distribution, { - if config.message_mask_length == 0 { + if !config.is_zk() { return Vec::new(); } // Lift ι parallel masks (total length source.mask_length × ι) and fold // chunks of length source.mask_length down to a single chunk. let raw = lift(config.source.embedding(), &source_witness.masks); - let mut mask = fold_chunks(&raw, config.source.mask_length, folding_randomness); + let mut mask = fold_chunks(&raw, config.source.mask_length(), folding_randomness); // Append fresh padding s of length message_mask_length - source.mask_length. mask.extend(random_vector::( rng, - config.message_mask_length - mask.len(), + config.message_mask_length() - mask.len(), )); mask } - fn mask_input(mask_msg: &[F]) -> MaskInput<'_, F> { - if mask_msg.is_empty() { - MaskInput::Disabled - } else { - MaskInput::Enabled(mask_msg) - } - } - fn test_config>(seed: u64, config: &Config>) where Standard: Distribution, @@ -510,7 +602,7 @@ mod tests { let mut rng = StdRng::seed_from_u64(seed); // Commit the full pre-fold vector of length source.vector_size // (= ι · message_length), which IRS encodes as ι parallel codewords. - let f_full: Vec = random_vector(&mut rng, config.source.vector_size); + let f_full: Vec = random_vector(&mut rng, config.source.vector_size()); let initial_sum: F = rng.gen(); let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); @@ -536,7 +628,7 @@ mod tests { &source_witness, &mut covector, &folding_randomness, - &mask_input(&mask_msg), + &mask_msg, ); let proof = prover_state.proof(); @@ -564,7 +656,7 @@ mod tests { Hash: ProverMessage<[u8]>, { let mut rng = StdRng::seed_from_u64(seed); - let f_full: Vec = random_vector(&mut rng, config.source.vector_size); + let f_full: Vec = random_vector(&mut rng, config.source.vector_size()); let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); @@ -602,7 +694,7 @@ mod tests { &source_witness, &mut covector, &folding_randomness, - &mask_input(&mask_msg), + &mask_msg, ); let proof = prover_state.proof(); @@ -635,7 +727,7 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut rng = StdRng::seed_from_u64(seed); - let f_full: Vec = random_vector(&mut rng, config.source.vector_size); + let f_full: Vec = random_vector(&mut rng, config.source.vector_size()); let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); @@ -659,7 +751,7 @@ mod tests { &source_witness, &mut covector, &folding_randomness, - &MaskInput::Disabled, + &[], ); let proof = prover_state.proof(); @@ -735,7 +827,7 @@ mod tests { crate::tests::init(); let configs = Config::arbitrary(Identity::::new()); proptest!(|(seed: u64, config in configs)| { - prop_assume!(config.source.in_domain_samples > 0); + prop_assume!(config.source.in_domain_samples() > 0); test_ior_identity_config(seed, &config); }); } @@ -743,14 +835,10 @@ mod tests { #[test] fn test_tampered_ood() { crate::tests::init(); - let configs = Config::arbitrary(Identity::::new()).prop_filter( - "non-ZK with ood > 0", - |config| { - config.message_mask_length == 0 - && config.source.mask_length == 0 - && config.out_domain_samples > 0 - }, - ); + let configs = Config::arbitrary(Identity::::new()) + .prop_filter("non-ZK", |config| { + !config.is_zk() && config.source.mask_length() == 0 + }); proptest!(|(seed: u64, config in configs)| { test_tampered_ood_config(seed, &config); }); diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index b5716f3a..2d0141c1 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -7,26 +7,16 @@ //! using an NTT friendly Reed-Solomon code to produce a `num_vectors * interleaving_depth` //! by `codeword_size` matrix. This matrix is committed using the [`matrix_commit`] protocol. //! -//! After committing the encoded matrix, the protocol generates a random Reed-Solomon code of -//! length `out_domain_samples` over an extension field `G` of `F` and encodes the original -//! matrix using this code to produce a `num_vectors` by `out_domain_samples` matrix over `G`. -//! Together, these two encoded matrices form a commitment to the original matrix. -//! //! On opening the commitment, the protocol randomly selects `in_domain_samples` rows and opens -//! it using the [`matrix_commit`] protocol. Sampling is done with replacement, so may produce -//! fewer than `in_domain_samples` distinct rows. This produces `in_domain_samples` evaluation -//! points in `F` and `in_domain_samples` by `num_vectors * interleaving_depth`. +//! them using the [`matrix_commit`] protocol. Sampling is done with replacement, so may produce +//! fewer than `in_domain_samples` distinct rows. //! -use std::{ - f64::{self, consts::LOG2_10}, - fmt, - ops::Neg, -}; +use std::{f64, fmt, num::NonZeroUsize}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; -use ordered_float::OrderedFloat; use serde::{Deserialize, Serialize}; +use thiserror::Error; #[cfg(feature = "tracing")] use tracing::instrument; @@ -37,76 +27,87 @@ use crate::{ }, engines::EngineId, hash::Hash, - protocols::{challenge_indices::challenge_indices, matrix_commit}, + protocols::{ + challenge_indices::challenge_indices, + matrix_commit, + params::{bounds::ood_per_sample_log2, regime::DecodingRegimeParams, spec::DecodingRegime}, + }, transcript::{ Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, VerificationResult, VerifierMessage, VerifierState, }, type_info::Typed, utils::{chunks_exact_or_empty, zip_strict}, - verify, }; +#[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] +pub enum IrsMode { + Standard, + ZeroKnowledge { mask_length: NonZeroUsize }, +} + +impl IrsMode { + /// Per-polynomial IRS randomness length. Returns 0 in Standard mode. + pub const fn mask_length(&self) -> usize { + match self { + Self::Standard => 0, + Self::ZeroKnowledge { mask_length } => mask_length.get(), + } + } +} + /// Commit to vectors over an fft-friendly field F +#[must_use] #[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config { /// Embedding into a (larger) field used for weights and drawing challenges. - pub embedding: Typed, + embedding: Typed, /// The number of vectors to commit to in one operation. - pub num_vectors: usize, + num_vectors: usize, /// The number of coefficients in each vector. - pub vector_size: usize, - - /// The number of masking values to add per codeword. - pub mask_length: usize, + vector_size: usize, /// The number of Reed-Solomon evaluation points. - pub codeword_length: usize, + codeword_length: usize, /// The number of independent codewords that are interleaved together. - pub interleaving_depth: usize, + interleaving_depth: usize, /// The matrix commitment configuration. - pub matrix_commit: matrix_commit::Config, + matrix_commit: matrix_commit::Config, - /// Slack to the Jonhnson bound in list decoding. - /// Zero indicates unique decoding. - pub johnson_slack: OrderedFloat, + /// Materialized Reed–Solomon decoding regime (Unique / Johnson w/ slack). + regime: DecodingRegimeParams, /// The number of in-domain samples. - pub in_domain_samples: usize, - - /// The number of out-of-domain samples. - pub out_domain_samples: usize, + in_domain_samples: usize, /// Whether to sort and deduplicate the in-domain samples. /// /// Deduplication can slightly reduce proof size and prover/verifier /// complexity, but it makes transcript pattern and control flow /// non-deterministic. - pub deduplicate_in_domain: bool, + deduplicate_in_domain: bool, + + /// Standard / ZeroKnowledge. + mode: IrsMode, } #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] #[must_use] -pub struct Witness -where - G: Field, -{ +pub struct Witness { pub masks: Vec, pub matrix: Vec, pub matrix_witness: matrix_commit::Witness, - pub out_of_domain: Evaluations, } #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] #[must_use] -pub struct Commitment { - matrix_commitment: matrix_commit::Commitment, - out_of_domain: Evaluations, +pub struct Commitment { + pub matrix_commitment: matrix_commit::Commitment, } /// Interleaved Reed-Solomon code. @@ -121,66 +122,82 @@ pub struct Evaluations { pub matrix: Vec, } +/// Named-field inputs to [`Config::new`] / [`Config::try_new`]. +/// +/// `num_vectors`, `vector_size`, and `interleaving_depth` share a primitive +/// type; named construction keeps call sites swap-proof. +#[derive(Debug, Clone)] +pub struct IrsParams { + pub security_target: f64, + pub decoding_regime: DecodingRegime, + pub hash_id: EngineId, + pub num_vectors: usize, + pub vector_size: usize, + pub interleaving_depth: usize, + pub rate: f64, + pub mode: IrsMode, +} + +/// The computed codeword length exceeds the NTT engine's supported order. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[error("codeword length {length} exceeds the NTT engine's supported order")] +pub struct CodewordLengthError { + pub length: usize, +} + impl Config { - pub fn new( - security_target: f64, - unique_decoding: bool, - hash_id: EngineId, - num_vectors: usize, - vector_size: usize, - interleaving_depth: usize, - rate: f64, - ) -> Self + /// Panicking version of [`Config::try_new`] for call sites that construct + /// from already-vetted parameters. + /// + /// # Panics + /// + /// If the codeword length exceeds the NTT engine's supported order. + pub fn new(params: IrsParams) -> Self where M: Default, { + Self::try_new(params).expect("IRS config construction failed") + } + + /// # Errors + /// + /// [`CodewordLengthError`] when `masked_message_length / rate` exceeds the + /// NTT engine's supported order. + pub fn try_new(params: IrsParams) -> Result + where + M: Default, + { + let IrsParams { + security_target, + decoding_regime, + hash_id, + num_vectors, + vector_size, + interleaving_depth, + rate, + mode, + } = params; assert!(vector_size.is_multiple_of(interleaving_depth)); assert!(rate > 0. && rate <= 1.); - let message_length = vector_size / interleaving_depth; - #[allow(clippy::cast_sign_loss)] - let codeword_length = (message_length as f64 / rate).ceil() as usize; - let rate = message_length as f64 / codeword_length as f64; - - // Pick in- and out-of-domain samples. - // η = slack to Johnson bound. We pick η = √ρ / 20. - // TODO: Optimize picking η. - let johnson_slack = if unique_decoding { - 0.0 - } else { - rate.sqrt() / 20. - }; - let out_domain_samples = { - let list_size = 1. / (2. * johnson_slack * rate.sqrt()); - num_ood_samples( - unique_decoding, - security_target, - M::Target::field_size_bits(), - list_size, - vector_size, - ) - }; + let masked_message_length = vector_size / interleaving_depth + mode.mask_length(); + // `interleaved_encode` requires `codeword_length` to divide the NTT root + // order. `masked_message_length` is allowed to be arbitrary (the coset + // NTT zero-extends internally), so we only round the codeword side here. #[allow(clippy::cast_sign_loss)] - let in_domain_samples = { - // Query error is (1 - δ)^q, so we compute 1 - δ - let per_sample = if unique_decoding { - // Unique decoding bound: δ = (1 - ρ) / 2 - f64::midpoint(1., rate) - } else { - // Johnson bound: δ = 1 - √ρ - η - rate.sqrt() + johnson_slack - }; - (security_target / (-per_sample.log2())).ceil() as usize - }; - debug_assert_eq!( - in_domain_samples, - num_in_domain_queries(unique_decoding, security_target, rate) - ); + let raw_codeword_length = (masked_message_length as f64 / rate).ceil() as usize; + let codeword_length = + ntt::next_order::(raw_codeword_length).ok_or(CodewordLengthError { + length: raw_codeword_length, + })?; + let rate = masked_message_length as f64 / codeword_length as f64; + + let regime = DecodingRegimeParams::from_policy(decoding_regime, rate); + let in_domain_samples = num_in_domain_queries(decoding_regime, security_target, rate).get(); - Self { + Ok(Self { embedding: Typed::::default(), num_vectors, vector_size, - mask_length: 0, codeword_length, interleaving_depth, matrix_commit: matrix_commit::Config::with_hash( @@ -188,11 +205,57 @@ impl Config { codeword_length, interleaving_depth * num_vectors, ), - johnson_slack: OrderedFloat(johnson_slack), + regime, in_domain_samples, - out_domain_samples, deduplicate_in_domain: false, - } + mode, + }) + } + + pub const fn num_vectors(&self) -> usize { + self.num_vectors + } + + pub const fn vector_size(&self) -> usize { + self.vector_size + } + + pub const fn codeword_length(&self) -> usize { + self.codeword_length + } + + pub const fn interleaving_depth(&self) -> usize { + self.interleaving_depth + } + + pub const fn matrix_commit(&self) -> &matrix_commit::Config { + &self.matrix_commit + } + + pub const fn regime(&self) -> DecodingRegimeParams { + self.regime + } + + pub const fn in_domain_samples(&self) -> usize { + self.in_domain_samples + } + + pub const fn deduplicate_in_domain(&self) -> bool { + self.deduplicate_in_domain + } + + pub const fn mode(&self) -> &IrsMode { + &self.mode + } + + #[cfg(test)] + pub(crate) const fn set_vector_size_for_test(&mut self, vector_size: usize) { + self.vector_size = vector_size; + } + + #[cfg(test)] + pub(crate) const fn set_in_domain_samples_for_test(&mut self, in_domain_samples: usize) { + self.in_domain_samples = in_domain_samples; } pub const fn num_cols(&self) -> usize { @@ -216,9 +279,14 @@ impl Config { self.vector_size / self.interleaving_depth } + /// Per-polynomial IRS randomness length. Returns 0 in Standard mode. + pub const fn mask_length(&self) -> usize { + self.mode.mask_length() + } + /// Message length including mask coefficients. pub fn masked_message_length(&self) -> usize { - self.message_length() + self.mask_length + self.message_length() + self.mask_length() } pub fn evaluation_points(&self, indices: &[usize]) -> Vec { @@ -233,58 +301,34 @@ impl Config { self.masked_message_length() as f64 / self.codeword_length as f64 } - pub fn unique_decoding(&self) -> bool { - self.out_domain_samples == 0 && self.johnson_slack == 0.0 + pub const fn unique_decoding(&self) -> bool { + self.regime.is_unique() } - /// Compute a list size bound. - pub fn list_size(&self) -> f64 { - if self.unique_decoding() { - 1. - } else { - // This is the Johnson bound $1 / (2 η √ρ)$. - 1. / (2. * self.johnson_slack.into_inner() * self.rate().sqrt()) - } + fn log_inv_rate(&self) -> f64 { + -self.rate().log2() } - /// Round-by-round soundness of the out-of-domain samples in bits. - pub fn rbr_ood_sample(&self) -> f64 { - let list_size = self.list_size(); - let log_field_size = M::Target::field_size_bits(); - // See [STIR] lemma 4.5. - let l_choose_2 = list_size * (list_size - 1.) / 2.; - let log_per_sample = ((self.vector_size - 1) as f64).log2() - log_field_size; - -l_choose_2.log2() - self.out_domain_samples as f64 * log_per_sample + /// Compute a list size bound. + pub fn list_size(&self) -> f64 { + let log_degree = (self.masked_message_length() as f64).log2(); + self.regime.list_size(log_degree, self.log_inv_rate()) } /// Round-by-round soundness of the in-domain queries in bits. pub fn rbr_queries(&self) -> f64 { - let per_sample = if self.unique_decoding() { - // 1 - δ = 1 - (1 + ρ) / 2 - f64::midpoint(1., self.rate()) - } else { - // 1 - δ = sqrt(ρ) + η - self.rate().sqrt() + self.johnson_slack.into_inner() - }; - self.in_domain_samples as f64 * per_sample.log2().neg() + // Query error is (1 - δ)^q in bits = -q · log2(1 - δ). + -(self.in_domain_samples as f64) * self.regime.one_minus_distance_log2(self.log_inv_rate()) } - // Compute the proximity gaps term of the fold + /// Round-by-round soundness of the proximity-gaps fold in bits. + /// See WHIR Theorem 4.8. pub fn rbr_soundness_fold_prox_gaps(&self) -> f64 { - let log_field_size = M::Target::field_size_bits(); - let log_inv_rate = self.rate().log2().neg(); - let log_k = (self.masked_message_length() as f64).log2(); - // See WHIR Theorem 4.8 - // Recall, at each round we are only folding by two at a time - let error = if self.unique_decoding() { - log_k + log_inv_rate - } else { - let log_eta = self.johnson_slack.into_inner().log2(); - // Make sure η hits the min bound. - assert!(log_eta >= -(0.5 * log_inv_rate + LOG2_10 + 1.0) - 1e-6); - 7. * LOG2_10 + 3.5 * log_inv_rate + 2. * log_k - }; - log_field_size - error + -self.regime.eps_mca_log2( + self.log_inv_rate(), + self.masked_message_length(), + M::Target::field_size_bits(), + ) } /// Commit to one or more vectors. @@ -293,7 +337,7 @@ impl Config { &self, prover_state: &mut ProverState, vectors: &[&[M::Source]], - ) -> Witness + ) -> Witness where Standard: Distribution, H: DuplexSpongeInterface, @@ -311,7 +355,7 @@ impl Config { assert!(vectors.iter().all(|p| p.len() == self.vector_size)); // Generate random mask - let masks = random_vector(prover_state.rng(), self.mask_length * self.num_messages()); + let masks = random_vector(prover_state.rng(), self.mask_length() * self.num_messages()); // Interleaved RS Encode the vectors let messages = vectors @@ -323,29 +367,10 @@ impl Config { // Commit to the matrix let matrix_witness = self.matrix_commit.commit(prover_state, &matrix); - // Handle out-of-domain points and values - // TODO : Remove this logic after main whir protocol is updated - // as this is not required in the new construction. This will be - // removed in next PR (Parameter Selection) - let oods_points: Vec = - prover_state.verifier_message_vec(self.out_domain_samples); - let mut oods_matrix = Vec::with_capacity(self.out_domain_samples * self.num_vectors); - for &point in &oods_points { - for &vector in vectors { - let value = mixed_univariate_evaluate(&*self.embedding, vector, point); - prover_state.prover_message(&value); - oods_matrix.push(value); - } - } - Witness { masks, matrix, matrix_witness, - out_of_domain: Evaluations { - points: oods_points, - matrix: oods_matrix, - }, } } @@ -354,24 +379,67 @@ impl Config { pub fn receive_commitment( &self, verifier_state: &mut VerifierState, - ) -> VerificationResult> + ) -> VerificationResult where H: DuplexSpongeInterface, Hash: ProverMessage<[H::U]>, M::Target: Codec<[H::U]>, { let matrix_commitment = self.matrix_commit.receive_commitment(verifier_state)?; - let oods_points: Vec = - verifier_state.verifier_message_vec(self.out_domain_samples); - let oods_matrix = - verifier_state.prover_messages_vec(self.out_domain_samples * self.num_vectors)?; - Ok(Commitment { - matrix_commitment, - out_of_domain: Evaluations { - points: oods_points, - matrix: oods_matrix, - }, - }) + Ok(Commitment { matrix_commitment }) + } + + /// Commit to vectors and run the legacy WHIR OOD step in one call. + /// + /// Layered helper bundling `commit` + the OOD message exchange (sample + /// `out_domain_samples` random points, send each vector's evaluation at + /// each point). Used by the legacy WHIR protocol while the OOD step + /// is still part of the per-commit protocol shape; the new construction + /// (Construction 9.7) handles OOD at the code-switch level instead. + #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] + pub fn commit_with_ood( + &self, + prover_state: &mut ProverState, + vectors: &[&[M::Source]], + out_domain_samples: usize, + ) -> (Witness, Evaluations) + where + Standard: Distribution, + H: DuplexSpongeInterface, + R: RngCore + CryptoRng, + M::Target: Codec<[H::U]>, + Hash: ProverMessage<[H::U]>, + { + let witness = self.commit(prover_state, vectors); + let points: Vec = prover_state.verifier_message_vec(out_domain_samples); + let mut matrix = Vec::with_capacity(out_domain_samples * vectors.len()); + for &point in &points { + for &vector in vectors { + let value = mixed_univariate_evaluate(&*self.embedding, vector, point); + prover_state.prover_message(&value); + matrix.push(value); + } + } + (witness, Evaluations { points, matrix }) + } + + /// Receive a commitment and the legacy WHIR OOD evaluations in one call. + /// Verifier mirror of `commit_with_ood`. + #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] + pub fn receive_commitment_with_ood( + &self, + verifier_state: &mut VerifierState, + out_domain_samples: usize, + ) -> VerificationResult<(Commitment, Evaluations)> + where + H: DuplexSpongeInterface, + Hash: ProverMessage<[H::U]>, + M::Target: Codec<[H::U]>, + { + let commitment = self.receive_commitment(verifier_state)?; + let points: Vec = verifier_state.verifier_message_vec(out_domain_samples); + let matrix = verifier_state.prover_messages_vec(out_domain_samples * self.num_vectors)?; + Ok((commitment, Evaluations { points, matrix })) } /// Opens the commitment and returns the evaluations of the vectors. @@ -385,7 +453,7 @@ impl Config { pub fn open( &self, prover_state: &mut ProverState, - witnesses: &[&Witness], + witnesses: &[&Witness], ) -> Evaluations where H: DuplexSpongeInterface, @@ -395,11 +463,6 @@ impl Config { { for witness in witnesses { assert_eq!(witness.matrix.len(), self.size()); - assert_eq!(witness.out_of_domain.points.len(), self.out_domain_samples); - assert_eq!( - witness.out_of_domain.matrix.len(), - self.out_domain_samples * self.num_vectors - ); } // Get in-domain openings @@ -438,20 +501,13 @@ impl Config { pub fn verify( &self, verifier_state: &mut VerifierState, - commitments: &[&Commitment], + commitments: &[&Commitment], ) -> VerificationResult> where H: DuplexSpongeInterface, u8: Decoding<[H::U]>, Hash: ProverMessage<[H::U]>, { - for commitment in commitments { - verify!(commitment.out_of_domain.points.len() == self.out_domain_samples); - verify!( - commitment.out_of_domain.matrix.len() == self.num_vectors * self.out_domain_samples - ); - } - // Get in-domain openings let (indices, points) = self.in_domain_challenges(verifier_state); @@ -500,28 +556,6 @@ impl Config { } } -impl Commitment { - /// Returns the out-of-domain evaluations. - pub const fn out_of_domain(&self) -> &Evaluations { - &self.out_of_domain - } - - pub fn num_vectors(&self) -> usize { - self.out_of_domain().num_columns() - } -} - -impl Witness { - /// Returns the out-of-domain evaluations. - pub const fn out_of_domain(&self) -> &Evaluations { - &self.out_of_domain - } - - pub fn num_vectors(&self) -> usize { - self.out_of_domain().num_columns() - } -} - impl Evaluations { pub const fn num_points(&self) -> usize { self.points.len() @@ -568,11 +602,7 @@ impl fmt::Display for Config { self.num_vectors, self.vector_size, self.interleaving_depth, )?; write!(f, " rate 2⁻{:.2}", -self.rate().log2())?; - write!( - f, - " samples {} in- {} out-domain", - self.in_domain_samples, self.out_domain_samples - ) + write!(f, " samples {} in-domain", self.in_domain_samples) } } @@ -582,19 +612,19 @@ impl fmt::Display for Config { /// where `L` is the list size and `degree` is the polynomial degree bound. /// See [STIR] Lemma 4.5. #[allow(clippy::cast_sign_loss)] -pub(crate) fn num_ood_samples( - unique_decoding: bool, +pub fn num_ood_samples( + decoding_regime: DecodingRegime, security_target: f64, field_size_bits: f64, list_size: f64, degree: usize, ) -> usize { - if unique_decoding { + if matches!(decoding_regime, DecodingRegime::Unique) { return 0; } - let l_choose_2 = list_size * (list_size - 1.) / 2.; - let log_per_sample = field_size_bits - ((degree - 1) as f64).log2(); + let log_per_sample = -ood_per_sample_log2(degree, field_size_bits); assert!(log_per_sample > 0.); + let l_choose_2 = list_size * (list_size - 1.) / 2.; ((security_target + l_choose_2.log2()) / log_per_sample) .ceil() .max(1.) as usize @@ -602,31 +632,22 @@ pub(crate) fn num_ood_samples( /// Return the number of in-domain queries. /// +/// Always ≥ 1 — the type carries that invariant so callers don't need to +/// re-prove it locally. +/// /// This is used by [`whir_zk`]. // TODO: A method with cleaner abstraction. #[allow(clippy::cast_sign_loss)] pub(crate) fn num_in_domain_queries( - unique_decoding: bool, + decoding_regime: DecodingRegime, security_target: f64, rate: f64, -) -> usize { - // Pick in- and out-of-domain samples. - // η = slack to Johnson bound. We pick η = √ρ / 20. - // TODO: Optimize picking η. - let johnson_slack = if unique_decoding { - 0.0 - } else { - rate.sqrt() / 20. - }; - // Query error is (1 - δ)^q, so we compute 1 - δ - let per_sample = if unique_decoding { - // Unique decoding bound: δ = (1 - ρ) / 2 - f64::midpoint(1., rate) - } else { - // Johnson bound: δ = 1 - √ρ - η - rate.sqrt() + johnson_slack - }; - (security_target / (-per_sample.log2())).ceil() as usize +) -> NonZeroUsize { + let regime = DecodingRegimeParams::from_policy(decoding_regime, rate); + // Query error is (1 - δ)^q in bits = -q · log2(1 - δ). + let log_one_minus_delta = regime.one_minus_distance_log2(-rate.log2()); + let q = (security_target / -log_one_minus_delta).ceil() as usize; + NonZeroUsize::new(q).unwrap_or(NonZeroUsize::MIN) } #[cfg(test)] @@ -683,24 +704,27 @@ pub(crate) mod tests { ) }); - (codeword_matrix, 0_usize..=10, 0_usize..=10, bool::ANY).prop_map( + (codeword_matrix, 0_usize..=10, bool::ANY).prop_map( move |( (codeword_length, matrix_commit), in_domain_samples, - out_domain_samples, - deduplicate_in_domain, - )| Self { - embedding: Typed::new(embedding.clone()), - num_vectors, - vector_size, - mask_length, - codeword_length, - interleaving_depth, - matrix_commit, - johnson_slack: OrderedFloat::default(), - in_domain_samples, - out_domain_samples, deduplicate_in_domain, + )| { + let mode = NonZeroUsize::new(mask_length).map_or(IrsMode::Standard, |n| { + IrsMode::ZeroKnowledge { mask_length: n } + }); + Self { + embedding: Typed::new(embedding.clone()), + num_vectors, + vector_size, + codeword_length, + interleaving_depth, + matrix_commit, + regime: DecodingRegimeParams::Unique, + in_domain_samples, + deduplicate_in_domain, + mode, + } }, ) } @@ -732,30 +756,6 @@ pub(crate) mod tests { &mut prover_state, &vectors.iter().map(|p| p.as_slice()).collect::>(), ); - assert_eq!( - witness.out_of_domain().points.len(), - config.out_domain_samples - ); - assert_eq!( - witness.out_of_domain().matrix.len(), - config.out_domain_samples * config.num_vectors - ); - if config.num_vectors > 0 { - for (point, evals) in zip_strict( - witness.out_of_domain().points.iter(), - witness - .out_of_domain() - .matrix - .chunks_exact(config.num_vectors), - ) { - for (vector, expected) in zip_strict(vectors.iter(), evals.iter()) { - assert_eq!( - mixed_univariate_evaluate(config.embedding(), vector, *point), - *expected - ); - } - } - } let in_domain_evals = config.open(&mut prover_state, &[&witness]); if config.deduplicate_in_domain { // Sorting is over index order, not points @@ -773,7 +773,11 @@ pub(crate) mod tests { in_domain_evals.matrix.len(), in_domain_evals.points.len() * config.num_vectors * config.interleaving_depth ); - if config.num_vectors > 0 { + // Value-correctness assertion only valid in non-ZK mode: in ZK the + // encoding is `Enc(f, r) = f(x) + x^ℓ · r(x)`, so opened values + // include the mask term. The lifecycle round-trip (open/verify + // agreement below) covers both modes. + if config.num_vectors > 0 && config.mask_length() == 0 { let base = config.vector_size / config.interleaving_depth; for (point, evals) in zip_strict( &in_domain_evals.points, @@ -799,7 +803,6 @@ pub(crate) mod tests { // Verifier let mut verifier_state = VerifierState::new_std(&ds, &proof); let commitment = config.receive_commitment(&mut verifier_state).unwrap(); - assert_eq!(commitment.out_of_domain(), witness.out_of_domain()); let verifier_in_domain_evals = config.verify(&mut verifier_state, &[&commitment]).unwrap(); assert_eq!(&verifier_in_domain_evals, &in_domain_evals); verifier_state.check_eof().unwrap(); @@ -816,13 +819,13 @@ pub(crate) mod tests { .collect::>(); let size = select(valid_sizes); - let config = (0_usize..=3, size, 1_usize..=10).prop_flat_map( - |(num_vectors, size, interleaving_depth)| { + let config = (0_usize..=3, size, 1_usize..=10, 0_usize..=8).prop_flat_map( + |(num_vectors, size, interleaving_depth, mask_length)| { Config::arbitrary( embedding.clone(), num_vectors, size * interleaving_depth, - 0, + mask_length, interleaving_depth, ) }, diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index 1d632ca0..5f285872 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -47,12 +47,13 @@ use serde::{Deserialize, Serialize}; use crate::{ algebra::{embedding::Identity, random_vector, scalar_mul_add_new, univariate_evaluate}, hash::Hash, - protocols::irs_commit::{ - Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness, + protocols::{ + irs_commit::{Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness}, + proof_of_work, }, transcript::{ - Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, VerificationResult, - VerifierMessage, VerifierState, + codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverMessage, ProverState, + VerificationResult, VerifierMessage, VerifierState, }, utils::zip_strict, verify, @@ -61,11 +62,13 @@ use crate::{ /// Mask proximity configuration. /// /// Wraps an IRS config for the shared mask tree and the number of mask pairs. +#[must_use] #[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config { - pub c_zk_commit: IrsConfig>, - pub num_masks: usize, + c_zk_commit: IrsConfig>, + num_masks: usize, + pow: proof_of_work::Config, } /// Prover output from the commit phase. @@ -76,33 +79,49 @@ pub struct Witness { } /// Verifier output from the commit phase. -pub type Commitment = IrsCommitment; +pub type Commitment = IrsCommitment; impl Config { - pub fn new(c_zk_commit: IrsConfig>, num_masks: usize) -> Self { + /// Required `c_zk.num_vectors` for `num_masks` originals: one fresh + /// mask-of-mask per original (Construction 7.2 originals + fresh pairs). + pub const fn num_vectors_for(num_masks: usize) -> usize { + 2 * num_masks + } + + pub fn new( + c_zk_commit: IrsConfig>, + num_masks: usize, + pow: proof_of_work::Config, + ) -> Self { assert_eq!( - c_zk_commit.num_vectors, - 2 * num_masks, + c_zk_commit.num_vectors(), + Self::num_vectors_for(num_masks), "c_zk.num_vectors must be 2 * num_masks" ); assert_eq!( - c_zk_commit.interleaving_depth, 1, + c_zk_commit.interleaving_depth(), + 1, "mask proximity requires interleaving_depth = 1" ); - // OOD evaluations are sent in the clear during IRS commit/receive, - // which would leak raw mask values before the γ-combination and - // break the ZK contract. The OOD path in irs_commit is slated for - // removal in the new construction; until then, enforce zero here. - assert_eq!( - c_zk_commit.out_domain_samples, 0, - "mask proximity requires out_domain_samples = 0 (OOD openings would leak raw mask evaluations)" - ); Self { c_zk_commit, num_masks, + pow, } } + pub const fn c_zk_commit(&self) -> &IrsConfig> { + &self.c_zk_commit + } + + pub const fn num_masks(&self) -> usize { + self.num_masks + } + + pub const fn pow(&self) -> proof_of_work::Config { + self.pow + } + /// Commit all masks and their mask-of-masks in a single shared tree. /// /// Samples n fresh mask-of-mask polynomials, combines them with the @@ -121,12 +140,12 @@ impl Config { { assert_eq!(original_msgs.len(), self.num_masks); for msg in original_msgs { - assert_eq!(msg.len(), self.c_zk_commit.vector_size); + assert_eq!(msg.len(), self.c_zk_commit.vector_size()); } // Sample fresh mask-of-masks let fresh_msgs: Vec> = (0..self.num_masks) - .map(|_| random_vector(prover_state.rng(), self.c_zk_commit.vector_size)) + .map(|_| random_vector(prover_state.rng(), self.c_zk_commit.vector_size())) .collect(); // Tree layout: [originals..., freshes...] @@ -148,7 +167,7 @@ impl Config { pub fn receive_commitment( &self, verifier_state: &mut VerifierState, - ) -> VerificationResult> + ) -> VerificationResult where F: Codec<[H::U]>, H: DuplexSpongeInterface, @@ -169,17 +188,22 @@ impl Config { R: RngCore + CryptoRng, Standard: Distribution, u8: Decoding<[H::U]>, + [u8; 32]: Decoding<[H::U]>, + U64: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { assert_eq!(original_msgs.len(), self.num_masks); assert_eq!(witness.fresh_msgs.len(), self.num_masks); + // Grind the Lemma 7.4 γ-combination gap before γ is sampled. + self.pow.prove(prover_state); + // Step 1: receive combination randomness γ let gamma: F = prover_state.verifier_message(); // Step 2: compute and send combined polynomials + IRS randomness let irs_masks_per_vector = - self.c_zk_commit.mask_length * self.c_zk_commit.interleaving_depth; + self.c_zk_commit.mask_length() * self.c_zk_commit.interleaving_depth(); assert_eq!( witness.mask_witness.masks.len(), 2 * self.num_masks * irs_masks_per_vector @@ -214,21 +238,26 @@ impl Config { pub fn verify( &self, verifier_state: &mut VerifierState, - commitment: &Commitment, + commitment: &Commitment, ) -> VerificationResult<()> where F: Codec<[H::U]>, H: DuplexSpongeInterface, u8: Decoding<[H::U]>, + [u8; 32]: Decoding<[H::U]>, + U64: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { + // Grind the Lemma 7.4 γ-combination gap before γ is sampled. + self.pow.verify(verifier_state)?; + // Step 1: send combination randomness γ let gamma: F = verifier_state.verifier_message(); // Step 2: read combined polynomials + IRS randomness let msg_len = self.c_zk_commit.message_length(); let irs_masks_per_vector = - self.c_zk_commit.mask_length * self.c_zk_commit.interleaving_depth; + self.c_zk_commit.mask_length() * self.c_zk_commit.interleaving_depth(); let has_irs_masks = irs_masks_per_vector > 0; let mut combined_msgs = Vec::with_capacity(self.num_masks); let mut combined_rs: Option>> = @@ -313,18 +342,16 @@ mod tests { .prop_flat_map(|(num_masks, vector_size, mask_length)| { let c_zk = IrsConfig::>::arbitrary( Identity::new(), - 2 * num_masks, + Self::num_vectors_for(num_masks), vector_size, mask_length, 1, ); (Just(num_masks), c_zk) }) - .prop_filter( - "mask proximity requires out_domain_samples = 0", - |(_, c_zk)| c_zk.out_domain_samples == 0, - ) - .prop_map(|(num_masks, c_zk)| Self::new(c_zk, num_masks)) + .prop_map(|(num_masks, c_zk)| { + Self::new(c_zk, num_masks, proof_of_work::Config::none()) + }) } } @@ -341,7 +368,7 @@ mod tests { let mut rng = StdRng::seed_from_u64(seed); let original_msgs: Vec> = (0..config.num_masks) - .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) + .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size())) .collect(); let mut prover_state = ProverState::new_std(&ds); @@ -415,7 +442,7 @@ mod tests { let mut rng = StdRng::seed_from_u64(seed); let original_msgs: Vec> = (0..config.num_masks) - .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) + .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size())) .collect(); let mut prover_state = ProverState::new_std(&ds); @@ -446,7 +473,7 @@ mod tests { let mut rng = StdRng::seed_from_u64(seed); let original_msgs: Vec> = (0..config.num_masks) - .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) + .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size())) .collect(); let mut prover_state = ProverState::new_std(&ds); @@ -454,7 +481,7 @@ mod tests { let gamma: F = prover_state.verifier_message(); let irs_masks_per_vector = - config.c_zk_commit.mask_length * config.c_zk_commit.interleaving_depth; + config.c_zk_commit.mask_length() * config.c_zk_commit.interleaving_depth(); for (i, (orig_msg, fresh_msg)) in original_msgs .iter() @@ -509,7 +536,7 @@ mod tests { fn test_tampered_mask_rejected() { crate::tests::init(); proptest!(|(seed: u64, config in Config::::arbitrary())| { - prop_assume!(config.c_zk_commit.in_domain_samples > 0); + prop_assume!(config.c_zk_commit.in_domain_samples() > 0); test_tampered_mask_config(seed, &config); }); } @@ -518,7 +545,7 @@ mod tests { fn test_tampered_combined_msg_rejected() { crate::tests::init(); proptest!(|(seed: u64, config in Config::::arbitrary())| { - prop_assume!(config.c_zk_commit.in_domain_samples > 0); + prop_assume!(config.c_zk_commit.in_domain_samples() > 0); test_tampered_combined_msg_config(seed, &config); }); } diff --git a/src/protocols/mod.rs b/src/protocols/mod.rs index 17e58040..64f1e0b9 100644 --- a/src/protocols/mod.rs +++ b/src/protocols/mod.rs @@ -16,6 +16,7 @@ pub mod irs_commit; pub mod mask_proximity; pub mod matrix_commit; pub mod merkle_tree; +pub mod params; pub mod proof_of_work; pub mod sumcheck; pub mod whir; diff --git a/src/protocols/params/basecase.rs b/src/protocols/params/basecase.rs new file mode 100644 index 00000000..87ce439b --- /dev/null +++ b/src/protocols/params/basecase.rs @@ -0,0 +1,198 @@ +//! Basecase (Construction 7.2, p.43) parameter selection + γ-combination bound. + +use ark_ff::Field; + +use crate::{ + algebra::{embedding::Identity, fields::FieldWithSize}, + bits::Bits, + protocols::{ + basecase::{self, Config as BasecaseConfig}, + irs_commit::Config as IrsConfig, + params::{ + error::{grind_to_at, DeriveError, Pow}, + irs_commit as irs_params, + protocol_config::BasecasePlan, + spec::{Mode as SpecMode, OodSampleBudget, RoundContext, SecuritySpec}, + sumcheck as sumcheck_params, + }, + proof_of_work::Config as PowConfig, + sumcheck::{self, Config as SumcheckConfig}, + }, +}; + +pub fn solve( + spec: &SecuritySpec, + vector_size: usize, + log_inv_rate: u32, +) -> Result, DeriveError> { + assert!(vector_size > 0, "basecase requires vector_size ≥ 1"); + + let ctx = RoundContext { + vector_size, + log_inv_rate, + folding_factor: 0, + }; + let commit = irs_params::solve(spec, &ctx, OodSampleBudget::ZERO)?; + + let sumcheck_analytic = sumcheck_params::analytic_error_bits(&commit, None); + let sumcheck_pow = grind_to_at(spec, sumcheck_analytic, Pow::BasecaseSumcheck)?; + let sumcheck = SumcheckConfig::new( + vector_size, + sumcheck_pow, + vector_size.next_power_of_two().trailing_zeros() as usize, + sumcheck::SumcheckMode::Standard, + ); + + let gamma_analytic = analytic_error_bits(&commit); + let (mode, pow) = match spec.mode { + SpecMode::Standard => (basecase::BasecaseMode::Standard, PowConfig::none()), + SpecMode::ZeroKnowledge => ( + basecase::BasecaseMode::ZeroKnowledge, + grind_to_at(spec, gamma_analytic, Pow::BasecaseGammaCombination)?, + ), + }; + + Ok(BasecasePlan::new( + BasecaseConfig::new(commit, sumcheck, mode, pow), + sumcheck_analytic, + gamma_analytic, + )) +} + +/// γ-combination soundness (Lemma 7.4 combination-randomness slot, paper p.45). +pub fn analytic_error_bits(commit: &IrsConfig>) -> Bits { + let field_bits = F::field_size_bits(); + let log_list = commit.list_size().log2(); + let prox_gaps = commit.rbr_soundness_fold_prox_gaps(); + let poly_id = field_bits - log_list; + Bits::new(prox_gaps.min(poly_id).max(0.0)) +} + +impl BasecaseConfig { + /// Analytic soundness bits (excluding PoW): `min(sumcheck round error, γ-slot error)`. + pub fn analytic_bits(&self) -> Bits { + let sumcheck_term = f64::from(sumcheck_params::analytic_error_bits(self.commit(), None)); + let min_bits = match self.mode() { + basecase::BasecaseMode::Standard => sumcheck_term, + basecase::BasecaseMode::ZeroKnowledge => { + sumcheck_term.min(f64::from(analytic_error_bits(self.commit()))) + } + }; + Bits::new(min_bits.max(0.0)) + } +} + +#[cfg(test)] +mod tests { + use proptest::prelude::*; + + use super::*; + use crate::protocols::params::test_utils::{ + arb_standard_spec, arb_zk_spec, assert_close, assert_pow_closes_gap, deterministic_spec, + TestField, TEST_TARGET_RANGE, + }; + + const FIXTURE_VECTOR_SIZE: usize = 16; + const FIXTURE_LOG_INV_RATE: u32 = 2; + + fn arb_dims() -> impl Strategy { + (1u32..=4, 1u32..=3) + } + + #[test] + fn analytic_error_formula() { + use crate::protocols::params::{ + irs_commit as irs_params, + spec::{Mode, OodSampleBudget, RoundContext}, + }; + + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = RoundContext { + vector_size: FIXTURE_VECTOR_SIZE, + log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: 0, + }; + let commit: IrsConfig> = + irs_params::solve(&spec, &ctx, OodSampleBudget::ZERO).expect("IRS fixture must solve"); + + let got = f64::from(analytic_error_bits(&commit)); + let field_bits = TestField::field_size_bits(); + let log_list = commit.list_size().log2(); + let prox_gaps = commit.rbr_soundness_fold_prox_gaps(); + let poly_id = field_bits - log_list; + let expected = prox_gaps.min(poly_id).max(0.0); + + assert_close(got, expected); + } + + #[test] + fn analytic_error_uses_eps_mca_when_limiting() { + use crate::protocols::params::{ + irs_commit as irs_params, + spec::{Mode, OodSampleBudget, RoundContext}, + }; + + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = RoundContext { + vector_size: FIXTURE_VECTOR_SIZE, + log_inv_rate: 1, + folding_factor: 0, + }; + let commit: IrsConfig> = + irs_params::solve(&spec, &ctx, OodSampleBudget::ZERO).expect("IRS fixture must solve"); + + let field_bits = TestField::field_size_bits(); + let log_list = commit.list_size().log2(); + let prox_gaps = commit.rbr_soundness_fold_prox_gaps(); + let poly_id = field_bits - log_list; + assert!( + prox_gaps < poly_id, + "fixture wants prox_gaps to bind: prox_gaps {prox_gaps} ≥ poly_id {poly_id}", + ); + + let got = f64::from(analytic_error_bits(&commit)); + assert_close(got, prox_gaps.max(0.0)); + } + + proptest! { + #[test] + fn solve_standard_assembles( + spec in arb_standard_spec(TEST_TARGET_RANGE), + (log_size, log_inv_rate) in arb_dims(), + ) { + let config = solve::(&spec, 1usize << log_size, log_inv_rate).unwrap(); + prop_assert!(matches!(config.mode(), basecase::BasecaseMode::Standard)); + prop_assert_eq!(config.commit().interleaving_depth(), 1); + prop_assert_eq!(config.commit().num_vectors(), 1); + prop_assert_eq!(config.commit().vector_size(), config.sumcheck().initial_size()); + } + + #[test] + fn solve_zk_assembles( + spec in arb_zk_spec(TEST_TARGET_RANGE), + (log_size, log_inv_rate) in arb_dims(), + ) { + let config = solve::(&spec, 1usize << log_size, log_inv_rate).unwrap(); + prop_assert!(matches!(config.mode(), basecase::BasecaseMode::ZeroKnowledge)); + prop_assert!(config.commit().mask_length() > 0); + } + + #[test] + fn pow_closes_gap_to_target_zk( + spec in arb_zk_spec(TEST_TARGET_RANGE), + (log_size, log_inv_rate) in arb_dims(), + ) { + let config = solve::(&spec, 1usize << log_size, log_inv_rate).unwrap(); + assert_pow_closes_gap(&spec, analytic_error_bits(config.commit()), &config.pow()); + } + + #[test] + fn standard_mode_has_no_pow( + spec in arb_standard_spec(TEST_TARGET_RANGE), + (log_size, log_inv_rate) in arb_dims(), + ) { + let config = solve::(&spec, 1usize << log_size, log_inv_rate).unwrap(); + prop_assert_eq!(config.pow(), PowConfig::none()); + } + } +} diff --git a/src/protocols/params/bounds.rs b/src/protocols/params/bounds.rs new file mode 100644 index 00000000..c8cad942 --- /dev/null +++ b/src/protocols/params/bounds.rs @@ -0,0 +1,40 @@ +//! Regime-agnostic analytic primitives shared across the params solvers. +//! +//! Regime-specific math (Unique / Johnson / Capacity branches) lives on +//! [`super::regime::DecodingRegimeParams`]. + +/// `ρ = 2^-log_inv_rate`. Centralized so the rate formula lives in one place. +pub(super) fn rate(log_inv_rate: f64) -> f64 { + 2_f64.powf(-log_inv_rate) +} + +/// Lossy `usize → f64` for analytic-error formulas. Named so individual call +/// sites can stay terse and intent-tagged. +pub(super) const fn usize_to_f64(x: usize) -> f64 { + x as f64 +} + +/// log2 of the per-OOD-sample Schwartz–Zippel error: `(k-1)/|F|`. +pub fn ood_per_sample_log2(message_length: usize, field_bits: f64) -> f64 { + ((message_length - 1) as f64).log2() - field_bits +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::params::test_utils::assert_close; + + /// OOD per-sample Schwartz–Zippel: `log₂((k−1) / |F|) = log₂(k−1) − field_bits`. + #[test] + fn ood_per_sample_log2_formula() { + // `k = 129` so `k − 1 = 128 = 2^7` for exact `log2`. + const K: usize = 129; + const FIELD_BITS: f64 = 64.0; + + let got = ood_per_sample_log2(K, FIELD_BITS); + let expected = ((K - 1) as f64).log2() - FIELD_BITS; + assert_close(got, expected); + // (k−1)/|F| < 1 for sane parameters ⇒ log is negative. + assert!(got < 0.0); + } +} diff --git a/src/protocols/params/branch.rs b/src/protocols/params/branch.rs new file mode 100644 index 00000000..7e2681ad --- /dev/null +++ b/src/protocols/params/branch.rs @@ -0,0 +1,65 @@ +//! Standard-vs-ZK branching used to thread mode through the build pipeline. +//! +//! [`Branch`] is the shared shape: a transient choice between the standard +//! path (no payload) and the zero-knowledge path (payload `T`). Concrete +//! pipeline stages alias it with the payload they carry: +//! +//! - [`RoundBuildMode`] — input to [`super::build_round::build_round_config`]. +//! - [`OodMode`] — input to OOD-bound helpers. +//! - [`SolveMode`] — input to the per-sub-protocol solvers (`sumcheck`, `code_switch`). +//! +//! Sharing one enum gives us a free [`Branch::map`] for stage-to-stage +//! payload conversions, replacing one-off `to_ood_mode`-style helpers. + +use crate::protocols::params::{ + protocol_config::MaskOracleInfo, + spec::{LogInvRate, ZkSpec}, +}; + +/// Standard (no payload) vs. zero-knowledge (payload `T`). +#[derive(Clone, Copy, Debug)] +pub enum Branch { + Standard, + ZeroKnowledge(T), +} + +impl Branch { + pub const fn is_zk(&self) -> bool { + matches!(self, Self::ZeroKnowledge(_)) + } + + /// Transform the ZK payload, leaving `Standard` unchanged. Replaces + /// per-stage `to_*` conversion helpers. + pub fn map(self, f: impl FnOnce(T) -> U) -> Branch { + match self { + Self::Standard => Branch::Standard, + Self::ZeroKnowledge(t) => Branch::ZeroKnowledge(f(t)), + } + } + + pub const fn as_ref(&self) -> Branch<&T> { + match self { + Self::Standard => Branch::Standard, + Self::ZeroKnowledge(t) => Branch::ZeroKnowledge(t), + } + } +} + +/// Payload carried by [`RoundBuildMode::ZeroKnowledge`] — references the +/// `SecuritySpec` (so its lifetime threads through) plus the planner-chosen +/// `C_zk` rate. +#[derive(Clone, Copy, Debug)] +pub struct RoundBuildPayload<'a> { + pub zk_spec: ZkSpec<'a>, + pub c_zk_log_inv_rate: LogInvRate, +} + +/// Mode-dispatch input for [`super::build_round::build_round_config`]. +pub type RoundBuildMode<'a> = Branch>; + +/// Mode flag for the OOD security bound. Payload is the `C_zk` log-inverse +/// rate; formulas coerce to `f64` at the point of use. +pub type OodMode = Branch; + +/// Solver-input mode for the per-round sumcheck and code-switch builders. +pub type SolveMode = Branch; diff --git a/src/protocols/params/build_round.rs b/src/protocols/params/build_round.rs new file mode 100644 index 00000000..e533815b --- /dev/null +++ b/src/protocols/params/build_round.rs @@ -0,0 +1,238 @@ +//! Per-round build: turns a [`RoundShape`] into a [`RoundConfig`]. +//! +//! Solves the `t_ood` fix-point, builds source/target IRS configs, and +//! (in ZK) assembles the per-round mask oracle. Consumed by +//! [`super::derive`], which drives the per-round loop. + +use crate::{ + algebra::{ + embedding::{Embedding, Identity}, + fields::FieldWithSize, + }, + protocols::{ + irs_commit::Config as IrsConfig, + mask_proximity::Config as MaskProximityConfig, + params::{ + bounds::usize_to_f64, + branch::{Branch, OodMode, RoundBuildMode, RoundBuildPayload, SolveMode}, + code_switch as code_switch_params, + error::{DeriveError, Pow}, + irs_commit as irs_params, + layout::{round_context, target_context, RoundShape}, + mask_proximity as mask_proximity_params, + protocol_config::{MaskOracleConfig, RoundConfig, RoundMode}, + spec::{ + DecodingRegime, LogInvRate, MaskCodeMessageLen, OodSampleBudget, RoundContext, + SecuritySpec, ZkSpec, + }, + sumcheck as sumcheck_params, + }, + }, +}; + +const T_OOD_MAX_ITER: usize = 32; + +pub(super) fn build_round_config( + spec: &SecuritySpec, + shape: &RoundShape, + mode: RoundBuildMode<'_>, +) -> Result, DeriveError> { + let ctx = round_context(shape); + let ood_mode = mode.map(|p| p.c_zk_log_inv_rate); + let (source, t_ood) = solve_round_source::(spec, shape, ood_mode)?; + + let (target_budget, solve_mode, round_mode) = match mode { + Branch::Standard => ( + OodSampleBudget::ZERO, + SolveMode::Standard, + RoundMode::Standard, + ), + Branch::ZeroKnowledge(RoundBuildPayload { + zk_spec, + c_zk_log_inv_rate, + }) => { + let num_masks = + sumcheck_params::masks_required(&ctx) + code_switch_params::masks_required(); + let mask_oracle = build_mask_oracle::( + zk_spec, + &source, + t_ood, + num_masks, + c_zk_log_inv_rate, + shape.round_index, + )?; + let solve_mode = SolveMode::ZeroKnowledge(mask_oracle.info()); + ( + OodSampleBudget::new(t_ood), + solve_mode, + RoundMode::ZeroKnowledge { + t_ood: OodSampleBudget::new(t_ood), + mask_oracle: Box::new(mask_oracle), + }, + ) + } + }; + + let target: IrsConfig> = + irs_params::solve(spec, &target_context(shape, &source), target_budget)?; + let sumcheck = sumcheck_params::solve( + spec, + &ctx, + &source, + solve_mode, + Pow::RoundSumcheck { + index: shape.round_index, + }, + )?; + let code_switch = + code_switch_params::solve(spec, source, target, t_ood, solve_mode, shape.round_index)?; + + Ok(RoundConfig::new( + shape.round_index, + sumcheck, + code_switch, + round_mode, + )) +} + +fn solve_round_source( + spec: &SecuritySpec, + shape: &RoundShape, + ood_mode: OodMode, +) -> Result<(IrsConfig, usize), DeriveError> { + let src_ctx = round_context(shape); + let target_log_inv_rate = f64::from( + shape + .source_log_inv_rate + .saturating_add(shape.source_folding_factor.saturating_sub(1)), + ); + let target_log_degree = f64::from( + shape + .source_vector_size + .trailing_zeros() + .saturating_sub(shape.source_folding_factor), + ); + let target_list_size = spec + .decoding_regime + .list_size_estimate(target_log_degree, target_log_inv_rate); + solve_t_ood::( + spec, + &src_ctx, + target_list_size, + ood_mode, + shape.round_index, + ) +} + +/// ZK-only: assemble the per-round mask oracle (C_zk codeword + mask-proximity +/// check). +fn build_mask_oracle( + zk_spec: ZkSpec<'_>, + source: &IrsConfig, + t_ood: usize, + num_masks: usize, + c_zk_log_inv_rate: LogInvRate, + round_index: usize, +) -> Result, DeriveError> { + let spec = zk_spec.as_inner(); + let l_zk = compute_l_zk(source, t_ood); + let c_zk: IrsConfig> = irs_params::solve_mask_code( + zk_spec, + l_zk, + source.mask_length(), + c_zk_log_inv_rate, + MaskProximityConfig::::num_vectors_for(num_masks), + )?; + let c_zk_list_size_estimate = spec.decoding_regime.list_size_estimate( + (l_zk.get() as f64).log2(), + f64::from(c_zk_log_inv_rate.get()), + ); + debug_assert!( + (c_zk.list_size() - c_zk_list_size_estimate).abs() + < 1e-9 * c_zk_list_size_estimate.max(1.0), + "c_zk.list_size() {} drifted from planner estimate {}", + c_zk.list_size(), + c_zk_list_size_estimate, + ); + let mask_proximity = mask_proximity_params::solve(spec, c_zk.clone(), num_masks, round_index)?; + Ok(MaskOracleConfig::new(c_zk, l_zk, mask_proximity)) +} + +/// `ℓ_zk = next_pow2(r + t_ood)` (Theorem 9.6 + Lemma 9.3). +pub(super) const fn compute_l_zk( + source: &IrsConfig, + t_ood: usize, +) -> MaskCodeMessageLen { + MaskCodeMessageLen::new( + source + .mask_length() + .saturating_add(t_ood) + .next_power_of_two(), + ) +} + +/// Per-round `(source, t_ood)`. +/// +/// Under `Unique`, `t_ood = 1` is pinned (the `log(L·(L−1)/2)` term degenerates +/// when `L = 1`, and Construction 9.7 requires `out_domain_samples ≥ 1`). +/// Otherwise linear search over `t_ood = 1..=T_OOD_MAX_ITER` for the smallest +/// value where [`ood_security_bits_at`] meets `protocol_security_target_bits`. +pub(super) fn solve_t_ood( + spec: &SecuritySpec, + src_ctx: &RoundContext, + target_list_size: f64, + ood_mode: OodMode, + round_index: usize, +) -> Result<(IrsConfig, usize), DeriveError> { + if matches!(spec.decoding_regime, DecodingRegime::Unique) { + let source = irs_params::solve(spec, src_ctx, OodSampleBudget::new(1))?; + return Ok((source, 1)); + } + + let security_target = f64::from(spec.protocol_security_target_bits()); + let field_bits = M::Target::field_size_bits(); + + for t_ood in 1..=T_OOD_MAX_ITER { + let source: IrsConfig = irs_params::solve(spec, src_ctx, OodSampleBudget::new(t_ood))?; + let bits = + ood_security_bits_at(spec, &source, t_ood, target_list_size, ood_mode, field_bits); + if bits >= security_target { + return Ok((source, t_ood)); + } + } + Err(DeriveError::FixedPointDidNotConverge { round_index }) +} + +/// OOD security bits at candidate `t_ood`, per STIR Lemma 4.5: +/// `bits = t · (|F| − log d) − log(L · (L − 1) / 2) ≈ t·(|F| − log d) − 2·log L + 1`. +fn ood_security_bits_at( + spec: &SecuritySpec, + source: &IrsConfig, + t_ood: usize, + target_list_size: f64, + ood_mode: OodMode, + field_bits: f64, +) -> f64 { + let (log_degree, log_combined_list) = match ood_mode { + Branch::Standard => ( + usize_to_f64(source.message_length()).log2(), + target_list_size.log2(), + ), + Branch::ZeroKnowledge(c_zk_log_inv_rate) => { + let l_zk = source + .mask_length() + .saturating_add(t_ood) + .next_power_of_two(); + let c_zk_list = spec.decoding_regime.list_size_estimate( + usize_to_f64(l_zk).log2(), + f64::from(c_zk_log_inv_rate.get()), + ); + ( + usize_to_f64(source.message_length().saturating_add(l_zk)).log2(), + (target_list_size * c_zk_list).log2(), + ) + } + }; + let ood = usize_to_f64(t_ood); + ood * (field_bits - log_degree) - 2.0 * log_combined_list + 1.0 +} diff --git a/src/protocols/params/code_switch.rs b/src/protocols/params/code_switch.rs new file mode 100644 index 00000000..cc116bd4 --- /dev/null +++ b/src/protocols/params/code_switch.rs @@ -0,0 +1,436 @@ +//! Code-switching IOR (Construction 9.7, p.55) builder + Lemma 9.9 OOD bound. + +use std::num::NonZeroUsize; + +use crate::{ + algebra::{ + embedding::{Embedding, Identity}, + fields::FieldWithSize, + }, + bits::Bits, + protocols::{ + code_switch::{self, Config as CodeSwitchConfig}, + irs_commit::Config as IrsConfig, + params::{ + bounds::usize_to_f64, + branch::SolveMode, + error::{grind_to_at, DeriveError, Pow}, + protocol_config::MaskOracleInfo, + solved::Solved, + spec::SecuritySpec, + }, + }, +}; + +/// Per-round code-switch builder. +pub fn solve( + spec: &SecuritySpec, + source: IrsConfig, + target: IrsConfig>, + t_ood: usize, + mode: SolveMode, + round_index: usize, +) -> Result>, DeriveError> { + let (mask_oracle, output_mode) = match mode { + SolveMode::Standard => (None, code_switch::CodeSwitchMode::Standard), + SolveMode::ZeroKnowledge(mask_oracle) => { + let l_zk = mask_oracle.l_zk.get(); + assert!( + l_zk >= source.mask_length().saturating_add(t_ood), + "ℓ_zk ({l_zk}) < r + t_ood ({} + {}) — violates Theorem 9.6 witness sizing", + source.mask_length(), + t_ood, + ); + ( + Some(mask_oracle), + code_switch::CodeSwitchMode::ZeroKnowledge { + message_mask_length: NonZeroUsize::new(l_zk).expect("ℓ_zk > 0"), + }, + ) + } + }; + + let analytic = analytic_error_bits(&source, &target, t_ood, mask_oracle); + let pow = grind_to_at(spec, analytic, Pow::RoundCodeSwitch { index: round_index })?; + + Ok(Solved::new( + CodeSwitchConfig::new(source, target, t_ood, output_mode, pow), + analytic, + )) +} + +/// Per-round code-switch soundness in bits: `min` over Lemma 9.9's three RBR +/// error slots (OOD, in-domain, combination). +pub fn analytic_error_bits( + source: &IrsConfig, + target: &IrsConfig>, + t_ood: usize, + mask_oracle: Option, +) -> Bits { + assert!(t_ood > 0, "code-switch requires t_ood ≥ 1"); + + let field_bits = M::Target::field_size_bits(); + let combined_list = + target.list_size() * mask_oracle.map_or(1.0, |info| info.c_zk_list_size.get()); + // OOD polynomial is over witness `[f; r_C; s]` of length `ℓ + ℓ_zk` (ZK) or + // `ℓ` (Standard). + let degree = mask_oracle.map_or_else( + || source.message_length(), + |info| source.message_length().saturating_add(info.l_zk.get()), + ); + let t_ood_f = usize_to_f64(t_ood); + + // OOD term — Lemma 9.9, term 1. + let log_degree_minus_1 = usize_to_f64(degree.saturating_sub(1)).log2(); + let log_l_choose_2 = (combined_list * (combined_list - 1.0) / 2.0).log2(); + let ood_term = t_ood_f * (field_bits - log_degree_minus_1) - log_l_choose_2; + + // In-domain term — Lemma 9.9, term 2. + let in_domain_term = source.rbr_queries(); + + // Combination term — Lemma 9.9, term 3 (γ-RLC, bounds doc §5.1). + let log_count = usize_to_f64( + t_ood.saturating_add(source.in_domain_samples() * source.interleaving_depth()), + ) + .log2(); + let combination_term = field_bits - log_count - combined_list.log2(); + + Bits::new(ood_term.min(in_domain_term).min(combination_term).max(0.0)) +} + +/// Number of `(r ‖ s)` mask polynomials code-switch contributes to C_zk per +/// round. +pub const fn masks_required() -> usize { + 1 +} + +#[cfg(test)] +mod tests { + use proptest::prelude::*; + + use super::*; + use crate::protocols::params::{ + branch::OodMode, + build_round::{compute_l_zk, solve_t_ood}, + irs_commit as irs_params, + spec::{ + ListSize, LogInvRate, MaskCodeMessageLen, Mode, OodSampleBudget, PowBudget, + RoundContext, SecuritySpec, ZkSpec, + }, + test_utils::{ + arb_standard_spec as utils_standard_spec, arb_zk_spec as utils_zk_spec, assert_close, + assert_pow_closes_gap, build_round_io, deterministic_spec, TestEmbedding, + TestExtensionField, TestField, TestNonIdentityEmbedding, TEST_TARGET_RANGE, + }, + }; + + type M = TestEmbedding; + + fn arb_zk_spec() -> impl Strategy { + utils_zk_spec(TEST_TARGET_RANGE) + } + + fn arb_standard_spec() -> impl Strategy { + utils_standard_spec(TEST_TARGET_RANGE) + } + + const NUM_VARS_HEADROOM: u32 = 4; + + fn arb_dims() -> impl Strategy { + (1u32..=3, 1u32..=2).prop_flat_map(|(log_inv_rate, folding_factor)| { + let min_num_vars = 2 * folding_factor; + ( + Just(log_inv_rate), + Just(folding_factor), + min_num_vars..=(min_num_vars + NUM_VARS_HEADROOM), + ) + }) + } + + const FORMULA_LOG_INV_RATE: u32 = 1; + const FORMULA_FOLDING_FACTOR: u32 = 2; + const FORMULA_NUM_VARS: u32 = 6; + + #[test] + fn analytic_error_standard_formula() { + let spec: SecuritySpec = deterministic_spec(Mode::Standard); + let (source, target, t_ood) = build_round_io::( + &spec, + FORMULA_LOG_INV_RATE, + FORMULA_FOLDING_FACTOR, + FORMULA_NUM_VARS, + None, + ); + let got = f64::from(analytic_error_bits(&source, &target, t_ood, None)); + + let field_bits = ::field_size_bits(); + let target_list = target.list_size(); + let degree = source.message_length(); + let log_deg_m1 = ((degree - 1) as f64).log2(); + let l_choose_2 = target_list * (target_list - 1.0) / 2.0; + let ood = (t_ood as f64) * (field_bits - log_deg_m1) - l_choose_2.log2(); + let in_domain = source.rbr_queries(); + let count = t_ood + source.in_domain_samples() * source.interleaving_depth(); + let comb = field_bits - (count as f64).log2() - target_list.log2(); + let expected = ood.min(in_domain).min(comb).max(0.0); + + assert_close(got, expected); + } + + #[test] + fn analytic_error_zk_formula() { + const C_ZK_LIST_SIZE: f64 = 4.0; + const L_ZK_USIZE: usize = 8; + + let spec: SecuritySpec = deterministic_spec(Mode::ZeroKnowledge); + let mask_oracle = MaskOracleInfo { + c_zk_list_size: ListSize::new(C_ZK_LIST_SIZE), + l_zk: MaskCodeMessageLen::new(L_ZK_USIZE), + }; + let (source, target, t_ood) = build_round_io::( + &spec, + FORMULA_LOG_INV_RATE, + FORMULA_FOLDING_FACTOR, + FORMULA_NUM_VARS, + Some(FORMULA_LOG_INV_RATE), + ); + let got = f64::from(analytic_error_bits( + &source, + &target, + t_ood, + Some(mask_oracle), + )); + + let field_bits = ::field_size_bits(); + let target_list = target.list_size(); + let combined_list = target_list * C_ZK_LIST_SIZE; + let degree = source.message_length() + L_ZK_USIZE; + let log_deg_m1 = ((degree - 1) as f64).log2(); + let l_choose_2 = combined_list * (combined_list - 1.0) / 2.0; + let ood = (t_ood as f64) * (field_bits - log_deg_m1) - l_choose_2.log2(); + let in_domain = source.rbr_queries(); + let count = t_ood + source.in_domain_samples() * source.interleaving_depth(); + let comb = field_bits - (count as f64).log2() - target_list.log2() - C_ZK_LIST_SIZE.log2(); + let expected = ood.min(in_domain).min(comb).max(0.0); + + assert_close(got, expected); + } + + #[test] + fn analytic_error_uses_in_domain_when_limiting() { + const LIMITING_TARGET_BITS: u32 = 16; + const LIMITING_LOG_INV_RATE: u32 = 1; + const LIMITING_FOLDING_FACTOR: u32 = 1; + const LIMITING_NUM_VARS: u32 = 4; + + let spec = SecuritySpec::new(LIMITING_TARGET_BITS).with_pow_budget(PowBudget::Forbidden); + let (source, target, t_ood) = build_round_io::( + &spec, + LIMITING_LOG_INV_RATE, + LIMITING_FOLDING_FACTOR, + LIMITING_NUM_VARS, + None, + ); + + let field_bits = ::field_size_bits(); + let target_list = target.list_size(); + let degree = source.message_length(); + let log_deg_m1 = ((degree - 1) as f64).log2(); + let l_choose_2 = target_list * (target_list - 1.0) / 2.0; + let ood = (t_ood as f64) * (field_bits - log_deg_m1) - l_choose_2.log2(); + let in_domain = source.rbr_queries(); + let count = t_ood + source.in_domain_samples() * source.interleaving_depth(); + let comb = field_bits - (count as f64).log2() - target_list.log2(); + assert!( + in_domain < ood && in_domain < comb, + "fixture wants in_domain to bind: in_domain {in_domain}, ood {ood}, comb {comb}", + ); + + let got = f64::from(analytic_error_bits(&source, &target, t_ood, None)); + assert_close(got, in_domain.max(0.0)); + } + + proptest! { + #[test] + fn solve_standard_assembles( + spec in arb_standard_spec(), + (log_inv_rate, folding_factor, num_vars) in arb_dims(), + ) { + let (source, target, t_ood) = + build_round_io::(&spec, log_inv_rate, folding_factor, num_vars, None); + let config = solve(&spec, source, target, t_ood, SolveMode::Standard, 0).unwrap(); + prop_assert!(matches!(config.mode(), code_switch::CodeSwitchMode::Standard)); + prop_assert!(config.out_domain_samples() >= 1); + } + + #[test] + fn solve_zk_mask_equals_padded_r_plus_t_ood( + spec in arb_zk_spec(), + (log_inv_rate, folding_factor, num_vars) in arb_dims(), + ) { + let (source, target, t_ood) = build_round_io::( + &spec, log_inv_rate, folding_factor, num_vars, Some(log_inv_rate), + ); + let r = source.mask_length(); + let l_zk = compute_l_zk(&source, t_ood); + let zk_spec = ZkSpec::try_new(&spec).expect("arb_zk_spec"); + let c_zk = irs_params::solve_mask_code::( + zk_spec, + l_zk, + r, + LogInvRate::new(log_inv_rate), + 2, + ) + .expect("C_zk fixture must solve"); + let mask_oracle = MaskOracleInfo { + c_zk_list_size: ListSize::new(c_zk.list_size()), + l_zk, + }; + let config = solve( + &spec, + source, + target, + t_ood, + SolveMode::ZeroKnowledge(mask_oracle), + 0, + ) + .unwrap(); + prop_assert_eq!(config.message_mask_length(), (r + t_ood).next_power_of_two()); + } + + #[test] + fn pow_closes_gap_to_target_standard( + spec in arb_standard_spec(), + (log_inv_rate, folding_factor, num_vars) in arb_dims(), + ) { + let (source, target, t_ood) = + build_round_io::(&spec, log_inv_rate, folding_factor, num_vars, None); + let error = analytic_error_bits(&source, &target, t_ood, None); + let config = solve(&spec, source, target, t_ood, SolveMode::Standard, 0).unwrap(); + assert_pow_closes_gap(&spec, error, &config.pow()); + } + } + + fn non_identity_smoke_ctxs() -> (RoundContext, RoundContext) { + const SOURCE_VECTOR_SIZE: usize = 64; + const SOURCE_LOG_INV_RATE: u32 = 1; + const FOLDING_FACTOR: u32 = 2; + + let source_ctx = RoundContext { + vector_size: SOURCE_VECTOR_SIZE, + log_inv_rate: SOURCE_LOG_INV_RATE, + folding_factor: FOLDING_FACTOR, + }; + let target_ctx = RoundContext { + vector_size: source_ctx.vector_size / (1 << source_ctx.folding_factor), + log_inv_rate: source_ctx.log_inv_rate + source_ctx.folding_factor - 1, + folding_factor: source_ctx.folding_factor, + }; + (source_ctx, target_ctx) + } + + #[test] + #[should_panic(expected = "violates Theorem 9.6")] + fn solve_zk_rejects_l_zk_below_r_plus_t_ood() { + const TOO_SMALL_L_ZK: usize = 1; + + let spec: SecuritySpec = deterministic_spec(Mode::ZeroKnowledge); + let (source, target, t_ood) = build_round_io::( + &spec, + FORMULA_LOG_INV_RATE, + FORMULA_FOLDING_FACTOR, + FORMULA_NUM_VARS, + Some(FORMULA_LOG_INV_RATE), + ); + assert!(source.mask_length() + t_ood > TOO_SMALL_L_ZK); + + let mask_oracle = MaskOracleInfo { + c_zk_list_size: ListSize::new(SMOKE_C_ZK_LIST_SIZE), + l_zk: MaskCodeMessageLen::new(TOO_SMALL_L_ZK), + }; + let _ = solve( + &spec, + source, + target, + t_ood, + SolveMode::ZeroKnowledge(mask_oracle), + 0, + ); + } + + #[test] + fn solve_works_with_basefield_embedding_standard() { + let spec: SecuritySpec = deterministic_spec(Mode::Standard); + let (source_ctx, target_ctx) = non_identity_smoke_ctxs(); + let target_log_degree = + f64::from((source_ctx.vector_size / (1 << source_ctx.folding_factor)).trailing_zeros()); + let target_list_size = spec + .decoding_regime + .list_size_estimate(target_log_degree, f64::from(target_ctx.log_inv_rate)); + let (source, t_ood) = solve_t_ood::( + &spec, + &source_ctx, + target_list_size, + OodMode::Standard, + 0, + ) + .unwrap(); + let target = irs_params::solve::>( + &spec, + &target_ctx, + OodSampleBudget::ZERO, + ) + .expect("target IRS fixture must solve"); + + let config = solve(&spec, source, target, t_ood, SolveMode::Standard, 0).unwrap(); + assert!(matches!( + config.mode(), + code_switch::CodeSwitchMode::Standard + )); + } + + const SMOKE_C_ZK_LIST_SIZE: f64 = 4.0; + + #[test] + fn solve_works_with_basefield_embedding_zk() { + let spec: SecuritySpec = deterministic_spec(Mode::ZeroKnowledge); + let (source_ctx, target_ctx) = non_identity_smoke_ctxs(); + let target_log_degree = + f64::from((source_ctx.vector_size / (1 << source_ctx.folding_factor)).trailing_zeros()); + let target_list_size = spec + .decoding_regime + .list_size_estimate(target_log_degree, f64::from(target_ctx.log_inv_rate)); + let (source, t_ood) = solve_t_ood::( + &spec, + &source_ctx, + target_list_size, + OodMode::ZeroKnowledge(LogInvRate::new(source_ctx.log_inv_rate)), + 0, + ) + .unwrap(); + let target = irs_params::solve::>( + &spec, + &target_ctx, + OodSampleBudget::new(t_ood), + ) + .expect("target IRS fixture must solve"); + + let mask_oracle = MaskOracleInfo { + c_zk_list_size: ListSize::new(SMOKE_C_ZK_LIST_SIZE), + l_zk: MaskCodeMessageLen::new((source.mask_length() + t_ood).next_power_of_two()), + }; + let config = solve( + &spec, + source, + target, + t_ood, + SolveMode::ZeroKnowledge(mask_oracle), + 0, + ) + .unwrap(); + assert!(matches!( + config.mode(), + code_switch::CodeSwitchMode::ZeroKnowledge { .. } + )); + } +} diff --git a/src/protocols/params/derive.rs b/src/protocols/params/derive.rs new file mode 100644 index 00000000..ca5cdd30 --- /dev/null +++ b/src/protocols/params/derive.rs @@ -0,0 +1,653 @@ +//! Derives a [`ProtocolConfig`] from a spec + tuning. + +use crate::{ + algebra::embedding::Embedding, + protocols::params::{ + basecase as basecase_params, + branch::{Branch, RoundBuildMode, RoundBuildPayload}, + build_round::build_round_config, + error::DeriveError, + layout::{round_layout, RoundLayout}, + protocol_config::{ProtocolConfig, RoundConfig}, + spec::{LogInvRate, SecuritySpec, TuningSpec}, + }, +}; + +impl ProtocolConfig { + /// Fails with [`DeriveError`] when the spec/tuning combination is + /// infeasible. + pub fn derive(spec: SecuritySpec, tuning: TuningSpec) -> Result { + let RoundLayout { + shapes, + basecase_vector_size, + basecase_log_inv_rate, + } = round_layout(&tuning)?; + + let mode: RoundBuildMode<'_> = spec.as_zk().map_or(Branch::Standard, |zk_spec| { + Branch::ZeroKnowledge(RoundBuildPayload { + zk_spec, + c_zk_log_inv_rate: LogInvRate::new(tuning.starting_log_inv_rate), + }) + }); + + let rounds: Vec> = shapes + .iter() + .map(|shape| build_round_config::(&spec, shape, mode)) + .collect::>()?; + + let basecase = basecase_params::solve(&spec, basecase_vector_size, basecase_log_inv_rate)?; + + let plan = Self::new(spec, tuning, rounds, basecase); + plan.validate()?; + Ok(plan) + } +} + +#[cfg(test)] +mod tests { + use proptest::prelude::*; + + use crate::{ + algebra::{ + embedding::Embedding, + fields::{Field64, FieldWithSize}, + }, + protocols::{ + basecase::BasecaseMode, + params::{ + error::{ChainSource, ChainTarget, DeriveError, Pow}, + protocol_config::{ProtocolConfig, RoundMode}, + spec::{DecodingRegime, FoldingFactor, Mode, PowBudget, SecuritySpec, TuningSpec}, + test_utils::{assert_close, assert_pow_closes_gap, TestEmbedding}, + }, + }, + }; + + fn arb_tuning() -> impl Strategy { + let folding = prop_oneof![ + (1usize..=3).prop_map(FoldingFactor::Constant), + (1usize..=3, 1usize..=3).prop_map(|(initial, rest)| { + FoldingFactor::ConstantFromSecondRound { initial, rest } + }), + ]; + (4u32..=8, 1u32..=3, folding).prop_map(|(log_size, log_inv_rate, folding_factor)| { + TuningSpec { + vector_size: 1usize << log_size, + starting_log_inv_rate: log_inv_rate, + folding_factor, + } + }) + } + + const FIXTURE_FOLDING_FACTOR: usize = 2; + const FIXTURE_LOG_INV_RATE: u32 = 1; + + const LOG_VECTOR_SIZE_NO_ROUNDS: u32 = 3; + const LOG_VECTOR_SIZE_MULTI_ROUND: u32 = 8; + + fn tuning_with(vector_size: usize) -> TuningSpec { + TuningSpec { + vector_size, + starting_log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: FoldingFactor::Constant(FIXTURE_FOLDING_FACTOR), + } + } + + const PLAN_FIXTURE_TARGET_BITS: u32 = 40; + + fn test_spec(mode: Mode) -> SecuritySpec { + SecuritySpec::new(PLAN_FIXTURE_TARGET_BITS) + .with_mode(mode) + .with_pow_budget(PowBudget::per_slot(LOOSE_POW_BUDGET_BITS)) + } + + #[test] + fn derive_standard_with_no_rounds_uses_basecase_only() { + let spec = test_spec(Mode::Standard); + let vector_size = 1usize << LOG_VECTOR_SIZE_NO_ROUNDS; + let plan = ProtocolConfig::::derive(spec, tuning_with(vector_size)).unwrap(); + assert!(plan.rounds().is_empty()); + assert_eq!(plan.basecase().commit().vector_size(), vector_size); + } + + #[test] + fn derive_zk_with_no_rounds_uses_zk_basecase_only() { + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_NO_ROUNDS), + ) + .unwrap(); + assert!(plan.rounds().is_empty()); + assert!(matches!( + plan.basecase().mode(), + BasecaseMode::ZeroKnowledge + )); + } + + #[test] + fn t_ood_nonzero_in_johnson_zk() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Johnson, + ..test_spec(Mode::ZeroKnowledge) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + for r in plan.rounds() { + let RoundMode::ZeroKnowledge { t_ood, .. } = r.mode() else { + panic!("expected ZK round") + }; + assert!(t_ood.get() >= 1); + } + } + + #[test] + fn t_ood_pinned_to_one_in_unique_zk() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Unique, + ..test_spec(Mode::ZeroKnowledge) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + for r in plan.rounds() { + let RoundMode::ZeroKnowledge { t_ood, .. } = r.mode() else { + panic!("expected ZK round") + }; + assert_eq!(t_ood.get(), 1); + } + } + + #[test] + fn c_zk_keeps_code_switch_mask_under_unique() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Unique, + ..test_spec(Mode::ZeroKnowledge) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + for r in plan.rounds() { + let mask_oracle = r.mask_oracle().expect("ZK round has a mask oracle"); + let k = r + .code_switch() + .source() + .interleaving_depth() + .trailing_zeros() as usize; + let expected_num_masks = k + 1; + assert_eq!(mask_oracle.c_zk().num_vectors(), 2 * expected_num_masks); + } + } + + #[test] + fn analytic_bits_finite_and_positive_standard() { + let spec = test_spec(Mode::Standard); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + let bits: f64 = plan.analytic_bits().into(); + assert!(bits.is_finite() && bits > 0.0, "bits = {bits}"); + let min_round = plan + .rounds() + .iter() + .map(|r| f64::from(r.analytic_bits())) + .fold(f64::INFINITY, f64::min); + let expected = min_round.min(f64::from(plan.basecase().analytic_bits())); + assert_close(bits, expected); + } + + #[test] + fn analytic_bits_includes_mask_oracle_in_zk() { + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + let plan_bits: f64 = plan.analytic_bits().into(); + let mo_floor = plan + .rounds() + .iter() + .filter_map(|r| r.mask_oracle().map(|mo| f64::from(mo.analytic_bits()))) + .fold(f64::INFINITY, f64::min); + assert!( + mo_floor.is_finite(), + "ZK plan must contribute mask-oracle bits" + ); + let min_round = plan + .rounds() + .iter() + .map(|r| f64::from(r.analytic_bits())) + .fold(f64::INFINITY, f64::min); + let expected = mo_floor + .min(min_round) + .min(f64::from(plan.basecase().analytic_bits())); + assert_close(plan_bits, expected); + } + + #[test] + fn derive_plans_basecase() { + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(matches!( + plan.basecase().mode(), + BasecaseMode::ZeroKnowledge + )); + assert_eq!(plan.basecase().commit().interleaving_depth(), 1); + assert_eq!(plan.basecase().sumcheck().final_size(), 1); + } + + const LOOSE_POW_BUDGET_BITS: u32 = 60; + const OVER_BUDGET_INJECTED_BITS: f64 = 50.0; + + /// Bounds doc §5.3 + §5.7: HVZK privacy error in bits matches the closed + /// form `−log Σ_r (t_ood_r² + t_ood_r) / (2|F|)` over ZK rounds. + #[test] + fn privacy_error_bits_matches_bound_3_sum() { + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + let field_bits = ::field_size_bits(); + let mut expected_total = 0.0_f64; + for r in plan.rounds() { + let RoundMode::ZeroKnowledge { t_ood, .. } = r.mode() else { + panic!("expected ZK round"); + }; + let t = t_ood.get() as f64; + expected_total += 2_f64.powf(f64::midpoint(t * t, t).log2() - field_bits); + } + let expected_bits = -expected_total.log2(); + let got = f64::from(plan.privacy_error_bits()); + assert_close(got, expected_bits); + } + + #[test] + fn privacy_error_bits_standard_returns_target_sentinel() { + let spec = test_spec(Mode::Standard); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert_close( + f64::from(plan.privacy_error_bits()), + f64::from(PLAN_FIXTURE_TARGET_BITS), + ); + } + + #[test] + fn check_pow_bits_passes_on_derived_plan() { + let plan = ProtocolConfig::::derive( + test_spec(Mode::ZeroKnowledge), + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(plan.check_pow_bits()); + } + + #[test] + fn check_pow_bits_detects_over_budget_slot() { + use crate::{bits::Bits, protocols::proof_of_work::Config as PowConfig}; + const MODERATE_POW_BUDGET_BITS: u32 = 30; + let spec = SecuritySpec { + pow_budget: PowBudget::per_slot(MODERATE_POW_BUDGET_BITS), + ..test_spec(Mode::ZeroKnowledge) + }; + let mut plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + plan.override_basecase_pow_for_test(PowConfig::from_difficulty(Bits::new( + OVER_BUDGET_INJECTED_BITS, + ))); + assert!(!plan.check_pow_bits()); + } + + #[test] + fn validate_round_chaining_detects_adjacent_round_mismatch() { + let spec = test_spec(Mode::ZeroKnowledge); + let mut plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + let n = plan.rounds().len(); + assert!(n >= 2, "need ≥ 2 rounds to break a mid-chain link"); + assert!(plan.check_all_invariants(), "fresh plan must validate"); + + let bad_size = plan.rounds()[0].code_switch().target().vector_size() + 1; + plan.corrupt_round_target_vector_size_for_test(0, bad_size); + + let err = plan + .validate_round_chaining() + .expect_err("adjacent-round mismatch must trip the chain check"); + assert!( + matches!( + err, + DeriveError::RoundChainBroken { + from: ChainSource::Round(0), + to: ChainTarget::NextRound(1), + .. + } + ), + "got {err:?}", + ); + assert!(!plan.check_all_invariants()); + } + + #[test] + fn validate_round_chaining_detects_basecase_mismatch() { + let spec = test_spec(Mode::ZeroKnowledge); + let mut plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + let n = plan.rounds().len(); + assert!(n >= 2, "need ≥ 2 rounds to break the chain by truncation"); + assert!(plan.check_all_invariants(), "fresh plan must validate"); + + plan.truncate_rounds_for_test(n - 1); + let err = plan + .validate_round_chaining() + .expect_err("truncated tail breaks basecase chaining"); + assert!( + matches!( + err, + DeriveError::RoundChainBroken { + to: ChainTarget::Basecase, + .. + } + ), + "got {err:?}", + ); + assert!(!plan.check_all_invariants()); + } + + #[test] + fn validate_security_target_met_passes_on_fresh_plan() { + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + plan.validate_security_target_met() + .expect("fresh plan must satisfy per-slot target check"); + } + + #[test] + fn validate_security_target_met_catches_recorded_analytic_drift() { + use crate::bits::Bits; + let spec = test_spec(Mode::ZeroKnowledge); + let mut plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(!plan.rounds().is_empty(), "need a round to corrupt"); + let recorded = plan + .rounds() + .first() + .map(|r| r.sumcheck().analytic()) + .expect("params solver records sumcheck analytic"); + // Bump the recorded value far from the recompute → triggers drift. + plan.corrupt_round_sumcheck_analytic_for_test(0, Bits::new(f64::from(recorded) + 10.0)); + let err = plan + .validate_security_target_met() + .expect_err("recorded vs recompute mismatch must trip drift check"); + assert!( + matches!( + err, + DeriveError::AnalyticDrift { + pow: Pow::RoundSumcheck { index: 0 }, + .. + } + ), + "got {err:?}", + ); + } + + #[test] + fn derive_reports_pow_ungrindable() { + const UNREACHABLE_TARGET_BITS: u32 = 200; + let spec = SecuritySpec { + target_security_bits: UNREACHABLE_TARGET_BITS, + ..test_spec(Mode::Standard) + }; + let err = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .expect_err("target above grind cap must fail"); + assert!( + matches!(err, DeriveError::PowUngrindable { .. }), + "got {err:?}", + ); + } + + #[test] + fn derive_reports_pow_budget_exceeded() { + const TIGHT_MAX_POW: u32 = 5; + let spec = SecuritySpec { + pow_budget: PowBudget::per_slot(TIGHT_MAX_POW), + ..test_spec(Mode::ZeroKnowledge) + }; + let err = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .expect_err("tight pow_budget must trip auto-validation"); + assert!( + matches!(err, DeriveError::PowBudgetExceeded { .. }), + "got {err:?}", + ); + } + + #[test] + fn derive_threads_unique_decoding_standard() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Unique, + ..test_spec(Mode::Standard) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_NO_ROUNDS), + ) + .unwrap(); + assert!(plan.rounds().is_empty()); + assert!(plan.basecase().commit().unique_decoding()); + } + + #[test] + fn derive_threads_unique_decoding_zk() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Unique, + ..test_spec(Mode::ZeroKnowledge) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_NO_ROUNDS), + ) + .unwrap(); + assert!(plan.rounds().is_empty()); + assert!(plan.basecase().commit().unique_decoding()); + } + + #[test] + fn derive_multi_round_unique_decoding_succeeds() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Unique, + ..test_spec(Mode::Standard) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(!plan.rounds().is_empty(), "expected multi-round plan"); + for r in plan.rounds() { + let cs = r.code_switch(); + assert!(cs.source().unique_decoding()); + assert!(cs.target().unique_decoding()); + assert!(cs.out_domain_samples() >= 1); + } + assert!(plan.basecase().commit().unique_decoding()); + } + + #[test] + fn derive_multi_round_unique_decoding_zk_succeeds() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Unique, + ..test_spec(Mode::ZeroKnowledge) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(!plan.rounds().is_empty(), "expected multi-round plan"); + for r in plan.rounds() { + let mo = r.mask_oracle().expect("ZK round must own a mask oracle"); + assert!(mo.c_zk().unique_decoding()); + assert!(r.code_switch().source().unique_decoding()); + assert!(r.code_switch().out_domain_samples() >= 1); + } + assert!(plan.basecase().commit().unique_decoding()); + } + + #[test] + fn derive_multi_round_capacity_decoding_succeeds() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Capacity, + ..test_spec(Mode::Standard) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(!plan.rounds().is_empty(), "expected multi-round plan"); + for r in plan.rounds() { + assert!(r.code_switch().out_domain_samples() >= 1); + } + } + + #[test] + fn derive_multi_round_capacity_decoding_zk_succeeds() { + let spec = SecuritySpec { + decoding_regime: DecodingRegime::Capacity, + ..test_spec(Mode::ZeroKnowledge) + }; + let plan = ProtocolConfig::::derive( + spec, + tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND), + ) + .unwrap(); + assert!(!plan.rounds().is_empty(), "expected multi-round plan"); + for r in plan.rounds() { + r.mask_oracle().expect("ZK round must own a mask oracle"); + assert!(r.code_switch().out_domain_samples() >= 1); + } + } + + /// Every slot from [`ProtocolConfig::pow_slots`] must close the gap to the + /// target with its configured grind, judged from a fresh recompute. + fn assert_plan_meets_target_per_slot( + spec: &SecuritySpec, + plan: &ProtocolConfig, + ) { + for slot in plan.pow_slots() { + assert_pow_closes_gap(spec, slot.recompute, &slot.pow); + } + } + + proptest! { + #[test] + fn derived_plan_meets_target_per_slot_standard(tuning in arb_tuning()) { + let spec = test_spec(Mode::Standard); + let plan = ProtocolConfig::::derive(spec.clone(), tuning).unwrap(); + assert_plan_meets_target_per_slot(&spec, &plan); + } + + #[test] + fn derived_plan_meets_target_per_slot_zk(tuning in arb_tuning()) { + let log_threshold = + tuning.folding_factor.at_round(0) + tuning.folding_factor.at_round(1); + prop_assume!(tuning.vector_size.trailing_zeros() as usize >= log_threshold); + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive(spec.clone(), tuning).unwrap(); + assert_plan_meets_target_per_slot(&spec, &plan); + } + + #[test] + fn derive_standard_succeeds_over_tunings(tuning in arb_tuning()) { + let spec = test_spec(Mode::Standard); + let plan = ProtocolConfig::::derive(spec, tuning).unwrap(); + for r in plan.rounds() { + prop_assert!(matches!(r.mode(), RoundMode::Standard)); + prop_assert!(r.mask_oracle().is_none()); + } + prop_assert!(matches!( + plan.basecase().mode(), + BasecaseMode::Standard + )); + prop_assert_eq!(plan.basecase().commit().interleaving_depth(), 1); + } + + #[test] + fn derive_zk_succeeds_over_tunings(tuning in arb_tuning()) { + let log_threshold = + tuning.folding_factor.at_round(0) + tuning.folding_factor.at_round(1); + prop_assume!(tuning.vector_size.trailing_zeros() as usize >= log_threshold); + + let spec = test_spec(Mode::ZeroKnowledge); + let plan = ProtocolConfig::::derive(spec, tuning).unwrap(); + for r in plan.rounds() { + let mask_oracle = r + .mask_oracle() + .expect("ZK round must have a mask oracle"); + let RoundMode::ZeroKnowledge { t_ood, .. } = r.mode() else { + panic!("expected ZK round"); + }; + let cs = r.code_switch(); + let k = cs.source().interleaving_depth().trailing_zeros() as usize; + let num_masks = k + 1; + prop_assert_eq!(mask_oracle.c_zk().num_vectors(), 2 * num_masks); + prop_assert_eq!(mask_oracle.mask_proximity().num_masks(), num_masks); + let source_mask = cs.source().mask_length(); + prop_assert!(mask_oracle.l_zk().get() >= source_mask + t_ood.get()); + } + prop_assert!(matches!( + plan.basecase().mode(), + BasecaseMode::ZeroKnowledge + )); + } + + #[test] + fn analytic_bits_finite_and_non_negative_standard(tuning in arb_tuning()) { + let spec = test_spec(Mode::Standard); + let plan = ProtocolConfig::::derive(spec, tuning).unwrap(); + let analytic = f64::from(plan.analytic_bits()); + prop_assert!(analytic.is_finite()); + prop_assert!(analytic >= 0.0); + } + } +} diff --git a/src/protocols/params/error.rs b/src/protocols/params/error.rs new file mode 100644 index 00000000..121a2e85 --- /dev/null +++ b/src/protocols/params/error.rs @@ -0,0 +1,196 @@ +//! Errors raised by [`super::derive::ProtocolConfig::derive`] and the +//! sub-protocol solvers. + +use std::fmt::{self, Display, Formatter}; + +use thiserror::Error; + +use crate::{ + bits::Bits, + protocols::{ + irs_commit::CodewordLengthError, + params::{layout::LayoutError, spec::SecuritySpec}, + proof_of_work::{Config as PowConfig, PowError}, + }, +}; + +/// Identifies a single PoW grind in the derived protocol. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Pow { + /// Basecase γ-RLC grind (Lemma 7.4) — ZK mode only. + BasecaseGammaCombination, + /// Basecase sumcheck grind. + BasecaseSumcheck, + /// Per-round sumcheck grind at `index`. + RoundSumcheck { index: usize }, + /// Per-round code-switch grind at `index`. + RoundCodeSwitch { index: usize }, + /// Per-round mask-proximity grind at `index` — ZK mode only. + RoundMaskProximity { index: usize }, +} + +impl Display for Pow { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::BasecaseGammaCombination => f.write_str("basecase γ-combination"), + Self::BasecaseSumcheck => f.write_str("basecase sumcheck"), + Self::RoundSumcheck { index } => write!(f, "round {index} sumcheck"), + Self::RoundCodeSwitch { index } => write!(f, "round {index} code-switch"), + Self::RoundMaskProximity { index } => write!(f, "round {index} mask-proximity"), + } + } +} + +/// Origin side of a [`DeriveError::RoundChainBroken`]: either a numbered round +/// or the pre-round `tuning` shape (for plans with no rounds at all). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChainSource { + Tuning, + Round(usize), +} + +impl Display for ChainSource { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Tuning => f.write_str("tuning"), + Self::Round(i) => write!(f, "round {i}"), + } + } +} + +/// Destination side of a [`DeriveError::RoundChainBroken`]: the next round in +/// sequence, or the basecase. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChainTarget { + NextRound(usize), + Basecase, +} + +impl Display for ChainTarget { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::NextRound(i) => write!(f, "round {i}"), + Self::Basecase => f.write_str("basecase"), + } + } +} + +/// Failure modes for [`super::derive::ProtocolConfig::derive`] and the +/// sub-protocol solvers it calls. +#[derive(Debug, Error, Clone, PartialEq, Eq)] +pub enum DeriveError { + /// The `t_ood` fixed-point in [`super::build_round::solve_t_ood`] ran out of + /// iterations — usually the field is too small for the security target. + #[error( + "t_ood fixed-point did not converge for round {round_index}; \ + lower target_security_bits or use a larger field" + )] + FixedPointDidNotConverge { round_index: usize }, + + /// A PoW grind cannot close the analytic-to-target gap — the spec is too + /// tight for any single grind to reach `target_security_bits`. + #[error( + "{pow} cannot be ground: {source}; lower target_security_bits or \ + switch to a less conservative decoding regime (Johnson/Capacity)" + )] + PowUngrindable { + pow: Pow, + #[source] + source: PowError, + }, + + /// A PoW grind fits the grind cap but exceeds the per-slot budget set by + /// [`super::spec::SecuritySpec::pow_budget`]. + #[error( + "{pow} requires {required} bits, exceeds spec.pow_budget = {max}; \ + raise pow_budget or lower target_security_bits" + )] + PowBudgetExceeded { pow: Pow, required: Bits, max: Bits }, + + /// Computed codeword length exceeds the NTT engine's supported order. + #[error( + "codeword length {length} exceeds the NTT engine's supported order; \ + reduce vector_size or starting_log_inv_rate" + )] + CodewordExceedsNtt { length: usize }, + + /// The tuning spec cannot produce a valid round layout. + #[error(transparent)] + Layout(#[from] LayoutError), + + /// Cross-round (or round → basecase) shape chain broken: the next + /// component's source `vector_size` does not match the previous + /// component's target `vector_size`. Surfaced by + /// [`super::protocol_config::ProtocolConfig::validate_round_chaining`]. + #[error("chain broken: {from} → {to} expected vector_size {expected}, found {found}")] + RoundChainBroken { + from: ChainSource, + to: ChainTarget, + expected: usize, + found: usize, + }, + + /// A PoW slot's `analytic + pow.difficulty()` is below `target_security_bits`. + /// `grind_to_at` guarantees this at construction; this error fires only if + /// the analytic-error formulas applied at validate-time disagree with what + /// the per-protocol `solve` functions consumed (e.g., a planner regression + /// drifts away from the actual configured IRS rate). Catches the case where + /// the rate-schedule plumbing under-reports a per-slot rate. + #[error("{pow} soundness gap: analytic {analytic} + pow {pow_bits} < target {target}")] + SecurityTargetNotMet { + pow: Pow, + analytic: Bits, + pow_bits: Bits, + target: Bits, + }, + + /// The analytic floor recorded at solve time disagrees with a fresh + /// recompute from the same config's state. Indicates that the inputs to + /// the `analytic_error_bits` formula drifted between solve and validate + /// (e.g., an IRS field was overwritten after construction). + #[error("{pow} analytic drift: recorded {recorded} vs recompute {recompute}")] + AnalyticDrift { + pow: Pow, + recorded: Bits, + recompute: Bits, + }, +} + +impl From for DeriveError { + fn from(e: CodewordLengthError) -> Self { + Self::CodewordExceedsNtt { length: e.length } + } +} + +/// Lift `Result` into `Result` by attaching a +/// [`Pow`] label. +pub trait PowResultExt { + fn at(self, pow: Pow) -> Result; +} + +impl PowResultExt for Result { + fn at(self, pow: Pow) -> Result { + self.map_err(|source| DeriveError::PowUngrindable { pow, source }) + } +} + +/// Grind `analytic → spec.target_security_bits`, then check the result against +/// `spec.pow_budget`. +pub fn grind_to_at( + spec: &SecuritySpec, + analytic: Bits, + pow_kind: Pow, +) -> Result { + let target = Bits::new(f64::from(spec.target_security_bits)); + let pow = PowConfig::grind_to(target, analytic, spec.hash_id).at(pow_kind)?; + let required = pow.difficulty(); + let max = Bits::new(f64::from(spec.pow_budget.bits())); + if required > max { + return Err(DeriveError::PowBudgetExceeded { + pow: pow_kind, + required, + max, + }); + } + Ok(pow) +} diff --git a/src/protocols/params/irs_commit.rs b/src/protocols/params/irs_commit.rs new file mode 100644 index 00000000..0d076cd6 --- /dev/null +++ b/src/protocols/params/irs_commit.rs @@ -0,0 +1,203 @@ +//! IRS-commit parameter selection. +//! +//! ZK mask sizing follows Construction 9.7 / Theorem 9.6: +//! `mask_length = in_domain + t_ood` (Proposition 3.19). + +use crate::{ + algebra::embedding::Embedding, + protocols::{ + irs_commit::{num_in_domain_queries, Config as IrsConfig, IrsMode, IrsParams}, + params::{ + bounds::rate, + error::DeriveError, + spec::{ + LogInvRate, MaskCodeMessageLen, Mode, OodSampleBudget, RoundContext, SecuritySpec, + ZkSpec, + }, + }, + }, +}; + +pub fn solve( + spec: &SecuritySpec, + ctx: &RoundContext, + out_domain_samples: OodSampleBudget, +) -> Result, DeriveError> { + let security_target = f64::from(spec.protocol_security_target_bits()); + let rate = rate(f64::from(ctx.log_inv_rate)); + let interleaving_depth = 1_usize << ctx.folding_factor; + + let mode = match spec.mode { + Mode::Standard => IrsMode::Standard, + Mode::ZeroKnowledge => { + let mask_length = num_in_domain_queries(spec.decoding_regime, security_target, rate) + .saturating_add(out_domain_samples.get()); + IrsMode::ZeroKnowledge { mask_length } + } + }; + + Ok(IrsConfig::try_new(IrsParams { + security_target, + decoding_regime: spec.decoding_regime, + hash_id: spec.hash_id, + num_vectors: 1, + vector_size: ctx.vector_size, + interleaving_depth, + rate, + mode, + })?) +} + +/// Shared C_zk IRS config for mask polynomials. +/// +/// - `l_zk`: message length, must be a power of 2. +/// - `source_mask_length`: `r` from Theorem 9.6. +/// - `num_vectors`: `2 * num_masks` (Construction 7.2: originals + fresh). +pub fn solve_mask_code( + spec: ZkSpec<'_>, + l_zk: MaskCodeMessageLen, + source_mask_length: usize, + log_inv_rate: LogInvRate, + num_vectors: usize, +) -> Result, DeriveError> { + let l_zk = l_zk.get(); + assert!( + l_zk >= source_mask_length, + "Theorem 9.6: ℓ_zk ({l_zk}) ≥ source mask length ({source_mask_length})", + ); + assert!(l_zk.is_power_of_two(), "ℓ_zk ({l_zk}) must be a power of 2"); + assert!( + num_vectors.is_multiple_of(2), + "num_vectors ({num_vectors}) must be even (mask-proximity original/fresh pairs)", + ); + + let security_target = f64::from(spec.protocol_security_target_bits()); + let rate = rate(f64::from(log_inv_rate.get())); + + Ok(IrsConfig::try_new(IrsParams { + security_target, + decoding_regime: spec.decoding_regime, + hash_id: spec.hash_id, + num_vectors, + vector_size: l_zk, + interleaving_depth: 1, + rate, + mode: IrsMode::Standard, + })?) +} + +#[cfg(test)] +mod tests { + use proptest::prelude::*; + + use super::*; + use crate::protocols::params::test_utils::{ + arb_round_ctx, arb_spec, arb_zk_spec, deterministic_spec, TestEmbedding, + TestNonIdentityEmbedding, + }; + + type M = TestEmbedding; + + #[test] + fn zk_spec_rejects_standard_mode() { + let spec: SecuritySpec = deterministic_spec(Mode::Standard); + assert!(ZkSpec::try_new(&spec).is_none()); + } + + #[test] + #[should_panic(expected = "must be a power of 2")] + fn solve_mask_code_rejects_non_pow2_l_zk() { + let spec: SecuritySpec = deterministic_spec(Mode::ZeroKnowledge); + let zk_spec = ZkSpec::try_new(&spec).unwrap(); + let _ = solve_mask_code::( + zk_spec, + MaskCodeMessageLen::new(3), + 0, + LogInvRate::new(1), + 2, + ); + } + + #[test] + #[should_panic(expected = "Theorem 9.6")] + fn solve_mask_code_rejects_l_zk_below_source_mask_length() { + let spec: SecuritySpec = deterministic_spec(Mode::ZeroKnowledge); + let zk_spec = ZkSpec::try_new(&spec).unwrap(); + let _ = solve_mask_code::( + zk_spec, + MaskCodeMessageLen::new(2), + 4, + LogInvRate::new(1), + 2, + ); + } + + #[test] + #[should_panic(expected = "must be even")] + fn solve_mask_code_rejects_odd_num_vectors() { + let spec: SecuritySpec = deterministic_spec(Mode::ZeroKnowledge); + let zk_spec = ZkSpec::try_new(&spec).unwrap(); + let _ = solve_mask_code::( + zk_spec, + MaskCodeMessageLen::new(2), + 0, + LogInvRate::new(1), + 3, + ); + } + + const IRS_TARGET_RANGE: std::ops::RangeInclusive = 80..=128; + + fn arb_zk_spec_default() -> impl Strategy { + arb_zk_spec(IRS_TARGET_RANGE) + } + + fn arb_standard_spec() -> impl Strategy { + arb_spec(Mode::Standard, IRS_TARGET_RANGE) + } + + proptest! { + #[test] + fn zk_mask_covers_in_domain_plus_ood( + spec in arb_zk_spec_default(), + ctx in arb_round_ctx(), + out_domain in 0usize..16, + ) { + let config = solve::(&spec, &ctx, OodSampleBudget::new(out_domain)).unwrap(); + prop_assert!( + config.mask_length() >= config.in_domain_samples() + out_domain, + "mask {} < in_domain {} + out_domain {}", + config.mask_length(), config.in_domain_samples(), out_domain, + ); + } + + #[test] + fn standard_has_no_mask( + spec in arb_standard_spec(), + ctx in arb_round_ctx(), + out_domain in 0usize..8, + ) { + let config = solve::(&spec, &ctx, OodSampleBudget::new(out_domain)).unwrap(); + prop_assert_eq!(config.mask_length(), 0); + } + } + + const SMOKE_VECTOR_SIZE: usize = 64; + const SMOKE_LOG_INV_RATE: u32 = 1; + const SMOKE_FOLDING_FACTOR: u32 = 2; + const SMOKE_OOD_BUDGET: usize = 2; + + #[test] + fn solve_works_with_basefield_embedding_zk() { + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = RoundContext { + vector_size: SMOKE_VECTOR_SIZE, + log_inv_rate: SMOKE_LOG_INV_RATE, + folding_factor: SMOKE_FOLDING_FACTOR, + }; + let config: IrsConfig = + solve(&spec, &ctx, OodSampleBudget::new(SMOKE_OOD_BUDGET)) + .expect("IRS fixture must solve"); + assert!(config.mask_length() > 0); + } +} diff --git a/src/protocols/params/layout.rs b/src/protocols/params/layout.rs new file mode 100644 index 00000000..8588b4b7 --- /dev/null +++ b/src/protocols/params/layout.rs @@ -0,0 +1,229 @@ +//! Round-skeleton layout: pure-data walk over the witness shape. +//! +//! Produces per-round shapes (vector size, log_inv_rate, folding factors) +//! independent of [`SecuritySpec`] and IRS solving. Consumed by +//! [`super::build_round`] to instantiate per-round configs and by +//! [`super::derive`] to drive the round/basecase split. + +use thiserror::Error; + +use crate::{ + algebra::embedding::Embedding, + protocols::{ + irs_commit::Config as IrsConfig, + params::spec::{RoundContext, TuningSpec}, + }, +}; + +/// Reasons a [`TuningSpec`] cannot produce a valid round layout. +/// +/// Nested into [`super::error::DeriveError`] via `#[from]`, so +/// [`round_layout`] failures propagate through +/// [`super::derive::ProtocolConfig::derive`] with `?`. +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +pub enum LayoutError { + /// `tuning.vector_size` must be a power of 2. + #[error("tuning.vector_size ({vector_size}) must be a power of 2; pad the vector")] + VectorSizeNotPowerOfTwo { vector_size: usize }, + + /// `tuning.folding_factor` must yield at least 1 at every round. + #[error("tuning.folding_factor min ({min}) must be ≥ 1")] + FoldingFactorBelowOne { min: usize }, +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct RoundShape { + pub(super) round_index: usize, + pub(super) source_vector_size: usize, + pub(super) source_log_inv_rate: u32, + pub(super) source_folding_factor: u32, + pub(super) target_folding_factor: u32, +} + +#[derive(Debug)] +pub(super) struct RoundLayout { + pub(super) shapes: Vec, + pub(super) basecase_vector_size: usize, + pub(super) basecase_log_inv_rate: u32, +} + +pub(super) fn round_layout(tuning: &TuningSpec) -> Result { + if !tuning.vector_size.is_power_of_two() { + return Err(LayoutError::VectorSizeNotPowerOfTwo { + vector_size: tuning.vector_size, + }); + } + let min_folding = tuning.folding_factor.min(); + if min_folding < 1 { + return Err(LayoutError::FoldingFactorBelowOne { min: min_folding }); + } + + let mut num_vars = tuning.vector_size.trailing_zeros() as usize; + let mut log_inv_rate = tuning.starting_log_inv_rate; + let mut shapes = Vec::new(); + + loop { + let round = shapes.len(); + let source_folding = tuning.folding_factor.at_round(round); + let target_folding = tuning.folding_factor.at_round(round.saturating_add(1)); + if num_vars < source_folding.saturating_add(target_folding) { + break; + } + shapes.push(RoundShape { + round_index: round, + source_vector_size: 1usize << num_vars, + source_log_inv_rate: log_inv_rate, + source_folding_factor: source_folding as u32, + target_folding_factor: target_folding as u32, + }); + num_vars = num_vars.saturating_sub(source_folding); + log_inv_rate = log_inv_rate.saturating_add((source_folding as u32).saturating_sub(1)); + } + + Ok(RoundLayout { + shapes, + basecase_vector_size: 1usize << num_vars, + basecase_log_inv_rate: log_inv_rate, + }) +} + +pub(super) const fn round_context(shape: &RoundShape) -> RoundContext { + RoundContext { + vector_size: shape.source_vector_size, + log_inv_rate: shape.source_log_inv_rate, + folding_factor: shape.source_folding_factor, + } +} + +pub(super) fn target_context( + shape: &RoundShape, + source: &IrsConfig, +) -> RoundContext { + RoundContext { + vector_size: source.message_length(), + log_inv_rate: shape + .source_log_inv_rate + .saturating_add(shape.source_folding_factor.saturating_sub(1)), + folding_factor: shape.target_folding_factor, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::params::spec::FoldingFactor; + + const FIXTURE_FOLDING_FACTOR: usize = 2; + const FIXTURE_LOG_INV_RATE: u32 = 1; + + const LOG_VECTOR_SIZE_NO_ROUNDS: u32 = 3; + const LOG_VECTOR_SIZE_MULTI_ROUND: u32 = 8; + + const VARIED_INITIAL_FOLDING: usize = 3; + const VARIED_STEADY_FOLDING: usize = 2; + + const RATE_STEPPING_STARTING_LOG_INV_RATE: u32 = 2; + const MIN_ROUNDS_FOR_CHAINING_TEST: usize = 2; + + fn tuning_with(vector_size: usize) -> TuningSpec { + TuningSpec { + vector_size, + starting_log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: FoldingFactor::Constant(FIXTURE_FOLDING_FACTOR), + } + } + + #[test] + fn round_layout_rate_steps_up_by_folding_minus_one() { + let tuning = TuningSpec { + vector_size: 1 << LOG_VECTOR_SIZE_MULTI_ROUND, + starting_log_inv_rate: RATE_STEPPING_STARTING_LOG_INV_RATE, + folding_factor: FoldingFactor::ConstantFromSecondRound { + initial: VARIED_INITIAL_FOLDING, + rest: VARIED_STEADY_FOLDING, + }, + }; + let layout = round_layout(&tuning).unwrap(); + + let mut expected_log_inv_rate = RATE_STEPPING_STARTING_LOG_INV_RATE; + for shape in &layout.shapes { + assert_eq!(shape.source_log_inv_rate, expected_log_inv_rate); + expected_log_inv_rate += shape.source_folding_factor.saturating_sub(1); + } + assert_eq!(layout.basecase_log_inv_rate, expected_log_inv_rate); + } + + #[test] + fn round_layout_chains_target_to_next_source_folding() { + let tuning = TuningSpec { + vector_size: 1 << LOG_VECTOR_SIZE_MULTI_ROUND, + starting_log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: FoldingFactor::ConstantFromSecondRound { + initial: VARIED_INITIAL_FOLDING, + rest: VARIED_STEADY_FOLDING, + }, + }; + let layout = round_layout(&tuning).unwrap(); + assert!( + layout.shapes.len() >= MIN_ROUNDS_FOR_CHAINING_TEST, + "need ≥ {MIN_ROUNDS_FOR_CHAINING_TEST} rounds to test chaining", + ); + for window in layout.shapes.windows(2) { + assert_eq!( + window[0].target_folding_factor, + window[1].source_folding_factor + ); + } + } + + #[test] + fn round_layout_basecase_size_consumes_remaining_num_vars() { + let tuning = tuning_with(1 << LOG_VECTOR_SIZE_MULTI_ROUND); + let layout = round_layout(&tuning).unwrap(); + let consumed: u32 = layout.shapes.iter().map(|s| s.source_folding_factor).sum(); + let initial_num_vars = tuning.vector_size.trailing_zeros(); + let remaining = initial_num_vars - consumed; + assert_eq!(layout.basecase_vector_size, 1usize << remaining); + } + + #[test] + fn round_layout_stops_when_no_room_for_source_plus_target() { + let vector_size = 1usize << LOG_VECTOR_SIZE_NO_ROUNDS; + let tuning = tuning_with(vector_size); + let layout = round_layout(&tuning).unwrap(); + assert!(layout.shapes.is_empty()); + assert_eq!(layout.basecase_vector_size, vector_size); + assert_eq!(layout.basecase_log_inv_rate, FIXTURE_LOG_INV_RATE); + } + + #[test] + fn round_layout_rejects_non_pow2_vector_size() { + let tuning = TuningSpec { + vector_size: 12, + starting_log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: FoldingFactor::Constant(FIXTURE_FOLDING_FACTOR), + }; + let err = round_layout(&tuning).expect_err("non-pow2 vector_size must fail"); + assert!( + matches!( + err, + LayoutError::VectorSizeNotPowerOfTwo { vector_size: 12 } + ), + "got {err:?}", + ); + } + + #[test] + fn round_layout_rejects_zero_folding_factor() { + let tuning = TuningSpec { + vector_size: 1 << LOG_VECTOR_SIZE_MULTI_ROUND, + starting_log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: FoldingFactor::Constant(0), + }; + let err = round_layout(&tuning).expect_err("folding_factor = 0 must fail"); + assert!( + matches!(err, LayoutError::FoldingFactorBelowOne { min: 0 }), + "got {err:?}", + ); + } +} diff --git a/src/protocols/params/mask_proximity.rs b/src/protocols/params/mask_proximity.rs new file mode 100644 index 00000000..9f9e367f --- /dev/null +++ b/src/protocols/params/mask_proximity.rs @@ -0,0 +1,167 @@ +//! Mask-proximity (Construction 7.2) builder + Lemma 7.4 γ-combination bound. +//! ZK-only. + +use ark_ff::Field; + +use crate::{ + algebra::{embedding::Identity, fields::FieldWithSize}, + bits::Bits, + protocols::{ + irs_commit::Config as IrsConfig, + mask_proximity::Config as MaskProximityConfig, + params::{ + bounds::usize_to_f64, + error::{grind_to_at, DeriveError, Pow}, + solved::Solved, + spec::SecuritySpec, + }, + }, +}; + +/// `c_zk.num_vectors` must equal `2 * num_masks` (originals + fresh). +pub fn solve( + spec: &SecuritySpec, + c_zk: IrsConfig>, + num_masks: usize, + round_index: usize, +) -> Result>, DeriveError> { + let analytic = analytic_error_bits(&c_zk, num_masks); + let pow = grind_to_at( + spec, + analytic, + Pow::RoundMaskProximity { index: round_index }, + )?; + Ok(Solved::new( + MaskProximityConfig::new(c_zk, num_masks, pow), + analytic, + )) +} + +/// γ-combination soundness (Lemma 7.4): +/// `log|F| − log(num_masks · (deg − 1))`, with `deg = c_zk.masked_message_length()`. +pub fn analytic_error_bits(c_zk: &IrsConfig>, num_masks: usize) -> Bits { + let field_bits = F::field_size_bits(); + let deg = c_zk.masked_message_length(); + if deg <= 1 || num_masks == 0 { + return Bits::new(field_bits.max(0.0)); + } + let log_combined = usize_to_f64(num_masks * deg.saturating_sub(1)).log2(); + Bits::new((field_bits - log_combined).max(0.0)) +} + +impl MaskProximityConfig { + /// Analytic soundness bits (excluding PoW) for the Lemma 7.4 γ-combination. + pub fn analytic_bits(&self) -> Bits { + analytic_error_bits(self.c_zk_commit(), self.num_masks()) + } +} + +#[cfg(test)] +mod tests { + use proptest::prelude::*; + + use super::*; + use crate::{ + algebra::fields::Field64, + hash, + protocols::{ + irs_commit::{IrsMode, IrsParams}, + params::{ + spec::{DecodingRegime, Mode}, + test_utils::{ + arb_zk_spec, assert_close, assert_pow_closes_gap, build_test_c_zk, + deterministic_spec, TEST_TARGET_RANGE, + }, + }, + }, + }; + + const FIXTURE_L_ZK: usize = 8; + const FIXTURE_NUM_MASKS: usize = 3; + const FIXTURE_LOG_INV_RATE: u32 = 1; + + #[test] + fn analytic_error_formula() { + let spec = deterministic_spec(Mode::ZeroKnowledge); + let c_zk = build_test_c_zk(&spec, FIXTURE_L_ZK, FIXTURE_LOG_INV_RATE, FIXTURE_NUM_MASKS); + + let got = f64::from(analytic_error_bits(&c_zk, FIXTURE_NUM_MASKS)); + + let field_bits = ::field_size_bits(); + let deg = c_zk.masked_message_length(); + let log_combined = ((FIXTURE_NUM_MASKS * (deg - 1)) as f64).log2(); + let expected = (field_bits - log_combined).max(0.0); + + assert_close(got, expected); + } + + #[test] + fn analytic_error_saturates_when_no_masks() { + let spec = deterministic_spec(Mode::ZeroKnowledge); + let c_zk = build_test_c_zk(&spec, 2, 1, 1); + let bits = f64::from(analytic_error_bits(&c_zk, 0)); + let field_bits = ::field_size_bits(); + assert_close(bits, field_bits.max(0.0)); + } + + proptest! { + #[test] + fn solve_assembles( + spec in arb_zk_spec(TEST_TARGET_RANGE), + log_inv_rate in 1u32..=3, + num_masks in 1usize..=8, + l_zk_log in 1u32..=5, + ) { + let c_zk = build_test_c_zk(&spec, 1usize << l_zk_log, log_inv_rate, num_masks); + let config = solve(&spec, c_zk, num_masks, 0).unwrap(); + prop_assert_eq!(config.num_masks(), num_masks); + prop_assert_eq!(config.c_zk_commit().num_vectors(), 2 * num_masks); + prop_assert_eq!(config.c_zk_commit().interleaving_depth(), 1); + } + + #[test] + fn pow_closes_gap_to_target( + spec in arb_zk_spec(TEST_TARGET_RANGE), + log_inv_rate in 1u32..=3, + num_masks in 1usize..=8, + l_zk_log in 1u32..=5, + ) { + let c_zk = build_test_c_zk(&spec, 1usize << l_zk_log, log_inv_rate, num_masks); + let analytic = analytic_error_bits(&c_zk, num_masks); + let config = solve(&spec, c_zk, num_masks, 0).unwrap(); + assert_pow_closes_gap(&spec, analytic, &config.pow()); + } + } + + #[test] + #[should_panic(expected = "c_zk.num_vectors must be 2 * num_masks")] + fn solve_rejects_mismatched_num_vectors() { + let spec = deterministic_spec(Mode::ZeroKnowledge); + let c_zk = build_test_c_zk(&spec, 2, 1, 2); + let _ = solve(&spec, c_zk, 3, 0); + } + + #[test] + #[should_panic(expected = "interleaving_depth = 1")] + fn solve_rejects_non_unit_interleaving() { + const SECURITY_TARGET_BITS: f64 = 80.0; + const NUM_VECTORS: usize = 2; + const VECTOR_SIZE: usize = 8; + const NON_UNIT_INTERLEAVING_DEPTH: usize = 2; + const RATE: f64 = 0.5; + const NUM_MASKS: usize = 1; + + let spec = deterministic_spec(Mode::ZeroKnowledge); + let c_zk = IrsConfig::>::new(IrsParams { + security_target: SECURITY_TARGET_BITS, + decoding_regime: DecodingRegime::Johnson, + hash_id: hash::BLAKE3, + num_vectors: NUM_VECTORS, + vector_size: VECTOR_SIZE, + interleaving_depth: NON_UNIT_INTERLEAVING_DEPTH, + rate: RATE, + mode: IrsMode::Standard, + }); + let _ = solve(&spec, c_zk, NUM_MASKS, 0); + } +} diff --git a/src/protocols/params/mod.rs b/src/protocols/params/mod.rs new file mode 100644 index 00000000..df9c026b --- /dev/null +++ b/src/protocols/params/mod.rs @@ -0,0 +1,37 @@ +//! Parameter selection for HVZK-WHIR. +//! +//! Soundness and ZK bound derivations (referred to in submodule comments as +//! "the bounds doc, §N") live at +//! . + +pub(crate) mod basecase; +pub(crate) mod bounds; +pub(crate) mod branch; +pub(crate) mod build_round; +pub(crate) mod code_switch; +pub(crate) mod derive; +pub(crate) mod error; +pub(crate) mod irs_commit; +pub(crate) mod layout; +pub(crate) mod mask_proximity; +pub(crate) mod protocol_config; +pub(crate) mod regime; +pub(crate) mod solved; +pub(crate) mod spec; +pub(crate) mod sumcheck; + +#[cfg(test)] +pub(crate) mod test_utils; + +pub use branch::{Branch, SolveMode}; +pub use error::{ChainSource, ChainTarget, DeriveError, Pow}; +pub use layout::LayoutError; +pub use protocol_config::{ + BasecasePlan, MaskOracleConfig, MaskOracleInfo, ProtocolConfig, RoundConfig, RoundMode, +}; +pub use solved::Solved; +pub use spec::{ + DecodingRegime, FoldingFactor, ListSize, LogInvRate, MaskCodeMessageLen, Mode, OodSampleBudget, + ParseDecodingRegimeError, PowBudget, RoundContext, SecuritySpec, TuningSpec, ZkSpec, + DEFAULT_POW_BUDGET_BITS, +}; diff --git a/src/protocols/params/protocol_config.rs b/src/protocols/params/protocol_config.rs new file mode 100644 index 00000000..42bdb4fb --- /dev/null +++ b/src/protocols/params/protocol_config.rs @@ -0,0 +1,522 @@ +//! Output of [`super::derive`]: the assembled per-round and basecase configs. + +use ark_ff::Field; + +use crate::{ + algebra::{ + embedding::{Embedding, Identity}, + fields::FieldWithSize, + }, + bits::Bits, + protocols::{ + basecase::Config as BasecaseConfig, + code_switch::Config as CodeSwitchConfig, + irs_commit::Config as IrsConfig, + mask_proximity::Config as MaskProximityConfig, + params::{ + basecase as basecase_params, + bounds::usize_to_f64, + code_switch as code_switch_params, + error::{ChainSource, ChainTarget, DeriveError, Pow}, + mask_proximity as mask_proximity_params, + solved::Solved, + spec::{ListSize, MaskCodeMessageLen, OodSampleBudget, SecuritySpec, TuningSpec}, + sumcheck as sumcheck_params, + }, + proof_of_work::Config as PowConfig, + sumcheck::Config as SumcheckConfig, + }, +}; + +#[derive(Clone, Debug)] +pub struct ProtocolConfig { + security: SecuritySpec, + tuning: TuningSpec, + rounds: Vec>, + basecase: BasecasePlan, +} + +impl ProtocolConfig { + pub(crate) const fn new( + security: SecuritySpec, + tuning: TuningSpec, + rounds: Vec>, + basecase: BasecasePlan, + ) -> Self { + Self { + security, + tuning, + rounds, + basecase, + } + } + + pub const fn security(&self) -> &SecuritySpec { + &self.security + } + + pub const fn tuning(&self) -> &TuningSpec { + &self.tuning + } + + pub fn rounds(&self) -> &[RoundConfig] { + &self.rounds + } + + pub const fn basecase(&self) -> &BasecaseConfig { + self.basecase.config() + } + + pub const fn basecase_plan(&self) -> &BasecasePlan { + &self.basecase + } + + /// `true` if every PoW slot's difficulty fits within `security.pow_budget`. + pub fn check_pow_bits(&self) -> bool { + self.validate_pow_budget().is_ok() + } + + /// Returns `true` if every post-construction invariant holds. + pub fn check_all_invariants(&self) -> bool { + self.validate().is_ok() + } + + /// Run every post-construction invariant check. + /// + /// Shape chaining runs first: the analytic recomputes behind the PoW + /// checks assume per-config shape invariants hold. + pub fn validate(&self) -> Result<(), DeriveError> { + self.validate_round_chaining()?; + self.validate_pow_budget()?; + self.validate_security_target_met()?; + Ok(()) + } + + /// Every PoW slot in the plan, in round order followed by the basecase. + /// + /// The single source of truth for slot enumeration: budget validation, + /// target validation, and the per-slot test assertions all iterate this. + /// A sub-protocol added here is automatically covered by every check. + pub(crate) fn pow_slots(&self) -> impl Iterator + '_ { + let round_slots = self.rounds.iter().flat_map(|r| { + let mask_info = r.mask_oracle_info(); + let index = r.round_index; + let cs = r.code_switch.config(); + let sumcheck = PowSlot { + kind: Pow::RoundSumcheck { index }, + pow: r.sumcheck.round_pow(), + recorded: r.sumcheck.analytic(), + recompute: sumcheck_params::analytic_error_bits(cs.source(), mask_info), + }; + let code_switch = PowSlot { + kind: Pow::RoundCodeSwitch { index }, + pow: cs.pow(), + recorded: r.code_switch.analytic(), + recompute: code_switch_params::analytic_error_bits( + cs.source(), + cs.target(), + cs.out_domain_samples(), + mask_info, + ), + }; + let mask_proximity = r.mask_oracle().map(|mo| PowSlot { + kind: Pow::RoundMaskProximity { index }, + pow: mo.mask_proximity.pow(), + recorded: mo.mask_proximity.analytic(), + recompute: mask_proximity_params::analytic_error_bits( + mo.mask_proximity.c_zk_commit(), + mo.mask_proximity.num_masks(), + ), + }); + [Some(sumcheck), Some(code_switch), mask_proximity] + .into_iter() + .flatten() + }); + + let basecase = self.basecase.config(); + let basecase_sumcheck = PowSlot { + kind: Pow::BasecaseSumcheck, + pow: basecase.sumcheck().round_pow(), + recorded: self.basecase.sumcheck_analytic(), + recompute: sumcheck_params::analytic_error_bits(basecase.commit(), None), + }; + let gamma = basecase.is_zk().then(|| PowSlot { + kind: Pow::BasecaseGammaCombination, + pow: basecase.pow(), + recorded: self.basecase.gamma_analytic(), + recompute: basecase_params::analytic_error_bits(basecase.commit()), + }); + + round_slots.chain([Some(basecase_sumcheck), gamma].into_iter().flatten()) + } + + /// For each PoW slot: verify (a) the analytic-bits floor recorded at + /// solve time still matches a fresh recompute from the config's current + /// state, and (b) `recorded + pow.difficulty() ≥ target_security_bits`. + /// + /// `grind_to_at` guarantees (b) at solve time. If (a) holds, (b) holds + /// trivially. If (a) drifts, (b) may fail — most often because a planner + /// regression overwrote an IRS field after the solver consumed it. + /// + /// `EPS` matches the `assert_pow_closes_gap` slack used by the per-slot + /// proptest helper, so validation stays consistent with test-time + /// assertions. + pub fn validate_security_target_met(&self) -> Result<(), DeriveError> { + const EPS: f64 = 1e-3; + let eps = Bits::new(EPS); + let target = Bits::new(f64::from(self.security.target_security_bits)); + for slot in self.pow_slots() { + if slot.recorded.abs_diff(slot.recompute) > eps { + return Err(DeriveError::AnalyticDrift { + pow: slot.kind, + recorded: slot.recorded, + recompute: slot.recompute, + }); + } + let pow_bits = slot.pow.difficulty(); + if slot.recorded + pow_bits + eps < target { + return Err(DeriveError::SecurityTargetNotMet { + pow: slot.kind, + analytic: slot.recorded, + pow_bits, + target, + }); + } + } + Ok(()) + } + + /// PoW slot difficulty ≤ `security.pow_budget` for every slot. + pub fn validate_pow_budget(&self) -> Result<(), DeriveError> { + let max = Bits::new(f64::from(self.security.pow_budget.bits())); + for slot in self.pow_slots() { + let required = slot.pow.difficulty(); + if required > max { + return Err(DeriveError::PowBudgetExceeded { + pow: slot.kind, + required, + max, + }); + } + } + Ok(()) + } + + /// Cross-round shape chaining: + /// - adjacent rounds: `round[i+1].source.vector_size == round[i].target.vector_size` + /// - last round → basecase: `basecase.commit.vector_size == last.target.vector_size` + /// - no rounds: `basecase.commit.vector_size == tuning.vector_size` + pub fn validate_round_chaining(&self) -> Result<(), DeriveError> { + for window in self.rounds.windows(2) { + let prev = &window[0]; + let next = &window[1]; + let expected = prev.code_switch.target().vector_size(); + let found = next.code_switch.source().vector_size(); + if expected != found { + return Err(DeriveError::RoundChainBroken { + from: ChainSource::Round(prev.round_index), + to: ChainTarget::NextRound(next.round_index), + expected, + found, + }); + } + } + + let basecase_vector_size = self.basecase.commit().vector_size(); + let expected = self.rounds.last().map_or(self.tuning.vector_size, |last| { + last.code_switch.target().vector_size() + }); + if expected != basecase_vector_size { + let from = self + .rounds + .last() + .map_or(ChainSource::Tuning, |r| ChainSource::Round(r.round_index)); + return Err(DeriveError::RoundChainBroken { + from, + to: ChainTarget::Basecase, + expected, + found: basecase_vector_size, + }); + } + + Ok(()) + } + + /// HVZK privacy error in bits, summed across ZK rounds: + /// `−log Σ_r (t_ood_r² + t_ood_r) / (2|F|)` (bounds doc, §5.3 + §5.7). + pub fn privacy_error_bits(&self) -> Bits { + let field_bits = ::field_size_bits(); + let mut total_error = 0.0_f64; + for r in &self.rounds { + if let RoundMode::ZeroKnowledge { t_ood, .. } = &r.mode { + let t = usize_to_f64(t_ood.get()); + let log_err = f64::midpoint(t * t, t).log2() - field_bits; + total_error += 2_f64.powf(log_err); + } + } + if total_error == 0.0 { + return Bits::new(f64::from(self.security.target_security_bits)); + } + Bits::new((-total_error.log2()).max(0.0)) + } +} + +impl ProtocolConfig { + /// Analytic soundness bits (excluding PoW). + pub fn analytic_bits(&self) -> Bits { + let mut min_bits = f64::from(self.basecase.config().analytic_bits()); + for round in &self.rounds { + min_bits = min_bits.min(f64::from(round.analytic_bits())); + } + Bits::new(min_bits.max(0.0)) + } +} + +/// One PoW grind slot in a derived plan: its identity, configured grind, the +/// analytic floor recorded at solve time, and a fresh recompute of that floor +/// for drift detection. +pub struct PowSlot { + pub kind: Pow, + pub pow: PowConfig, + pub recorded: Bits, + pub recompute: Bits, +} + +/// Output of [`super::basecase::solve`]: the runtime basecase config plus the +/// analytic floors its two PoW slots were ground against. +/// +/// `gamma_analytic` is always computed (the Lemma 7.4 formula is well-defined +/// in both modes) but only enters validation when the basecase is ZK — the +/// γ slot does not exist otherwise. +#[derive(Clone, Debug)] +pub struct BasecasePlan { + config: BasecaseConfig, + sumcheck_analytic: Bits, + gamma_analytic: Bits, +} + +impl BasecasePlan { + pub(crate) const fn new( + config: BasecaseConfig, + sumcheck_analytic: Bits, + gamma_analytic: Bits, + ) -> Self { + Self { + config, + sumcheck_analytic, + gamma_analytic, + } + } + + pub const fn config(&self) -> &BasecaseConfig { + &self.config + } + + pub const fn sumcheck_analytic(&self) -> Bits { + self.sumcheck_analytic + } + + pub const fn gamma_analytic(&self) -> Bits { + self.gamma_analytic + } +} + +impl std::ops::Deref for BasecasePlan { + type Target = BasecaseConfig; + + fn deref(&self) -> &BasecaseConfig { + &self.config + } +} + +#[cfg(test)] +impl ProtocolConfig { + pub(crate) const fn override_basecase_pow_for_test(&mut self, pow: PowConfig) { + self.basecase.config.set_pow_for_test(pow); + } + + pub(crate) fn truncate_rounds_for_test(&mut self, len: usize) { + self.rounds.truncate(len); + } + + pub(crate) fn corrupt_round_target_vector_size_for_test( + &mut self, + round_idx: usize, + new_size: usize, + ) { + self.rounds[round_idx] + .code_switch + .config_mut_for_test() + .target_mut_for_test() + .set_vector_size_for_test(new_size); + } + + pub(crate) fn corrupt_round_sumcheck_analytic_for_test( + &mut self, + round_idx: usize, + new_value: Bits, + ) { + self.rounds[round_idx] + .sumcheck + .set_analytic_for_test(new_value); + } +} + +#[derive(Clone, Debug)] +pub struct RoundConfig { + round_index: usize, + sumcheck: Solved>, + code_switch: Solved>, + mode: RoundMode, +} + +impl RoundConfig { + pub(crate) const fn new( + round_index: usize, + sumcheck: Solved>, + code_switch: Solved>, + mode: RoundMode, + ) -> Self { + Self { + round_index, + sumcheck, + code_switch, + mode, + } + } + + pub const fn round_index(&self) -> usize { + self.round_index + } + + pub const fn sumcheck(&self) -> &Solved> { + &self.sumcheck + } + + pub const fn code_switch(&self) -> &Solved> { + &self.code_switch + } + + pub const fn mode(&self) -> &RoundMode { + &self.mode + } + + /// Borrow the round's mask oracle if this is a ZK round. + pub fn mask_oracle(&self) -> Option<&MaskOracleConfig> { + match &self.mode { + RoundMode::Standard => None, + RoundMode::ZeroKnowledge { mask_oracle, .. } => Some(mask_oracle), + } + } + + /// Slim mask-oracle view derived from `mask_oracle()`. + pub fn mask_oracle_info(&self) -> Option { + self.mask_oracle().map(MaskOracleConfig::info) + } +} + +/// Standard vs. ZK round. +/// +/// The ZK payload owns the round's mask oracle, so a mode/oracle mismatch is +/// unrepresentable. Boxed to keep the `Standard` variant from paying the +/// oracle's footprint. +#[derive(Clone, Debug)] +pub enum RoundMode { + Standard, + ZeroKnowledge { + /// Lemma 9.9 OOD-sample budget (bounds doc §5.2). + t_ood: OodSampleBudget, + /// Sized for this round's `k + 1` masks. + mask_oracle: Box>, + }, +} + +impl RoundMode { + pub const fn is_zk(&self) -> bool { + matches!(self, Self::ZeroKnowledge { .. }) + } +} + +impl RoundConfig { + /// Round-level analytic floor: the smallest of `sumcheck`, `code_switch`, + /// and (when present) the per-round mask-oracle proximity check. + pub fn analytic_bits(&self) -> Bits { + let source = &self.code_switch.source(); + let target = &self.code_switch.target(); + let mask_info = self.mask_oracle_info(); + + let sumcheck_term = f64::from(sumcheck_params::analytic_error_bits(source, mask_info)); + let code_switch_term = f64::from(code_switch_params::analytic_error_bits( + source, + target, + self.code_switch.out_domain_samples(), + mask_info, + )); + let mask_oracle_term = self + .mask_oracle() + .map_or(f64::INFINITY, |mo| f64::from(mo.analytic_bits())); + + Bits::new( + sumcheck_term + .min(code_switch_term) + .min(mask_oracle_term) + .max(0.0), + ) + } +} + +/// One round's mask oracle: a C_zk codeword + ℓ_zk + mask-proximity check. +#[derive(Clone, Debug)] +pub struct MaskOracleConfig { + c_zk: IrsConfig>, + /// `next_pow2(r + t_ood)` (Theorem 9.6 + Lemma 9.3). + l_zk: MaskCodeMessageLen, + mask_proximity: Solved>, +} + +impl MaskOracleConfig { + pub(crate) const fn new( + c_zk: IrsConfig>, + l_zk: MaskCodeMessageLen, + mask_proximity: Solved>, + ) -> Self { + Self { + c_zk, + l_zk, + mask_proximity, + } + } + + pub const fn c_zk(&self) -> &IrsConfig> { + &self.c_zk + } + + pub const fn l_zk(&self) -> MaskCodeMessageLen { + self.l_zk + } + + pub const fn mask_proximity(&self) -> &Solved> { + &self.mask_proximity + } + + pub fn info(&self) -> MaskOracleInfo { + MaskOracleInfo { + c_zk_list_size: ListSize::new(self.c_zk.list_size()), + l_zk: self.l_zk, + } + } +} + +/// Slim mask-oracle view (C_zk's list size + ℓ_zk). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct MaskOracleInfo { + pub c_zk_list_size: ListSize, + pub l_zk: MaskCodeMessageLen, +} + +impl MaskOracleConfig { + /// Analytic soundness bits (excluding PoW) for this round's mask oracle. + pub fn analytic_bits(&self) -> Bits { + self.mask_proximity.analytic_bits() + } +} diff --git a/src/protocols/params/regime.rs b/src/protocols/params/regime.rs new file mode 100644 index 00000000..3504461d --- /dev/null +++ b/src/protocols/params/regime.rs @@ -0,0 +1,358 @@ +//! Reed–Solomon decoding regime — materialized per-round parameters and the +//! analytic helpers that depend on them. +//! +//! # References +//! +//! - Johnson proximity-gap error follows the BCSS25 improvement +//! (`O(n/η^5)`, m=10 at canonical slack) over BCIKS '20. +//! - Capacity bound follows STIR Conjecture 5.6: `(1 − ρ − η, d/(ρ·η))`-list +//! decodability for RS codes. + +use std::f64::consts::LOG2_10; + +use ordered_float::OrderedFloat; +use serde::{Deserialize, Serialize}; + +use crate::protocols::params::{ + bounds::{rate, usize_to_f64}, + spec::DecodingRegime, +}; + +/// Materialized decoding-regime parameters at a known rate. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum DecodingRegimeParams { + Unique, + Johnson { slack: OrderedFloat }, + Capacity { slack: OrderedFloat }, +} + +impl DecodingRegimeParams { + /// Materialize spec policy at a known rate. Canonical slacks: `√ρ/20` for + /// Johnson, `ρ/20` for Capacity. + // TODO: Optimize picking η. + pub fn from_policy(policy: DecodingRegime, rate: f64) -> Self { + match policy { + DecodingRegime::Unique => Self::Unique, + DecodingRegime::Johnson => Self::johnson_canonical(rate), + DecodingRegime::Capacity => Self::capacity_canonical(rate), + } + } + + /// Johnson regime with the canonical `η = √ρ / 20` slack. + pub fn johnson_canonical(rate: f64) -> Self { + Self::Johnson { + slack: OrderedFloat(rate.sqrt() / 20.0), + } + } + + /// Capacity regime with the canonical `η = ρ / 20` slack. + pub fn capacity_canonical(rate: f64) -> Self { + Self::Capacity { + slack: OrderedFloat(rate / 20.0), + } + } + + pub const fn is_unique(self) -> bool { + matches!(self, Self::Unique) + } + + /// `log₂ |Λ(C, δ)|`. + pub fn list_size_log2(self, log_degree: f64, log_inv_rate: f64) -> f64 { + match self { + Self::Unique => 0.0, + // Johnson: |Λ| = 1 / (2 η √ρ). + Self::Johnson { slack } => -1.0 - slack.into_inner().log2() + 0.5 * log_inv_rate, + // Capacity (STIR Conj 5.6): |Λ| = d / (ρ · η). + Self::Capacity { slack } => log_degree + log_inv_rate - slack.into_inner().log2(), + } + } + + /// `|Λ(C, δ)|`. + pub fn list_size(self, log_degree: f64, log_inv_rate: f64) -> f64 { + 2_f64.powf(self.list_size_log2(log_degree, log_inv_rate)) + } + + /// `log₂(1 − δ)`. + pub fn one_minus_distance_log2(self, log_inv_rate: f64) -> f64 { + let one_minus_delta = match self { + Self::Unique => f64::midpoint(1.0, rate(log_inv_rate)), + Self::Johnson { slack } => rate(log_inv_rate).sqrt() + slack.into_inner(), + Self::Capacity { slack } => rate(log_inv_rate) + slack.into_inner(), + }; + one_minus_delta.log2() + } + + /// Bits of security delivered by `ood_samples` OOD challenges (STIR Lemma 4.5): + /// `ood · (|F| − log d) − 2·log|Λ| + 1`. Returns `0` under `Unique`. + pub fn ood_security_bits( + self, + log_degree: f64, + log_inv_rate: f64, + field_bits: f64, + ood_samples: usize, + ) -> f64 { + if self.is_unique() { + return 0.0; + } + let log_list = self.list_size_log2(log_degree, log_inv_rate); + let ood = usize_to_f64(ood_samples); + ood * (field_bits - log_degree) - 2.0 * log_list + 1.0 + } + + /// `log₂ ε_mca(C, δ)` for the per-step proximity-gaps error. + /// + /// - Unique: `(k − 1) / |F|`, log = `log k − |F|` (with `+ log ρ⁻¹`). + /// - Johnson: BCSS25 Theorem 1.5 at canonical `η = √ρ/20`, `m = 10`: + /// `ε ≈ (2·10.5⁵/3) · n · ρ^{−3/2} / |F|`. + /// - Capacity: STIR Conj 5.6, `ε ≈ d / (η · ρ²) / |F|`. + pub fn eps_mca_log2(self, log_inv_rate: f64, message_length: usize, field_bits: f64) -> f64 { + let log_k = usize_to_f64(message_length).log2(); + let error = match self { + Self::Unique => log_k + log_inv_rate, + Self::Johnson { slack } => { + debug_assert!( + slack.into_inner().log2() >= -(0.5 * log_inv_rate + LOG2_10 + 1.0) - 1e-6 + ); + // BCSS25 with m = 10: log_2(2·10.5⁵/3) + log n + 1.5·log ρ⁻¹. + let bcss25_const = (2.0 * 10.5_f64.powi(5) / 3.0).log2(); + bcss25_const + log_k + 2.5 * log_inv_rate + } + Self::Capacity { slack } => { + debug_assert!(slack.into_inner().log2() >= -(log_inv_rate + LOG2_10 + 1.0) - 1e-6); + // d / (η · ρ²) at canonical η = ρ/20: log d + log 20 + 3·log ρ⁻¹. + log_k + 3.0 * log_inv_rate + LOG2_10 + 1.0 + } + }; + error - field_bits + } +} + +impl DecodingRegime { + /// `|Λ|` at canonical slack, before an IRS config exists. + pub fn list_size_estimate(self, log_degree: f64, log_inv_rate: f64) -> f64 { + DecodingRegimeParams::from_policy(self, rate(log_inv_rate)) + .list_size(log_degree, log_inv_rate) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::params::test_utils::assert_close; + + const TIGHT_EPS: f64 = 1e-12; + + fn johnson(slack: f64) -> DecodingRegimeParams { + DecodingRegimeParams::Johnson { + slack: OrderedFloat(slack), + } + } + + fn capacity(slack: f64) -> DecodingRegimeParams { + DecodingRegimeParams::Capacity { + slack: OrderedFloat(slack), + } + } + + /// Johnson list size: `|Λ| = 1 / (2η√ρ)`, log₂ form. + #[test] + fn list_size_log2_johnson_formula() { + let got = johnson(0.1).list_size_log2(/* log_degree */ 4.0, 2.0); + let expected = -1.0 - 0.1_f64.log2() + 0.5 * 2.0; + assert_close(got, expected); + } + + /// Capacity list size: `|Λ| = d / (ρ · η)`, log₂ form. + #[test] + fn list_size_log2_capacity_formula() { + let got = capacity(0.125).list_size_log2(4.0, 2.0); + let expected = 4.0 + 2.0 - 0.125_f64.log2(); + assert_close(got, expected); + } + + /// Unique-decoding regime gives `|Λ| = 1`, i.e. log = 0. + #[test] + fn list_size_log2_unique_decoding_is_zero() { + assert_close(DecodingRegimeParams::Unique.list_size_log2(4.0, 2.0), 0.0); + } + + /// `η = √ρ/20` ⇒ `|Λ| = 10/ρ` ⇒ `list_size_estimate(_, b) = 10 · 2^b`. + #[test] + fn johnson_list_size_closed_form() { + for b in [1.0, 2.0, 3.0, 5.0] { + let got = DecodingRegime::Johnson.list_size_estimate(/* log_degree */ 4.0, b); + let expected = 10.0 * 2_f64.powf(b); + assert!( + (got - expected).abs() / expected < TIGHT_EPS, + "log_inv_rate={b}: got {got} vs {expected}", + ); + } + } + + /// `η = ρ/20` ⇒ `|Λ| = 20 · d / ρ²`. + #[test] + fn capacity_list_size_closed_form() { + for (log_d, b) in [(4.0, 1.0), (6.0, 2.0), (8.0, 3.0)] { + let got = DecodingRegime::Capacity.list_size_estimate(log_d, b); + let expected = 20.0 * 2_f64.powf(log_d) * 2_f64.powf(2.0 * b); + assert!( + (got - expected).abs() / expected < TIGHT_EPS, + "log_d={log_d}, log_inv_rate={b}: got {got} vs {expected}", + ); + } + } + + #[test] + fn johnson_list_size_matches_config_list_size() { + use crate::{ + algebra::{embedding::Identity, fields::Field64}, + hash, + protocols::irs_commit::{Config, IrsMode, IrsParams}, + }; + const PLACEHOLDER_SECURITY_TARGET_BITS: f64 = 80.0; + const PLACEHOLDER_NUM_VECTORS: usize = 2; + const PLACEHOLDER_VECTOR_SIZE: usize = 8; + const PLACEHOLDER_INTERLEAVING_DEPTH: usize = 1; + const LOG_INV_RATE: u32 = 2; + + let config: Config> = Config::new(IrsParams { + security_target: PLACEHOLDER_SECURITY_TARGET_BITS, + decoding_regime: DecodingRegime::Johnson, + hash_id: hash::BLAKE3, + num_vectors: PLACEHOLDER_NUM_VECTORS, + vector_size: PLACEHOLDER_VECTOR_SIZE, + interleaving_depth: PLACEHOLDER_INTERLEAVING_DEPTH, + rate: 2_f64.powf(-f64::from(LOG_INV_RATE)), + mode: IrsMode::Standard, + }); + let log_degree = (config.masked_message_length() as f64).log2(); + let got = DecodingRegime::Johnson.list_size_estimate(log_degree, f64::from(LOG_INV_RATE)); + let expected = config.list_size(); + assert!( + (got - expected).abs() / expected < TIGHT_EPS, + "regime helper ({got}) vs Config::list_size ({expected})", + ); + } + + /// `1 − δ` in unique-decoding mode: midpoint of 1 and ρ. + #[test] + fn one_minus_distance_log2_unique() { + let log_inv_rate = 2.0; + let got = DecodingRegimeParams::Unique.one_minus_distance_log2(log_inv_rate); + let rho = 2_f64.powf(-log_inv_rate); + let expected = f64::midpoint(1.0, rho).log2(); + assert_close(got, expected); + } + + /// `1 − δ` in Johnson regime: `√ρ + η`. + #[test] + fn one_minus_distance_log2_johnson() { + let log_inv_rate = 2.0; + let eta = 0.1; + let got = johnson(eta).one_minus_distance_log2(log_inv_rate); + let rho = 2_f64.powf(-log_inv_rate); + let expected = (rho.sqrt() + eta).log2(); + assert_close(got, expected); + } + + /// `1 − δ` in Capacity regime: `ρ + η`. + #[test] + fn one_minus_distance_log2_capacity() { + let log_inv_rate = 2.0; + let eta = 0.05; + let got = capacity(eta).one_minus_distance_log2(log_inv_rate); + let rho = 2_f64.powf(-log_inv_rate); + let expected = (rho + eta).log2(); + assert_close(got, expected); + } + + /// `ood_security_bits = t · (|F| − log d) − 2·log|Λ| + 1`. Returns 0 + /// under Unique. + #[test] + fn ood_security_bits_formula() { + const LOG_DEGREE: f64 = 6.0; + const LOG_INV_RATE: f64 = 2.0; + const FIELD_BITS: f64 = 64.0; + const OOD: usize = 3; + + // Unique → 0 (no soundness from OOD). + let unique = DecodingRegimeParams::Unique.ood_security_bits( + LOG_DEGREE, + LOG_INV_RATE, + FIELD_BITS, + OOD, + ); + assert_close(unique, 0.0); + + let slack = 2_f64.powf(-LOG_INV_RATE).sqrt() / 20.0; + let got = johnson(slack).ood_security_bits(LOG_DEGREE, LOG_INV_RATE, FIELD_BITS, OOD); + let log_list = johnson(slack).list_size_log2(LOG_DEGREE, LOG_INV_RATE); + let expected = (OOD as f64) * (FIELD_BITS - LOG_DEGREE) - 2.0 * log_list + 1.0; + assert_close(got, expected); + } + + const MCA_MESSAGE_LENGTH: usize = 16; + const MCA_LOG_INV_RATE: f64 = 2.0; + const MCA_FIELD_BITS: f64 = 64.0; + + /// MCA error, unique-decoding branch: `log k + log_inv_rate − field_bits`. + #[test] + fn eps_mca_log2_unique_decoding_formula() { + let got = DecodingRegimeParams::Unique.eps_mca_log2( + MCA_LOG_INV_RATE, + MCA_MESSAGE_LENGTH, + MCA_FIELD_BITS, + ); + let expected = (MCA_MESSAGE_LENGTH as f64).log2() + MCA_LOG_INV_RATE - MCA_FIELD_BITS; + assert_close(got, expected); + } + + /// MCA error, Johnson (BCSS25): `log₂(2·10.5⁵/3) + log k + 2.5·log_inv_rate − field_bits`. + #[test] + fn eps_mca_log2_johnson_formula() { + let canonical_slack = 2_f64.powf(-MCA_LOG_INV_RATE).sqrt() / 20.0; + + let got = johnson(canonical_slack).eps_mca_log2( + MCA_LOG_INV_RATE, + MCA_MESSAGE_LENGTH, + MCA_FIELD_BITS, + ); + let bcss25_const = (2.0 * 10.5_f64.powi(5) / 3.0).log2(); + let expected = bcss25_const + (MCA_MESSAGE_LENGTH as f64).log2() + 2.5 * MCA_LOG_INV_RATE + - MCA_FIELD_BITS; + assert_close(got, expected); + } + + /// MCA error, Capacity (STIR Conj 5.6 at canonical η): + /// `log k + 3·log_inv_rate + log₂10 + 1 − field_bits`. + #[test] + fn eps_mca_log2_capacity_formula() { + let canonical_slack = 2_f64.powf(-MCA_LOG_INV_RATE) / 20.0; + + let got = capacity(canonical_slack).eps_mca_log2( + MCA_LOG_INV_RATE, + MCA_MESSAGE_LENGTH, + MCA_FIELD_BITS, + ); + let expected = (MCA_MESSAGE_LENGTH as f64).log2() + 3.0 * MCA_LOG_INV_RATE + LOG2_10 + 1.0 + - MCA_FIELD_BITS; + assert_close(got, expected); + } + + /// `from_policy` dispatches to the canonical constructor for each regime. + #[test] + fn from_policy_matches_canonical() { + assert_eq!( + DecodingRegimeParams::from_policy(DecodingRegime::Unique, 0.25), + DecodingRegimeParams::Unique, + ); + assert_eq!( + DecodingRegimeParams::from_policy(DecodingRegime::Johnson, 0.25), + DecodingRegimeParams::johnson_canonical(0.25), + ); + assert_eq!( + DecodingRegimeParams::from_policy(DecodingRegime::Capacity, 0.25), + DecodingRegimeParams::capacity_canonical(0.25), + ); + } +} diff --git a/src/protocols/params/solved.rs b/src/protocols/params/solved.rs new file mode 100644 index 00000000..f5d6aed6 --- /dev/null +++ b/src/protocols/params/solved.rs @@ -0,0 +1,57 @@ +//! Solver-recorded analytic floors. + +use std::ops::Deref; + +use crate::bits::Bits; + +/// A sub-protocol config paired with the analytic-error floor the params +/// solver consumed when grinding its PoW slot. +/// +/// Being wrapped in `Solved` is a type-level guarantee that the config came +/// out of a params solver; ad-hoc construction paths only produce the bare +/// config. Drift checks in +/// [`super::protocol_config::ProtocolConfig::validate`] compare the recorded +/// floor against a fresh recompute. +#[derive(Clone, Debug)] +pub struct Solved { + config: C, + analytic: Bits, +} + +impl Solved { + pub(crate) const fn new(config: C, analytic: Bits) -> Self { + Self { config, analytic } + } + + /// Analytic-error floor recorded at solve time. + pub const fn analytic(&self) -> Bits { + self.analytic + } + + pub const fn config(&self) -> &C { + &self.config + } + + pub fn into_config(self) -> C { + self.config + } +} + +#[cfg(test)] +impl Solved { + pub(crate) const fn config_mut_for_test(&mut self) -> &mut C { + &mut self.config + } + + pub(crate) const fn set_analytic_for_test(&mut self, analytic: Bits) { + self.analytic = analytic; + } +} + +impl Deref for Solved { + type Target = C; + + fn deref(&self) -> &C { + &self.config + } +} diff --git a/src/protocols/params/spec.rs b/src/protocols/params/spec.rs new file mode 100644 index 00000000..7673acdb --- /dev/null +++ b/src/protocols/params/spec.rs @@ -0,0 +1,415 @@ +use std::{ + fmt::{self, Debug, Display, Formatter}, + hash::{Hash, Hasher}, + marker::PhantomData, + num::NonZeroU32, + ops::Deref, + str::FromStr, +}; + +use ordered_float::OrderedFloat; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::{bits::Bits, engines::EngineId, hash}; + +/// Default per-slot PoW budget when the user does not specify one. +/// +/// 16 bits balances prover grinding cost against the security credit it buys: +/// higher values slow the prover on every slot; lower values shrink the PoW +/// contribution and push the analytic floor (and thus proof size) up. +pub const DEFAULT_POW_BUDGET_BITS: u32 = 16; + +/// Per-slot proof-of-work policy. +/// +/// `bits` plays two roles: +/// - **Planning credit**: subtracted from `target_security_bits` so solvers +/// know the analytic floor they must reach. +/// - **Validation cap**: rejects any per-slot PoW that exceeds `bits`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum PowBudget { + Forbidden, + PerSlot { bits: NonZeroU32 }, +} + +impl PowBudget { + /// `Forbidden` when `bits == 0`, else `PerSlot { bits }`. + pub const fn per_slot(bits: u32) -> Self { + match NonZeroU32::new(bits) { + Some(bits) => Self::PerSlot { bits }, + None => Self::Forbidden, + } + } + + /// Bits of grinding allowed per slot. `0` for [`PowBudget::Forbidden`]. + pub const fn bits(self) -> u32 { + match self { + Self::Forbidden => 0, + Self::PerSlot { bits } => bits.get(), + } + } +} + +/// Phantom-typed newtype. +/// +/// Trait impls are written by hand so that bounds apply to `T` only — tag +/// types stay bare, uninhabited enums. +pub struct Tagged(T, PhantomData); + +impl Tagged { + pub const fn new(v: T) -> Self { + Self(v, PhantomData) + } + + pub const fn get(self) -> T { + self.0 + } +} + +impl Debug for Tagged { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_tuple("Tagged").field(&self.0).finish() + } +} + +impl Clone for Tagged { + fn clone(&self) -> Self { + Self(self.0.clone(), PhantomData) + } +} + +impl Copy for Tagged {} + +impl PartialEq for Tagged { + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} + +impl Eq for Tagged {} + +impl Hash for Tagged { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SecuritySpec { + pub mode: Mode, + pub decoding_regime: DecodingRegime, + pub target_security_bits: u32, + pub pow_budget: PowBudget, + pub hash_id: EngineId, +} + +impl SecuritySpec { + /// Spec with canonical defaults: [`Mode::Standard`], + /// [`DecodingRegime::Johnson`], BLAKE3, and a + /// [`DEFAULT_POW_BUDGET_BITS`]-bit per-slot PoW budget. + /// + /// Override individual choices with the `with_*` methods or struct-update + /// syntax. + pub const fn new(target_security_bits: u32) -> Self { + Self { + mode: Mode::Standard, + decoding_regime: DecodingRegime::Johnson, + target_security_bits, + pow_budget: PowBudget::per_slot(DEFAULT_POW_BUDGET_BITS), + hash_id: hash::BLAKE3, + } + } + + #[must_use] + pub const fn with_mode(mut self, mode: Mode) -> Self { + self.mode = mode; + self + } + + #[must_use] + pub const fn with_decoding_regime(mut self, decoding_regime: DecodingRegime) -> Self { + self.decoding_regime = decoding_regime; + self + } + + #[must_use] + pub const fn with_pow_budget(mut self, pow_budget: PowBudget) -> Self { + self.pow_budget = pow_budget; + self + } + + #[must_use] + pub const fn with_hash(mut self, hash_id: EngineId) -> Self { + self.hash_id = hash_id; + self + } + + pub fn protocol_security_target_bits(&self) -> Bits { + let pow = self.pow_budget.bits(); + Bits::new(f64::from(self.target_security_bits.saturating_sub(pow))) + } + + /// Borrow this spec as a [`ZkSpec`] proof, or `None` in standard mode. + /// Prefer branching on this over matching `mode` and re-proving with + /// [`ZkSpec::try_new`]. + pub fn as_zk(&self) -> Option> { + ZkSpec::try_new(self) + } +} + +/// Per-round folding strategy. `at_round(i)` returns the factor for round `i`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum FoldingFactor { + /// Same folding factor across all rounds. + Constant(usize), + /// `at_round(0) = initial`; `at_round(i) = rest` for `i ≥ 1`. + ConstantFromSecondRound { initial: usize, rest: usize }, +} + +impl FoldingFactor { + pub const fn at_round(&self, round: usize) -> usize { + match self { + Self::Constant(f) => *f, + Self::ConstantFromSecondRound { initial, rest } => { + if round == 0 { + *initial + } else { + *rest + } + } + } + } + + /// Smallest factor across rounds. + pub const fn min(&self) -> usize { + match self { + Self::Constant(f) => *f, + Self::ConstantFromSecondRound { initial, rest } => { + if *initial < *rest { + *initial + } else { + *rest + } + } + } + } +} + +/// Proof-size / prover-time / soundness-margin tradeoffs. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TuningSpec { + pub vector_size: usize, + pub starting_log_inv_rate: u32, + pub folding_factor: FoldingFactor, +} + +/// Per-round context handed to a sub-protocol builder. +#[derive(Debug, Clone)] +pub struct RoundContext { + pub vector_size: usize, + pub log_inv_rate: u32, + pub folding_factor: u32, +} + +/// Standard vs. zero-knowledge selection. Orthogonal to [`DecodingRegime`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum Mode { + Standard, + ZeroKnowledge, +} + +/// A `SecuritySpec` borrow proven to be in [`Mode::ZeroKnowledge`]. +#[derive(Debug, Clone, Copy)] +pub struct ZkSpec<'a>(&'a SecuritySpec); + +impl<'a> ZkSpec<'a> { + pub fn try_new(spec: &'a SecuritySpec) -> Option { + matches!(spec.mode, Mode::ZeroKnowledge).then_some(Self(spec)) + } + + pub const fn as_inner(self) -> &'a SecuritySpec { + self.0 + } +} + +impl Deref for ZkSpec<'_> { + type Target = SecuritySpec; + fn deref(&self) -> &SecuritySpec { + self.0 + } +} + +/// Reed–Solomon decoding regime selection. +/// +/// - `Unique`: `δ < (1 − ρ)/2`, list size 1, no conjectures. +/// - `Johnson`: `δ < 1 − √ρ − η`, canonical `η = √ρ/20`. Proximity-gap error +/// per the BCSS25 improvement to BCIKS '20. +/// - `Capacity`: `δ < 1 − ρ − η`, canonical `η = ρ/20`. Conjectured list size +/// `d/(ρ·η)` and proximity-gap error per STIR Conjecture 5.6. +/// +/// WHIR's rate stepping (each round bumps `log_inv_rate` by +/// `folding_factor − 1`) pushes ρ → 1, shrinking the unique-decoding +/// radius. At high security targets or deep folding, `Unique` may exceed +/// the grind cap on per-round PoW and [`super::derive::ProtocolConfig::derive`] +/// will return `PowUngrindable` — pick `Johnson` or `Capacity` for those. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum DecodingRegime { + Unique, + Johnson, + Capacity, +} + +impl Display for DecodingRegime { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Unique => f.write_str("Unique"), + Self::Johnson => f.write_str("Johnson"), + Self::Capacity => f.write_str("Capacity"), + } + } +} + +/// Error returned by [`DecodingRegime`]'s [`FromStr`] impl. +#[derive(Debug, Clone, PartialEq, Eq, Error)] +#[error("invalid decoding regime: {0}, options are: Unique, Johnson, Capacity")] +pub struct ParseDecodingRegimeError(String); + +impl FromStr for DecodingRegime { + type Err = ParseDecodingRegimeError; + + fn from_str(s: &str) -> Result { + match s { + "Unique" => Ok(Self::Unique), + "Johnson" => Ok(Self::Johnson), + "Capacity" => Ok(Self::Capacity), + _ => Err(ParseDecodingRegimeError(s.to_owned())), + } + } +} + +#[cfg(test)] +mod decoding_regime_tests { + use super::*; + + #[test] + fn from_str_round_trips_display() { + for r in [ + DecodingRegime::Unique, + DecodingRegime::Johnson, + DecodingRegime::Capacity, + ] { + assert_eq!(r.to_string().parse::().unwrap(), r); + } + } + + #[test] + fn from_str_rejects_unknown() { + assert!("johnson".parse::().is_err()); // case-sensitive + assert!("".parse::().is_err()); + assert!("capacity".parse::().is_err()); + } +} + +pub enum OodSampleBudgetTag {} +pub enum MaskCodeMessageLenTag {} +pub enum LogInvRateTag {} + +/// OOD-sample budget (Lemma 9.9 / bounds doc §5.2). +pub type OodSampleBudget = Tagged; + +impl Tagged { + /// Sentinel for "no OOD samples". + pub const ZERO: Self = Self::new(0); +} + +/// C_zk message length (Theorem 9.6: `ℓ_zk ≥ source mask length`). +pub type MaskCodeMessageLen = Tagged; + +/// `rate = 2^-log_inv_rate`. +pub type LogInvRate = Tagged; + +/// Reed–Solomon list-decoding ball size `|Λ(C, δ)|`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct ListSize(OrderedFloat); + +impl ListSize { + pub const fn new(v: f64) -> Self { + Self(OrderedFloat(v)) + } + + pub const fn get(self) -> f64 { + self.0 .0 + } + + /// `log₂ |Λ|` — the form every analytic-error formula consumes. + pub fn log2(self) -> f64 { + self.get().log2() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hash; + + const TARGET_BITS: u32 = 100; + + fn spec(pow_budget: PowBudget) -> SecuritySpec { + SecuritySpec::new(TARGET_BITS) + .with_mode(Mode::ZeroKnowledge) + .with_pow_budget(pow_budget) + } + + #[test] + fn new_uses_documented_defaults() { + let spec = SecuritySpec::new(128); + assert_eq!(spec.mode, Mode::Standard); + assert_eq!(spec.decoding_regime, DecodingRegime::Johnson); + assert_eq!(spec.target_security_bits, 128); + assert_eq!( + spec.pow_budget, + PowBudget::per_slot(DEFAULT_POW_BUDGET_BITS) + ); + assert_eq!(spec.hash_id, hash::BLAKE3); + } + + #[test] + fn forbidden_means_no_pow_credit() { + assert_eq!( + spec(PowBudget::Forbidden).protocol_security_target_bits(), + Bits::new(f64::from(TARGET_BITS)), + ); + } + + #[test] + fn per_slot_zero_collapses_to_forbidden() { + assert_eq!(PowBudget::per_slot(0), PowBudget::Forbidden); + } + + #[test] + fn per_slot_bits_round_trip() { + assert_eq!(PowBudget::per_slot(20).bits(), 20); + assert_eq!(PowBudget::Forbidden.bits(), 0); + } + + #[test] + fn pow_credit_shifts_analytic_floor() { + assert_eq!( + spec(PowBudget::per_slot(20)).protocol_security_target_bits(), + Bits::new(80.0), + ); + assert_eq!( + spec(PowBudget::per_slot(60)).protocol_security_target_bits(), + Bits::new(40.0), + ); + } + + #[test] + fn pow_exceeding_target_saturates_to_zero() { + let pow_over_target = TARGET_BITS + 100; + assert_eq!( + spec(PowBudget::per_slot(pow_over_target)).protocol_security_target_bits(), + Bits::new(0.0), + ); + } +} diff --git a/src/protocols/params/sumcheck.rs b/src/protocols/params/sumcheck.rs new file mode 100644 index 00000000..a34cca86 --- /dev/null +++ b/src/protocols/params/sumcheck.rs @@ -0,0 +1,285 @@ +//! Sumcheck parameter selection. ZK mode adds a degree-2 mask per round +//! (Lemma 6.4, p.38). + +use crate::{ + algebra::{embedding::Embedding, fields::FieldWithSize}, + bits::Bits, + protocols::{ + irs_commit::Config as IrsConfig, + params::{ + bounds::usize_to_f64, + branch::SolveMode, + error::{grind_to_at, DeriveError, Pow}, + protocol_config::MaskOracleInfo, + solved::Solved, + spec::{RoundContext, SecuritySpec}, + }, + sumcheck::{self, Config as SumcheckConfig, SumcheckMaskLen}, + }, +}; + +/// Per-round sumcheck builder. +pub fn solve( + spec: &SecuritySpec, + ctx: &RoundContext, + source_irs: &IrsConfig, + mode: SolveMode, + pow: Pow, +) -> Result>, DeriveError> { + let (mask_oracle, output_mode) = match mode { + SolveMode::Standard => (None, sumcheck::SumcheckMode::Standard), + SolveMode::ZeroKnowledge(mask_oracle) => ( + Some(mask_oracle), + sumcheck::SumcheckMode::ZeroKnowledge { + mask_length: zk_mask_length(), + }, + ), + }; + let analytic = analytic_error_bits(source_irs, mask_oracle); + let round_pow = grind_to_at(spec, analytic, pow)?; + Ok(Solved::new( + SumcheckConfig::new( + ctx.vector_size, + round_pow, + num_sumcheck_rounds(ctx), + output_mode, + ), + analytic, + )) +} + +/// Per-sumcheck-round soundness in bits: `min(ε_mca, poly_identity_term)`. +/// +/// - Standard (degree-2): `log|F| − log|Λ(C)| − 1`. +/// - ZK (Lemma 6.5, p.40): `log|F| − log|Λ(C)| − log|Λ(C_zk)| − log ℓ_zk`. +pub fn analytic_error_bits( + source_irs: &IrsConfig, + mask_oracle: Option, +) -> Bits { + let field_bits = M::Target::field_size_bits(); + let log_list_size = source_irs.list_size().log2(); + let prox_gaps = source_irs.rbr_soundness_fold_prox_gaps(); + + let poly_id = mask_oracle.map_or(field_bits - log_list_size - 1.0, |info| { + let log_list_size_c_zk = info.c_zk_list_size.log2(); + let log_l_zk = usize_to_f64(info.l_zk.get()).log2(); + field_bits - log_list_size - log_list_size_c_zk - log_l_zk + }); + + Bits::new(prox_gaps.min(poly_id).max(0.0)) +} + +/// Number of degree-2 round-polynomial masks sumcheck contributes to C_zk +/// per round (Lemma 6.4). +pub const fn masks_required(ctx: &RoundContext) -> usize { + num_sumcheck_rounds(ctx) +} + +const fn num_sumcheck_rounds(ctx: &RoundContext) -> usize { + ctx.folding_factor as usize +} + +/// Construction 6.3 step 4(a) sends `h_j ∈ F^{ SumcheckMaskLen { + SumcheckMaskLen::new(3) +} + +#[cfg(test)] +mod tests { + use proptest::prelude::*; + + use super::*; + use crate::protocols::params::{ + irs_commit as irs_params, + spec::{ListSize, MaskCodeMessageLen, Mode, OodSampleBudget}, + test_utils::{ + arb_round_ctx, arb_standard_spec, arb_zk_spec, assert_close, assert_pow_closes_gap, + build_minimal_mask_oracle, deterministic_spec, TestEmbedding, TestField, + TestNonIdentityEmbedding, EPS, TEST_TARGET_RANGE, + }, + }; + + const FIXTURE_C_ZK_LIST_SIZE: f64 = 4.0; + const FIXTURE_L_ZK: usize = 8; + + fn build_source_irs(spec: &SecuritySpec, ctx: &RoundContext) -> IrsConfig { + irs_params::solve(spec, ctx, OodSampleBudget::ZERO).expect("source IRS fixture must solve") + } + + const FIXTURE_LOG_VECTOR_SIZE: u32 = 4; + const FIXTURE_LOG_INV_RATE: u32 = 1; + const FIXTURE_FOLDING_FACTOR: u32 = 2; + + fn fixture_ctx() -> RoundContext { + RoundContext { + vector_size: 1 << FIXTURE_LOG_VECTOR_SIZE, + log_inv_rate: FIXTURE_LOG_INV_RATE, + folding_factor: FIXTURE_FOLDING_FACTOR, + } + } + + #[test] + fn zk_mode_has_three_mask_coefficients() { + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = fixture_ctx(); + let source_irs = build_source_irs(&spec, &ctx); + let mask_oracle = + build_minimal_mask_oracle(&spec).expect("ZK spec must produce a mask oracle"); + let config = solve( + &spec, + &ctx, + &source_irs, + SolveMode::ZeroKnowledge(mask_oracle), + Pow::RoundSumcheck { index: 0 }, + ) + .unwrap(); + match config.mode() { + sumcheck::SumcheckMode::ZeroKnowledge { mask_length } => { + assert_eq!(mask_length.get(), 3); + } + sumcheck::SumcheckMode::Standard => panic!("expected ZK"), + } + } + + #[test] + fn analytic_error_standard_formula() { + let spec = deterministic_spec(Mode::Standard); + let ctx = fixture_ctx(); + let irs = build_source_irs(&spec, &ctx); + + let got = f64::from(analytic_error_bits::(&irs, None)); + + let field_bits = TestField::field_size_bits(); + let log_list = irs.list_size().log2(); + let prox = irs.rbr_soundness_fold_prox_gaps(); + let expected = prox.min(field_bits - log_list - 1.0).max(0.0); + + assert_close(got, expected); + } + + #[test] + fn analytic_error_zk_formula() { + let log_c_zk_list = FIXTURE_C_ZK_LIST_SIZE.log2(); + let log_l_zk = (FIXTURE_L_ZK as f64).log2(); + + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = fixture_ctx(); + let irs = build_source_irs(&spec, &ctx); + let info = MaskOracleInfo { + c_zk_list_size: ListSize::new(FIXTURE_C_ZK_LIST_SIZE), + l_zk: MaskCodeMessageLen::new(FIXTURE_L_ZK), + }; + + let got = f64::from(analytic_error_bits::(&irs, Some(info))); + + let field_bits = TestField::field_size_bits(); + let log_list = irs.list_size().log2(); + let prox = irs.rbr_soundness_fold_prox_gaps(); + let expected = prox + .min(field_bits - log_list - log_c_zk_list - log_l_zk) + .max(0.0); + + assert_close(got, expected); + } + + #[test] + fn analytic_error_clamps_to_zero() { + const OVERSIZED_LOG_C_ZK_LIST: i32 = 60; + const OVERSIZED_LOG_L_ZK: u32 = 30; + + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = fixture_ctx(); + let irs = build_source_irs(&spec, &ctx); + let huge = MaskOracleInfo { + c_zk_list_size: ListSize::new(2_f64.powi(OVERSIZED_LOG_C_ZK_LIST)), + l_zk: MaskCodeMessageLen::new(1 << OVERSIZED_LOG_L_ZK), + }; + let bits = f64::from(analytic_error_bits::(&irs, Some(huge))); + assert_close(bits, 0.0); + } + + proptest! { + #[test] + fn standard_mode_propagates( + spec in arb_standard_spec(TEST_TARGET_RANGE), + ctx in arb_round_ctx(), + ) { + let source_irs = build_source_irs(&spec, &ctx); + let pow = Pow::RoundSumcheck { index: 0 }; + let config = solve(&spec, &ctx, &source_irs, SolveMode::Standard, pow).unwrap(); + prop_assert!(matches!(config.mode(), sumcheck::SumcheckMode::Standard)); + } + + #[test] + fn num_rounds_matches_folding_factor( + spec in prop_oneof![ + arb_standard_spec(TEST_TARGET_RANGE), + arb_zk_spec(TEST_TARGET_RANGE), + ], + ctx in arb_round_ctx(), + ) { + let source_irs = build_source_irs(&spec, &ctx); + let pow = Pow::RoundSumcheck { index: 0 }; + let mode = build_minimal_mask_oracle(&spec) + .map_or(SolveMode::Standard, SolveMode::ZeroKnowledge); + let config = solve(&spec, &ctx, &source_irs, mode, pow).unwrap(); + prop_assert_eq!(config.num_rounds(), ctx.folding_factor as usize); + } + + #[test] + fn zk_error_le_standard_error( + spec in arb_zk_spec(TEST_TARGET_RANGE), + ctx in arb_round_ctx(), + ) { + let irs = build_source_irs(&spec, &ctx); + let mo = build_minimal_mask_oracle(&spec); + let zk = f64::from(analytic_error_bits::(&irs, mo)); + let standard = f64::from(analytic_error_bits::(&irs, None)); + prop_assert!(zk <= standard + EPS, "zk {} > standard {}", zk, standard); + } + + #[test] + fn round_pow_closes_gap_to_target( + spec in prop_oneof![ + arb_standard_spec(TEST_TARGET_RANGE), + arb_zk_spec(TEST_TARGET_RANGE), + ], + ctx in arb_round_ctx(), + ) { + let source_irs = build_source_irs(&spec, &ctx); + let mask_oracle = build_minimal_mask_oracle(&spec); + let error = analytic_error_bits(&source_irs, mask_oracle); + let pow = Pow::RoundSumcheck { index: 0 }; + let mode = mask_oracle.map_or(SolveMode::Standard, SolveMode::ZeroKnowledge); + let config = solve(&spec, &ctx, &source_irs, mode, pow).unwrap(); + assert_pow_closes_gap(&spec, error, &config.round_pow()); + } + } + + #[test] + fn solve_works_with_basefield_embedding_zk() { + let spec = deterministic_spec(Mode::ZeroKnowledge); + let ctx = fixture_ctx(); + let source_irs: IrsConfig = + irs_params::solve(&spec, &ctx, OodSampleBudget::ZERO) + .expect("source IRS fixture must solve"); + let info = MaskOracleInfo { + c_zk_list_size: ListSize::new(FIXTURE_C_ZK_LIST_SIZE), + l_zk: MaskCodeMessageLen::new(FIXTURE_L_ZK), + }; + let config = solve( + &spec, + &ctx, + &source_irs, + SolveMode::ZeroKnowledge(info), + Pow::RoundSumcheck { index: 0 }, + ) + .unwrap(); + assert!(matches!( + config.mode(), + sumcheck::SumcheckMode::ZeroKnowledge { .. } + )); + } +} diff --git a/src/protocols/params/test_utils.rs b/src/protocols/params/test_utils.rs new file mode 100644 index 00000000..bcf38665 --- /dev/null +++ b/src/protocols/params/test_utils.rs @@ -0,0 +1,171 @@ +//! Shared test fixtures. + +use std::ops::RangeInclusive; + +use proptest::prelude::*; + +use crate::{ + algebra::{ + embedding::{Basefield, Embedding, Identity}, + fields::{Field64, Field64_2}, + }, + bits::Bits, + protocols::{ + irs_commit::Config as IrsConfig, + mask_proximity::Config as MaskProximityConfig, + params::{ + branch::OodMode, + build_round::solve_t_ood, + irs_commit as irs_params, + protocol_config::MaskOracleInfo, + spec::{ + DecodingRegime, ListSize, LogInvRate, MaskCodeMessageLen, Mode, OodSampleBudget, + PowBudget, RoundContext, SecuritySpec, ZkSpec, + }, + }, + proof_of_work::Config as PowConfig, + }, +}; + +pub type TestField = Field64; +pub type TestEmbedding = Identity; +pub type TestExtensionField = Field64_2; +/// `Source = Field64, Target = Field64_2`. +pub type TestNonIdentityEmbedding = Basefield; + +pub const TEST_TARGET_RANGE: RangeInclusive = 30..=50; + +pub const FIXTURE_TARGET_BITS: u32 = 80; + +pub const EPS: f64 = 1e-9; + +pub const FIXTURE_POW_BUDGET_BITS: u32 = 60; + +pub fn deterministic_spec(mode: Mode) -> SecuritySpec { + SecuritySpec::new(FIXTURE_TARGET_BITS) + .with_mode(mode) + .with_pow_budget(PowBudget::per_slot(FIXTURE_POW_BUDGET_BITS)) +} + +fn arb_decoding_regime() -> impl Strategy { + prop_oneof![ + Just(DecodingRegime::Johnson), + Just(DecodingRegime::Unique), + Just(DecodingRegime::Capacity), + ] +} + +pub fn arb_spec( + mode: Mode, + target_range: RangeInclusive, +) -> impl Strategy { + (target_range, arb_decoding_regime()).prop_map(move |(target, decoding_regime)| { + SecuritySpec::new(target) + .with_mode(mode) + .with_decoding_regime(decoding_regime) + .with_pow_budget(PowBudget::per_slot(FIXTURE_POW_BUDGET_BITS)) + }) +} + +pub fn arb_zk_spec(target_range: RangeInclusive) -> impl Strategy { + arb_spec(Mode::ZeroKnowledge, target_range) +} + +pub fn arb_standard_spec(target_range: RangeInclusive) -> impl Strategy { + arb_spec(Mode::Standard, target_range) +} + +pub fn arb_round_ctx() -> impl Strategy { + (4u32..=8, 1u32..=4, 1u32..=3).prop_map(|(log_size, log_inv_rate, folding_factor)| { + RoundContext { + vector_size: 1usize << log_size, + log_inv_rate, + folding_factor, + } + }) +} + +/// `None` in Standard; `Some(ℓ_zk=2, c_zk rate 1/2)` in ZK. +pub fn build_minimal_mask_oracle(spec: &SecuritySpec) -> Option { + let zk_spec = ZkSpec::try_new(spec)?; + let l_zk = MaskCodeMessageLen::new(2); + let c_zk: IrsConfig = + irs_params::solve_mask_code(zk_spec, l_zk, 0, LogInvRate::new(1), 2) + .expect("minimal mask-code fixture must solve"); + Some(MaskOracleInfo { + c_zk_list_size: ListSize::new(c_zk.list_size()), + l_zk, + }) +} + +/// `analytic_error_bits + pow.difficulty() ≥ target_security_bits`. +pub fn assert_pow_closes_gap(spec: &SecuritySpec, analytic: Bits, pow: &PowConfig) { + let error = f64::from(analytic); + let pow_bits = f64::from(pow.difficulty()); + let target = f64::from(spec.target_security_bits); + assert!( + error + pow_bits >= target - 1e-3, + "error {error} + pow {pow_bits} < target {target}", + ); +} + +/// `|got − expected| < EPS`. +pub fn assert_close(got: f64, expected: f64) { + assert!( + (got - expected).abs() < EPS, + "got {got} vs expected {expected}", + ); +} + +/// C_zk fixture for `mask_proximity` tests. +pub fn build_test_c_zk( + spec: &SecuritySpec, + l_zk: usize, + log_inv_rate: u32, + num_masks: usize, +) -> IrsConfig { + let zk_spec = ZkSpec::try_new(spec).expect("build_test_c_zk requires a ZK spec"); + irs_params::solve_mask_code( + zk_spec, + MaskCodeMessageLen::new(l_zk), + 0, + LogInvRate::new(log_inv_rate), + MaskProximityConfig::::num_vectors_for(num_masks), + ) + .expect("C_zk fixture must solve") +} + +/// Builds a self-consistent `(source, target, t_ood)` triplet matching the +/// per-round shape that `code_switch::solve` expects. +pub fn build_round_io( + spec: &SecuritySpec, + log_inv_rate: u32, + folding_factor: u32, + num_vars: u32, + c_zk_log_inv_rate: Option, +) -> (IrsConfig, IrsConfig>, usize) { + let source_ctx = RoundContext { + vector_size: 1usize << num_vars, + log_inv_rate, + folding_factor, + }; + let target_log_inv_rate = log_inv_rate + folding_factor - 1; + let target_log_degree = f64::from(num_vars - folding_factor); + let target_list_size = spec + .decoding_regime + .list_size_estimate(target_log_degree, f64::from(target_log_inv_rate)); + let ood_mode = c_zk_log_inv_rate.map_or(OodMode::Standard, |rate| { + OodMode::ZeroKnowledge(LogInvRate::new(rate)) + }); + let (source, t_ood) = solve_t_ood::(spec, &source_ctx, target_list_size, ood_mode, 0) + .expect("solve_t_ood diverged in test fixture"); + + let target_ctx = RoundContext { + vector_size: source.message_length(), + log_inv_rate: target_log_inv_rate, + folding_factor, + }; + let target = irs_params::solve(spec, &target_ctx, OodSampleBudget::new(t_ood)) + .expect("target IRS fixture must solve"); + (source, target, t_ood) +} diff --git a/src/protocols/proof_of_work.rs b/src/protocols/proof_of_work.rs index a2a8aabc..ccdc3a60 100644 --- a/src/protocols/proof_of_work.rs +++ b/src/protocols/proof_of_work.rs @@ -26,8 +26,20 @@ pub struct Config { pub threshold: u64, } +/// Failure modes for [`Config::grind_to`]. +#[derive(Debug, thiserror::Error, Clone, Copy, PartialEq, Eq)] +pub enum PowError { + /// `target − analytic_error` exceeds what a single grind slot can deliver + /// (`MAX_DIFFICULTY = 60` bits, matching the grinding engine's capacity). + #[error("required {required} bits exceeds the {max} grind cap")] + GapExceedsGrindCap { required: Bits, max: Bits }, +} + +/// Largest gap a single grind slot can close (in bits). +pub const MAX_DIFFICULTY: f64 = 60.0; + pub fn threshold(difficulty: Bits) -> u64 { - assert!((0.0..=60.0).contains(&difficulty.into())); + assert!((0.0..=MAX_DIFFICULTY).contains(&difficulty.into())); let threshold = (64.0 - f64::from(difficulty)).exp2().ceil(); #[allow(clippy::cast_sign_loss)] @@ -64,6 +76,34 @@ impl Config { difficulty(self.threshold) } + /// Build a PoW config whose difficulty closes the gap `target - analytic_error`, + /// clamped at zero. + /// + /// Used by parameter solvers: each PoW slot independently lifts its own + /// soundness up to `target`. The caller is responsible for ensuring + /// `analytic_error` is computed from the local protocol step (see e.g. + /// `params::sumcheck`). + /// + /// Returns [`PowError::GapExceedsGrindCap`] if the required difficulty + /// exceeds [`MAX_DIFFICULTY`] — the spec is too tight for any single slot. + pub fn grind_to( + target: Bits, + analytic_error: Bits, + hash_id: EngineId, + ) -> Result { + let gap = (f64::from(target) - f64::from(analytic_error)).max(0.0); + if gap > MAX_DIFFICULTY { + return Err(PowError::GapExceedsGrindCap { + required: Bits::new(gap), + max: Bits::new(MAX_DIFFICULTY), + }); + } + Ok(Self { + hash_id, + threshold: threshold(Bits::new(gap)), + }) + } + #[cfg_attr(feature = "tracing", instrument(skip_all, fields(engine)))] pub fn prove(&self, prover_state: &mut ProverState) where diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 5018e22d..06e2b2e0 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -30,20 +30,102 @@ pub struct SumcheckOpening { pub mask_rlc: F, } +/// ZK sumcheck mask polynomial dimension. +/// +/// Validated at construction to be at least `MIN = 3` — the round polynomial +/// has 3 coefficients (degree-2), so the mask must have at least as many to +/// hide it. Lemma 6.4 itself only requires `ℓ_zk ≥ 2`; the `3` floor is a +/// WHIR design choice tied to the degree-2 round polynomial (see +/// `params::sumcheck::zk_mask_length`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct SumcheckMaskLen(usize); + +impl SumcheckMaskLen { + pub const MIN: usize = 3; + + pub const fn new(n: usize) -> Self { + assert!(n >= Self::MIN); + Self(n) + } + + pub const fn get(self) -> usize { + self.0 + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum SumcheckMode { + Standard, + ZeroKnowledge { mask_length: SumcheckMaskLen }, +} + +#[must_use] #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(bound = "")] pub struct Config where F: Field, { - pub field: Type, - pub initial_size: usize, - pub round_pow: proof_of_work::Config, - pub num_rounds: usize, - pub mask_length: usize, + field: Type, + initial_size: usize, + round_pow: proof_of_work::Config, + num_rounds: usize, + mode: SumcheckMode, } impl Config { + pub fn new( + initial_size: usize, + round_pow: proof_of_work::Config, + num_rounds: usize, + mode: SumcheckMode, + ) -> Self { + assert!(num_rounds == 0 || initial_size.next_power_of_two() >= 1 << num_rounds); + // `SumcheckMaskLen::new` already enforces the ≥ 3 floor at construction; + // here we only need the field-characteristic precondition from Lemma 6.4. + if matches!(mode, SumcheckMode::ZeroKnowledge { .. }) { + assert!( + !F::ONE.double().is_zero(), + "ZK sumcheck requires char(F) ≠ 2" + ); + } + Self { + field: Type::new(), + initial_size, + round_pow, + num_rounds, + mode, + } + } + + pub const fn initial_size(&self) -> usize { + self.initial_size + } + + pub const fn round_pow(&self) -> proof_of_work::Config { + self.round_pow + } + + pub const fn num_rounds(&self) -> usize { + self.num_rounds + } + + pub const fn mode(&self) -> &SumcheckMode { + &self.mode + } + + const fn mask_length(&self) -> usize { + match &self.mode { + SumcheckMode::Standard => 0, + SumcheckMode::ZeroKnowledge { mask_length } => mask_length.get(), + } + } + + #[cfg(test)] + pub(crate) const fn override_round_pow_for_test(&mut self, round_pow: proof_of_work::Config) { + self.round_pow = round_pow; + } + pub fn final_size(&self) -> usize { assert!( self.num_rounds == 0 || self.initial_size.next_power_of_two() >= 1 << self.num_rounds @@ -84,29 +166,21 @@ impl Config { assert!( self.num_rounds == 0 || self.initial_size.next_power_of_two() >= 1 << self.num_rounds ); - // Mask must cover all 3 sumcheck polynomial coefficients (c0, c1, c2) - assert!(self.mask_length == 0 || self.mask_length >= 3); assert_eq!(a.len(), self.initial_size); assert_eq!(b.len(), self.initial_size); debug_assert_eq!(dot(a, b), *sum); - assert_eq!(masks.len(), self.num_rounds * self.mask_length); + assert_eq!(masks.len(), self.num_rounds * self.mask_length()); + let half = F::from(2).inverse().unwrap(); + let polynomial_len = self.mask_length().max(3); - let mut mask_sum = F::ZERO; - let mut mask_rlc = F::ONE; - if self.mask_length > 0 && self.num_rounds > 0 { - let sum_multiple = F::from(1 << self.num_rounds.saturating_sub(1)); - mask_sum = masks.chunks_exact(self.mask_length).map(eval_01).sum::() * sum_multiple; - prover_state.prover_message(&mask_sum); - mask_rlc = prover_state.verifier_message(); - } + let (mut mask_sum, mask_rlc) = self.maybe_send_initial_mask_sum(prover_state, masks); - // We do a staggered Sumcheck loop so we can merge the inner fold+compute loops. - let mut univariate = Vec::new(); + let mut univariate = Vec::with_capacity(polynomial_len); let mut round_challenges = Vec::with_capacity(self.num_rounds); let mut prev_round_challenge = None; for (round, mask) in - chunks_exact_or_empty(masks, self.mask_length, self.num_rounds).enumerate() + chunks_exact_or_empty(masks, self.mask_length(), self.num_rounds).enumerate() { // Fold and compute sumcheck polynomial in one pass. let (c0, c2) = if let Some(w) = prev_round_challenge { @@ -116,40 +190,33 @@ impl Config { }; let c1 = *sum - c0.double() - c2; - // Optionally mask with univariate - if mask.is_empty() { - prover_state.prover_messages(&[c0, c2]); - } else { - // Initialize to round masking univariate polynomial. - univariate.clear(); - let sum_multiple = F::from(1 << self.num_rounds.saturating_sub(round + 1)); - univariate.extend(mask.iter().map(|m| sum_multiple * *m)); - - // Add constant term from previous and future masks. - univariate[0] += (mask_sum - sum_multiple * eval_01(mask)) * half; - - // Add plain sumcheck polynomial - univariate[0] += mask_rlc * c0; - univariate[1] += mask_rlc * c1; - univariate[2] += mask_rlc * c2; - - prover_state.prover_message(&univariate[0]); - prover_state.prover_messages(&univariate[2..]); + // Build round polynomial. In Standard (`mask = []`, `mask_rlc = 1`, + // `mask_sum = 0`) this collapses to `[c0, c1, c2]`. + univariate.clear(); + univariate.resize(polynomial_len, F::ZERO); + let sum_multiple = F::from(1 << self.num_rounds.saturating_sub(round + 1)); + for (u, m) in univariate.iter_mut().zip(mask.iter()) { + *u = sum_multiple * *m; } + univariate[0] += (mask_sum - sum_multiple * eval_01(mask)) * half; + univariate[0] += mask_rlc * c0; + univariate[1] += mask_rlc * c1; + univariate[2] += mask_rlc * c2; - // Receive the random evaluation point and update the sum + prover_state.prover_message(&univariate[0]); + prover_state.prover_messages(&univariate[2..]); + + // Receive the random evaluation point and update the sum. self.round_pow.prove(prover_state); let r = prover_state.verifier_message::(); round_challenges.push(r); *sum = (c2 * r + c1) * r + c0; - if self.mask_length > 0 && self.num_rounds > 0 { - let masked_sum = univariate_evaluate(&univariate, r); - mask_sum = masked_sum - mask_rlc * *sum; - } + + mask_sum = univariate_evaluate(&univariate, r) - mask_rlc * *sum; prev_round_challenge = Some(r); } if let Some(w) = prev_round_challenge { - // Final fold of the inputs (no polynomial computation) + // Final fold of the inputs (no polynomial computation). fold(a, w); fold(b, w); } @@ -161,6 +228,35 @@ impl Config { } } + fn maybe_send_initial_mask_sum( + &self, + prover_state: &mut ProverState, + masks: &[F], + ) -> (F, F) + where + H: DuplexSpongeInterface, + R: CryptoRng + RngCore, + F: Codec<[H::U]>, + { + match &self.mode { + SumcheckMode::Standard => (F::ZERO, F::ONE), + SumcheckMode::ZeroKnowledge { mask_length } => { + if self.num_rounds == 0 { + return (F::ZERO, F::ONE); + } + let sum_multiple = F::from(1 << self.num_rounds.saturating_sub(1)); + let mask_sum = masks + .chunks_exact(mask_length.get()) + .map(eval_01) + .sum::() + * sum_multiple; + prover_state.prover_message(&mask_sum); + let mask_rlc = prover_state.verifier_message(); + (mask_sum, mask_rlc) + } + } + } + #[cfg_attr(feature = "tracing", instrument(skip_all))] pub fn verify( &self, @@ -176,17 +272,10 @@ impl Config { assert!( self.num_rounds == 0 || self.initial_size.next_power_of_two() >= 1 << self.num_rounds ); - // Mask must cover all 3 sumcheck polynomial coefficients (c0, c1, c2) - assert!(self.mask_length == 0 || self.mask_length >= 3); - - let mut mask_rlc = F::ONE; - if self.mask_length > 0 && self.num_rounds > 0 { - let mask_sum: F = verifier_state.prover_message()?; - mask_rlc = verifier_state.verifier_message(); - *sum = mask_sum + mask_rlc * *sum; - } - let mut univariate = vec![F::ZERO; self.mask_length.max(3)]; + let mask_rlc = self.maybe_receive_initial_mask_sum(verifier_state, sum)?; + + let mut univariate = vec![F::ZERO; self.mask_length().max(3)]; let mut round_challenges = Vec::with_capacity(self.num_rounds); for _ in 0..self.num_rounds { // Receive all but linear coefficient. @@ -195,17 +284,17 @@ impl Config { *c = verifier_state.prover_message()?; } - // Derive linear coefficient from relation `univariate(0) + univariate(1) = sum` + // Derive linear coefficient from relation `univariate(0) + univariate(1) = sum`. univariate[1] = *sum - univariate[0].double() - univariate[2..].iter().sum::(); - // Check proof of work (if any) + // Check proof of work (if any). self.round_pow.verify(verifier_state)?; - // Receive the random evaluation point + // Receive the random evaluation point. let round_challenge = verifier_state.verifier_message::(); round_challenges.push(round_challenge); - // Update the sum + // Update the sum. *sum = univariate_evaluate(&univariate, round_challenge); } Ok(SumcheckOpening { @@ -213,17 +302,46 @@ impl Config { mask_rlc, }) } + + fn maybe_receive_initial_mask_sum( + &self, + verifier_state: &mut VerifierState, + sum: &mut F, + ) -> VerificationResult + where + H: DuplexSpongeInterface, + F: Codec<[H::U]>, + { + match &self.mode { + SumcheckMode::Standard => Ok(F::ONE), + SumcheckMode::ZeroKnowledge { .. } => { + if self.num_rounds == 0 { + return Ok(F::ONE); + } + let mask_sum: F = verifier_state.prover_message()?; + let mask_rlc = verifier_state.verifier_message(); + *sum = mask_sum + mask_rlc * *sum; + Ok(mask_rlc) + } + } + } } impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mode_str = match &self.mode { + SumcheckMode::Standard => "standard".to_string(), + SumcheckMode::ZeroKnowledge { mask_length } => { + format!("zk ℓ_zk={}", mask_length.get()) + } + }; write!( f, - "size {} rounds {} pow {:.2} ℓ_zk {}", + "size {} rounds {} pow {:.2} {}", self.initial_size, self.num_rounds, self.round_pow.difficulty(), - self.mask_length + mode_str, ) } } @@ -261,21 +379,22 @@ mod tests { Standard: Distribution, { pub fn arbitrary() -> impl Strategy { - let mask_length = prop_oneof![ - 3 => Just(0_usize), - 7 => 3_usize..20, + let mode_strategy = prop_oneof![ + 3 => Just(SumcheckMode::Standard), + 7 => (3_usize..20).prop_map(|n| SumcheckMode::ZeroKnowledge { + mask_length: SumcheckMaskLen::new(n), + }), ]; - (0_usize..(1 << 12), 0_usize..12, mask_length).prop_map( - |(initial_size, num_rounds, mask_length)| { + (0_usize..(1 << 12), 0_usize..12, mode_strategy).prop_map( + |(initial_size, num_rounds, mode)| { let num_rounds = num_rounds.min(initial_size.next_power_of_two().trailing_zeros() as usize); - Self { - field: Type::new(), + Self::new( initial_size, + proof_of_work::Config::none(), num_rounds, - round_pow: proof_of_work::Config::none(), - mask_length, - } + mode, + ) }, ) } @@ -296,7 +415,7 @@ mod tests { let initial_vector = random_vector(&mut rng, config.initial_size); let initial_covector = random_vector(&mut rng, config.initial_size); let initial_sum = dot(&initial_vector, &initial_covector); - let masks = random_vector(&mut rng, config.mask_length * config.num_rounds); + let masks = random_vector(&mut rng, config.mask_length() * config.num_rounds); // Prover let mut vector = initial_vector.clone(); @@ -323,7 +442,7 @@ mod tests { } let expected_mask_sum: F = - chunks_exact_or_empty(&masks, config.mask_length, config.num_rounds) + chunks_exact_or_empty(&masks, config.mask_length(), config.num_rounds) .zip(&point) .map(|(m, x)| univariate_evaluate(m, *x)) .sum(); @@ -345,8 +464,8 @@ mod tests { assert_eq!(verifier_sum, sum); verifier_state.check_eof().unwrap(); - // Non-ZK path: mask_rlc defaults to ONE (no combination randomness sampled) - if config.mask_length == 0 || config.num_rounds == 0 { + // Standard path: mask_rlc defaults to ONE (no combination randomness sampled). + if matches!(config.mode, SumcheckMode::Standard) || config.num_rounds == 0 { assert_eq!(mask_rlc, F::ONE); } } @@ -365,13 +484,14 @@ mod tests { fn test_single_round() { test_config( 0, - &Config:: { - field: Type::new(), - initial_size: 2, - round_pow: proof_of_work::Config::none(), - num_rounds: 1, - mask_length: 3, - }, + &Config::::new( + 2, + proof_of_work::Config::none(), + 1, + SumcheckMode::ZeroKnowledge { + mask_length: SumcheckMaskLen::new(3), + }, + ), ); } @@ -379,13 +499,14 @@ mod tests { fn test_two_rounds() { test_config( 0, - &Config:: { - field: Type::new(), - initial_size: 3, - round_pow: proof_of_work::Config::none(), - num_rounds: 2, - mask_length: 3, - }, + &Config::::new( + 3, + proof_of_work::Config::none(), + 2, + SumcheckMode::ZeroKnowledge { + mask_length: SumcheckMaskLen::new(3), + }, + ), ); } @@ -393,13 +514,14 @@ mod tests { fn test_three_rounds() { test_config( 0, - &Config:: { - field: Type::new(), - initial_size: 5, - round_pow: proof_of_work::Config::none(), - num_rounds: 3, - mask_length: 3, - }, + &Config::::new( + 5, + proof_of_work::Config::none(), + 3, + SumcheckMode::ZeroKnowledge { + mask_length: SumcheckMaskLen::new(3), + }, + ), ); } diff --git a/src/protocols/whir/config.rs b/src/protocols/whir/config.rs index 89892425..9ce1b166 100644 --- a/src/protocols/whir/config.rs +++ b/src/protocols/whir/config.rs @@ -7,10 +7,26 @@ use crate::{ algebra::{embedding::Embedding, fields::FieldWithSize}, bits::Bits, parameters::ProtocolParameters, - protocols::{irs_commit, proof_of_work, sumcheck}, - type_info::Type, + protocols::{ + irs_commit::{self, num_ood_samples, IrsMode}, + proof_of_work, sumcheck, + }, }; +/// log2 round-by-round soundness of `t_ood` OOD samples against a code with +/// the given list size — formerly `irs_commit::Config::rbr_ood_sample`. +fn rbr_ood_sample( + list_size: f64, + log_field_size: f64, + vector_size: usize, + out_domain_samples: usize, +) -> f64 { + // [STIR] Lemma 4.5. + let l_choose_2 = list_size * (list_size - 1.) / 2.; + let log_per_sample = ((vector_size - 1) as f64).log2() - log_field_size; + -l_choose_2.log2() - out_domain_samples as f64 * log_per_sample +} + impl Config { #[allow(clippy::too_many_lines)] pub fn new(size: usize, whir_parameters: &ProtocolParameters) -> Self @@ -37,14 +53,22 @@ impl Config { let mut num_variables = size.trailing_zeros() as usize; #[allow(clippy::cast_possible_wrap)] - let initial_committer = irs_commit::Config::new( + let initial_committer = irs_commit::Config::new(irs_commit::IrsParams { + security_target: protocol_security_level, + decoding_regime: whir_parameters.decoding_regime, + hash_id: whir_parameters.hash_id, + num_vectors: whir_parameters.batch_size, + vector_size: size, + interleaving_depth: 1 << whir_parameters.initial_folding_factor, + rate: 0.5_f64.powi(whir_parameters.starting_log_inv_rate as i32), + mode: IrsMode::Standard, + }); + let initial_out_domain_samples = num_ood_samples( + whir_parameters.decoding_regime, protocol_security_level, - whir_parameters.unique_decoding, - whir_parameters.hash_id, - whir_parameters.batch_size, + field_size_bits, + initial_committer.list_size(), size, - 1 << whir_parameters.initial_folding_factor, - 0.5_f64.powi(whir_parameters.starting_log_inv_rate as i32), ); // Initial sumcheck round pow bits. @@ -64,7 +88,7 @@ impl Config { let mut round_configs = Vec::new(); let mut round = 0; - let mut in_domain_samples = initial_committer.in_domain_samples; + let mut in_domain_samples = initial_committer.in_domain_samples(); let mut query_error = initial_committer.rbr_queries(); num_variables -= whir_parameters.initial_folding_factor; while num_variables >= whir_parameters.folding_factor { @@ -77,18 +101,26 @@ impl Config { let next_rate = log_inv_rate + (round_folding_factor - 1); #[allow(clippy::cast_possible_wrap)] - let irs_committer = irs_commit::Config::new( + let irs_committer = irs_commit::Config::new(irs_commit::IrsParams { + security_target: protocol_security_level, + decoding_regime: whir_parameters.decoding_regime, + hash_id: whir_parameters.hash_id, + num_vectors: 1, + vector_size: 1 << num_variables, + interleaving_depth: 1 << whir_parameters.folding_factor, + rate: 0.5_f64.powi(next_rate as i32), + mode: IrsMode::Standard, + }); + let round_out_domain_samples = num_ood_samples( + whir_parameters.decoding_regime, protocol_security_level, - whir_parameters.unique_decoding, - whir_parameters.hash_id, - 1, + field_size_bits, + irs_committer.list_size(), 1 << num_variables, - 1 << whir_parameters.folding_factor, - 0.5_f64.powi(next_rate as i32), ); let combination_error = { let log_list_size = irs_committer.list_size().log2(); - let count = irs_committer.out_domain_samples + in_domain_samples; + let count = round_out_domain_samples + in_domain_samples; let log_combination = (count as f64).log2(); field_size_bits - (log_combination + log_list_size + 1.) }; @@ -103,20 +135,20 @@ impl Config { let config = RoundConfig { irs_committer, - sumcheck: sumcheck::Config { - field: Type::new(), - initial_size: 1 << num_variables, - round_pow: pow(folding_pow_bits), - num_rounds: whir_parameters.folding_factor, - mask_length: 0, - }, + out_domain_samples: round_out_domain_samples, + sumcheck: sumcheck::Config::new( + 1 << num_variables, + pow(folding_pow_bits), + whir_parameters.folding_factor, + sumcheck::SumcheckMode::Standard, + ), pow: pow(pow_bits), }; round += 1; num_variables -= whir_parameters.folding_factor; log_inv_rate = next_rate; - in_domain_samples = config.irs_committer.in_domain_samples; + in_domain_samples = config.irs_committer.in_domain_samples(); query_error = config.irs_committer.rbr_queries(); round_configs.push(config); } @@ -131,22 +163,21 @@ impl Config { Self { initial_committer, - initial_sumcheck: sumcheck::Config { - field: Type::new(), - initial_size: size, - round_pow: pow(starting_folding_pow_bits), - num_rounds: whir_parameters.initial_folding_factor, - mask_length: 0, - }, + initial_out_domain_samples, + initial_sumcheck: sumcheck::Config::new( + size, + pow(starting_folding_pow_bits), + whir_parameters.initial_folding_factor, + sumcheck::SumcheckMode::Standard, + ), initial_skip_pow: pow(initial_skip_pow_bits), round_configs, - final_sumcheck: sumcheck::Config { - field: Type::new(), - initial_size: 1 << num_variables, - round_pow: pow(final_folding_pow_bits), - num_rounds: num_variables, - mask_length: 0, - }, + final_sumcheck: sumcheck::Config::new( + 1 << num_variables, + pow(final_folding_pow_bits), + num_variables, + sumcheck::SumcheckMode::Standard, + ), final_pow: pow(final_pow_bits), } } @@ -171,11 +202,15 @@ impl Config { security_level = security_level.min(field_size_bits - ((num_linear_forms - 1) as f64).log2()); } - let has_initial_constraints = - num_linear_forms > 0 || self.initial_committer.out_domain_samples > 0; + let has_initial_constraints = num_linear_forms > 0 || self.initial_out_domain_samples > 0; if !self.initial_committer.unique_decoding() { - security_level = security_level.min(self.initial_committer.rbr_ood_sample()); + security_level = security_level.min(rbr_ood_sample( + self.initial_committer.list_size(), + field_size_bits, + self.initial_committer.vector_size(), + self.initial_out_domain_samples, + )); } // Initial sumcheck error (or the skipped version for LDT). @@ -184,29 +219,34 @@ impl Config { let log_list_size = self.initial_committer.list_size().log2(); let initial_sumcheck_error = field_size_bits - (log_list_size + 1.); let initial_fold_error = initial_prox_gaps_error.min(initial_sumcheck_error) - + f64::from(self.initial_sumcheck.round_pow.difficulty()); + + f64::from(self.initial_sumcheck.round_pow().difficulty()); security_level = security_level.min(initial_fold_error); } else { let skipped_initial_fold_error = initial_prox_gaps_error - + (self.initial_sumcheck.num_rounds as f64).log2() + + (self.initial_sumcheck.num_rounds() as f64).log2() + f64::from(self.initial_skip_pow.difficulty()); security_level = security_level.min(skipped_initial_fold_error); } let mut rbr_queries = self.initial_committer.rbr_queries(); - let mut old_in_domain_samples = self.initial_committer.in_domain_samples; + let mut old_in_domain_samples = self.initial_committer.in_domain_samples(); for round in &self.round_configs { // Query soundness is computed at the old rate, while all fold and OOD terms use the new rate. let new_unique_decoding = round.irs_committer.unique_decoding(); if !new_unique_decoding { - let ood_error = round.irs_committer.rbr_ood_sample(); + let ood_error = rbr_ood_sample( + round.irs_committer.list_size(), + field_size_bits, + round.irs_committer.vector_size(), + round.out_domain_samples, + ); security_level = security_level.min(ood_error); } let log_list_size = round.irs_committer.list_size().log2(); let combination_error = { - let count = round.irs_committer.out_domain_samples + old_in_domain_samples; + let count = round.out_domain_samples + old_in_domain_samples; let log_combination = (count as f64).log2(); field_size_bits - (log_combination + log_list_size + 1.) }; @@ -217,19 +257,19 @@ impl Config { let prox_gaps_error = round.irs_committer.rbr_soundness_fold_prox_gaps(); let sumcheck_error = field_size_bits - (log_list_size + 1.); let round_fold_error = prox_gaps_error.min(sumcheck_error) - + f64::from(round.sumcheck.round_pow.difficulty()); + + f64::from(round.sumcheck.round_pow().difficulty()); security_level = security_level.min(round_fold_error); - old_in_domain_samples = round.irs_committer.in_domain_samples; + old_in_domain_samples = round.irs_committer.in_domain_samples(); rbr_queries = round.irs_committer.rbr_queries(); } let final_query_error = rbr_queries + f64::from(self.final_pow.difficulty()); security_level = security_level.min(final_query_error); - if self.final_sumcheck.num_rounds > 0 { + if self.final_sumcheck.num_rounds() > 0 { let final_combination_error = - field_size_bits - 1. + f64::from(self.final_sumcheck.round_pow.difficulty()); + field_size_bits - 1. + f64::from(self.final_sumcheck.round_pow().difficulty()); security_level = security_level.min(final_combination_error); } @@ -241,21 +281,21 @@ impl Config { } pub fn check_max_pow_bits(&self, max_bits: Bits) -> bool { - if self.initial_sumcheck.round_pow.difficulty() > max_bits { + if self.initial_sumcheck.round_pow().difficulty() > max_bits { return false; } for round_config in &self.round_configs { if round_config.pow.difficulty() > max_bits { return false; } - if round_config.sumcheck.round_pow.difficulty() > max_bits { + if round_config.sumcheck.round_pow().difficulty() > max_bits { return false; } } if self.final_pow.difficulty() > max_bits { return false; } - if self.final_sumcheck.round_pow.difficulty() > max_bits { + if self.final_sumcheck.round_pow().difficulty() > max_bits { return false; } true @@ -266,7 +306,7 @@ impl Config { } pub const fn initial_size(&self) -> usize { - self.initial_committer.vector_size + self.initial_committer.vector_size() } pub fn initial_num_variables(&self) -> usize { @@ -289,7 +329,7 @@ impl Display for Config { writeln!( f, "Security level: {:.2} bits using {} decoding", - self.security_level(self.initial_committer.num_vectors, 1), + self.security_level(self.initial_committer.num_vectors(), 1), if self.unique_decoding() { "unique" } else { @@ -319,7 +359,7 @@ impl Display for Config { writeln!(f, "------------------------------------")?; let field_size_bits = M::Target::field_size_bits(); - let num_vectors = self.initial_committer.num_vectors; + let num_vectors = self.initial_committer.num_vectors(); let num_linear_forms = 10; // TODO if num_vectors > 1 { let rlc_error = field_size_bits - ((num_vectors - 1) as f64).log2(); @@ -347,7 +387,12 @@ impl Display for Config { writeln!( f, "{:.1} bits -- OOD commitment", - self.initial_committer.rbr_ood_sample() + rbr_ood_sample( + self.initial_committer.list_size(), + field_size_bits, + self.initial_committer.vector_size(), + self.initial_out_domain_samples, + ) )?; } let prox_gaps_error = self.initial_committer.rbr_soundness_fold_prox_gaps(); @@ -357,28 +402,33 @@ impl Display for Config { f, "{:.1} bits -- (x{}) prox gaps: {:.1}, sumcheck: {:.1}, pow: {:.1}, list size 2^{:.1}", prox_gaps_error.min(sumcheck_error) - + f64::from(self.initial_sumcheck.round_pow.difficulty()), - self.initial_sumcheck.num_rounds, + + f64::from(self.initial_sumcheck.round_pow().difficulty()), + self.initial_sumcheck.num_rounds(), prox_gaps_error, sumcheck_error, - self.initial_sumcheck.round_pow.difficulty(), + self.initial_sumcheck.round_pow().difficulty(), log_list_size, )?; let mut query_error = self.initial_committer.rbr_queries(); - let mut old_in_domain_samples = self.initial_committer.in_domain_samples; + let mut old_in_domain_samples = self.initial_committer.in_domain_samples(); for r in &self.round_configs { if !r.irs_committer.unique_decoding() { writeln!( f, "{:.1} bits -- OOD sample", - r.irs_committer.rbr_ood_sample() + rbr_ood_sample( + r.irs_committer.list_size(), + field_size_bits, + r.irs_committer.vector_size(), + r.out_domain_samples, + ) )?; } let log_list_size = r.irs_committer.list_size().log2(); let combination_error = { - let count = r.irs_committer.out_domain_samples + old_in_domain_samples; + let count = r.out_domain_samples + old_in_domain_samples; let log_combination = (count as f64).log2(); field_size_bits - (log_combination + log_list_size + 1.) }; @@ -396,15 +446,15 @@ impl Display for Config { writeln!( f, "{:.1} bits -- (x{}) prox gaps: {:.1}, sumcheck: {:.1}, pow: {:.1}, list size 2^{:.1}", - prox_gaps_error.min(sumcheck_error) + f64::from(r.sumcheck.round_pow.difficulty()), - r.sumcheck.num_rounds, + prox_gaps_error.min(sumcheck_error) + f64::from(r.sumcheck.round_pow().difficulty()), + r.sumcheck.num_rounds(), prox_gaps_error, sumcheck_error, - r.sumcheck.round_pow.difficulty(), + r.sumcheck.round_pow().difficulty(), log_list_size )?; - old_in_domain_samples = r.irs_committer.in_domain_samples; + old_in_domain_samples = r.irs_committer.in_domain_samples(); query_error = r.irs_committer.rbr_queries(); } @@ -416,15 +466,15 @@ impl Display for Config { self.final_pow.difficulty(), )?; - if self.final_sumcheck.num_rounds > 0 { + if self.final_sumcheck.num_rounds() > 0 { let combination_error = field_size_bits - 1.; writeln!( f, "{:.1} bits -- (x{}) combination: {:.1}, pow: {:.1}", - combination_error + f64::from(self.final_sumcheck.round_pow.difficulty()), - self.final_sumcheck.num_rounds, + combination_error + f64::from(self.final_sumcheck.round_pow().difficulty()), + self.final_sumcheck.num_rounds(), combination_error, - self.final_sumcheck.round_pow.difficulty(), + self.final_sumcheck.round_pow().difficulty(), )?; } @@ -434,8 +484,11 @@ impl Display for Config { impl RoundConfig { pub fn initial_size(&self) -> usize { - assert_eq!(self.irs_committer.vector_size, self.sumcheck.initial_size); - self.sumcheck.initial_size + assert_eq!( + self.irs_committer.vector_size(), + self.sumcheck.initial_size() + ); + self.sumcheck.initial_size() } pub fn final_size(&self) -> usize { @@ -443,12 +496,12 @@ impl RoundConfig { } pub fn initial_num_variables(&self) -> usize { - assert!(self.irs_committer.vector_size.is_power_of_two()); - self.irs_committer.vector_size.ilog2() as usize + assert!(self.irs_committer.vector_size().is_power_of_two()); + self.irs_committer.vector_size().ilog2() as usize } pub fn final_num_variables(&self) -> usize { - self.initial_num_variables() - self.sumcheck.num_rounds + self.initial_num_variables() - self.sumcheck.num_rounds() } } @@ -462,8 +515,6 @@ impl Display for RoundConfig { #[cfg(test)] mod tests { - use ordered_float::OrderedFloat; - use super::*; use crate::{ algebra::{ @@ -472,11 +523,24 @@ mod tests { }, bits::Bits, hash, - protocols::matrix_commit, - type_info::Typed, utils::test_serde, }; + /// IRS committer fixture for pow-bit tests; `check_max_pow_bits` never + /// inspects it, so any well-formed config works. + fn test_irs_committer(log_inv_rate: i32) -> irs_commit::Config> { + irs_commit::Config::new(irs_commit::IrsParams { + security_target: 60.0, + decoding_regime: crate::protocols::params::DecodingRegime::Unique, + hash_id: hash::BLAKE3, + num_vectors: 1, + vector_size: 1 << 10, + interleaving_depth: 1 << 2, + rate: 0.5_f64.powi(log_inv_rate), + mode: IrsMode::Standard, + }) + } + /// Generates default WHIR parameters fn default_whir_params() -> ProtocolParameters { ProtocolParameters { @@ -484,7 +548,7 @@ mod tests { pow_bits: 20, initial_folding_factor: 4, folding_factor: 4, - unique_decoding: false, + decoding_regime: crate::protocols::params::DecodingRegime::Johnson, starting_log_inv_rate: 1, batch_size: 1, hash_id: hash::BLAKE3, @@ -519,56 +583,36 @@ mod tests { let mut config = Config::>::new(1 << 10, ¶ms); // Set all values within limits - config.initial_sumcheck.round_pow = proof_of_work::Config::from_difficulty(Bits::new(15.0)); + config + .initial_sumcheck + .override_round_pow_for_test(proof_of_work::Config::from_difficulty(Bits::new(15.0))); config.final_pow = proof_of_work::Config::from_difficulty(Bits::new(18.0)); - config.final_sumcheck.round_pow = proof_of_work::Config::from_difficulty(Bits::new(19.5)); + config + .final_sumcheck + .override_round_pow_for_test(proof_of_work::Config::from_difficulty(Bits::new(19.5))); // Ensure all rounds are within limits config.round_configs = vec![ RoundConfig { - irs_committer: irs_commit::Config { - embedding: Typed::new(embedding::Identity::new()), - num_vectors: 1, - vector_size: 1 << 10, - mask_length: 0, - codeword_length: 1 << (10 + 3 - 2), - interleaving_depth: 1 << 2, - matrix_commit: matrix_commit::Config::::new(0, 0), - johnson_slack: OrderedFloat::default(), - in_domain_samples: 5, - out_domain_samples: 2, - deduplicate_in_domain: true, - }, - sumcheck: sumcheck::Config { - field: Type::::new(), - initial_size: 1 << 10, - round_pow: proof_of_work::Config::from_difficulty(Bits::new(19.0)), - num_rounds: 2, - mask_length: 0, - }, + irs_committer: test_irs_committer(3), + out_domain_samples: 2, + sumcheck: sumcheck::Config::::new( + 1 << 10, + proof_of_work::Config::from_difficulty(Bits::new(19.0)), + 2, + sumcheck::SumcheckMode::Standard, + ), pow: proof_of_work::Config::from_difficulty(Bits::new(17.0)), }, RoundConfig { - irs_committer: irs_commit::Config { - embedding: Typed::new(embedding::Identity::new()), - num_vectors: 1, - vector_size: 1 << 10, - mask_length: 0, - codeword_length: 1 << (10 + 4 - 2), - interleaving_depth: 1 << 2, - matrix_commit: matrix_commit::Config::::new(0, 0), - johnson_slack: OrderedFloat::default(), - in_domain_samples: 6, - out_domain_samples: 2, - deduplicate_in_domain: true, - }, - sumcheck: sumcheck::Config { - field: Type::::new(), - initial_size: 1 << 10, - round_pow: proof_of_work::Config::from_difficulty(Bits::new(19.5)), - num_rounds: 2, - mask_length: 0, - }, + irs_committer: test_irs_committer(4), + out_domain_samples: 2, + sumcheck: sumcheck::Config::::new( + 1 << 10, + proof_of_work::Config::from_difficulty(Bits::new(19.5)), + 2, + sumcheck::SumcheckMode::Standard, + ), pow: proof_of_work::Config::from_difficulty(Bits::new(18.0)), }, ]; @@ -584,9 +628,13 @@ mod tests { let params = default_whir_params(); let mut config = Config::>::new(1 << 10, ¶ms); - config.initial_sumcheck.round_pow = proof_of_work::Config::from_difficulty(Bits::new(21.0)); + config + .initial_sumcheck + .override_round_pow_for_test(proof_of_work::Config::from_difficulty(Bits::new(21.0))); config.final_pow = proof_of_work::Config::from_difficulty(Bits::new(18.0)); - config.final_sumcheck.round_pow = proof_of_work::Config::from_difficulty(Bits::new(19.5)); + config + .final_sumcheck + .override_round_pow_for_test(proof_of_work::Config::from_difficulty(Bits::new(19.5))); assert!( !config.check_max_pow_bits(Bits::new(20.0)), diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 805a206e..74912af0 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -30,6 +30,9 @@ use crate::{ #[serde(bound = "")] pub struct Config { pub initial_committer: irs_commit::Config, + /// OOD samples on the initial commit (Construction 9.7-style OOD step, + /// formerly inside `irs_commit::commit`). + pub initial_out_domain_samples: usize, pub initial_sumcheck: sumcheck::Config, pub initial_skip_pow: proof_of_work::Config, pub round_configs: Vec>, @@ -41,12 +44,42 @@ pub struct Config { #[serde(bound = "")] pub struct RoundConfig { pub irs_committer: irs_commit::Config>, + /// OOD samples for this round's commit. + pub out_domain_samples: usize, pub sumcheck: sumcheck::Config, pub pow: proof_of_work::Config, } -pub type Witness> = irs_commit::Witness; -pub type Commitment = irs_commit::Commitment; +/// WHIR-level witness: IRS witness + OOD evaluations sampled at commit time. +#[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] +#[serde(bound( + serialize = "M::Source: Serialize, F: Serialize", + deserialize = "M::Source: Deserialize<'de>, F: Deserialize<'de>" +))] +pub struct Witness> { + pub irs: irs_commit::Witness, + pub out_of_domain: irs_commit::Evaluations, +} + +/// WHIR-level commitment: IRS commitment + OOD evaluations received at commit time. +#[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] +#[serde(bound(serialize = "F: Serialize", deserialize = "F: Deserialize<'de>"))] +pub struct Commitment { + pub irs: irs_commit::Commitment, + pub out_of_domain: irs_commit::Evaluations, +} + +impl> Witness { + pub fn num_vectors(&self) -> usize { + self.out_of_domain.num_columns() + } +} + +impl Commitment { + pub fn num_vectors(&self) -> usize { + self.out_of_domain.num_columns() + } +} #[must_use = "The final claim must be checked if there where any linear forms."] #[derive(Debug, Clone, PartialEq, Eq, Default)] @@ -75,6 +108,10 @@ impl FinalClaim { impl Config { /// Commit to one or more vectors. + /// + /// After the IRS commit, runs the legacy WHIR OOD step: samples + /// `initial_out_domain_samples` random points from the verifier and sends + /// each vector's evaluation at each point. #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().unwrap().len())))] pub fn commit( &self, @@ -88,7 +125,12 @@ impl Config { M::Target: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { - self.initial_committer.commit(prover_state, vectors) + let (irs, out_of_domain) = self.initial_committer.commit_with_ood( + prover_state, + vectors, + self.initial_out_domain_samples, + ); + Witness { irs, out_of_domain } } /// Receive a commitment to vectors. @@ -101,19 +143,22 @@ impl Config { M::Target: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { - self.initial_committer.receive_commitment(verifier_state) + let (irs, out_of_domain) = self + .initial_committer + .receive_commitment_with_ood(verifier_state, self.initial_out_domain_samples)?; + Ok(Commitment { irs, out_of_domain }) } /// Disable proof-of-work for test. #[cfg(test)] pub(crate) fn disable_pow(&mut self) { - self.initial_sumcheck.round_pow.threshold = u64::MAX; + self.initial_sumcheck.round_pow().threshold = u64::MAX; self.initial_skip_pow.threshold = u64::MAX; for round in &mut self.round_configs { - round.sumcheck.round_pow.threshold = u64::MAX; + round.sumcheck.round_pow().threshold = u64::MAX; round.pow.threshold = u64::MAX; } - self.final_sumcheck.round_pow.threshold = u64::MAX; + self.final_sumcheck.round_pow().threshold = u64::MAX; self.final_pow.threshold = u64::MAX; } } @@ -135,6 +180,7 @@ mod tests { }, hash, parameters::ProtocolParameters, + protocols::params::DecodingRegime, transcript::{codecs::Empty, DomainSeparator, ProverState, VerifierState}, utils::test_serde, }; @@ -178,7 +224,7 @@ mod tests { initial_folding_factor: usize, folding_factor: usize, num_points: usize, - unique_decoding: bool, + decoding_regime: DecodingRegime, pow_bits: usize, ) { // Number of coefficients in the multilinear polynomial (2^num_variables) @@ -190,7 +236,7 @@ mod tests { pow_bits, initial_folding_factor, folding_factor, - unique_decoding, + decoding_regime, starting_log_inv_rate: 1, batch_size: 1, hash_id: hash::SHA2, @@ -276,14 +322,14 @@ mod tests { let num_variables = folding_factor..=3 * folding_factor; for num_variable in num_variables { for num_points in [0, 1, 2] { - for unique_decoding in [true, false] { + for decoding_regime in [DecodingRegime::Unique, DecodingRegime::Johnson] { for pow_bits in [0, 5, 10] { eprintln!(); dbg!( folding_factor, num_variable, num_points, - unique_decoding, + decoding_regime, pow_bits ); @@ -292,7 +338,7 @@ mod tests { folding_factor, folding_factor, num_points, - unique_decoding, + decoding_regime, pow_bits, ); } @@ -304,7 +350,7 @@ mod tests { #[test] fn test_fail() { - make_whir_things(3, 2, 2, 0, false, 0); + make_whir_things(3, 2, 2, 0, DecodingRegime::Johnson, 0); } #[test] @@ -334,7 +380,7 @@ mod tests { initial_folding_factor, folding_factor, num_points, - false, + DecodingRegime::Johnson, 5, ); } @@ -353,7 +399,7 @@ mod tests { folding_factor: usize, num_points_per_poly: usize, num_vectors: usize, - unique_decoding: bool, + decoding_regime: DecodingRegime, pow_bits: usize, ) { let num_coeffs = 1 << num_variables; @@ -363,7 +409,7 @@ mod tests { pow_bits, initial_folding_factor, folding_factor, - unique_decoding, + decoding_regime, starting_log_inv_rate: 1, batch_size: 1, hash_id: hash::SHA2, @@ -484,7 +530,7 @@ mod tests { folding_factor, num_points_per_poly, num_polys, - false, + DecodingRegime::Johnson, 0, // pow_bits ); } @@ -503,7 +549,8 @@ mod tests { 2, // folding_factor 2, // num_points_per_poly 1, // num_polynomials (single!) - false, 0, + DecodingRegime::Johnson, + 0, ); } @@ -531,7 +578,7 @@ mod tests { pow_bits: 0, initial_folding_factor, folding_factor, - unique_decoding: false, + decoding_regime: DecodingRegime::Johnson, starting_log_inv_rate: 1, batch_size: 1, hash_id: hash::SHA2, @@ -627,7 +674,7 @@ mod tests { num_points_per_poly: usize, num_witnesses: usize, batch_size: usize, - unique_decoding: bool, + decoding_regime: DecodingRegime, pow_bits: usize, ) { let num_coeffs = 1 << num_variables; @@ -637,7 +684,7 @@ mod tests { pow_bits, initial_folding_factor, folding_factor, - unique_decoding, + decoding_regime, starting_log_inv_rate: 1, batch_size, // KEY: batch_size > 1 hash_id: hash::SHA2, @@ -743,7 +790,7 @@ mod tests { 1, // num_points_per_poly num_witness, batch_size, - false, + DecodingRegime::Johnson, 0, // pow_bits ); } @@ -758,7 +805,7 @@ mod tests { initial_folding_factor: usize, folding_factor: usize, num_points: usize, - unique_decoding: bool, + decoding_regime: DecodingRegime, pow_bits: usize, ) { eprintln!("\n---------------------"); @@ -768,7 +815,7 @@ mod tests { eprintln!(" initial_folding : {initial_folding_factor}"); eprintln!(" folding_factor : {folding_factor}"); eprintln!(" num_points : {num_points:?}"); - eprintln!(" unique_decoding : {unique_decoding:?}"); + eprintln!(" decoding_regime : {decoding_regime:?}"); eprintln!(" pow_bits : {pow_bits}"); // Number of coefficients in the multilinear polynomial (2^num_variables) @@ -780,7 +827,7 @@ mod tests { pow_bits, initial_folding_factor, folding_factor, - unique_decoding, + decoding_regime, starting_log_inv_rate: 1, batch_size, hash_id: hash::SHA2, @@ -865,7 +912,7 @@ mod tests { #[test] fn test_batched_whir() { let folding_factors = [1, 4]; - let unique_decoding_options = [false, true]; + let decoding_regime_options = [DecodingRegime::Johnson, DecodingRegime::Unique]; let num_points = [0, 2]; let pow_bits = [0, 10]; @@ -873,7 +920,7 @@ mod tests { let num_variables = (2 * folding_factor)..=3 * folding_factor; for num_variable in num_variables { for num_points in num_points { - for unique_decoding in unique_decoding_options { + for decoding_regime in decoding_regime_options { for pow_bits in pow_bits { for batch_size in 1..=4 { make_batched_whir_things( @@ -882,7 +929,7 @@ mod tests { folding_factor, folding_factor, num_points, - unique_decoding, + decoding_regime, pow_bits, ); } diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 04b4473e..598ed31d 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -26,8 +26,8 @@ use crate::{ }; enum RoundWitness<'a, F: Field, M: Embedding> { - Initial(Vec>>), - Round(irs_commit::Witness), + Initial(Vec>>), + Round(irs_commit::Witness), } impl Config { @@ -70,7 +70,7 @@ impl Config { // Input validation assert_eq!( num_vectors, - witnesses.len() * self.initial_committer.num_vectors + witnesses.len() * self.initial_committer.num_vectors() ); assert_eq!(evaluations.len(), num_vectors * linear_forms.len()); for vector in &vectors { @@ -103,8 +103,8 @@ impl Config { let mut vector_offset = 0; for witness in &witnesses { for (oods_eval, oods_row) in zip_strict( - witness.out_of_domain().evaluators(self.initial_size()), - witness.out_of_domain().rows(), + witness.out_of_domain.evaluators(self.initial_size()), + witness.out_of_domain.rows(), ) { for (j, vector) in vectors.iter().enumerate() { if j >= vector_offset && j < oods_row.len() + vector_offset { @@ -201,7 +201,7 @@ impl Config { // There are no constraints yet, so we can skip the sumcheck. // (If we did run it, all sumcheck vectors would be constant zero) // TODO: Don't compute evaluations and constraints in the first place. - let folding_randomness = (0..self.initial_sumcheck.num_rounds) + let folding_randomness = (0..self.initial_sumcheck.num_rounds()) .map(|_| prover_state.verifier_message()) .collect(); self.initial_skip_pow.prove(prover_state); @@ -219,8 +219,12 @@ impl Config { // Execute standard WHIR rounds on the batched vectors for (round_index, round_config) in self.round_configs.iter().enumerate() { - // Commit to the vector, this generates out-of-domain evaluations. - let new_witness = round_config.irs_committer.commit(prover_state, &[&vector]); + // Commit to the folded vector and run the per-round OOD step. + let (new_witness, out_of_domain) = round_config.irs_committer.commit_with_ood( + prover_state, + &[&vector], + round_config.out_domain_samples, + ); // Proof of work before in-domain challenges round_config.pow.prove(prover_state); @@ -228,9 +232,9 @@ impl Config { // Open the previous round's witness. let in_domain = match prev_witness { RoundWitness::Initial(init_witnesses) => { - let witness_refs: Vec<&_> = init_witnesses.iter().map(|c| &**c).collect(); + let irs_refs: Vec<&_> = init_witnesses.iter().map(|c| &c.irs).collect(); self.initial_committer - .open(prover_state, &witness_refs) + .open(prover_state, &irs_refs) .lift(self.embedding()) } RoundWitness::Round(old_witness) => { @@ -242,13 +246,11 @@ impl Config { }; // Collect constraints for this round and RLC them in - let stir_challenges = new_witness - .out_of_domain() + let stir_challenges = out_of_domain .evaluators(round_config.initial_size()) .chain(in_domain.evaluators(round_config.initial_size())) .collect::>(); - let stir_evaluations = new_witness - .out_of_domain() + let stir_evaluations = out_of_domain .values(&[M::Target::ONE]) .chain(in_domain.values(&tensor_product( &vector_rlc_coeffs, @@ -278,7 +280,7 @@ impl Config { } // Directly send the vector to the verifier. - assert_eq!(vector.len(), self.final_sumcheck.initial_size); + assert_eq!(vector.len(), self.final_sumcheck.initial_size()); for coeff in &vector { prover_state.prover_message(coeff); } @@ -289,8 +291,8 @@ impl Config { // Open and consume the final previous witness. match prev_witness { RoundWitness::Initial(init_witnesses) => { - let witness_refs: Vec<&_> = init_witnesses.iter().map(|c| &**c).collect(); - let _in_domain = self.initial_committer.open(prover_state, &witness_refs); + let irs_refs: Vec<&_> = init_witnesses.iter().map(|c| &c.irs).collect(); + let _in_domain = self.initial_committer.open(prover_state, &irs_refs); } RoundWitness::Round(old_witness) => { let prev_config = self.round_configs.last().unwrap(); diff --git a/src/protocols/whir/verifier.rs b/src/protocols/whir/verifier.rs index 8421ceab..de1736d5 100644 --- a/src/protocols/whir/verifier.rs +++ b/src/protocols/whir/verifier.rs @@ -23,11 +23,11 @@ use crate::{ enum RoundCommitment<'a, F: Field> { Initial { - commitments: &'a [&'a irs_commit::Commitment], + commitments: &'a [&'a Commitment], batching_weights: Vec, }, Round { - commitment: irs_commit::Commitment, + commitment: irs_commit::Commitment, }, } @@ -56,13 +56,20 @@ impl Config { U64: Codec<[H::U]>, Hash: ProverMessage<[H::U]>, { - let num_vectors = commitments.len() * self.initial_committer.num_vectors; + let num_vectors = commitments.len() * self.initial_committer.num_vectors(); verify!(evaluations.len().is_multiple_of(num_vectors)); let num_linear_forms = evaluations.len() / num_vectors; if num_vectors == 0 { return Ok(FinalClaim::default()); } + let expected_matrix_len = + self.initial_out_domain_samples * self.initial_committer.num_vectors(); + for commitment in commitments { + verify!(commitment.out_of_domain.points.len() == self.initial_out_domain_samples); + verify!(commitment.out_of_domain.matrix.len() == expected_matrix_len); + } + // Complete the constraint and evaluation matrix with OODs and their cross-terms. let (oods_evals, oods_matrix) = { let mut oods_evals = Vec::new(); @@ -72,8 +79,8 @@ impl Config { let mut vector_offset = 0; for commitment in commitments { for (weights, oods_row) in zip_strict( - commitment.out_of_domain().evaluators(self.initial_size()), - commitment.out_of_domain().rows(), + commitment.out_of_domain.evaluators(self.initial_size()), + commitment.out_of_domain.rows(), ) { for j in 0..num_vectors { if j >= vector_offset && j < oods_row.len() + vector_offset { @@ -122,7 +129,7 @@ impl Config { // (If we did run it, all sumcheck polynomials would be constant zero) assert_eq!(the_sum, M::Target::ZERO); let folding_randomness = - verifier_state.verifier_message_vec(self.initial_sumcheck.num_rounds); + verifier_state.verifier_message_vec(self.initial_sumcheck.num_rounds()); self.initial_skip_pow.verify(verifier_state)?; folding_randomness } else { @@ -133,10 +140,10 @@ impl Config { round_folding_randomness.push(folding_randomness); for (round_index, round_config) in self.round_configs.iter().enumerate() { - // Receive commitment to the folded vector, plus out-of-domain constraints - let commitment = round_config + // Receive commitment to the folded vector and the per-round OOD evaluations. + let (commitment, out_of_domain) = round_config .irs_committer - .receive_commitment(verifier_state)?; + .receive_commitment_with_ood(verifier_state, round_config.out_domain_samples)?; // Proof of work before in-domain challenges round_config.pow.verify(verifier_state)?; @@ -147,7 +154,8 @@ impl Config { commitments, batching_weights, } => { - let in_domain = self.initial_committer.verify(verifier_state, commitments)?; + let irs_refs: Vec<&_> = commitments.iter().map(|c| &c.irs).collect(); + let in_domain = self.initial_committer.verify(verifier_state, &irs_refs)?; // TODO: Skip lift and keep initial in-domain in subfield for evaluation. // This should be every so slightly more performant. (in_domain.lift(self.embedding()), batching_weights) @@ -162,13 +170,11 @@ impl Config { }; // Random linear combination of out- and in-domain constraints - let constraint_weights = commitment - .out_of_domain() + let constraint_weights = out_of_domain .evaluators(round_config.initial_size()) .chain(in_domain.evaluators(round_config.initial_size())) .collect::>(); - let constraint_values = commitment - .out_of_domain() + let constraint_values = out_of_domain .values(&[M::Target::ONE]) .chain(in_domain.values(&tensor_product( &poly_rlc, @@ -191,7 +197,8 @@ impl Config { } // Final round (we receive the full vector instead of a commitment) - let final_vector = verifier_state.prover_messages_vec(self.final_sumcheck.initial_size)?; + let final_vector = + verifier_state.prover_messages_vec(self.final_sumcheck.initial_size())?; // Final proof of work. self.final_pow.verify(verifier_state)?; @@ -202,7 +209,8 @@ impl Config { commitments, batching_weights, } => { - let in_domain = self.initial_committer.verify(verifier_state, commitments)?; + let irs_refs: Vec<&_> = commitments.iter().map(|c| &c.irs).collect(); + let in_domain = self.initial_committer.verify(verifier_state, &irs_refs)?; (in_domain.lift(self.embedding()), batching_weights) } RoundCommitment::Round { commitment } => { diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index 942b87e1..e895bf85 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -5,8 +5,9 @@ use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; use crate::{ + algebra::embedding::Identity, hash::Hash, - protocols::{irs_commit, whir}, + protocols::whir, transcript::{ Codec, DuplexSpongeInterface, ProverMessage, ProverState, VerificationResult, VerifierState, }, @@ -26,10 +27,10 @@ pub struct Commitment { #[derive(Clone, Debug)] pub struct Witness { pub f_hat_vectors: Vec>, - pub f_hat_witnesses: Vec>, + pub f_hat_witnesses: Vec>>, pub blinding_polynomials: Vec>, pub blinding_vectors: Vec>, - pub blinding_witness: irs_commit::Witness, + pub blinding_witness: whir::Witness>, } impl Config { @@ -53,7 +54,8 @@ impl Config { Hash: ProverMessage<[H::U]>, { assert_eq!( - self.blinded_commitment.initial_committer.num_vectors, 1, + self.blinded_commitment.initial_committer.num_vectors(), + 1, "zkWHIR currently expects one vector per commitment" ); @@ -94,7 +96,7 @@ impl Config { blinding_polynomials.push(blinding); } - let blinding_num_vectors = self.blinding_commitment.initial_committer.num_vectors; + let blinding_num_vectors = self.blinding_commitment.initial_committer.num_vectors(); assert_eq!( blinding_num_vectors, polynomials.len() * (num_witness_variables + 1), diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index b5e41c6e..8275226e 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -43,16 +43,18 @@ impl BlindingSizePolicy { .saturating_sub(main_whir_params.pow_bits); #[allow(clippy::cast_possible_wrap)] let q_delta_1 = irs_commit::num_in_domain_queries( - main_whir_params.unique_decoding, + main_whir_params.decoding_regime, protocol_security_level_main as f64, 0.5_f64.powi(main_whir_params.starting_log_inv_rate as i32), - ); + ) + .get(); #[allow(clippy::cast_possible_wrap)] let q_delta_2 = irs_commit::num_in_domain_queries( - main_whir_params.unique_decoding, + main_whir_params.decoding_regime, main_whir_params.security_level as f64, 0.5_f64.powi(main_whir_params.starting_log_inv_rate as i32), - ); + ) + .get(); // Default send-in-clear thresholds match query complexities. Self { @@ -132,7 +134,7 @@ impl Config { size_policy.t2 >= size_policy.q_delta_2, "invalid blinding size policy: T2 must satisfy T2 >= q(delta2)" ); - let k1 = 1usize << blinded.initial_sumcheck.num_rounds; + let k1 = 1usize << blinded.initial_sumcheck.num_rounds(); let k2 = blinding_first_round_interleaving_depth; let sumcheck_coeff_leakage = size_policy .sumcheck_round_degree @@ -171,12 +173,14 @@ impl Config { /// Interleaving depth of the initial IRS commitment (= 2^folding_factor). pub(crate) const fn interleaving_depth(&self) -> usize { - self.blinded_commitment.initial_committer.interleaving_depth + self.blinded_commitment + .initial_committer + .interleaving_depth() } /// Generator ω of the full NTT domain (size = num_rows × interleaving_depth). pub(crate) fn omega_full(&self) -> F { - let codeword_length = self.blinded_commitment.initial_committer.codeword_length; + let codeword_length = self.blinded_commitment.initial_committer.codeword_length(); let full_domain_size = codeword_length * self.interleaving_depth(); crate::algebra::ntt::generator(full_domain_size) } @@ -185,19 +189,19 @@ impl Config { // #[deprecated = "RS codes do not necessarily have a generator."] fn omega_sub(&self) -> F { // Assume it has a generator, and the evaluation order is 1, g, g^2, ... - let codeword_length = self.blinded_commitment.initial_committer.codeword_length; + let codeword_length = self.blinded_commitment.initial_committer.codeword_length(); crate::algebra::ntt::generator(codeword_length) } /// ζ = ω^num_rows — the interleaving_depth-th root of unity. pub(crate) fn zeta(&self) -> F { - let codeword_length = self.blinded_commitment.initial_committer.codeword_length; + let codeword_length = self.blinded_commitment.initial_committer.codeword_length(); self.omega_full().pow([codeword_length as u64]) } /// Precomputed sub-domain powers [1, ω_sub, ω_sub², ..., ω_sub^(num_rows-1)]. pub(crate) fn omega_powers(&self) -> Vec { - let codeword_length = self.blinded_commitment.initial_committer.codeword_length; + let codeword_length = self.blinded_commitment.initial_committer.codeword_length(); crate::algebra::geometric_sequence(self.omega_sub(), codeword_length) } @@ -255,6 +259,7 @@ mod tests { }, hash, parameters::ProtocolParameters, + protocols::params::DecodingRegime, transcript::{codecs::Empty, DomainSeparator, ProverState, VerifierState}, }; @@ -275,7 +280,7 @@ mod tests { fn make_test_config(num_polynomials: usize) -> Config { let whir_params = ProtocolParameters { - unique_decoding: false, + decoding_regime: DecodingRegime::Johnson, security_level: 16, pow_bits: 0, initial_folding_factor: 2, diff --git a/src/protocols/whir_zk/prover.rs b/src/protocols/whir_zk/prover.rs index d1527c13..549e5f36 100644 --- a/src/protocols/whir_zk/prover.rs +++ b/src/protocols/whir_zk/prover.rs @@ -225,7 +225,8 @@ impl Config { Hash: ProverMessage<[H::U]>, { assert_eq!( - self.blinded_commitment.initial_committer.num_vectors, 1, + self.blinded_commitment.initial_committer.num_vectors(), + 1, "zkWHIR currently expects one vector per commitment" ); assert_eq!( @@ -240,7 +241,7 @@ impl Config { ); assert_eq!( witness.blinding_vectors.len(), - self.blinding_commitment.initial_committer.num_vectors, + self.blinding_commitment.initial_committer.num_vectors(), "blinding vectors/witness mismatch" ); assert_eq!( @@ -320,7 +321,7 @@ impl Config { let initial_in_domain = { #[cfg(feature = "tracing")] let _span = tracing::info_span!("open_f_hat").entered(); - let witness_refs: Vec<_> = f_hat_witnesses.iter().collect(); + let witness_refs: Vec<_> = f_hat_witnesses.iter().map(|w| &w.irs).collect(); self.blinded_commitment .initial_committer .open(prover_state, &witness_refs) diff --git a/src/protocols/whir_zk/verifier.rs b/src/protocols/whir_zk/verifier.rs index 7df12023..6021fd7b 100644 --- a/src/protocols/whir_zk/verifier.rs +++ b/src/protocols/whir_zk/verifier.rs @@ -53,7 +53,8 @@ impl Config { Hash: ProverMessage<[H::U]>, { assert_eq!( - self.blinded_commitment.initial_committer.num_vectors, 1, + self.blinded_commitment.initial_committer.num_vectors(), + 1, "zkWHIR currently expects one vector per commitment" ); let num_polynomials = commitment.f_hat.len(); @@ -73,10 +74,11 @@ impl Config { let masking_challenge: F = verifier_state.verifier_message(); verify!(masking_challenge != F::ZERO); let commitments = commitment.f_hat.iter().collect::>(); + let irs_commitments = commitments.iter().map(|c| &c.irs).collect::>(); let initial_in_domain = self .blinded_commitment .initial_committer - .verify(verifier_state, &commitments)?; + .verify(verifier_state, &irs_commitments)?; // Expand base queries into coset points for the first folding round. let h_gammas = self.all_gammas(&initial_in_domain.points); From 2ff88f4377e568c897a10b8c2a1a8c3df9883ced Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Tue, 16 Jun 2026 09:01:59 +0200 Subject: [PATCH 13/35] wip: rs interface --- benches/expand_from_coeff.rs | 9 +- src/algebra/buffer.rs | 41 +--- src/algebra/metal_buffer.rs | 78 ++++++-- src/algebra/ntt/cooley_tukey.rs | 343 ++++++++++++++++++++++---------- src/algebra/ntt/mod.rs | 99 ++++----- src/protocols/irs_commit.rs | 17 +- 6 files changed, 371 insertions(+), 216 deletions(-) diff --git a/benches/expand_from_coeff.rs b/benches/expand_from_coeff.rs index c244514e..309a4287 100644 --- a/benches/expand_from_coeff.rs +++ b/benches/expand_from_coeff.rs @@ -1,5 +1,5 @@ use divan::{black_box, AllocProfiler, Bencher}; -use whir::algebra::{fields::Field64, ntt, random_vector}; +use whir::algebra::{fields::Field64, ntt::NttEngine, random_vector}; #[global_allocator] static ALLOC: AllocProfiler = AllocProfiler::system(); @@ -31,11 +31,12 @@ fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) { let coeffs: Vec> = (0..num_messages) .map(|_| random_vector(&mut rng, message_length)) .collect(); - (coeffs, expansion, coset_sz) + let engine = NttEngine::::new_from_fftfield(); + (engine, coeffs, expansion) }) - .bench_values(|(coeffs, expansion, _coset_sz)| { + .bench_values(|(engine, coeffs, expansion)| { let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::>(); - black_box(ntt::interleaved_rs_encode( + black_box(engine.interleaved_encode_slices( &coeffs_refs, &[], coeffs[0].len() * expansion, diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index 7430e66b..5925db29 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -3,14 +3,11 @@ use std::{any::Any, mem}; use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; -use crate::{ - algebra::{ - embedding::{Embedding, Identity}, - linear_form::{Covector, LinearForm, UnivariateEvaluation}, - mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, ntt, - sumcheck::{compute_sumcheck_polynomial, fold}, - }, - utils::chunks_exact_or_empty, +use crate::algebra::{ + embedding::{Embedding, Identity}, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, + mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, + sumcheck::{compute_sumcheck_polynomial, fold}, }; #[cfg(all(feature = "metal", target_os = "macos"))] @@ -95,13 +92,6 @@ pub trait FieldOps: Clone { accumulator: &mut Self::TargetBuffer, weight: M::Target, ); - fn interleaved_rs_encode( - vectors: &[&Self], - masks: &Self, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - ) -> Self; } #[derive( @@ -116,7 +106,7 @@ pub trait FieldOps: Clone { serde::Serialize, serde::Deserialize, )] -pub struct CpuBuffer { +pub struct CpuBuffer { data: Vec, } @@ -279,23 +269,4 @@ impl FieldOps for CpuBuffer { ) { mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.as_slice()); } - - fn interleaved_rs_encode( - vectors: &[&Self], - masks: &Self, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - ) -> Self { - let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); - let messages = vectors - .iter() - .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) - .collect::>(); - Self::from_vec(ntt::interleaved_rs_encode( - &messages, - masks.as_slice(), - codeword_length, - )) - } } diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs index 2e5c225b..c9cc2d91 100644 --- a/src/algebra/metal_buffer.rs +++ b/src/algebra/metal_buffer.rs @@ -6,7 +6,7 @@ use std::{ hash::{Hash, Hasher}, marker::PhantomData, os::raw::c_void, - sync::{Mutex, OnceLock}, + sync::{Arc, Mutex, OnceLock}, time::Instant, }; @@ -24,6 +24,7 @@ use crate::{ embedding::{Embedding, Identity}, fields::{BN254Config, Field256}, linear_form::{Covector, LinearForm, UnivariateEvaluation}, + ntt::{ReedSolomon, RsDomain}, }, hash::{metal_profile, Hash as Digest, MetalSha2}, }; @@ -970,7 +971,7 @@ struct MetalHashBuffer { } #[derive(Clone, Debug)] -pub struct MetalBuffer { +pub struct MetalBuffer { len: usize, host_cache: OnceCell>, field: Option, @@ -1518,27 +1519,70 @@ impl FieldOps for MetalBuffer { ); accumulator.invalidate_host_cache(); } +} - fn interleaved_rs_encode( - vectors: &[&Self], - masks: &Self, +/// Metal (GPU) Reed-Solomon encoder. +/// +/// Wraps the shared [`RsDomain`] core: scalar methods delegate to the domain, and the coset +/// layout used by the GPU encode is taken from [`RsDomain::coset_params`] so it can never +/// drift from [`RsDomain::evaluation_points`]. Unlike the CPU encoder, it does not need the +/// host NTT engine at all. +#[derive(Debug, Clone)] +pub struct MetalRs { + domain: Arc>, +} + +impl MetalRs { + pub fn new(domain: Arc>) -> Self { + Self { domain } + } +} + +impl ReedSolomon for MetalRs { + fn next_order(&self, size: usize) -> Option { + self.domain.next_order(size) + } + + fn generator(&self, codeword_length: usize) -> F { + self.domain.generator(codeword_length) + } + + fn evaluation_points( + &self, + masked_message_length: usize, + codeword_length: usize, + indices: &[usize], + ) -> Vec { + self.domain + .evaluation_points(masked_message_length, codeword_length, indices) + } + + fn interleaved_encode( + &self, + vectors: &[&MetalBuffer], + masks: &MetalBuffer, message_length: usize, interleaving_depth: usize, codeword_length: usize, - ) -> Self { + ) -> MetalBuffer { assert_bn254::(); let num_messages = vectors.len() * interleaving_depth; if num_messages == 0 { - return Self::from_vec(Vec::new()); + return MetalBuffer::from_vec(Vec::new()); } assert!(masks.len().is_multiple_of(num_messages)); let mask_length = masks.len() / num_messages; if vectors.len() == 1 && mask_length == 0 { + // Single source of the coset layout: derived from the shared domain rather + // than recomputed on the device. + let (coset_size, _num_cosets) = + self.domain.coset_params(message_length, codeword_length); return encode_single_vector_coset_ntt( vectors[0], message_length, interleaving_depth, codeword_length, + coset_size, ); } @@ -2156,12 +2200,12 @@ fn encode_single_vector_coset_ntt( message_length: usize, interleaving_depth: usize, codeword_length: usize, + coset_size: usize, ) -> MetalBuffer { assert!(codeword_length.is_power_of_two()); assert!(Field256::get_root_of_unity(codeword_length as u64).is_some()); assert_eq!(vector.len(), message_length * interleaving_depth); - let coset_size = message_length.next_power_of_two(); assert!(codeword_length.is_multiple_of(coset_size)); let num_cosets = codeword_length / coset_size; let total_elements = interleaving_depth @@ -2607,7 +2651,7 @@ fn as_field256_slice(values: &[T]) -> &[Field256] { #[cfg(test)] mod tests { use super::*; - use crate::algebra::buffer::CpuBuffer; + use crate::algebra::{buffer::CpuBuffer, ntt::NttEngine}; fn values(len: usize, offset: u64) -> Vec { (0..len) @@ -2683,13 +2727,21 @@ mod tests { #[test] fn metal_bn254_interleaved_rs_encode_matches_cpu() { + // The GPU encoder and the CPU reference share one coset layout: the engine and the + // domain are derived from the same field, and the GPU encode reads its layout from + // `RsDomain::coset_params` (which the engine's slice encode also uses). + let engine = NttEngine::::new_from_fftfield(); + let gpu_rs = MetalRs::new(Arc::new(RsDomain::::from_fftfield())); + let a = values(8, 1); - let cpu_a = CpuBuffer::from_slice(&a); - let cpu_masks = CpuBuffer::from_slice(&[]); let gpu_a = MetalBuffer::from_slice(&a); let gpu_masks = MetalBuffer::from_slice(&[]); - let cpu = CpuBuffer::interleaved_rs_encode(&[&cpu_a], &cpu_masks, 4, 2, 8); - let gpu = MetalBuffer::interleaved_rs_encode(&[&gpu_a], &gpu_masks, 4, 2, 8); + + // CPU reference straight from the engine's slice API: `a` is one vector of two + // length-4 messages, no masks, codeword length 8. + let messages = a.chunks_exact(4).collect::>(); + let cpu = engine.interleaved_encode_slices(&messages, &[], 8); + let gpu = gpu_rs.interleaved_encode(&[&gpu_a], &gpu_masks, 4, 2, 8); assert_eq!(gpu.as_slice(), cpu.as_slice()); } diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index fde9dc5d..8a4f782b 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -4,6 +4,9 @@ //! A global cache is used for twiddle factors. use std::sync::{RwLock, RwLockReadGuard}; +// `Arc` is only used by the CPU encoder, which is gated to non-Metal builds. +#[cfg(not(all(feature = "metal", target_os = "macos")))] +use std::sync::Arc; use ark_ff::{FftField, Field}; #[cfg(feature = "tracing")] @@ -11,11 +14,13 @@ use tracing::instrument; #[cfg(feature = "parallel")] use {crate::utils::workload_size, rayon::prelude::*, std::cmp::max}; -use super::{ - transpose, - utils::{lcm, sqrt_factor}, - ReedSolomon, -}; +use super::transpose; +use super::utils::{lcm, sqrt_factor}; +// The CPU `ReedSolomon` impl and `CpuBuffer` are only used when CPU buffers are the active backend. +#[cfg(not(all(feature = "metal", target_os = "macos")))] +use super::ReedSolomon; +#[cfg(not(all(feature = "metal", target_os = "macos")))] +use crate::algebra::buffer::CpuBuffer; #[cfg(not(feature = "rs_in_order"))] use crate::algebra::ntt::transpose::transpose_permute; use crate::{ @@ -26,13 +31,26 @@ use crate::{ // Supported primes const PRIMES: [usize; 2] = [2, 3]; -/// Engine for computing NTTs over arbitrary fields. -/// Assumes the field has large two-adicity. +/// Reed-Solomon code domain for a field `F`: the minimal, backend-neutral core. +/// +/// Holds only the NTT subgroup `order`, its `divisors`, and a primitive root, together with +/// the coset convention (`next_order` / `generator` / `coset_params` / `evaluation_points`). +/// Both the CPU engine ([`NttEngine`]) and the Metal encoder share this, so the coset layout +/// is defined in exactly one place. Assumes the field has large two-adicity. +#[derive(Debug)] +pub struct RsDomain { + pub(crate) order: usize, // order of omega_order + pub(crate) divisors: Vec, // divisors of the order. + pub(crate) omega_order: F, // primitive order'th root. +} + +/// CPU Cooley-Tukey NTT engine for a field `F`. +/// +/// Wraps the shared [`RsDomain`] and adds the host-side NTT machinery: small-order roots +/// for the butterflies and an on-demand roots-of-unity cache. CPU only. #[derive(Debug)] pub struct NttEngine { - order: usize, // order of omega_orger - divisors: Vec, // divisors of the order. - omega_order: F, // primitive order'th root. + domain: RsDomain, // Roots of small order (zero if unavailable). The naming convention is that omega_foo has order foo. half_omega_3_1_plus_2: F, // ½(ω₃ + ω₃²) @@ -48,9 +66,9 @@ pub struct NttEngine { roots: RwLock>, } -impl NttEngine { - /// Construct a new engine from the field's `FftField` trait. - pub fn new_from_fftfield() -> Self { +impl RsDomain { + /// Construct the domain from the field's `FftField` parameters. + pub fn from_fftfield() -> Self { let (mut omega, mut order) = if let (Some(mut omega), Some(b), Some(k)) = ( F::LARGE_SUBGROUP_ROOT_OF_UNITY, F::SMALL_SUBGROUP_BASE, @@ -79,8 +97,8 @@ impl NttEngine { } } -/// Creates a new NttEngine. `omega_order` must be a primitive root of unity of even order `omega`. -impl NttEngine { +impl RsDomain { + /// Creates a new domain. `omega_order` must be a primitive root of unity of even order `order`. pub fn new(order: usize, omega_order: F) -> Self { // Make sure `omega_order` is a primitive root of unity. assert_eq!(omega_order.pow([order as u64]), F::ONE); @@ -89,11 +107,99 @@ impl NttEngine { assert_ne!(omega_order.pow([(order / prime) as u64]), F::ONE); } } - - let mut res = Self { + Self { order, divisors: divisors(order, &PRIMES), omega_order, + } + } + + pub fn next_order(&self, size: usize) -> Option { + match self.divisors.binary_search(&size) { + Ok(index) | Err(index) => self.divisors.get(index).copied(), + } + } + + pub fn generator(&self, codeword_length: usize) -> F { + self.omega_order + .pow([(self.order / codeword_length) as u64]) + } + + pub fn checked_root(&self, order: usize) -> Option { + if order == 0 { + return Some(F::ONE); + } + self.order + .is_multiple_of(order) + .then(|| self.omega_order.pow([(self.order / order) as u64])) + } + + pub fn root(&self, order: usize) -> F { + self.checked_root(order) + .expect("Subgroup of requested order does not exist.") + } + + /// The single definition of the coset layout used by Reed-Solomon encoding. + /// + /// Returns `(coset_size, num_cosets)` for the given masked message length and + /// codeword length. Both [`Self::evaluation_points`] and every backend's encode + /// derive their layout from here, so the index ordering can never drift between them. + pub fn coset_params( + &self, + masked_message_length: usize, + codeword_length: usize, + ) -> (usize, usize) { + let mut coset_size = self.next_order(masked_message_length).unwrap(); + while !codeword_length.is_multiple_of(coset_size) { + coset_size = self.next_order(coset_size + 1).unwrap(); + } + (coset_size, codeword_length / coset_size) + } + + pub fn evaluation_points( + &self, + masked_message_length: usize, + codeword_length: usize, + indices: &[usize], + ) -> Vec { + assert!(masked_message_length <= codeword_length); + assert!(self.order.is_multiple_of(codeword_length)); + let mut result = Vec::new(); + let generator = self.generator(codeword_length); + + let (coset_size, num_cosets) = self.coset_params(masked_message_length, codeword_length); + #[cfg(feature = "rs_in_order")] + let _ = (coset_size, num_cosets); + + for &index in indices { + assert!(index < codeword_length); + + #[cfg(not(feature = "rs_in_order"))] + let index = transpose_permute(index, num_cosets, coset_size); + result.push(generator.pow([index as u64])); + } + result + } +} + +impl NttEngine { + /// Construct a new engine from the field's `FftField` trait. + pub fn new_from_fftfield() -> Self { + Self::from_domain(RsDomain::from_fftfield()) + } +} + +/// Creates a new NttEngine. `omega_order` must be a primitive root of unity of even order `omega`. +impl NttEngine { + pub fn new(order: usize, omega_order: F) -> Self { + Self::from_domain(RsDomain::new(order, omega_order)) + } + + /// Build the CPU engine on top of a shared [`RsDomain`], precomputing small-order roots. + pub fn from_domain(domain: RsDomain) -> Self { + let order = domain.order; + let mut res = Self { + domain, half_omega_3_1_plus_2: F::ZERO, half_omega_3_1_min_2: F::ZERO, omega_4_1: F::ZERO, @@ -105,27 +211,32 @@ impl NttEngine { roots: RwLock::new(Vec::new()), }; if order.is_multiple_of(3) { - let omega_3_1 = res.root(3); + let omega_3_1 = res.domain.root(3); let omega_3_2 = omega_3_1 * omega_3_1; // Note: char F cannot be 2 and so division by 2 works, because primitive roots of unity with even order exist. res.half_omega_3_1_min_2 = (omega_3_1 - omega_3_2) / F::from(2u64); res.half_omega_3_1_plus_2 = (omega_3_1 + omega_3_2) / F::from(2u64); } if order.is_multiple_of(4) { - res.omega_4_1 = res.root(4); + res.omega_4_1 = res.domain.root(4); } if order.is_multiple_of(8) { - res.omega_8_1 = res.root(8); + res.omega_8_1 = res.domain.root(8); res.omega_8_3 = res.omega_8_1.pow([3]); } if order.is_multiple_of(16) { - res.omega_16_1 = res.root(16); + res.omega_16_1 = res.domain.root(16); res.omega_16_3 = res.omega_16_1.pow([3]); res.omega_16_9 = res.omega_16_1.pow([9]); } res } + /// The shared field/domain core this engine is built on. + pub fn domain(&self) -> &RsDomain { + &self.domain + } + pub fn ntt(&self, values: &mut [F]) { self.ntt_batch(values, values.len()); } @@ -168,24 +279,10 @@ impl NttEngine { self.ntt_batch(values, size); } - pub fn checked_root(&self, order: usize) -> Option { - if order == 0 { - return Some(F::ONE); - } - self.order - .is_multiple_of(order) - .then(|| self.omega_order.pow([(self.order / order) as u64])) - } - - pub fn root(&self, order: usize) -> F { - self.checked_root(order) - .expect("Subgroup of requested order does not exist.") - } - /// Returns a cached table of roots of unity of the given order. fn roots_table(&self, order: usize) -> RwLockReadGuard<'_, Vec> { assert!( - self.order.is_multiple_of(order), + self.domain.order.is_multiple_of(order), "No subgroup of order {order}." ); @@ -208,7 +305,7 @@ impl NttEngine { roots.reserve_exact(size); // Compute powers of roots of unity. - let root = self.root(size); + let root = self.domain.root(size); #[cfg(not(feature = "parallel"))] { let mut root_i = F::ONE; @@ -360,47 +457,12 @@ impl NttEngine { } } -impl ReedSolomon for NttEngine { - fn next_order(&self, size: usize) -> Option { - match self.divisors.binary_search(&size) { - Ok(index) | Err(index) => self.divisors.get(index).copied(), - } - } - - fn generator(&self, codeword_length: usize) -> F { - self.omega_order - .pow([(self.order / codeword_length) as u64]) - } - - fn evaluation_points( - &self, - masked_message_length: usize, - codeword_length: usize, - indices: &[usize], - ) -> Vec { - assert!(masked_message_length <= codeword_length); - assert!(self.order.is_multiple_of(codeword_length)); - let mut result = Vec::new(); - let generator = self.generator(codeword_length); - - // Coset transformation - let mut coset_size = self.next_order(masked_message_length).unwrap(); - while !codeword_length.is_multiple_of(coset_size) { - coset_size = self.next_order(coset_size + 1).unwrap(); - } - #[cfg(not(feature = "rs_in_order"))] - let num_cosets = codeword_length / coset_size; - - for &index in indices { - assert!(index < codeword_length); - - #[cfg(not(feature = "rs_in_order"))] - let index = transpose_permute(index, num_cosets, coset_size); - result.push(generator.pow([index as u64])); - } - result - } - +impl NttEngine { + /// Masked interleaved Reed-Solomon encoding over CPU slices. + /// + /// `messages` are `num_messages` already-chunked slices of `message_length` elements. + /// Used directly by benchmarks and tests; backend encoders call this (CPU) or reimplement + /// it on the device (Metal), all driven by [`Self::coset_params`]. #[cfg_attr(feature = "tracing", instrument(skip(self, messages, masks), fields( num_messages = messages.len(), message_len = messages.first().map(|c| c.len()), @@ -408,8 +470,13 @@ impl ReedSolomon for NttEngine { mask_len = masks.len().checked_div(messages.len()) )))] - fn interleaved_encode(&self, messages: &[&[F]], masks: &[F], codeword_length: usize) -> Vec { - assert!(self.order.is_multiple_of(codeword_length)); + pub fn interleaved_encode_slices( + &self, + messages: &[&[F]], + masks: &[F], + codeword_length: usize, + ) -> Vec { + assert!(self.domain.order.is_multiple_of(codeword_length)); if messages.is_empty() { assert!(masks.is_empty()); return Vec::new(); @@ -433,11 +500,8 @@ impl ReedSolomon for NttEngine { // You can also see this as applying a first round of Cooley-Tukey with // N = coset_size × num_cosets, and solving it directly by observing that // only the first coset is non-zero. - let mut coset_size = self.next_order(masked_message_length).unwrap(); - while !codeword_length.is_multiple_of(coset_size) { - coset_size = self.next_order(coset_size + 1).unwrap(); - } - let num_cosets = codeword_length / coset_size; + let (coset_size, num_cosets) = + self.domain.coset_params(masked_message_length, codeword_length); let coset_padding = coset_size - masked_message_length; // Lay out twisted coefficients in contiguous coset blocks of length @@ -477,6 +541,69 @@ impl ReedSolomon for NttEngine { } } +/// CPU Reed-Solomon encoder. +/// +/// Thin wrapper over a CPU [`NttEngine`]: the scalar methods delegate to the engine's shared +/// [`RsDomain`], and `interleaved_encode` chunks the input vectors and runs the host coset-NTT. +/// +/// Only built when CPU buffers are the active backend. Under a Metal build +/// `ActiveBuffer = MetalBuffer`, so its `ReedSolomon` impl would not type-check; code that needs +/// a CPU codeword there reaches it through [`NttEngine::interleaved_encode_slices`] instead. +#[cfg(not(all(feature = "metal", target_os = "macos")))] +#[derive(Debug, Clone)] +pub struct CpuRs { + engine: Arc>, +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +impl CpuRs { + pub fn new(engine: Arc>) -> Self { + Self { engine } + } +} + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +impl ReedSolomon for CpuRs { + fn next_order(&self, size: usize) -> Option { + self.engine.domain().next_order(size) + } + + fn generator(&self, codeword_length: usize) -> F { + self.engine.domain().generator(codeword_length) + } + + fn evaluation_points( + &self, + masked_message_length: usize, + codeword_length: usize, + indices: &[usize], + ) -> Vec { + self.engine + .domain() + .evaluation_points(masked_message_length, codeword_length, indices) + } + + fn interleaved_encode( + &self, + vectors: &[&CpuBuffer], + masks: &CpuBuffer, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + ) -> CpuBuffer { + let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); + let messages = vectors + .iter() + .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) + .collect::>(); + CpuBuffer::from_vec(self.engine.interleaved_encode_slices( + &messages, + masks.as_slice(), + codeword_length, + )) + } +} + /// Applies twiddle factors to a slice of field elements in-place. /// /// This is part of the six-step Cooley-Tukey NTT algorithm, @@ -585,12 +712,12 @@ mod tests { let engine = NttEngine::::new_from_fftfield(); // Verify that the order of the engine is correctly set - assert_eq!(engine.order, 3 << 32); + assert_eq!(engine.domain.order, 3 << 32); // Verify that the root of unity is correctly initialized assert_eq!( - engine.omega_order, - Field64::GENERATOR.pow([(18_446_744_069_414_584_320 / engine.order) as u64]) + engine.domain.omega_order, + Field64::GENERATOR.pow([(18_446_744_069_414_584_320 / engine.domain.order) as u64]) ); } @@ -599,31 +726,31 @@ mod tests { let engine = NttEngine::::new_from_fftfield(); // Ensure the root exponentiates correctly - assert_eq!(engine.root(8).pow([8]), Field64::ONE); - assert_eq!(engine.root(4).pow([4]), Field64::ONE); - assert_eq!(engine.root(2).pow([2]), Field64::ONE); + assert_eq!(engine.domain().root(8).pow([8]), Field64::ONE); + assert_eq!(engine.domain().root(4).pow([4]), Field64::ONE); + assert_eq!(engine.domain().root(2).pow([2]), Field64::ONE); // Ensure it's not a lower-order root - assert_ne!(engine.root(8).pow([4]), Field64::ONE); - assert_ne!(engine.root(4).pow([2]), Field64::ONE); + assert_ne!(engine.domain().root(8).pow([4]), Field64::ONE); + assert_ne!(engine.domain().root(4).pow([2]), Field64::ONE); } #[test] fn test_root_of_unity_multiplication() { let engine = NttEngine::::new_from_fftfield(); - let root = engine.root(16); + let root = engine.domain().root(16); // Multiply root by itself repeatedly and verify expected outcomes - assert_eq!(root.pow([2]), engine.root(8)); - assert_eq!(root.pow([4]), engine.root(4)); - assert_eq!(root.pow([8]), engine.root(2)); + assert_eq!(root.pow([2]), engine.domain().root(8)); + assert_eq!(root.pow([4]), engine.domain().root(4)); + assert_eq!(root.pow([8]), engine.domain().root(2)); } #[test] fn test_root_of_unity_inversion() { let engine = NttEngine::::new_from_fftfield(); - let root = engine.root(16); + let root = engine.domain().root(16); // The inverse of ω is ω^{-1}, computed as ω^(p-2) in Field64. let p: u64 = u64::from_be_bytes(Field64::MODULUS.to_bytes_be().try_into().unwrap()); @@ -649,9 +776,9 @@ mod tests { let engine2 = NttEngine::::new_from_fftfield(); // Ensure that multiple instances yield the same results - assert_eq!(engine1.root(8), engine2.root(8)); - assert_eq!(engine1.root(4), engine2.root(4)); - assert_eq!(engine1.root(2), engine2.root(2)); + assert_eq!(engine1.domain().root(8), engine2.domain().root(8)); + assert_eq!(engine1.domain().root(4), engine2.domain().root(4)); + assert_eq!(engine1.domain().root(2), engine2.domain().root(2)); } #[test] @@ -661,9 +788,9 @@ mod tests { // Check hardcoded expected values (ω^i) assert_eq!(roots_4[0], Field::ONE); - assert_eq!(roots_4[1], engine.root(4)); - assert_eq!(roots_4[2], engine.root(4).pow([2])); - assert_eq!(roots_4[3], engine.root(4).pow([3])); + assert_eq!(roots_4[1], engine.domain().root(4)); + assert_eq!(roots_4[2], engine.domain().root(4).pow([2])); + assert_eq!(roots_4[3], engine.domain().root(4).pow([3])); } #[test] @@ -675,7 +802,7 @@ mod tests { // Must contain only ω^0 and ω^1 assert_eq!(roots_2.len(), 2); assert_eq!(roots_2[0], Field64::ONE); - assert_eq!(roots_2[1], engine.root(2)); + assert_eq!(roots_2[1], engine.domain().root(2)); } #[test] @@ -686,7 +813,7 @@ mod tests { // Ensure the sequence follows expected powers of the root of unity for i in 0..4 { - assert_eq!(roots_4[i], engine.root(4).pow([i as u64])); + assert_eq!(roots_4[i], engine.domain().root(4).pow([i as u64])); } } @@ -732,7 +859,7 @@ mod tests { let roots = vec![r1]; // Ensure the root of unity is correct - assert_eq!(engine.root(4).pow([4]), Field64::ONE); + assert_eq!(engine.domain().root(4).pow([4]), Field64::ONE); apply_twiddles(&mut values, &roots, 2, 4); @@ -1051,7 +1178,7 @@ mod tests { // // ω is the 32nd root of unity: ω³² = 1. - let omega = engine.root(32); + let omega = engine.domain().root(32); let mut expected_values = vec![Field64::ZERO; 32]; for (k, expected_value) in expected_values.iter_mut().enumerate().take(32) { let omega_k = omega.pow([k as u64]); diff --git a/src/algebra/ntt/mod.rs b/src/algebra/ntt/mod.rs index 9e3fc74d..f3ab5e08 100644 --- a/src/algebra/ntt/mod.rs +++ b/src/algebra/ntt/mod.rs @@ -11,46 +11,45 @@ use std::{ sync::{Arc, LazyLock}, }; -use ark_ff::Field; +use ark_ff::{Field, FftField}; use static_assertions::assert_obj_safe; use self::matrix::MatrixMut; +pub use self::cooley_tukey::{NttEngine, RsDomain}; +// The CPU encoder is only the active backend (and only constructible) on non-Metal builds. +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub use self::cooley_tukey::CpuRs; pub use self::{ - cooley_tukey::NttEngine, transpose::transpose, wavelet::{inverse_wavelet_transform, wavelet_transform}, }; use crate::{ - algebra::fields, + algebra::{buffer::ActiveBuffer, fields}, type_map::{self, TypeMap}, }; pub static NTT: LazyLock> = LazyLock::new(|| { let map = TypeMap::new(); - map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> - ); - map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> - ); - map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> - ); - map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> - ); - map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc>, - ); - map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc>, - ); - map.insert(Arc::new( - NttEngine::<::BasePrimeField>::new_from_fftfield(), - ) as Arc>); - map.insert(Arc::new( - NttEngine::<::BasePrimeField>::new_from_fftfield(), - ) as Arc>); + fn register(map: &TypeMap) { + // Both backends share the same `RsDomain` coset convention and differ only in how + // `interleaved_encode` runs: the CPU encoder needs the full NTT engine, while the + // Metal encoder only needs the shared domain (the NTT itself runs on the GPU). + #[cfg(all(feature = "metal", target_os = "macos"))] + let encoder = crate::algebra::metal_buffer::MetalRs::new(Arc::new( + RsDomain::::from_fftfield(), + )); + #[cfg(not(all(feature = "metal", target_os = "macos")))] + let encoder = CpuRs::new(Arc::new(NttEngine::::new_from_fftfield())); + map.insert(Arc::new(encoder) as Arc>); + } + register::(&map); + register::(&map); + register::(&map); + register::(&map); + register::(&map); + register::(&map); + register::<::BasePrimeField>(&map); + register::<::BasePrimeField>(&map); map }); @@ -62,6 +61,10 @@ impl type_map::Family for NttFamily { } /// Trait for a Reed-Solomon encoder implementation for a given field `F`. +/// +/// The scalar methods (`next_order`, `generator`, `evaluation_points`) describe the +/// field/domain structure and are backend independent. `interleaved_encode` produces a +/// codeword over the active backend buffer ([`ActiveBuffer`]). pub trait ReedSolomon: Debug + Send + Sync { /// Returns the next supported order equal or larger than `size`. /// @@ -88,15 +91,22 @@ pub trait ReedSolomon: Debug + Send + Sync { /// Compute a masked interleaved Reed-Solomon encoding. /// - /// `messages` are `num_messages` slices of `message_length` elements. - /// `masks` is a `num_messages` × `mask_length` matrix of blinding coefficients. + /// `vectors` are `num_vectors` buffers, each holding `interleaving_depth` messages of + /// `message_length` elements laid out contiguously. `masks` is a + /// `num_vectors * interleaving_depth` × `mask_length` matrix of blinding coefficients. /// `codeword_length` must be an NTT-smooth number >= `message_length + mask_length`. - /// returns an `codeword_length × num_messages` matrix. + /// Returns a `codeword_length × (num_vectors * interleaving_depth)` matrix. /// /// Each output value is the univariate polynomial evaluation in the evaluation point /// corresponding with the index of a coefficient list formed by concatenating message and mask. - /// - fn interleaved_encode(&self, messages: &[&[F]], masks: &[F], codeword_length: usize) -> Vec; + fn interleaved_encode( + &self, + vectors: &[&ActiveBuffer], + masks: &ActiveBuffer, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + ) -> ActiveBuffer; } assert_obj_safe!(ReedSolomon); @@ -117,16 +127,6 @@ pub fn evaluation_points( .evaluation_points(masked_message_length, codeword_length, indices) } -pub fn interleaved_rs_encode( - messages: &[&[F]], - masks: &[F], - codeword_length: usize, -) -> Vec { - NTT.get::() - .expect("Unsupported NTT field.") - .interleaved_encode(messages, masks, codeword_length) -} - pub fn generator(codeword_length: usize) -> F { NTT.get::() .expect("Unsupported NTT field.") @@ -148,14 +148,14 @@ mod tests { utils::{chunks_exact_or_empty, zip_strict}, }; - fn valid_codeword_lengths(size: usize, count: usize) -> Vec { - let ntt = NTT.get::().expect("No NTT engine for field."); - iter::successors(ntt.next_order(size), |size| ntt.next_order(*size + 1)) + fn valid_codeword_lengths(size: usize, count: usize) -> Vec { + let domain = RsDomain::::from_fftfield(); + iter::successors(domain.next_order(size), |size| domain.next_order(*size + 1)) .take(count) .collect() } - fn test(ntt: &dyn ReedSolomon) + fn test(engine: &NttEngine) where Standard: Distribution, { @@ -189,7 +189,7 @@ mod tests { .collect::>(); let masks = random_vector(&mut rng, mask_length * num_messages); let message_refs = messages.iter().map(|v| v.as_slice()).collect::>(); - let codeword = ntt.interleaved_encode( + let codeword = engine.interleaved_encode_slices( &message_refs, &masks, codeword_length, @@ -199,7 +199,7 @@ mod tests { assert_eq!(codeword.len(), codeword_length * num_messages); // Output values are polynomial evaluations in the evaluation points. - let mut evaluation_points = ntt.evaluation_points(message_length + mask_length, codeword_length, &sampled_indices); + let mut evaluation_points = engine.domain().evaluation_points(message_length + mask_length, codeword_length, &sampled_indices); for (&index, &evaluation_point) in zip_strict(&sampled_indices, &evaluation_points) { let evaluations = &codeword[index * num_messages.. (index + 1) * num_messages]; let masks = chunks_exact_or_empty(&masks, mask_length, num_messages); @@ -223,6 +223,7 @@ mod tests { #[test] fn test_field64_1() { - test::(NTT.get().unwrap().as_ref()); + let engine = NttEngine::::new_from_fftfield(); + test::(&engine); } } diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 8bc722f9..a30be48a 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -323,13 +323,16 @@ impl Config { prover_state.rng(), self.mask_length * self.num_messages(), ); - let matrix = ActiveBuffer::interleaved_rs_encode( - vectors, - &masks, - self.message_length(), - self.interleaving_depth, - self.codeword_length, - ); + let matrix = ntt::NTT + .get::() + .expect("Unsupported NTT field.") + .interleaved_encode( + vectors, + &masks, + self.message_length(), + self.interleaving_depth, + self.codeword_length, + ); // Commit to the matrix let matrix_witness = self.matrix_commit.commit(prover_state, &matrix); From bf57962eae7bb7ab9f99900496add7b5a83294ac Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Tue, 16 Jun 2026 16:42:12 +0200 Subject: [PATCH 14/35] feat: add merklize to BufferOps --- src/algebra/buffer.rs | 48 ++++++++++++++++++++++++++---- src/algebra/metal_buffer.rs | 54 +++++++++++++++++++++++++++++++++- src/protocols/matrix_commit.rs | 50 +++---------------------------- src/protocols/merkle_tree.rs | 22 ++++++-------- 4 files changed, 109 insertions(+), 65 deletions(-) diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index 5925db29..d03cc3e7 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -3,11 +3,16 @@ use std::{any::Any, mem}; use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; -use crate::algebra::{ - embedding::{Embedding, Identity}, - linear_form::{Covector, LinearForm, UnivariateEvaluation}, - mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, - sumcheck::{compute_sumcheck_polynomial, fold}, +use crate::{ + algebra::{ + embedding::{Embedding, Identity}, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, + mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, + sumcheck::{compute_sumcheck_polynomial, fold}, + }, + engines::EngineId, + hash::{self, Hash}, + protocols::{matrix_commit::Encodable, merkle_tree}, }; #[cfg(all(feature = "metal", target_os = "macos"))] @@ -31,6 +36,14 @@ pub trait BufferOps { fn is_empty(&self) -> bool { self.len() == 0 } + fn merklize( + &self, + num_cols: usize, + hash_engine_id: EngineId, + merkle_config: &merkle_tree::Config, + ) -> (ActiveBuffer, Hash) + where + T: Encodable + Send + Sync; } pub trait FieldOps: Clone { @@ -150,6 +163,31 @@ impl BufferOps for CpuBuffer { fn from_slice(source: &[T]) -> Self { Self::from_slice(source) } + + fn merklize( + &self, + num_cols: usize, + hash_engine_id: EngineId, + merkle_config: &merkle_tree::Config, + ) -> (ActiveBuffer, Hash) + where + T: Encodable + Send + Sync, + { + let _ = num_cols; // CPU leaf hashing derives the column count from the row count. + let engine = hash::ENGINES + .retrieve(hash_engine_id) + .expect("Failed to retrieve hash engine"); + #[cfg(feature = "tracing")] + tracing::Span::current().record("engine", engine.name().as_ref()); + + // Hash each row into a leaf, then build the full Merkle node array. + let mut leaves = vec![Hash::default(); merkle_config.num_leaves]; + crate::protocols::matrix_commit::hash_rows(&*engine, self.as_slice(), &mut leaves); + let nodes = merkle_config.build_nodes(leaves); + let root = nodes[nodes.len() - 1]; + + (ActiveBuffer::::from_vec(nodes), root) + } } impl FieldOps for CpuBuffer { diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs index c9cc2d91..c0bdb3b8 100644 --- a/src/algebra/metal_buffer.rs +++ b/src/algebra/metal_buffer.rs @@ -26,7 +26,12 @@ use crate::{ linear_form::{Covector, LinearForm, UnivariateEvaluation}, ntt::{ReedSolomon, RsDomain}, }, - hash::{metal_profile, Hash as Digest, MetalSha2}, + engines::EngineId, + hash::{self, metal_profile, Hash as Digest, MetalSha2, SHA2}, + protocols::{ + matrix_commit::{hash_rows, Encodable}, + merkle_tree, + }, }; const METAL_SOURCE: &str = r#" @@ -1187,6 +1192,53 @@ impl BufferOps for MetalBuffer { fn from_slice(source: &[T]) -> Self { Self::from_slice(source) } + + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (MetalBuffer, Digest) + where + T: Encodable + Send + Sync, + { + let num_rows = merkle.num_leaves; + let layers = merkle.layers.len(); + + // Fast path: build the whole tree on the GPU when the leaf and every node layer is SHA2. + if leaf_hash == SHA2 && merkle.layers.iter().all(|layer| layer.hash_id == SHA2) { + if let Some(nodes) = self.commit_bn254_rows_sha2_merkle(num_cols, num_rows, layers) { + let root = nodes + .read_hash_at(merkle.num_nodes() - 1) + .expect("missing Metal Merkle root"); + return (nodes, root); + } + } + + // CPU fallback: hash each row on the CPU, then build the tree. + let cpu_nodes = || { + let engine = hash::ENGINES + .retrieve(leaf_hash) + .expect("Failed to retrieve hash engine"); + let mut leaves = vec![Digest::default(); num_rows]; + hash_rows(&*engine, self.as_slice(), &mut leaves); + merkle.build_nodes(leaves) + }; + + // Otherwise hash leaves (on the GPU if SHA2, else CPU) and build the tree on the CPU. + let nodes = if leaf_hash == SHA2 { + let mut leaves = vec![Digest::default(); num_rows]; + if self.hash_bn254_rows_sha2(num_cols, &mut leaves) { + merkle.build_nodes(leaves) + } else { + cpu_nodes() + } + } else { + cpu_nodes() + }; + let root = nodes[merkle.num_nodes() - 1]; + (MetalBuffer::from_vec(nodes), root) + } } impl FieldOps for MetalBuffer { diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index ef13bc7b..0533e518 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -226,51 +226,9 @@ impl Config { { assert_eq!(matrix.len(), self.num_rows() * self.num_cols); - let engine = hash::ENGINES - .retrieve(self.leaf_hash_id) - .expect("Failed to retrieve hash engine"); - #[cfg(feature = "tracing")] - tracing::Span::current().record("engine", engine.name().as_ref()); - - #[cfg(all(feature = "metal", target_os = "macos"))] - if self.leaf_hash_id == hash::SHA2 - && self - .merkle_tree - .layers - .iter() - .all(|layer| layer.hash_id == hash::SHA2) - { - if let Some(nodes) = matrix.commit_bn254_rows_sha2_merkle( - self.num_cols, - self.num_rows(), - self.merkle_tree.layers.len(), - ) { - let root = nodes - .read_hash_at(self.merkle_tree.num_nodes() - 1) - .expect("missing Metal Merkle root"); - prover_state.prover_message(&root); - return BufferWitness { nodes }; - } - } - - // Compute leaf hashes - let mut leaves = Vec::with_capacity(self.merkle_tree.num_nodes()); - leaves.resize(self.merkle_tree.num_leaves, Hash::default()); - #[cfg(all(feature = "metal", target_os = "macos"))] - let hashed_on_metal = self.leaf_hash_id == hash::SHA2 - && matrix.hash_bn254_rows_sha2(self.num_cols, &mut leaves[..self.num_rows()]); - #[cfg(not(all(feature = "metal", target_os = "macos")))] - let hashed_on_metal = false; - if !hashed_on_metal { - hash_rows(&*engine, matrix.as_slice(), &mut leaves[..self.num_rows()]); - } - - // Commit the leaf hashes - BufferWitness { - nodes: ActiveBuffer::::from_vec( - self.merkle_tree.commit(prover_state, leaves).nodes, - ), - } + let (nodes, root) = matrix.merklize(self.num_cols, self.leaf_hash_id, &self.merkle_tree); + prover_state.prover_message(&root); + BufferWitness { nodes } } #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] @@ -412,7 +370,7 @@ fn hash_rows( } #[cfg(feature = "parallel")] -fn hash_rows( +pub fn hash_rows( engine: &dyn hash::HashEngine, matrix: &[T], out: &mut [Hash], diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index 3c1a7ea7..111b998d 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -78,13 +78,10 @@ impl Config { (1 << (self.layers.len() + 1)) - 1 } - #[cfg_attr(feature = "tracing", instrument(skip(prover_state, leaves), fields(self = %self)))] - pub fn commit(&self, prover_state: &mut ProverState, leaves: Vec) -> Witness - where - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - Hash: ProverMessage<[H::U]>, - { + /// Build the full node array from leaf hashes, without touching the transcript. + /// + /// The returned vector holds the leaf layer first and the root last (at `num_nodes() - 1`). + pub fn build_nodes(&self, leaves: Vec) -> Vec { assert_eq!( leaves.len(), self.num_leaves, @@ -120,10 +117,7 @@ impl Config { remaining = next_remaining; } - // Commit to the root hash. - prover_state.prover_message(&previous[0]); - - Witness { nodes } + nodes } pub fn receive_commitment( @@ -340,8 +334,10 @@ pub(crate) mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); - let tree = config.commit(&mut prover_state, leaves); - config.open(&mut prover_state, &tree, &[13, 42]); + let nodes = config.build_nodes(leaves); + prover_state.prover_message(&nodes[config.num_nodes() - 1]); + let witness = Witness { nodes }; + config.open(&mut prover_state, &witness, &[13, 42]); let proof = prover_state.proof(); // Verifier From 9f29bb2d42f07b161f4a0b836ee67eba079f656f Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Tue, 16 Jun 2026 17:20:47 +0200 Subject: [PATCH 15/35] refactor: change trait bound from Clone to Copy for Config implementation --- src/protocols/matrix_commit.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index 0533e518..bada92b6 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -174,7 +174,7 @@ impl Encoder for ZeroCopyEncoder { } } -impl Config { +impl Config { /// Create a new matrix commit configuration with the recommended hash function. pub fn new(num_rows: usize, num_cols: usize) -> Self { // Select a leaf hash function. @@ -354,7 +354,7 @@ fn metal_merkle_opening_hints( nodes.read_hash_indices(&node_indices) } -impl fmt::Display for Config { +impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "MatrixCommit({} x {})", self.num_rows(), self.num_cols) } @@ -468,7 +468,7 @@ pub(crate) mod tests { num_cols: usize, indices: &[usize], ) where - T: Clone + TypeInfo + Encodable + Send + Sync, + T: Copy + TypeInfo + Encodable + Send + Sync, Standard: Distribution, { crate::tests::init(); @@ -511,7 +511,7 @@ pub(crate) mod tests { fn proptest() where - T: Clone + TypeInfo + Encodable + Send + Sync, + T: Copy + TypeInfo + Encodable + Send + Sync, Standard: Distribution, { let hashes = [hash::COPY, hash::SHA2, hash::SHA3, hash::BLAKE3]; From b55662dea7b48c34846e83002ca613b94a57ab8b Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 17 Jun 2026 09:26:12 +0200 Subject: [PATCH 16/35] feat: use ActiveBuffer in basecase --- src/protocols/basecase.rs | 67 ++++++++++++++++++--------------------- 1 file changed, 31 insertions(+), 36 deletions(-) diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 77eeeabb..ebe613c8 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -1,4 +1,5 @@ -// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. +// NOTE ON BUFFER ABSTRACTION: since verifier is not succinct, full vector readbacks are forced +// consider alternative approaches here //! Base Case Linear Opening Protocol //! @@ -13,8 +14,9 @@ use spongefish::{Decoding, VerificationResult}; use crate::{ algebra::{ - buffer::ActiveBuffer, dot, embedding::Identity, multilinear_extend, random_vector, - scalar_mul_add_new, univariate_evaluate, + buffer::{ActiveBuffer, BufferOps, FieldOps}, + embedding::Identity, + multilinear_extend, scalar_mul_add_new, univariate_evaluate, }, hash::Hash, protocols::{irs_commit, sumcheck}, @@ -51,9 +53,9 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - vector: Vec, + mut vector: ActiveBuffer, witness: &irs_commit::Witness, - covector: Vec, + mut covector: ActiveBuffer, mut sum: F, ) -> Opening where @@ -70,7 +72,7 @@ impl Config { assert_eq!(self.commit.num_vectors, 1); assert_eq!(self.commit.vector_size, self.sumcheck.initial_size); assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size)); - debug_assert_eq!(dot(&vector, &covector), sum); + debug_assert_eq!(vector.dot(&covector), sum); if self.size() == 0 { return Opening { evaluation_points: Vec::new(), @@ -80,44 +82,40 @@ impl Config { // Even more trivial non-zk protocol: send f and r directly. if !self.masked { - prover_state.prover_messages(&vector); + // TODO: avoid these big readbacks even though they are transcript-sized + prover_state.prover_messages(&*vector.as_slice()); prover_state.prover_messages(witness.masks.as_slice()); let _ = self.commit.open(prover_state, &[witness]); - let mut vector_buffer = ActiveBuffer::from_vec(vector); - let mut covector_buffer = ActiveBuffer::from_vec(covector); let point = self .sumcheck - .prove( - prover_state, - &mut vector_buffer, - &mut covector_buffer, - &mut sum, - &[], - ) + .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) .round_challenges; - assert!(!vector_buffer.as_slice()[0].is_zero(), "Proof failed"); + assert!( + !vector.at_index(0).unwrap_or(F::one()).is_zero(), + "Proof failed" + ); return Opening { evaluation_points: point, - linear_form_evaluation: covector_buffer.as_slice()[0], + linear_form_evaluation: covector.at_index(0).expect("Proof failed"), }; } // Create masking vector. - let mask = random_vector(prover_state.rng(), vector.len()); - let mask_buffer = ActiveBuffer::from_slice(&mask); + let mask = ActiveBuffer::random(prover_state.rng(), vector.len()); // Commit to the masking vector. - let mask_witness = self.commit.commit(prover_state, &[&mask_buffer]); + let mask_witness = self.commit.commit(prover_state, &[&mask]); // Compute and send linear form of mask (μ' in paper). - let mask_sum = dot(&mask, &covector); + let mask_sum = mask.dot(&covector); prover_state.prover_message(&mask_sum); // RLC the mask with the vector let mask_rlc = prover_state.verifier_message::(); assert!(!mask_rlc.is_zero(), "Proof failed"); - let masked_vector = scalar_mul_add_new(&mask, mask_rlc, &vector); - prover_state.prover_messages(&masked_vector); + let mut masked_vector = mask.clone(); + vector.mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_vector, mask_rlc); + prover_state.prover_messages(&*masked_vector.as_slice()); // Send combined IRS randomness. (r^* in paper) let masked_masks = scalar_mul_add_new( @@ -132,14 +130,12 @@ impl Config { // Run sumcheck to reduce linear form claim let mut masked_sum = mask_sum + mask_rlc * sum; - let mut masked_vector_buffer = ActiveBuffer::from_vec(masked_vector); - let mut covector_buffer = ActiveBuffer::from_vec(covector); let point = self .sumcheck .prove( prover_state, - &mut masked_vector_buffer, - &mut covector_buffer, + &mut masked_vector, + &mut covector, &mut masked_sum, &[], ) @@ -150,13 +146,13 @@ impl Config { // no constraints on l(r) that the verifier can return. // This event is cryptographically unlikely as `F` is challenge sized. assert!( - !masked_vector_buffer.as_slice()[0].is_zero(), + !masked_vector.at_index(0).unwrap_or(F::one()).is_zero(), "Proof failed" ); Opening { evaluation_points: point, - linear_form_evaluation: covector_buffer.as_slice()[0], + linear_form_evaluation: covector.at_index(0).expect("Proof failed"), } } @@ -298,14 +294,13 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut rng = StdRng::seed_from_u64(seed); - let vector = random_vector(&mut rng, config.size()); - let covector = random_vector(&mut rng, config.size()); - let sum = dot(&vector, &covector); + let vector = ActiveBuffer::random(&mut rng, config.size()); + let covector = ActiveBuffer::random(&mut rng, config.size()); + let sum = vector.dot(&covector); // Prover let mut prover_state = ProverState::new_std(&ds); - let vector_buffer = ActiveBuffer::from_slice(&vector); - let witness = config.commit.commit(&mut prover_state, &[&vector_buffer]); + let witness = config.commit.commit(&mut prover_state, &[&vector]); let prover_result = config.prove( &mut prover_state, vector.clone(), @@ -314,7 +309,7 @@ mod tests { sum, ); assert_eq!( - multilinear_extend(&covector, &prover_result.evaluation_points), + covector.mixed_extend(&Identity::::new(), &prover_result.evaluation_points), prover_result.linear_form_evaluation ); let proof = prover_state.proof(); From 9b8a9b95aee72f8d3faedd1073392072c6bf3d70 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 17 Jun 2026 09:54:19 +0200 Subject: [PATCH 17/35] feat: use ActiveBuffer in code switch --- src/protocols/code_switch.rs | 132 +++++++++++++++++++++-------------- 1 file changed, 78 insertions(+), 54 deletions(-) diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index ec457cb1..4c22c79b 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -15,10 +15,12 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::ActiveBuffer, + buffer::{ActiveBuffer, BufferOps, FieldOps, WindowOps}, dot, embedding::{Embedding, Identity}, - eq_weights, geometric_accumulate, lift, mixed_dot, scalar_mul, univariate_evaluate, + eq_weights, lift, + linear_form::UnivariateEvaluation, + mixed_dot, }, hash::Hash, protocols::{ @@ -46,7 +48,7 @@ pub struct Config { #[must_use] #[derive(Clone, Debug)] pub struct Witness { - pub message: Vec, + pub message: ActiveBuffer, pub target_witness: IrsWitness, } @@ -55,9 +57,18 @@ pub type Commitment = IrsCommitment; /// Mask input for the code-switch prover. // TODO : This may be removed after parameter selection PR -pub enum MaskInput<'a, F> { +pub enum MaskInput { Disabled, - Enabled(&'a [F]), + Enabled(ActiveBuffer), +} + +// helper function to get UnivariateEvaluations +#[inline] +fn evaluators(points: &[F], size: usize) -> Vec> { + points + .iter() + .map(|&point| UnivariateEvaluation::new(point, size)) + .collect() } impl Config { @@ -156,11 +167,11 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - message: Vec, + message: ActiveBuffer, witness: &IrsWitness, - covector: &mut [M::Target], - folding_randomness: &[M::Target], - mask_input: &MaskInput<'_, M::Target>, + covector: &mut ActiveBuffer, + folding_randomness: &mut ActiveBuffer, + mask_input: &MaskInput, ) -> Witness where H: DuplexSpongeInterface, @@ -179,7 +190,7 @@ impl Config { folding_randomness.len(), self.source.interleaving_depth, ); - let mask_msg: Option<&[M::Target]> = match &mask_input { + let mask_msg: Option<&ActiveBuffer> = match &mask_input { MaskInput::Disabled => { assert_eq!( self.message_mask_length, 0, @@ -198,17 +209,17 @@ impl Config { }; // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 - let message_buffer = ActiveBuffer::from_slice(&message); - let target_witness = self.target.commit(prover_state, &[&message_buffer]); + let target_witness = self.target.commit(prover_state, &[&message]); // Step 2-3: OOD challenge + answers — Construction 9.7 Steps 2-3, p.55 // y := ze_ood(ρ) · [f; r; s] = f(α) + α^ℓ · (r,s)(α) let ood_points: Vec = prover_state.verifier_message_vec(self.out_domain_samples); let msg_len = message.len(); for &point in &ood_points { - let f_eval = univariate_evaluate(&message, point); + let f_eval = message.mixed_univariate_evaluate(&Identity::::new(), point); if let Some(mask) = mask_msg { - let mask_eval = univariate_evaluate(mask, point); + let mask_eval = + mask.mixed_univariate_evaluate(&Identity::::new(), point); let shift = point.pow([msg_len as u64]); prover_state.prover_message(&(f_eval + shift * mask_eval)); } else { @@ -227,26 +238,25 @@ impl Config { let (&original_sl_coeff, constraint_rlc_coeffs) = batching_coeffs.split_first().unwrap(); let (ood_rlc_coeffs, in_domain_rlc_coeffs) = constraint_rlc_coeffs.split_at(num_ood); - // Covector update — sl' from Completeness proof (p.55-56) + // Covector update — sl' from Completeness proof (p.55-56). + covector.scale(original_sl_coeff); let eval_points = lift(self.source.embedding(), &source_evaluations.points); - scalar_mul(covector, original_sl_coeff); if self.message_mask_length == 0 { - // Non-ZK: single accumulate over all points - let all_points: Vec<_> = ood_points.iter().chain(&eval_points).copied().collect(); - let pows: Vec<_> = ood_rlc_coeffs - .iter() - .chain(in_domain_rlc_coeffs) - .copied() - .collect(); - geometric_accumulate(covector, pows, &all_points); + // Non-ZK: single accumulate over all points. Scalars are the + // contiguous challenge tail (ood ‖ in-domain), so no copy is needed. + let mut all_evaluators = evaluators(&ood_points, covector.len()); + all_evaluators.extend(evaluators(&eval_points, covector.len())); + covector.accumulate_univariate_evaluations(&all_evaluators, constraint_rlc_coeffs); } else { - // ZK: OOD contributes to full [f; r; s], in-domain only to [f; r] - geometric_accumulate(covector, ood_rlc_coeffs.to_vec(), &ood_points); - geometric_accumulate( - &mut covector[..self.source.masked_message_length()], - in_domain_rlc_coeffs.to_vec(), - &eval_points, - ); + // ZK: OOD contributes to full [f; r; s], in-domain only to [f; r]. + let ood_evaluators = evaluators(&ood_points, covector.len()); + covector.accumulate_univariate_evaluations(&ood_evaluators, ood_rlc_coeffs); + + let masked = self.source.masked_message_length(); + let in_domain_evaluators = evaluators(&eval_points, masked); + covector + .window_mut(0, masked) + .accumulate_univariate_evaluations(&in_domain_evaluators, in_domain_rlc_coeffs); } Witness { @@ -324,8 +334,8 @@ impl Config { let num_ood = self.out_domain_samples; let num_in_domain = source_evaluations.points.len(); let coeffs = geometric_challenge(verifier_state, 1 + num_ood + num_in_domain); - let (&original_sl_coeff, all_rlc_coeffs) = coeffs.split_first().unwrap(); - let (ood_rlc_coeffs, in_domain_rlc_coeffs) = all_rlc_coeffs.split_at(num_ood); + let (&original_sl_coeff, constraint_rlc_coeffs) = coeffs.split_first().unwrap(); + let (ood_rlc_coeffs, in_domain_rlc_coeffs) = constraint_rlc_coeffs.split_at(num_ood); *sum = original_sl_coeff * *sum + dot(ood_rlc_coeffs, &ood_answers) @@ -498,7 +508,7 @@ mod tests { mask } - fn mask_input(mask_msg: &[F]) -> MaskInput<'_, F> { + fn mask_input(mask_msg: ActiveBuffer) -> MaskInput { if mask_msg.is_empty() { MaskInput::Disabled } else { @@ -519,6 +529,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let covector = ActiveBuffer::from_vec(covector); let instance = U64(seed); let ds = DomainSeparator::protocol(config) @@ -531,17 +542,20 @@ mod tests { // Sample γ for sumcheck folding (length log2(ι)). let folding_randomness = sample_folding_randomness(config, &mut rng); // Post-fold message Fold(f_full, γ) of length message_length. - let folded_message = - fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + let folded_message = ActiveBuffer::from_vec(fold_chunks( + &f_full, + config.source.message_length(), + &folding_randomness, + )); let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); let witness = config.prove( &mut prover_state, folded_message.clone(), &source_witness, - &mut covector, - &folding_randomness, - &mask_input(&mask_msg), + &mut covector.clone(), + &mut ActiveBuffer::from_vec(folding_randomness.clone()), + &mask_input(ActiveBuffer::from_vec(mask_msg)), ); let proof = prover_state.proof(); @@ -573,6 +587,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let mut covector = ActiveBuffer::from_vec(covector); let instance = U64(seed); let ds = DomainSeparator::protocol(config) @@ -583,32 +598,38 @@ mod tests { let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); - let folded_message = - fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + let folded_message = ActiveBuffer::from_vec(fold_chunks( + &f_full, + config.source.message_length(), + &folding_randomness, + )); let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); // h is the post-fold polynomial whose inner product with covector // should equal the verifier sum: // - non-ZK: h = folded_message (length message_length) // - ZK: h = [folded_message; mask_msg] (length message_length + l_zk) - let h: Vec = if mask_msg.is_empty() { + let h: ActiveBuffer = if mask_msg.is_empty() { folded_message.clone() } else { - folded_message - .iter() - .chain(mask_msg.iter()) - .copied() - .collect() + ActiveBuffer::from_vec( + folded_message + .as_slice() + .iter() + .chain(mask_msg.iter()) + .copied() + .collect(), + ) }; - let initial_mu = dot(&h, &covector); + let initial_mu = h.dot(&covector); let _witness = config.prove( &mut prover_state, folded_message, &source_witness, &mut covector, - &folding_randomness, - &mask_input(&mask_msg), + &mut ActiveBuffer::from_vec(folding_randomness.clone()), + &mask_input(ActiveBuffer::from_vec(mask_msg)), ); let proof = prover_state.proof(); @@ -628,7 +649,7 @@ mod tests { .unwrap(); verifier_state.check_eof().unwrap(); - assert_eq!(dot(&h, &covector), verifier_sum); + assert_eq!(h.dot(&covector), verifier_sum); } fn test_tampered_ood_config>(seed: u64, config: &Config>) @@ -645,6 +666,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let mut covector = ActiveBuffer::from_vec(covector); // Commit honest f_full, fold to get the honest post-fold message. let mut prover_state = ProverState::new_std(&ds); @@ -655,17 +677,19 @@ mod tests { fold_chunks(&f_full, config.source.message_length(), &folding_randomness); // For non-ZK and source.mask_length == 0, h = folded_message and identity holds. - let initial_mu = dot(&folded_message, &covector); + let folded_message_buffer = ActiveBuffer::from_vec(folded_message.clone()); + let initial_mu = folded_message_buffer.dot(&covector); // Tamper the post-fold message before proving. let mut tampered = folded_message.clone(); tampered[0] += F::ONE; + let tampered = ActiveBuffer::from_vec(tampered); let _witness = config.prove( &mut prover_state, tampered, &source_witness, &mut covector, - &folding_randomness, + &mut ActiveBuffer::from_vec(folding_randomness.clone()), &MaskInput::Disabled, ); let proof = prover_state.proof(); @@ -687,7 +711,7 @@ mod tests { verifier_state.check_eof().unwrap(); // Sum diverges — downstream sumcheck would reject - assert_ne!(dot(&folded_message, &covector), verifier_sum); + assert_ne!(folded_message_buffer.dot(&covector), verifier_sum); } fn test + 'static>() From 0f5f0dd75bee6e6d594622b1a06d1aed8c34f2f0 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 18 Jun 2026 09:55:24 +0200 Subject: [PATCH 18/35] feat: add BufferRead and BufferWrite traits --- benches/metal_buffer.rs | 2 +- src/algebra/buffer.rs | 583 +++++++++++++++++++--- src/algebra/metal_buffer.rs | 938 ++++++++++++++++++++++++++++------- src/protocols/basecase.rs | 2 +- src/protocols/code_switch.rs | 4 +- src/protocols/irs_commit.rs | 2 +- src/protocols/sumcheck.rs | 2 +- src/protocols/whir/prover.rs | 15 +- 8 files changed, 1297 insertions(+), 251 deletions(-) diff --git a/benches/metal_buffer.rs b/benches/metal_buffer.rs index d91b8f42..a3e85b65 100644 --- a/benches/metal_buffer.rs +++ b/benches/metal_buffer.rs @@ -3,7 +3,7 @@ mod bench { use divan::{black_box, AllocProfiler, Bencher}; use whir::{ algebra::{ - buffer::{CpuBuffer, FieldOps, MetalBuffer}, + buffer::{BufferRead, CpuBuffer, MetalBuffer}, fields::Field256 as F, }, hash::{Hash, HashEngine, MetalSha2, Sha2}, diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index d03cc3e7..329a684e 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -1,11 +1,12 @@ -use std::{any::Any, mem}; +use std::{any::Any, mem, ops::RangeBounds}; use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; use crate::{ algebra::{ - embedding::{Embedding, Identity}, + dot, + embedding::Embedding, linear_form::{Covector, LinearForm, UnivariateEvaluation}, mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, sumcheck::{compute_sumcheck_polynomial, fold}, @@ -16,13 +17,21 @@ use crate::{ }; #[cfg(all(feature = "metal", target_os = "macos"))] -pub use super::metal_buffer::MetalBuffer; +pub use super::metal_buffer::{MetalBuffer, MetalSlice, MetalSliceMut}; #[cfg(all(feature = "metal", target_os = "macos"))] pub type ActiveBuffer = MetalBuffer; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type ActiveSlice<'a, T> = MetalSlice<'a, T>; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type ActiveSliceMut<'a, T> = MetalSliceMut<'a, T>; #[cfg(not(all(feature = "metal", target_os = "macos")))] pub type ActiveBuffer = CpuBuffer; +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type ActiveSlice<'a, T> = CpuSlice<'a, T>; +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type ActiveSliceMut<'a, T> = CpuSliceMut<'a, T>; pub trait BufferOps { fn from_vec(source: Vec) -> Self; @@ -32,10 +41,12 @@ pub trait BufferOps { self.len() / num_cols } fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; + fn at_index(&self, index: usize) -> Option; fn len(&self) -> usize; fn is_empty(&self) -> bool { self.len() == 0 } + fn extend(&self, other: &Self) -> Self; fn merklize( &self, num_cols: usize, @@ -46,65 +57,351 @@ pub trait BufferOps { T: Encodable + Send + Sync; } -pub trait FieldOps: Clone { - type TargetBuffer: FieldOps; +/// Read-only operations over a contiguous run of field elements — the buffer +/// analogue of `&[F]`. +/// +/// Implemented by the owned buffers ([`CpuBuffer`]/`MetalBuffer`, as the +/// full-range case) and by the borrowed read views ([`CpuSlice`]/`MetalSlice`). +/// Every op here works identically on a whole buffer or any sub-range obtained +/// via [`Self::slice`]. On `MetalBuffer` a view aliases the parent's GPU +/// allocation (byte-offset binding), so no data is copied. +pub trait BufferRead { + /// A same-backend owned buffer over another field, produced by the + /// mixed-field ops (e.g. `mixed_dot` against a target-field buffer). For an + /// owned buffer this is the buffer itself over `T`; views report the owning + /// buffer type. + type TargetBuffer: Buffer; + + /// Read-only view type produced by slicing this buffer/view. + type Slice<'a>: BufferRead + where + Self: 'a, + F: 'a; + + /// Number of elements in this view. + fn read_len(&self) -> usize; + fn read_is_empty(&self) -> bool { + self.read_len() == 0 + } + + /// Inner product with another view of the same length. + fn dot(&self, other: &Self) -> F + where + Self: Sized; + + /// Sumcheck round coefficients `(c0, c2)` for `dot(self, other)`. + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) + where + Self: Sized; + + /// Multilinear extension evaluated at a target-field point. + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target; + + /// Inner product with a target-field buffer. + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &Self::TargetBuffer, + ) -> M::Target; + + /// Univariate evaluation at a target-field point. + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target; + + /// `accumulator += weight * self`, lifted into the target field. + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut Self::TargetBuffer, + weight: M::Target, + ); + + /// Borrow `self[range]` as a read-only view (analogous to `&v[range]`). + fn slice(&self, range: impl RangeBounds) -> Self::Slice<'_>; +} +/// In-place, length-preserving operations over a contiguous run of field +/// elements — the buffer analogue of `&mut [F]`. +/// +/// Implemented by the owned buffers (full-range) and the borrowed mutable views +/// ([`CpuSliceMut`]/`MetalSliceMut`). A sub-range obtained via [`Self::slice_mut`] +/// is the same view type as the whole buffer, so the ops compose just like +/// slicing a `Vec`. +/// +/// Constructors (`zeros`, `random`, …) and length-changing operations (`fold`, +/// `zero_pad`) live on [`Buffer`], the owned-buffer trait, because a borrowed +/// view cannot construct or resize its backing storage. +pub trait BufferWrite: BufferRead { + /// Mutable view type produced by slicing this buffer/view. + type SliceMut<'a>: BufferWrite + where + Self: 'a, + F: 'a; + + /// In-place scalar multiplication: `self[i] *= weight`. + fn scale(&mut self, weight: F); + + /// Accumulate `Σ_j scalars[j] · evaluators[j].point^i` into entry `i`. + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ); + + /// Borrow `self[range]` as a mutable, zero-copy view (analogous to + /// `&mut v[range]`). + fn slice_mut(&mut self, range: impl RangeBounds) -> Self::SliceMut<'_>; + + /// Split into two disjoint mutable views `[0, mid)` and `[mid, len)`. + fn split_at_mut(&mut self, mid: usize) -> (Self::SliceMut<'_>, Self::SliceMut<'_>); +} + +/// Owned field-buffer operations — the `Vec` half of the field API. +/// +/// `BufferOps` stays generic over any element type, so hashes and digests can +/// use it. `BufferRead`/`BufferWrite` are the slice-like field ops implemented +/// by owners and views. This trait is only for owned field buffers: operations +/// here construct storage or change its length, so borrowed views cannot +/// implement them. +pub trait Buffer: BufferOps + BufferWrite + Clone { fn zeros(length: usize) -> Self; - // consider removing this to avoid dependency on ark for Buffer design + /// Geometric sequence `[1, base, base², …, base^(length-1)]`. + fn geometric_sequence(base: F, length: usize) -> Self; + fn random(rng: &mut R, length: usize) -> Self where R: RngCore + CryptoRng, Standard: Distribution; + fn zero_pad(&mut self); - fn dot(&self, other: &Self) -> F; fn fold(&mut self, weight: F); + fn fold_pair(&mut self, other: &mut Self, weight: F) { self.fold(weight); other.fold(weight); } - fn sumcheck_polynomial(&self, other: &Self) -> (F, F); + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { self.fold_pair(other, weight); self.sumcheck_polynomial(other) } - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ); + fn linear_forms_rlc( size: usize, linear_forms: &mut [Box>], rlc_coeffs: &[F], ) -> Self; + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer; +} + +/// Resolve any range expression (`a..b`, `..b`, `a..`, `..`) against `len`, +/// returning `(start, end)` and bounds-checking like slice indexing. +pub(crate) fn resolve_range(range: impl RangeBounds, len: usize) -> (usize, usize) { + use std::ops::Bound::{Excluded, Included, Unbounded}; + let start = match range.start_bound() { + Included(&s) => s, + Excluded(&s) => s + 1, + Unbounded => 0, + }; + let end = match range.end_bound() { + Included(&e) => e + 1, + Excluded(&e) => e, + Unbounded => len, + }; + assert!( + start <= end && end <= len, + "slice range {start}..{end} out of bounds for length {len}" + ); + (start, end) +} + +/// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. +pub struct CpuSlice<'a, F> { + data: &'a [F], +} + +/// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. +pub struct CpuSliceMut<'a, F> { + data: &'a mut [F], +} + +impl CpuSlice<'_, F> { + pub fn len(&self) -> usize { + self.data.len() + } + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + pub fn as_slice(&self) -> &[F] { + self.data + } +} + +impl CpuSliceMut<'_, F> { + pub fn len(&self) -> usize { + self.data.len() + } + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + pub fn as_slice(&self) -> &[F] { + self.data + } +} + +impl BufferRead for CpuSlice<'_, F> { + type TargetBuffer = CpuBuffer; + type Slice<'a> + = CpuSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.data.len() + } + fn dot(&self, other: &Self) -> F { + crate::algebra::dot(self.data, other.data) + } + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + compute_sumcheck_polynomial(self.data, other.data) + } fn mixed_extend, T: Field>( &self, embedding: &M, point: &[M::Target], - ) -> M::Target; + ) -> M::Target { + mixed_multilinear_extend(embedding, self.data, point) + } fn mixed_dot, T: Field>( &self, embedding: &M, - other: &Self::TargetBuffer, - ) -> M::Target; + other: &CpuBuffer, + ) -> M::Target { + mixed_dot(embedding, other.as_slice(), self.data) + } fn mixed_univariate_evaluate>( &self, embedding: &M, point: M::Target, - ) -> M::Target; - fn mixed_linear_combination>( + ) -> M::Target { + mixed_univariate_evaluate(embedding, self.data, point) + } + fn mixed_scalar_mul_add_to>( + &self, embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer; + accumulator: &mut CpuBuffer, + weight: M::Target, + ) { + mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.data); + } + + fn slice(&self, range: impl RangeBounds) -> CpuSlice<'_, F> { + let (start, end) = resolve_range(range, self.data.len()); + CpuSlice { + data: &self.data[start..end], + } + } +} + +impl BufferRead for CpuSliceMut<'_, F> { + type TargetBuffer = CpuBuffer; + type Slice<'a> + = CpuSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.data.len() + } + fn dot(&self, other: &Self) -> F { + crate::algebra::dot(self.data, other.data) + } + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + compute_sumcheck_polynomial(self.data, other.data) + } + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + mixed_multilinear_extend(embedding, self.data, point) + } + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &CpuBuffer, + ) -> M::Target { + mixed_dot(embedding, other.as_slice(), self.data) + } + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + mixed_univariate_evaluate(embedding, self.data, point) + } fn mixed_scalar_mul_add_to>( &self, embedding: &M, - accumulator: &mut Self::TargetBuffer, + accumulator: &mut CpuBuffer, weight: M::Target, - ); + ) { + mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.data); + } + + fn slice(&self, range: impl RangeBounds) -> CpuSlice<'_, F> { + let (start, end) = resolve_range(range, self.data.len()); + CpuSlice { + data: &self.data[start..end], + } + } +} + +impl BufferWrite for CpuSliceMut<'_, F> { + type SliceMut<'a> + = CpuSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scale(&mut self, weight: F) { + crate::algebra::scalar_mul(self.data, weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + UnivariateEvaluation::accumulate_many(evaluators, self.data, scalars); + } + + fn slice_mut(&mut self, range: impl RangeBounds) -> CpuSliceMut<'_, F> { + let (start, end) = resolve_range(range, self.data.len()); + CpuSliceMut { + data: &mut self.data[start..end], + } + } + + fn split_at_mut(&mut self, mid: usize) -> (CpuSliceMut<'_, F>, CpuSliceMut<'_, F>) { + let (lo, hi) = self.data.split_at_mut(mid); + (CpuSliceMut { data: lo }, CpuSliceMut { data: hi }) + } } #[derive( @@ -139,7 +436,15 @@ impl CpuBuffer { } } -impl BufferOps for CpuBuffer { +impl BufferOps for CpuBuffer { + fn at_index(&self, index: usize) -> Option { + if index >= self.len() { + None + } else { + Some(self.data[index]) + } + } + fn as_slice(&self) -> &[T] { &self.data } @@ -156,6 +461,13 @@ impl BufferOps for CpuBuffer { result } + fn extend(&self, other: &Self) -> Self { + let mut data = Vec::with_capacity(self.data.len() + other.data.len()); + data.extend_from_slice(&self.data); + data.extend_from_slice(&other.data); + Self { data } + } + fn from_vec(source: Vec) -> Self { Self::from_vec(source) } @@ -190,15 +502,110 @@ impl BufferOps for CpuBuffer { } } -impl FieldOps for CpuBuffer { +impl BufferRead for CpuBuffer { type TargetBuffer = CpuBuffer; + type Slice<'a> + = CpuSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.data.len() + } + + fn dot(&self, other: &Self) -> F { + dot(&self.data, &other.data) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + compute_sumcheck_polynomial(&self.data, other.as_slice()) + } + + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + mixed_multilinear_extend(embedding, &self.data, point) + } + + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &CpuBuffer, + ) -> M::Target { + mixed_dot(embedding, other.as_slice(), self.as_slice()) + } + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + mixed_univariate_evaluate(embedding, &self.data, point) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut CpuBuffer, + weight: M::Target, + ) { + mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.as_slice()); + } + + fn slice(&self, range: impl RangeBounds) -> CpuSlice<'_, F> { + let (start, end) = resolve_range(range, self.data.len()); + CpuSlice { + data: &self.data[start..end], + } + } +} + +impl BufferWrite for CpuBuffer { + type SliceMut<'a> + = CpuSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scale(&mut self, weight: F) { + crate::algebra::scalar_mul(&mut self.data, weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + UnivariateEvaluation::accumulate_many(evaluators, &mut self.data, scalars); + } + fn slice_mut(&mut self, range: impl RangeBounds) -> CpuSliceMut<'_, F> { + let (start, end) = resolve_range(range, self.data.len()); + CpuSliceMut { + data: &mut self.data[start..end], + } + } + + fn split_at_mut(&mut self, mid: usize) -> (CpuSliceMut<'_, F>, CpuSliceMut<'_, F>) { + let (lo, hi) = self.data.split_at_mut(mid); + (CpuSliceMut { data: lo }, CpuSliceMut { data: hi }) + } +} + +impl Buffer for CpuBuffer { fn zeros(length: usize) -> Self { Self { data: vec![F::ZERO; length], } } + fn geometric_sequence(base: F, length: usize) -> Self { + Self::from_vec(crate::algebra::geometric_sequence(base, length)) + } + fn random(rng: &mut R, length: usize) -> Self where R: RngCore + CryptoRng, @@ -242,44 +649,14 @@ impl FieldOps for CpuBuffer { fold(&mut self.data, weight); } - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - compute_sumcheck_polynomial(&self.data, other.as_slice()) - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - UnivariateEvaluation::accumulate_many(evaluators, &mut self.data, scalars); - } - - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - mixed_multilinear_extend(embedding, &self.data, point) - } - - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &Self::TargetBuffer, - ) -> M::Target { - mixed_dot(embedding, other.as_slice(), self.as_slice()) - } - - fn dot(&self, other: &Self) -> F { - self.mixed_dot(&Identity::new(), other) + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); } - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - mixed_univariate_evaluate(embedding, &self.data, point) + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) } fn mixed_linear_combination>( @@ -298,13 +675,83 @@ impl FieldOps for CpuBuffer { } CpuBuffer::from_vec(accumulator) } +} - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut Self::TargetBuffer, - weight: M::Target, - ) { - mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.as_slice()); +#[cfg(test)] +mod tests { + use ark_ff::AdditiveGroup; + + use super::*; + use crate::algebra::{ + fields::Field64, geometric_accumulate, geometric_sequence as geometric_sequence_vec, + }; + + type F = Field64; + + #[test] + fn geometric_sequence_matches_free_function() { + let base = F::from(7u64); + for length in 0..6 { + let buffer = CpuBuffer::::geometric_sequence(base, length); + assert_eq!( + BufferOps::as_slice(&buffer), + geometric_sequence_vec(base, length).as_slice(), + "geometric_sequence mismatch at length {length}" + ); + } + } + + #[test] + fn scale_multiplies_in_place() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.scale(weight); + let expected: Vec = values.iter().map(|&v| v * weight).collect(); + assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); + } + + #[test] + fn slice_accumulate_matches_restricted_geometric_accumulate() { + let len = 8usize; + let points = vec![F::from(3u64), F::from(5u64)]; + let scalars = vec![F::from(2u64), F::from(9u64)]; + + // Try a prefix slice (offset 0) and an interior slice (offset > 0). + for (offset, window_len) in [(0usize, 5usize), (2usize, 4usize)] { + let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); + let evaluators: Vec<_> = points + .iter() + .map(|&point| UnivariateEvaluation::new(point, window_len)) + .collect(); + buffer + .slice_mut(offset..offset + window_len) + .accumulate_univariate_evaluations(&evaluators, &scalars); + + // Reference: accumulate Σ_j scalars[j]·points[j]^i into the same + // sub-slice of a plain vector. + let mut reference = vec![F::ZERO; len]; + geometric_accumulate( + &mut reference[offset..offset + window_len], + scalars.clone(), + &points, + ); + + assert_eq!( + BufferOps::as_slice(&buffer), + reference.as_slice(), + "slice accumulate mismatch at offset {offset}, len {window_len}" + ); + } + } + + #[test] + fn slice_scale_multiplies_only_the_range() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.slice_mut(1..3).scale(weight); + let expected = vec![values[0], values[1] * weight, values[2] * weight, values[3]]; + assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); } } diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs index c0bdb3b8..6a07de37 100644 --- a/src/algebra/metal_buffer.rs +++ b/src/algebra/metal_buffer.rs @@ -5,6 +5,7 @@ use std::{ collections::HashMap, hash::{Hash, Hasher}, marker::PhantomData, + ops::RangeBounds, os::raw::c_void, sync::{Arc, Mutex, OnceLock}, time::Instant, @@ -20,7 +21,7 @@ use zerocopy::IntoBytes; use crate::{ algebra::{ - buffer::{BufferOps, FieldOps}, + buffer::{BufferOps, BufferRead, BufferWrite}, embedding::{Embedding, Identity}, fields::{BN254Config, Field256}, linear_form::{Covector, LinearForm, UnivariateEvaluation}, @@ -354,6 +355,17 @@ kernel void bn254_scalar_mul_add( store_f(acc, gid, add_f(load_f(acc, gid), mul_f(weight, load_f(vector, gid)))); } +kernel void bn254_scalar_mul( + device ulong *acc [[buffer(0)]], + device const ulong *weight_buf [[buffer(1)]], + constant uint &len [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F weight = load_f(weight_buf, 0); + store_f(acc, gid, mul_f(weight, load_f(acc, gid))); +} + kernel void bn254_dot( device const ulong *a [[buffer(0)]], device const ulong *b [[buffer(1)]], @@ -1162,6 +1174,10 @@ impl BufferOps for MetalBuffer { self.as_slice() } + fn at_index(&self, index: usize) -> Option { + self.as_slice().get(index).cloned() + } + fn len(&self) -> usize { self.len } @@ -1193,6 +1209,13 @@ impl BufferOps for MetalBuffer { Self::from_slice(source) } + fn extend(&self, other: &Self) -> Self { + let mut data = Vec::with_capacity(self.len() + other.len()); + data.extend_from_slice(self.as_slice()); + data.extend_from_slice(other.as_slice()); + Self::from_vec(data) + } + fn merklize( &self, num_cols: usize, @@ -1241,9 +1264,7 @@ impl BufferOps for MetalBuffer { } } -impl FieldOps for MetalBuffer { - type TargetBuffer = MetalBuffer; - +impl crate::algebra::buffer::Buffer for MetalBuffer { fn zeros(length: usize) -> Self { assert_bn254::(); // Montgomery zero is all-zero bytes, so a device-side fill suffices. @@ -1256,6 +1277,11 @@ impl FieldOps for MetalBuffer { } } + fn geometric_sequence(base: F, length: usize) -> Self { + assert_bn254::(); + Self::from_vec(crate::algebra::geometric_sequence(base, length)) + } + fn random(rng: &mut R, length: usize) -> Self where R: RngCore + CryptoRng, @@ -1274,14 +1300,6 @@ impl FieldOps for MetalBuffer { } } - fn dot(&self, other: &Self) -> F { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - let this = self.bn254_buffer(); - let other = other.bn254_buffer(); - field256_to_f::(parallel_dot(&this, &other, self.len())) - } - fn fold(&mut self, weight: F) { assert_bn254::(); if self.len() <= 1 { @@ -1324,22 +1342,6 @@ impl FieldOps for MetalBuffer { other.invalidate_host_cache(); } - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - assert_bn254::(); - let len = self.len().min(other.len()); - if len == 0 { - return (F::ZERO, F::ZERO); - } - if len == 1 { - return (self.as_slice()[0] * other.as_slice()[0], F::ZERO); - } - let fold_half = len.next_power_of_two() >> 1; - let this = self.bn254_buffer(); - let other = other.bn254_buffer(); - let (c0, c2) = parallel_sumcheck(&this, &other, len, fold_half); - (field256_to_f::(c0), field256_to_f::(c2)) - } - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { assert_bn254::(); assert_eq!(self.len(), other.len()); @@ -1363,100 +1365,6 @@ impl FieldOps for MetalBuffer { (field256_to_f::(c0), field256_to_f::(c2)) } - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - assert_bn254::(); - assert_eq!(evaluators.len(), scalars.len()); - let Some(size) = evaluators.first().map(|e| e.size) else { - return; - }; - assert_eq!(self.len(), size); - for evaluator in evaluators { - assert_eq!(evaluator.size, size); - } - let points = evaluators - .iter() - .map(|e| f_to_field256(e.point)) - .collect::>(); - let scalars = scalars - .iter() - .copied() - .map(f_to_field256) - .collect::>(); - let points = upload_field(&points); - let scalars = upload_field(&scalars); - let field = self.bn254_buffer(); - let chunk_size = geometric_accumulate_chunk_size(self.len()); - if std::env::var_os("WHIR_METAL_TRACE").is_some() { - eprintln!( - "metal geometric shape len={} points={} chunk={} chunks={}", - self.len(), - evaluators.len(), - chunk_size, - self.len().div_ceil(chunk_size) - ); - } - if chunk_size <= 1 { - run_in_place( - "bn254_geometric_accumulate", - &[&field.limbs, &points.limbs, &scalars.limbs], - &[self.len() as u32, evaluators.len() as u32], - self.len(), - ); - } else { - let point_steps = evaluators - .iter() - .map(|e| f_to_field256(e.point.pow([chunk_size as u64]))) - .collect::>(); - let point_steps = upload_field(&point_steps); - if should_use_geometric_point_blocks(self.len(), evaluators.len(), chunk_size) { - parallel_geometric_accumulate_point_blocks( - &field, - &points, - &point_steps, - &scalars, - self.len(), - evaluators.len(), - chunk_size, - ); - } else if should_use_geometric_point_blocks_batched( - self.len(), - evaluators.len(), - chunk_size, - ) { - parallel_geometric_accumulate_point_blocks_batched( - &field, - &points, - &point_steps, - &scalars, - self.len(), - evaluators.len(), - chunk_size, - ); - } else { - run_in_place( - "bn254_geometric_accumulate_chunks_strided", - &[ - &field.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[ - self.len() as u32, - evaluators.len() as u32, - chunk_size as u32, - ], - self.len().div_ceil(chunk_size), - ); - } - } - self.invalidate_host_cache(); - } - fn linear_forms_rlc( size: usize, linear_forms: &mut [Box>], @@ -1481,6 +1389,164 @@ impl FieldOps for MetalBuffer { accumulator } + fn mixed_linear_combination>( + _embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer { + assert_bn254::(); + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return MetalBuffer::from_vec(Vec::new()); + }; + let mut accumulator = MetalBuffer:: { + len: first.len(), + host_cache: OnceCell::new(), + field: Some(copy_field_buffer(&first.bn254_buffer(), first.len())), + hash: None, + _marker: PhantomData, + }; + for (coeff, vector) in coeffs[1..].iter().copied().zip(vectors) { + vector.mixed_scalar_mul_add_to(_embedding, &mut accumulator, coeff); + } + accumulator + } +} + +/// Core geometric accumulate over `field[0..len]` (offset 0). Shared by the +/// owned buffer and a full-range view so the optimized chunk strategies are +/// written once. +fn geometric_accumulate_full( + field: &MetalFieldBuffer, + len: usize, + evaluators: &[UnivariateEvaluation], + scalars: &[F], +) { + let points = evaluators + .iter() + .map(|e| f_to_field256(e.point)) + .collect::>(); + let scalars = scalars + .iter() + .copied() + .map(f_to_field256) + .collect::>(); + let points = upload_field(&points); + let scalars = upload_field(&scalars); + let chunk_size = geometric_accumulate_chunk_size(len); + if std::env::var_os("WHIR_METAL_TRACE").is_some() { + eprintln!( + "metal geometric shape len={} points={} chunk={} chunks={}", + len, + evaluators.len(), + chunk_size, + len.div_ceil(chunk_size) + ); + } + if chunk_size <= 1 { + run_in_place( + "bn254_geometric_accumulate", + &[&field.limbs, &points.limbs, &scalars.limbs], + &[len as u32, evaluators.len() as u32], + len, + ); + } else { + let point_steps = evaluators + .iter() + .map(|e| f_to_field256(e.point.pow([chunk_size as u64]))) + .collect::>(); + let point_steps = upload_field(&point_steps); + if should_use_geometric_point_blocks(len, evaluators.len(), chunk_size) { + parallel_geometric_accumulate_point_blocks( + field, + &points, + &point_steps, + &scalars, + len, + evaluators.len(), + chunk_size, + ); + } else if should_use_geometric_point_blocks_batched(len, evaluators.len(), chunk_size) { + parallel_geometric_accumulate_point_blocks_batched( + field, + &points, + &point_steps, + &scalars, + len, + evaluators.len(), + chunk_size, + ); + } else { + run_in_place( + "bn254_geometric_accumulate_chunks_strided", + &[ + &field.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[len as u32, evaluators.len() as u32, chunk_size as u32], + len.div_ceil(chunk_size), + ); + } + } +} + +/// Validate that every evaluator targets `len` entries; returns `false` when +/// there is nothing to accumulate. +fn check_univariate_evaluators( + evaluators: &[UnivariateEvaluation], + scalars: &[F], + len: usize, +) -> bool { + assert_bn254::(); + assert_eq!(evaluators.len(), scalars.len()); + let Some(size) = evaluators.first().map(|e| e.size) else { + return false; + }; + assert_eq!(len, size); + for evaluator in evaluators { + assert_eq!(evaluator.size, size); + } + true +} + +impl BufferRead for MetalBuffer { + type TargetBuffer = MetalBuffer; + type Slice<'a> + = MetalSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.len() + } + + fn dot(&self, other: &Self) -> F { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + let this = self.bn254_buffer(); + let other = other.bn254_buffer(); + field256_to_f::(parallel_dot(&this, &other, self.len())) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + assert_bn254::(); + let len = self.len().min(other.len()); + if len == 0 { + return (F::ZERO, F::ZERO); + } + if len == 1 { + return (self.as_slice()[0] * other.as_slice()[0], F::ZERO); + } + let fold_half = len.next_power_of_two() >> 1; + let this = self.bn254_buffer(); + let other = other.bn254_buffer(); + let (c0, c2) = parallel_sumcheck(&this, &other, len, fold_half); + (field256_to_f::(c0), field256_to_f::(c2)) + } + fn mixed_extend, T: Field>( &self, _embedding: &M, @@ -1496,19 +1562,14 @@ impl FieldOps for MetalBuffer { .collect::>(); let point = upload_field(&point); let this = self.bn254_buffer(); - let out = run_reduce( - "bn254_multilinear_extend", - &[&this.limbs, &point.limbs], - &[self.len() as u32, num_vars as u32], - 1, - ); - field256_to_target::(out[0]) + let value = parallel_multilinear_extend_at(&this, 0, self.len(), &point, num_vars); + field256_to_target::(value) } fn mixed_dot, T: Field>( &self, _embedding: &M, - other: &Self::TargetBuffer, + other: &MetalBuffer, ) -> M::Target { assert_bn254::(); assert_bn254::(); @@ -1530,46 +1591,91 @@ impl FieldOps for MetalBuffer { field256_to_target::(parallel_univariate_evaluate(&this, &point, self.len())) } - fn mixed_linear_combination>( - _embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer { - assert_bn254::(); - assert_eq!(vectors.len(), coeffs.len()); - let Some((first, vectors)) = vectors.split_first() else { - return MetalBuffer::from_vec(Vec::new()); - }; - let mut accumulator = MetalBuffer:: { - len: first.len(), - host_cache: OnceCell::new(), - field: Some(copy_field_buffer(&first.bn254_buffer(), first.len())), - hash: None, - _marker: PhantomData, - }; - for (coeff, vector) in coeffs[1..].iter().copied().zip(vectors) { - vector.mixed_scalar_mul_add_to(_embedding, &mut accumulator, coeff); - } - accumulator - } - fn mixed_scalar_mul_add_to>( &self, _embedding: &M, - accumulator: &mut Self::TargetBuffer, + accumulator: &mut MetalBuffer, weight: M::Target, ) { assert_bn254::(); let weight = upload_field(&[target_to_field256(weight)]); let vector = self.bn254_buffer(); let acc = accumulator.bn254_buffer_target(); + scalar_mul_add_at(&acc, 0, &vector, 0, &weight, self.len()); + accumulator.field = Some(acc); + accumulator.invalidate_host_cache(); + } + + fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { + assert_bn254::(); + let (start, end) = crate::algebra::buffer::resolve_range(range, self.len()); + MetalSlice { + field: self.bn254_buffer(), + offset: start, + len: end - start, + _parent: PhantomData, + } + } +} + +impl BufferWrite for MetalBuffer { + type SliceMut<'a> + = MetalSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scale(&mut self, weight: F) { + assert_bn254::(); + if self.is_empty() { + return; + } + let weight = upload_field(&[f_to_field256(weight)]); + let field = self.bn254_buffer(); run_in_place( - "bn254_scalar_mul_add", - &[&acc.limbs, &vector.limbs, &weight.limbs], + "bn254_scalar_mul", + &[&field.limbs, &weight.limbs], &[self.len() as u32], self.len(), ); - accumulator.invalidate_host_cache(); + self.field = Some(field); + self.invalidate_host_cache(); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + if !check_univariate_evaluators(evaluators, scalars, self.len()) { + return; + } + let field = self.bn254_buffer(); + geometric_accumulate_full(&field, self.len(), evaluators, scalars); + self.field = Some(field); + self.invalidate_host_cache(); + } + + fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { + assert_bn254::(); + let (start, end) = crate::algebra::buffer::resolve_range(range, self.len()); + MetalSliceMut::new(self, start, end - start) + } + + fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { + assert_bn254::(); + let len = self.len(); + assert!(mid <= len, "split_at_mut mid {mid} out of bounds for {len}"); + // Materialize the parent's GPU allocation and invalidate its host + // cache once up front; the returned views share the handle and write + // disjoint ranges, so no per-op parent borrow is needed. + let field = self.bn254_buffer(); + self.field = Some(field.clone()); + self.invalidate_host_cache(); + ( + MetalSliceMut::from_field(field.clone(), 0, mid), + MetalSliceMut::from_field(field, mid, len - mid), + ) } } @@ -1684,6 +1790,338 @@ impl MetalBuffer { } } +/// Read-only, zero-copy view into a [`MetalBuffer`]'s GPU allocation. +/// +/// Holds a clone of the parent's allocation handle plus an `(offset, len)` +/// window. Reductions bind the handle at the byte offset, so they read +/// `parent[offset + i]` without copying. +pub struct MetalSlice<'a, F> { + field: MetalFieldBuffer, + offset: usize, + len: usize, + _parent: PhantomData<&'a MetalBuffer>, +} + +/// Mutable, zero-copy view into a [`MetalBuffer`]'s GPU allocation. +/// +/// Like [`MetalSlice`] but exclusive: the parent's host cache is invalidated +/// up front when the view is created, and writes go through the shared handle +/// at the byte offset, landing in the parent's own memory. +pub struct MetalSliceMut<'a, F> { + field: MetalFieldBuffer, + offset: usize, + len: usize, + _parent: PhantomData<&'a mut MetalBuffer>, +} + +impl<'a, F: Field + Clone> MetalSliceMut<'a, F> { + /// Borrow `parent[offset..offset+len]`, materializing the parent's GPU + /// allocation and invalidating its host cache once up front. + fn new(parent: &'a mut MetalBuffer, offset: usize, len: usize) -> Self { + let field = parent.bn254_buffer(); + parent.field = Some(field.clone()); + parent.invalidate_host_cache(); + Self { + field, + offset, + len, + _parent: PhantomData, + } + } + + /// Build a view directly from an already-materialized handle (used by + /// `split_at_mut`, where the parent cache was invalidated by the caller). + fn from_field(field: MetalFieldBuffer, offset: usize, len: usize) -> Self { + Self { + field, + offset, + len, + _parent: PhantomData, + } + } + + fn as_read(&self) -> MetalSlice<'_, F> { + MetalSlice { + field: self.field.clone(), + offset: self.offset, + len: self.len, + _parent: PhantomData, + } + } +} + +impl MetalSlice<'_, F> { + pub fn len(&self) -> usize { + self.len + } + pub fn is_empty(&self) -> bool { + self.len == 0 + } +} + +impl MetalSliceMut<'_, F> { + pub fn len(&self) -> usize { + self.len + } + pub fn is_empty(&self) -> bool { + self.len == 0 + } +} + +impl BufferRead for MetalSlice<'_, F> { + type TargetBuffer = MetalBuffer; + type Slice<'a> + = MetalSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.len + } + + fn dot(&self, other: &Self) -> F { + assert_bn254::(); + assert_eq!(self.len, other.len); + field256_to_f::(parallel_dot_at( + &self.field, + self.offset, + &other.field, + other.offset, + self.len, + )) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + assert_bn254::(); + let len = self.len.min(other.len); + if len == 0 { + return (F::ZERO, F::ZERO); + } + let fold_half = len.next_power_of_two() >> 1; + let (c0, c2) = parallel_sumcheck_at( + &self.field, + self.offset, + &other.field, + other.offset, + len, + fold_half, + ); + (field256_to_f::(c0), field256_to_f::(c2)) + } + + fn mixed_extend, T: Field>( + &self, + _embedding: &M, + point: &[M::Target], + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let num_vars = point.len(); + let point = point + .iter() + .copied() + .map(target_to_field256) + .collect::>(); + let point = upload_field(&point); + let value = + parallel_multilinear_extend_at(&self.field, self.offset, self.len, &point, num_vars); + field256_to_target::(value) + } + + fn mixed_dot, T: Field>( + &self, + _embedding: &M, + other: &MetalBuffer, + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let other = other.bn254_buffer_target(); + let value = field256_to_f::(parallel_dot_at( + &self.field, + self.offset, + &other, + 0, + self.len, + )); + field256_to_target::(f_to_field256(value)) + } + + fn mixed_univariate_evaluate>( + &self, + _embedding: &M, + point: M::Target, + ) -> M::Target { + assert_bn254::(); + let point = upload_field(&[target_to_field256(point)]); + field256_to_target::(parallel_univariate_evaluate_at( + &self.field, + self.offset, + &point, + self.len, + )) + } + + fn mixed_scalar_mul_add_to>( + &self, + _embedding: &M, + accumulator: &mut MetalBuffer, + weight: M::Target, + ) { + assert_bn254::(); + let weight = upload_field(&[target_to_field256(weight)]); + let acc = accumulator.bn254_buffer_target(); + scalar_mul_add_at(&acc, 0, &self.field, self.offset, &weight, self.len); + accumulator.field = Some(acc); + accumulator.invalidate_host_cache(); + } + + fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { + let (start, end) = crate::algebra::buffer::resolve_range(range, self.len); + MetalSlice { + field: self.field.clone(), + offset: self.offset + start, + len: end - start, + _parent: PhantomData, + } + } +} + +impl BufferRead for MetalSliceMut<'_, F> { + type TargetBuffer = MetalBuffer; + type Slice<'a> + = MetalSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.len + } + + fn dot(&self, other: &Self) -> F { + self.as_read().dot(&other.as_read()) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + self.as_read().sumcheck_polynomial(&other.as_read()) + } + + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + self.as_read().mixed_extend(embedding, point) + } + + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &MetalBuffer, + ) -> M::Target { + self.as_read().mixed_dot(embedding, other) + } + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + self.as_read().mixed_univariate_evaluate(embedding, point) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut MetalBuffer, + weight: M::Target, + ) { + self.as_read() + .mixed_scalar_mul_add_to(embedding, accumulator, weight) + } + + fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { + let (start, end) = crate::algebra::buffer::resolve_range(range, self.len); + MetalSlice { + field: self.field.clone(), + offset: self.offset + start, + len: end - start, + _parent: PhantomData, + } + } +} + +impl BufferWrite for MetalSliceMut<'_, F> { + type SliceMut<'a> + = MetalSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scale(&mut self, weight: F) { + assert_bn254::(); + if self.len == 0 { + return; + } + let weight = upload_field(&[f_to_field256(weight)]); + scalar_mul_at_offset(&self.field, self.offset, self.len, &weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + if !check_univariate_evaluators(evaluators, scalars, self.len) { + return; + } + if self.offset == 0 { + // Offset 0: reuse the optimized full-buffer accumulate path. + geometric_accumulate_full(&self.field, self.len, evaluators, scalars); + } else { + // Arbitrary offset: bind the buffer at a byte offset so the + // kernel's gid-based indexing addresses field[offset + gid]. + let points = evaluators + .iter() + .map(|e| f_to_field256(e.point)) + .collect::>(); + let scalars = scalars + .iter() + .copied() + .map(f_to_field256) + .collect::>(); + let points = upload_field(&points); + let scalars = upload_field(&scalars); + geometric_accumulate_at_offset( + &self.field, + self.offset, + self.len, + &points, + &scalars, + evaluators.len(), + ); + } + } + + fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { + let (start, end) = crate::algebra::buffer::resolve_range(range, self.len); + MetalSliceMut::from_field(self.field.clone(), self.offset + start, end - start) + } + + fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { + assert!( + mid <= self.len, + "split_at_mut mid {mid} out of bounds for {}", + self.len + ); + ( + MetalSliceMut::from_field(self.field.clone(), self.offset, mid), + MetalSliceMut::from_field(self.field.clone(), self.offset + mid, self.len - mid), + ) + } +} + impl MetalBuffer { fn bn254_buffer_target(&self) -> MetalFieldBuffer { self.field @@ -1698,6 +2136,7 @@ struct MetalRuntime { fold: ComputePipelineState, fold_pair: ComputePipelineState, scalar_mul_add: ComputePipelineState, + scalar_mul: ComputePipelineState, dot: ComputePipelineState, sumcheck: ComputePipelineState, dot_chunks: ComputePipelineState, @@ -1749,6 +2188,7 @@ fn runtime() -> &'static MetalRuntime { fold: pipeline("bn254_fold"), fold_pair: pipeline("bn254_fold_pair"), scalar_mul_add: pipeline("bn254_scalar_mul_add"), + scalar_mul: pipeline("bn254_scalar_mul"), dot: pipeline("bn254_dot"), sumcheck: pipeline("bn254_sumcheck"), dot_chunks: pipeline("bn254_dot_chunks"), @@ -1796,6 +2236,7 @@ fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { "bn254_fold" => &rt.fold, "bn254_fold_pair" => &rt.fold_pair, "bn254_scalar_mul_add" => &rt.scalar_mul_add, + "bn254_scalar_mul" => &rt.scalar_mul, "bn254_dot" => &rt.dot, "bn254_sumcheck" => &rt.sumcheck, "bn254_dot_chunks" => &rt.dot_chunks, @@ -1934,6 +2375,17 @@ struct ApplyCosetTwiddlesParams { } fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field256 { + parallel_dot_at(a, 0, b, 0, len) +} + +/// Inner product of `a[a_off..a_off+len]` and `b[b_off..b_off+len]`. +fn parallel_dot_at( + a: &MetalFieldBuffer, + a_off: usize, + b: &MetalFieldBuffer, + b_off: usize, + len: usize, +) -> Field256 { if len == 0 { return Field256::ZERO; } @@ -1943,10 +2395,11 @@ fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), }; let command = rt.queue.new_command_buffer(); - encode_u32_kernel( + encode_u32_kernel_with_offsets( command, pipeline(rt, "bn254_dot_chunks"), &[&a.limbs, &b.limbs, &partials.limbs], + &[field_byte_offset(a_off), field_byte_offset(b_off), 0], &[len as u32, REDUCTION_CHUNK_SIZE as u32], partial_count, ); @@ -1960,6 +2413,18 @@ fn parallel_sumcheck( b: &MetalFieldBuffer, len: usize, fold_half: usize, +) -> (Field256, Field256) { + parallel_sumcheck_at(a, 0, b, 0, len, fold_half) +} + +/// Sumcheck `(c0, c2)` over `a[a_off..]` and `b[b_off..]` (logical length `len`). +fn parallel_sumcheck_at( + a: &MetalFieldBuffer, + a_off: usize, + b: &MetalFieldBuffer, + b_off: usize, + len: usize, + fold_half: usize, ) -> (Field256, Field256) { let partial_count = fold_half.div_ceil(REDUCTION_CHUNK_SIZE); let rt = runtime(); @@ -1967,10 +2432,11 @@ fn parallel_sumcheck( limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), }; let command = rt.queue.new_command_buffer(); - encode_u32_kernel( + encode_u32_kernel_with_offsets( command, pipeline(rt, "bn254_sumcheck_chunks"), &[&a.limbs, &b.limbs, &partials.limbs], + &[field_byte_offset(a_off), field_byte_offset(b_off), 0], &[ len as u32, fold_half as u32, @@ -2021,6 +2487,16 @@ fn parallel_univariate_evaluate( coeffs: &MetalFieldBuffer, point: &MetalFieldBuffer, len: usize, +) -> Field256 { + parallel_univariate_evaluate_at(coeffs, 0, point, len) +} + +/// Univariate evaluation of `coeffs[coeff_off..coeff_off+len]` at `point`. +fn parallel_univariate_evaluate_at( + coeffs: &MetalFieldBuffer, + coeff_off: usize, + point: &MetalFieldBuffer, + len: usize, ) -> Field256 { if len == 0 { return Field256::ZERO; @@ -2031,10 +2507,11 @@ fn parallel_univariate_evaluate( limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), }; let command = rt.queue.new_command_buffer(); - encode_u32_kernel( + encode_u32_kernel_with_offsets( command, pipeline(rt, "bn254_univariate_eval_chunks"), &[&coeffs.limbs, &point.limbs, &partials.limbs], + &[field_byte_offset(coeff_off), 0, 0], &[len as u32, REDUCTION_CHUNK_SIZE as u32], partial_count, ); @@ -2535,6 +3012,111 @@ fn maybe_upload_bn254(values: &[T]) -> Option { (type_name::() == type_name::()).then(|| upload_field(as_field256_slice(values))) } +/// Geometric accumulate into `field[offset .. offset+len]` by binding the +/// field buffer at a byte offset; the kernel itself indexes from `gid == 0`. +fn geometric_accumulate_at_offset( + field: &MetalFieldBuffer, + offset: usize, + len: usize, + points: &MetalFieldBuffer, + scalars: &MetalFieldBuffer, + num_points: usize, +) { + if len == 0 || num_points == 0 { + return; + } + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + let encoder = command.new_compute_command_encoder(); + let pipe = pipeline(rt, "bn254_geometric_accumulate"); + encoder.set_compute_pipeline_state(pipe); + let byte_offset = (offset * size_of::()) as u64; + encoder.set_buffer(0, Some(&field.limbs), byte_offset); + encoder.set_buffer(1, Some(&points.limbs), 0); + encoder.set_buffer(2, Some(&scalars.limbs), 0); + let len_u32 = len as u32; + let num_u32 = num_points as u32; + encoder.set_bytes(3, size_of::() as u64, (&len_u32 as *const u32).cast()); + encoder.set_bytes(4, size_of::() as u64, (&num_u32 as *const u32).cast()); + dispatch(&encoder, pipe, len); + encoder.end_encoding(); + wait_for_command_named(command, "bn254_geometric_accumulate_window"); +} + +/// In-place scalar multiply of `field[offset .. offset+len]` by binding the +/// field buffer at a byte offset; the kernel indexes from `gid == 0`. +fn scalar_mul_at_offset( + field: &MetalFieldBuffer, + offset: usize, + len: usize, + weight: &MetalFieldBuffer, +) { + if len == 0 { + return; + } + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + let encoder = command.new_compute_command_encoder(); + let pipe = pipeline(rt, "bn254_scalar_mul"); + encoder.set_compute_pipeline_state(pipe); + let byte_offset = (offset * size_of::()) as u64; + encoder.set_buffer(0, Some(&field.limbs), byte_offset); + encoder.set_buffer(1, Some(&weight.limbs), 0); + let len_u32 = len as u32; + encoder.set_bytes(2, size_of::() as u64, (&len_u32 as *const u32).cast()); + dispatch(&encoder, pipe, len); + encoder.end_encoding(); + wait_for_command_named(command, "bn254_scalar_mul_window"); +} + +/// Multilinear extension of `field[off..off+len]` evaluated at `point`. +fn parallel_multilinear_extend_at( + field: &MetalFieldBuffer, + off: usize, + len: usize, + point: &MetalFieldBuffer, + num_vars: usize, +) -> Field256 { + let rt = runtime(); + let out = new_shared_buffer(rt, (4 * size_of::()) as u64); + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_multilinear_extend"), + &[&field.limbs, &point.limbs, &out], + &[field_byte_offset(off), 0, 0], + &[len as u32, num_vars as u32], + 1, + ); + wait_for_command_named(command, "bn254_multilinear_extend"); + download_field(&out, 1)[0] +} + +/// `acc[acc_off..] += weight * vector[vec_off..]` over `len` elements. +fn scalar_mul_add_at( + acc: &MetalFieldBuffer, + acc_off: usize, + vector: &MetalFieldBuffer, + vec_off: usize, + weight: &MetalFieldBuffer, + len: usize, +) { + if len == 0 { + return; + } + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_scalar_mul_add"), + &[&acc.limbs, &vector.limbs, &weight.limbs], + &[field_byte_offset(acc_off), field_byte_offset(vec_off), 0], + &[len as u32], + len, + ); + wait_for_command_named(command, "bn254_scalar_mul_add"); +} + fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer { let rt = runtime(); let byte_len = (len * 4 * size_of::()) as u64; @@ -2555,12 +3137,29 @@ fn encode_u32_kernel( buffers: &[&Buffer], constants: &[u32], threads: usize, +) { + encode_u32_kernel_with_offsets(command, pipeline, buffers, &[], constants, threads); +} + +/// Like [`encode_u32_kernel`], but binds `buffers[i]` at byte offset +/// `offsets[i]` (defaulting to 0 when `offsets` is shorter). This is the +/// mechanism that lets a view dispatch a kernel over `parent[offset..]` +/// without copying: only the input binding shifts, the kernel still indexes +/// from `gid == 0`. +fn encode_u32_kernel_with_offsets( + command: &metal::CommandBufferRef, + pipeline: &ComputePipelineState, + buffers: &[&Buffer], + offsets: &[u64], + constants: &[u32], + threads: usize, ) { let encoder = command.new_compute_command_encoder(); encoder.set_compute_pipeline_state(pipeline); let mut index = 0; - for buffer in buffers { - encoder.set_buffer(index, Some(buffer), 0); + for (i, buffer) in buffers.iter().enumerate() { + let byte_offset = offsets.get(i).copied().unwrap_or(0); + encoder.set_buffer(index, Some(buffer), byte_offset); index += 1; } for constant in constants { @@ -2575,6 +3174,12 @@ fn encode_u32_kernel( encoder.end_encoding(); } +/// Byte offset of element `offset` within a `Field256` buffer. +#[inline] +fn field_byte_offset(offset: usize) -> u64 { + (offset * size_of::()) as u64 +} + fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { let rt = runtime(); let command = rt.queue.new_command_buffer(); @@ -2582,15 +3187,6 @@ fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usi wait_for_command_named(command, name); } -fn run_reduce(name: &str, buffers: &[&Buffer], constants: &[u32], out_len: usize) -> Vec { - let rt = runtime(); - let out = new_shared_buffer(rt, (out_len * 4 * size_of::()) as u64); - let mut all = buffers.to_vec(); - all.push(&out); - run_in_place(name, &all, constants, 1); - download_field(&out, out_len) -} - fn dispatch( encoder: &metal::ComputeCommandEncoderRef, pipeline: &ComputePipelineState, diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index ebe613c8..220ea28d 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -14,7 +14,7 @@ use spongefish::{Decoding, VerificationResult}; use crate::{ algebra::{ - buffer::{ActiveBuffer, BufferOps, FieldOps}, + buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, embedding::Identity, multilinear_extend, scalar_mul_add_new, univariate_evaluate, }, diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 4c22c79b..00b23e2c 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -15,7 +15,7 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::{ActiveBuffer, BufferOps, FieldOps, WindowOps}, + buffer::{ActiveBuffer, BufferOps, BufferRead, BufferWrite}, dot, embedding::{Embedding, Identity}, eq_weights, lift, @@ -255,7 +255,7 @@ impl Config { let masked = self.source.masked_message_length(); let in_domain_evaluators = evaluators(&eval_points, masked); covector - .window_mut(0, masked) + .slice_mut(..masked) .accumulate_univariate_evaluations(&in_domain_evaluators, in_domain_rlc_coeffs); } diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index a30be48a..36a00d32 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -24,7 +24,7 @@ use std::{ ops::Neg, }; -use crate::algebra::buffer::FieldOps; +use crate::algebra::buffer::{Buffer, BufferRead}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use ordered_float::OrderedFloat; diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 4bef3bbc..ae83e348 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -2,7 +2,7 @@ use std::fmt; -use crate::algebra::buffer::FieldOps; +use crate::algebra::buffer::{Buffer, BufferRead}; use ark_ff::Field; use ark_std::rand::{CryptoRng, RngCore}; use serde::{Deserialize, Serialize}; diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 55649e62..d8c09d2b 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -1,6 +1,6 @@ use std::borrow::Cow; -use crate::algebra::buffer::FieldOps; +use crate::algebra::buffer::{Buffer, BufferRead, BufferWrite}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; #[cfg(feature = "tracing")] @@ -141,8 +141,11 @@ impl Config { // Random linear combination of the vectors. let mut vector_rlc_coeffs: Vec = geometric_challenge(prover_state, num_vectors); assert_eq!(vector_rlc_coeffs[0], M::Target::ONE); - let mut vector = - FieldOps::mixed_linear_combination(self.embedding(), vectors, &vector_rlc_coeffs); + let mut vector = ActiveBuffer::::mixed_linear_combination( + self.embedding(), + vectors, + &vector_rlc_coeffs, + ); let mut prev_witness: RoundWitness<'a, M::Target, M> = RoundWitness::Initial(witnesses); @@ -154,13 +157,13 @@ impl Config { constraint_rlc_coeffs.split_at(linear_forms.len()); let mut linear_forms = linear_forms; let mut covector = if has_constraints { - FieldOps::linear_forms_rlc( + ActiveBuffer::::linear_forms_rlc( self.initial_size(), &mut linear_forms, initial_forms_rlc_coeffs, ) } else { - FieldOps::zeros(0) + ActiveBuffer::::zeros(0) }; drop(linear_forms); @@ -203,7 +206,7 @@ impl Config { vector.fold(f); } // Covector must be all zeros. - covector = FieldOps::zeros(self.initial_sumcheck.final_size()); + covector = ActiveBuffer::::zeros(self.initial_sumcheck.final_size()); folding_randomness }; let mut evaluation_point = folding_randomness.clone(); From b08064c250652a9a299b29ec31c1a1a247acc0a9 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 18 Jun 2026 13:15:13 +0200 Subject: [PATCH 19/35] refactor: move buffer abstraction into its own crate --- benches/metal_buffer.rs | 1 + src/algebra/buffer.rs | 742 +---------------------------- src/algebra/metal_buffer.rs | 101 ++-- src/algebra/ntt/cooley_tukey.rs | 6 +- src/algebra/ntt/mod.rs | 44 +- src/bin/benchmark.rs | 2 + src/bin/gpu_proof_cpu_verify.rs | 24 +- src/bin/main.rs | 4 +- src/buffer/cpu/mod.rs | 307 ++++++++++++ src/buffer/cpu/traits.rs | 117 +++++ src/buffer/mod.rs | 193 ++++++++ src/lib.rs | 1 + src/protocols/basecase.rs | 8 +- src/protocols/code_switch.rs | 31 +- src/protocols/mask_proximity.rs | 71 +-- src/protocols/matrix_commit.rs | 4 +- src/protocols/whir/mod.rs | 1 + src/protocols/whir/prover.rs | 3 +- src/protocols/whir_zk/committer.rs | 1 + src/protocols/whir_zk/mod.rs | 17 +- src/protocols/whir_zk/prover.rs | 26 +- 21 files changed, 805 insertions(+), 899 deletions(-) create mode 100644 src/buffer/cpu/mod.rs create mode 100644 src/buffer/cpu/traits.rs create mode 100644 src/buffer/mod.rs diff --git a/benches/metal_buffer.rs b/benches/metal_buffer.rs index a3e85b65..e1fdb39d 100644 --- a/benches/metal_buffer.rs +++ b/benches/metal_buffer.rs @@ -6,6 +6,7 @@ mod bench { buffer::{BufferRead, CpuBuffer, MetalBuffer}, fields::Field256 as F, }, + buffer::BufferOps, hash::{Hash, HashEngine, MetalSha2, Sha2}, }; diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs index 329a684e..aa11ed17 100644 --- a/src/algebra/buffer.rs +++ b/src/algebra/buffer.rs @@ -1,20 +1,5 @@ -use std::{any::Any, mem, ops::RangeBounds}; - -use ark_ff::Field; -use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; - -use crate::{ - algebra::{ - dot, - embedding::Embedding, - linear_form::{Covector, LinearForm, UnivariateEvaluation}, - mixed_dot, mixed_multilinear_extend, mixed_scalar_mul_add, mixed_univariate_evaluate, - sumcheck::{compute_sumcheck_polynomial, fold}, - }, - engines::EngineId, - hash::{self, Hash}, - protocols::{matrix_commit::Encodable, merkle_tree}, -}; +pub use crate::buffer::cpu::{CpuBuffer, CpuSlice, CpuSliceMut}; +pub use crate::buffer::{Buffer, BufferOps, BufferRead, BufferWrite}; #[cfg(all(feature = "metal", target_os = "macos"))] pub use super::metal_buffer::{MetalBuffer, MetalSlice, MetalSliceMut}; @@ -32,726 +17,3 @@ pub type ActiveBuffer = CpuBuffer; pub type ActiveSlice<'a, T> = CpuSlice<'a, T>; #[cfg(not(all(feature = "metal", target_os = "macos")))] pub type ActiveSliceMut<'a, T> = CpuSliceMut<'a, T>; - -pub trait BufferOps { - fn from_vec(source: Vec) -> Self; - fn from_slice(source: &[T]) -> Self; - fn as_slice(&self) -> &[T]; - fn num_rows(&self, num_cols: usize) -> usize { - self.len() / num_cols - } - fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; - fn at_index(&self, index: usize) -> Option; - fn len(&self) -> usize; - fn is_empty(&self) -> bool { - self.len() == 0 - } - fn extend(&self, other: &Self) -> Self; - fn merklize( - &self, - num_cols: usize, - hash_engine_id: EngineId, - merkle_config: &merkle_tree::Config, - ) -> (ActiveBuffer, Hash) - where - T: Encodable + Send + Sync; -} - -/// Read-only operations over a contiguous run of field elements — the buffer -/// analogue of `&[F]`. -/// -/// Implemented by the owned buffers ([`CpuBuffer`]/`MetalBuffer`, as the -/// full-range case) and by the borrowed read views ([`CpuSlice`]/`MetalSlice`). -/// Every op here works identically on a whole buffer or any sub-range obtained -/// via [`Self::slice`]. On `MetalBuffer` a view aliases the parent's GPU -/// allocation (byte-offset binding), so no data is copied. -pub trait BufferRead { - /// A same-backend owned buffer over another field, produced by the - /// mixed-field ops (e.g. `mixed_dot` against a target-field buffer). For an - /// owned buffer this is the buffer itself over `T`; views report the owning - /// buffer type. - type TargetBuffer: Buffer; - - /// Read-only view type produced by slicing this buffer/view. - type Slice<'a>: BufferRead - where - Self: 'a, - F: 'a; - - /// Number of elements in this view. - fn read_len(&self) -> usize; - fn read_is_empty(&self) -> bool { - self.read_len() == 0 - } - - /// Inner product with another view of the same length. - fn dot(&self, other: &Self) -> F - where - Self: Sized; - - /// Sumcheck round coefficients `(c0, c2)` for `dot(self, other)`. - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) - where - Self: Sized; - - /// Multilinear extension evaluated at a target-field point. - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target; - - /// Inner product with a target-field buffer. - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &Self::TargetBuffer, - ) -> M::Target; - - /// Univariate evaluation at a target-field point. - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target; - - /// `accumulator += weight * self`, lifted into the target field. - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut Self::TargetBuffer, - weight: M::Target, - ); - - /// Borrow `self[range]` as a read-only view (analogous to `&v[range]`). - fn slice(&self, range: impl RangeBounds) -> Self::Slice<'_>; -} - -/// In-place, length-preserving operations over a contiguous run of field -/// elements — the buffer analogue of `&mut [F]`. -/// -/// Implemented by the owned buffers (full-range) and the borrowed mutable views -/// ([`CpuSliceMut`]/`MetalSliceMut`). A sub-range obtained via [`Self::slice_mut`] -/// is the same view type as the whole buffer, so the ops compose just like -/// slicing a `Vec`. -/// -/// Constructors (`zeros`, `random`, …) and length-changing operations (`fold`, -/// `zero_pad`) live on [`Buffer`], the owned-buffer trait, because a borrowed -/// view cannot construct or resize its backing storage. -pub trait BufferWrite: BufferRead { - /// Mutable view type produced by slicing this buffer/view. - type SliceMut<'a>: BufferWrite - where - Self: 'a, - F: 'a; - - /// In-place scalar multiplication: `self[i] *= weight`. - fn scale(&mut self, weight: F); - - /// Accumulate `Σ_j scalars[j] · evaluators[j].point^i` into entry `i`. - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ); - - /// Borrow `self[range]` as a mutable, zero-copy view (analogous to - /// `&mut v[range]`). - fn slice_mut(&mut self, range: impl RangeBounds) -> Self::SliceMut<'_>; - - /// Split into two disjoint mutable views `[0, mid)` and `[mid, len)`. - fn split_at_mut(&mut self, mid: usize) -> (Self::SliceMut<'_>, Self::SliceMut<'_>); -} - -/// Owned field-buffer operations — the `Vec` half of the field API. -/// -/// `BufferOps` stays generic over any element type, so hashes and digests can -/// use it. `BufferRead`/`BufferWrite` are the slice-like field ops implemented -/// by owners and views. This trait is only for owned field buffers: operations -/// here construct storage or change its length, so borrowed views cannot -/// implement them. -pub trait Buffer: BufferOps + BufferWrite + Clone { - fn zeros(length: usize) -> Self; - - /// Geometric sequence `[1, base, base², …, base^(length-1)]`. - fn geometric_sequence(base: F, length: usize) -> Self; - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution; - - fn zero_pad(&mut self); - fn fold(&mut self, weight: F); - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - self.fold(weight); - other.fold(weight); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - self.fold_pair(other, weight); - self.sumcheck_polynomial(other) - } - - fn linear_forms_rlc( - size: usize, - linear_forms: &mut [Box>], - rlc_coeffs: &[F], - ) -> Self; - - fn mixed_linear_combination>( - embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer; -} - -/// Resolve any range expression (`a..b`, `..b`, `a..`, `..`) against `len`, -/// returning `(start, end)` and bounds-checking like slice indexing. -pub(crate) fn resolve_range(range: impl RangeBounds, len: usize) -> (usize, usize) { - use std::ops::Bound::{Excluded, Included, Unbounded}; - let start = match range.start_bound() { - Included(&s) => s, - Excluded(&s) => s + 1, - Unbounded => 0, - }; - let end = match range.end_bound() { - Included(&e) => e + 1, - Excluded(&e) => e, - Unbounded => len, - }; - assert!( - start <= end && end <= len, - "slice range {start}..{end} out of bounds for length {len}" - ); - (start, end) -} - -/// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. -pub struct CpuSlice<'a, F> { - data: &'a [F], -} - -/// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. -pub struct CpuSliceMut<'a, F> { - data: &'a mut [F], -} - -impl CpuSlice<'_, F> { - pub fn len(&self) -> usize { - self.data.len() - } - pub fn is_empty(&self) -> bool { - self.data.is_empty() - } - pub fn as_slice(&self) -> &[F] { - self.data - } -} - -impl CpuSliceMut<'_, F> { - pub fn len(&self) -> usize { - self.data.len() - } - pub fn is_empty(&self) -> bool { - self.data.is_empty() - } - pub fn as_slice(&self) -> &[F] { - self.data - } -} - -impl BufferRead for CpuSlice<'_, F> { - type TargetBuffer = CpuBuffer; - type Slice<'a> - = CpuSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.data.len() - } - fn dot(&self, other: &Self) -> F { - crate::algebra::dot(self.data, other.data) - } - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - compute_sumcheck_polynomial(self.data, other.data) - } - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - mixed_multilinear_extend(embedding, self.data, point) - } - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &CpuBuffer, - ) -> M::Target { - mixed_dot(embedding, other.as_slice(), self.data) - } - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - mixed_univariate_evaluate(embedding, self.data, point) - } - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut CpuBuffer, - weight: M::Target, - ) { - mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.data); - } - - fn slice(&self, range: impl RangeBounds) -> CpuSlice<'_, F> { - let (start, end) = resolve_range(range, self.data.len()); - CpuSlice { - data: &self.data[start..end], - } - } -} - -impl BufferRead for CpuSliceMut<'_, F> { - type TargetBuffer = CpuBuffer; - type Slice<'a> - = CpuSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.data.len() - } - fn dot(&self, other: &Self) -> F { - crate::algebra::dot(self.data, other.data) - } - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - compute_sumcheck_polynomial(self.data, other.data) - } - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - mixed_multilinear_extend(embedding, self.data, point) - } - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &CpuBuffer, - ) -> M::Target { - mixed_dot(embedding, other.as_slice(), self.data) - } - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - mixed_univariate_evaluate(embedding, self.data, point) - } - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut CpuBuffer, - weight: M::Target, - ) { - mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.data); - } - - fn slice(&self, range: impl RangeBounds) -> CpuSlice<'_, F> { - let (start, end) = resolve_range(range, self.data.len()); - CpuSlice { - data: &self.data[start..end], - } - } -} - -impl BufferWrite for CpuSliceMut<'_, F> { - type SliceMut<'a> - = CpuSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scale(&mut self, weight: F) { - crate::algebra::scalar_mul(self.data, weight); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - UnivariateEvaluation::accumulate_many(evaluators, self.data, scalars); - } - - fn slice_mut(&mut self, range: impl RangeBounds) -> CpuSliceMut<'_, F> { - let (start, end) = resolve_range(range, self.data.len()); - CpuSliceMut { - data: &mut self.data[start..end], - } - } - - fn split_at_mut(&mut self, mid: usize) -> (CpuSliceMut<'_, F>, CpuSliceMut<'_, F>) { - let (lo, hi) = self.data.split_at_mut(mid); - (CpuSliceMut { data: lo }, CpuSliceMut { data: hi }) - } -} - -#[derive( - Clone, - PartialEq, - Eq, - PartialOrd, - Ord, - Hash, - Debug, - Default, - serde::Serialize, - serde::Deserialize, -)] -pub struct CpuBuffer { - data: Vec, -} - -impl CpuBuffer { - pub fn from_vec(source: Vec) -> Self { - Self { data: source } - } - - pub fn from_slice(source: &[T]) -> Self { - Self { - data: Vec::from(source), - } - } - - pub(crate) fn as_slice(&self) -> &[T] { - self.data.as_slice() - } -} - -impl BufferOps for CpuBuffer { - fn at_index(&self, index: usize) -> Option { - if index >= self.len() { - None - } else { - Some(self.data[index]) - } - } - - fn as_slice(&self) -> &[T] { - &self.data - } - - fn len(&self) -> usize { - self.data.len() - } - - fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { - let mut result = Vec::with_capacity(indices.len() * num_cols); - for i in indices { - result.extend_from_slice(&self.data[i * num_cols..(i + 1) * num_cols]); - } - result - } - - fn extend(&self, other: &Self) -> Self { - let mut data = Vec::with_capacity(self.data.len() + other.data.len()); - data.extend_from_slice(&self.data); - data.extend_from_slice(&other.data); - Self { data } - } - - fn from_vec(source: Vec) -> Self { - Self::from_vec(source) - } - - fn from_slice(source: &[T]) -> Self { - Self::from_slice(source) - } - - fn merklize( - &self, - num_cols: usize, - hash_engine_id: EngineId, - merkle_config: &merkle_tree::Config, - ) -> (ActiveBuffer, Hash) - where - T: Encodable + Send + Sync, - { - let _ = num_cols; // CPU leaf hashing derives the column count from the row count. - let engine = hash::ENGINES - .retrieve(hash_engine_id) - .expect("Failed to retrieve hash engine"); - #[cfg(feature = "tracing")] - tracing::Span::current().record("engine", engine.name().as_ref()); - - // Hash each row into a leaf, then build the full Merkle node array. - let mut leaves = vec![Hash::default(); merkle_config.num_leaves]; - crate::protocols::matrix_commit::hash_rows(&*engine, self.as_slice(), &mut leaves); - let nodes = merkle_config.build_nodes(leaves); - let root = nodes[nodes.len() - 1]; - - (ActiveBuffer::::from_vec(nodes), root) - } -} - -impl BufferRead for CpuBuffer { - type TargetBuffer = CpuBuffer; - type Slice<'a> - = CpuSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.data.len() - } - - fn dot(&self, other: &Self) -> F { - dot(&self.data, &other.data) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - compute_sumcheck_polynomial(&self.data, other.as_slice()) - } - - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - mixed_multilinear_extend(embedding, &self.data, point) - } - - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &CpuBuffer, - ) -> M::Target { - mixed_dot(embedding, other.as_slice(), self.as_slice()) - } - - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - mixed_univariate_evaluate(embedding, &self.data, point) - } - - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut CpuBuffer, - weight: M::Target, - ) { - mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, self.as_slice()); - } - - fn slice(&self, range: impl RangeBounds) -> CpuSlice<'_, F> { - let (start, end) = resolve_range(range, self.data.len()); - CpuSlice { - data: &self.data[start..end], - } - } -} - -impl BufferWrite for CpuBuffer { - type SliceMut<'a> - = CpuSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scale(&mut self, weight: F) { - crate::algebra::scalar_mul(&mut self.data, weight); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - UnivariateEvaluation::accumulate_many(evaluators, &mut self.data, scalars); - } - - fn slice_mut(&mut self, range: impl RangeBounds) -> CpuSliceMut<'_, F> { - let (start, end) = resolve_range(range, self.data.len()); - CpuSliceMut { - data: &mut self.data[start..end], - } - } - - fn split_at_mut(&mut self, mid: usize) -> (CpuSliceMut<'_, F>, CpuSliceMut<'_, F>) { - let (lo, hi) = self.data.split_at_mut(mid); - (CpuSliceMut { data: lo }, CpuSliceMut { data: hi }) - } -} - -impl Buffer for CpuBuffer { - fn zeros(length: usize) -> Self { - Self { - data: vec![F::ZERO; length], - } - } - - fn geometric_sequence(base: F, length: usize) -> Self { - Self::from_vec(crate::algebra::geometric_sequence(base, length)) - } - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution, - { - Self { - data: (0..length).map(|_| rng.gen()).collect(), - } - } - - fn linear_forms_rlc( - size: usize, - linear_forms: &mut [Box>], - rlc_coeffs: &[F], - ) -> Self { - assert_eq!(linear_forms.len(), rlc_coeffs.len()); - let mut covector = vec![F::ZERO; size]; - if let Some((first, linear_forms)) = linear_forms.split_first_mut() { - debug_assert_eq!(rlc_coeffs[0], F::ONE); - if let Some(covector_form) = - (first.as_mut() as &mut dyn Any).downcast_mut::>() - { - mem::swap(&mut covector, &mut covector_form.vector); - } else { - first.accumulate(&mut covector, F::ONE); - } - for (rlc_coeff, linear_form) in rlc_coeffs[1..].iter().zip(linear_forms) { - linear_form.accumulate(&mut covector, *rlc_coeff); - } - } - Self { data: covector } - } - - fn zero_pad(&mut self) { - if !self.is_empty() { - self.data.resize(self.len().next_power_of_two(), F::ZERO); - } - } - - fn fold(&mut self, weight: F) { - fold(&mut self.data, weight); - } - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - self.fold(weight); - other.fold(weight); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - self.fold_pair(other, weight); - self.sumcheck_polynomial(other) - } - - fn mixed_linear_combination>( - embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer { - assert_eq!(vectors.len(), coeffs.len()); - let Some((first, vectors)) = vectors.split_first() else { - return CpuBuffer::from_vec(Vec::new()); - }; - debug_assert_eq!(coeffs[0], M::Target::ONE); - let mut accumulator = crate::algebra::lift(embedding, first.as_slice()); - for (coeff, vector) in coeffs[1..].iter().zip(vectors) { - mixed_scalar_mul_add(embedding, &mut accumulator, *coeff, vector.as_slice()); - } - CpuBuffer::from_vec(accumulator) - } -} - -#[cfg(test)] -mod tests { - use ark_ff::AdditiveGroup; - - use super::*; - use crate::algebra::{ - fields::Field64, geometric_accumulate, geometric_sequence as geometric_sequence_vec, - }; - - type F = Field64; - - #[test] - fn geometric_sequence_matches_free_function() { - let base = F::from(7u64); - for length in 0..6 { - let buffer = CpuBuffer::::geometric_sequence(base, length); - assert_eq!( - BufferOps::as_slice(&buffer), - geometric_sequence_vec(base, length).as_slice(), - "geometric_sequence mismatch at length {length}" - ); - } - } - - #[test] - fn scale_multiplies_in_place() { - let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; - let weight = F::from(5u64); - let mut buffer = CpuBuffer::from_vec(values.clone()); - buffer.scale(weight); - let expected: Vec = values.iter().map(|&v| v * weight).collect(); - assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); - } - - #[test] - fn slice_accumulate_matches_restricted_geometric_accumulate() { - let len = 8usize; - let points = vec![F::from(3u64), F::from(5u64)]; - let scalars = vec![F::from(2u64), F::from(9u64)]; - - // Try a prefix slice (offset 0) and an interior slice (offset > 0). - for (offset, window_len) in [(0usize, 5usize), (2usize, 4usize)] { - let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); - let evaluators: Vec<_> = points - .iter() - .map(|&point| UnivariateEvaluation::new(point, window_len)) - .collect(); - buffer - .slice_mut(offset..offset + window_len) - .accumulate_univariate_evaluations(&evaluators, &scalars); - - // Reference: accumulate Σ_j scalars[j]·points[j]^i into the same - // sub-slice of a plain vector. - let mut reference = vec![F::ZERO; len]; - geometric_accumulate( - &mut reference[offset..offset + window_len], - scalars.clone(), - &points, - ); - - assert_eq!( - BufferOps::as_slice(&buffer), - reference.as_slice(), - "slice accumulate mismatch at offset {offset}, len {window_len}" - ); - } - } - - #[test] - fn slice_scale_multiplies_only_the_range() { - let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; - let weight = F::from(5u64); - let mut buffer = CpuBuffer::from_vec(values.clone()); - buffer.slice_mut(1..3).scale(weight); - let expected = vec![values[0], values[1] * weight, values[2] * weight, values[3]]; - assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); - } -} diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs index 6a07de37..6334797e 100644 --- a/src/algebra/metal_buffer.rs +++ b/src/algebra/metal_buffer.rs @@ -28,7 +28,7 @@ use crate::{ ntt::{ReedSolomon, RsDomain}, }, engines::EngineId, - hash::{self, metal_profile, Hash as Digest, MetalSha2, SHA2}, + hash::{self, metal_profile, Hash as Digest, MetalSha2}, protocols::{ matrix_commit::{hash_rows, Encodable}, merkle_tree, @@ -1055,7 +1055,7 @@ where D: Deserializer<'de>, { let data = Vec::::deserialize(deserializer)?; - Ok(Self::from_vec(data)) + Ok(BufferOps::from_vec(data)) } } @@ -1064,27 +1064,6 @@ impl MetalBuffer { let _ = runtime(); } - pub fn from_vec(source: Vec) -> Self { - let len = source.len(); - let field = maybe_upload_bn254(&source); - let host_cache = if field.is_some() { - OnceCell::new() - } else { - OnceCell::from(source) - }; - Self { - len, - host_cache, - field, - hash: None, - _marker: PhantomData, - } - } - - pub fn from_slice(source: &[T]) -> Self { - Self::from_vec(Vec::from(source)) - } - pub(crate) fn as_slice(&self) -> &[T] { self.host_cache .get_or_init(|| self.download_host_cache()) @@ -1170,6 +1149,8 @@ impl MetalBuffer { } impl BufferOps for MetalBuffer { + type Nodes = MetalBuffer; + fn as_slice(&self) -> &[T] { self.as_slice() } @@ -1202,18 +1183,24 @@ impl BufferOps for MetalBuffer { } fn from_vec(source: Vec) -> Self { - Self::from_vec(source) + let len = source.len(); + let field = maybe_upload_bn254(&source); + let host_cache = if field.is_some() { + OnceCell::new() + } else { + OnceCell::from(source) + }; + Self { + len, + host_cache, + field, + hash: None, + _marker: PhantomData, + } } fn from_slice(source: &[T]) -> Self { - Self::from_slice(source) - } - - fn extend(&self, other: &Self) -> Self { - let mut data = Vec::with_capacity(self.len() + other.len()); - data.extend_from_slice(self.as_slice()); - data.extend_from_slice(other.as_slice()); - Self::from_vec(data) + Self::from_vec(Vec::from(source)) } fn merklize( @@ -1221,15 +1208,21 @@ impl BufferOps for MetalBuffer { num_cols: usize, leaf_hash: EngineId, merkle: &merkle_tree::Config, - ) -> (MetalBuffer, Digest) + ) -> (Self::Nodes, Digest) where T: Encodable + Send + Sync, { let num_rows = merkle.num_leaves; let layers = merkle.layers.len(); + assert_eq!(self.len(), num_cols * num_rows); // Fast path: build the whole tree on the GPU when the leaf and every node layer is SHA2. - if leaf_hash == SHA2 && merkle.layers.iter().all(|layer| layer.hash_id == SHA2) { + if leaf_hash == hash::SHA2 + && merkle + .layers + .iter() + .all(|layer| layer.hash_id == hash::SHA2) + { if let Some(nodes) = self.commit_bn254_rows_sha2_merkle(num_cols, num_rows, layers) { let root = nodes .read_hash_at(merkle.num_nodes() - 1) @@ -1238,7 +1231,6 @@ impl BufferOps for MetalBuffer { } } - // CPU fallback: hash each row on the CPU, then build the tree. let cpu_nodes = || { let engine = hash::ENGINES .retrieve(leaf_hash) @@ -1248,8 +1240,7 @@ impl BufferOps for MetalBuffer { merkle.build_nodes(leaves) }; - // Otherwise hash leaves (on the GPU if SHA2, else CPU) and build the tree on the CPU. - let nodes = if leaf_hash == SHA2 { + let nodes = if leaf_hash == hash::SHA2 { let mut leaves = vec![Digest::default(); num_rows]; if self.hash_bn254_rows_sha2(num_cols, &mut leaves) { merkle.build_nodes(leaves) @@ -1260,7 +1251,7 @@ impl BufferOps for MetalBuffer { cpu_nodes() }; let root = nodes[merkle.num_nodes() - 1]; - (MetalBuffer::from_vec(nodes), root) + (BufferOps::from_vec(nodes), root) } } @@ -1279,7 +1270,7 @@ impl crate::algebra::buffer::Buffer for MetalBuffer { fn geometric_sequence(base: F, length: usize) -> Self { assert_bn254::(); - Self::from_vec(crate::algebra::geometric_sequence(base, length)) + BufferOps::from_vec(crate::algebra::geometric_sequence(base, length)) } fn random(rng: &mut R, length: usize) -> Self @@ -1288,7 +1279,7 @@ impl crate::algebra::buffer::Buffer for MetalBuffer { Standard: Distribution, { assert_bn254::(); - Self::from_vec((0..length).map(|_| rng.gen()).collect()) + BufferOps::from_vec((0..length).map(|_| rng.gen()).collect()) } fn zero_pad(&mut self) { @@ -1296,7 +1287,7 @@ impl crate::algebra::buffer::Buffer for MetalBuffer { if !self.is_empty() { let mut data = self.as_slice().to_vec(); data.resize(self.len().next_power_of_two(), F::ZERO); - *self = Self::from_vec(data); + *self = BufferOps::from_vec(data); } } @@ -1378,12 +1369,12 @@ impl crate::algebra::buffer::Buffer for MetalBuffer { let first = (first.as_mut() as &mut dyn std::any::Any) .downcast_mut::>() .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); - let mut accumulator = Self::from_slice(&first.vector); + let mut accumulator = >::from_slice(&first.vector); for (coeff, linear_form) in rlc_coeffs[1..].iter().zip(rest) { let covector = (linear_form.as_mut() as &mut dyn std::any::Any) .downcast_mut::>() .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); - let vector = Self::from_slice(&covector.vector); + let vector = >::from_slice(&covector.vector); vector.mixed_scalar_mul_add_to(&Identity::new(), &mut accumulator, *coeff); } accumulator @@ -1397,7 +1388,7 @@ impl crate::algebra::buffer::Buffer for MetalBuffer { assert_bn254::(); assert_eq!(vectors.len(), coeffs.len()); let Some((first, vectors)) = vectors.split_first() else { - return MetalBuffer::from_vec(Vec::new()); + return BufferOps::from_vec(Vec::new()); }; let mut accumulator = MetalBuffer:: { len: first.len(), @@ -1608,7 +1599,7 @@ impl BufferRead for MetalBuffer { fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { assert_bn254::(); - let (start, end) = crate::algebra::buffer::resolve_range(range, self.len()); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); MetalSlice { field: self.bn254_buffer(), offset: start, @@ -1625,7 +1616,7 @@ impl BufferWrite for MetalBuffer { Self: 'a, F: 'a; - fn scale(&mut self, weight: F) { + fn scalar_mul(&mut self, weight: F) { assert_bn254::(); if self.is_empty() { return; @@ -1658,7 +1649,7 @@ impl BufferWrite for MetalBuffer { fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { assert_bn254::(); - let (start, end) = crate::algebra::buffer::resolve_range(range, self.len()); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); MetalSliceMut::new(self, start, end - start) } @@ -1726,7 +1717,7 @@ impl ReedSolomon for MetalRs { assert_bn254::(); let num_messages = vectors.len() * interleaving_depth; if num_messages == 0 { - return MetalBuffer::from_vec(Vec::new()); + return BufferOps::from_vec(Vec::new()); } assert!(masks.len().is_multiple_of(num_messages)); let mask_length = masks.len() / num_messages; @@ -1977,7 +1968,7 @@ impl BufferRead for MetalSlice<'_, F> { } fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - let (start, end) = crate::algebra::buffer::resolve_range(range, self.len); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); MetalSlice { field: self.field.clone(), offset: self.offset + start, @@ -2042,7 +2033,7 @@ impl BufferRead for MetalSliceMut<'_, F> { } fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - let (start, end) = crate::algebra::buffer::resolve_range(range, self.len); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); MetalSlice { field: self.field.clone(), offset: self.offset + start, @@ -2059,7 +2050,7 @@ impl BufferWrite for MetalSliceMut<'_, F> { Self: 'a, F: 'a; - fn scale(&mut self, weight: F) { + fn scalar_mul(&mut self, weight: F) { assert_bn254::(); if self.len == 0 { return; @@ -2105,7 +2096,7 @@ impl BufferWrite for MetalSliceMut<'_, F> { } fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { - let (start, end) = crate::algebra::buffer::resolve_range(range, self.len); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); MetalSliceMut::from_field(self.field.clone(), self.offset + start, end - start) } @@ -3299,7 +3290,11 @@ fn as_field256_slice(values: &[T]) -> &[Field256] { #[cfg(test)] mod tests { use super::*; - use crate::algebra::{buffer::CpuBuffer, ntt::NttEngine}; + use crate::algebra::{ + buffer::{Buffer, CpuBuffer}, + ntt::NttEngine, + }; + use crate::buffer::BufferOps; fn values(len: usize, offset: u64) -> Vec { (0..len) diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index 8a4f782b..a01c957f 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -20,7 +20,7 @@ use super::utils::{lcm, sqrt_factor}; #[cfg(not(all(feature = "metal", target_os = "macos")))] use super::ReedSolomon; #[cfg(not(all(feature = "metal", target_os = "macos")))] -use crate::algebra::buffer::CpuBuffer; +use crate::algebra::buffer::{BufferOps, CpuBuffer}; #[cfg(not(feature = "rs_in_order"))] use crate::algebra::ntt::transpose::transpose_permute; use crate::{ @@ -233,7 +233,7 @@ impl NttEngine { } /// The shared field/domain core this engine is built on. - pub fn domain(&self) -> &RsDomain { + pub const fn domain(&self) -> &RsDomain { &self.domain } @@ -557,7 +557,7 @@ pub struct CpuRs { #[cfg(not(all(feature = "metal", target_os = "macos")))] impl CpuRs { - pub fn new(engine: Arc>) -> Self { + pub const fn new(engine: Arc>) -> Self { Self { engine } } } diff --git a/src/algebra/ntt/mod.rs b/src/algebra/ntt/mod.rs index f3ab5e08..50782dd6 100644 --- a/src/algebra/ntt/mod.rs +++ b/src/algebra/ntt/mod.rs @@ -11,11 +11,11 @@ use std::{ sync::{Arc, LazyLock}, }; -use ark_ff::{Field, FftField}; +use ark_ff::{FftField, Field}; use static_assertions::assert_obj_safe; -use self::matrix::MatrixMut; pub use self::cooley_tukey::{NttEngine, RsDomain}; +use self::matrix::MatrixMut; // The CPU encoder is only the active backend (and only constructible) on non-Metal builds. #[cfg(not(all(feature = "metal", target_os = "macos")))] pub use self::cooley_tukey::CpuRs; @@ -30,29 +30,29 @@ use crate::{ pub static NTT: LazyLock> = LazyLock::new(|| { let map = TypeMap::new(); - fn register(map: &TypeMap) { - // Both backends share the same `RsDomain` coset convention and differ only in how - // `interleaved_encode` runs: the CPU encoder needs the full NTT engine, while the - // Metal encoder only needs the shared domain (the NTT itself runs on the GPU). - #[cfg(all(feature = "metal", target_os = "macos"))] - let encoder = crate::algebra::metal_buffer::MetalRs::new(Arc::new( - RsDomain::::from_fftfield(), - )); - #[cfg(not(all(feature = "metal", target_os = "macos")))] - let encoder = CpuRs::new(Arc::new(NttEngine::::new_from_fftfield())); - map.insert(Arc::new(encoder) as Arc>); - } - register::(&map); - register::(&map); - register::(&map); - register::(&map); - register::(&map); - register::(&map); - register::<::BasePrimeField>(&map); - register::<::BasePrimeField>(&map); + register_ntt::(&map); + register_ntt::(&map); + register_ntt::(&map); + register_ntt::(&map); + register_ntt::(&map); + register_ntt::(&map); + register_ntt::<::BasePrimeField>(&map); + register_ntt::<::BasePrimeField>(&map); map }); +fn register_ntt(map: &TypeMap) { + // Both backends share the same `RsDomain` coset convention and differ only in how + // `interleaved_encode` runs: the CPU encoder needs the full NTT engine, while the + // Metal encoder only needs the shared domain (the NTT itself runs on the GPU). + #[cfg(all(feature = "metal", target_os = "macos"))] + let encoder = + crate::algebra::metal_buffer::MetalRs::new(Arc::new(RsDomain::::from_fftfield())); + #[cfg(not(all(feature = "metal", target_os = "macos")))] + let encoder = CpuRs::new(Arc::new(NttEngine::::new_from_fftfield())); + map.insert(Arc::new(encoder) as Arc>); +} + #[derive(Default)] pub struct NttFamily; diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index f8204c2b..e2bccebd 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -22,6 +22,8 @@ use whir::{ transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; +use whir::buffer::BufferOps; + #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { diff --git a/src/bin/gpu_proof_cpu_verify.rs b/src/bin/gpu_proof_cpu_verify.rs index 3a72cf66..a1f86362 100644 --- a/src/bin/gpu_proof_cpu_verify.rs +++ b/src/bin/gpu_proof_cpu_verify.rs @@ -1,4 +1,8 @@ -use std::{borrow::Cow, fs, path::PathBuf}; +use std::{ + borrow::Cow, + fs, + path::{Path, PathBuf}, +}; use clap::{Parser, Subcommand}; use serde::{Deserialize, Serialize}; @@ -12,6 +16,8 @@ use whir::{ transcript::{codecs::Empty, DomainSeparator, Proof, ProverState, VerifierState}, }; +use whir::buffer::BufferOps; + type F = Field256; type M = Identity; @@ -70,12 +76,12 @@ struct Artifact { fn main() { let args = Args::parse(); match args.command { - Command::Prove(args) => prove(args), - Command::Verify { input } => verify(input), + Command::Prove(args) => prove(&args), + Command::Verify { input } => verify(&input), } } -fn prove(args: RoundtripArgs) { +fn prove(args: &RoundtripArgs) { assert!(args.fold <= args.log_size, "fold must be <= log_size"); let size = 1usize << args.log_size; let whir_params = protocol_parameters(args.security_level, args.pow_bits, args.fold, args.rate); @@ -120,8 +126,8 @@ fn prove(args: RoundtripArgs) { ); } -fn verify(input: PathBuf) { - let bytes = fs::read(&input).expect("read proof artifact"); +fn verify(input: &Path) { + let bytes = fs::read(input).expect("read proof artifact"); let artifact: Artifact = serde_json::from_slice(&bytes).expect("decode proof artifact"); let size = 1usize << artifact.log_size; let whir_params = protocol_parameters( @@ -156,7 +162,7 @@ fn verify(input: PathBuf) { ); } -fn protocol_parameters( +const fn protocol_parameters( security_level: usize, pow_bits: usize, fold: usize, @@ -179,11 +185,11 @@ fn input_vector(size: usize) -> Vec { } #[cfg(all(feature = "metal", target_os = "macos"))] -fn backend_name() -> &'static str { +const fn backend_name() -> &'static str { "gpu-metal" } #[cfg(not(all(feature = "metal", target_os = "macos")))] -fn backend_name() -> &'static str { +const fn backend_name() -> &'static str { "cpu" } diff --git a/src/bin/main.rs b/src/bin/main.rs index 1c19922b..dc9dacae 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -18,6 +18,8 @@ use whir::{ transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; +use whir::buffer::BufferOps; + #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { @@ -324,7 +326,7 @@ where let _ = params.prove( &mut prover_state, &[&vector_buffer], - vec![&witness], + witness, prove_linear_forms, Cow::Borrowed(&evaluations), ); diff --git a/src/buffer/cpu/mod.rs b/src/buffer/cpu/mod.rs new file mode 100644 index 00000000..cf77b25c --- /dev/null +++ b/src/buffer/cpu/mod.rs @@ -0,0 +1,307 @@ +mod traits; + +use std::{any::Any, mem}; + +use ark_ff::Field; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; +use traits::{impl_cpu_read, impl_cpu_write}; + +use crate::{ + algebra::{ + embedding::Embedding, + linear_form::{Covector, LinearForm}, + }, + buffer::{Buffer, BufferOps, BufferRead, BufferWrite}, + engines::EngineId, + hash::{self, Hash}, + protocols::{ + matrix_commit::{hash_rows, Encodable}, + merkle_tree, + }, +}; + +#[derive( + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + Debug, + Default, + serde::Serialize, + serde::Deserialize, +)] +pub struct CpuBuffer { + pub(super) data: Vec, +} + +/// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. +pub struct CpuSlice<'a, F> { + pub(super) data: &'a [F], +} + +/// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. +pub struct CpuSliceMut<'a, F> { + pub(super) data: &'a mut [F], +} + +impl_cpu_read!(CpuSlice<'_, F>); +impl_cpu_read!(CpuSliceMut<'_, F>); +impl_cpu_read!(CpuBuffer); +impl_cpu_write!(CpuSliceMut<'_, F>); +impl_cpu_write!(CpuBuffer); + +impl BufferOps for CpuBuffer { + type Nodes = CpuBuffer; + + fn at_index(&self, index: usize) -> Option { + if index >= self.len() { + None + } else { + Some(self.data[index]) + } + } + + fn as_slice(&self) -> &[T] { + &self.data + } + + fn len(&self) -> usize { + self.data.len() + } + + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { + let mut result = Vec::with_capacity(indices.len() * num_cols); + for i in indices { + result.extend_from_slice(&self.data[i * num_cols..(i + 1) * num_cols]); + } + result + } + + fn from_vec(source: Vec) -> Self { + Self { data: source } + } + + fn from_slice(source: &[T]) -> Self { + Self { + data: Vec::from(source), + } + } + + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (Self::Nodes, Hash) + where + T: Encodable + Send + Sync, + { + assert_eq!(self.len(), num_cols * merkle.num_leaves); + let engine = hash::ENGINES + .retrieve(leaf_hash) + .expect("Failed to retrieve hash engine"); + #[cfg(feature = "tracing")] + tracing::Span::current().record("engine", engine.name().as_ref()); + + let mut leaves = vec![Hash::default(); merkle.num_leaves]; + hash_rows(&*engine, &self.data, &mut leaves); + let nodes = merkle.build_nodes(leaves); + let root = nodes[merkle.num_nodes() - 1]; + (CpuBuffer { data: nodes }, root) + } +} + +impl Buffer for CpuBuffer { + fn zeros(length: usize) -> Self { + Self { + data: vec![F::ZERO; length], + } + } + + fn geometric_sequence(base: F, length: usize) -> Self { + Self { + data: crate::algebra::geometric_sequence(base, length), + } + } + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + Self { + data: (0..length).map(|_| rng.gen()).collect(), + } + } + + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self { + assert_eq!(linear_forms.len(), rlc_coeffs.len()); + let mut covector = vec![F::ZERO; size]; + if let Some((first, linear_forms)) = linear_forms.split_first_mut() { + debug_assert_eq!(rlc_coeffs[0], F::ONE); + if let Some(covector_form) = + (first.as_mut() as &mut dyn Any).downcast_mut::>() + { + mem::swap(&mut covector, &mut covector_form.vector); + } else { + first.accumulate(&mut covector, F::ONE); + } + for (rlc_coeff, linear_form) in rlc_coeffs[1..].iter().zip(linear_forms) { + linear_form.accumulate(&mut covector, *rlc_coeff); + } + } + Self { data: covector } + } + + fn zero_pad(&mut self) { + if !self.is_empty() { + self.data.resize(self.len().next_power_of_two(), F::ZERO); + } + } + + fn fold(&mut self, weight: F) { + crate::algebra::sumcheck::fold(&mut self.data, weight); + } + + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); + } + + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) + } + + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer { + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return CpuBuffer { data: Vec::new() }; + }; + debug_assert_eq!(coeffs[0], M::Target::ONE); + let mut accumulator = crate::algebra::lift(embedding, first.as_slice()); + for (coeff, vector) in coeffs[1..].iter().zip(vectors) { + crate::algebra::mixed_scalar_mul_add( + embedding, + &mut accumulator, + *coeff, + vector.as_slice(), + ); + } + CpuBuffer { data: accumulator } + } +} + +/// Resolve any range expression (`a..b`, `..b`, `a..`, `..`) against `len`, +/// returning `(start, end)` and bounds-checking like slice indexing. +pub(crate) fn resolve_range( + range: impl std::ops::RangeBounds, + len: usize, +) -> (usize, usize) { + use std::ops::Bound::{Excluded, Included, Unbounded}; + let start = match range.start_bound() { + Included(&s) => s, + Excluded(&s) => s + 1, + Unbounded => 0, + }; + let end = match range.end_bound() { + Included(&e) => e + 1, + Excluded(&e) => e, + Unbounded => len, + }; + assert!( + start <= end && end <= len, + "slice range {start}..{end} out of bounds for length {len}" + ); + (start, end) +} + +#[cfg(test)] +mod tests { + use ark_ff::AdditiveGroup; + + use super::*; + use crate::algebra::{ + fields::Field64, geometric_accumulate, geometric_sequence as geometric_sequence_vec, + linear_form::UnivariateEvaluation, + }; + + type F = Field64; + + #[test] + fn geometric_sequence_matches_free_function() { + let base = F::from(7u64); + for length in 0..6 { + let buffer = CpuBuffer::::geometric_sequence(base, length); + assert_eq!( + BufferOps::as_slice(&buffer), + geometric_sequence_vec(base, length).as_slice(), + "geometric_sequence mismatch at length {length}" + ); + } + } + + #[test] + fn scale_multiplies_in_place() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.scalar_mul(weight); + let expected: Vec = values.iter().map(|&v| v * weight).collect(); + assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); + } + + #[test] + fn slice_accumulate_matches_restricted_geometric_accumulate() { + let len = 8usize; + let points = vec![F::from(3u64), F::from(5u64)]; + let scalars = vec![F::from(2u64), F::from(9u64)]; + + // Try a prefix slice (offset 0) and an interior slice (offset > 0). + for (offset, window_len) in [(0usize, 5usize), (2usize, 4usize)] { + let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); + let evaluators: Vec<_> = points + .iter() + .map(|&point| UnivariateEvaluation::new(point, window_len)) + .collect(); + buffer + .slice_mut(offset..offset + window_len) + .accumulate_univariate_evaluations(&evaluators, &scalars); + + // Reference: accumulate Σ_j scalars[j]·points[j]^i into the same + // sub-slice of a plain vector. + let mut reference = vec![F::ZERO; len]; + geometric_accumulate( + &mut reference[offset..offset + window_len], + scalars.clone(), + &points, + ); + + assert_eq!( + BufferOps::as_slice(&buffer), + reference.as_slice(), + "slice accumulate mismatch at offset {offset}, len {window_len}" + ); + } + } + + #[test] + fn slice_scale_multiplies_only_the_range() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.slice_mut(1..3).scalar_mul(weight); + let expected = vec![values[0], values[1] * weight, values[2] * weight, values[3]]; + assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); + } +} diff --git a/src/buffer/cpu/traits.rs b/src/buffer/cpu/traits.rs new file mode 100644 index 00000000..fdeec2d6 --- /dev/null +++ b/src/buffer/cpu/traits.rs @@ -0,0 +1,117 @@ +macro_rules! impl_cpu_read { + ($ty:ty) => { + impl BufferRead for $ty { + type TargetBuffer = CpuBuffer; + type Slice<'a> + = CpuSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.data.len() + } + + fn dot(&self, other: &Self) -> F { + crate::algebra::dot(&*self.data, &*other.data) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + crate::algebra::sumcheck::compute_sumcheck_polynomial(&*self.data, &*other.data) + } + + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + crate::algebra::mixed_multilinear_extend(embedding, &*self.data, point) + } + + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &CpuBuffer, + ) -> M::Target { + crate::algebra::mixed_dot(embedding, other.as_slice(), &*self.data) + } + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + crate::algebra::mixed_univariate_evaluate(embedding, &*self.data, point) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut CpuBuffer, + weight: M::Target, + ) { + crate::algebra::mixed_scalar_mul_add( + embedding, + &mut accumulator.data, + weight, + &*self.data, + ); + } + + fn slice(&self, range: impl std::ops::RangeBounds) -> CpuSlice<'_, F> { + let data = &*self.data; + let (start, end) = $crate::buffer::cpu::resolve_range(range, data.len()); + CpuSlice { + data: &data[start..end], + } + } + } + }; +} + +macro_rules! impl_cpu_write { + ($ty:ty) => { + impl BufferWrite for $ty { + type SliceMut<'a> + = CpuSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scalar_mul(&mut self, weight: F) { + crate::algebra::scalar_mul(&mut *self.data, weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[crate::algebra::linear_form::UnivariateEvaluation], + scalars: &[F], + ) { + crate::algebra::linear_form::UnivariateEvaluation::accumulate_many( + evaluators, + &mut *self.data, + scalars, + ); + } + + fn slice_mut( + &mut self, + range: impl std::ops::RangeBounds, + ) -> CpuSliceMut<'_, F> { + let data = &mut *self.data; + let (start, end) = $crate::buffer::cpu::resolve_range(range, data.len()); + CpuSliceMut { + data: &mut data[start..end], + } + } + + fn split_at_mut(&mut self, mid: usize) -> (CpuSliceMut<'_, F>, CpuSliceMut<'_, F>) { + let (lo, hi) = self.data.split_at_mut(mid); + (CpuSliceMut { data: lo }, CpuSliceMut { data: hi }) + } + } + }; +} + +pub(crate) use impl_cpu_read; +pub(crate) use impl_cpu_write; diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs new file mode 100644 index 00000000..32b87b8f --- /dev/null +++ b/src/buffer/mod.rs @@ -0,0 +1,193 @@ +pub mod cpu; +use std::ops::RangeBounds; + +use ark_ff::Field; +use ark_std::rand::{ + distributions::{Distribution, Standard}, + CryptoRng, RngCore, +}; + +use crate::{ + algebra::{ + embedding::Embedding, + linear_form::{LinearForm, UnivariateEvaluation}, + }, + engines::EngineId, + hash::Hash, + protocols::{matrix_commit::Encodable, merkle_tree}, +}; + +pub trait BufferOps { + type Nodes: BufferOps; + + fn from_vec(source: Vec) -> Self; + fn from_slice(source: &[T]) -> Self; + fn as_slice(&self) -> &[T]; + fn num_rows(&self, num_cols: usize) -> usize { + self.len() / num_cols + } + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; + fn at_index(&self, index: usize) -> Option; + fn len(&self) -> usize; + fn is_empty(&self) -> bool { + self.len() == 0 + } + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (Self::Nodes, Hash) + where + T: Encodable + Send + Sync; +} + +/// Read-only operations over a contiguous run of field elements — the buffer +/// analogue of `&[F]`. +/// +/// Implemented by the owned buffers ([`CpuBuffer`]/`MetalBuffer`, as the +/// full-range case) and by the borrowed read views ([`CpuSlice`]/`MetalSlice`). +/// Every op here works identically on a whole buffer or any sub-range obtained +/// via [`Self::slice`]. On `MetalBuffer` a view aliases the parent's GPU +/// allocation (byte-offset binding), so no data is copied. +pub trait BufferRead { + /// A same-backend owned buffer over another field, produced by the + /// mixed-field ops (e.g. `mixed_dot` against a target-field buffer). For an + /// owned buffer this is the buffer itself over `T`; views report the owning + /// buffer type. + type TargetBuffer: Buffer; + + /// Read-only view type produced by slicing this buffer/view. + type Slice<'a>: BufferRead + where + Self: 'a, + F: 'a; + + /// Number of elements in this view. + fn read_len(&self) -> usize; + fn read_is_empty(&self) -> bool { + self.read_len() == 0 + } + + /// Inner product with another view of the same length. + fn dot(&self, other: &Self) -> F + where + Self: Sized; + + /// Sumcheck round coefficients `(c0, c2)` for `dot(self, other)`. + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) + where + Self: Sized; + + /// Multilinear extension evaluated at a target-field point. + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target; + + /// Inner product with a target-field buffer. + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &Self::TargetBuffer, + ) -> M::Target; + + /// Univariate evaluation at a target-field point. + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target; + + /// `accumulator += weight * self`, lifted into the target field. + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut Self::TargetBuffer, + weight: M::Target, + ); + + /// Borrow `self[range]` as a read-only view (analogous to `&v[range]`). + fn slice(&self, range: impl RangeBounds) -> Self::Slice<'_>; +} + +/// In-place, length-preserving operations over a contiguous run of field +/// elements — the buffer analogue of `&mut [F]`. +/// +/// Implemented by the owned buffers (full-range) and the borrowed mutable views +/// ([`CpuSliceMut`]/`MetalSliceMut`). A sub-range obtained via [`Self::slice_mut`] +/// is the same view type as the whole buffer, so the ops compose just like +/// slicing a `Vec`. +/// +/// Constructors (`zeros`, `random`, …) and length-changing operations (`fold`, +/// `zero_pad`) live on [`Buffer`], the owned-buffer trait, because a borrowed +/// view cannot construct or resize its backing storage. +pub trait BufferWrite: BufferRead { + /// Mutable view type produced by slicing this buffer/view. + type SliceMut<'a>: BufferWrite + where + Self: 'a, + F: 'a; + + /// In-place scalar multiplication: `self[i] *= weight`. + fn scalar_mul(&mut self, weight: F); + + /// Accumulate `Σ_j scalars[j] · evaluators[j].point^i` into entry `i`. + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ); + + /// Borrow `self[range]` as a mutable, zero-copy view (analogous to + /// `&mut v[range]`). + fn slice_mut(&mut self, range: impl RangeBounds) -> Self::SliceMut<'_>; + + /// Split into two disjoint mutable views `[0, mid)` and `[mid, len)`. + fn split_at_mut(&mut self, mid: usize) -> (Self::SliceMut<'_>, Self::SliceMut<'_>); +} + +/// Owned field-buffer operations — the `Vec` half of the field API. +/// +/// `BufferOps` stays generic over any element type, so hashes and digests can +/// use it. `BufferRead`/`BufferWrite` are the slice-like field ops implemented +/// by owners and views. This trait is only for owned field buffers: operations +/// here construct storage or change its length, so borrowed views cannot +/// implement them. +pub trait Buffer: BufferOps + BufferWrite + Clone { + fn zeros(length: usize) -> Self; + + /// Geometric sequence `[1, base, base², …, base^(length-1)]`. + fn geometric_sequence(base: F, length: usize) -> Self; + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution; + + fn zero_pad(&mut self); + fn fold(&mut self, weight: F); + + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); + } + + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) + } + + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self; + + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer; +} diff --git a/src/lib.rs b/src/lib.rs index 7d1d8b45..4c1a5c3d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ pub mod algebra; pub mod ark_serde; pub mod bits; +pub mod buffer; pub mod cmdline_utils; pub mod engines; pub mod hash; diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 220ea28d..4cefdd25 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -83,7 +83,7 @@ impl Config { // Even more trivial non-zk protocol: send f and r directly. if !self.masked { // TODO: avoid these big readbacks even though they are transcript-sized - prover_state.prover_messages(&*vector.as_slice()); + prover_state.prover_messages(vector.as_slice()); prover_state.prover_messages(witness.masks.as_slice()); let _ = self.commit.open(prover_state, &[witness]); let point = self @@ -91,7 +91,7 @@ impl Config { .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) .round_challenges; assert!( - !vector.at_index(0).unwrap_or(F::one()).is_zero(), + !vector.at_index(0).unwrap_or_else(F::one).is_zero(), "Proof failed" ); return Opening { @@ -115,7 +115,7 @@ impl Config { assert!(!mask_rlc.is_zero(), "Proof failed"); let mut masked_vector = mask.clone(); vector.mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_vector, mask_rlc); - prover_state.prover_messages(&*masked_vector.as_slice()); + prover_state.prover_messages(masked_vector.as_slice()); // Send combined IRS randomness. (r^* in paper) let masked_masks = scalar_mul_add_new( @@ -146,7 +146,7 @@ impl Config { // no constraints on l(r) that the verifier can return. // This event is cryptographically unlikely as `F` is challenge sized. assert!( - !masked_vector.at_index(0).unwrap_or(F::one()).is_zero(), + !masked_vector.at_index(0).unwrap_or_else(F::one).is_zero(), "Proof failed" ); diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 00b23e2c..e1108081 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -1,5 +1,3 @@ -// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. - //! Code-switching IOR: R_{C, C_zk, sl} → R_{C', C_zk, sl'} //! //! Reduces a proximity claim about oracle f (source code C) to a proximity @@ -62,9 +60,8 @@ pub enum MaskInput { Enabled(ActiveBuffer), } -// helper function to get UnivariateEvaluations #[inline] -fn evaluators(points: &[F], size: usize) -> Vec> { +fn univariate_evaluators(points: &[F], size: usize) -> Vec> { points .iter() .map(|&point| UnivariateEvaluation::new(point, size)) @@ -170,7 +167,7 @@ impl Config { message: ActiveBuffer, witness: &IrsWitness, covector: &mut ActiveBuffer, - folding_randomness: &mut ActiveBuffer, + folding_randomness: &ActiveBuffer, mask_input: &MaskInput, ) -> Witness where @@ -239,21 +236,21 @@ impl Config { let (ood_rlc_coeffs, in_domain_rlc_coeffs) = constraint_rlc_coeffs.split_at(num_ood); // Covector update — sl' from Completeness proof (p.55-56). - covector.scale(original_sl_coeff); + covector.scalar_mul(original_sl_coeff); let eval_points = lift(self.source.embedding(), &source_evaluations.points); if self.message_mask_length == 0 { // Non-ZK: single accumulate over all points. Scalars are the // contiguous challenge tail (ood ‖ in-domain), so no copy is needed. - let mut all_evaluators = evaluators(&ood_points, covector.len()); - all_evaluators.extend(evaluators(&eval_points, covector.len())); + let mut all_evaluators = univariate_evaluators(&ood_points, covector.len()); + all_evaluators.extend(univariate_evaluators(&eval_points, covector.len())); covector.accumulate_univariate_evaluations(&all_evaluators, constraint_rlc_coeffs); } else { // ZK: OOD contributes to full [f; r; s], in-domain only to [f; r]. - let ood_evaluators = evaluators(&ood_points, covector.len()); + let ood_evaluators = univariate_evaluators(&ood_points, covector.len()); covector.accumulate_univariate_evaluations(&ood_evaluators, ood_rlc_coeffs); let masked = self.source.masked_message_length(); - let in_domain_evaluators = evaluators(&eval_points, masked); + let in_domain_evaluators = univariate_evaluators(&eval_points, masked); covector .slice_mut(..masked) .accumulate_univariate_evaluations(&in_domain_evaluators, in_domain_rlc_coeffs); @@ -529,7 +526,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); - let covector = ActiveBuffer::from_vec(covector); + let mut covector = ActiveBuffer::from_vec(covector); let instance = U64(seed); let ds = DomainSeparator::protocol(config) @@ -553,8 +550,8 @@ mod tests { &mut prover_state, folded_message.clone(), &source_witness, - &mut covector.clone(), - &mut ActiveBuffer::from_vec(folding_randomness.clone()), + &mut covector, + &ActiveBuffer::from_vec(folding_randomness.clone()), &mask_input(ActiveBuffer::from_vec(mask_msg)), ); let proof = prover_state.proof(); @@ -628,7 +625,7 @@ mod tests { folded_message, &source_witness, &mut covector, - &mut ActiveBuffer::from_vec(folding_randomness.clone()), + &ActiveBuffer::from_vec(folding_randomness.clone()), &mask_input(ActiveBuffer::from_vec(mask_msg)), ); let proof = prover_state.proof(); @@ -677,11 +674,11 @@ mod tests { fold_chunks(&f_full, config.source.message_length(), &folding_randomness); // For non-ZK and source.mask_length == 0, h = folded_message and identity holds. - let folded_message_buffer = ActiveBuffer::from_vec(folded_message.clone()); + let folded_message_buffer = ActiveBuffer::from_slice(&folded_message); let initial_mu = folded_message_buffer.dot(&covector); // Tamper the post-fold message before proving. - let mut tampered = folded_message.clone(); + let mut tampered = folded_message; tampered[0] += F::ONE; let tampered = ActiveBuffer::from_vec(tampered); let _witness = config.prove( @@ -689,7 +686,7 @@ mod tests { tampered, &source_witness, &mut covector, - &mut ActiveBuffer::from_vec(folding_randomness.clone()), + &ActiveBuffer::from_vec(folding_randomness.clone()), &MaskInput::Disabled, ); let proof = prover_state.proof(); diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index ae88e428..63ecb37a 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -1,5 +1,3 @@ -// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. - //! Mask proximity verification via γ-combination. //! //! Implements Construction 7.2 (p.43-44) specialized for zero-constraint mask @@ -48,8 +46,9 @@ use serde::{Deserialize, Serialize}; use crate::{ algebra::{ - buffer::ActiveBuffer, embedding::Identity, random_vector, scalar_mul_add_new, - univariate_evaluate, + buffer::{ActiveBuffer, Buffer, BufferOps}, + embedding::Identity, + scalar_mul_add_new, univariate_evaluate, }, hash::Hash, protocols::irs_commit::{ @@ -77,7 +76,7 @@ pub struct Config { #[must_use] pub struct Witness { pub(crate) mask_witness: IrsWitness, - pub(crate) fresh_msgs: Vec>, + pub(crate) fresh_msgs: Vec>, } /// Verifier output from the commit phase. @@ -115,7 +114,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - original_msgs: &[Vec], + original_msgs: &[&ActiveBuffer], ) -> Witness where F: Codec<[H::U]>, @@ -130,23 +129,15 @@ impl Config { } // Sample fresh mask-of-masks - let fresh_msgs: Vec> = (0..self.num_masks) - .map(|_| random_vector(prover_state.rng(), self.c_zk_commit.vector_size)) + let fresh_msgs: Vec> = (0..self.num_masks) + .map(|_| ActiveBuffer::random(prover_state.rng(), self.c_zk_commit.vector_size)) .collect(); - let original_buffers = original_msgs - .iter() - .map(|msg| ActiveBuffer::from_slice(msg)) - .collect::>(); - let fresh_buffers = fresh_msgs - .iter() - .map(|msg| ActiveBuffer::from_slice(msg)) - .collect::>(); - // Tree layout: [originals..., freshes...] - let all_vectors: Vec<&ActiveBuffer> = original_buffers + let all_vectors: Vec<&ActiveBuffer> = original_msgs .iter() - .chain(fresh_buffers.iter()) + .copied() + .chain(fresh_msgs.iter()) .collect(); let mask_witness = self.c_zk_commit.commit(prover_state, &all_vectors); @@ -175,7 +166,7 @@ impl Config { &self, prover_state: &mut ProverState, witness: &Witness, - original_msgs: &[Vec], + original_msgs: &[&ActiveBuffer], ) where F: Codec<[H::U]>, H: DuplexSpongeInterface, @@ -197,11 +188,13 @@ impl Config { assert_eq!(irs_masks.len(), 2 * self.num_masks * irs_masks_per_vector); for (i, (orig_msg, fresh_msg)) in original_msgs .iter() + .copied() .zip(witness.fresh_msgs.iter()) .enumerate() { // ξ*_i = s_i + γ · ξ_i - let combined_msg = scalar_mul_add_new(fresh_msg, gamma, orig_msg); + // TODO: port to buffer backend + let combined_msg = scalar_mul_add_new(fresh_msg.as_slice(), gamma, orig_msg.as_slice()); prover_state.prover_messages(&combined_msg); // r*_i = r'_i + γ · r_i @@ -352,10 +345,15 @@ mod tests { let original_msgs: Vec> = (0..config.num_masks) .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let original_refs = original_buffers.iter().collect::>(); let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit(&mut prover_state, &original_msgs); - config.prove(&mut prover_state, &witness, &original_msgs); + let witness = config.commit(&mut prover_state, &original_refs); + config.prove(&mut prover_state, &witness, &original_refs); let proof = prover_state.proof(); let mut verifier_state = VerifierState::new_std(&ds, &proof); @@ -426,13 +424,23 @@ mod tests { let original_msgs: Vec> = (0..config.num_masks) .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let original_refs = original_buffers.iter().collect::>(); let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit(&mut prover_state, &original_msgs); + let witness = config.commit(&mut prover_state, &original_refs); let mut tampered_msgs = original_msgs; tampered_msgs[0][0] += F::ONE; - config.prove(&mut prover_state, &witness, &tampered_msgs); + let tampered_buffers = tampered_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let tampered_refs = tampered_buffers.iter().collect::>(); + config.prove(&mut prover_state, &witness, &tampered_refs); let proof = prover_state.proof(); let mut verifier_state = VerifierState::new_std(&ds, &proof); @@ -457,21 +465,28 @@ mod tests { let original_msgs: Vec> = (0..config.num_masks) .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size)) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let original_refs = original_buffers.iter().collect::>(); let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit(&mut prover_state, &original_msgs); + let witness = config.commit(&mut prover_state, &original_refs); let gamma: F = prover_state.verifier_message(); let irs_masks_per_vector = config.c_zk_commit.mask_length * config.c_zk_commit.interleaving_depth; let irs_masks = witness.mask_witness.masks.as_slice(); - for (i, (orig_msg, fresh_msg)) in original_msgs + for (i, (orig_msg, fresh_msg)) in original_refs .iter() + .copied() .zip(witness.fresh_msgs.iter()) .enumerate() { - let mut combined_msg = scalar_mul_add_new(fresh_msg, gamma, orig_msg); + let mut combined_msg = + scalar_mul_add_new(fresh_msg.as_slice(), gamma, orig_msg.as_slice()); if i == 0 { combined_msg[0] += F::ONE; } diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index bada92b6..d725ed41 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -12,7 +12,8 @@ use tracing::instrument; use zerocopy::{Immutable, IntoBytes}; use crate::{ - algebra::buffer::{ActiveBuffer, BufferOps}, + algebra::buffer::ActiveBuffer, + buffer::BufferOps, engines::EngineId, hash::{self, Hash}, protocols::merkle_tree, @@ -223,6 +224,7 @@ impl Config { H: DuplexSpongeInterface, R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>, + ActiveBuffer: BufferOps>, { assert_eq!(matrix.len(), self.num_rows() * self.num_cols); diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 2a8506f7..80f64237 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -131,6 +131,7 @@ mod tests { use ark_std::rand::thread_rng; use super::*; + use crate::buffer::BufferOps; use crate::{ algebra::{ buffer::ActiveBuffer, diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index d8c09d2b..d84f7555 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -1,6 +1,5 @@ use std::borrow::Cow; -use crate::algebra::buffer::{Buffer, BufferRead, BufferWrite}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; #[cfg(feature = "tracing")] @@ -9,7 +8,7 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ algebra::{ - buffer::{ActiveBuffer, BufferOps}, + buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead, BufferWrite}, dot, embedding::Embedding, eq_weights, diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index c8c73c07..3d525a91 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -6,6 +6,7 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution}; use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; +use crate::buffer::BufferOps; use crate::{ algebra::buffer::ActiveBuffer, hash::Hash, diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index 0f2ed6f2..960f5a3c 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -361,10 +361,7 @@ mod tests { let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|&v| Cow::Borrowed(v)) - .collect::>(), + &vector_refs, witness, prove_forms, Cow::Borrowed(evaluations), @@ -458,10 +455,7 @@ mod tests { let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|&v| Cow::Borrowed(v)) - .collect::>(), + &vector_refs, witness, prove_forms, Cow::Borrowed(&evaluations), @@ -517,10 +511,7 @@ mod tests { let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|&v| Cow::Borrowed(v)) - .collect::>(), + &vector_refs, witness, prove_forms, Cow::Borrowed(&evaluations), @@ -579,7 +570,7 @@ mod tests { let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(&vector)], + &[&vector_buffer], witness, prove_forms, Cow::Owned(vec![wrong_evaluation]), diff --git a/src/protocols/whir_zk/prover.rs b/src/protocols/whir_zk/prover.rs index d1527c13..a7512534 100644 --- a/src/protocols/whir_zk/prover.rs +++ b/src/protocols/whir_zk/prover.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + use std::borrow::Cow; use ark_ff::Field; @@ -8,8 +10,10 @@ use rayon::prelude::*; use tracing::instrument; use super::{Config, Witness}; +use crate::buffer::BufferOps; use crate::{ algebra::{ + buffer::ActiveBuffer, embedding::Identity, linear_form::{Covector, Evaluate, LinearForm}, mixed_dot, scalar_mul_add, @@ -209,7 +213,7 @@ impl Config { pub fn prove<'a, H, R>( &self, prover_state: &mut ProverState, - vectors: Vec>, + vectors: &[&ActiveBuffer], witness: Witness, linear_forms: Vec>>, evaluations: Cow<'a, [F]>, @@ -265,7 +269,6 @@ impl Config { let num_witness_variables = self.num_witness_variables(); let num_blinding_variables = self.num_blinding_variables(); let num_witness_variables_plus_1 = num_witness_variables + 1; - drop(vectors); // TODO: These are never touched? // Compute w_folded evaluations of all blinding vectors before rho for binding. let (w_folded_weights, m_evals, w_folded_blinding_evals) = { @@ -396,10 +399,16 @@ impl Config { let result = { #[cfg(feature = "tracing")] let _span = tracing::info_span!("inner_blinded_prove").entered(); + let f_hat_buffers = f_hat_vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let f_hat_refs = f_hat_buffers.iter().collect::>(); + let f_hat_witness_refs = f_hat_witnesses.iter().collect::>(); self.blinded_commitment.prove( prover_state, - f_hat_vectors.into_iter().map(Cow::Owned).collect(), - f_hat_witnesses.into_iter().map(Cow::Owned).collect(), + &f_hat_refs, + f_hat_witness_refs, linear_forms, Cow::Owned(modified_evaluations), ) @@ -423,10 +432,15 @@ impl Config { .concat(); // Blinding sub-proof result is discarded: the blinding WHIR's // evaluation point is not needed by the outer protocol. + let blinding_buffers = blinding_vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let blinding_refs = blinding_buffers.iter().collect::>(); let _ = self.blinding_commitment.prove( prover_state, - blinding_vectors.into_iter().map(Cow::Owned).collect(), - vec![Cow::Owned(blinding_witness)], + &blinding_refs, + vec![&blinding_witness], blinding_forms, Cow::Owned(all_blinding_claims), ); From e5bba9b80207ae178005dd93f8f493e7bc710816 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Fri, 19 Jun 2026 14:59:19 +0200 Subject: [PATCH 20/35] refactor: overall code reorganization --- benches/expand_from_coeff.rs | 21 +- proptest-regressions/protocols/irs_commit.txt | 8 - src/algebra/buffer.rs | 19 - src/algebra/metal_buffer.rs | 3402 ----------------- src/algebra/mod.rs | 3 - src/algebra/ntt/cooley_tukey.rs | 370 +- src/algebra/ntt/mod.rs | 136 +- src/bin/benchmark.rs | 3 +- src/bin/gpu_proof_cpu_verify.rs | 6 +- src/bin/main.rs | 3 +- src/buffer/cpu/buffer.rs | 312 ++ src/buffer/cpu/mod.rs | 310 +- src/buffer/cpu/{traits.rs => read_write.rs} | 4 + src/buffer/metal/buffer.rs | 1291 +++++++ src/buffer/metal/kernels.rs | 943 +++++ src/buffer/metal/mod.rs | 16 + .../metal/profile.rs} | 7 +- src/buffer/metal/rs.rs | 173 + src/buffer/metal/runtime.rs | 1218 ++++++ .../metal/sha2.rs} | 21 +- src/buffer/mod.rs | 92 +- src/hash/mod.rs | 11 +- src/protocols/basecase.rs | 7 +- src/protocols/code_switch.rs | 2 +- src/protocols/irs_commit.rs | 31 +- src/protocols/mask_proximity.rs | 51 +- src/protocols/matrix_commit.rs | 65 +- src/protocols/merkle_tree.rs | 69 +- src/protocols/sumcheck.rs | 9 +- src/protocols/whir/mod.rs | 10 +- src/protocols/whir/prover.rs | 10 +- src/protocols/whir_zk/committer.rs | 2 +- 32 files changed, 4375 insertions(+), 4250 deletions(-) delete mode 100644 proptest-regressions/protocols/irs_commit.txt delete mode 100644 src/algebra/buffer.rs delete mode 100644 src/algebra/metal_buffer.rs create mode 100644 src/buffer/cpu/buffer.rs rename src/buffer/cpu/{traits.rs => read_write.rs} (96%) create mode 100644 src/buffer/metal/buffer.rs create mode 100644 src/buffer/metal/kernels.rs create mode 100644 src/buffer/metal/mod.rs rename src/{hash/metal_profile.rs => buffer/metal/profile.rs} (97%) create mode 100644 src/buffer/metal/rs.rs create mode 100644 src/buffer/metal/runtime.rs rename src/{hash/metal_sha2_engine.rs => buffer/metal/sha2.rs} (97%) diff --git a/benches/expand_from_coeff.rs b/benches/expand_from_coeff.rs index 309a4287..862da207 100644 --- a/benches/expand_from_coeff.rs +++ b/benches/expand_from_coeff.rs @@ -1,5 +1,10 @@ use divan::{black_box, AllocProfiler, Bencher}; -use whir::algebra::{fields::Field64, ntt::NttEngine, random_vector}; +use whir::algebra::{ + buffer::{ActiveBuffer, BufferOps}, + fields::Field64, + ntt::{Messages, NttEngine, ReedSolomon}, + random_vector, +}; #[global_allocator] static ALLOC: AllocProfiler = AllocProfiler::system(); @@ -28,19 +33,17 @@ fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) { let message_length = 1 << (exp - coset_sz); let num_messages = 1 << coset_sz; let mut rng = ark_std::rand::thread_rng(); - let coeffs: Vec> = (0..num_messages) - .map(|_| random_vector(&mut rng, message_length)) + let coeffs: Vec> = (0..num_messages) + .map(|_| ActiveBuffer::from_vec(random_vector(&mut rng, message_length))) .collect(); let engine = NttEngine::::new_from_fftfield(); (engine, coeffs, expansion) }) .bench_values(|(engine, coeffs, expansion)| { - let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::>(); - black_box(engine.interleaved_encode_slices( - &coeffs_refs, - &[], - coeffs[0].len() * expansion, - )) + let coeffs_refs = coeffs.iter().collect::>(); + let messages = Messages::new(&coeffs_refs, coeffs[0].len(), 1); + let masks = ActiveBuffer::from_slice(&[]); + black_box(engine.interleaved_encode(messages, &masks, coeffs[0].len() * expansion)) }); } diff --git a/proptest-regressions/protocols/irs_commit.txt b/proptest-regressions/protocols/irs_commit.txt deleted file mode 100644 index c02b8454..00000000 --- a/proptest-regressions/protocols/irs_commit.txt +++ /dev/null @@ -1,8 +0,0 @@ -# Seeds for failure cases proptest has generated in the past. It is -# automatically read and these particular cases re-run before any -# novel cases are generated. -# -# It is recommended to check this file in to source control so that -# everyone who runs the test benefits from these saved cases. -cc 33aef60251379abf1e94f5821e0fbd4f69d0b742576f047a92921eeb9438c0f8 # shrinks to seed = 0, config = Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 0, vector_size: 1, mask_length: 0, codeword_length: 1, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 0, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 1, layers: [] } }, johnson_slack: 0.0, in_domain_samples: 0, out_domain_samples: 0, deduplicate_in_domain: false } -cc 2c8600bcbca6f088a51752815a2baef31e96a3c8ae7e516c28b3a7c407904c3f # shrinks to seed = 0, config = Config { embedding: Basefield { extension: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 2 } }, num_vectors: 0, vector_size: 1, mask_length: 0, codeword_length: 1, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 0, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 1, layers: [] } }, johnson_slack: 0.0, in_domain_samples: 0, out_domain_samples: 0, deduplicate_in_domain: false } diff --git a/src/algebra/buffer.rs b/src/algebra/buffer.rs deleted file mode 100644 index aa11ed17..00000000 --- a/src/algebra/buffer.rs +++ /dev/null @@ -1,19 +0,0 @@ -pub use crate::buffer::cpu::{CpuBuffer, CpuSlice, CpuSliceMut}; -pub use crate::buffer::{Buffer, BufferOps, BufferRead, BufferWrite}; - -#[cfg(all(feature = "metal", target_os = "macos"))] -pub use super::metal_buffer::{MetalBuffer, MetalSlice, MetalSliceMut}; - -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type ActiveBuffer = MetalBuffer; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type ActiveSlice<'a, T> = MetalSlice<'a, T>; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type ActiveSliceMut<'a, T> = MetalSliceMut<'a, T>; - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -pub type ActiveBuffer = CpuBuffer; -#[cfg(not(all(feature = "metal", target_os = "macos")))] -pub type ActiveSlice<'a, T> = CpuSlice<'a, T>; -#[cfg(not(all(feature = "metal", target_os = "macos")))] -pub type ActiveSliceMut<'a, T> = CpuSliceMut<'a, T>; diff --git a/src/algebra/metal_buffer.rs b/src/algebra/metal_buffer.rs deleted file mode 100644 index 6334797e..00000000 --- a/src/algebra/metal_buffer.rs +++ /dev/null @@ -1,3402 +0,0 @@ -use std::{ - any::type_name, - cell::OnceCell, - cmp::Ordering, - collections::HashMap, - hash::{Hash, Hasher}, - marker::PhantomData, - ops::RangeBounds, - os::raw::c_void, - sync::{Arc, Mutex, OnceLock}, - time::Instant, -}; - -use ark_ff::{AdditiveGroup, BigInt, FftField, Field, Fp, MontBackend}; -use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; -use metal::{ - Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, MTLResourceOptions, MTLSize, -}; -use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use zerocopy::IntoBytes; - -use crate::{ - algebra::{ - buffer::{BufferOps, BufferRead, BufferWrite}, - embedding::{Embedding, Identity}, - fields::{BN254Config, Field256}, - linear_form::{Covector, LinearForm, UnivariateEvaluation}, - ntt::{ReedSolomon, RsDomain}, - }, - engines::EngineId, - hash::{self, metal_profile, Hash as Digest, MetalSha2}, - protocols::{ - matrix_commit::{hash_rows, Encodable}, - merkle_tree, - }, -}; - -const METAL_SOURCE: &str = r#" -#include -using namespace metal; - -struct F { ulong v[4]; }; - -constant ulong MODULUS[4] = { - 0x43e1f593f0000001UL, - 0x2833e84879b97091UL, - 0xb85045b68181585dUL, - 0x30644e72e131a029UL, -}; - -constant ulong BN254_N0PRIME = 0xc2e1f593efffffffUL; -constant F CANONICAL_ONE = {{1UL, 0UL, 0UL, 0UL}}; -constant F MONT_ONE = {{ - 0xac96341c4ffffffbUL, - 0x36fc76959f60cd29UL, - 0x666ea36f7879462eUL, - 0xe0a77c19a07df2fUL, -}}; - -struct StageConfig { - uint row_len; - uint half_m; - uint twiddle_offset; - uint _pad0; -}; - -struct BitReverseParams { - uint row_len; - uint log_n; - uint total_elements; - uint _pad0; -}; - -struct TransposeParams { - uint rows; - uint cols; - uint total_elements; -}; - -struct ReplicateCosetsParams { - uint row_len; - uint coset_size; - uint trailing_elements; -}; - -struct PackSingleVectorParams { - uint row_count; - uint message_length; - uint codeword_length; - uint coset_size; - uint total_elements; -}; - -struct FieldBytesParams { - uint rows; - uint cols; -}; - -struct ApplyCosetTwiddlesParams { - uint row_count; - uint num_cosets; - uint coset_size; - uint codeword_length; - uint total_elements; -}; - -inline F zero_f() { - F r; - r.v[0] = 0; r.v[1] = 0; r.v[2] = 0; r.v[3] = 0; - return r; -} - -inline F one_f() { - return MONT_ONE; -} - -inline F load_f(device const ulong *data, uint idx) { - F r; - uint base = idx * 4; - r.v[0] = data[base + 0]; - r.v[1] = data[base + 1]; - r.v[2] = data[base + 2]; - r.v[3] = data[base + 3]; - return r; -} - -inline void store_f(device ulong *data, uint idx, F x) { - uint base = idx * 4; - data[base + 0] = x.v[0]; - data[base + 1] = x.v[1]; - data[base + 2] = x.v[2]; - data[base + 3] = x.v[3]; -} - -inline bool ge_mod(F a) { - for (int i = 3; i >= 0; i--) { - if (a.v[i] > MODULUS[i]) return true; - if (a.v[i] < MODULUS[i]) return false; - } - return true; -} - -inline bool ge_f(F a, F b) { - for (int i = 3; i >= 0; i--) { - if (a.v[i] > b.v[i]) return true; - if (a.v[i] < b.v[i]) return false; - } - return true; -} - -inline F sub_raw(F a, thread const ulong *b, thread bool &borrow) { - F r; - borrow = false; - for (uint i = 0; i < 4; i++) { - ulong bi = b[i] + (borrow ? 1UL : 0UL); - bool bcarry = borrow && bi == 0; - r.v[i] = a.v[i] - bi; - borrow = bcarry || a.v[i] < bi; - } - return r; -} - -inline F sub_modulus(F a) { - F r; - bool borrow = false; - for (uint i = 0; i < 4; i++) { - ulong bi = MODULUS[i] + (borrow ? 1UL : 0UL); - bool bcarry = borrow && bi == 0; - r.v[i] = a.v[i] - bi; - borrow = bcarry || a.v[i] < bi; - } - return r; -} - -inline F add_modulus(F a) { - F r; - bool carry = false; - for (uint i = 0; i < 4; i++) { - ulong mi = MODULUS[i]; - ulong sum = a.v[i] + mi; - bool c0 = sum < a.v[i]; - ulong sum2 = sum + (carry ? 1UL : 0UL); - bool c1 = carry && sum2 == 0; - r.v[i] = sum2; - carry = c0 || c1; - } - return r; -} - -inline F add_f(F a, F b) { - F r; - bool carry = false; - for (uint i = 0; i < 4; i++) { - ulong sum = a.v[i] + b.v[i]; - bool c0 = sum < a.v[i]; - ulong sum2 = sum + (carry ? 1UL : 0UL); - bool c1 = carry && sum2 == 0; - r.v[i] = sum2; - carry = c0 || c1; - } - if (carry || ge_mod(r)) { - r = sub_modulus(r); - } - return r; -} - -inline F sub_f(F a, F b) { - ulong bv[4] = { b.v[0], b.v[1], b.v[2], b.v[3] }; - bool borrow = false; - F r = sub_raw(a, bv, borrow); - if (borrow) { - r = add_modulus(r); - } - return r; -} - -inline F double_f(F a) { - return add_f(a, a); -} - -inline ulong add_with_carry(ulong a, ulong b, thread ulong &carry) { - ulong sum = a + b; - ulong c1 = sum < a ? 1UL : 0UL; - ulong sum_with_carry = sum + carry; - ulong c2 = sum_with_carry < sum ? 1UL : 0UL; - carry = c1 + c2; - return sum_with_carry; -} - -inline void add_scaled_step(thread ulong &dst, ulong s, ulong a, thread ulong &carry) { - ulong product_lo = s * a; - ulong product_hi = mulhi(s, a); - - ulong sum = dst + product_lo; - ulong carry0 = sum < dst ? 1UL : 0UL; - ulong sum_with_carry = sum + carry; - ulong carry1 = sum_with_carry < sum ? 1UL : 0UL; - - dst = sum_with_carry; - carry = product_hi + carry0 + carry1; -} - -inline void add_scaled(thread ulong *dst, ulong s, ulong a0, ulong a1, ulong a2, ulong a3) { - ulong carry = 0; - add_scaled_step(dst[0], s, a0, carry); - add_scaled_step(dst[1], s, a1, carry); - add_scaled_step(dst[2], s, a2, carry); - add_scaled_step(dst[3], s, a3, carry); - dst[4] += carry; -} - -inline F mont_mul(F lhs, F rhs) { - ulong buf[9] = {0}; - uint off = 0; - -#pragma clang loop unroll(enable) - for (uint i = 0; i < 4; i++) { - add_scaled(&buf[off], lhs.v[i], rhs.v[0], rhs.v[1], rhs.v[2], rhs.v[3]); - - ulong m = buf[off] * BN254_N0PRIME; - add_scaled( - &buf[off], - m, - MODULUS[0], - MODULUS[1], - MODULUS[2], - MODULUS[3] - ); - - off += 1; - buf[off + 4] = 0; - } - - F result; - result.v[0] = buf[off + 0]; - result.v[1] = buf[off + 1]; - result.v[2] = buf[off + 2]; - result.v[3] = buf[off + 3]; - if (ge_mod(result)) { - result = sub_modulus(result); - } - return result; -} - -inline F from_mont(F value) { - F result = mont_mul(value, CANONICAL_ONE); - if (ge_mod(result)) { - result = sub_modulus(result); - } - return result; -} - -inline F mul_f(F a, F b) { - return mont_mul(a, b); -} - -inline F pow_f(F base, uint exp) { - F acc = one_f(); - F x = base; - while (exp != 0) { - if ((exp & 1) != 0) { - acc = mul_f(acc, x); - } - exp >>= 1; - if (exp != 0) { - x = mul_f(x, x); - } - } - return acc; -} - -kernel void bn254_fold( - device ulong *values [[buffer(0)]], - device const ulong *weight_buf [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &fold_half [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= fold_half) return; - F low = gid < len ? load_f(values, gid) : zero_f(); - F high = gid + fold_half < len ? load_f(values, gid + fold_half) : zero_f(); - F weight = load_f(weight_buf, 0); - store_f(values, gid, add_f(low, mul_f(sub_f(high, low), weight))); -} - -inline F fold_value_at(device const ulong *values, uint len, uint fold_half, uint idx, F weight) { - F low = idx < len ? load_f(values, idx) : zero_f(); - F high = idx + fold_half < len ? load_f(values, idx + fold_half) : zero_f(); - return add_f(low, mul_f(sub_f(high, low), weight)); -} - -kernel void bn254_fold_pair( - device ulong *a [[buffer(0)]], - device ulong *b [[buffer(1)]], - device const ulong *weight_buf [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &fold_half [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= fold_half) return; - F weight = load_f(weight_buf, 0); - store_f(a, gid, fold_value_at(a, len, fold_half, gid, weight)); - store_f(b, gid, fold_value_at(b, len, fold_half, gid, weight)); -} - -kernel void bn254_scalar_mul_add( - device ulong *acc [[buffer(0)]], - device const ulong *vector [[buffer(1)]], - device const ulong *weight_buf [[buffer(2)]], - constant uint &len [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F weight = load_f(weight_buf, 0); - store_f(acc, gid, add_f(load_f(acc, gid), mul_f(weight, load_f(vector, gid)))); -} - -kernel void bn254_scalar_mul( - device ulong *acc [[buffer(0)]], - device const ulong *weight_buf [[buffer(1)]], - constant uint &len [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F weight = load_f(weight_buf, 0); - store_f(acc, gid, mul_f(weight, load_f(acc, gid))); -} - -kernel void bn254_dot( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - F acc = zero_f(); - for (uint i = 0; i < len; i++) { - acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); - } - store_f(out, 0, acc); -} - -kernel void bn254_sumcheck( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &fold_half [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = 0; i < fold_half; i++) { - F a0 = i < len ? load_f(a, i) : zero_f(); - F b0 = i < len ? load_f(b, i) : zero_f(); - F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); - F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); - c0 = add_f(c0, mul_f(a0, b0)); - c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); - } - store_f(out, 0, c0); - store_f(out, 1, c2); -} - -kernel void bn254_dot_chunks( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &chunk_size [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - F acc = zero_f(); - for (uint i = start; i < end; i++) { - acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); - } - store_f(out, gid, acc); -} - -kernel void bn254_sum_chunks( - device const ulong *input [[buffer(0)]], - device ulong *out [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &offset [[buffer(3)]], - constant uint &chunk_size [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - F acc = zero_f(); - for (uint i = start; i < end; i++) { - acc = add_f(acc, load_f(input, offset + i)); - } - store_f(out, gid, acc); -} - -kernel void bn254_sumcheck_reduce_chunks( - device const ulong *input [[buffer(0)]], - device ulong *out [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &chunk_size [[buffer(3)]], - constant uint &out_len [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= out_len) return; - uint start = gid * chunk_size; - uint end = min(start + chunk_size, len); - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = start; i < end; i++) { - c0 = add_f(c0, load_f(input, i)); - c2 = add_f(c2, load_f(input, len + i)); - } - store_f(out, gid, c0); - store_f(out, out_len + gid, c2); -} - -kernel void bn254_sumcheck_chunks( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &fold_half [[buffer(4)]], - constant uint &chunk_size [[buffer(5)]], - constant uint &partial_count [[buffer(6)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= fold_half) return; - uint end = min(start + chunk_size, fold_half); - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = start; i < end; i++) { - F a0 = i < len ? load_f(a, i) : zero_f(); - F b0 = i < len ? load_f(b, i) : zero_f(); - F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); - F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); - c0 = add_f(c0, mul_f(a0, b0)); - c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); - } - store_f(out, gid, c0); - store_f(out, partial_count + gid, c2); -} - -kernel void bn254_fold_pair_sumcheck_chunks( - device ulong *a [[buffer(0)]], - device ulong *b [[buffer(1)]], - device const ulong *weight_buf [[buffer(2)]], - device ulong *out [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &fold_half [[buffer(5)]], - constant uint &sum_half [[buffer(6)]], - constant uint &chunk_size [[buffer(7)]], - constant uint &partial_count [[buffer(8)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= sum_half) return; - uint end = min(start + chunk_size, sum_half); - F weight = load_f(weight_buf, 0); - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = start; i < end; i++) { - uint right = i + sum_half; - F a0 = fold_value_at(a, len, fold_half, i, weight); - F b0 = fold_value_at(b, len, fold_half, i, weight); - F a1 = right < fold_half ? fold_value_at(a, len, fold_half, right, weight) : zero_f(); - F b1 = right < fold_half ? fold_value_at(b, len, fold_half, right, weight) : zero_f(); - store_f(a, i, a0); - store_f(b, i, b0); - if (right < fold_half) { - store_f(a, right, a1); - store_f(b, right, b1); - } - c0 = add_f(c0, mul_f(a0, b0)); - c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); - } - store_f(out, gid, c0); - store_f(out, partial_count + gid, c2); -} - -kernel void bn254_geometric_accumulate( - device ulong *acc [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *scalars [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &num_points [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F value = load_f(acc, gid); - for (uint j = 0; j < num_points; j++) { - value = add_f(value, mul_f(load_f(scalars, j), pow_f(load_f(points, j), gid))); - } - store_f(acc, gid, value); -} - -kernel void bn254_geometric_accumulate_chunks( - device ulong *acc [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *scalars [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &num_points [[buffer(4)]], - constant uint &chunk_size [[buffer(5)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - for (uint j = 0; j < num_points; j++) { - F point = load_f(points, j); - F scalar = load_f(scalars, j); - F power = pow_f(point, start); - for (uint i = start; i < end; i++) { - F value = load_f(acc, i); - value = add_f(value, mul_f(scalar, power)); - store_f(acc, i, value); - power = mul_f(power, point); - } - } -} - -kernel void bn254_geometric_accumulate_chunks_strided( - device ulong *acc [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *point_steps [[buffer(2)]], - device const ulong *scalars [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &num_points [[buffer(5)]], - constant uint &chunk_size [[buffer(6)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - uint count = end - start; - - // Accumulate all points into registers so `acc` is read and written - // once per element instead of once per (element, point). - F sums[32]; - for (uint k = 0; k < count; k++) { - sums[k] = zero_f(); - } - for (uint j = 0; j < num_points; j++) { - F point = load_f(points, j); - // Fold the scalar into the running power so the inner loop needs a - // single multiplication per element. - F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), gid)); - for (uint k = 0; k < count; k++) { - sums[k] = add_f(sums[k], power); - power = mul_f(power, point); - } - } - for (uint i = start, k = 0; i < end; i++, k++) { - store_f(acc, i, add_f(load_f(acc, i), sums[k])); - } -} - -kernel void bn254_geometric_accumulate_point_blocks( - device ulong *partials [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *point_steps [[buffer(2)]], - device const ulong *scalars [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &num_points [[buffer(5)]], - constant uint &chunk_size [[buffer(6)]], - constant uint &point_block_size [[buffer(7)]], - constant uint &point_blocks [[buffer(8)]], - uint gid [[thread_position_in_grid]] -) { - uint chunk = gid / point_blocks; - uint point_block = gid - chunk * point_blocks; - uint start = chunk * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - uint point_start = point_block * point_block_size; - if (point_start >= num_points) return; - uint point_end = min(point_start + point_block_size, num_points); - - F sums[32]; - for (uint k = 0; k < chunk_size; k++) { - sums[k] = zero_f(); - } - - for (uint j = point_start; j < point_end; j++) { - F point = load_f(points, j); - F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); - for (uint i = start, k = 0; i < end; i++, k++) { - sums[k] = add_f(sums[k], power); - power = mul_f(power, point); - } - } - - uint partial_offset = point_block * len; - for (uint i = start, k = 0; i < end; i++, k++) { - store_f(partials, partial_offset + i, sums[k]); - } -} - -kernel void bn254_geometric_accumulate_point_blocks_range( - device ulong *partials [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *point_steps [[buffer(2)]], - device const ulong *scalars [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &num_points [[buffer(5)]], - constant uint &chunk_size [[buffer(6)]], - constant uint &point_block_size [[buffer(7)]], - constant uint &point_block_offset [[buffer(8)]], - constant uint &batch_point_blocks [[buffer(9)]], - uint gid [[thread_position_in_grid]] -) { - uint chunk = gid / batch_point_blocks; - uint local_point_block = gid - chunk * batch_point_blocks; - uint global_point_block = point_block_offset + local_point_block; - uint start = chunk * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - uint point_start = global_point_block * point_block_size; - if (point_start >= num_points) return; - uint point_end = min(point_start + point_block_size, num_points); - - F sums[32]; - for (uint k = 0; k < chunk_size; k++) { - sums[k] = zero_f(); - } - - for (uint j = point_start; j < point_end; j++) { - F point = load_f(points, j); - F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); - for (uint i = start, k = 0; i < end; i++, k++) { - sums[k] = add_f(sums[k], power); - power = mul_f(power, point); - } - } - - uint partial_offset = local_point_block * len; - for (uint i = start, k = 0; i < end; i++, k++) { - store_f(partials, partial_offset + i, sums[k]); - } -} - -kernel void bn254_geometric_accumulate_reduce_point_blocks( - device ulong *acc [[buffer(0)]], - device const ulong *partials [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &point_blocks [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F value = load_f(acc, gid); - for (uint block = 0; block < point_blocks; block++) { - value = add_f(value, load_f(partials, block * len + gid)); - } - store_f(acc, gid, value); -} - -kernel void bn254_univariate_evaluate( - device const ulong *coeffs [[buffer(0)]], - device const ulong *point_buf [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - if (len == 0) { - store_f(out, 0, zero_f()); - return; - } - F point = load_f(point_buf, 0); - F acc = load_f(coeffs, len - 1); - for (uint i = len - 1; i > 0; i--) { - acc = add_f(mul_f(acc, point), load_f(coeffs, i - 1)); - } - store_f(out, 0, acc); -} - -kernel void bn254_univariate_eval_chunks( - device const ulong *coeffs [[buffer(0)]], - device const ulong *point_buf [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &chunk_size [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - F point = load_f(point_buf, 0); - F power = pow_f(point, start); - F acc = zero_f(); - for (uint i = start; i < end; i++) { - acc = add_f(acc, mul_f(load_f(coeffs, i), power)); - power = mul_f(power, point); - } - store_f(out, gid, acc); -} - -kernel void bn254_interleaved_rs_encode( - device const ulong *coeffs [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &num_messages [[buffer(3)]], - constant uint &coeff_len [[buffer(4)]], - constant uint &total [[buffer(5)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= total) return; - uint row = gid / num_messages; - uint message = gid - row * num_messages; - uint base = message * coeff_len; - F point = load_f(points, row); - F acc = load_f(coeffs, base + coeff_len - 1); - for (uint i = coeff_len - 1; i > 0; i--) { - acc = add_f(mul_f(acc, point), load_f(coeffs, base + i - 1)); - } - store_f(out, gid, acc); -} - -kernel void bn254_interleaved_rs_encode_single_vector( - device const ulong *vector [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &interleaving_depth [[buffer(3)]], - constant uint &message_length [[buffer(4)]], - constant uint &total [[buffer(5)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= total) return; - uint row = gid / interleaving_depth; - uint message = gid - row * interleaving_depth; - uint base = message * message_length; - F point = load_f(points, row); - F acc = load_f(vector, base + message_length - 1); - for (uint i = message_length - 1; i > 0; i--) { - acc = add_f(mul_f(acc, point), load_f(vector, base + i - 1)); - } - store_f(out, gid, acc); -} - -inline uint reverse_bits_width(uint value, uint width) { - return reverse_bits(value) >> (32u - width); -} - -kernel void bn254_pack_single_vector_cosets( - device const ulong *vector [[buffer(0)]], - device ulong *out [[buffer(1)]], - constant PackSingleVectorParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint row = gid / params.codeword_length; - uint col = gid - row * params.codeword_length; - if (col < params.message_length) { - store_f(out, gid, load_f(vector, row * params.message_length + col)); - } else { - store_f(out, gid, zero_f()); - } -} - -kernel void bn254_replicate_first_coset( - device ulong *buffer [[buffer(0)]], - constant ReplicateCosetsParams ¶ms [[buffer(1)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.trailing_elements) return; - uint repeats_per_row = params.row_len - params.coset_size; - uint row = gid / repeats_per_row; - uint within = gid - row * repeats_per_row; - uint dst = row * params.row_len + params.coset_size + within; - uint src = row * params.row_len + (within % params.coset_size); - store_f(buffer, dst, load_f(buffer, src)); -} - -kernel void bn254_bit_reverse_permute_rows_in_place( - device ulong *values [[buffer(0)]], - constant BitReverseParams &config [[buffer(1)]], - uint index [[thread_position_in_grid]] -) { - if (index >= config.total_elements || config.row_len <= 1u) return; - uint row = index / config.row_len; - uint within = index - row * config.row_len; - uint reversed = reverse_bits_width(within, config.log_n); - if (reversed <= within) return; - - uint row_base = row * config.row_len; - uint mate = row_base + reversed; - uint current = row_base + within; - F tmp = load_f(values, current); - store_f(values, current, load_f(values, mate)); - store_f(values, mate, tmp); -} - -kernel void bn254_radix2_ntt_stage_rows_in_place( - device ulong *values [[buffer(0)]], - device const ulong *twiddles [[buffer(1)]], - constant StageConfig &config [[buffer(2)]], - uint index [[thread_position_in_grid]] -) { - uint butterflies_per_row = config.row_len >> 1u; - uint row = index / butterflies_per_row; - uint local = index - row * butterflies_per_row; - uint half_m = config.half_m; - uint pair_in_group = local % half_m; - uint group = local / half_m; - uint row_base = row * config.row_len; - uint base = row_base + group * (half_m << 1u) + pair_in_group; - uint mate = base + half_m; - - F even = load_f(values, base); - F odd = load_f(values, mate); - F twiddle = load_f(twiddles, config.twiddle_offset + pair_in_group); - F t = mul_f(twiddle, odd); - - store_f(values, base, add_f(even, t)); - store_f(values, mate, sub_f(even, t)); -} - -kernel void bn254_transpose_matrix_reverse_rows( - device const ulong *input [[buffer(0)]], - device ulong *output [[buffer(1)]], - constant TransposeParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint row = gid / params.cols; - uint col = gid - row * params.cols; - uint row_bits = 31u - clz(params.cols); - uint dst_row = reverse_bits_width(col, row_bits); - uint dst = dst_row * params.rows + row; - store_f(output, dst, load_f(input, gid)); -} - -kernel void bn254_transpose_matrix( - device const ulong *input [[buffer(0)]], - device ulong *output [[buffer(1)]], - constant TransposeParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint row = gid / params.cols; - uint col = gid - row * params.cols; - uint dst = col * params.rows + row; - store_f(output, dst, load_f(input, gid)); -} - -kernel void bn254_apply_coset_twiddles( - device ulong *values [[buffer(0)]], - device const ulong *root_powers [[buffer(1)]], - constant ApplyCosetTwiddlesParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint within_codeword = gid % params.codeword_length; - uint coset = within_codeword / params.coset_size; - uint col = within_codeword - coset * params.coset_size; - if (coset == 0 || col == 0) return; - uint root_index = (coset * col) % params.codeword_length; - store_f(values, gid, mul_f(load_f(values, gid), load_f(root_powers, root_index))); -} - -kernel void bn254_encode_field_rows_le( - device const ulong *input [[buffer(0)]], - device uchar *output [[buffer(1)]], - constant FieldBytesParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - uint total_elements = params.rows * params.cols; - if (gid >= total_elements) return; - - F canonical = from_mont(load_f(input, gid)); - uint byte_offset = gid * 32u; - for (uint limb = 0; limb < 4; ++limb) { - ulong value = canonical.v[limb]; - for (uint byte = 0; byte < 8; ++byte) { - output[byte_offset + limb * 8u + byte] = uchar((value >> (byte * 8u)) & 0xffUL); - } - } -} - -kernel void bn254_read_rows( - device const ulong *input [[buffer(0)]], - device const uint *indices [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &num_cols [[buffer(3)]], - constant uint &total [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= total) return; - uint row = gid / num_cols; - uint col = gid - row * num_cols; - uint src = indices[row] * num_cols + col; - store_f(out, gid, load_f(input, src)); -} - -kernel void bn254_multilinear_extend( - device const ulong *values [[buffer(0)]], - device const ulong *point [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &num_vars [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - F acc = zero_f(); - F one = one_f(); - for (uint i = 0; i < len; i++) { - F weight = one; - for (uint j = 0; j < num_vars; j++) { - F r = load_f(point, num_vars - 1 - j); - if (((i >> j) & 1) != 0) { - weight = mul_f(weight, r); - } else { - weight = mul_f(weight, sub_f(one, r)); - } - } - acc = add_f(acc, mul_f(load_f(values, i), weight)); - } - store_f(out, 0, acc); -} -"#; - -const REDUCTION_CHUNK_SIZE: usize = 64; -const SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 8; -const LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; -/// Must match the `sums` register array size in the strided kernel. -const MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; -const LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD: usize = 1 << 18; -const GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE: usize = 16; -const GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD: usize = 1 << 17; -const GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS: usize = 64; -const GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES: usize = 256 << 20; - -#[derive(Clone, Debug)] -struct MetalFieldBuffer { - limbs: Buffer, -} - -#[derive(Clone, Debug)] -struct MetalHashBuffer { - bytes: Buffer, -} - -#[derive(Clone, Debug)] -pub struct MetalBuffer { - len: usize, - host_cache: OnceCell>, - field: Option, - hash: Option, - _marker: PhantomData, -} - -impl PartialEq for MetalBuffer -where - T: PartialEq, -{ - fn eq(&self, other: &Self) -> bool { - self.as_slice() == other.as_slice() - } -} - -impl Eq for MetalBuffer {} - -impl PartialOrd for MetalBuffer { - fn partial_cmp(&self, other: &Self) -> Option { - self.as_slice().partial_cmp(other.as_slice()) - } -} - -impl Ord for MetalBuffer { - fn cmp(&self, other: &Self) -> Ordering { - self.as_slice().cmp(other.as_slice()) - } -} - -impl Hash for MetalBuffer { - fn hash(&self, state: &mut H) { - self.as_slice().hash(state); - } -} - -impl Default for MetalBuffer { - fn default() -> Self { - Self { - len: 0, - host_cache: OnceCell::new(), - field: None, - hash: None, - _marker: PhantomData, - } - } -} - -impl Serialize for MetalBuffer { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - self.as_slice().serialize(serializer) - } -} - -impl<'de, T> Deserialize<'de> for MetalBuffer -where - T: Clone + Deserialize<'de>, -{ - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let data = Vec::::deserialize(deserializer)?; - Ok(BufferOps::from_vec(data)) - } -} - -impl MetalBuffer { - pub fn warmup() { - let _ = runtime(); - } - - pub(crate) fn as_slice(&self) -> &[T] { - self.host_cache - .get_or_init(|| self.download_host_cache()) - .as_slice() - } - - pub(crate) fn hash_bn254_rows_sha2(&self, num_cols: usize, out: &mut [Digest]) -> bool { - if type_name::() != type_name::() || self.field.is_none() { - return false; - } - assert_eq!(self.len(), num_cols * out.len()); - let message_size = num_cols * size_of::(); - let encoded = encode_field_rows_le( - &self - .field - .as_ref() - .expect("missing Metal field buffer") - .limbs, - out.len(), - num_cols, - ); - MetalSha2::new().hash_many_buffer(message_size, &encoded, out.len(), out); - true - } - - pub(crate) fn commit_bn254_rows_sha2_merkle( - &self, - num_cols: usize, - num_rows: usize, - layers: usize, - ) -> Option> { - if type_name::() != type_name::() || self.field.is_none() { - return None; - } - if num_rows != (1usize << layers) { - return None; - } - assert_eq!(self.len(), num_cols * num_rows); - let message_size = num_cols * size_of::(); - let encoded = encode_field_rows_le( - &self - .field - .as_ref() - .expect("missing Metal field buffer") - .limbs, - num_rows, - num_cols, - ); - let sha = MetalSha2::new(); - let nodes = sha.build_merkle_tree_buffer_from_messages_buffer( - message_size, - &encoded, - num_rows, - layers, - ); - Some(MetalBuffer::::from_digest_buffer( - nodes, - (1usize << (layers + 1)) - 1, - )) - } -} - -impl MetalBuffer { - pub(crate) fn from_digest_buffer(bytes: Buffer, len: usize) -> Self { - Self { - len, - host_cache: OnceCell::new(), - field: None, - hash: Some(MetalHashBuffer { bytes }), - _marker: PhantomData, - } - } - - pub(crate) fn read_hash_at(&self, index: usize) -> Option { - self.read_hash_indices(&[index]) - .map(|mut values| values.pop().expect("missing hash")) - } - - pub(crate) fn read_hash_indices(&self, indices: &[usize]) -> Option> { - let buffer = self.hash.as_ref()?; - Some(download_hash_indices(&buffer.bytes, self.len, indices)) - } -} - -impl BufferOps for MetalBuffer { - type Nodes = MetalBuffer; - - fn as_slice(&self) -> &[T] { - self.as_slice() - } - - fn at_index(&self, index: usize) -> Option { - self.as_slice().get(index).cloned() - } - - fn len(&self) -> usize { - self.len - } - - fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { - if type_name::() == type_name::() && self.field.is_some() { - return read_bn254_rows( - self.field.as_ref().expect("missing Metal field buffer"), - num_cols, - indices, - ) - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - let data = self.as_slice(); - let mut result = Vec::with_capacity(indices.len() * num_cols); - for i in indices { - result.extend_from_slice(&data[i * num_cols..(i + 1) * num_cols]); - } - result - } - - fn from_vec(source: Vec) -> Self { - let len = source.len(); - let field = maybe_upload_bn254(&source); - let host_cache = if field.is_some() { - OnceCell::new() - } else { - OnceCell::from(source) - }; - Self { - len, - host_cache, - field, - hash: None, - _marker: PhantomData, - } - } - - fn from_slice(source: &[T]) -> Self { - Self::from_vec(Vec::from(source)) - } - - fn merklize( - &self, - num_cols: usize, - leaf_hash: EngineId, - merkle: &merkle_tree::Config, - ) -> (Self::Nodes, Digest) - where - T: Encodable + Send + Sync, - { - let num_rows = merkle.num_leaves; - let layers = merkle.layers.len(); - assert_eq!(self.len(), num_cols * num_rows); - - // Fast path: build the whole tree on the GPU when the leaf and every node layer is SHA2. - if leaf_hash == hash::SHA2 - && merkle - .layers - .iter() - .all(|layer| layer.hash_id == hash::SHA2) - { - if let Some(nodes) = self.commit_bn254_rows_sha2_merkle(num_cols, num_rows, layers) { - let root = nodes - .read_hash_at(merkle.num_nodes() - 1) - .expect("missing Metal Merkle root"); - return (nodes, root); - } - } - - let cpu_nodes = || { - let engine = hash::ENGINES - .retrieve(leaf_hash) - .expect("Failed to retrieve hash engine"); - let mut leaves = vec![Digest::default(); num_rows]; - hash_rows(&*engine, self.as_slice(), &mut leaves); - merkle.build_nodes(leaves) - }; - - let nodes = if leaf_hash == hash::SHA2 { - let mut leaves = vec![Digest::default(); num_rows]; - if self.hash_bn254_rows_sha2(num_cols, &mut leaves) { - merkle.build_nodes(leaves) - } else { - cpu_nodes() - } - } else { - cpu_nodes() - }; - let root = nodes[merkle.num_nodes() - 1]; - (BufferOps::from_vec(nodes), root) - } -} - -impl crate::algebra::buffer::Buffer for MetalBuffer { - fn zeros(length: usize) -> Self { - assert_bn254::(); - // Montgomery zero is all-zero bytes, so a device-side fill suffices. - Self { - len: length, - host_cache: OnceCell::new(), - field: Some(zeroed_field_buffer(length)), - hash: None, - _marker: PhantomData, - } - } - - fn geometric_sequence(base: F, length: usize) -> Self { - assert_bn254::(); - BufferOps::from_vec(crate::algebra::geometric_sequence(base, length)) - } - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution, - { - assert_bn254::(); - BufferOps::from_vec((0..length).map(|_| rng.gen()).collect()) - } - - fn zero_pad(&mut self) { - assert_bn254::(); - if !self.is_empty() { - let mut data = self.as_slice().to_vec(); - data.resize(self.len().next_power_of_two(), F::ZERO); - *self = BufferOps::from_vec(data); - } - } - - fn fold(&mut self, weight: F) { - assert_bn254::(); - if self.len() <= 1 { - return; - } - let len = self.len(); - let fold_half = len.next_power_of_two() >> 1; - let weight = upload_field(&[f_to_field256(weight)]); - let field = self.bn254_buffer(); - run_in_place( - "bn254_fold", - &[&field.limbs, &weight.limbs], - &[len as u32, fold_half as u32], - fold_half, - ); - self.len = fold_half; - self.invalidate_host_cache(); - } - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - if self.len() <= 1 { - return; - } - let len = self.len(); - let fold_half = len.next_power_of_two() >> 1; - let weight = upload_field(&[f_to_field256(weight)]); - let this = self.bn254_buffer(); - let other_buffer = other.bn254_buffer(); - run_in_place( - "bn254_fold_pair", - &[&this.limbs, &other_buffer.limbs, &weight.limbs], - &[len as u32, fold_half as u32], - fold_half, - ); - self.len = fold_half; - other.len = fold_half; - self.invalidate_host_cache(); - other.invalidate_host_cache(); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - if self.len() <= 1 { - return self.sumcheck_polynomial(other); - } - let len = self.len(); - let fold_half = len.next_power_of_two() >> 1; - if fold_half == 1 { - self.fold_pair(other, weight); - return self.sumcheck_polynomial(other); - } - let weight = upload_field(&[f_to_field256(weight)]); - let this = self.bn254_buffer(); - let other_buffer = other.bn254_buffer(); - let (c0, c2) = parallel_fold_pair_sumcheck(&this, &other_buffer, &weight, len, fold_half); - self.len = fold_half; - other.len = fold_half; - self.invalidate_host_cache(); - other.invalidate_host_cache(); - (field256_to_f::(c0), field256_to_f::(c2)) - } - - fn linear_forms_rlc( - size: usize, - linear_forms: &mut [Box>], - rlc_coeffs: &[F], - ) -> Self { - assert_bn254::(); - assert_eq!(linear_forms.len(), rlc_coeffs.len()); - let Some((first, rest)) = linear_forms.split_first_mut() else { - return Self::zeros(size); - }; - let first = (first.as_mut() as &mut dyn std::any::Any) - .downcast_mut::>() - .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); - let mut accumulator = >::from_slice(&first.vector); - for (coeff, linear_form) in rlc_coeffs[1..].iter().zip(rest) { - let covector = (linear_form.as_mut() as &mut dyn std::any::Any) - .downcast_mut::>() - .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); - let vector = >::from_slice(&covector.vector); - vector.mixed_scalar_mul_add_to(&Identity::new(), &mut accumulator, *coeff); - } - accumulator - } - - fn mixed_linear_combination>( - _embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer { - assert_bn254::(); - assert_eq!(vectors.len(), coeffs.len()); - let Some((first, vectors)) = vectors.split_first() else { - return BufferOps::from_vec(Vec::new()); - }; - let mut accumulator = MetalBuffer:: { - len: first.len(), - host_cache: OnceCell::new(), - field: Some(copy_field_buffer(&first.bn254_buffer(), first.len())), - hash: None, - _marker: PhantomData, - }; - for (coeff, vector) in coeffs[1..].iter().copied().zip(vectors) { - vector.mixed_scalar_mul_add_to(_embedding, &mut accumulator, coeff); - } - accumulator - } -} - -/// Core geometric accumulate over `field[0..len]` (offset 0). Shared by the -/// owned buffer and a full-range view so the optimized chunk strategies are -/// written once. -fn geometric_accumulate_full( - field: &MetalFieldBuffer, - len: usize, - evaluators: &[UnivariateEvaluation], - scalars: &[F], -) { - let points = evaluators - .iter() - .map(|e| f_to_field256(e.point)) - .collect::>(); - let scalars = scalars - .iter() - .copied() - .map(f_to_field256) - .collect::>(); - let points = upload_field(&points); - let scalars = upload_field(&scalars); - let chunk_size = geometric_accumulate_chunk_size(len); - if std::env::var_os("WHIR_METAL_TRACE").is_some() { - eprintln!( - "metal geometric shape len={} points={} chunk={} chunks={}", - len, - evaluators.len(), - chunk_size, - len.div_ceil(chunk_size) - ); - } - if chunk_size <= 1 { - run_in_place( - "bn254_geometric_accumulate", - &[&field.limbs, &points.limbs, &scalars.limbs], - &[len as u32, evaluators.len() as u32], - len, - ); - } else { - let point_steps = evaluators - .iter() - .map(|e| f_to_field256(e.point.pow([chunk_size as u64]))) - .collect::>(); - let point_steps = upload_field(&point_steps); - if should_use_geometric_point_blocks(len, evaluators.len(), chunk_size) { - parallel_geometric_accumulate_point_blocks( - field, - &points, - &point_steps, - &scalars, - len, - evaluators.len(), - chunk_size, - ); - } else if should_use_geometric_point_blocks_batched(len, evaluators.len(), chunk_size) { - parallel_geometric_accumulate_point_blocks_batched( - field, - &points, - &point_steps, - &scalars, - len, - evaluators.len(), - chunk_size, - ); - } else { - run_in_place( - "bn254_geometric_accumulate_chunks_strided", - &[ - &field.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[len as u32, evaluators.len() as u32, chunk_size as u32], - len.div_ceil(chunk_size), - ); - } - } -} - -/// Validate that every evaluator targets `len` entries; returns `false` when -/// there is nothing to accumulate. -fn check_univariate_evaluators( - evaluators: &[UnivariateEvaluation], - scalars: &[F], - len: usize, -) -> bool { - assert_bn254::(); - assert_eq!(evaluators.len(), scalars.len()); - let Some(size) = evaluators.first().map(|e| e.size) else { - return false; - }; - assert_eq!(len, size); - for evaluator in evaluators { - assert_eq!(evaluator.size, size); - } - true -} - -impl BufferRead for MetalBuffer { - type TargetBuffer = MetalBuffer; - type Slice<'a> - = MetalSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.len() - } - - fn dot(&self, other: &Self) -> F { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - let this = self.bn254_buffer(); - let other = other.bn254_buffer(); - field256_to_f::(parallel_dot(&this, &other, self.len())) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - assert_bn254::(); - let len = self.len().min(other.len()); - if len == 0 { - return (F::ZERO, F::ZERO); - } - if len == 1 { - return (self.as_slice()[0] * other.as_slice()[0], F::ZERO); - } - let fold_half = len.next_power_of_two() >> 1; - let this = self.bn254_buffer(); - let other = other.bn254_buffer(); - let (c0, c2) = parallel_sumcheck(&this, &other, len, fold_half); - (field256_to_f::(c0), field256_to_f::(c2)) - } - - fn mixed_extend, T: Field>( - &self, - _embedding: &M, - point: &[M::Target], - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let num_vars = point.len(); - let point = point - .iter() - .copied() - .map(target_to_field256) - .collect::>(); - let point = upload_field(&point); - let this = self.bn254_buffer(); - let value = parallel_multilinear_extend_at(&this, 0, self.len(), &point, num_vars); - field256_to_target::(value) - } - - fn mixed_dot, T: Field>( - &self, - _embedding: &M, - other: &MetalBuffer, - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let this = self.bn254_buffer(); - let other = other.bn254_buffer_target(); - let value = field256_to_f::(parallel_dot(&this, &other, self.len())); - field256_to_target::(f_to_field256(value)) - } - - fn mixed_univariate_evaluate>( - &self, - _embedding: &M, - point: M::Target, - ) -> M::Target { - assert_bn254::(); - let point = target_to_field256(point); - let point = upload_field(&[point]); - let this = self.bn254_buffer(); - field256_to_target::(parallel_univariate_evaluate(&this, &point, self.len())) - } - - fn mixed_scalar_mul_add_to>( - &self, - _embedding: &M, - accumulator: &mut MetalBuffer, - weight: M::Target, - ) { - assert_bn254::(); - let weight = upload_field(&[target_to_field256(weight)]); - let vector = self.bn254_buffer(); - let acc = accumulator.bn254_buffer_target(); - scalar_mul_add_at(&acc, 0, &vector, 0, &weight, self.len()); - accumulator.field = Some(acc); - accumulator.invalidate_host_cache(); - } - - fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - assert_bn254::(); - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); - MetalSlice { - field: self.bn254_buffer(), - offset: start, - len: end - start, - _parent: PhantomData, - } - } -} - -impl BufferWrite for MetalBuffer { - type SliceMut<'a> - = MetalSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scalar_mul(&mut self, weight: F) { - assert_bn254::(); - if self.is_empty() { - return; - } - let weight = upload_field(&[f_to_field256(weight)]); - let field = self.bn254_buffer(); - run_in_place( - "bn254_scalar_mul", - &[&field.limbs, &weight.limbs], - &[self.len() as u32], - self.len(), - ); - self.field = Some(field); - self.invalidate_host_cache(); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - if !check_univariate_evaluators(evaluators, scalars, self.len()) { - return; - } - let field = self.bn254_buffer(); - geometric_accumulate_full(&field, self.len(), evaluators, scalars); - self.field = Some(field); - self.invalidate_host_cache(); - } - - fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { - assert_bn254::(); - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); - MetalSliceMut::new(self, start, end - start) - } - - fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { - assert_bn254::(); - let len = self.len(); - assert!(mid <= len, "split_at_mut mid {mid} out of bounds for {len}"); - // Materialize the parent's GPU allocation and invalidate its host - // cache once up front; the returned views share the handle and write - // disjoint ranges, so no per-op parent borrow is needed. - let field = self.bn254_buffer(); - self.field = Some(field.clone()); - self.invalidate_host_cache(); - ( - MetalSliceMut::from_field(field.clone(), 0, mid), - MetalSliceMut::from_field(field, mid, len - mid), - ) - } -} - -/// Metal (GPU) Reed-Solomon encoder. -/// -/// Wraps the shared [`RsDomain`] core: scalar methods delegate to the domain, and the coset -/// layout used by the GPU encode is taken from [`RsDomain::coset_params`] so it can never -/// drift from [`RsDomain::evaluation_points`]. Unlike the CPU encoder, it does not need the -/// host NTT engine at all. -#[derive(Debug, Clone)] -pub struct MetalRs { - domain: Arc>, -} - -impl MetalRs { - pub fn new(domain: Arc>) -> Self { - Self { domain } - } -} - -impl ReedSolomon for MetalRs { - fn next_order(&self, size: usize) -> Option { - self.domain.next_order(size) - } - - fn generator(&self, codeword_length: usize) -> F { - self.domain.generator(codeword_length) - } - - fn evaluation_points( - &self, - masked_message_length: usize, - codeword_length: usize, - indices: &[usize], - ) -> Vec { - self.domain - .evaluation_points(masked_message_length, codeword_length, indices) - } - - fn interleaved_encode( - &self, - vectors: &[&MetalBuffer], - masks: &MetalBuffer, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - ) -> MetalBuffer { - assert_bn254::(); - let num_messages = vectors.len() * interleaving_depth; - if num_messages == 0 { - return BufferOps::from_vec(Vec::new()); - } - assert!(masks.len().is_multiple_of(num_messages)); - let mask_length = masks.len() / num_messages; - if vectors.len() == 1 && mask_length == 0 { - // Single source of the coset layout: derived from the shared domain rather - // than recomputed on the device. - let (coset_size, _num_cosets) = - self.domain.coset_params(message_length, codeword_length); - return encode_single_vector_coset_ntt( - vectors[0], - message_length, - interleaving_depth, - codeword_length, - coset_size, - ); - } - - panic!("MetalBuffer BN254 RS encoding supports only one unmasked vector") - } -} - -impl MetalBuffer { - fn bn254_buffer(&self) -> MetalFieldBuffer { - self.field - .clone() - .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) - } - - fn invalidate_host_cache(&mut self) { - let _ = self.host_cache.take(); - } - - fn download_host_cache(&self) -> Vec { - if self.field.is_some() && type_name::() == type_name::() { - return download_field( - &self - .field - .as_ref() - .expect("missing Metal field buffer") - .limbs, - self.len, - ) - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - if self.hash.is_some() && type_name::() == type_name::() { - return download_hash_indices( - &self.hash.as_ref().expect("missing Metal hash buffer").bytes, - self.len, - &(0..self.len).collect::>(), - ) - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - panic!( - "MetalBuffer<{}> has no host cache and cannot be materialized", - type_name::() - ); - } -} - -/// Read-only, zero-copy view into a [`MetalBuffer`]'s GPU allocation. -/// -/// Holds a clone of the parent's allocation handle plus an `(offset, len)` -/// window. Reductions bind the handle at the byte offset, so they read -/// `parent[offset + i]` without copying. -pub struct MetalSlice<'a, F> { - field: MetalFieldBuffer, - offset: usize, - len: usize, - _parent: PhantomData<&'a MetalBuffer>, -} - -/// Mutable, zero-copy view into a [`MetalBuffer`]'s GPU allocation. -/// -/// Like [`MetalSlice`] but exclusive: the parent's host cache is invalidated -/// up front when the view is created, and writes go through the shared handle -/// at the byte offset, landing in the parent's own memory. -pub struct MetalSliceMut<'a, F> { - field: MetalFieldBuffer, - offset: usize, - len: usize, - _parent: PhantomData<&'a mut MetalBuffer>, -} - -impl<'a, F: Field + Clone> MetalSliceMut<'a, F> { - /// Borrow `parent[offset..offset+len]`, materializing the parent's GPU - /// allocation and invalidating its host cache once up front. - fn new(parent: &'a mut MetalBuffer, offset: usize, len: usize) -> Self { - let field = parent.bn254_buffer(); - parent.field = Some(field.clone()); - parent.invalidate_host_cache(); - Self { - field, - offset, - len, - _parent: PhantomData, - } - } - - /// Build a view directly from an already-materialized handle (used by - /// `split_at_mut`, where the parent cache was invalidated by the caller). - fn from_field(field: MetalFieldBuffer, offset: usize, len: usize) -> Self { - Self { - field, - offset, - len, - _parent: PhantomData, - } - } - - fn as_read(&self) -> MetalSlice<'_, F> { - MetalSlice { - field: self.field.clone(), - offset: self.offset, - len: self.len, - _parent: PhantomData, - } - } -} - -impl MetalSlice<'_, F> { - pub fn len(&self) -> usize { - self.len - } - pub fn is_empty(&self) -> bool { - self.len == 0 - } -} - -impl MetalSliceMut<'_, F> { - pub fn len(&self) -> usize { - self.len - } - pub fn is_empty(&self) -> bool { - self.len == 0 - } -} - -impl BufferRead for MetalSlice<'_, F> { - type TargetBuffer = MetalBuffer; - type Slice<'a> - = MetalSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.len - } - - fn dot(&self, other: &Self) -> F { - assert_bn254::(); - assert_eq!(self.len, other.len); - field256_to_f::(parallel_dot_at( - &self.field, - self.offset, - &other.field, - other.offset, - self.len, - )) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - assert_bn254::(); - let len = self.len.min(other.len); - if len == 0 { - return (F::ZERO, F::ZERO); - } - let fold_half = len.next_power_of_two() >> 1; - let (c0, c2) = parallel_sumcheck_at( - &self.field, - self.offset, - &other.field, - other.offset, - len, - fold_half, - ); - (field256_to_f::(c0), field256_to_f::(c2)) - } - - fn mixed_extend, T: Field>( - &self, - _embedding: &M, - point: &[M::Target], - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let num_vars = point.len(); - let point = point - .iter() - .copied() - .map(target_to_field256) - .collect::>(); - let point = upload_field(&point); - let value = - parallel_multilinear_extend_at(&self.field, self.offset, self.len, &point, num_vars); - field256_to_target::(value) - } - - fn mixed_dot, T: Field>( - &self, - _embedding: &M, - other: &MetalBuffer, - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let other = other.bn254_buffer_target(); - let value = field256_to_f::(parallel_dot_at( - &self.field, - self.offset, - &other, - 0, - self.len, - )); - field256_to_target::(f_to_field256(value)) - } - - fn mixed_univariate_evaluate>( - &self, - _embedding: &M, - point: M::Target, - ) -> M::Target { - assert_bn254::(); - let point = upload_field(&[target_to_field256(point)]); - field256_to_target::(parallel_univariate_evaluate_at( - &self.field, - self.offset, - &point, - self.len, - )) - } - - fn mixed_scalar_mul_add_to>( - &self, - _embedding: &M, - accumulator: &mut MetalBuffer, - weight: M::Target, - ) { - assert_bn254::(); - let weight = upload_field(&[target_to_field256(weight)]); - let acc = accumulator.bn254_buffer_target(); - scalar_mul_add_at(&acc, 0, &self.field, self.offset, &weight, self.len); - accumulator.field = Some(acc); - accumulator.invalidate_host_cache(); - } - - fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); - MetalSlice { - field: self.field.clone(), - offset: self.offset + start, - len: end - start, - _parent: PhantomData, - } - } -} - -impl BufferRead for MetalSliceMut<'_, F> { - type TargetBuffer = MetalBuffer; - type Slice<'a> - = MetalSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.len - } - - fn dot(&self, other: &Self) -> F { - self.as_read().dot(&other.as_read()) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - self.as_read().sumcheck_polynomial(&other.as_read()) - } - - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - self.as_read().mixed_extend(embedding, point) - } - - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &MetalBuffer, - ) -> M::Target { - self.as_read().mixed_dot(embedding, other) - } - - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - self.as_read().mixed_univariate_evaluate(embedding, point) - } - - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut MetalBuffer, - weight: M::Target, - ) { - self.as_read() - .mixed_scalar_mul_add_to(embedding, accumulator, weight) - } - - fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); - MetalSlice { - field: self.field.clone(), - offset: self.offset + start, - len: end - start, - _parent: PhantomData, - } - } -} - -impl BufferWrite for MetalSliceMut<'_, F> { - type SliceMut<'a> - = MetalSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scalar_mul(&mut self, weight: F) { - assert_bn254::(); - if self.len == 0 { - return; - } - let weight = upload_field(&[f_to_field256(weight)]); - scalar_mul_at_offset(&self.field, self.offset, self.len, &weight); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - if !check_univariate_evaluators(evaluators, scalars, self.len) { - return; - } - if self.offset == 0 { - // Offset 0: reuse the optimized full-buffer accumulate path. - geometric_accumulate_full(&self.field, self.len, evaluators, scalars); - } else { - // Arbitrary offset: bind the buffer at a byte offset so the - // kernel's gid-based indexing addresses field[offset + gid]. - let points = evaluators - .iter() - .map(|e| f_to_field256(e.point)) - .collect::>(); - let scalars = scalars - .iter() - .copied() - .map(f_to_field256) - .collect::>(); - let points = upload_field(&points); - let scalars = upload_field(&scalars); - geometric_accumulate_at_offset( - &self.field, - self.offset, - self.len, - &points, - &scalars, - evaluators.len(), - ); - } - } - - fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); - MetalSliceMut::from_field(self.field.clone(), self.offset + start, end - start) - } - - fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { - assert!( - mid <= self.len, - "split_at_mut mid {mid} out of bounds for {}", - self.len - ); - ( - MetalSliceMut::from_field(self.field.clone(), self.offset, mid), - MetalSliceMut::from_field(self.field.clone(), self.offset + mid, self.len - mid), - ) - } -} - -impl MetalBuffer { - fn bn254_buffer_target(&self) -> MetalFieldBuffer { - self.field - .clone() - .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) - } -} - -struct MetalRuntime { - device: Device, - queue: CommandQueue, - fold: ComputePipelineState, - fold_pair: ComputePipelineState, - scalar_mul_add: ComputePipelineState, - scalar_mul: ComputePipelineState, - dot: ComputePipelineState, - sumcheck: ComputePipelineState, - dot_chunks: ComputePipelineState, - sum_chunks: ComputePipelineState, - sumcheck_reduce_chunks: ComputePipelineState, - sumcheck_chunks: ComputePipelineState, - fold_pair_sumcheck_chunks: ComputePipelineState, - geometric_accumulate: ComputePipelineState, - geometric_accumulate_chunks: ComputePipelineState, - geometric_accumulate_chunks_strided: ComputePipelineState, - geometric_accumulate_point_blocks: ComputePipelineState, - geometric_accumulate_point_blocks_range: ComputePipelineState, - geometric_accumulate_reduce_point_blocks: ComputePipelineState, - univariate_evaluate: ComputePipelineState, - univariate_eval_chunks: ComputePipelineState, - interleaved_rs_encode: ComputePipelineState, - interleaved_rs_encode_single_vector: ComputePipelineState, - multilinear_extend: ComputePipelineState, - pack_single_vector_cosets: ComputePipelineState, - apply_coset_twiddles: ComputePipelineState, - replicate_first_coset: ComputePipelineState, - bit_reverse_rows: ComputePipelineState, - ntt_stage_rows: ComputePipelineState, - transpose: ComputePipelineState, - transpose_reverse_rows: ComputePipelineState, - encode_field_rows_le: ComputePipelineState, - read_rows: ComputePipelineState, - ntt_roots: Mutex>, - root_powers: Mutex>, -} - -fn runtime() -> &'static MetalRuntime { - static RUNTIME: OnceLock = OnceLock::new(); - RUNTIME.get_or_init(|| { - let device = Device::system_default().expect("Metal device is not available"); - let library = device - .new_library_with_source(METAL_SOURCE, &CompileOptions::new()) - .expect("failed to compile Metal BN254 kernels"); - let pipeline = |name: &str| { - let function = library - .get_function(name, None) - .unwrap_or_else(|_| panic!("missing Metal kernel {name}")); - device - .new_compute_pipeline_state_with_function(&function) - .unwrap_or_else(|err| panic!("failed to compile Metal kernel {name}: {err}")) - }; - MetalRuntime { - queue: device.new_command_queue(), - fold: pipeline("bn254_fold"), - fold_pair: pipeline("bn254_fold_pair"), - scalar_mul_add: pipeline("bn254_scalar_mul_add"), - scalar_mul: pipeline("bn254_scalar_mul"), - dot: pipeline("bn254_dot"), - sumcheck: pipeline("bn254_sumcheck"), - dot_chunks: pipeline("bn254_dot_chunks"), - sum_chunks: pipeline("bn254_sum_chunks"), - sumcheck_reduce_chunks: pipeline("bn254_sumcheck_reduce_chunks"), - sumcheck_chunks: pipeline("bn254_sumcheck_chunks"), - fold_pair_sumcheck_chunks: pipeline("bn254_fold_pair_sumcheck_chunks"), - geometric_accumulate: pipeline("bn254_geometric_accumulate"), - geometric_accumulate_chunks: pipeline("bn254_geometric_accumulate_chunks"), - geometric_accumulate_chunks_strided: pipeline( - "bn254_geometric_accumulate_chunks_strided", - ), - geometric_accumulate_point_blocks: pipeline("bn254_geometric_accumulate_point_blocks"), - geometric_accumulate_point_blocks_range: pipeline( - "bn254_geometric_accumulate_point_blocks_range", - ), - geometric_accumulate_reduce_point_blocks: pipeline( - "bn254_geometric_accumulate_reduce_point_blocks", - ), - univariate_evaluate: pipeline("bn254_univariate_evaluate"), - univariate_eval_chunks: pipeline("bn254_univariate_eval_chunks"), - interleaved_rs_encode: pipeline("bn254_interleaved_rs_encode"), - interleaved_rs_encode_single_vector: pipeline( - "bn254_interleaved_rs_encode_single_vector", - ), - multilinear_extend: pipeline("bn254_multilinear_extend"), - pack_single_vector_cosets: pipeline("bn254_pack_single_vector_cosets"), - apply_coset_twiddles: pipeline("bn254_apply_coset_twiddles"), - replicate_first_coset: pipeline("bn254_replicate_first_coset"), - bit_reverse_rows: pipeline("bn254_bit_reverse_permute_rows_in_place"), - ntt_stage_rows: pipeline("bn254_radix2_ntt_stage_rows_in_place"), - transpose: pipeline("bn254_transpose_matrix"), - transpose_reverse_rows: pipeline("bn254_transpose_matrix_reverse_rows"), - encode_field_rows_le: pipeline("bn254_encode_field_rows_le"), - read_rows: pipeline("bn254_read_rows"), - ntt_roots: Mutex::new(HashMap::new()), - root_powers: Mutex::new(HashMap::new()), - device, - } - }) -} - -fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { - match name { - "bn254_fold" => &rt.fold, - "bn254_fold_pair" => &rt.fold_pair, - "bn254_scalar_mul_add" => &rt.scalar_mul_add, - "bn254_scalar_mul" => &rt.scalar_mul, - "bn254_dot" => &rt.dot, - "bn254_sumcheck" => &rt.sumcheck, - "bn254_dot_chunks" => &rt.dot_chunks, - "bn254_sum_chunks" => &rt.sum_chunks, - "bn254_sumcheck_reduce_chunks" => &rt.sumcheck_reduce_chunks, - "bn254_sumcheck_chunks" => &rt.sumcheck_chunks, - "bn254_fold_pair_sumcheck_chunks" => &rt.fold_pair_sumcheck_chunks, - "bn254_geometric_accumulate" => &rt.geometric_accumulate, - "bn254_geometric_accumulate_chunks" => &rt.geometric_accumulate_chunks, - "bn254_geometric_accumulate_chunks_strided" => &rt.geometric_accumulate_chunks_strided, - "bn254_geometric_accumulate_point_blocks" => &rt.geometric_accumulate_point_blocks, - "bn254_geometric_accumulate_point_blocks_range" => { - &rt.geometric_accumulate_point_blocks_range - } - "bn254_geometric_accumulate_reduce_point_blocks" => { - &rt.geometric_accumulate_reduce_point_blocks - } - "bn254_univariate_evaluate" => &rt.univariate_evaluate, - "bn254_univariate_eval_chunks" => &rt.univariate_eval_chunks, - "bn254_interleaved_rs_encode" => &rt.interleaved_rs_encode, - "bn254_interleaved_rs_encode_single_vector" => &rt.interleaved_rs_encode_single_vector, - "bn254_multilinear_extend" => &rt.multilinear_extend, - "bn254_pack_single_vector_cosets" => &rt.pack_single_vector_cosets, - "bn254_apply_coset_twiddles" => &rt.apply_coset_twiddles, - "bn254_replicate_first_coset" => &rt.replicate_first_coset, - "bn254_bit_reverse_permute_rows_in_place" => &rt.bit_reverse_rows, - "bn254_radix2_ntt_stage_rows_in_place" => &rt.ntt_stage_rows, - "bn254_transpose_matrix" => &rt.transpose, - "bn254_transpose_matrix_reverse_rows" => &rt.transpose_reverse_rows, - "bn254_encode_field_rows_le" => &rt.encode_field_rows_le, - "bn254_read_rows" => &rt.read_rows, - _ => panic!("unknown Metal kernel {name}"), - } -} - -fn new_shared_buffer(rt: &MetalRuntime, bytes: u64) -> Buffer { - metal_profile::record_alloc(bytes); - let buffer = rt - .device - .new_buffer(bytes, MTLResourceOptions::StorageModeShared); - metal_profile::record_device_allocated(rt.device.current_allocated_size()); - buffer -} - -fn new_shared_buffer_with_data(rt: &MetalRuntime, data: *const c_void, bytes: u64) -> Buffer { - let start = Instant::now(); - let buffer = rt - .device - .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); - metal_profile::record_alloc(bytes); - metal_profile::record_upload(bytes, start.elapsed()); - metal_profile::record_device_allocated(rt.device.current_allocated_size()); - buffer -} - -fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { - command.commit(); - let start = Instant::now(); - command.wait_until_completed(); - let elapsed = start.elapsed(); - if std::env::var_os("WHIR_METAL_TRACE").is_some() { - eprintln!( - "metal command {label} {:.3} ms", - elapsed.as_secs_f64() * 1_000.0 - ); - } - metal_profile::record_command_wait(elapsed); -} - -fn wait_for_blit(command: &metal::CommandBufferRef, bytes: u64) { - command.commit(); - let start = Instant::now(); - command.wait_until_completed(); - metal_profile::record_blit(bytes, start.elapsed()); -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct BitReverseParams { - row_len: u32, - log_n: u32, - total_elements: u32, - _pad0: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct NttStageParams { - row_len: u32, - half_m: u32, - twiddle_offset: u32, - _pad0: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct TransposeParams { - rows: u32, - cols: u32, - total_elements: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct ReplicateCosetsParams { - row_len: u32, - coset_size: u32, - trailing_elements: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct PackSingleVectorParams { - row_count: u32, - message_length: u32, - codeword_length: u32, - coset_size: u32, - total_elements: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct FieldBytesParams { - rows: u32, - cols: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct ApplyCosetTwiddlesParams { - row_count: u32, - num_cosets: u32, - coset_size: u32, - codeword_length: u32, - total_elements: u32, -} - -fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field256 { - parallel_dot_at(a, 0, b, 0, len) -} - -/// Inner product of `a[a_off..a_off+len]` and `b[b_off..b_off+len]`. -fn parallel_dot_at( - a: &MetalFieldBuffer, - a_off: usize, - b: &MetalFieldBuffer, - b_off: usize, - len: usize, -) -> Field256 { - if len == 0 { - return Field256::ZERO; - } - let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_dot_chunks"), - &[&a.limbs, &b.limbs, &partials.limbs], - &[field_byte_offset(a_off), field_byte_offset(b_off), 0], - &[len as u32, REDUCTION_CHUNK_SIZE as u32], - partial_count, - ); - let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); - wait_for_command_named(command, "bn254_dot_chunks"); - download_field_at(&result.limbs, offset) -} - -fn parallel_sumcheck( - a: &MetalFieldBuffer, - b: &MetalFieldBuffer, - len: usize, - fold_half: usize, -) -> (Field256, Field256) { - parallel_sumcheck_at(a, 0, b, 0, len, fold_half) -} - -/// Sumcheck `(c0, c2)` over `a[a_off..]` and `b[b_off..]` (logical length `len`). -fn parallel_sumcheck_at( - a: &MetalFieldBuffer, - a_off: usize, - b: &MetalFieldBuffer, - b_off: usize, - len: usize, - fold_half: usize, -) -> (Field256, Field256) { - let partial_count = fold_half.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_sumcheck_chunks"), - &[&a.limbs, &b.limbs, &partials.limbs], - &[field_byte_offset(a_off), field_byte_offset(b_off), 0], - &[ - len as u32, - fold_half as u32, - REDUCTION_CHUNK_SIZE as u32, - partial_count as u32, - ], - partial_count, - ); - let result = encode_sumcheck_reduction(command, partials, partial_count); - wait_for_command_named(command, "bn254_sumcheck_chunks"); - download_sumcheck_pair(&result) -} - -fn parallel_fold_pair_sumcheck( - a: &MetalFieldBuffer, - b: &MetalFieldBuffer, - weight: &MetalFieldBuffer, - len: usize, - fold_half: usize, -) -> (Field256, Field256) { - let sum_half = fold_half.next_power_of_two() >> 1; - debug_assert!(sum_half > 0); - let partial_count = sum_half.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel( - command, - pipeline(rt, "bn254_fold_pair_sumcheck_chunks"), - &[&a.limbs, &b.limbs, &weight.limbs, &partials.limbs], - &[ - len as u32, - fold_half as u32, - sum_half as u32, - REDUCTION_CHUNK_SIZE as u32, - partial_count as u32, - ], - partial_count, - ); - let result = encode_sumcheck_reduction(command, partials, partial_count); - wait_for_command_named(command, "bn254_fold_pair_sumcheck_chunks"); - download_sumcheck_pair(&result) -} - -fn parallel_univariate_evaluate( - coeffs: &MetalFieldBuffer, - point: &MetalFieldBuffer, - len: usize, -) -> Field256 { - parallel_univariate_evaluate_at(coeffs, 0, point, len) -} - -/// Univariate evaluation of `coeffs[coeff_off..coeff_off+len]` at `point`. -fn parallel_univariate_evaluate_at( - coeffs: &MetalFieldBuffer, - coeff_off: usize, - point: &MetalFieldBuffer, - len: usize, -) -> Field256 { - if len == 0 { - return Field256::ZERO; - } - let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_univariate_eval_chunks"), - &[&coeffs.limbs, &point.limbs, &partials.limbs], - &[field_byte_offset(coeff_off), 0, 0], - &[len as u32, REDUCTION_CHUNK_SIZE as u32], - partial_count, - ); - let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); - wait_for_command_named(command, "bn254_univariate_eval_chunks"); - download_field_at(&result.limbs, offset) -} - -fn parallel_geometric_accumulate_point_blocks( - acc: &MetalFieldBuffer, - points: &MetalFieldBuffer, - point_steps: &MetalFieldBuffer, - scalars: &MetalFieldBuffer, - len: usize, - num_points: usize, - chunk_size: usize, -) { - assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); - let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); - let chunks = len.div_ceil(chunk_size); - let partial_len = point_blocks - .checked_mul(len) - .expect("Metal geometric partial size overflow"); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), - }; - run_in_place( - "bn254_geometric_accumulate_point_blocks", - &[ - &partials.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[ - len as u32, - num_points as u32, - chunk_size as u32, - GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, - point_blocks as u32, - ], - chunks * point_blocks, - ); - run_in_place( - "bn254_geometric_accumulate_reduce_point_blocks", - &[&acc.limbs, &partials.limbs], - &[len as u32, point_blocks as u32], - len, - ); -} - -fn parallel_geometric_accumulate_point_blocks_batched( - acc: &MetalFieldBuffer, - points: &MetalFieldBuffer, - point_steps: &MetalFieldBuffer, - scalars: &MetalFieldBuffer, - len: usize, - num_points: usize, - chunk_size: usize, -) { - assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); - let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); - let chunks = len.div_ceil(chunk_size); - let bytes_per_point_block = len - .checked_mul(size_of::()) - .expect("Metal geometric partial size overflow"); - let default_batch_blocks = - (GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES / bytes_per_point_block).max(1); - let batch_blocks = std::env::var("WHIR_METAL_GEOM_POINT_BLOCK_BATCH") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|&value| value > 0) - .unwrap_or(default_batch_blocks) - .min(point_blocks); - let rt = runtime(); - for point_block_offset in (0..point_blocks).step_by(batch_blocks) { - let current_batch = batch_blocks.min(point_blocks - point_block_offset); - let partial_len = current_batch - .checked_mul(len) - .expect("Metal geometric partial size overflow"); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), - }; - run_in_place( - "bn254_geometric_accumulate_point_blocks_range", - &[ - &partials.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[ - len as u32, - num_points as u32, - chunk_size as u32, - GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, - point_block_offset as u32, - current_batch as u32, - ], - chunks * current_batch, - ); - run_in_place( - "bn254_geometric_accumulate_reduce_point_blocks", - &[&acc.limbs, &partials.limbs], - &[len as u32, current_batch as u32], - len, - ); - } -} - -fn geometric_accumulate_chunk_size(len: usize) -> usize { - std::env::var("WHIR_METAL_GEOM_CHUNK") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|&value| value > 0) - .unwrap_or(if len >= LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD { - LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - } else { - SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - }) - .min(MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE) -} - -fn should_use_geometric_point_blocks(len: usize, num_points: usize, chunk_size: usize) -> bool { - chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS - && len <= GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD -} - -fn should_use_geometric_point_blocks_batched( - len: usize, - num_points: usize, - chunk_size: usize, -) -> bool { - chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS - && len > GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD -} - -/// Encodes all tree-reduction levels into `command` and returns the buffer -/// and offset holding the final scalar. Runs with a single command wait. -fn encode_field_reduction( - command: &metal::CommandBufferRef, - input: MetalFieldBuffer, - len: usize, - offset: usize, -) -> (MetalFieldBuffer, usize) { - let rt = runtime(); - let mut current = input; - let mut current_len = len; - let mut current_offset = offset; - while current_len > 1 { - let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); - let next = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (next_len * size_of::()) as u64), - }; - encode_u32_kernel( - command, - pipeline(rt, "bn254_sum_chunks"), - &[¤t.limbs, &next.limbs], - &[ - current_len as u32, - current_offset as u32, - REDUCTION_CHUNK_SIZE as u32, - ], - next_len, - ); - current = next; - current_len = next_len; - current_offset = 0; - } - (current, current_offset) -} - -/// Encodes all (c0, c2) tree-reduction levels into `command` and returns the -/// buffer holding the final pair. Runs with a single command wait. -fn encode_sumcheck_reduction( - command: &metal::CommandBufferRef, - input: MetalFieldBuffer, - len: usize, -) -> MetalFieldBuffer { - let rt = runtime(); - let mut current = input; - let mut current_len = len; - while current_len > 1 { - let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); - let next = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (next_len * 2 * size_of::()) as u64), - }; - encode_u32_kernel( - command, - pipeline(rt, "bn254_sumcheck_reduce_chunks"), - &[¤t.limbs, &next.limbs], - &[ - current_len as u32, - REDUCTION_CHUNK_SIZE as u32, - next_len as u32, - ], - next_len, - ); - current = next; - current_len = next_len; - } - current -} - -fn download_sumcheck_pair(buffer: &MetalFieldBuffer) -> (Field256, Field256) { - let values = download_field(&buffer.limbs, 2); - (values[0], values[1]) -} - -fn encode_single_vector_coset_ntt( - vector: &MetalBuffer, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - coset_size: usize, -) -> MetalBuffer { - assert!(codeword_length.is_power_of_two()); - assert!(Field256::get_root_of_unity(codeword_length as u64).is_some()); - assert_eq!(vector.len(), message_length * interleaving_depth); - - assert!(codeword_length.is_multiple_of(coset_size)); - let num_cosets = codeword_length / coset_size; - let total_elements = interleaving_depth - .checked_mul(codeword_length) - .expect("Metal RS encode size overflow"); - assert!(total_elements <= u32::MAX as usize); - - let rt = runtime(); - let source = vector.bn254_buffer(); - let current = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); - let transposed = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); - let codeword_root_powers = root_powers_buffer(codeword_length); - let coset_roots = roots_buffer(coset_size); - - let command = rt.queue.new_command_buffer(); - - encode_kernel( - &command, - &rt.pack_single_vector_cosets, - &[&source.limbs, ¤t], - &PackSingleVectorParams { - row_count: interleaving_depth as u32, - message_length: message_length as u32, - codeword_length: codeword_length as u32, - coset_size: coset_size as u32, - total_elements: total_elements as u32, - }, - total_elements, - ); - - let trailing_elements = interleaving_depth.saturating_mul(codeword_length - coset_size); - if trailing_elements != 0 { - encode_kernel( - &command, - &rt.replicate_first_coset, - &[¤t], - &ReplicateCosetsParams { - row_len: codeword_length as u32, - coset_size: coset_size as u32, - trailing_elements: trailing_elements as u32, - }, - trailing_elements, - ); - } - - encode_kernel( - &command, - &rt.apply_coset_twiddles, - &[¤t, &codeword_root_powers.limbs], - &ApplyCosetTwiddlesParams { - row_count: interleaving_depth as u32, - num_cosets: num_cosets as u32, - coset_size: coset_size as u32, - codeword_length: codeword_length as u32, - total_elements: total_elements as u32, - }, - total_elements, - ); - - let stage_count = coset_size.trailing_zeros() as usize; - encode_kernel( - &command, - &rt.bit_reverse_rows, - &[¤t], - &BitReverseParams { - row_len: coset_size as u32, - log_n: stage_count as u32, - total_elements: total_elements as u32, - _pad0: 0, - }, - total_elements, - ); - - let total_butterflies = total_elements / 2; - let mut twiddle_offset = 0usize; - for stage in 0..stage_count { - let half_m = 1usize << stage; - encode_kernel( - &command, - &rt.ntt_stage_rows, - &[¤t, &coset_roots.limbs], - &NttStageParams { - row_len: coset_size as u32, - half_m: half_m as u32, - twiddle_offset: twiddle_offset as u32, - _pad0: 0, - }, - total_butterflies, - ); - twiddle_offset += 1usize << stage; - } - - encode_kernel( - &command, - &rt.transpose, - &[¤t, &transposed], - &TransposeParams { - rows: interleaving_depth as u32, - cols: codeword_length as u32, - total_elements: total_elements as u32, - }, - total_elements, - ); - - wait_for_command_named(&command, "bn254_rs_encode"); - - MetalBuffer { - len: total_elements, - host_cache: OnceCell::new(), - field: Some(MetalFieldBuffer { limbs: transposed }), - hash: None, - _marker: PhantomData, - } -} - -fn encode_kernel( - command: &metal::CommandBufferRef, - pipeline: &ComputePipelineState, - buffers: &[&Buffer], - params: &P, - threads: usize, -) { - let encoder = command.new_compute_command_encoder(); - encoder.set_compute_pipeline_state(pipeline); - for (index, buffer) in buffers.iter().enumerate() { - encoder.set_buffer(index as u64, Some(buffer), 0); - } - encoder.set_bytes( - buffers.len() as u64, - size_of::

() as u64, - (params as *const P).cast::(), - ); - dispatch(&encoder, pipeline, threads.max(1)); - encoder.end_encoding(); -} - -fn roots_buffer(codeword_length: usize) -> MetalFieldBuffer { - let rt = runtime(); - if let Some(buffer) = rt.ntt_roots.lock().unwrap().get(&codeword_length).cloned() { - return buffer; - } - let root = Field256::get_root_of_unity(codeword_length as u64) - .expect("BN254 root of unity unavailable for Metal NTT"); - let stage_count = codeword_length.trailing_zeros() as usize; - let mut roots = Vec::with_capacity(codeword_length.saturating_sub(1)); - for stage in 0..stage_count { - let stage_size = 1usize << (stage + 1); - let half_stage = stage_size >> 1; - let stage_root = root.pow([(codeword_length / stage_size) as u64]); - let mut current = Field256::ONE; - for _ in 0..half_stage { - roots.push(current); - current *= stage_root; - } - } - let buffer = upload_field(&roots); - rt.ntt_roots - .lock() - .unwrap() - .insert(codeword_length, buffer.clone()); - buffer -} - -fn root_powers_buffer(codeword_length: usize) -> MetalFieldBuffer { - let rt = runtime(); - if let Some(buffer) = rt - .root_powers - .lock() - .unwrap() - .get(&codeword_length) - .cloned() - { - return buffer; - } - let root = Field256::get_root_of_unity(codeword_length as u64) - .expect("BN254 root of unity unavailable for Metal RS twiddles"); - let mut powers = Vec::with_capacity(codeword_length); - let mut current = Field256::ONE; - for _ in 0..codeword_length { - powers.push(current); - current *= root; - } - let buffer = upload_field(&powers); - rt.root_powers - .lock() - .unwrap() - .insert(codeword_length, buffer.clone()); - buffer -} - -fn encode_field_rows_le(input: &Buffer, rows: usize, cols: usize) -> Buffer { - assert!(rows <= u32::MAX as usize); - assert!(cols <= u32::MAX as usize); - let rt = runtime(); - let output = new_shared_buffer(rt, (rows * cols * size_of::()) as u64); - let command = rt.queue.new_command_buffer(); - encode_kernel( - &command, - &rt.encode_field_rows_le, - &[input, &output], - &FieldBytesParams { - rows: rows as u32, - cols: cols as u32, - }, - rows * cols, - ); - wait_for_command_named(&command, "bn254_encode_field_rows_le"); - output -} - -fn read_bn254_rows(source: &MetalFieldBuffer, num_cols: usize, indices: &[usize]) -> Vec { - if indices.is_empty() || num_cols == 0 { - return Vec::new(); - } - assert!(num_cols <= u32::MAX as usize); - assert!(indices.iter().all(|&index| index <= u32::MAX as usize)); - let total = indices - .len() - .checked_mul(num_cols) - .expect("Metal read_rows size overflow"); - assert!(total <= u32::MAX as usize); - - let rt = runtime(); - let indices = indices - .iter() - .copied() - .map(|index| index as u32) - .collect::>(); - let indices_buffer = new_shared_buffer_with_data( - rt, - indices.as_ptr().cast(), - (indices.len() * size_of::()) as u64, - ); - let out = new_shared_buffer(rt, (total * size_of::()) as u64); - run_in_place( - "bn254_read_rows", - &[&source.limbs, &indices_buffer, &out], - &[num_cols as u32, total as u32], - total, - ); - download_field(&out, total) -} - -fn upload_field(values: &[Field256]) -> MetalFieldBuffer { - let rt = runtime(); - let buffer = if values.is_empty() { - new_shared_buffer(rt, 0) - } else { - // Field256 is 4 contiguous u64 limbs; upload directly without - // flattening into an intermediate Vec. - new_shared_buffer_with_data( - rt, - values.as_ptr().cast(), - std::mem::size_of_val(values) as u64, - ) - }; - MetalFieldBuffer { limbs: buffer } -} - -fn zeroed_field_buffer(len: usize) -> MetalFieldBuffer { - let rt = runtime(); - let bytes = (len * size_of::()) as u64; - let buffer = new_shared_buffer(rt, bytes); - if bytes > 0 { - let command = rt.queue.new_command_buffer(); - let blit = command.new_blit_command_encoder(); - blit.fill_buffer(&buffer, metal::NSRange::new(0, bytes), 0); - blit.end_encoding(); - wait_for_blit(command, bytes); - } - MetalFieldBuffer { limbs: buffer } -} - -fn maybe_upload_bn254(values: &[T]) -> Option { - (type_name::() == type_name::()).then(|| upload_field(as_field256_slice(values))) -} - -/// Geometric accumulate into `field[offset .. offset+len]` by binding the -/// field buffer at a byte offset; the kernel itself indexes from `gid == 0`. -fn geometric_accumulate_at_offset( - field: &MetalFieldBuffer, - offset: usize, - len: usize, - points: &MetalFieldBuffer, - scalars: &MetalFieldBuffer, - num_points: usize, -) { - if len == 0 || num_points == 0 { - return; - } - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - let encoder = command.new_compute_command_encoder(); - let pipe = pipeline(rt, "bn254_geometric_accumulate"); - encoder.set_compute_pipeline_state(pipe); - let byte_offset = (offset * size_of::()) as u64; - encoder.set_buffer(0, Some(&field.limbs), byte_offset); - encoder.set_buffer(1, Some(&points.limbs), 0); - encoder.set_buffer(2, Some(&scalars.limbs), 0); - let len_u32 = len as u32; - let num_u32 = num_points as u32; - encoder.set_bytes(3, size_of::() as u64, (&len_u32 as *const u32).cast()); - encoder.set_bytes(4, size_of::() as u64, (&num_u32 as *const u32).cast()); - dispatch(&encoder, pipe, len); - encoder.end_encoding(); - wait_for_command_named(command, "bn254_geometric_accumulate_window"); -} - -/// In-place scalar multiply of `field[offset .. offset+len]` by binding the -/// field buffer at a byte offset; the kernel indexes from `gid == 0`. -fn scalar_mul_at_offset( - field: &MetalFieldBuffer, - offset: usize, - len: usize, - weight: &MetalFieldBuffer, -) { - if len == 0 { - return; - } - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - let encoder = command.new_compute_command_encoder(); - let pipe = pipeline(rt, "bn254_scalar_mul"); - encoder.set_compute_pipeline_state(pipe); - let byte_offset = (offset * size_of::()) as u64; - encoder.set_buffer(0, Some(&field.limbs), byte_offset); - encoder.set_buffer(1, Some(&weight.limbs), 0); - let len_u32 = len as u32; - encoder.set_bytes(2, size_of::() as u64, (&len_u32 as *const u32).cast()); - dispatch(&encoder, pipe, len); - encoder.end_encoding(); - wait_for_command_named(command, "bn254_scalar_mul_window"); -} - -/// Multilinear extension of `field[off..off+len]` evaluated at `point`. -fn parallel_multilinear_extend_at( - field: &MetalFieldBuffer, - off: usize, - len: usize, - point: &MetalFieldBuffer, - num_vars: usize, -) -> Field256 { - let rt = runtime(); - let out = new_shared_buffer(rt, (4 * size_of::()) as u64); - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_multilinear_extend"), - &[&field.limbs, &point.limbs, &out], - &[field_byte_offset(off), 0, 0], - &[len as u32, num_vars as u32], - 1, - ); - wait_for_command_named(command, "bn254_multilinear_extend"); - download_field(&out, 1)[0] -} - -/// `acc[acc_off..] += weight * vector[vec_off..]` over `len` elements. -fn scalar_mul_add_at( - acc: &MetalFieldBuffer, - acc_off: usize, - vector: &MetalFieldBuffer, - vec_off: usize, - weight: &MetalFieldBuffer, - len: usize, -) { - if len == 0 { - return; - } - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_scalar_mul_add"), - &[&acc.limbs, &vector.limbs, &weight.limbs], - &[field_byte_offset(acc_off), field_byte_offset(vec_off), 0], - &[len as u32], - len, - ); - wait_for_command_named(command, "bn254_scalar_mul_add"); -} - -fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer { - let rt = runtime(); - let byte_len = (len * 4 * size_of::()) as u64; - let target = new_shared_buffer(rt, byte_len); - let command = rt.queue.new_command_buffer(); - let blit = command.new_blit_command_encoder(); - blit.copy_from_buffer(&source.limbs, 0, &target, 0, byte_len); - blit.end_encoding(); - wait_for_blit(&command, byte_len); - MetalFieldBuffer { limbs: target } -} - -/// Encodes a kernel dispatch with `u32` constants bound as inline bytes -/// (no per-constant buffer allocations). -fn encode_u32_kernel( - command: &metal::CommandBufferRef, - pipeline: &ComputePipelineState, - buffers: &[&Buffer], - constants: &[u32], - threads: usize, -) { - encode_u32_kernel_with_offsets(command, pipeline, buffers, &[], constants, threads); -} - -/// Like [`encode_u32_kernel`], but binds `buffers[i]` at byte offset -/// `offsets[i]` (defaulting to 0 when `offsets` is shorter). This is the -/// mechanism that lets a view dispatch a kernel over `parent[offset..]` -/// without copying: only the input binding shifts, the kernel still indexes -/// from `gid == 0`. -fn encode_u32_kernel_with_offsets( - command: &metal::CommandBufferRef, - pipeline: &ComputePipelineState, - buffers: &[&Buffer], - offsets: &[u64], - constants: &[u32], - threads: usize, -) { - let encoder = command.new_compute_command_encoder(); - encoder.set_compute_pipeline_state(pipeline); - let mut index = 0; - for (i, buffer) in buffers.iter().enumerate() { - let byte_offset = offsets.get(i).copied().unwrap_or(0); - encoder.set_buffer(index, Some(buffer), byte_offset); - index += 1; - } - for constant in constants { - encoder.set_bytes( - index, - size_of::() as u64, - (constant as *const u32).cast(), - ); - index += 1; - } - dispatch(&encoder, pipeline, threads.max(1)); - encoder.end_encoding(); -} - -/// Byte offset of element `offset` within a `Field256` buffer. -#[inline] -fn field_byte_offset(offset: usize) -> u64 { - (offset * size_of::()) as u64 -} - -fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - encode_u32_kernel(command, pipeline(rt, name), buffers, constants, threads); - wait_for_command_named(command, name); -} - -fn dispatch( - encoder: &metal::ComputeCommandEncoderRef, - pipeline: &ComputePipelineState, - threads: usize, -) { - // Use full threadgroups (capped at 256) instead of a single execution - // width; the pipeline limit already accounts for register pressure. - let width = pipeline.max_total_threads_per_threadgroup().clamp(1, 256); - let group_width = width.min(threads as u64).max(1); - encoder.dispatch_threads( - MTLSize { - width: threads as u64, - height: 1, - depth: 1, - }, - MTLSize { - width: group_width, - height: 1, - depth: 1, - }, - ); -} - -fn download_field(buffer: &Buffer, len: usize) -> Vec { - if len == 0 { - return Vec::new(); - } - let start = Instant::now(); - let limbs = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 4) }; - let result = limbs - .chunks_exact(4) - .map(|chunk| { - Fp::, 4>( - BigInt([chunk[0], chunk[1], chunk[2], chunk[3]]), - PhantomData, - ) - }) - .collect(); - metal_profile::record_readback((len * size_of::()) as u64, start.elapsed()); - result -} - -fn download_field_at(buffer: &Buffer, index: usize) -> Field256 { - let start = Instant::now(); - let limbs = - unsafe { std::slice::from_raw_parts(buffer.contents().cast::().add(index * 4), 4) }; - let result = Fp::, 4>( - BigInt([limbs[0], limbs[1], limbs[2], limbs[3]]), - PhantomData, - ); - metal_profile::record_readback(size_of::() as u64, start.elapsed()); - result -} - -fn download_hash_indices(buffer: &Buffer, len: usize, indices: &[usize]) -> Vec { - let start = Instant::now(); - let bytes = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 32) }; - let mut result = Vec::with_capacity(indices.len()); - for &index in indices { - assert!(index < len, "Metal hash index out of bounds"); - let mut hash = Digest::default(); - hash.as_mut_bytes() - .copy_from_slice(&bytes[index * 32..(index + 1) * 32]); - result.push(hash); - } - metal_profile::record_readback( - (indices.len() * size_of::()) as u64, - start.elapsed(), - ); - result -} - -fn assert_bn254() { - assert_eq!( - type_name::(), - type_name::(), - "MetalBuffer only supports BN254 Field256 field operations" - ); -} - -fn f_to_field256(value: F) -> Field256 { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -fn field256_to_f(value: Field256) -> F { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -fn target_to_field256(value: T) -> Field256 { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -fn field256_to_target(value: Field256) -> T { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -fn as_field256_slice(values: &[T]) -> &[Field256] { - assert_eq!( - type_name::(), - type_name::(), - "MetalBuffer only supports BN254 Field256 buffers" - ); - unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::algebra::{ - buffer::{Buffer, CpuBuffer}, - ntt::NttEngine, - }; - use crate::buffer::BufferOps; - - fn values(len: usize, offset: u64) -> Vec { - (0..len) - .map(|i| Field256::from(i as u64 + offset)) - .collect() - } - - #[test] - fn metal_bn254_dot_matches_cpu() { - let a = values(33, 1); - let b = values(33, 9); - let cpu_a = CpuBuffer::from_slice(&a); - let cpu_b = CpuBuffer::from_slice(&b); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_b = MetalBuffer::from_slice(&b); - assert_eq!(gpu_a.dot(&gpu_b), cpu_a.dot(&cpu_b)); - } - - #[test] - fn metal_bn254_fold_matches_cpu() { - let mut cpu = CpuBuffer::from_vec(values(31, 2)); - let mut gpu = MetalBuffer::from_vec(values(31, 2)); - let weight = Field256::from(42); - cpu.fold(weight); - gpu.fold(weight); - assert_eq!(gpu.as_slice(), cpu.as_slice()); - } - - #[test] - fn metal_bn254_sumcheck_matches_cpu() { - for len in [1, 2, 27, 64, 65] { - let a = values(len, 3); - let b = values(len, 11); - let cpu_a = CpuBuffer::from_slice(&a); - let cpu_b = CpuBuffer::from_slice(&b); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_b = MetalBuffer::from_slice(&b); - assert_eq!( - gpu_a.sumcheck_polynomial(&gpu_b), - cpu_a.sumcheck_polynomial(&cpu_b) - ); - } - } - - #[test] - fn metal_bn254_fold_pair_sumcheck_matches_cpu() { - for len in [2, 3, 27, 64, 65] { - let mut cpu_a = CpuBuffer::from_vec(values(len, 3)); - let mut cpu_b = CpuBuffer::from_vec(values(len, 11)); - let mut gpu_a = MetalBuffer::from_vec(values(len, 3)); - let mut gpu_b = MetalBuffer::from_vec(values(len, 11)); - let weight = Field256::from(42); - let cpu_result = cpu_a.fold_pair_sumcheck_polynomial(&mut cpu_b, weight); - let gpu_result = gpu_a.fold_pair_sumcheck_polynomial(&mut gpu_b, weight); - assert_eq!(gpu_result, cpu_result); - assert_eq!(gpu_a.as_slice(), cpu_a.as_slice()); - assert_eq!(gpu_b.as_slice(), cpu_b.as_slice()); - } - } - - #[test] - fn metal_bn254_scalar_mul_add_matches_cpu() { - let mut cpu = CpuBuffer::from_vec(values(19, 1)); - let mut gpu = MetalBuffer::from_vec(values(19, 1)); - let vector = values(19, 5); - let cpu_vector = CpuBuffer::from_slice(&vector); - let gpu_vector = MetalBuffer::from_slice(&vector); - let weight = Field256::from(7); - cpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut cpu, weight); - gpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut gpu, weight); - assert_eq!(gpu.as_slice(), cpu.as_slice()); - } - - #[test] - fn metal_bn254_interleaved_rs_encode_matches_cpu() { - // The GPU encoder and the CPU reference share one coset layout: the engine and the - // domain are derived from the same field, and the GPU encode reads its layout from - // `RsDomain::coset_params` (which the engine's slice encode also uses). - let engine = NttEngine::::new_from_fftfield(); - let gpu_rs = MetalRs::new(Arc::new(RsDomain::::from_fftfield())); - - let a = values(8, 1); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_masks = MetalBuffer::from_slice(&[]); - - // CPU reference straight from the engine's slice API: `a` is one vector of two - // length-4 messages, no masks, codeword length 8. - let messages = a.chunks_exact(4).collect::>(); - let cpu = engine.interleaved_encode_slices(&messages, &[], 8); - let gpu = gpu_rs.interleaved_encode(&[&gpu_a], &gpu_masks, 4, 2, 8); - assert_eq!(gpu.as_slice(), cpu.as_slice()); - } - - #[test] - fn metal_bn254_mixed_extend_matches_cpu() { - let values = values(8, 3); - let point = vec![Field256::from(2), Field256::from(5), Field256::from(9)]; - let cpu = CpuBuffer::from_slice(&values); - let gpu = MetalBuffer::from_slice(&values); - assert_eq!( - gpu.mixed_extend(&Identity::new(), &point), - cpu.mixed_extend(&Identity::new(), &point) - ); - } -} diff --git a/src/algebra/mod.rs b/src/algebra/mod.rs index 592a9109..2c7e4838 100644 --- a/src/algebra/mod.rs +++ b/src/algebra/mod.rs @@ -1,9 +1,6 @@ -pub mod buffer; pub mod embedding; pub mod fields; pub mod linear_form; -#[cfg(all(feature = "metal", target_os = "macos"))] -mod metal_buffer; mod multilinear; pub mod ntt; pub mod sumcheck; diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index a01c957f..4970b1a3 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -4,9 +4,6 @@ //! A global cache is used for twiddle factors. use std::sync::{RwLock, RwLockReadGuard}; -// `Arc` is only used by the CPU encoder, which is gated to non-Metal builds. -#[cfg(not(all(feature = "metal", target_os = "macos")))] -use std::sync::Arc; use ark_ff::{FftField, Field}; #[cfg(feature = "tracing")] @@ -14,13 +11,14 @@ use tracing::instrument; #[cfg(feature = "parallel")] use {crate::utils::workload_size, rayon::prelude::*, std::cmp::max}; -use super::transpose; -use super::utils::{lcm, sqrt_factor}; -// The CPU `ReedSolomon` impl and `CpuBuffer` are only used when CPU buffers are the active backend. -#[cfg(not(all(feature = "metal", target_os = "macos")))] -use super::ReedSolomon; -#[cfg(not(all(feature = "metal", target_os = "macos")))] -use crate::algebra::buffer::{BufferOps, CpuBuffer}; +use super::{ + transpose, + utils::{lcm, sqrt_factor}, + Messages, ReedSolomon, +}; + +use crate::buffer::{ActiveBuffer, BufferOps}; + #[cfg(not(feature = "rs_in_order"))] use crate::algebra::ntt::transpose::transpose_permute; use crate::{ @@ -31,26 +29,13 @@ use crate::{ // Supported primes const PRIMES: [usize; 2] = [2, 3]; -/// Reed-Solomon code domain for a field `F`: the minimal, backend-neutral core. -/// -/// Holds only the NTT subgroup `order`, its `divisors`, and a primitive root, together with -/// the coset convention (`next_order` / `generator` / `coset_params` / `evaluation_points`). -/// Both the CPU engine ([`NttEngine`]) and the Metal encoder share this, so the coset layout -/// is defined in exactly one place. Assumes the field has large two-adicity. -#[derive(Debug)] -pub struct RsDomain { - pub(crate) order: usize, // order of omega_order - pub(crate) divisors: Vec, // divisors of the order. - pub(crate) omega_order: F, // primitive order'th root. -} - -/// CPU Cooley-Tukey NTT engine for a field `F`. -/// -/// Wraps the shared [`RsDomain`] and adds the host-side NTT machinery: small-order roots -/// for the butterflies and an on-demand roots-of-unity cache. CPU only. +/// Engine for computing NTTs over arbitrary fields. +/// Assumes the field has large two-adicity. #[derive(Debug)] pub struct NttEngine { - domain: RsDomain, + order: usize, // order of omega_orger + divisors: Vec, // divisors of the order. + omega_order: F, // primitive order'th root. // Roots of small order (zero if unavailable). The naming convention is that omega_foo has order foo. half_omega_3_1_plus_2: F, // ½(ω₃ + ω₃²) @@ -66,9 +51,9 @@ pub struct NttEngine { roots: RwLock>, } -impl RsDomain { - /// Construct the domain from the field's `FftField` parameters. - pub fn from_fftfield() -> Self { +impl NttEngine { + /// Construct a new engine from the field's `FftField` trait. + pub fn new_from_fftfield() -> Self { let (mut omega, mut order) = if let (Some(mut omega), Some(b), Some(k)) = ( F::LARGE_SUBGROUP_ROOT_OF_UNITY, F::SMALL_SUBGROUP_BASE, @@ -97,8 +82,8 @@ impl RsDomain { } } -impl RsDomain { - /// Creates a new domain. `omega_order` must be a primitive root of unity of even order `order`. +/// Creates a new NttEngine. `omega_order` must be a primitive root of unity of even order `omega`. +impl NttEngine { pub fn new(order: usize, omega_order: F) -> Self { // Make sure `omega_order` is a primitive root of unity. assert_eq!(omega_order.pow([order as u64]), F::ONE); @@ -107,99 +92,11 @@ impl RsDomain { assert_ne!(omega_order.pow([(order / prime) as u64]), F::ONE); } } - Self { + + let mut res = Self { order, divisors: divisors(order, &PRIMES), omega_order, - } - } - - pub fn next_order(&self, size: usize) -> Option { - match self.divisors.binary_search(&size) { - Ok(index) | Err(index) => self.divisors.get(index).copied(), - } - } - - pub fn generator(&self, codeword_length: usize) -> F { - self.omega_order - .pow([(self.order / codeword_length) as u64]) - } - - pub fn checked_root(&self, order: usize) -> Option { - if order == 0 { - return Some(F::ONE); - } - self.order - .is_multiple_of(order) - .then(|| self.omega_order.pow([(self.order / order) as u64])) - } - - pub fn root(&self, order: usize) -> F { - self.checked_root(order) - .expect("Subgroup of requested order does not exist.") - } - - /// The single definition of the coset layout used by Reed-Solomon encoding. - /// - /// Returns `(coset_size, num_cosets)` for the given masked message length and - /// codeword length. Both [`Self::evaluation_points`] and every backend's encode - /// derive their layout from here, so the index ordering can never drift between them. - pub fn coset_params( - &self, - masked_message_length: usize, - codeword_length: usize, - ) -> (usize, usize) { - let mut coset_size = self.next_order(masked_message_length).unwrap(); - while !codeword_length.is_multiple_of(coset_size) { - coset_size = self.next_order(coset_size + 1).unwrap(); - } - (coset_size, codeword_length / coset_size) - } - - pub fn evaluation_points( - &self, - masked_message_length: usize, - codeword_length: usize, - indices: &[usize], - ) -> Vec { - assert!(masked_message_length <= codeword_length); - assert!(self.order.is_multiple_of(codeword_length)); - let mut result = Vec::new(); - let generator = self.generator(codeword_length); - - let (coset_size, num_cosets) = self.coset_params(masked_message_length, codeword_length); - #[cfg(feature = "rs_in_order")] - let _ = (coset_size, num_cosets); - - for &index in indices { - assert!(index < codeword_length); - - #[cfg(not(feature = "rs_in_order"))] - let index = transpose_permute(index, num_cosets, coset_size); - result.push(generator.pow([index as u64])); - } - result - } -} - -impl NttEngine { - /// Construct a new engine from the field's `FftField` trait. - pub fn new_from_fftfield() -> Self { - Self::from_domain(RsDomain::from_fftfield()) - } -} - -/// Creates a new NttEngine. `omega_order` must be a primitive root of unity of even order `omega`. -impl NttEngine { - pub fn new(order: usize, omega_order: F) -> Self { - Self::from_domain(RsDomain::new(order, omega_order)) - } - - /// Build the CPU engine on top of a shared [`RsDomain`], precomputing small-order roots. - pub fn from_domain(domain: RsDomain) -> Self { - let order = domain.order; - let mut res = Self { - domain, half_omega_3_1_plus_2: F::ZERO, half_omega_3_1_min_2: F::ZERO, omega_4_1: F::ZERO, @@ -211,32 +108,27 @@ impl NttEngine { roots: RwLock::new(Vec::new()), }; if order.is_multiple_of(3) { - let omega_3_1 = res.domain.root(3); + let omega_3_1 = res.root(3); let omega_3_2 = omega_3_1 * omega_3_1; // Note: char F cannot be 2 and so division by 2 works, because primitive roots of unity with even order exist. res.half_omega_3_1_min_2 = (omega_3_1 - omega_3_2) / F::from(2u64); res.half_omega_3_1_plus_2 = (omega_3_1 + omega_3_2) / F::from(2u64); } if order.is_multiple_of(4) { - res.omega_4_1 = res.domain.root(4); + res.omega_4_1 = res.root(4); } if order.is_multiple_of(8) { - res.omega_8_1 = res.domain.root(8); + res.omega_8_1 = res.root(8); res.omega_8_3 = res.omega_8_1.pow([3]); } if order.is_multiple_of(16) { - res.omega_16_1 = res.domain.root(16); + res.omega_16_1 = res.root(16); res.omega_16_3 = res.omega_16_1.pow([3]); res.omega_16_9 = res.omega_16_1.pow([9]); } res } - /// The shared field/domain core this engine is built on. - pub const fn domain(&self) -> &RsDomain { - &self.domain - } - pub fn ntt(&self, values: &mut [F]) { self.ntt_batch(values, values.len()); } @@ -279,10 +171,24 @@ impl NttEngine { self.ntt_batch(values, size); } + pub fn checked_root(&self, order: usize) -> Option { + if order == 0 { + return Some(F::ONE); + } + self.order + .is_multiple_of(order) + .then(|| self.omega_order.pow([(self.order / order) as u64])) + } + + pub fn root(&self, order: usize) -> F { + self.checked_root(order) + .expect("Subgroup of requested order does not exist.") + } + /// Returns a cached table of roots of unity of the given order. fn roots_table(&self, order: usize) -> RwLockReadGuard<'_, Vec> { assert!( - self.domain.order.is_multiple_of(order), + self.order.is_multiple_of(order), "No subgroup of order {order}." ); @@ -305,7 +211,7 @@ impl NttEngine { roots.reserve_exact(size); // Compute powers of roots of unity. - let root = self.domain.root(size); + let root = self.root(size); #[cfg(not(feature = "parallel"))] { let mut root_i = F::ONE; @@ -457,29 +363,75 @@ impl NttEngine { } } -impl NttEngine { - /// Masked interleaved Reed-Solomon encoding over CPU slices. - /// - /// `messages` are `num_messages` already-chunked slices of `message_length` elements. - /// Used directly by benchmarks and tests; backend encoders call this (CPU) or reimplement - /// it on the device (Metal), all driven by [`Self::coset_params`]. +impl ReedSolomon for NttEngine { + fn next_order(&self, size: usize) -> Option { + match self.divisors.binary_search(&size) { + Ok(index) | Err(index) => self.divisors.get(index).copied(), + } + } + + fn generator(&self, codeword_length: usize) -> F { + self.omega_order + .pow([(self.order / codeword_length) as u64]) + } + + fn evaluation_points( + &self, + masked_message_length: usize, + codeword_length: usize, + indices: &[usize], + ) -> Vec { + assert!(masked_message_length <= codeword_length); + assert!(self.order.is_multiple_of(codeword_length)); + let mut result = Vec::new(); + let generator = self.generator(codeword_length); + + // Coset transformation + let mut coset_size = self.next_order(masked_message_length).unwrap(); + while !codeword_length.is_multiple_of(coset_size) { + coset_size = self.next_order(coset_size + 1).unwrap(); + } + #[cfg(not(feature = "rs_in_order"))] + let num_cosets = codeword_length / coset_size; + + for &index in indices { + assert!(index < codeword_length); + + #[cfg(not(feature = "rs_in_order"))] + let index = transpose_permute(index, num_cosets, coset_size); + result.push(generator.pow([index as u64])); + } + result + } + #[cfg_attr(feature = "tracing", instrument(skip(self, messages, masks), fields( - num_messages = messages.len(), - message_len = messages.first().map(|c| c.len()), + num_messages = messages.vectors.len() * messages.interleaving_depth, + message_len = messages.message_length, codeword_length = codeword_length, - mask_len = masks.len().checked_div(messages.len()) + mask_len = masks.len().checked_div(messages.vectors.len() * messages.interleaving_depth) )))] - pub fn interleaved_encode_slices( + fn interleaved_encode( &self, - messages: &[&[F]], - masks: &[F], + messages: Messages<'_, F>, + masks: &ActiveBuffer, codeword_length: usize, - ) -> Vec { - assert!(self.domain.order.is_multiple_of(codeword_length)); + ) -> ActiveBuffer { + let vectors = messages + .vectors + .iter() + .map(|v| v.as_slice()) + .collect::>(); + let messages = vectors + .iter() + .flat_map(|v| { + chunks_exact_or_empty(v, messages.message_length, messages.interleaving_depth) + }) + .collect::>(); + assert!(self.order.is_multiple_of(codeword_length)); if messages.is_empty() { assert!(masks.is_empty()); - return Vec::new(); + return ActiveBuffer::from_vec(Vec::new()); } let num_messages = messages.len(); let message_len = messages[0].len(); @@ -500,8 +452,11 @@ impl NttEngine { // You can also see this as applying a first round of Cooley-Tukey with // N = coset_size × num_cosets, and solving it directly by observing that // only the first coset is non-zero. - let (coset_size, num_cosets) = - self.domain.coset_params(masked_message_length, codeword_length); + let mut coset_size = self.next_order(masked_message_length).unwrap(); + while !codeword_length.is_multiple_of(coset_size) { + coset_size = self.next_order(coset_size + 1).unwrap(); + } + let num_cosets = codeword_length / coset_size; let coset_padding = coset_size - masked_message_length; // Lay out twisted coefficients in contiguous coset blocks of length @@ -509,7 +464,7 @@ impl NttEngine { let mut result = Vec::with_capacity(num_messages * codeword_length); for (message, mask) in zip_strict( messages, - chunks_exact_or_empty(masks, mask_length, num_messages), + chunks_exact_or_empty(masks.as_slice(), mask_length, num_messages), ) { // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension. for _ in 0..num_cosets { @@ -537,70 +492,7 @@ impl NttEngine { // Transpose to row-major order with vectors stacked horizontally. transpose(&mut result, num_messages, codeword_length); - result - } -} - -/// CPU Reed-Solomon encoder. -/// -/// Thin wrapper over a CPU [`NttEngine`]: the scalar methods delegate to the engine's shared -/// [`RsDomain`], and `interleaved_encode` chunks the input vectors and runs the host coset-NTT. -/// -/// Only built when CPU buffers are the active backend. Under a Metal build -/// `ActiveBuffer = MetalBuffer`, so its `ReedSolomon` impl would not type-check; code that needs -/// a CPU codeword there reaches it through [`NttEngine::interleaved_encode_slices`] instead. -#[cfg(not(all(feature = "metal", target_os = "macos")))] -#[derive(Debug, Clone)] -pub struct CpuRs { - engine: Arc>, -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -impl CpuRs { - pub const fn new(engine: Arc>) -> Self { - Self { engine } - } -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -impl ReedSolomon for CpuRs { - fn next_order(&self, size: usize) -> Option { - self.engine.domain().next_order(size) - } - - fn generator(&self, codeword_length: usize) -> F { - self.engine.domain().generator(codeword_length) - } - - fn evaluation_points( - &self, - masked_message_length: usize, - codeword_length: usize, - indices: &[usize], - ) -> Vec { - self.engine - .domain() - .evaluation_points(masked_message_length, codeword_length, indices) - } - - fn interleaved_encode( - &self, - vectors: &[&CpuBuffer], - masks: &CpuBuffer, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - ) -> CpuBuffer { - let vectors = vectors.iter().map(|v| v.as_slice()).collect::>(); - let messages = vectors - .iter() - .flat_map(|v| chunks_exact_or_empty(v, message_length, interleaving_depth)) - .collect::>(); - CpuBuffer::from_vec(self.engine.interleaved_encode_slices( - &messages, - masks.as_slice(), - codeword_length, - )) + ActiveBuffer::from_vec(result) } } @@ -712,12 +604,12 @@ mod tests { let engine = NttEngine::::new_from_fftfield(); // Verify that the order of the engine is correctly set - assert_eq!(engine.domain.order, 3 << 32); + assert_eq!(engine.order, 3 << 32); // Verify that the root of unity is correctly initialized assert_eq!( - engine.domain.omega_order, - Field64::GENERATOR.pow([(18_446_744_069_414_584_320 / engine.domain.order) as u64]) + engine.omega_order, + Field64::GENERATOR.pow([(18_446_744_069_414_584_320 / engine.order) as u64]) ); } @@ -726,31 +618,31 @@ mod tests { let engine = NttEngine::::new_from_fftfield(); // Ensure the root exponentiates correctly - assert_eq!(engine.domain().root(8).pow([8]), Field64::ONE); - assert_eq!(engine.domain().root(4).pow([4]), Field64::ONE); - assert_eq!(engine.domain().root(2).pow([2]), Field64::ONE); + assert_eq!(engine.root(8).pow([8]), Field64::ONE); + assert_eq!(engine.root(4).pow([4]), Field64::ONE); + assert_eq!(engine.root(2).pow([2]), Field64::ONE); // Ensure it's not a lower-order root - assert_ne!(engine.domain().root(8).pow([4]), Field64::ONE); - assert_ne!(engine.domain().root(4).pow([2]), Field64::ONE); + assert_ne!(engine.root(8).pow([4]), Field64::ONE); + assert_ne!(engine.root(4).pow([2]), Field64::ONE); } #[test] fn test_root_of_unity_multiplication() { let engine = NttEngine::::new_from_fftfield(); - let root = engine.domain().root(16); + let root = engine.root(16); // Multiply root by itself repeatedly and verify expected outcomes - assert_eq!(root.pow([2]), engine.domain().root(8)); - assert_eq!(root.pow([4]), engine.domain().root(4)); - assert_eq!(root.pow([8]), engine.domain().root(2)); + assert_eq!(root.pow([2]), engine.root(8)); + assert_eq!(root.pow([4]), engine.root(4)); + assert_eq!(root.pow([8]), engine.root(2)); } #[test] fn test_root_of_unity_inversion() { let engine = NttEngine::::new_from_fftfield(); - let root = engine.domain().root(16); + let root = engine.root(16); // The inverse of ω is ω^{-1}, computed as ω^(p-2) in Field64. let p: u64 = u64::from_be_bytes(Field64::MODULUS.to_bytes_be().try_into().unwrap()); @@ -776,9 +668,9 @@ mod tests { let engine2 = NttEngine::::new_from_fftfield(); // Ensure that multiple instances yield the same results - assert_eq!(engine1.domain().root(8), engine2.domain().root(8)); - assert_eq!(engine1.domain().root(4), engine2.domain().root(4)); - assert_eq!(engine1.domain().root(2), engine2.domain().root(2)); + assert_eq!(engine1.root(8), engine2.root(8)); + assert_eq!(engine1.root(4), engine2.root(4)); + assert_eq!(engine1.root(2), engine2.root(2)); } #[test] @@ -788,9 +680,9 @@ mod tests { // Check hardcoded expected values (ω^i) assert_eq!(roots_4[0], Field::ONE); - assert_eq!(roots_4[1], engine.domain().root(4)); - assert_eq!(roots_4[2], engine.domain().root(4).pow([2])); - assert_eq!(roots_4[3], engine.domain().root(4).pow([3])); + assert_eq!(roots_4[1], engine.root(4)); + assert_eq!(roots_4[2], engine.root(4).pow([2])); + assert_eq!(roots_4[3], engine.root(4).pow([3])); } #[test] @@ -802,7 +694,7 @@ mod tests { // Must contain only ω^0 and ω^1 assert_eq!(roots_2.len(), 2); assert_eq!(roots_2[0], Field64::ONE); - assert_eq!(roots_2[1], engine.domain().root(2)); + assert_eq!(roots_2[1], engine.root(2)); } #[test] @@ -813,7 +705,7 @@ mod tests { // Ensure the sequence follows expected powers of the root of unity for i in 0..4 { - assert_eq!(roots_4[i], engine.domain().root(4).pow([i as u64])); + assert_eq!(roots_4[i], engine.root(4).pow([i as u64])); } } @@ -859,7 +751,7 @@ mod tests { let roots = vec![r1]; // Ensure the root of unity is correct - assert_eq!(engine.domain().root(4).pow([4]), Field64::ONE); + assert_eq!(engine.root(4).pow([4]), Field64::ONE); apply_twiddles(&mut values, &roots, 2, 4); @@ -1178,7 +1070,7 @@ mod tests { // // ω is the 32nd root of unity: ω³² = 1. - let omega = engine.domain().root(32); + let omega = engine.root(32); let mut expected_values = vec![Field64::ZERO; 32]; for (k, expected_value) in expected_values.iter_mut().enumerate().take(32) { let omega_k = omega.pow([k as u64]); diff --git a/src/algebra/ntt/mod.rs b/src/algebra/ntt/mod.rs index 50782dd6..1e5d356d 100644 --- a/src/algebra/ntt/mod.rs +++ b/src/algebra/ntt/mod.rs @@ -11,48 +11,50 @@ use std::{ sync::{Arc, LazyLock}, }; -use ark_ff::{FftField, Field}; +use ark_ff::Field; use static_assertions::assert_obj_safe; -pub use self::cooley_tukey::{NttEngine, RsDomain}; use self::matrix::MatrixMut; -// The CPU encoder is only the active backend (and only constructible) on non-Metal builds. -#[cfg(not(all(feature = "metal", target_os = "macos")))] -pub use self::cooley_tukey::CpuRs; pub use self::{ + cooley_tukey::NttEngine, transpose::transpose, wavelet::{inverse_wavelet_transform, wavelet_transform}, }; use crate::{ - algebra::{buffer::ActiveBuffer, fields}, + algebra::fields, + buffer::{ActiveBuffer, BufferOps, DefaultRs}, type_map::{self, TypeMap}, }; pub static NTT: LazyLock> = LazyLock::new(|| { let map = TypeMap::new(); - register_ntt::(&map); - register_ntt::(&map); - register_ntt::(&map); - register_ntt::(&map); - register_ntt::(&map); - register_ntt::(&map); - register_ntt::<::BasePrimeField>(&map); - register_ntt::<::BasePrimeField>(&map); + map.insert( + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> + ); + map.insert( + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> + ); + map.insert( + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> + ); + map.insert( + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> + ); + map.insert( + Arc::new(DefaultRs::::new_from_fftfield()) as Arc>, + ); + map.insert( + Arc::new(DefaultRs::::new_from_fftfield()) as Arc>, + ); + map.insert(Arc::new( + DefaultRs::<::BasePrimeField>::new_from_fftfield(), + ) as Arc>); + map.insert(Arc::new( + DefaultRs::<::BasePrimeField>::new_from_fftfield(), + ) as Arc>); map }); -fn register_ntt(map: &TypeMap) { - // Both backends share the same `RsDomain` coset convention and differ only in how - // `interleaved_encode` runs: the CPU encoder needs the full NTT engine, while the - // Metal encoder only needs the shared domain (the NTT itself runs on the GPU). - #[cfg(all(feature = "metal", target_os = "macos"))] - let encoder = - crate::algebra::metal_buffer::MetalRs::new(Arc::new(RsDomain::::from_fftfield())); - #[cfg(not(all(feature = "metal", target_os = "macos")))] - let encoder = CpuRs::new(Arc::new(NttEngine::::new_from_fftfield())); - map.insert(Arc::new(encoder) as Arc>); -} - #[derive(Default)] pub struct NttFamily; @@ -60,11 +62,38 @@ impl type_map::Family for NttFamily { type Dyn = dyn ReedSolomon; } -/// Trait for a Reed-Solomon encoder implementation for a given field `F`. +/// Interleaved Reed-Solomon input buffers. /// -/// The scalar methods (`next_order`, `generator`, `evaluation_points`) describe the -/// field/domain structure and are backend independent. `interleaved_encode` produces a -/// codeword over the active backend buffer ([`ActiveBuffer`]). +/// Chunking buffer-backed vectors into per-message views can require backend work. This descriptor +/// keeps the encoding input compact and lets each backend derive the chunked messages internally. +/// +/// Each vector stores `interleaving_depth` consecutive messages of `message_length` field +/// elements. +#[derive(Clone, Copy)] +pub struct Messages<'a, F> { + pub(crate) vectors: &'a [&'a ActiveBuffer], + pub(crate) message_length: usize, + pub(crate) interleaving_depth: usize, +} + +impl<'a, F: Field> Messages<'a, F> { + pub fn new( + vectors: &'a [&'a ActiveBuffer], + message_length: usize, + interleaving_depth: usize, + ) -> Self { + assert!(vectors + .iter() + .all(|vector| vector.len() == message_length * interleaving_depth)); + Self { + vectors, + message_length, + interleaving_depth, + } + } +} + +/// Trait for a Reed-Solomon encoder implementation for a given field `F`. pub trait ReedSolomon: Debug + Send + Sync { /// Returns the next supported order equal or larger than `size`. /// @@ -91,9 +120,9 @@ pub trait ReedSolomon: Debug + Send + Sync { /// Compute a masked interleaved Reed-Solomon encoding. /// - /// `vectors` are `num_vectors` buffers, each holding `interleaving_depth` messages of + /// `messages` contains `num_vectors` buffers, each holding `interleaving_depth` messages of /// `message_length` elements laid out contiguously. `masks` is a - /// `num_vectors * interleaving_depth` × `mask_length` matrix of blinding coefficients. + /// `(num_vectors * interleaving_depth) × mask_length` matrix of blinding coefficients. /// `codeword_length` must be an NTT-smooth number >= `message_length + mask_length`. /// Returns a `codeword_length × (num_vectors * interleaving_depth)` matrix. /// @@ -101,10 +130,8 @@ pub trait ReedSolomon: Debug + Send + Sync { /// corresponding with the index of a coefficient list formed by concatenating message and mask. fn interleaved_encode( &self, - vectors: &[&ActiveBuffer], + messages: Messages<'_, F>, masks: &ActiveBuffer, - message_length: usize, - interleaving_depth: usize, codeword_length: usize, ) -> ActiveBuffer; } @@ -127,6 +154,16 @@ pub fn evaluation_points( .evaluation_points(masked_message_length, codeword_length, indices) } +pub fn interleaved_rs_encode( + messages: Messages<'_, F>, + masks: &ActiveBuffer, + codeword_length: usize, +) -> ActiveBuffer { + NTT.get::() + .expect("Unsupported NTT field.") + .interleaved_encode(messages, masks, codeword_length) +} + pub fn generator(codeword_length: usize) -> F { NTT.get::() .expect("Unsupported NTT field.") @@ -145,17 +182,18 @@ mod tests { use super::*; use crate::{ algebra::{random_vector, univariate_evaluate}, + buffer::BufferOps, utils::{chunks_exact_or_empty, zip_strict}, }; - fn valid_codeword_lengths(size: usize, count: usize) -> Vec { - let domain = RsDomain::::from_fftfield(); - iter::successors(domain.next_order(size), |size| domain.next_order(*size + 1)) + fn valid_codeword_lengths(size: usize, count: usize) -> Vec { + let ntt = NTT.get::().expect("No NTT engine for field."); + iter::successors(ntt.next_order(size), |size| ntt.next_order(*size + 1)) .take(count) .collect() } - fn test(engine: &NttEngine) + fn test(ntt: &dyn ReedSolomon) where Standard: Distribution, { @@ -188,10 +226,16 @@ mod tests { .map(|_| random_vector(&mut rng, message_length)) .collect::>(); let masks = random_vector(&mut rng, mask_length * num_messages); - let message_refs = messages.iter().map(|v| v.as_slice()).collect::>(); - let codeword = engine.interleaved_encode_slices( - &message_refs, - &masks, + let vectors = messages + .iter() + .map(|message| ActiveBuffer::from_slice(message)) + .collect::>(); + let vector_refs = vectors.iter().collect::>(); + let mask_buffer = ActiveBuffer::from_slice(&masks); + let rs_messages = Messages::new(&vector_refs, message_length, 1); + let codeword = ntt.interleaved_encode( + rs_messages, + &mask_buffer, codeword_length, ); @@ -199,7 +243,8 @@ mod tests { assert_eq!(codeword.len(), codeword_length * num_messages); // Output values are polynomial evaluations in the evaluation points. - let mut evaluation_points = engine.domain().evaluation_points(message_length + mask_length, codeword_length, &sampled_indices); + let mut evaluation_points = ntt.evaluation_points(message_length + mask_length, codeword_length, &sampled_indices); + let codeword = codeword.as_slice(); for (&index, &evaluation_point) in zip_strict(&sampled_indices, &evaluation_points) { let evaluations = &codeword[index * num_messages.. (index + 1) * num_messages]; let masks = chunks_exact_or_empty(&masks, mask_length, num_messages); @@ -223,7 +268,6 @@ mod tests { #[test] fn test_field64_1() { - let engine = NttEngine::::new_from_fftfield(); - test::(&engine); + test::(NTT.get().unwrap().as_ref()); } } diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index e2bccebd..58f0451b 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -10,7 +10,6 @@ use clap::Parser; use serde::Serialize; use whir::{ algebra::{ - buffer::ActiveBuffer, embedding::{Basefield, Embedding, Identity}, fields::{Field128, Field192, Field256, Field64, Field64_2, Field64_3}, linear_form::{Evaluate, LinearForm, MultilinearExtension}, @@ -22,7 +21,7 @@ use whir::{ transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; -use whir::buffer::BufferOps; +use whir::buffer::{ActiveBuffer, BufferOps}; #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] diff --git a/src/bin/gpu_proof_cpu_verify.rs b/src/bin/gpu_proof_cpu_verify.rs index a1f86362..d9c27dc4 100644 --- a/src/bin/gpu_proof_cpu_verify.rs +++ b/src/bin/gpu_proof_cpu_verify.rs @@ -7,16 +7,14 @@ use std::{ use clap::{Parser, Subcommand}; use serde::{Deserialize, Serialize}; use whir::{ - algebra::{ - buffer::ActiveBuffer, embedding::Identity, fields::Field256, linear_form::LinearForm, - }, + algebra::{embedding::Identity, fields::Field256, linear_form::LinearForm}, hash, parameters::ProtocolParameters, protocols::whir::Config as WhirConfig, transcript::{codecs::Empty, DomainSeparator, Proof, ProverState, VerifierState}, }; -use whir::buffer::BufferOps; +use whir::buffer::{ActiveBuffer, BufferOps}; type F = Field256; type M = Identity; diff --git a/src/bin/main.rs b/src/bin/main.rs index dc9dacae..552ea58b 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -6,7 +6,6 @@ use ark_std::rand::distributions::{Distribution, Standard}; use clap::Parser; use whir::{ algebra::{ - buffer::ActiveBuffer, embedding::{Basefield, Embedding, Identity}, fields::{Field128, Field192, Field256, Field64, Field64_2, Field64_3}, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, @@ -18,7 +17,7 @@ use whir::{ transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; -use whir::buffer::BufferOps; +use whir::buffer::{ActiveBuffer, BufferOps}; #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] diff --git a/src/buffer/cpu/buffer.rs b/src/buffer/cpu/buffer.rs new file mode 100644 index 00000000..1f9d145c --- /dev/null +++ b/src/buffer/cpu/buffer.rs @@ -0,0 +1,312 @@ +use std::{any::Any, mem}; + +use super::read_write::{impl_cpu_read, impl_cpu_write}; +use ark_ff::Field; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; + +use crate::{ + algebra::{ + embedding::Embedding, + linear_form::{Covector, LinearForm}, + }, + buffer::{Buffer, BufferOps, BufferRead, BufferWrite}, + engines::EngineId, + hash::{self, Hash}, + protocols::{ + matrix_commit::{hash_rows, Encodable}, + merkle_tree, + }, +}; + +#[derive( + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + Debug, + Default, + serde::Serialize, + serde::Deserialize, +)] +pub struct CpuBuffer { + pub(super) data: Vec, +} + +/// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. +pub struct CpuSlice<'a, F> { + pub(super) data: &'a [F], +} + +/// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. +pub struct CpuSliceMut<'a, F> { + pub(super) data: &'a mut [F], +} + +impl BufferOps for CpuBuffer { + type Nodes = CpuBuffer; + + fn at_index(&self, index: usize) -> Option { + if index >= self.len() { + None + } else { + Some(self.data[index]) + } + } + + fn as_slice(&self) -> &[T] { + &self.data + } + + fn len(&self) -> usize { + self.data.len() + } + + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { + let mut result = Vec::with_capacity(indices.len() * num_cols); + for i in indices { + result.extend_from_slice(&self.data[i * num_cols..(i + 1) * num_cols]); + } + result + } + + fn gather_at_indices(&self, indices: &[usize]) -> Vec { + indices + .iter() + .map(|&i| self.data[i]) + .collect() + } + + fn from_vec(source: Vec) -> Self { + Self { data: source } + } + + fn from_slice(source: &[T]) -> Self { + Self { + data: Vec::from(source), + } + } + + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (Self::Nodes, Hash) + where + T: Encodable + Send + Sync, + { + assert_eq!(self.len(), num_cols * merkle.num_leaves); + let engine = hash::ENGINES + .retrieve(leaf_hash) + .expect("Failed to retrieve hash engine"); + #[cfg(feature = "tracing")] + tracing::Span::current().record("engine", engine.name().as_ref()); + + let mut leaves = vec![Hash::default(); merkle.num_leaves]; + hash_rows(&*engine, &self.data, &mut leaves); + let nodes = merkle.build_nodes(leaves); + let root = nodes[merkle.num_nodes() - 1]; + (CpuBuffer { data: nodes }, root) + } +} + +impl Buffer for CpuBuffer { + fn zeros(length: usize) -> Self { + Self { + data: vec![F::ZERO; length], + } + } + + fn geometric_sequence(base: F, length: usize) -> Self { + Self { + data: crate::algebra::geometric_sequence(base, length), + } + } + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + Self { + data: (0..length).map(|_| rng.gen()).collect(), + } + } + + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self { + assert_eq!(linear_forms.len(), rlc_coeffs.len()); + let mut covector = vec![F::ZERO; size]; + if let Some((first, linear_forms)) = linear_forms.split_first_mut() { + debug_assert_eq!(rlc_coeffs[0], F::ONE); + if let Some(covector_form) = + (first.as_mut() as &mut dyn Any).downcast_mut::>() + { + mem::swap(&mut covector, &mut covector_form.vector); + } else { + first.accumulate(&mut covector, F::ONE); + } + for (rlc_coeff, linear_form) in rlc_coeffs[1..].iter().zip(linear_forms) { + linear_form.accumulate(&mut covector, *rlc_coeff); + } + } + Self { data: covector } + } + + fn zero_pad(&mut self) { + if !self.is_empty() { + self.data.resize(self.len().next_power_of_two(), F::ZERO); + } + } + + fn fold(&mut self, weight: F) { + crate::algebra::sumcheck::fold(&mut self.data, weight); + } + + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); + } + + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) + } + + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer { + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return CpuBuffer { data: Vec::new() }; + }; + debug_assert_eq!(coeffs[0], M::Target::ONE); + let mut accumulator = crate::algebra::lift(embedding, first.as_slice()); + for (coeff, vector) in coeffs[1..].iter().zip(vectors) { + crate::algebra::mixed_scalar_mul_add( + embedding, + &mut accumulator, + *coeff, + vector.as_slice(), + ); + } + CpuBuffer { data: accumulator } + } +} + +impl_cpu_read!(CpuSlice<'_, F>); +impl_cpu_read!(CpuSliceMut<'_, F>); +impl_cpu_read!(CpuBuffer); +impl_cpu_write!(CpuSliceMut<'_, F>); +impl_cpu_write!(CpuBuffer); + +/// Resolve any range expression (`a..b`, `..b`, `a..`, `..`) against `len`, +/// returning `(start, end)` and bounds-checking like slice indexing. +pub(crate) fn resolve_range( + range: impl std::ops::RangeBounds, + len: usize, +) -> (usize, usize) { + use std::ops::Bound::{Excluded, Included, Unbounded}; + let start = match range.start_bound() { + Included(&s) => s, + Excluded(&s) => s + 1, + Unbounded => 0, + }; + let end = match range.end_bound() { + Included(&e) => e + 1, + Excluded(&e) => e, + Unbounded => len, + }; + assert!( + start <= end && end <= len, + "slice range {start}..{end} out of bounds for length {len}" + ); + (start, end) +} + +#[cfg(test)] +mod tests { + use ark_ff::AdditiveGroup; + + use super::*; + use crate::algebra::{ + fields::Field64, geometric_accumulate, geometric_sequence as geometric_sequence_vec, + linear_form::UnivariateEvaluation, + }; + + type F = Field64; + + #[test] + fn geometric_sequence_matches_free_function() { + let base = F::from(7u64); + for length in 0..6 { + let buffer = CpuBuffer::::geometric_sequence(base, length); + assert_eq!( + BufferOps::as_slice(&buffer), + geometric_sequence_vec(base, length).as_slice(), + "geometric_sequence mismatch at length {length}" + ); + } + } + + #[test] + fn scale_multiplies_in_place() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.scalar_mul(weight); + let expected: Vec = values.iter().map(|&v| v * weight).collect(); + assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); + } + + #[test] + fn slice_accumulate_matches_restricted_geometric_accumulate() { + let len = 8usize; + let points = vec![F::from(3u64), F::from(5u64)]; + let scalars = vec![F::from(2u64), F::from(9u64)]; + + // Try a prefix slice (offset 0) and an interior slice (offset > 0). + for (offset, window_len) in [(0usize, 5usize), (2usize, 4usize)] { + let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); + let evaluators: Vec<_> = points + .iter() + .map(|&point| UnivariateEvaluation::new(point, window_len)) + .collect(); + buffer + .slice_mut(offset..offset + window_len) + .accumulate_univariate_evaluations(&evaluators, &scalars); + + // Reference: accumulate Σ_j scalars[j]·points[j]^i into the same + // sub-slice of a plain vector. + let mut reference = vec![F::ZERO; len]; + geometric_accumulate( + &mut reference[offset..offset + window_len], + scalars.clone(), + &points, + ); + + assert_eq!( + BufferOps::as_slice(&buffer), + reference.as_slice(), + "slice accumulate mismatch at offset {offset}, len {window_len}" + ); + } + } + + #[test] + fn slice_scale_multiplies_only_the_range() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.slice_mut(1..3).scalar_mul(weight); + let expected = vec![values[0], values[1] * weight, values[2] * weight, values[3]]; + assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); + } +} diff --git a/src/buffer/cpu/mod.rs b/src/buffer/cpu/mod.rs index cf77b25c..d8c36c4a 100644 --- a/src/buffer/cpu/mod.rs +++ b/src/buffer/cpu/mod.rs @@ -1,307 +1,5 @@ -mod traits; +mod buffer; +mod read_write; -use std::{any::Any, mem}; - -use ark_ff::Field; -use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; -use traits::{impl_cpu_read, impl_cpu_write}; - -use crate::{ - algebra::{ - embedding::Embedding, - linear_form::{Covector, LinearForm}, - }, - buffer::{Buffer, BufferOps, BufferRead, BufferWrite}, - engines::EngineId, - hash::{self, Hash}, - protocols::{ - matrix_commit::{hash_rows, Encodable}, - merkle_tree, - }, -}; - -#[derive( - Clone, - PartialEq, - Eq, - PartialOrd, - Ord, - Hash, - Debug, - Default, - serde::Serialize, - serde::Deserialize, -)] -pub struct CpuBuffer { - pub(super) data: Vec, -} - -/// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. -pub struct CpuSlice<'a, F> { - pub(super) data: &'a [F], -} - -/// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. -pub struct CpuSliceMut<'a, F> { - pub(super) data: &'a mut [F], -} - -impl_cpu_read!(CpuSlice<'_, F>); -impl_cpu_read!(CpuSliceMut<'_, F>); -impl_cpu_read!(CpuBuffer); -impl_cpu_write!(CpuSliceMut<'_, F>); -impl_cpu_write!(CpuBuffer); - -impl BufferOps for CpuBuffer { - type Nodes = CpuBuffer; - - fn at_index(&self, index: usize) -> Option { - if index >= self.len() { - None - } else { - Some(self.data[index]) - } - } - - fn as_slice(&self) -> &[T] { - &self.data - } - - fn len(&self) -> usize { - self.data.len() - } - - fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { - let mut result = Vec::with_capacity(indices.len() * num_cols); - for i in indices { - result.extend_from_slice(&self.data[i * num_cols..(i + 1) * num_cols]); - } - result - } - - fn from_vec(source: Vec) -> Self { - Self { data: source } - } - - fn from_slice(source: &[T]) -> Self { - Self { - data: Vec::from(source), - } - } - - fn merklize( - &self, - num_cols: usize, - leaf_hash: EngineId, - merkle: &merkle_tree::Config, - ) -> (Self::Nodes, Hash) - where - T: Encodable + Send + Sync, - { - assert_eq!(self.len(), num_cols * merkle.num_leaves); - let engine = hash::ENGINES - .retrieve(leaf_hash) - .expect("Failed to retrieve hash engine"); - #[cfg(feature = "tracing")] - tracing::Span::current().record("engine", engine.name().as_ref()); - - let mut leaves = vec![Hash::default(); merkle.num_leaves]; - hash_rows(&*engine, &self.data, &mut leaves); - let nodes = merkle.build_nodes(leaves); - let root = nodes[merkle.num_nodes() - 1]; - (CpuBuffer { data: nodes }, root) - } -} - -impl Buffer for CpuBuffer { - fn zeros(length: usize) -> Self { - Self { - data: vec![F::ZERO; length], - } - } - - fn geometric_sequence(base: F, length: usize) -> Self { - Self { - data: crate::algebra::geometric_sequence(base, length), - } - } - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution, - { - Self { - data: (0..length).map(|_| rng.gen()).collect(), - } - } - - fn linear_forms_rlc( - size: usize, - linear_forms: &mut [Box>], - rlc_coeffs: &[F], - ) -> Self { - assert_eq!(linear_forms.len(), rlc_coeffs.len()); - let mut covector = vec![F::ZERO; size]; - if let Some((first, linear_forms)) = linear_forms.split_first_mut() { - debug_assert_eq!(rlc_coeffs[0], F::ONE); - if let Some(covector_form) = - (first.as_mut() as &mut dyn Any).downcast_mut::>() - { - mem::swap(&mut covector, &mut covector_form.vector); - } else { - first.accumulate(&mut covector, F::ONE); - } - for (rlc_coeff, linear_form) in rlc_coeffs[1..].iter().zip(linear_forms) { - linear_form.accumulate(&mut covector, *rlc_coeff); - } - } - Self { data: covector } - } - - fn zero_pad(&mut self) { - if !self.is_empty() { - self.data.resize(self.len().next_power_of_two(), F::ZERO); - } - } - - fn fold(&mut self, weight: F) { - crate::algebra::sumcheck::fold(&mut self.data, weight); - } - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - self.fold(weight); - other.fold(weight); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - self.fold_pair(other, weight); - self.sumcheck_polynomial(other) - } - - fn mixed_linear_combination>( - embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer { - assert_eq!(vectors.len(), coeffs.len()); - let Some((first, vectors)) = vectors.split_first() else { - return CpuBuffer { data: Vec::new() }; - }; - debug_assert_eq!(coeffs[0], M::Target::ONE); - let mut accumulator = crate::algebra::lift(embedding, first.as_slice()); - for (coeff, vector) in coeffs[1..].iter().zip(vectors) { - crate::algebra::mixed_scalar_mul_add( - embedding, - &mut accumulator, - *coeff, - vector.as_slice(), - ); - } - CpuBuffer { data: accumulator } - } -} - -/// Resolve any range expression (`a..b`, `..b`, `a..`, `..`) against `len`, -/// returning `(start, end)` and bounds-checking like slice indexing. -pub(crate) fn resolve_range( - range: impl std::ops::RangeBounds, - len: usize, -) -> (usize, usize) { - use std::ops::Bound::{Excluded, Included, Unbounded}; - let start = match range.start_bound() { - Included(&s) => s, - Excluded(&s) => s + 1, - Unbounded => 0, - }; - let end = match range.end_bound() { - Included(&e) => e + 1, - Excluded(&e) => e, - Unbounded => len, - }; - assert!( - start <= end && end <= len, - "slice range {start}..{end} out of bounds for length {len}" - ); - (start, end) -} - -#[cfg(test)] -mod tests { - use ark_ff::AdditiveGroup; - - use super::*; - use crate::algebra::{ - fields::Field64, geometric_accumulate, geometric_sequence as geometric_sequence_vec, - linear_form::UnivariateEvaluation, - }; - - type F = Field64; - - #[test] - fn geometric_sequence_matches_free_function() { - let base = F::from(7u64); - for length in 0..6 { - let buffer = CpuBuffer::::geometric_sequence(base, length); - assert_eq!( - BufferOps::as_slice(&buffer), - geometric_sequence_vec(base, length).as_slice(), - "geometric_sequence mismatch at length {length}" - ); - } - } - - #[test] - fn scale_multiplies_in_place() { - let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; - let weight = F::from(5u64); - let mut buffer = CpuBuffer::from_vec(values.clone()); - buffer.scalar_mul(weight); - let expected: Vec = values.iter().map(|&v| v * weight).collect(); - assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); - } - - #[test] - fn slice_accumulate_matches_restricted_geometric_accumulate() { - let len = 8usize; - let points = vec![F::from(3u64), F::from(5u64)]; - let scalars = vec![F::from(2u64), F::from(9u64)]; - - // Try a prefix slice (offset 0) and an interior slice (offset > 0). - for (offset, window_len) in [(0usize, 5usize), (2usize, 4usize)] { - let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); - let evaluators: Vec<_> = points - .iter() - .map(|&point| UnivariateEvaluation::new(point, window_len)) - .collect(); - buffer - .slice_mut(offset..offset + window_len) - .accumulate_univariate_evaluations(&evaluators, &scalars); - - // Reference: accumulate Σ_j scalars[j]·points[j]^i into the same - // sub-slice of a plain vector. - let mut reference = vec![F::ZERO; len]; - geometric_accumulate( - &mut reference[offset..offset + window_len], - scalars.clone(), - &points, - ); - - assert_eq!( - BufferOps::as_slice(&buffer), - reference.as_slice(), - "slice accumulate mismatch at offset {offset}, len {window_len}" - ); - } - } - - #[test] - fn slice_scale_multiplies_only_the_range() { - let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; - let weight = F::from(5u64); - let mut buffer = CpuBuffer::from_vec(values.clone()); - buffer.slice_mut(1..3).scalar_mul(weight); - let expected = vec![values[0], values[1] * weight, values[2] * weight, values[3]]; - assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); - } -} +pub(crate) use buffer::resolve_range; +pub use buffer::{CpuBuffer, CpuSlice, CpuSliceMut}; diff --git a/src/buffer/cpu/traits.rs b/src/buffer/cpu/read_write.rs similarity index 96% rename from src/buffer/cpu/traits.rs rename to src/buffer/cpu/read_write.rs index fdeec2d6..b191e86a 100644 --- a/src/buffer/cpu/traits.rs +++ b/src/buffer/cpu/read_write.rs @@ -65,6 +65,10 @@ macro_rules! impl_cpu_read { data: &data[start..end], } } + + fn copy_to_owned(&self) -> CpuBuffer { + CpuBuffer::from_slice(&*self.data) + } } }; } diff --git a/src/buffer/metal/buffer.rs b/src/buffer/metal/buffer.rs new file mode 100644 index 00000000..4a73184d --- /dev/null +++ b/src/buffer/metal/buffer.rs @@ -0,0 +1,1291 @@ +// NOTE: 100% AI GENERATED + +use std::{ + any::type_name, + cell::OnceCell, + cmp::Ordering, + hash::{Hash, Hasher}, + marker::PhantomData, + ops::RangeBounds, +}; + +use ark_ff::{AdditiveGroup, Field}; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; +use metal::Buffer; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +use crate::{ + algebra::{ + embedding::{Embedding, Identity}, + fields::Field256, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, + }, + buffer::{BufferOps, BufferRead, BufferWrite}, + engines::EngineId, + hash::{self, Hash as Digest}, + protocols::{ + matrix_commit::{hash_rows, Encodable}, + merkle_tree, + }, +}; + +use super::runtime::{ + as_field256_slice, assert_bn254, copy_field_buffer, copy_field_buffer_at, download_field, + download_hash_indices, + encode_field_rows_le, field256_to_f, field256_to_target, f_to_field256, + geometric_accumulate_at_offset, geometric_accumulate_chunk_size, maybe_upload_bn254, + parallel_dot, parallel_dot_at, parallel_fold_pair_sumcheck, parallel_geometric_accumulate_point_blocks, + parallel_geometric_accumulate_point_blocks_batched, parallel_multilinear_extend_at, + parallel_sumcheck, parallel_sumcheck_at, parallel_univariate_evaluate, + parallel_univariate_evaluate_at, read_bn254_rows, run_in_place, runtime, scalar_mul_add_at, + scalar_mul_at_offset, should_use_geometric_point_blocks, should_use_geometric_point_blocks_batched, + target_to_field256, upload_field, zeroed_field_buffer, +}; +use super::sha2::MetalSha2; + +#[derive(Clone, Debug)] +pub(crate) struct MetalFieldBuffer { + pub(crate) limbs: Buffer, +} + +#[derive(Clone, Debug)] +pub(crate) struct MetalHashBuffer { + pub(crate) bytes: Buffer, +} + +#[derive(Clone, Debug)] +pub struct MetalBuffer { + len: usize, + host_cache: OnceCell>, + field: Option, + hash: Option, + _marker: PhantomData, +} + +impl PartialEq for MetalBuffer +where + T: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for MetalBuffer {} + +impl PartialOrd for MetalBuffer { + fn partial_cmp(&self, other: &Self) -> Option { + self.as_slice().partial_cmp(other.as_slice()) + } +} + +impl Ord for MetalBuffer { + fn cmp(&self, other: &Self) -> Ordering { + self.as_slice().cmp(other.as_slice()) + } +} + +impl Hash for MetalBuffer { + fn hash(&self, state: &mut H) { + self.as_slice().hash(state); + } +} + +impl Default for MetalBuffer { + fn default() -> Self { + Self { + len: 0, + host_cache: OnceCell::new(), + field: None, + hash: None, + _marker: PhantomData, + } + } +} + +impl Serialize for MetalBuffer { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.as_slice().serialize(serializer) + } +} + +impl<'de, T> Deserialize<'de> for MetalBuffer +where + T: Clone + Deserialize<'de>, +{ + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let data = Vec::::deserialize(deserializer)?; + Ok(BufferOps::from_vec(data)) + } +} + +impl MetalBuffer { + pub fn warmup() { + super::runtime::init(); + } + + pub(crate) fn as_slice(&self) -> &[T] { + self.host_cache + .get_or_init(|| self.download_host_cache()) + .as_slice() + } + + pub(crate) fn hash_bn254_rows_sha2(&self, num_cols: usize, out: &mut [Digest]) -> bool { + if type_name::() != type_name::() || self.field.is_none() { + return false; + } + assert_eq!(self.len(), num_cols * out.len()); + let message_size = num_cols * size_of::(); + let encoded = encode_field_rows_le( + &self + .field + .as_ref() + .expect("missing Metal field buffer") + .limbs, + out.len(), + num_cols, + ); + MetalSha2::new().hash_many_buffer(message_size, &encoded, out.len(), out); + true + } + + pub(crate) fn commit_bn254_rows_sha2_merkle( + &self, + num_cols: usize, + num_rows: usize, + layers: usize, + ) -> Option> { + if type_name::() != type_name::() || self.field.is_none() { + return None; + } + if num_rows != (1usize << layers) { + return None; + } + assert_eq!(self.len(), num_cols * num_rows); + let message_size = num_cols * size_of::(); + let encoded = encode_field_rows_le( + &self + .field + .as_ref() + .expect("missing Metal field buffer") + .limbs, + num_rows, + num_cols, + ); + let sha = MetalSha2::new(); + let nodes = sha.build_merkle_tree_buffer_from_messages_buffer( + message_size, + &encoded, + num_rows, + layers, + ); + Some(MetalBuffer::::from_digest_buffer( + nodes, + (1usize << (layers + 1)) - 1, + )) + } + + pub(crate) fn read_hash_indices(&self, indices: &[usize]) -> Option> { + let buffer = self.hash.as_ref()?; + Some(download_hash_indices(&buffer.bytes, self.len, indices)) + } +} + +impl MetalBuffer { + pub(crate) fn from_digest_buffer(bytes: Buffer, len: usize) -> Self { + Self { + len, + host_cache: OnceCell::new(), + field: None, + hash: Some(MetalHashBuffer { bytes }), + _marker: PhantomData, + } + } + + pub(crate) fn read_hash_at(&self, index: usize) -> Option { + self.read_hash_indices(&[index]) + .map(|mut values| values.pop().expect("missing hash")) + } +} + +impl BufferOps for MetalBuffer { + type Nodes = MetalBuffer; + + fn as_slice(&self) -> &[T] { + self.as_slice() + } + + fn at_index(&self, index: usize) -> Option { + self.as_slice().get(index).cloned() + } + + fn gather_at_indices(&self, indices: &[usize]) -> Vec + where + T: Copy, + { + if let Some(values) = self.read_hash_indices(indices) { + return values + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + indices + .iter() + .map(|&i| self.at_index(i).expect("index out of bounds")) + .collect() + } + + fn len(&self) -> usize { + self.len + } + + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { + if type_name::() == type_name::() && self.field.is_some() { + return read_bn254_rows( + self.field.as_ref().expect("missing Metal field buffer"), + num_cols, + indices, + ) + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + let data = self.as_slice(); + let mut result = Vec::with_capacity(indices.len() * num_cols); + for i in indices { + result.extend_from_slice(&data[i * num_cols..(i + 1) * num_cols]); + } + result + } + + fn from_vec(source: Vec) -> Self { + let len = source.len(); + let field = maybe_upload_bn254(&source); + let host_cache = if field.is_some() { + OnceCell::new() + } else { + OnceCell::from(source) + }; + Self { + len, + host_cache, + field, + hash: None, + _marker: PhantomData, + } + } + + fn from_slice(source: &[T]) -> Self { + Self::from_vec(Vec::from(source)) + } + + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (Self::Nodes, Digest) + where + T: Encodable + Send + Sync, + { + let num_rows = merkle.num_leaves; + let layers = merkle.layers.len(); + assert_eq!(self.len(), num_cols * num_rows); + + // Fast path: build the whole tree on the GPU when the leaf and every node layer is SHA2. + if leaf_hash == hash::SHA2 + && merkle + .layers + .iter() + .all(|layer| layer.hash_id == hash::SHA2) + { + if let Some(nodes) = self.commit_bn254_rows_sha2_merkle(num_cols, num_rows, layers) { + let root = nodes + .read_hash_at(merkle.num_nodes() - 1) + .expect("missing Metal Merkle root"); + return (nodes, root); + } + } + + let cpu_nodes = || { + let engine = hash::ENGINES + .retrieve(leaf_hash) + .expect("Failed to retrieve hash engine"); + let mut leaves = vec![Digest::default(); num_rows]; + hash_rows(&*engine, self.as_slice(), &mut leaves); + merkle.build_nodes(leaves) + }; + + let nodes = if leaf_hash == hash::SHA2 { + let mut leaves = vec![Digest::default(); num_rows]; + if self.hash_bn254_rows_sha2(num_cols, &mut leaves) { + merkle.build_nodes(leaves) + } else { + cpu_nodes() + } + } else { + cpu_nodes() + }; + let root = nodes[merkle.num_nodes() - 1]; + (BufferOps::from_vec(nodes), root) + } +} + +impl crate::buffer::Buffer for MetalBuffer { + fn zeros(length: usize) -> Self { + assert_bn254::(); + // Montgomery zero is all-zero bytes, so a device-side fill suffices. + Self { + len: length, + host_cache: OnceCell::new(), + field: Some(zeroed_field_buffer(length)), + hash: None, + _marker: PhantomData, + } + } + + fn geometric_sequence(base: F, length: usize) -> Self { + assert_bn254::(); + BufferOps::from_vec(crate::algebra::geometric_sequence(base, length)) + } + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + assert_bn254::(); + BufferOps::from_vec((0..length).map(|_| rng.gen()).collect()) + } + + fn zero_pad(&mut self) { + assert_bn254::(); + if !self.is_empty() { + let mut data = self.as_slice().to_vec(); + data.resize(self.len().next_power_of_two(), F::ZERO); + *self = BufferOps::from_vec(data); + } + } + + fn fold(&mut self, weight: F) { + assert_bn254::(); + if self.len() <= 1 { + return; + } + let len = self.len(); + let fold_half = len.next_power_of_two() >> 1; + let weight = upload_field(&[f_to_field256(weight)]); + let field = self.bn254_buffer(); + run_in_place( + "bn254_fold", + &[&field.limbs, &weight.limbs], + &[len as u32, fold_half as u32], + fold_half, + ); + self.len = fold_half; + self.invalidate_host_cache(); + } + + fn fold_pair(&mut self, other: &mut Self, weight: F) { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + if self.len() <= 1 { + return; + } + let len = self.len(); + let fold_half = len.next_power_of_two() >> 1; + let weight = upload_field(&[f_to_field256(weight)]); + let this = self.bn254_buffer(); + let other_buffer = other.bn254_buffer(); + run_in_place( + "bn254_fold_pair", + &[&this.limbs, &other_buffer.limbs, &weight.limbs], + &[len as u32, fold_half as u32], + fold_half, + ); + self.len = fold_half; + other.len = fold_half; + self.invalidate_host_cache(); + other.invalidate_host_cache(); + } + + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + if self.len() <= 1 { + return self.sumcheck_polynomial(other); + } + let len = self.len(); + let fold_half = len.next_power_of_two() >> 1; + if fold_half == 1 { + self.fold_pair(other, weight); + return self.sumcheck_polynomial(other); + } + let weight = upload_field(&[f_to_field256(weight)]); + let this = self.bn254_buffer(); + let other_buffer = other.bn254_buffer(); + let (c0, c2) = parallel_fold_pair_sumcheck(&this, &other_buffer, &weight, len, fold_half); + self.len = fold_half; + other.len = fold_half; + self.invalidate_host_cache(); + other.invalidate_host_cache(); + (field256_to_f::(c0), field256_to_f::(c2)) + } + + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self { + assert_bn254::(); + assert_eq!(linear_forms.len(), rlc_coeffs.len()); + let Some((first, rest)) = linear_forms.split_first_mut() else { + return Self::zeros(size); + }; + let first = (first.as_mut() as &mut dyn std::any::Any) + .downcast_mut::>() + .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); + let mut accumulator = >::from_slice(&first.vector); + for (coeff, linear_form) in rlc_coeffs[1..].iter().zip(rest) { + let covector = (linear_form.as_mut() as &mut dyn std::any::Any) + .downcast_mut::>() + .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); + let vector = >::from_slice(&covector.vector); + vector.mixed_scalar_mul_add_to(&Identity::new(), &mut accumulator, *coeff); + } + accumulator + } + + fn mixed_linear_combination>( + _embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer { + assert_bn254::(); + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return BufferOps::from_vec(Vec::new()); + }; + let mut accumulator = MetalBuffer:: { + len: first.len(), + host_cache: OnceCell::new(), + field: Some(copy_field_buffer(&first.bn254_buffer(), first.len())), + hash: None, + _marker: PhantomData, + }; + for (coeff, vector) in coeffs[1..].iter().copied().zip(vectors) { + vector.mixed_scalar_mul_add_to(_embedding, &mut accumulator, coeff); + } + accumulator + } +} + +/// Core geometric accumulate over `field[0..len]` (offset 0). Shared by the +/// owned buffer and a full-range view so the optimized chunk strategies are +/// written once. +fn geometric_accumulate_full( + field: &MetalFieldBuffer, + len: usize, + evaluators: &[UnivariateEvaluation], + scalars: &[F], +) { + let points = evaluators + .iter() + .map(|e| f_to_field256(e.point)) + .collect::>(); + let scalars = scalars + .iter() + .copied() + .map(f_to_field256) + .collect::>(); + let points = upload_field(&points); + let scalars = upload_field(&scalars); + let chunk_size = geometric_accumulate_chunk_size(len); + if std::env::var_os("WHIR_METAL_TRACE").is_some() { + eprintln!( + "metal geometric shape len={} points={} chunk={} chunks={}", + len, + evaluators.len(), + chunk_size, + len.div_ceil(chunk_size) + ); + } + if chunk_size <= 1 { + run_in_place( + "bn254_geometric_accumulate", + &[&field.limbs, &points.limbs, &scalars.limbs], + &[len as u32, evaluators.len() as u32], + len, + ); + } else { + let point_steps = evaluators + .iter() + .map(|e| f_to_field256(e.point.pow([chunk_size as u64]))) + .collect::>(); + let point_steps = upload_field(&point_steps); + if should_use_geometric_point_blocks(len, evaluators.len(), chunk_size) { + parallel_geometric_accumulate_point_blocks( + field, + &points, + &point_steps, + &scalars, + len, + evaluators.len(), + chunk_size, + ); + } else if should_use_geometric_point_blocks_batched(len, evaluators.len(), chunk_size) { + parallel_geometric_accumulate_point_blocks_batched( + field, + &points, + &point_steps, + &scalars, + len, + evaluators.len(), + chunk_size, + ); + } else { + run_in_place( + "bn254_geometric_accumulate_chunks_strided", + &[ + &field.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[len as u32, evaluators.len() as u32, chunk_size as u32], + len.div_ceil(chunk_size), + ); + } + } +} + +/// Validate that every evaluator targets `len` entries; returns `false` when +/// there is nothing to accumulate. +fn check_univariate_evaluators( + evaluators: &[UnivariateEvaluation], + scalars: &[F], + len: usize, +) -> bool { + assert_bn254::(); + assert_eq!(evaluators.len(), scalars.len()); + let Some(size) = evaluators.first().map(|e| e.size) else { + return false; + }; + assert_eq!(len, size); + for evaluator in evaluators { + assert_eq!(evaluator.size, size); + } + true +} + +impl BufferRead for MetalBuffer { + type TargetBuffer = MetalBuffer; + type Slice<'a> + = MetalSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.len() + } + + fn dot(&self, other: &Self) -> F { + assert_bn254::(); + assert_eq!(self.len(), other.len()); + let this = self.bn254_buffer(); + let other = other.bn254_buffer(); + field256_to_f::(parallel_dot(&this, &other, self.len())) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + assert_bn254::(); + let len = self.len().min(other.len()); + if len == 0 { + return (F::ZERO, F::ZERO); + } + if len == 1 { + return (self.as_slice()[0] * other.as_slice()[0], F::ZERO); + } + let fold_half = len.next_power_of_two() >> 1; + let this = self.bn254_buffer(); + let other = other.bn254_buffer(); + let (c0, c2) = parallel_sumcheck(&this, &other, len, fold_half); + (field256_to_f::(c0), field256_to_f::(c2)) + } + + fn mixed_extend, T: Field>( + &self, + _embedding: &M, + point: &[M::Target], + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let num_vars = point.len(); + let point = point + .iter() + .copied() + .map(target_to_field256) + .collect::>(); + let point = upload_field(&point); + let this = self.bn254_buffer(); + let value = parallel_multilinear_extend_at(&this, 0, self.len(), &point, num_vars); + field256_to_target::(value) + } + + fn mixed_dot, T: Field>( + &self, + _embedding: &M, + other: &MetalBuffer, + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let this = self.bn254_buffer(); + let other = other.bn254_buffer_target(); + let value = field256_to_f::(parallel_dot(&this, &other, self.len())); + field256_to_target::(f_to_field256(value)) + } + + fn mixed_univariate_evaluate>( + &self, + _embedding: &M, + point: M::Target, + ) -> M::Target { + assert_bn254::(); + let point = target_to_field256(point); + let point = upload_field(&[point]); + let this = self.bn254_buffer(); + field256_to_target::(parallel_univariate_evaluate(&this, &point, self.len())) + } + + fn mixed_scalar_mul_add_to>( + &self, + _embedding: &M, + accumulator: &mut MetalBuffer, + weight: M::Target, + ) { + assert_bn254::(); + let weight = upload_field(&[target_to_field256(weight)]); + let vector = self.bn254_buffer(); + let acc = accumulator.bn254_buffer_target(); + scalar_mul_add_at(&acc, 0, &vector, 0, &weight, self.len()); + accumulator.field = Some(acc); + accumulator.invalidate_host_cache(); + } + + fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { + assert_bn254::(); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); + MetalSlice { + field: self.bn254_buffer(), + offset: start, + len: end - start, + _parent: PhantomData, + } + } + + fn copy_to_owned(&self) -> MetalBuffer { + assert_bn254::(); + MetalBuffer { + len: self.len, + host_cache: OnceCell::new(), + field: Some(copy_field_buffer(&self.bn254_buffer(), self.len)), + hash: None, + _marker: PhantomData, + } + } +} + +impl BufferWrite for MetalBuffer { + type SliceMut<'a> + = MetalSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scalar_mul(&mut self, weight: F) { + assert_bn254::(); + if self.is_empty() { + return; + } + let weight = upload_field(&[f_to_field256(weight)]); + let field = self.bn254_buffer(); + run_in_place( + "bn254_scalar_mul", + &[&field.limbs, &weight.limbs], + &[self.len() as u32], + self.len(), + ); + self.field = Some(field); + self.invalidate_host_cache(); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + if !check_univariate_evaluators(evaluators, scalars, self.len()) { + return; + } + let field = self.bn254_buffer(); + geometric_accumulate_full(&field, self.len(), evaluators, scalars); + self.field = Some(field); + self.invalidate_host_cache(); + } + + fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { + assert_bn254::(); + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); + MetalSliceMut::new(self, start, end - start) + } + + fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { + assert_bn254::(); + let len = self.len(); + assert!(mid <= len, "split_at_mut mid {mid} out of bounds for {len}"); + // Materialize the parent's GPU allocation and invalidate its host + // cache once up front; the returned views share the handle and write + // disjoint ranges, so no per-op parent borrow is needed. + let field = self.bn254_buffer(); + self.field = Some(field.clone()); + self.invalidate_host_cache(); + ( + MetalSliceMut::from_field(field.clone(), 0, mid), + MetalSliceMut::from_field(field, mid, len - mid), + ) + } +} +impl MetalBuffer { + pub(crate) fn bn254_buffer(&self) -> MetalFieldBuffer { + self.field + .clone() + .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) + } + + fn invalidate_host_cache(&mut self) { + let _ = self.host_cache.take(); + } + + fn download_host_cache(&self) -> Vec { + if self.field.is_some() && type_name::() == type_name::() { + return download_field( + &self + .field + .as_ref() + .expect("missing Metal field buffer") + .limbs, + self.len, + ) + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + if self.hash.is_some() && type_name::() == type_name::() { + return download_hash_indices( + &self.hash.as_ref().expect("missing Metal hash buffer").bytes, + self.len, + &(0..self.len).collect::>(), + ) + .into_iter() + .map(|value| unsafe { std::mem::transmute_copy(&value) }) + .collect(); + } + panic!( + "MetalBuffer<{}> has no host cache and cannot be materialized", + type_name::() + ); + } +} + +impl MetalBuffer { + pub(crate) fn from_field_limb_buffer(limbs: Buffer, len: usize) -> Self { + Self { + len, + host_cache: OnceCell::new(), + field: Some(MetalFieldBuffer { limbs }), + hash: None, + _marker: PhantomData, + } + } +} + +impl MetalBuffer { + fn bn254_buffer_target(&self) -> MetalFieldBuffer { + self.field + .clone() + .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) + } +} + +/// Read-only, zero-copy view into a [`MetalBuffer`]'s GPU allocation. +/// +/// Holds a clone of the parent's allocation handle plus an `(offset, len)` +/// window. Reductions bind the handle at the byte offset, so they read +/// `parent[offset + i]` without copying. +pub struct MetalSlice<'a, F> { + field: MetalFieldBuffer, + offset: usize, + len: usize, + _parent: PhantomData<&'a MetalBuffer>, +} + +/// Mutable, zero-copy view into a [`MetalBuffer`]'s GPU allocation. +/// +/// Like [`MetalSlice`] but exclusive: the parent's host cache is invalidated +/// up front when the view is created, and writes go through the shared handle +/// at the byte offset, landing in the parent's own memory. +pub struct MetalSliceMut<'a, F> { + field: MetalFieldBuffer, + offset: usize, + len: usize, + _parent: PhantomData<&'a mut MetalBuffer>, +} + +impl<'a, F: Field + Clone> MetalSliceMut<'a, F> { + /// Borrow `parent[offset..offset+len]`, materializing the parent's GPU + /// allocation and invalidating its host cache once up front. + fn new(parent: &'a mut MetalBuffer, offset: usize, len: usize) -> Self { + let field = parent.bn254_buffer(); + parent.field = Some(field.clone()); + parent.invalidate_host_cache(); + Self { + field, + offset, + len, + _parent: PhantomData, + } + } + + /// Build a view directly from an already-materialized handle (used by + /// `split_at_mut`, where the parent cache was invalidated by the caller). + fn from_field(field: MetalFieldBuffer, offset: usize, len: usize) -> Self { + Self { + field, + offset, + len, + _parent: PhantomData, + } + } + + fn as_read(&self) -> MetalSlice<'_, F> { + MetalSlice { + field: self.field.clone(), + offset: self.offset, + len: self.len, + _parent: PhantomData, + } + } +} + +impl MetalSlice<'_, F> { + pub fn len(&self) -> usize { + self.len + } + pub fn is_empty(&self) -> bool { + self.len == 0 + } +} + +impl MetalSliceMut<'_, F> { + pub fn len(&self) -> usize { + self.len + } + pub fn is_empty(&self) -> bool { + self.len == 0 + } +} + +impl BufferRead for MetalSlice<'_, F> { + type TargetBuffer = MetalBuffer; + type Slice<'a> + = MetalSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.len + } + + fn dot(&self, other: &Self) -> F { + assert_bn254::(); + assert_eq!(self.len, other.len); + field256_to_f::(parallel_dot_at( + &self.field, + self.offset, + &other.field, + other.offset, + self.len, + )) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + assert_bn254::(); + let len = self.len.min(other.len); + if len == 0 { + return (F::ZERO, F::ZERO); + } + let fold_half = len.next_power_of_two() >> 1; + let (c0, c2) = parallel_sumcheck_at( + &self.field, + self.offset, + &other.field, + other.offset, + len, + fold_half, + ); + (field256_to_f::(c0), field256_to_f::(c2)) + } + + fn mixed_extend, T: Field>( + &self, + _embedding: &M, + point: &[M::Target], + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let num_vars = point.len(); + let point = point + .iter() + .copied() + .map(target_to_field256) + .collect::>(); + let point = upload_field(&point); + let value = + parallel_multilinear_extend_at(&self.field, self.offset, self.len, &point, num_vars); + field256_to_target::(value) + } + + fn mixed_dot, T: Field>( + &self, + _embedding: &M, + other: &MetalBuffer, + ) -> M::Target { + assert_bn254::(); + assert_bn254::(); + let other = other.bn254_buffer_target(); + let value = field256_to_f::(parallel_dot_at( + &self.field, + self.offset, + &other, + 0, + self.len, + )); + field256_to_target::(f_to_field256(value)) + } + + fn mixed_univariate_evaluate>( + &self, + _embedding: &M, + point: M::Target, + ) -> M::Target { + assert_bn254::(); + let point = upload_field(&[target_to_field256(point)]); + field256_to_target::(parallel_univariate_evaluate_at( + &self.field, + self.offset, + &point, + self.len, + )) + } + + fn mixed_scalar_mul_add_to>( + &self, + _embedding: &M, + accumulator: &mut MetalBuffer, + weight: M::Target, + ) { + assert_bn254::(); + let weight = upload_field(&[target_to_field256(weight)]); + let acc = accumulator.bn254_buffer_target(); + scalar_mul_add_at(&acc, 0, &self.field, self.offset, &weight, self.len); + accumulator.field = Some(acc); + accumulator.invalidate_host_cache(); + } + + fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); + MetalSlice { + field: self.field.clone(), + offset: self.offset + start, + len: end - start, + _parent: PhantomData, + } + } + + fn copy_to_owned(&self) -> MetalBuffer { + assert_bn254::(); + MetalBuffer { + len: self.len, + host_cache: OnceCell::new(), + field: Some(copy_field_buffer_at(&self.field, self.offset, self.len)), + hash: None, + _marker: PhantomData, + } + } +} + +impl BufferRead for MetalSliceMut<'_, F> { + type TargetBuffer = MetalBuffer; + type Slice<'a> + = MetalSlice<'a, F> + where + Self: 'a, + F: 'a; + + fn read_len(&self) -> usize { + self.len + } + + fn dot(&self, other: &Self) -> F { + self.as_read().dot(&other.as_read()) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + self.as_read().sumcheck_polynomial(&other.as_read()) + } + + fn mixed_extend, T: Field>( + &self, + embedding: &M, + point: &[M::Target], + ) -> M::Target { + self.as_read().mixed_extend(embedding, point) + } + + fn mixed_dot, T: Field>( + &self, + embedding: &M, + other: &MetalBuffer, + ) -> M::Target { + self.as_read().mixed_dot(embedding, other) + } + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + self.as_read().mixed_univariate_evaluate(embedding, point) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut MetalBuffer, + weight: M::Target, + ) { + self.as_read() + .mixed_scalar_mul_add_to(embedding, accumulator, weight) + } + + fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); + MetalSlice { + field: self.field.clone(), + offset: self.offset + start, + len: end - start, + _parent: PhantomData, + } + } + + fn copy_to_owned(&self) -> MetalBuffer { + self.as_read().copy_to_owned() + } +} + +impl BufferWrite for MetalSliceMut<'_, F> { + type SliceMut<'a> + = MetalSliceMut<'a, F> + where + Self: 'a, + F: 'a; + + fn scalar_mul(&mut self, weight: F) { + assert_bn254::(); + if self.len == 0 { + return; + } + let weight = upload_field(&[f_to_field256(weight)]); + scalar_mul_at_offset(&self.field, self.offset, self.len, &weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + if !check_univariate_evaluators(evaluators, scalars, self.len) { + return; + } + if self.offset == 0 { + // Offset 0: reuse the optimized full-buffer accumulate path. + geometric_accumulate_full(&self.field, self.len, evaluators, scalars); + } else { + // Arbitrary offset: bind the buffer at a byte offset so the + // kernel's gid-based indexing addresses field[offset + gid]. + let points = evaluators + .iter() + .map(|e| f_to_field256(e.point)) + .collect::>(); + let scalars = scalars + .iter() + .copied() + .map(f_to_field256) + .collect::>(); + let points = upload_field(&points); + let scalars = upload_field(&scalars); + geometric_accumulate_at_offset( + &self.field, + self.offset, + self.len, + &points, + &scalars, + evaluators.len(), + ); + } + } + + fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { + let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); + MetalSliceMut::from_field(self.field.clone(), self.offset + start, end - start) + } + + fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { + assert!( + mid <= self.len, + "split_at_mut mid {mid} out of bounds for {}", + self.len + ); + ( + MetalSliceMut::from_field(self.field.clone(), self.offset, mid), + MetalSliceMut::from_field(self.field.clone(), self.offset + mid, self.len - mid), + ) + } +} +#[cfg(test)] +mod tests { + use super::*; + use crate::algebra::ntt::{Messages, NttEngine, ReedSolomon}; + use crate::algebra::univariate_evaluate; + use crate::buffer::metal::rs::rs_transpose_permute; + use crate::buffer::{Buffer, BufferOps, CpuBuffer, MetalRs}; + + fn values(len: usize, offset: u64) -> Vec { + (0..len) + .map(|i| Field256::from(i as u64 + offset)) + .collect() + } + + #[test] + fn metal_bn254_dot_matches_cpu() { + let a = values(33, 1); + let b = values(33, 9); + let cpu_a = CpuBuffer::from_slice(&a); + let cpu_b = CpuBuffer::from_slice(&b); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_b = MetalBuffer::from_slice(&b); + assert_eq!(gpu_a.dot(&gpu_b), cpu_a.dot(&cpu_b)); + } + + #[test] + fn metal_bn254_fold_matches_cpu() { + let mut cpu = CpuBuffer::from_vec(values(31, 2)); + let mut gpu = MetalBuffer::from_vec(values(31, 2)); + let weight = Field256::from(42); + cpu.fold(weight); + gpu.fold(weight); + assert_eq!(gpu.as_slice(), cpu.as_slice()); + } + + #[test] + fn metal_bn254_sumcheck_matches_cpu() { + for len in [1, 2, 27, 64, 65] { + let a = values(len, 3); + let b = values(len, 11); + let cpu_a = CpuBuffer::from_slice(&a); + let cpu_b = CpuBuffer::from_slice(&b); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_b = MetalBuffer::from_slice(&b); + assert_eq!( + gpu_a.sumcheck_polynomial(&gpu_b), + cpu_a.sumcheck_polynomial(&cpu_b) + ); + } + } + + #[test] + fn metal_bn254_fold_pair_sumcheck_matches_cpu() { + for len in [2, 3, 27, 64, 65] { + let mut cpu_a = CpuBuffer::from_vec(values(len, 3)); + let mut cpu_b = CpuBuffer::from_vec(values(len, 11)); + let mut gpu_a = MetalBuffer::from_vec(values(len, 3)); + let mut gpu_b = MetalBuffer::from_vec(values(len, 11)); + let weight = Field256::from(42); + let cpu_result = cpu_a.fold_pair_sumcheck_polynomial(&mut cpu_b, weight); + let gpu_result = gpu_a.fold_pair_sumcheck_polynomial(&mut gpu_b, weight); + assert_eq!(gpu_result, cpu_result); + assert_eq!(gpu_a.as_slice(), cpu_a.as_slice()); + assert_eq!(gpu_b.as_slice(), cpu_b.as_slice()); + } + } + + #[test] + fn metal_bn254_scalar_mul_add_matches_cpu() { + let mut cpu = CpuBuffer::from_vec(values(19, 1)); + let mut gpu = MetalBuffer::from_vec(values(19, 1)); + let vector = values(19, 5); + let cpu_vector = CpuBuffer::from_slice(&vector); + let gpu_vector = MetalBuffer::from_slice(&vector); + let weight = Field256::from(7); + cpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut cpu, weight); + gpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut gpu, weight); + assert_eq!(gpu.as_slice(), cpu.as_slice()); + } + + #[test] + fn metal_bn254_interleaved_rs_encode_matches_cpu() { + // The GPU encoder and the CPU reference derive the same coset layout from the field. + let engine = NttEngine::::new_from_fftfield(); + let gpu_rs = MetalRs::::new_from_fftfield(); + + let a = values(8, 1); + let gpu_a = MetalBuffer::from_slice(&a); + let gpu_masks = MetalBuffer::from_slice(&[]); + + let messages = a.chunks_exact(4).collect::>(); + let generator = engine.root(8); + let mut cpu = Vec::with_capacity(16); + for index in 0..8 { + #[cfg(not(feature = "rs_in_order"))] + let index = rs_transpose_permute(index, 2, 4); + let point = generator.pow([index as u64]); + for message in &messages { + cpu.push(univariate_evaluate(message, point)); + } + } + let gpu_vectors = [&gpu_a]; + let gpu_messages = Messages::new(&gpu_vectors, 4, 2); + let gpu = gpu_rs.interleaved_encode(gpu_messages, &gpu_masks, 8); + assert_eq!(gpu.as_slice(), cpu.as_slice()); + } + + #[test] + fn metal_bn254_mixed_extend_matches_cpu() { + let values = values(8, 3); + let point = vec![Field256::from(2), Field256::from(5), Field256::from(9)]; + let cpu = CpuBuffer::from_slice(&values); + let gpu = MetalBuffer::from_slice(&values); + assert_eq!( + gpu.mixed_extend(&Identity::new(), &point), + cpu.mixed_extend(&Identity::new(), &point) + ); + } +} diff --git a/src/buffer/metal/kernels.rs b/src/buffer/metal/kernels.rs new file mode 100644 index 00000000..d12b4a59 --- /dev/null +++ b/src/buffer/metal/kernels.rs @@ -0,0 +1,943 @@ +// NOTE: 100% AI GENERATED + +pub(crate) const METAL_SOURCE: &str = r#" +#include +using namespace metal; + +struct F { ulong v[4]; }; + +constant ulong MODULUS[4] = { + 0x43e1f593f0000001UL, + 0x2833e84879b97091UL, + 0xb85045b68181585dUL, + 0x30644e72e131a029UL, +}; + +constant ulong BN254_N0PRIME = 0xc2e1f593efffffffUL; +constant F CANONICAL_ONE = {{1UL, 0UL, 0UL, 0UL}}; +constant F MONT_ONE = {{ + 0xac96341c4ffffffbUL, + 0x36fc76959f60cd29UL, + 0x666ea36f7879462eUL, + 0xe0a77c19a07df2fUL, +}}; + +struct StageConfig { + uint row_len; + uint half_m; + uint twiddle_offset; + uint _pad0; +}; + +struct BitReverseParams { + uint row_len; + uint log_n; + uint total_elements; + uint _pad0; +}; + +struct TransposeParams { + uint rows; + uint cols; + uint total_elements; +}; + +struct ReplicateCosetsParams { + uint row_len; + uint coset_size; + uint trailing_elements; +}; + +struct PackSingleVectorParams { + uint row_count; + uint message_length; + uint codeword_length; + uint coset_size; + uint total_elements; +}; + +struct FieldBytesParams { + uint rows; + uint cols; +}; + +struct ApplyCosetTwiddlesParams { + uint row_count; + uint num_cosets; + uint coset_size; + uint codeword_length; + uint total_elements; +}; + +inline F zero_f() { + F r; + r.v[0] = 0; r.v[1] = 0; r.v[2] = 0; r.v[3] = 0; + return r; +} + +inline F one_f() { + return MONT_ONE; +} + +inline F load_f(device const ulong *data, uint idx) { + F r; + uint base = idx * 4; + r.v[0] = data[base + 0]; + r.v[1] = data[base + 1]; + r.v[2] = data[base + 2]; + r.v[3] = data[base + 3]; + return r; +} + +inline void store_f(device ulong *data, uint idx, F x) { + uint base = idx * 4; + data[base + 0] = x.v[0]; + data[base + 1] = x.v[1]; + data[base + 2] = x.v[2]; + data[base + 3] = x.v[3]; +} + +inline bool ge_mod(F a) { + for (int i = 3; i >= 0; i--) { + if (a.v[i] > MODULUS[i]) return true; + if (a.v[i] < MODULUS[i]) return false; + } + return true; +} + +inline bool ge_f(F a, F b) { + for (int i = 3; i >= 0; i--) { + if (a.v[i] > b.v[i]) return true; + if (a.v[i] < b.v[i]) return false; + } + return true; +} + +inline F sub_raw(F a, thread const ulong *b, thread bool &borrow) { + F r; + borrow = false; + for (uint i = 0; i < 4; i++) { + ulong bi = b[i] + (borrow ? 1UL : 0UL); + bool bcarry = borrow && bi == 0; + r.v[i] = a.v[i] - bi; + borrow = bcarry || a.v[i] < bi; + } + return r; +} + +inline F sub_modulus(F a) { + F r; + bool borrow = false; + for (uint i = 0; i < 4; i++) { + ulong bi = MODULUS[i] + (borrow ? 1UL : 0UL); + bool bcarry = borrow && bi == 0; + r.v[i] = a.v[i] - bi; + borrow = bcarry || a.v[i] < bi; + } + return r; +} + +inline F add_modulus(F a) { + F r; + bool carry = false; + for (uint i = 0; i < 4; i++) { + ulong mi = MODULUS[i]; + ulong sum = a.v[i] + mi; + bool c0 = sum < a.v[i]; + ulong sum2 = sum + (carry ? 1UL : 0UL); + bool c1 = carry && sum2 == 0; + r.v[i] = sum2; + carry = c0 || c1; + } + return r; +} + +inline F add_f(F a, F b) { + F r; + bool carry = false; + for (uint i = 0; i < 4; i++) { + ulong sum = a.v[i] + b.v[i]; + bool c0 = sum < a.v[i]; + ulong sum2 = sum + (carry ? 1UL : 0UL); + bool c1 = carry && sum2 == 0; + r.v[i] = sum2; + carry = c0 || c1; + } + if (carry || ge_mod(r)) { + r = sub_modulus(r); + } + return r; +} + +inline F sub_f(F a, F b) { + ulong bv[4] = { b.v[0], b.v[1], b.v[2], b.v[3] }; + bool borrow = false; + F r = sub_raw(a, bv, borrow); + if (borrow) { + r = add_modulus(r); + } + return r; +} + +inline F double_f(F a) { + return add_f(a, a); +} + +inline ulong add_with_carry(ulong a, ulong b, thread ulong &carry) { + ulong sum = a + b; + ulong c1 = sum < a ? 1UL : 0UL; + ulong sum_with_carry = sum + carry; + ulong c2 = sum_with_carry < sum ? 1UL : 0UL; + carry = c1 + c2; + return sum_with_carry; +} + +inline void add_scaled_step(thread ulong &dst, ulong s, ulong a, thread ulong &carry) { + ulong product_lo = s * a; + ulong product_hi = mulhi(s, a); + + ulong sum = dst + product_lo; + ulong carry0 = sum < dst ? 1UL : 0UL; + ulong sum_with_carry = sum + carry; + ulong carry1 = sum_with_carry < sum ? 1UL : 0UL; + + dst = sum_with_carry; + carry = product_hi + carry0 + carry1; +} + +inline void add_scaled(thread ulong *dst, ulong s, ulong a0, ulong a1, ulong a2, ulong a3) { + ulong carry = 0; + add_scaled_step(dst[0], s, a0, carry); + add_scaled_step(dst[1], s, a1, carry); + add_scaled_step(dst[2], s, a2, carry); + add_scaled_step(dst[3], s, a3, carry); + dst[4] += carry; +} + +inline F mont_mul(F lhs, F rhs) { + ulong buf[9] = {0}; + uint off = 0; + +#pragma clang loop unroll(enable) + for (uint i = 0; i < 4; i++) { + add_scaled(&buf[off], lhs.v[i], rhs.v[0], rhs.v[1], rhs.v[2], rhs.v[3]); + + ulong m = buf[off] * BN254_N0PRIME; + add_scaled( + &buf[off], + m, + MODULUS[0], + MODULUS[1], + MODULUS[2], + MODULUS[3] + ); + + off += 1; + buf[off + 4] = 0; + } + + F result; + result.v[0] = buf[off + 0]; + result.v[1] = buf[off + 1]; + result.v[2] = buf[off + 2]; + result.v[3] = buf[off + 3]; + if (ge_mod(result)) { + result = sub_modulus(result); + } + return result; +} + +inline F from_mont(F value) { + F result = mont_mul(value, CANONICAL_ONE); + if (ge_mod(result)) { + result = sub_modulus(result); + } + return result; +} + +inline F mul_f(F a, F b) { + return mont_mul(a, b); +} + +inline F pow_f(F base, uint exp) { + F acc = one_f(); + F x = base; + while (exp != 0) { + if ((exp & 1) != 0) { + acc = mul_f(acc, x); + } + exp >>= 1; + if (exp != 0) { + x = mul_f(x, x); + } + } + return acc; +} + +kernel void bn254_fold( + device ulong *values [[buffer(0)]], + device const ulong *weight_buf [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &fold_half [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= fold_half) return; + F low = gid < len ? load_f(values, gid) : zero_f(); + F high = gid + fold_half < len ? load_f(values, gid + fold_half) : zero_f(); + F weight = load_f(weight_buf, 0); + store_f(values, gid, add_f(low, mul_f(sub_f(high, low), weight))); +} + +inline F fold_value_at(device const ulong *values, uint len, uint fold_half, uint idx, F weight) { + F low = idx < len ? load_f(values, idx) : zero_f(); + F high = idx + fold_half < len ? load_f(values, idx + fold_half) : zero_f(); + return add_f(low, mul_f(sub_f(high, low), weight)); +} + +kernel void bn254_fold_pair( + device ulong *a [[buffer(0)]], + device ulong *b [[buffer(1)]], + device const ulong *weight_buf [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &fold_half [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= fold_half) return; + F weight = load_f(weight_buf, 0); + store_f(a, gid, fold_value_at(a, len, fold_half, gid, weight)); + store_f(b, gid, fold_value_at(b, len, fold_half, gid, weight)); +} + +kernel void bn254_scalar_mul_add( + device ulong *acc [[buffer(0)]], + device const ulong *vector [[buffer(1)]], + device const ulong *weight_buf [[buffer(2)]], + constant uint &len [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F weight = load_f(weight_buf, 0); + store_f(acc, gid, add_f(load_f(acc, gid), mul_f(weight, load_f(vector, gid)))); +} + +kernel void bn254_scalar_mul( + device ulong *acc [[buffer(0)]], + device const ulong *weight_buf [[buffer(1)]], + constant uint &len [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F weight = load_f(weight_buf, 0); + store_f(acc, gid, mul_f(weight, load_f(acc, gid))); +} + +kernel void bn254_dot( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + F acc = zero_f(); + for (uint i = 0; i < len; i++) { + acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); + } + store_f(out, 0, acc); +} + +kernel void bn254_sumcheck( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &fold_half [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = 0; i < fold_half; i++) { + F a0 = i < len ? load_f(a, i) : zero_f(); + F b0 = i < len ? load_f(b, i) : zero_f(); + F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); + F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); + c0 = add_f(c0, mul_f(a0, b0)); + c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); + } + store_f(out, 0, c0); + store_f(out, 1, c2); +} + +kernel void bn254_dot_chunks( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &chunk_size [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + F acc = zero_f(); + for (uint i = start; i < end; i++) { + acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); + } + store_f(out, gid, acc); +} + +kernel void bn254_sum_chunks( + device const ulong *input [[buffer(0)]], + device ulong *out [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &offset [[buffer(3)]], + constant uint &chunk_size [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + F acc = zero_f(); + for (uint i = start; i < end; i++) { + acc = add_f(acc, load_f(input, offset + i)); + } + store_f(out, gid, acc); +} + +kernel void bn254_sumcheck_reduce_chunks( + device const ulong *input [[buffer(0)]], + device ulong *out [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &chunk_size [[buffer(3)]], + constant uint &out_len [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= out_len) return; + uint start = gid * chunk_size; + uint end = min(start + chunk_size, len); + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = start; i < end; i++) { + c0 = add_f(c0, load_f(input, i)); + c2 = add_f(c2, load_f(input, len + i)); + } + store_f(out, gid, c0); + store_f(out, out_len + gid, c2); +} + +kernel void bn254_sumcheck_chunks( + device const ulong *a [[buffer(0)]], + device const ulong *b [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &fold_half [[buffer(4)]], + constant uint &chunk_size [[buffer(5)]], + constant uint &partial_count [[buffer(6)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= fold_half) return; + uint end = min(start + chunk_size, fold_half); + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = start; i < end; i++) { + F a0 = i < len ? load_f(a, i) : zero_f(); + F b0 = i < len ? load_f(b, i) : zero_f(); + F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); + F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); + c0 = add_f(c0, mul_f(a0, b0)); + c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); + } + store_f(out, gid, c0); + store_f(out, partial_count + gid, c2); +} + +kernel void bn254_fold_pair_sumcheck_chunks( + device ulong *a [[buffer(0)]], + device ulong *b [[buffer(1)]], + device const ulong *weight_buf [[buffer(2)]], + device ulong *out [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &fold_half [[buffer(5)]], + constant uint &sum_half [[buffer(6)]], + constant uint &chunk_size [[buffer(7)]], + constant uint &partial_count [[buffer(8)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= sum_half) return; + uint end = min(start + chunk_size, sum_half); + F weight = load_f(weight_buf, 0); + F c0 = zero_f(); + F c2 = zero_f(); + for (uint i = start; i < end; i++) { + uint right = i + sum_half; + F a0 = fold_value_at(a, len, fold_half, i, weight); + F b0 = fold_value_at(b, len, fold_half, i, weight); + F a1 = right < fold_half ? fold_value_at(a, len, fold_half, right, weight) : zero_f(); + F b1 = right < fold_half ? fold_value_at(b, len, fold_half, right, weight) : zero_f(); + store_f(a, i, a0); + store_f(b, i, b0); + if (right < fold_half) { + store_f(a, right, a1); + store_f(b, right, b1); + } + c0 = add_f(c0, mul_f(a0, b0)); + c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); + } + store_f(out, gid, c0); + store_f(out, partial_count + gid, c2); +} + +kernel void bn254_geometric_accumulate( + device ulong *acc [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *scalars [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &num_points [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F value = load_f(acc, gid); + for (uint j = 0; j < num_points; j++) { + value = add_f(value, mul_f(load_f(scalars, j), pow_f(load_f(points, j), gid))); + } + store_f(acc, gid, value); +} + +kernel void bn254_geometric_accumulate_chunks( + device ulong *acc [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *scalars [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &num_points [[buffer(4)]], + constant uint &chunk_size [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + for (uint j = 0; j < num_points; j++) { + F point = load_f(points, j); + F scalar = load_f(scalars, j); + F power = pow_f(point, start); + for (uint i = start; i < end; i++) { + F value = load_f(acc, i); + value = add_f(value, mul_f(scalar, power)); + store_f(acc, i, value); + power = mul_f(power, point); + } + } +} + +kernel void bn254_geometric_accumulate_chunks_strided( + device ulong *acc [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *point_steps [[buffer(2)]], + device const ulong *scalars [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &num_points [[buffer(5)]], + constant uint &chunk_size [[buffer(6)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + uint count = end - start; + + // Accumulate all points into registers so `acc` is read and written + // once per element instead of once per (element, point). + F sums[32]; + for (uint k = 0; k < count; k++) { + sums[k] = zero_f(); + } + for (uint j = 0; j < num_points; j++) { + F point = load_f(points, j); + // Fold the scalar into the running power so the inner loop needs a + // single multiplication per element. + F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), gid)); + for (uint k = 0; k < count; k++) { + sums[k] = add_f(sums[k], power); + power = mul_f(power, point); + } + } + for (uint i = start, k = 0; i < end; i++, k++) { + store_f(acc, i, add_f(load_f(acc, i), sums[k])); + } +} + +kernel void bn254_geometric_accumulate_point_blocks( + device ulong *partials [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *point_steps [[buffer(2)]], + device const ulong *scalars [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &num_points [[buffer(5)]], + constant uint &chunk_size [[buffer(6)]], + constant uint &point_block_size [[buffer(7)]], + constant uint &point_blocks [[buffer(8)]], + uint gid [[thread_position_in_grid]] +) { + uint chunk = gid / point_blocks; + uint point_block = gid - chunk * point_blocks; + uint start = chunk * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + uint point_start = point_block * point_block_size; + if (point_start >= num_points) return; + uint point_end = min(point_start + point_block_size, num_points); + + F sums[32]; + for (uint k = 0; k < chunk_size; k++) { + sums[k] = zero_f(); + } + + for (uint j = point_start; j < point_end; j++) { + F point = load_f(points, j); + F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); + for (uint i = start, k = 0; i < end; i++, k++) { + sums[k] = add_f(sums[k], power); + power = mul_f(power, point); + } + } + + uint partial_offset = point_block * len; + for (uint i = start, k = 0; i < end; i++, k++) { + store_f(partials, partial_offset + i, sums[k]); + } +} + +kernel void bn254_geometric_accumulate_point_blocks_range( + device ulong *partials [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device const ulong *point_steps [[buffer(2)]], + device const ulong *scalars [[buffer(3)]], + constant uint &len [[buffer(4)]], + constant uint &num_points [[buffer(5)]], + constant uint &chunk_size [[buffer(6)]], + constant uint &point_block_size [[buffer(7)]], + constant uint &point_block_offset [[buffer(8)]], + constant uint &batch_point_blocks [[buffer(9)]], + uint gid [[thread_position_in_grid]] +) { + uint chunk = gid / batch_point_blocks; + uint local_point_block = gid - chunk * batch_point_blocks; + uint global_point_block = point_block_offset + local_point_block; + uint start = chunk * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + uint point_start = global_point_block * point_block_size; + if (point_start >= num_points) return; + uint point_end = min(point_start + point_block_size, num_points); + + F sums[32]; + for (uint k = 0; k < chunk_size; k++) { + sums[k] = zero_f(); + } + + for (uint j = point_start; j < point_end; j++) { + F point = load_f(points, j); + F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); + for (uint i = start, k = 0; i < end; i++, k++) { + sums[k] = add_f(sums[k], power); + power = mul_f(power, point); + } + } + + uint partial_offset = local_point_block * len; + for (uint i = start, k = 0; i < end; i++, k++) { + store_f(partials, partial_offset + i, sums[k]); + } +} + +kernel void bn254_geometric_accumulate_reduce_point_blocks( + device ulong *acc [[buffer(0)]], + device const ulong *partials [[buffer(1)]], + constant uint &len [[buffer(2)]], + constant uint &point_blocks [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= len) return; + F value = load_f(acc, gid); + for (uint block = 0; block < point_blocks; block++) { + value = add_f(value, load_f(partials, block * len + gid)); + } + store_f(acc, gid, value); +} + +kernel void bn254_univariate_evaluate( + device const ulong *coeffs [[buffer(0)]], + device const ulong *point_buf [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + if (len == 0) { + store_f(out, 0, zero_f()); + return; + } + F point = load_f(point_buf, 0); + F acc = load_f(coeffs, len - 1); + for (uint i = len - 1; i > 0; i--) { + acc = add_f(mul_f(acc, point), load_f(coeffs, i - 1)); + } + store_f(out, 0, acc); +} + +kernel void bn254_univariate_eval_chunks( + device const ulong *coeffs [[buffer(0)]], + device const ulong *point_buf [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &chunk_size [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint start = gid * chunk_size; + if (start >= len) return; + uint end = min(start + chunk_size, len); + F point = load_f(point_buf, 0); + F power = pow_f(point, start); + F acc = zero_f(); + for (uint i = start; i < end; i++) { + acc = add_f(acc, mul_f(load_f(coeffs, i), power)); + power = mul_f(power, point); + } + store_f(out, gid, acc); +} + +kernel void bn254_interleaved_rs_encode( + device const ulong *coeffs [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &num_messages [[buffer(3)]], + constant uint &coeff_len [[buffer(4)]], + constant uint &total [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= total) return; + uint row = gid / num_messages; + uint message = gid - row * num_messages; + uint base = message * coeff_len; + F point = load_f(points, row); + F acc = load_f(coeffs, base + coeff_len - 1); + for (uint i = coeff_len - 1; i > 0; i--) { + acc = add_f(mul_f(acc, point), load_f(coeffs, base + i - 1)); + } + store_f(out, gid, acc); +} + +kernel void bn254_interleaved_rs_encode_single_vector( + device const ulong *vector [[buffer(0)]], + device const ulong *points [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &interleaving_depth [[buffer(3)]], + constant uint &message_length [[buffer(4)]], + constant uint &total [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= total) return; + uint row = gid / interleaving_depth; + uint message = gid - row * interleaving_depth; + uint base = message * message_length; + F point = load_f(points, row); + F acc = load_f(vector, base + message_length - 1); + for (uint i = message_length - 1; i > 0; i--) { + acc = add_f(mul_f(acc, point), load_f(vector, base + i - 1)); + } + store_f(out, gid, acc); +} + +inline uint reverse_bits_width(uint value, uint width) { + return reverse_bits(value) >> (32u - width); +} + +kernel void bn254_pack_single_vector_cosets( + device const ulong *vector [[buffer(0)]], + device ulong *out [[buffer(1)]], + constant PackSingleVectorParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint row = gid / params.codeword_length; + uint col = gid - row * params.codeword_length; + if (col < params.message_length) { + store_f(out, gid, load_f(vector, row * params.message_length + col)); + } else { + store_f(out, gid, zero_f()); + } +} + +kernel void bn254_replicate_first_coset( + device ulong *buffer [[buffer(0)]], + constant ReplicateCosetsParams ¶ms [[buffer(1)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.trailing_elements) return; + uint repeats_per_row = params.row_len - params.coset_size; + uint row = gid / repeats_per_row; + uint within = gid - row * repeats_per_row; + uint dst = row * params.row_len + params.coset_size + within; + uint src = row * params.row_len + (within % params.coset_size); + store_f(buffer, dst, load_f(buffer, src)); +} + +kernel void bn254_bit_reverse_permute_rows_in_place( + device ulong *values [[buffer(0)]], + constant BitReverseParams &config [[buffer(1)]], + uint index [[thread_position_in_grid]] +) { + if (index >= config.total_elements || config.row_len <= 1u) return; + uint row = index / config.row_len; + uint within = index - row * config.row_len; + uint reversed = reverse_bits_width(within, config.log_n); + if (reversed <= within) return; + + uint row_base = row * config.row_len; + uint mate = row_base + reversed; + uint current = row_base + within; + F tmp = load_f(values, current); + store_f(values, current, load_f(values, mate)); + store_f(values, mate, tmp); +} + +kernel void bn254_radix2_ntt_stage_rows_in_place( + device ulong *values [[buffer(0)]], + device const ulong *twiddles [[buffer(1)]], + constant StageConfig &config [[buffer(2)]], + uint index [[thread_position_in_grid]] +) { + uint butterflies_per_row = config.row_len >> 1u; + uint row = index / butterflies_per_row; + uint local = index - row * butterflies_per_row; + uint half_m = config.half_m; + uint pair_in_group = local % half_m; + uint group = local / half_m; + uint row_base = row * config.row_len; + uint base = row_base + group * (half_m << 1u) + pair_in_group; + uint mate = base + half_m; + + F even = load_f(values, base); + F odd = load_f(values, mate); + F twiddle = load_f(twiddles, config.twiddle_offset + pair_in_group); + F t = mul_f(twiddle, odd); + + store_f(values, base, add_f(even, t)); + store_f(values, mate, sub_f(even, t)); +} + +kernel void bn254_transpose_matrix_reverse_rows( + device const ulong *input [[buffer(0)]], + device ulong *output [[buffer(1)]], + constant TransposeParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint row = gid / params.cols; + uint col = gid - row * params.cols; + uint row_bits = 31u - clz(params.cols); + uint dst_row = reverse_bits_width(col, row_bits); + uint dst = dst_row * params.rows + row; + store_f(output, dst, load_f(input, gid)); +} + +kernel void bn254_transpose_matrix( + device const ulong *input [[buffer(0)]], + device ulong *output [[buffer(1)]], + constant TransposeParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint row = gid / params.cols; + uint col = gid - row * params.cols; + uint dst = col * params.rows + row; + store_f(output, dst, load_f(input, gid)); +} + +kernel void bn254_apply_coset_twiddles( + device ulong *values [[buffer(0)]], + device const ulong *root_powers [[buffer(1)]], + constant ApplyCosetTwiddlesParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= params.total_elements) return; + uint within_codeword = gid % params.codeword_length; + uint coset = within_codeword / params.coset_size; + uint col = within_codeword - coset * params.coset_size; + if (coset == 0 || col == 0) return; + uint root_index = (coset * col) % params.codeword_length; + store_f(values, gid, mul_f(load_f(values, gid), load_f(root_powers, root_index))); +} + +kernel void bn254_encode_field_rows_le( + device const ulong *input [[buffer(0)]], + device uchar *output [[buffer(1)]], + constant FieldBytesParams ¶ms [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + uint total_elements = params.rows * params.cols; + if (gid >= total_elements) return; + + F canonical = from_mont(load_f(input, gid)); + uint byte_offset = gid * 32u; + for (uint limb = 0; limb < 4; ++limb) { + ulong value = canonical.v[limb]; + for (uint byte = 0; byte < 8; ++byte) { + output[byte_offset + limb * 8u + byte] = uchar((value >> (byte * 8u)) & 0xffUL); + } + } +} + +kernel void bn254_read_rows( + device const ulong *input [[buffer(0)]], + device const uint *indices [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &num_cols [[buffer(3)]], + constant uint &total [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid >= total) return; + uint row = gid / num_cols; + uint col = gid - row * num_cols; + uint src = indices[row] * num_cols + col; + store_f(out, gid, load_f(input, src)); +} + +kernel void bn254_multilinear_extend( + device const ulong *values [[buffer(0)]], + device const ulong *point [[buffer(1)]], + device ulong *out [[buffer(2)]], + constant uint &len [[buffer(3)]], + constant uint &num_vars [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid != 0) return; + F acc = zero_f(); + F one = one_f(); + for (uint i = 0; i < len; i++) { + F weight = one; + for (uint j = 0; j < num_vars; j++) { + F r = load_f(point, num_vars - 1 - j); + if (((i >> j) & 1) != 0) { + weight = mul_f(weight, r); + } else { + weight = mul_f(weight, sub_f(one, r)); + } + } + acc = add_f(acc, mul_f(load_f(values, i), weight)); + } + store_f(out, 0, acc); +} +"#; + +pub(crate) const REDUCTION_CHUNK_SIZE: usize = 64; +pub(crate) const SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 8; +pub(crate) const LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; +/// Must match the `sums` register array size in the strided kernel. +pub(crate) const MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; +pub(crate) const LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD: usize = 1 << 18; +pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE: usize = 16; +pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD: usize = 1 << 17; +pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS: usize = 64; +pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES: usize = 256 << 20; diff --git a/src/buffer/metal/mod.rs b/src/buffer/metal/mod.rs new file mode 100644 index 00000000..a796bf9e --- /dev/null +++ b/src/buffer/metal/mod.rs @@ -0,0 +1,16 @@ +// NOTE: 100% AI GENERATED + +mod buffer; +mod kernels; +mod profile; +mod rs; +mod runtime; +mod sha2; + +pub use buffer::{MetalBuffer, MetalSlice, MetalSliceMut}; +pub use profile::{ + reset_device_peak as metal_reset_device_peak, snapshot as metal_profile_snapshot, + MetalProfileSnapshot, +}; +pub use rs::MetalRs; +pub use sha2::MetalSha2; diff --git a/src/hash/metal_profile.rs b/src/buffer/metal/profile.rs similarity index 97% rename from src/hash/metal_profile.rs rename to src/buffer/metal/profile.rs index afdf4cb6..cebd654d 100644 --- a/src/hash/metal_profile.rs +++ b/src/buffer/metal/profile.rs @@ -1,3 +1,5 @@ +// NOTE: 100% AI GENERATED + use std::{ sync::atomic::{AtomicU64, Ordering}, time::Duration, @@ -112,7 +114,10 @@ pub fn record_device_allocated(bytes: u64) { /// Reset the peak gauge to the current value (e.g. at the start of a /// measured phase). pub fn reset_device_peak() { - DEVICE_PEAK_BYTES.store(DEVICE_CURRENT_BYTES.load(Ordering::Relaxed), Ordering::Relaxed); + DEVICE_PEAK_BYTES.store( + DEVICE_CURRENT_BYTES.load(Ordering::Relaxed), + Ordering::Relaxed, + ); } pub fn record_upload(bytes: u64, duration: Duration) { diff --git a/src/buffer/metal/rs.rs b/src/buffer/metal/rs.rs new file mode 100644 index 00000000..7a1dc4d5 --- /dev/null +++ b/src/buffer/metal/rs.rs @@ -0,0 +1,173 @@ +// NOTE: 100% AI GENERATED + +use std::any::type_name; + +use ark_ff::{FftField, Field}; + +use crate::{ + algebra::{ + fields::Field256, + ntt::{Messages, NttEngine, ReedSolomon}, + }, + buffer::{BufferOps, MetalBuffer}, +}; + +use super::runtime::{assert_bn254, encode_single_vector_coset_ntt}; + +const RS_PRIMES: [usize; 2] = [2, 3]; + +fn rs_divisors(n: usize) -> Vec { + let mut result = vec![1usize]; + let mut remaining = n; + for &p in &RS_PRIMES { + let mut pk = 1usize; + let existing = result.clone(); + while remaining.is_multiple_of(p) { + pk *= p; + remaining /= p; + result.extend(existing.iter().map(|d| d * pk)); + } + } + assert_eq!(remaining, 1); + result.sort_unstable(); + result +} + +#[cfg(not(feature = "rs_in_order"))] +pub(crate) fn rs_transpose_permute(index: usize, rows: usize, cols: usize) -> usize { + debug_assert!(index < rows * cols); + let (row, col) = (index / cols, index % cols); + row + col * rows +} +/// Metal (GPU) Reed-Solomon encoder. +/// +/// Keeps the scalar Reed-Solomon parameters directly and runs encoding on Metal buffers. +#[derive(Debug, Clone)] +pub struct MetalRs { + order: usize, + divisors: Vec, + omega_order: F, +} + +impl MetalRs { + pub fn new(order: usize, omega_order: F) -> Self { + assert_eq!(omega_order.pow([order as u64]), F::ONE); + for prime in RS_PRIMES { + if order.is_multiple_of(prime) { + assert_ne!(omega_order.pow([(order / prime) as u64]), F::ONE); + } + } + Self { + order, + divisors: rs_divisors(order), + omega_order, + } + } +} + +impl MetalRs { + pub fn new_from_fftfield() -> Self { + let (mut omega, mut order) = if let (Some(mut omega), Some(b), Some(k)) = ( + F::LARGE_SUBGROUP_ROOT_OF_UNITY, + F::SMALL_SUBGROUP_BASE, + F::SMALL_SUBGROUP_BASE_ADICITY, + ) { + let mut order = 1; + let mut remaining = (b as usize).checked_pow(k).expect("Small group too large."); + for p in RS_PRIMES { + while remaining.is_multiple_of(p) { + order *= p; + remaining /= p; + } + } + omega = omega.pow([remaining as u64]); + (omega, order) + } else { + (F::TWO_ADIC_ROOT_OF_UNITY, 1) + }; + let twos = F::TWO_ADICITY.min(order.leading_zeros()) as usize; + for _ in 0..(F::TWO_ADICITY as usize - twos) { + omega.square_in_place(); + } + order <<= twos; + Self::new(order, omega) + } +} + +impl ReedSolomon for MetalRs { + fn next_order(&self, size: usize) -> Option { + match self.divisors.binary_search(&size) { + Ok(index) | Err(index) => self.divisors.get(index).copied(), + } + } + + fn generator(&self, codeword_length: usize) -> F { + self.omega_order + .pow([(self.order / codeword_length) as u64]) + } + + fn evaluation_points( + &self, + masked_message_length: usize, + codeword_length: usize, + indices: &[usize], + ) -> Vec { + assert!(masked_message_length <= codeword_length); + assert!(self.order.is_multiple_of(codeword_length)); + let mut result = Vec::new(); + let generator = self.generator(codeword_length); + + let mut coset_size = self.next_order(masked_message_length).unwrap(); + while !codeword_length.is_multiple_of(coset_size) { + coset_size = self.next_order(coset_size + 1).unwrap(); + } + let num_cosets = codeword_length / coset_size; + #[cfg(feature = "rs_in_order")] + let _ = (coset_size, num_cosets); + + for &index in indices { + assert!(index < codeword_length); + + #[cfg(not(feature = "rs_in_order"))] + let index = rs_transpose_permute(index, num_cosets, coset_size); + result.push(generator.pow([index as u64])); + } + result + } + + fn interleaved_encode( + &self, + messages: Messages<'_, F>, + masks: &MetalBuffer, + codeword_length: usize, + ) -> MetalBuffer { + let vectors = messages.vectors; + let num_messages = vectors.len() * messages.interleaving_depth; + if num_messages == 0 { + return BufferOps::from_vec(Vec::new()); + } + assert!(masks.len().is_multiple_of(num_messages)); + let mask_length = masks.len() / num_messages; + + if type_name::() == type_name::() && vectors.len() == 1 && mask_length == 0 { + assert_bn254::(); + let mut coset_size = self.next_order(messages.message_length).unwrap(); + while !codeword_length.is_multiple_of(coset_size) { + coset_size = self.next_order(coset_size + 1).unwrap(); + } + return encode_single_vector_coset_ntt( + vectors[0], + messages.message_length, + messages.interleaving_depth, + codeword_length, + coset_size, + ); + } + + NttEngine::new(self.order, self.omega_order).interleaved_encode( + messages, + masks, + codeword_length, + ) + } +} diff --git a/src/buffer/metal/runtime.rs b/src/buffer/metal/runtime.rs new file mode 100644 index 00000000..0fc7265a --- /dev/null +++ b/src/buffer/metal/runtime.rs @@ -0,0 +1,1218 @@ +// NOTE: 100% AI GENERATED + +use std::{ + any::type_name, + cell::OnceCell, + collections::HashMap, + marker::PhantomData, + os::raw::c_void, + sync::{Mutex, OnceLock}, + time::Instant, +}; + +use ark_ff::{AdditiveGroup, BigInt, FftField, Field, Fp, MontBackend}; +use metal::{ + Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, MTLResourceOptions, MTLSize, +}; +use zerocopy::IntoBytes; + +use crate::{ + algebra::fields::{BN254Config, Field256}, + buffer::{BufferOps, MetalBuffer}, + hash::Hash as Digest, +}; + +use super::buffer::MetalFieldBuffer; +use super::kernels::{ + GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES, GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS, + GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE, GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD, + LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE, LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD, + MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE, METAL_SOURCE, REDUCTION_CHUNK_SIZE, + SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE, +}; +use super::profile; + +pub(crate) struct MetalRuntime { + device: Device, + queue: CommandQueue, + fold: ComputePipelineState, + fold_pair: ComputePipelineState, + scalar_mul_add: ComputePipelineState, + scalar_mul: ComputePipelineState, + dot: ComputePipelineState, + sumcheck: ComputePipelineState, + dot_chunks: ComputePipelineState, + sum_chunks: ComputePipelineState, + sumcheck_reduce_chunks: ComputePipelineState, + sumcheck_chunks: ComputePipelineState, + fold_pair_sumcheck_chunks: ComputePipelineState, + geometric_accumulate: ComputePipelineState, + geometric_accumulate_chunks: ComputePipelineState, + geometric_accumulate_chunks_strided: ComputePipelineState, + geometric_accumulate_point_blocks: ComputePipelineState, + geometric_accumulate_point_blocks_range: ComputePipelineState, + geometric_accumulate_reduce_point_blocks: ComputePipelineState, + univariate_evaluate: ComputePipelineState, + univariate_eval_chunks: ComputePipelineState, + interleaved_rs_encode: ComputePipelineState, + interleaved_rs_encode_single_vector: ComputePipelineState, + multilinear_extend: ComputePipelineState, + pack_single_vector_cosets: ComputePipelineState, + apply_coset_twiddles: ComputePipelineState, + replicate_first_coset: ComputePipelineState, + bit_reverse_rows: ComputePipelineState, + ntt_stage_rows: ComputePipelineState, + transpose: ComputePipelineState, + transpose_reverse_rows: ComputePipelineState, + encode_field_rows_le: ComputePipelineState, + read_rows: ComputePipelineState, + ntt_roots: Mutex>, + root_powers: Mutex>, +} + +pub fn init() { + let _ = runtime(); +} + +pub(crate) fn runtime() -> &'static MetalRuntime { + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + let device = Device::system_default().expect("Metal device is not available"); + let library = device + .new_library_with_source(METAL_SOURCE, &CompileOptions::new()) + .expect("failed to compile Metal BN254 kernels"); + let pipeline = |name: &str| { + let function = library + .get_function(name, None) + .unwrap_or_else(|_| panic!("missing Metal kernel {name}")); + device + .new_compute_pipeline_state_with_function(&function) + .unwrap_or_else(|err| panic!("failed to compile Metal kernel {name}: {err}")) + }; + MetalRuntime { + queue: device.new_command_queue(), + fold: pipeline("bn254_fold"), + fold_pair: pipeline("bn254_fold_pair"), + scalar_mul_add: pipeline("bn254_scalar_mul_add"), + scalar_mul: pipeline("bn254_scalar_mul"), + dot: pipeline("bn254_dot"), + sumcheck: pipeline("bn254_sumcheck"), + dot_chunks: pipeline("bn254_dot_chunks"), + sum_chunks: pipeline("bn254_sum_chunks"), + sumcheck_reduce_chunks: pipeline("bn254_sumcheck_reduce_chunks"), + sumcheck_chunks: pipeline("bn254_sumcheck_chunks"), + fold_pair_sumcheck_chunks: pipeline("bn254_fold_pair_sumcheck_chunks"), + geometric_accumulate: pipeline("bn254_geometric_accumulate"), + geometric_accumulate_chunks: pipeline("bn254_geometric_accumulate_chunks"), + geometric_accumulate_chunks_strided: pipeline( + "bn254_geometric_accumulate_chunks_strided", + ), + geometric_accumulate_point_blocks: pipeline("bn254_geometric_accumulate_point_blocks"), + geometric_accumulate_point_blocks_range: pipeline( + "bn254_geometric_accumulate_point_blocks_range", + ), + geometric_accumulate_reduce_point_blocks: pipeline( + "bn254_geometric_accumulate_reduce_point_blocks", + ), + univariate_evaluate: pipeline("bn254_univariate_evaluate"), + univariate_eval_chunks: pipeline("bn254_univariate_eval_chunks"), + interleaved_rs_encode: pipeline("bn254_interleaved_rs_encode"), + interleaved_rs_encode_single_vector: pipeline( + "bn254_interleaved_rs_encode_single_vector", + ), + multilinear_extend: pipeline("bn254_multilinear_extend"), + pack_single_vector_cosets: pipeline("bn254_pack_single_vector_cosets"), + apply_coset_twiddles: pipeline("bn254_apply_coset_twiddles"), + replicate_first_coset: pipeline("bn254_replicate_first_coset"), + bit_reverse_rows: pipeline("bn254_bit_reverse_permute_rows_in_place"), + ntt_stage_rows: pipeline("bn254_radix2_ntt_stage_rows_in_place"), + transpose: pipeline("bn254_transpose_matrix"), + transpose_reverse_rows: pipeline("bn254_transpose_matrix_reverse_rows"), + encode_field_rows_le: pipeline("bn254_encode_field_rows_le"), + read_rows: pipeline("bn254_read_rows"), + ntt_roots: Mutex::new(HashMap::new()), + root_powers: Mutex::new(HashMap::new()), + device, + } + }) +} + +pub(crate) fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { + match name { + "bn254_fold" => &rt.fold, + "bn254_fold_pair" => &rt.fold_pair, + "bn254_scalar_mul_add" => &rt.scalar_mul_add, + "bn254_scalar_mul" => &rt.scalar_mul, + "bn254_dot" => &rt.dot, + "bn254_sumcheck" => &rt.sumcheck, + "bn254_dot_chunks" => &rt.dot_chunks, + "bn254_sum_chunks" => &rt.sum_chunks, + "bn254_sumcheck_reduce_chunks" => &rt.sumcheck_reduce_chunks, + "bn254_sumcheck_chunks" => &rt.sumcheck_chunks, + "bn254_fold_pair_sumcheck_chunks" => &rt.fold_pair_sumcheck_chunks, + "bn254_geometric_accumulate" => &rt.geometric_accumulate, + "bn254_geometric_accumulate_chunks" => &rt.geometric_accumulate_chunks, + "bn254_geometric_accumulate_chunks_strided" => &rt.geometric_accumulate_chunks_strided, + "bn254_geometric_accumulate_point_blocks" => &rt.geometric_accumulate_point_blocks, + "bn254_geometric_accumulate_point_blocks_range" => { + &rt.geometric_accumulate_point_blocks_range + } + "bn254_geometric_accumulate_reduce_point_blocks" => { + &rt.geometric_accumulate_reduce_point_blocks + } + "bn254_univariate_evaluate" => &rt.univariate_evaluate, + "bn254_univariate_eval_chunks" => &rt.univariate_eval_chunks, + "bn254_interleaved_rs_encode" => &rt.interleaved_rs_encode, + "bn254_interleaved_rs_encode_single_vector" => &rt.interleaved_rs_encode_single_vector, + "bn254_multilinear_extend" => &rt.multilinear_extend, + "bn254_pack_single_vector_cosets" => &rt.pack_single_vector_cosets, + "bn254_apply_coset_twiddles" => &rt.apply_coset_twiddles, + "bn254_replicate_first_coset" => &rt.replicate_first_coset, + "bn254_bit_reverse_permute_rows_in_place" => &rt.bit_reverse_rows, + "bn254_radix2_ntt_stage_rows_in_place" => &rt.ntt_stage_rows, + "bn254_transpose_matrix" => &rt.transpose, + "bn254_transpose_matrix_reverse_rows" => &rt.transpose_reverse_rows, + "bn254_encode_field_rows_le" => &rt.encode_field_rows_le, + "bn254_read_rows" => &rt.read_rows, + _ => panic!("unknown Metal kernel {name}"), + } +} + +pub(crate) fn new_shared_buffer(rt: &MetalRuntime, bytes: u64) -> Buffer { + profile::record_alloc(bytes); + let buffer = rt + .device + .new_buffer(bytes, MTLResourceOptions::StorageModeShared); + profile::record_device_allocated(rt.device.current_allocated_size()); + buffer +} + +pub(crate) fn new_shared_buffer_with_data( + rt: &MetalRuntime, + data: *const c_void, + bytes: u64, +) -> Buffer { + let start = Instant::now(); + let buffer = rt + .device + .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); + profile::record_alloc(bytes); + profile::record_upload(bytes, start.elapsed()); + profile::record_device_allocated(rt.device.current_allocated_size()); + buffer +} + +pub(crate) fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { + command.commit(); + let start = Instant::now(); + command.wait_until_completed(); + let elapsed = start.elapsed(); + if std::env::var_os("WHIR_METAL_TRACE").is_some() { + eprintln!( + "metal command {label} {:.3} ms", + elapsed.as_secs_f64() * 1_000.0 + ); + } + profile::record_command_wait(elapsed); +} + +pub(crate) fn wait_for_blit(command: &metal::CommandBufferRef, bytes: u64) { + command.commit(); + let start = Instant::now(); + command.wait_until_completed(); + profile::record_blit(bytes, start.elapsed()); +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct BitReverseParams { + row_len: u32, + log_n: u32, + total_elements: u32, + _pad0: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct NttStageParams { + row_len: u32, + half_m: u32, + twiddle_offset: u32, + _pad0: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct TransposeParams { + rows: u32, + cols: u32, + total_elements: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct ReplicateCosetsParams { + row_len: u32, + coset_size: u32, + trailing_elements: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct PackSingleVectorParams { + row_count: u32, + message_length: u32, + codeword_length: u32, + coset_size: u32, + total_elements: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct FieldBytesParams { + rows: u32, + cols: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct ApplyCosetTwiddlesParams { + row_count: u32, + num_cosets: u32, + coset_size: u32, + codeword_length: u32, + total_elements: u32, +} + +pub(crate) fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field256 { + parallel_dot_at(a, 0, b, 0, len) +} + +/// Inner product of `a[a_off..a_off+len]` and `b[b_off..b_off+len]`. +pub(crate) fn parallel_dot_at( + a: &MetalFieldBuffer, + a_off: usize, + b: &MetalFieldBuffer, + b_off: usize, + len: usize, +) -> Field256 { + if len == 0 { + return Field256::ZERO; + } + let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), + }; + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_dot_chunks"), + &[&a.limbs, &b.limbs, &partials.limbs], + &[field_byte_offset(a_off), field_byte_offset(b_off), 0], + &[len as u32, REDUCTION_CHUNK_SIZE as u32], + partial_count, + ); + let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); + wait_for_command_named(command, "bn254_dot_chunks"); + download_field_at(&result.limbs, offset) +} + +pub(crate) fn parallel_sumcheck( + a: &MetalFieldBuffer, + b: &MetalFieldBuffer, + len: usize, + fold_half: usize, +) -> (Field256, Field256) { + parallel_sumcheck_at(a, 0, b, 0, len, fold_half) +} + +/// Sumcheck `(c0, c2)` over `a[a_off..]` and `b[b_off..]` (logical length `len`). +pub(crate) fn parallel_sumcheck_at( + a: &MetalFieldBuffer, + a_off: usize, + b: &MetalFieldBuffer, + b_off: usize, + len: usize, + fold_half: usize, +) -> (Field256, Field256) { + let partial_count = fold_half.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), + }; + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_sumcheck_chunks"), + &[&a.limbs, &b.limbs, &partials.limbs], + &[field_byte_offset(a_off), field_byte_offset(b_off), 0], + &[ + len as u32, + fold_half as u32, + REDUCTION_CHUNK_SIZE as u32, + partial_count as u32, + ], + partial_count, + ); + let result = encode_sumcheck_reduction(command, partials, partial_count); + wait_for_command_named(command, "bn254_sumcheck_chunks"); + download_sumcheck_pair(&result) +} + +pub(crate) fn parallel_fold_pair_sumcheck( + a: &MetalFieldBuffer, + b: &MetalFieldBuffer, + weight: &MetalFieldBuffer, + len: usize, + fold_half: usize, +) -> (Field256, Field256) { + let sum_half = fold_half.next_power_of_two() >> 1; + debug_assert!(sum_half > 0); + let partial_count = sum_half.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), + }; + let command = rt.queue.new_command_buffer(); + encode_u32_kernel( + command, + pipeline(rt, "bn254_fold_pair_sumcheck_chunks"), + &[&a.limbs, &b.limbs, &weight.limbs, &partials.limbs], + &[ + len as u32, + fold_half as u32, + sum_half as u32, + REDUCTION_CHUNK_SIZE as u32, + partial_count as u32, + ], + partial_count, + ); + let result = encode_sumcheck_reduction(command, partials, partial_count); + wait_for_command_named(command, "bn254_fold_pair_sumcheck_chunks"); + download_sumcheck_pair(&result) +} + +pub(crate) fn parallel_univariate_evaluate( + coeffs: &MetalFieldBuffer, + point: &MetalFieldBuffer, + len: usize, +) -> Field256 { + parallel_univariate_evaluate_at(coeffs, 0, point, len) +} + +/// Univariate evaluation of `coeffs[coeff_off..coeff_off+len]` at `point`. +pub(crate) fn parallel_univariate_evaluate_at( + coeffs: &MetalFieldBuffer, + coeff_off: usize, + point: &MetalFieldBuffer, + len: usize, +) -> Field256 { + if len == 0 { + return Field256::ZERO; + } + let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), + }; + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_univariate_eval_chunks"), + &[&coeffs.limbs, &point.limbs, &partials.limbs], + &[field_byte_offset(coeff_off), 0, 0], + &[len as u32, REDUCTION_CHUNK_SIZE as u32], + partial_count, + ); + let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); + wait_for_command_named(command, "bn254_univariate_eval_chunks"); + download_field_at(&result.limbs, offset) +} + +pub(crate) fn parallel_geometric_accumulate_point_blocks( + acc: &MetalFieldBuffer, + points: &MetalFieldBuffer, + point_steps: &MetalFieldBuffer, + scalars: &MetalFieldBuffer, + len: usize, + num_points: usize, + chunk_size: usize, +) { + assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); + let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); + let chunks = len.div_ceil(chunk_size); + let partial_len = point_blocks + .checked_mul(len) + .expect("Metal geometric partial size overflow"); + let rt = runtime(); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), + }; + run_in_place( + "bn254_geometric_accumulate_point_blocks", + &[ + &partials.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[ + len as u32, + num_points as u32, + chunk_size as u32, + GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, + point_blocks as u32, + ], + chunks * point_blocks, + ); + run_in_place( + "bn254_geometric_accumulate_reduce_point_blocks", + &[&acc.limbs, &partials.limbs], + &[len as u32, point_blocks as u32], + len, + ); +} + +pub(crate) fn parallel_geometric_accumulate_point_blocks_batched( + acc: &MetalFieldBuffer, + points: &MetalFieldBuffer, + point_steps: &MetalFieldBuffer, + scalars: &MetalFieldBuffer, + len: usize, + num_points: usize, + chunk_size: usize, +) { + assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); + let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); + let chunks = len.div_ceil(chunk_size); + let bytes_per_point_block = len + .checked_mul(size_of::()) + .expect("Metal geometric partial size overflow"); + let default_batch_blocks = + (GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES / bytes_per_point_block).max(1); + let batch_blocks = std::env::var("WHIR_METAL_GEOM_POINT_BLOCK_BATCH") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&value| value > 0) + .unwrap_or(default_batch_blocks) + .min(point_blocks); + let rt = runtime(); + for point_block_offset in (0..point_blocks).step_by(batch_blocks) { + let current_batch = batch_blocks.min(point_blocks - point_block_offset); + let partial_len = current_batch + .checked_mul(len) + .expect("Metal geometric partial size overflow"); + let partials = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), + }; + run_in_place( + "bn254_geometric_accumulate_point_blocks_range", + &[ + &partials.limbs, + &points.limbs, + &point_steps.limbs, + &scalars.limbs, + ], + &[ + len as u32, + num_points as u32, + chunk_size as u32, + GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, + point_block_offset as u32, + current_batch as u32, + ], + chunks * current_batch, + ); + run_in_place( + "bn254_geometric_accumulate_reduce_point_blocks", + &[&acc.limbs, &partials.limbs], + &[len as u32, current_batch as u32], + len, + ); + } +} + +pub(crate) fn geometric_accumulate_chunk_size(len: usize) -> usize { + std::env::var("WHIR_METAL_GEOM_CHUNK") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&value| value > 0) + .unwrap_or(if len >= LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD { + LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + } else { + SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + }) + .min(MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE) +} + +pub(crate) fn should_use_geometric_point_blocks( + len: usize, + num_points: usize, + chunk_size: usize, +) -> bool { + chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS + && len <= GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD +} + +pub(crate) fn should_use_geometric_point_blocks_batched( + len: usize, + num_points: usize, + chunk_size: usize, +) -> bool { + chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE + && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS + && len > GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD +} + +/// Encodes all tree-reduction levels into `command` and returns the buffer +/// and offset holding the final scalar. Runs with a single command wait. +pub(crate) fn encode_field_reduction( + command: &metal::CommandBufferRef, + input: MetalFieldBuffer, + len: usize, + offset: usize, +) -> (MetalFieldBuffer, usize) { + let rt = runtime(); + let mut current = input; + let mut current_len = len; + let mut current_offset = offset; + while current_len > 1 { + let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); + let next = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (next_len * size_of::()) as u64), + }; + encode_u32_kernel( + command, + pipeline(rt, "bn254_sum_chunks"), + &[¤t.limbs, &next.limbs], + &[ + current_len as u32, + current_offset as u32, + REDUCTION_CHUNK_SIZE as u32, + ], + next_len, + ); + current = next; + current_len = next_len; + current_offset = 0; + } + (current, current_offset) +} + +/// Encodes all (c0, c2) tree-reduction levels into `command` and returns the +/// buffer holding the final pair. Runs with a single command wait. +pub(crate) fn encode_sumcheck_reduction( + command: &metal::CommandBufferRef, + input: MetalFieldBuffer, + len: usize, +) -> MetalFieldBuffer { + let rt = runtime(); + let mut current = input; + let mut current_len = len; + while current_len > 1 { + let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); + let next = MetalFieldBuffer { + limbs: new_shared_buffer(rt, (next_len * 2 * size_of::()) as u64), + }; + encode_u32_kernel( + command, + pipeline(rt, "bn254_sumcheck_reduce_chunks"), + &[¤t.limbs, &next.limbs], + &[ + current_len as u32, + REDUCTION_CHUNK_SIZE as u32, + next_len as u32, + ], + next_len, + ); + current = next; + current_len = next_len; + } + current +} + +pub(crate) fn download_sumcheck_pair(buffer: &MetalFieldBuffer) -> (Field256, Field256) { + let values = download_field(&buffer.limbs, 2); + (values[0], values[1]) +} + +pub(crate) fn encode_single_vector_coset_ntt( + vector: &MetalBuffer, + message_length: usize, + interleaving_depth: usize, + codeword_length: usize, + coset_size: usize, +) -> MetalBuffer { + assert!(codeword_length.is_power_of_two()); + assert!(Field256::get_root_of_unity(codeword_length as u64).is_some()); + assert_eq!(vector.len(), message_length * interleaving_depth); + + assert!(codeword_length.is_multiple_of(coset_size)); + let num_cosets = codeword_length / coset_size; + let total_elements = interleaving_depth + .checked_mul(codeword_length) + .expect("Metal RS encode size overflow"); + assert!(total_elements <= u32::MAX as usize); + + let rt = runtime(); + let source = vector.bn254_buffer(); + let current = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); + let transposed = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); + let codeword_root_powers = root_powers_buffer(codeword_length); + let coset_roots = roots_buffer(coset_size); + + let command = rt.queue.new_command_buffer(); + + encode_kernel( + &command, + &rt.pack_single_vector_cosets, + &[&source.limbs, ¤t], + &PackSingleVectorParams { + row_count: interleaving_depth as u32, + message_length: message_length as u32, + codeword_length: codeword_length as u32, + coset_size: coset_size as u32, + total_elements: total_elements as u32, + }, + total_elements, + ); + + let trailing_elements = interleaving_depth.saturating_mul(codeword_length - coset_size); + if trailing_elements != 0 { + encode_kernel( + &command, + &rt.replicate_first_coset, + &[¤t], + &ReplicateCosetsParams { + row_len: codeword_length as u32, + coset_size: coset_size as u32, + trailing_elements: trailing_elements as u32, + }, + trailing_elements, + ); + } + + encode_kernel( + &command, + &rt.apply_coset_twiddles, + &[¤t, &codeword_root_powers.limbs], + &ApplyCosetTwiddlesParams { + row_count: interleaving_depth as u32, + num_cosets: num_cosets as u32, + coset_size: coset_size as u32, + codeword_length: codeword_length as u32, + total_elements: total_elements as u32, + }, + total_elements, + ); + + let stage_count = coset_size.trailing_zeros() as usize; + encode_kernel( + &command, + &rt.bit_reverse_rows, + &[¤t], + &BitReverseParams { + row_len: coset_size as u32, + log_n: stage_count as u32, + total_elements: total_elements as u32, + _pad0: 0, + }, + total_elements, + ); + + let total_butterflies = total_elements / 2; + let mut twiddle_offset = 0usize; + for stage in 0..stage_count { + let half_m = 1usize << stage; + encode_kernel( + &command, + &rt.ntt_stage_rows, + &[¤t, &coset_roots.limbs], + &NttStageParams { + row_len: coset_size as u32, + half_m: half_m as u32, + twiddle_offset: twiddle_offset as u32, + _pad0: 0, + }, + total_butterflies, + ); + twiddle_offset += 1usize << stage; + } + + encode_kernel( + &command, + &rt.transpose, + &[¤t, &transposed], + &TransposeParams { + rows: interleaving_depth as u32, + cols: codeword_length as u32, + total_elements: total_elements as u32, + }, + total_elements, + ); + + wait_for_command_named(&command, "bn254_rs_encode"); + + MetalBuffer::from_field_limb_buffer(transposed, total_elements) +} + +pub(crate) fn encode_kernel( + command: &metal::CommandBufferRef, + pipeline: &ComputePipelineState, + buffers: &[&Buffer], + params: &P, + threads: usize, +) { + let encoder = command.new_compute_command_encoder(); + encoder.set_compute_pipeline_state(pipeline); + for (index, buffer) in buffers.iter().enumerate() { + encoder.set_buffer(index as u64, Some(buffer), 0); + } + encoder.set_bytes( + buffers.len() as u64, + size_of::

() as u64, + (params as *const P).cast::(), + ); + dispatch(&encoder, pipeline, threads.max(1)); + encoder.end_encoding(); +} + +pub(crate) fn roots_buffer(codeword_length: usize) -> MetalFieldBuffer { + let rt = runtime(); + if let Some(buffer) = rt.ntt_roots.lock().unwrap().get(&codeword_length).cloned() { + return buffer; + } + let root = Field256::get_root_of_unity(codeword_length as u64) + .expect("BN254 root of unity unavailable for Metal NTT"); + let stage_count = codeword_length.trailing_zeros() as usize; + let mut roots = Vec::with_capacity(codeword_length.saturating_sub(1)); + for stage in 0..stage_count { + let stage_size = 1usize << (stage + 1); + let half_stage = stage_size >> 1; + let stage_root = root.pow([(codeword_length / stage_size) as u64]); + let mut current = Field256::ONE; + for _ in 0..half_stage { + roots.push(current); + current *= stage_root; + } + } + let buffer = upload_field(&roots); + rt.ntt_roots + .lock() + .unwrap() + .insert(codeword_length, buffer.clone()); + buffer +} + +pub(crate) fn root_powers_buffer(codeword_length: usize) -> MetalFieldBuffer { + let rt = runtime(); + if let Some(buffer) = rt + .root_powers + .lock() + .unwrap() + .get(&codeword_length) + .cloned() + { + return buffer; + } + let root = Field256::get_root_of_unity(codeword_length as u64) + .expect("BN254 root of unity unavailable for Metal RS twiddles"); + let mut powers = Vec::with_capacity(codeword_length); + let mut current = Field256::ONE; + for _ in 0..codeword_length { + powers.push(current); + current *= root; + } + let buffer = upload_field(&powers); + rt.root_powers + .lock() + .unwrap() + .insert(codeword_length, buffer.clone()); + buffer +} + +pub(crate) fn encode_field_rows_le(input: &Buffer, rows: usize, cols: usize) -> Buffer { + assert!(rows <= u32::MAX as usize); + assert!(cols <= u32::MAX as usize); + let rt = runtime(); + let output = new_shared_buffer(rt, (rows * cols * size_of::()) as u64); + let command = rt.queue.new_command_buffer(); + encode_kernel( + &command, + &rt.encode_field_rows_le, + &[input, &output], + &FieldBytesParams { + rows: rows as u32, + cols: cols as u32, + }, + rows * cols, + ); + wait_for_command_named(&command, "bn254_encode_field_rows_le"); + output +} + +pub(crate) fn read_bn254_rows( + source: &MetalFieldBuffer, + num_cols: usize, + indices: &[usize], +) -> Vec { + if indices.is_empty() || num_cols == 0 { + return Vec::new(); + } + assert!(num_cols <= u32::MAX as usize); + assert!(indices.iter().all(|&index| index <= u32::MAX as usize)); + let total = indices + .len() + .checked_mul(num_cols) + .expect("Metal read_rows size overflow"); + assert!(total <= u32::MAX as usize); + + let rt = runtime(); + let indices = indices + .iter() + .copied() + .map(|index| index as u32) + .collect::>(); + let indices_buffer = new_shared_buffer_with_data( + rt, + indices.as_ptr().cast(), + (indices.len() * size_of::()) as u64, + ); + let out = new_shared_buffer(rt, (total * size_of::()) as u64); + run_in_place( + "bn254_read_rows", + &[&source.limbs, &indices_buffer, &out], + &[num_cols as u32, total as u32], + total, + ); + download_field(&out, total) +} + +pub(crate) fn upload_field(values: &[Field256]) -> MetalFieldBuffer { + let rt = runtime(); + let buffer = if values.is_empty() { + new_shared_buffer(rt, 0) + } else { + // Field256 is 4 contiguous u64 limbs; upload directly without + // flattening into an intermediate Vec. + new_shared_buffer_with_data( + rt, + values.as_ptr().cast(), + std::mem::size_of_val(values) as u64, + ) + }; + MetalFieldBuffer { limbs: buffer } +} + +pub(crate) fn zeroed_field_buffer(len: usize) -> MetalFieldBuffer { + let rt = runtime(); + let bytes = (len * size_of::()) as u64; + let buffer = new_shared_buffer(rt, bytes); + if bytes > 0 { + let command = rt.queue.new_command_buffer(); + let blit = command.new_blit_command_encoder(); + blit.fill_buffer(&buffer, metal::NSRange::new(0, bytes), 0); + blit.end_encoding(); + wait_for_blit(command, bytes); + } + MetalFieldBuffer { limbs: buffer } +} + +pub(crate) fn maybe_upload_bn254(values: &[T]) -> Option { + (type_name::() == type_name::()).then(|| upload_field(as_field256_slice(values))) +} + +/// Geometric accumulate into `field[offset .. offset+len]` by binding the +/// field buffer at a byte offset; the kernel itself indexes from `gid == 0`. +pub(crate) fn geometric_accumulate_at_offset( + field: &MetalFieldBuffer, + offset: usize, + len: usize, + points: &MetalFieldBuffer, + scalars: &MetalFieldBuffer, + num_points: usize, +) { + if len == 0 || num_points == 0 { + return; + } + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + let encoder = command.new_compute_command_encoder(); + let pipe = pipeline(rt, "bn254_geometric_accumulate"); + encoder.set_compute_pipeline_state(pipe); + let byte_offset = (offset * size_of::()) as u64; + encoder.set_buffer(0, Some(&field.limbs), byte_offset); + encoder.set_buffer(1, Some(&points.limbs), 0); + encoder.set_buffer(2, Some(&scalars.limbs), 0); + let len_u32 = len as u32; + let num_u32 = num_points as u32; + encoder.set_bytes(3, size_of::() as u64, (&len_u32 as *const u32).cast()); + encoder.set_bytes(4, size_of::() as u64, (&num_u32 as *const u32).cast()); + dispatch(&encoder, pipe, len); + encoder.end_encoding(); + wait_for_command_named(command, "bn254_geometric_accumulate_window"); +} + +/// In-place scalar multiply of `field[offset .. offset+len]` by binding the +/// field buffer at a byte offset; the kernel indexes from `gid == 0`. +pub(crate) fn scalar_mul_at_offset( + field: &MetalFieldBuffer, + offset: usize, + len: usize, + weight: &MetalFieldBuffer, +) { + if len == 0 { + return; + } + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + let encoder = command.new_compute_command_encoder(); + let pipe = pipeline(rt, "bn254_scalar_mul"); + encoder.set_compute_pipeline_state(pipe); + let byte_offset = (offset * size_of::()) as u64; + encoder.set_buffer(0, Some(&field.limbs), byte_offset); + encoder.set_buffer(1, Some(&weight.limbs), 0); + let len_u32 = len as u32; + encoder.set_bytes(2, size_of::() as u64, (&len_u32 as *const u32).cast()); + dispatch(&encoder, pipe, len); + encoder.end_encoding(); + wait_for_command_named(command, "bn254_scalar_mul_window"); +} + +/// Multilinear extension of `field[off..off+len]` evaluated at `point`. +pub(crate) fn parallel_multilinear_extend_at( + field: &MetalFieldBuffer, + off: usize, + len: usize, + point: &MetalFieldBuffer, + num_vars: usize, +) -> Field256 { + let rt = runtime(); + let out = new_shared_buffer(rt, (4 * size_of::()) as u64); + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_multilinear_extend"), + &[&field.limbs, &point.limbs, &out], + &[field_byte_offset(off), 0, 0], + &[len as u32, num_vars as u32], + 1, + ); + wait_for_command_named(command, "bn254_multilinear_extend"); + download_field(&out, 1)[0] +} + +/// `acc[acc_off..] += weight * vector[vec_off..]` over `len` elements. +pub(crate) fn scalar_mul_add_at( + acc: &MetalFieldBuffer, + acc_off: usize, + vector: &MetalFieldBuffer, + vec_off: usize, + weight: &MetalFieldBuffer, + len: usize, +) { + if len == 0 { + return; + } + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + encode_u32_kernel_with_offsets( + command, + pipeline(rt, "bn254_scalar_mul_add"), + &[&acc.limbs, &vector.limbs, &weight.limbs], + &[field_byte_offset(acc_off), field_byte_offset(vec_off), 0], + &[len as u32], + len, + ); + wait_for_command_named(command, "bn254_scalar_mul_add"); +} + +pub(crate) fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer { + copy_field_buffer_at(source, 0, len) +} + +pub(crate) fn copy_field_buffer_at( + source: &MetalFieldBuffer, + offset: usize, + len: usize, +) -> MetalFieldBuffer { + let rt = runtime(); + let byte_len = (len * 4 * size_of::()) as u64; + let source_offset = (offset * 4 * size_of::()) as u64; + let target = new_shared_buffer(rt, byte_len); + let command = rt.queue.new_command_buffer(); + let blit = command.new_blit_command_encoder(); + blit.copy_from_buffer(&source.limbs, source_offset, &target, 0, byte_len); + blit.end_encoding(); + wait_for_blit(&command, byte_len); + MetalFieldBuffer { limbs: target } +} + +/// Encodes a kernel dispatch with `u32` constants bound as inline bytes +/// (no per-constant buffer allocations). +pub(crate) fn encode_u32_kernel( + command: &metal::CommandBufferRef, + pipeline: &ComputePipelineState, + buffers: &[&Buffer], + constants: &[u32], + threads: usize, +) { + encode_u32_kernel_with_offsets(command, pipeline, buffers, &[], constants, threads); +} + +/// Like [`encode_u32_kernel`], but binds `buffers[i]` at byte offset +/// `offsets[i]` (defaulting to 0 when `offsets` is shorter). This is the +/// mechanism that lets a view dispatch a kernel over `parent[offset..]` +/// without copying: only the input binding shifts, the kernel still indexes +/// from `gid == 0`. +pub(crate) fn encode_u32_kernel_with_offsets( + command: &metal::CommandBufferRef, + pipeline: &ComputePipelineState, + buffers: &[&Buffer], + offsets: &[u64], + constants: &[u32], + threads: usize, +) { + let encoder = command.new_compute_command_encoder(); + encoder.set_compute_pipeline_state(pipeline); + let mut index = 0; + for (i, buffer) in buffers.iter().enumerate() { + let byte_offset = offsets.get(i).copied().unwrap_or(0); + encoder.set_buffer(index, Some(buffer), byte_offset); + index += 1; + } + for constant in constants { + encoder.set_bytes( + index, + size_of::() as u64, + (constant as *const u32).cast(), + ); + index += 1; + } + dispatch(&encoder, pipeline, threads.max(1)); + encoder.end_encoding(); +} + +/// Byte offset of element `offset` within a `Field256` buffer. +#[inline] +pub(crate) fn field_byte_offset(offset: usize) -> u64 { + (offset * size_of::()) as u64 +} + +pub(crate) fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { + let rt = runtime(); + let command = rt.queue.new_command_buffer(); + encode_u32_kernel(command, pipeline(rt, name), buffers, constants, threads); + wait_for_command_named(command, name); +} + +pub(crate) fn dispatch( + encoder: &metal::ComputeCommandEncoderRef, + pipeline: &ComputePipelineState, + threads: usize, +) { + // Use full threadgroups (capped at 256) instead of a single execution + // width; the pipeline limit already accounts for register pressure. + let width = pipeline.max_total_threads_per_threadgroup().clamp(1, 256); + let group_width = width.min(threads as u64).max(1); + encoder.dispatch_threads( + MTLSize { + width: threads as u64, + height: 1, + depth: 1, + }, + MTLSize { + width: group_width, + height: 1, + depth: 1, + }, + ); +} + +pub(crate) fn download_field(buffer: &Buffer, len: usize) -> Vec { + if len == 0 { + return Vec::new(); + } + let start = Instant::now(); + let limbs = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 4) }; + let result = limbs + .chunks_exact(4) + .map(|chunk| { + Fp::, 4>( + BigInt([chunk[0], chunk[1], chunk[2], chunk[3]]), + PhantomData, + ) + }) + .collect(); + profile::record_readback((len * size_of::()) as u64, start.elapsed()); + result +} + +pub(crate) fn download_field_at(buffer: &Buffer, index: usize) -> Field256 { + let start = Instant::now(); + let limbs = + unsafe { std::slice::from_raw_parts(buffer.contents().cast::().add(index * 4), 4) }; + let result = Fp::, 4>( + BigInt([limbs[0], limbs[1], limbs[2], limbs[3]]), + PhantomData, + ); + profile::record_readback(size_of::() as u64, start.elapsed()); + result +} + +pub(crate) fn download_hash_indices(buffer: &Buffer, len: usize, indices: &[usize]) -> Vec { + let start = Instant::now(); + let bytes = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 32) }; + let mut result = Vec::with_capacity(indices.len()); + for &index in indices { + assert!(index < len, "Metal hash index out of bounds"); + let mut hash = Digest::default(); + hash.as_mut_bytes() + .copy_from_slice(&bytes[index * 32..(index + 1) * 32]); + result.push(hash); + } + profile::record_readback( + (indices.len() * size_of::()) as u64, + start.elapsed(), + ); + result +} + +pub(crate) fn assert_bn254() { + assert_eq!( + type_name::(), + type_name::(), + "MetalBuffer only supports BN254 Field256 field operations" + ); +} + +pub(crate) fn f_to_field256(value: F) -> Field256 { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +pub(crate) fn field256_to_f(value: Field256) -> F { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +pub(crate) fn target_to_field256(value: T) -> Field256 { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +pub(crate) fn field256_to_target(value: Field256) -> T { + assert_bn254::(); + unsafe { std::mem::transmute_copy(&value) } +} + +pub(crate) fn as_field256_slice(values: &[T]) -> &[Field256] { + assert_eq!( + type_name::(), + type_name::(), + "MetalBuffer only supports BN254 Field256 buffers" + ); + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } +} diff --git a/src/hash/metal_sha2_engine.rs b/src/buffer/metal/sha2.rs similarity index 97% rename from src/hash/metal_sha2_engine.rs rename to src/buffer/metal/sha2.rs index 1ad6cb65..f2862738 100644 --- a/src/hash/metal_sha2_engine.rs +++ b/src/buffer/metal/sha2.rs @@ -1,3 +1,5 @@ +// NOTE: 100% AI GENERATED + use std::{borrow::Cow, sync::OnceLock, time::Instant}; use const_oid::ObjectIdentifier; @@ -8,7 +10,8 @@ use metal::{ use sha2::Sha256; use zerocopy::IntoBytes; -use super::{metal_profile, Hash, HashEngine, HASH_COUNTER}; +use super::profile; +use crate::hash::{Hash, HashEngine, HASH_COUNTER}; const USE_SHA256_64_SPECIALIZATION: bool = false; const POW_MAX_WINDOW: u32 = 1 << 20; @@ -301,11 +304,11 @@ struct MetalSha2Runtime { } fn new_shared_buffer(rt: &MetalSha2Runtime, bytes: u64) -> Buffer { - metal_profile::record_alloc(bytes); + profile::record_alloc(bytes); let buffer = rt .device .new_buffer(bytes, MTLResourceOptions::StorageModeShared); - metal_profile::record_device_allocated(rt.device.current_allocated_size()); + profile::record_device_allocated(rt.device.current_allocated_size()); buffer } @@ -318,9 +321,9 @@ fn new_shared_buffer_with_data( let buffer = rt .device .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); - metal_profile::record_alloc(bytes); - metal_profile::record_upload(bytes, start.elapsed()); - metal_profile::record_device_allocated(rt.device.current_allocated_size()); + profile::record_alloc(bytes); + profile::record_upload(bytes, start.elapsed()); + profile::record_device_allocated(rt.device.current_allocated_size()); buffer } @@ -347,7 +350,7 @@ fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { elapsed.as_secs_f64() * 1_000.0 ); } - metal_profile::record_command_wait(elapsed); + profile::record_command_wait(elapsed); } fn runtime() -> &'static MetalSha2Runtime { @@ -440,7 +443,7 @@ impl MetalSha2 { ) }; let nonce = candidates.iter().copied().min().unwrap_or(u64::MAX); - metal_profile::record_readback(candidate_bytes, read_start.elapsed()); + profile::record_readback(candidate_bytes, read_start.elapsed()); if nonce != u64::MAX { HASH_COUNTER.add((nonce - start_nonce + 1) as usize); return nonce; @@ -611,7 +614,7 @@ impl MetalSha2 { for (out, bytes) in output.iter_mut().zip(bytes.chunks_exact(32)) { out.as_mut_bytes().copy_from_slice(bytes); } - metal_profile::record_readback(output_bytes, start.elapsed()); + profile::record_readback(output_bytes, start.elapsed()); } } diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index 32b87b8f..daaad4ab 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -1,4 +1,20 @@ +//! Backend-agnostic buffers for protocol data. +//! +//! Protocol code uses [`ActiveBuffer`], [`ActiveSlice`], and [`ActiveSliceMut`] +//! to select the CPU or GPU backend at compile time. Owned buffers manage +//! storage; slices are zero-copy views into a contiguous range. +//! +//! The trait split follows the ownership model. [`BufferOps`] is generic over +//! any element type and is used for field elements, [`struct@Hash`] digests, and +//! Merkle nodes. [`BufferRead`] and [`BufferWrite`] provide field operations for +//! owned buffers and views. [`Buffer`] contains owned-only field operations such +//! as construction, folding, and zero-padding. +//! +//! [`DefaultRs`] selects the Reed-Solomon encoder for the active backend. + pub mod cpu; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub mod metal; use std::ops::RangeBounds; use ark_ff::Field; @@ -17,21 +33,56 @@ use crate::{ protocols::{matrix_commit::Encodable, merkle_tree}, }; +pub use cpu::{CpuBuffer, CpuSlice, CpuSliceMut}; + +#[cfg(all(feature = "metal", target_os = "macos"))] +pub use metal::{MetalBuffer, MetalRs, MetalSlice, MetalSliceMut}; + +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type ActiveBuffer = MetalBuffer; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type ActiveSlice<'a, T> = MetalSlice<'a, T>; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type ActiveSliceMut<'a, T> = MetalSliceMut<'a, T>; +#[cfg(all(feature = "metal", target_os = "macos"))] +pub type DefaultRs = MetalRs; + +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type ActiveBuffer = CpuBuffer; +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type ActiveSlice<'a, T> = CpuSlice<'a, T>; +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type ActiveSliceMut<'a, T> = CpuSliceMut<'a, T>; +#[cfg(not(all(feature = "metal", target_os = "macos")))] +pub type DefaultRs = crate::algebra::ntt::NttEngine; + +/// Owned buffer operations over any element type. +/// +/// This trait is not field-specific, so it also covers hash buffers and Merkle +/// node layers. Field arithmetic lives on [`BufferRead`] and [`BufferWrite`]. pub trait BufferOps { + /// Same-backend buffer type used for Merkle tree nodes. type Nodes: BufferOps; fn from_vec(source: Vec) -> Self; fn from_slice(source: &[T]) -> Self; fn as_slice(&self) -> &[T]; + /// Number of rows when the buffer is laid out with `num_cols` columns. fn num_rows(&self, num_cols: usize) -> usize { self.len() / num_cols } + /// Gather full rows `indices[i] * num_cols .. (indices[i] + 1) * num_cols`. fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; fn at_index(&self, index: usize) -> Option; + /// Gather elements at arbitrary indices. + fn gather_at_indices(&self, indices: &[usize]) -> Vec + where + T: Copy; fn len(&self) -> usize; fn is_empty(&self) -> bool { self.len() == 0 } + /// Hash rows of width `num_cols` and build a Merkle tree. fn merklize( &self, num_cols: usize, @@ -42,19 +93,12 @@ pub trait BufferOps { T: Encodable + Send + Sync; } -/// Read-only operations over a contiguous run of field elements — the buffer -/// analogue of `&[F]`. +/// Read-only operations over a contiguous run of field elements. /// -/// Implemented by the owned buffers ([`CpuBuffer`]/`MetalBuffer`, as the -/// full-range case) and by the borrowed read views ([`CpuSlice`]/`MetalSlice`). -/// Every op here works identically on a whole buffer or any sub-range obtained -/// via [`Self::slice`]. On `MetalBuffer` a view aliases the parent's GPU -/// allocation (byte-offset binding), so no data is copied. +/// Implemented by owned buffers and borrowed views. A view produced by +/// [`Self::slice`] aliases the original storage. pub trait BufferRead { - /// A same-backend owned buffer over another field, produced by the - /// mixed-field ops (e.g. `mixed_dot` against a target-field buffer). For an - /// owned buffer this is the buffer itself over `T`; views report the owning - /// buffer type. + /// Same-backend owned buffer over another field. type TargetBuffer: Buffer; /// Read-only view type produced by slicing this buffer/view. @@ -110,19 +154,18 @@ pub trait BufferRead { /// Borrow `self[range]` as a read-only view (analogous to `&v[range]`). fn slice(&self, range: impl RangeBounds) -> Self::Slice<'_>; + + /// Copy this view into a new owned buffer. + fn copy_to_owned(&self) -> Self::TargetBuffer; } -/// In-place, length-preserving operations over a contiguous run of field -/// elements — the buffer analogue of `&mut [F]`. +/// In-place operations over a contiguous run of field elements. /// -/// Implemented by the owned buffers (full-range) and the borrowed mutable views -/// ([`CpuSliceMut`]/`MetalSliceMut`). A sub-range obtained via [`Self::slice_mut`] -/// is the same view type as the whole buffer, so the ops compose just like -/// slicing a `Vec`. +/// Implemented by owned buffers and mutable views. A view produced by +/// [`Self::slice_mut`] aliases the original storage. /// -/// Constructors (`zeros`, `random`, …) and length-changing operations (`fold`, -/// `zero_pad`) live on [`Buffer`], the owned-buffer trait, because a borrowed -/// view cannot construct or resize its backing storage. +/// Construction and length-changing operations live on [`Buffer`], because +/// borrowed views cannot construct or resize their backing storage. pub trait BufferWrite: BufferRead { /// Mutable view type produced by slicing this buffer/view. type SliceMut<'a>: BufferWrite @@ -148,13 +191,10 @@ pub trait BufferWrite: BufferRead { fn split_at_mut(&mut self, mid: usize) -> (Self::SliceMut<'_>, Self::SliceMut<'_>); } -/// Owned field-buffer operations — the `Vec` half of the field API. +/// Owned field-buffer operations. /// -/// `BufferOps` stays generic over any element type, so hashes and digests can -/// use it. `BufferRead`/`BufferWrite` are the slice-like field ops implemented -/// by owners and views. This trait is only for owned field buffers: operations -/// here construct storage or change its length, so borrowed views cannot -/// implement them. +/// This trait contains operations that construct storage or change its length, +/// so it is implemented only by owned buffers. pub trait Buffer: BufferOps + BufferWrite + Clone { fn zeros(length: usize) -> Self; diff --git a/src/hash/mod.rs b/src/hash/mod.rs index 94c6570f..2618beda 100644 --- a/src/hash/mod.rs +++ b/src/hash/mod.rs @@ -2,10 +2,6 @@ mod blake3_engine; mod copy_engine; mod digest_engine; mod hash_counter; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub(crate) mod metal_profile; -#[cfg(all(feature = "metal", target_os = "macos"))] -mod metal_sha2_engine; use core::fmt; use std::{ @@ -19,12 +15,9 @@ use static_assertions::{assert_impl_all, assert_obj_safe}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned}; #[cfg(all(feature = "metal", target_os = "macos"))] -pub use self::metal_profile::{ - reset_device_peak as metal_reset_device_peak, snapshot as metal_profile_snapshot, - MetalProfileSnapshot, +pub use crate::buffer::metal::{ + metal_profile_snapshot, metal_reset_device_peak, MetalProfileSnapshot, MetalSha2, }; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub use self::metal_sha2_engine::MetalSha2; pub use self::{ blake3_engine::{Blake3, BLAKE3}, copy_engine::{Copy, COPY}, diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 4cefdd25..e1e4adeb 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -13,11 +13,8 @@ use serde::{Deserialize, Serialize}; use spongefish::{Decoding, VerificationResult}; use crate::{ - algebra::{ - buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, - embedding::Identity, - multilinear_extend, scalar_mul_add_new, univariate_evaluate, - }, + algebra::{embedding::Identity, multilinear_extend, scalar_mul_add_new, univariate_evaluate}, + buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, hash::Hash, protocols::{irs_commit, sumcheck}, transcript::{ diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index e1108081..0d7153aa 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -13,13 +13,13 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::{ActiveBuffer, BufferOps, BufferRead, BufferWrite}, dot, embedding::{Embedding, Identity}, eq_weights, lift, linear_form::UnivariateEvaluation, mixed_dot, }, + buffer::{ActiveBuffer, BufferOps, BufferRead, BufferWrite}, hash::Hash, protocols::{ geometric_challenge::geometric_challenge, diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 36a00d32..ad345b4b 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -20,27 +20,20 @@ use std::{ f64::{self, consts::LOG2_10}, fmt, - marker::PhantomData, ops::Neg, }; -use crate::algebra::buffer::{Buffer, BufferRead}; +use crate::buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use ordered_float::OrderedFloat; use serde::{Deserialize, Serialize}; - #[cfg(feature = "tracing")] use tracing::instrument; use crate::{ algebra::{ - buffer::{ActiveBuffer, BufferOps}, - dot, - embedding::Embedding, - fields::FieldWithSize, - lift, - linear_form::UnivariateEvaluation, + dot, embedding::Embedding, fields::FieldWithSize, lift, linear_form::UnivariateEvaluation, ntt, }, engines::EngineId, @@ -108,7 +101,6 @@ where pub matrix: ActiveBuffer, pub matrix_witness: matrix_commit::Witness, pub out_of_domain: Evaluations, - source_field: PhantomData, } #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] @@ -323,16 +315,8 @@ impl Config { prover_state.rng(), self.mask_length * self.num_messages(), ); - let matrix = ntt::NTT - .get::() - .expect("Unsupported NTT field.") - .interleaved_encode( - vectors, - &masks, - self.message_length(), - self.interleaving_depth, - self.codeword_length, - ); + let messages = ntt::Messages::new(vectors, self.message_length(), self.interleaving_depth); + let matrix = ntt::interleaved_rs_encode(messages, &masks, self.codeword_length); // Commit to the matrix let matrix_witness = self.matrix_commit.commit(prover_state, &matrix); @@ -360,7 +344,6 @@ impl Config { points: oods_points, matrix: oods_matrix, }, - source_field: PhantomData, } } @@ -523,11 +506,7 @@ impl Commitment { } } -impl Witness -where - F: Field, - G: Field, -{ +impl Witness { /// Returns the out-of-domain evaluations. pub const fn out_of_domain(&self) -> &Evaluations { &self.out_of_domain diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index 63ecb37a..abce0b82 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -45,11 +45,8 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, R use serde::{Deserialize, Serialize}; use crate::{ - algebra::{ - buffer::{ActiveBuffer, Buffer, BufferOps}, - embedding::Identity, - scalar_mul_add_new, univariate_evaluate, - }, + algebra::{embedding::Identity, univariate_evaluate}, + buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, hash::Hash, protocols::irs_commit::{ Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness, @@ -184,7 +181,7 @@ impl Config { // Step 2: compute and send combined polynomials + IRS randomness let irs_masks_per_vector = self.c_zk_commit.mask_length * self.c_zk_commit.interleaving_depth; - let irs_masks = witness.mask_witness.masks.as_slice(); + let irs_masks = &witness.mask_witness.masks; assert_eq!(irs_masks.len(), 2 * self.num_masks * irs_masks_per_vector); for (i, (orig_msg, fresh_msg)) in original_msgs .iter() @@ -193,17 +190,19 @@ impl Config { .enumerate() { // ξ*_i = s_i + γ · ξ_i - // TODO: port to buffer backend - let combined_msg = scalar_mul_add_new(fresh_msg.as_slice(), gamma, orig_msg.as_slice()); - prover_state.prover_messages(&combined_msg); + let mut combined_msg = fresh_msg.copy_to_owned(); + orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); + prover_state.prover_messages(combined_msg.as_slice()); // r*_i = r'_i + γ · r_i if irs_masks_per_vector > 0 { - let orig_r = &irs_masks[i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; - let fresh_r = &irs_masks[(self.num_masks + i) * irs_masks_per_vector - ..(self.num_masks + i + 1) * irs_masks_per_vector]; - let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); - prover_state.prover_messages(&combined_r); + let base = i * irs_masks_per_vector; + let fresh_base = (self.num_masks + i) * irs_masks_per_vector; + let orig_r = irs_masks.slice(base..base + irs_masks_per_vector); + let fresh_r = irs_masks.slice(fresh_base..fresh_base + irs_masks_per_vector); + let mut combined_r = fresh_r.copy_to_owned(); + orig_r.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_r, gamma); + prover_state.prover_messages(combined_r.as_slice()); } } @@ -477,7 +476,7 @@ mod tests { let gamma: F = prover_state.verifier_message(); let irs_masks_per_vector = config.c_zk_commit.mask_length * config.c_zk_commit.interleaving_depth; - let irs_masks = witness.mask_witness.masks.as_slice(); + let irs_masks = &witness.mask_witness.masks; for (i, (orig_msg, fresh_msg)) in original_refs .iter() @@ -485,19 +484,23 @@ mod tests { .zip(witness.fresh_msgs.iter()) .enumerate() { - let mut combined_msg = - scalar_mul_add_new(fresh_msg.as_slice(), gamma, orig_msg.as_slice()); + let mut combined_msg = fresh_msg.copy_to_owned(); + orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); if i == 0 { - combined_msg[0] += F::ONE; + let mut tampered = combined_msg.as_slice().to_vec(); + tampered[0] += F::ONE; + combined_msg = ActiveBuffer::from_slice(&tampered); } - prover_state.prover_messages(&combined_msg); + prover_state.prover_messages(combined_msg.as_slice()); if irs_masks_per_vector > 0 { - let orig_r = &irs_masks[i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; - let fresh_r = &irs_masks[(config.num_masks + i) * irs_masks_per_vector - ..(config.num_masks + i + 1) * irs_masks_per_vector]; - let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); - prover_state.prover_messages(&combined_r); + let base = i * irs_masks_per_vector; + let fresh_base = (config.num_masks + i) * irs_masks_per_vector; + let orig_r = irs_masks.slice(base..base + irs_masks_per_vector); + let fresh_r = irs_masks.slice(fresh_base..fresh_base + irs_masks_per_vector); + let mut combined_r = fresh_r.copy_to_owned(); + orig_r.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_r, gamma); + prover_state.prover_messages(combined_r.as_slice()); } } diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index d725ed41..ceb966fc 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -12,7 +12,7 @@ use tracing::instrument; use zerocopy::{Immutable, IntoBytes}; use crate::{ - algebra::buffer::ActiveBuffer, + buffer::ActiveBuffer, buffer::BufferOps, engines::EngineId, hash::{self, Hash}, @@ -260,26 +260,14 @@ impl Config { R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>, { - #[cfg(all(feature = "metal", target_os = "macos"))] - if let Some(hints) = metal_merkle_opening_hints( - &witness.nodes, + let node_indices = merkle_tree::opening_sibling_indices( self.merkle_tree.num_leaves, self.merkle_tree.layers.len(), indices, - ) { - for hint in hints { - prover_state.prover_hint(&hint); - } - return; - } - - self.merkle_tree.open( - prover_state, - &merkle_tree::Witness { - nodes: Vec::from(witness.nodes.as_slice()), - }, - indices, ); + for hint in witness.nodes.gather_at_indices(&node_indices) { + prover_state.prover_hint(&hint); + } } /// Verifies the commitment at the provided row indices. @@ -313,49 +301,6 @@ impl Config { } } -#[cfg(all(feature = "metal", target_os = "macos"))] -fn metal_merkle_opening_hints( - nodes: &ActiveBuffer, - num_leaves: usize, - layers: usize, - indices: &[usize], -) -> Option> { - if num_leaves != (1usize << layers) { - return None; - } - assert!(indices.iter().all(|&i| i < num_leaves)); - - let mut indices = indices.to_vec(); - indices.sort_unstable(); - indices.dedup(); - - let mut node_indices = Vec::new(); - let mut layer_offset = 0usize; - let mut layer_len = 1usize << layers; - while layer_len > 1 { - let mut next_indices = Vec::with_capacity(indices.len()); - let mut iter = indices.iter().copied().peekable(); - loop { - match (iter.next(), iter.peek()) { - (Some(a), Some(&b)) if b == a ^ 1 => { - next_indices.push(a >> 1); - iter.next(); - } - (Some(a), _) => { - node_indices.push(layer_offset + (a ^ 1)); - next_indices.push(a >> 1); - } - (None, _) => break, - } - } - indices = next_indices; - layer_offset += layer_len; - layer_len /= 2; - } - - nodes.read_hash_indices(&node_indices) -} - impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "MatrixCommit({} x {})", self.num_rows(), self.num_cols) diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index 111b998d..006a2bca 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -149,33 +149,10 @@ impl Config { assert_eq!(witness.nodes.len(), self.num_nodes()); assert!(indices.iter().all(|&i| i < self.num_leaves)); - // Abstract execution of verify algorithm writing required hashes. - let mut indices = indices.to_vec(); - indices.sort_unstable(); - indices.dedup(); - let (mut layer, mut remaining) = witness.nodes.split_at(1 << self.layers.len()); - while layer.len() > 1 { - let mut next_indices = Vec::with_capacity(indices.len()); - let mut iter = indices.iter().copied().peekable(); - loop { - match (iter.next(), iter.peek()) { - (Some(a), Some(&b)) if b == a ^ 1 => { - // Neighboring indices, merging branches. - next_indices.push(a >> 1); - iter.next(); // Skip the next index. - } - (Some(a), _) => { - // Single index, pushing the neighbor hash. - prover_state.prover_hint(&layer[a ^ 1]); - next_indices.push(a >> 1); - } - (None, _) => break, - } - } - indices = next_indices; - let (next_layer, next_remaining) = remaining.split_at(layer.len() / 2); - layer = next_layer; - remaining = next_remaining; + let node_indices = + opening_sibling_indices(self.num_leaves, self.layers.len(), indices); + for &i in &node_indices { + prover_state.prover_hint(&witness.nodes[i]); } } @@ -366,3 +343,41 @@ pub(crate) mod tests { assert_eq!(layers_for_size(8), 3); } } + +/// Flat node indices of sibling hashes required to open at `indices`. +pub fn opening_sibling_indices( + num_leaves: usize, + layers: usize, + indices: &[usize], +) -> Vec { + debug_assert!(indices.iter().all(|&i| i < num_leaves)); + + let mut indices = indices.to_vec(); + indices.sort_unstable(); + indices.dedup(); + + let mut node_indices = Vec::new(); + let mut layer_offset = 0usize; + let mut layer_len = 1usize << layers; + while layer_len > 1 { + let mut next_indices = Vec::with_capacity(indices.len()); + let mut iter = indices.iter().copied().peekable(); + loop { + match (iter.next(), iter.peek()) { + (Some(a), Some(&b)) if b == a ^ 1 => { + next_indices.push(a >> 1); + iter.next(); + } + (Some(a), _) => { + node_indices.push(layer_offset + (a ^ 1)); + next_indices.push(a >> 1); + } + (None, _) => break, + } + } + indices = next_indices; + layer_offset += layer_len; + layer_len /= 2; + } + node_indices +} diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index ae83e348..7b4b36a9 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -2,7 +2,7 @@ use std::fmt; -use crate::algebra::buffer::{Buffer, BufferRead}; +use crate::buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}; use ark_ff::Field; use ark_std::rand::{CryptoRng, RngCore}; use serde::{Deserialize, Serialize}; @@ -10,10 +10,7 @@ use serde::{Deserialize, Serialize}; use tracing::instrument; use crate::{ - algebra::{ - buffer::{ActiveBuffer, BufferOps}, - univariate_evaluate, - }, + algebra::univariate_evaluate, protocols::proof_of_work, transcript::{ codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverState, VerificationResult, @@ -240,11 +237,11 @@ mod tests { use super::*; use crate::{ algebra::{ - buffer::ActiveBuffer, dot, fields::{self, Field64}, multilinear_extend, random_vector, }, + buffer::ActiveBuffer, transcript::DomainSeparator, }; use ark_std::rand::{ diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 80f64237..67497b30 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -14,10 +14,10 @@ use tracing::instrument; use crate::{ algebra::{ - buffer::ActiveBuffer, embedding::{Embedding, Identity}, linear_form::LinearForm, }, + buffer::ActiveBuffer, hash::Hash, protocols::{irs_commit, proof_of_work, sumcheck}, transcript::{ @@ -77,10 +77,7 @@ impl FinalClaim { impl Config { /// Commit to one or more vectors. - #[cfg_attr( - feature = "tracing", - instrument(skip_all, fields(size = vectors.first().unwrap().len())) - )] + #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().unwrap().len())))] pub fn commit( &self, prover_state: &mut ProverState, @@ -131,10 +128,9 @@ mod tests { use ark_std::rand::thread_rng; use super::*; - use crate::buffer::BufferOps; + use crate::buffer::{ActiveBuffer, BufferOps}; use crate::{ algebra::{ - buffer::ActiveBuffer, embedding::Basefield, fields::{Field64, Field64_3}, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index d84f7555..5b66145f 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -7,14 +7,8 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ - algebra::{ - buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead, BufferWrite}, - dot, - embedding::Embedding, - eq_weights, - linear_form::LinearForm, - tensor_product, - }, + algebra::{dot, embedding::Embedding, eq_weights, linear_form::LinearForm, tensor_product}, + buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead, BufferWrite}, hash::Hash, protocols::{geometric_challenge::geometric_challenge, irs_commit, whir::FinalClaim}, transcript::{ diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index 3d525a91..94f390c2 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -8,7 +8,7 @@ use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; use crate::buffer::BufferOps; use crate::{ - algebra::buffer::ActiveBuffer, + buffer::ActiveBuffer, hash::Hash, protocols::{irs_commit, whir}, transcript::{ From 85f44c1764fedf778b8efd0ab0c0014ffa0770dd Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Fri, 19 Jun 2026 16:10:35 +0200 Subject: [PATCH 21/35] refactor: remove experimental Metal backend --- Cargo.lock | 94 --- Cargo.toml | 8 - benches/metal_buffer.rs | 87 --- src/bin/gpu_proof_cpu_verify.rs | 193 ----- src/buffer/metal/buffer.rs | 1291 ------------------------------- src/buffer/metal/kernels.rs | 943 ---------------------- src/buffer/metal/mod.rs | 16 - src/buffer/metal/profile.rs | 153 ---- src/buffer/metal/rs.rs | 173 ----- src/buffer/metal/runtime.rs | 1218 ----------------------------- src/buffer/metal/sha2.rs | 722 ----------------- src/buffer/mod.rs | 18 - src/hash/mod.rs | 6 - src/protocols/matrix_commit.rs | 14 - src/protocols/proof_of_work.rs | 10 - 15 files changed, 4946 deletions(-) delete mode 100644 benches/metal_buffer.rs delete mode 100644 src/bin/gpu_proof_cpu_verify.rs delete mode 100644 src/buffer/metal/buffer.rs delete mode 100644 src/buffer/metal/kernels.rs delete mode 100644 src/buffer/metal/mod.rs delete mode 100644 src/buffer/metal/profile.rs delete mode 100644 src/buffer/metal/rs.rs delete mode 100644 src/buffer/metal/runtime.rs delete mode 100644 src/buffer/metal/sha2.rs diff --git a/Cargo.lock b/Cargo.lock index e1bd056b..acc6ca37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -195,12 +195,6 @@ dependencies = [ "digest", ] -[[package]] -name = "block" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" - [[package]] name = "block-buffer" version = "0.10.4" @@ -394,33 +388,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" -[[package]] -name = "core-foundation" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" -dependencies = [ - "core-foundation-sys", - "libc", -] - -[[package]] -name = "core-foundation-sys" -version = "0.8.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" - -[[package]] -name = "core-graphics-types" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb" -dependencies = [ - "bitflags", - "core-foundation", - "libc", -] - [[package]] name = "cpufeatures" version = "0.2.17" @@ -577,33 +544,6 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" -[[package]] -name = "foreign-types" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" -dependencies = [ - "foreign-types-macros", - "foreign-types-shared", -] - -[[package]] -name = "foreign-types-macros" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "foreign-types-shared" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" - [[package]] name = "generic-array" version = "0.14.7" @@ -760,15 +700,6 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" -[[package]] -name = "malloc_buf" -version = "0.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" -dependencies = [ - "libc", -] - [[package]] name = "matchers" version = "0.2.0" @@ -784,21 +715,6 @@ version = "2.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" -[[package]] -name = "metal" -version = "0.33.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7047791b5bc903b8cd963014b355f71dc9864a9a0b727057676c1dcae5cbc15" -dependencies = [ - "bitflags", - "block", - "core-graphics-types", - "foreign-types", - "log", - "objc", - "paste", -] - [[package]] name = "nix" version = "0.30.1" @@ -848,15 +764,6 @@ dependencies = [ "autocfg", ] -[[package]] -name = "objc" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" -dependencies = [ - "malloc_buf", -] - [[package]] name = "once_cell" version = "1.21.3" @@ -1640,7 +1547,6 @@ dependencies = [ "digest", "hex", "hex-literal", - "metal", "ordered-float", "proptest", "rayon", diff --git a/Cargo.toml b/Cargo.toml index 8ab53a45..515ee2b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -56,9 +56,6 @@ arrayvec = "0.7.6" derive-where = { version = "1.6.0", features = ["safe"] } ordered-float = { version = "5.1.0", features = ["serde"] } -[target.'cfg(target_os = "macos")'.dependencies] -metal = { version = "0.33.0", optional = true } - [dev-dependencies] proptest = "1.0" serde_json = "1.0" @@ -77,7 +74,6 @@ parallel = ["dep:rayon"] rayon = ["dep:rayon"] asm = ["ark-ff/asm"] tracing = ["dep:tracing"] -metal = ["dep:metal"] # Do not permute evaluation pointsin RS # This flag is to be removed after whir_zk supports it. rs_in_order = [] @@ -90,10 +86,6 @@ harness = false name = "sumcheck" harness = false -[[bench]] -name = "metal_buffer" -harness = false - # Disable untill fixed. # [[bench]] # name = "zk_whir" diff --git a/benches/metal_buffer.rs b/benches/metal_buffer.rs deleted file mode 100644 index e1fdb39d..00000000 --- a/benches/metal_buffer.rs +++ /dev/null @@ -1,87 +0,0 @@ -#[cfg(all(feature = "metal", target_os = "macos"))] -mod bench { - use divan::{black_box, AllocProfiler, Bencher}; - use whir::{ - algebra::{ - buffer::{BufferRead, CpuBuffer, MetalBuffer}, - fields::Field256 as F, - }, - buffer::BufferOps, - hash::{Hash, HashEngine, MetalSha2, Sha2}, - }; - - #[global_allocator] - static ALLOC: AllocProfiler = AllocProfiler::system(); - - const FIELD_SIZES: &[usize] = &[1 << 8, 1 << 10]; - const HASH_ROWS: &[usize] = &[1 << 10, 1 << 12]; - const HASH_ROW_SIZE: usize = 128; - - #[divan::bench(args = FIELD_SIZES)] - fn cpu_bn254_dot(bencher: Bencher, size: usize) { - bencher - .with_inputs(|| { - let a = CpuBuffer::from_vec((0..size).map(|i| F::from(i as u64)).collect()); - let b = CpuBuffer::from_vec((0..size).map(|i| F::from((i + 7) as u64)).collect()); - (a, b) - }) - .bench_values(|(a, b)| black_box(a.dot(&b))); - } - - #[divan::bench(args = FIELD_SIZES)] - fn metal_bn254_dot(bencher: Bencher, size: usize) { - bencher - .with_inputs(|| { - let a = MetalBuffer::from_vec((0..size).map(|i| F::from(i as u64)).collect()); - let b = MetalBuffer::from_vec((0..size).map(|i| F::from((i + 7) as u64)).collect()); - (a, b) - }) - .bench_values(|(a, b)| black_box(a.dot(&b))); - } - - #[divan::bench(args = HASH_ROWS)] - fn cpu_sha256_hash_many(bencher: Bencher, rows: usize) { - bencher - .with_inputs(|| { - let input = (0..rows * HASH_ROW_SIZE) - .map(|i| (i & 0xff) as u8) - .collect::>(); - let output = vec![Hash::default(); rows]; - (input, output) - }) - .bench_values(|(input, mut output)| { - Sha2::new().hash_many(HASH_ROW_SIZE, &input, &mut output); - black_box(output) - }); - } - - #[divan::bench(args = HASH_ROWS)] - fn metal_sha256_hash_many(bencher: Bencher, rows: usize) { - bencher - .with_inputs(|| { - let input = (0..rows * HASH_ROW_SIZE) - .map(|i| (i & 0xff) as u8) - .collect::>(); - let output = vec![Hash::default(); rows]; - (input, output) - }) - .bench_values(|(input, mut output)| { - MetalSha2::new().hash_many(HASH_ROW_SIZE, &input, &mut output); - black_box(output) - }); - } - - pub fn main() { - divan::main(); - } -} - -#[cfg(all(feature = "metal", target_os = "macos"))] -fn main() { - bench::main(); -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -fn main() { - eprintln!("metal_buffer benchmark requires macOS and --features metal"); -} diff --git a/src/bin/gpu_proof_cpu_verify.rs b/src/bin/gpu_proof_cpu_verify.rs deleted file mode 100644 index d9c27dc4..00000000 --- a/src/bin/gpu_proof_cpu_verify.rs +++ /dev/null @@ -1,193 +0,0 @@ -use std::{ - borrow::Cow, - fs, - path::{Path, PathBuf}, -}; - -use clap::{Parser, Subcommand}; -use serde::{Deserialize, Serialize}; -use whir::{ - algebra::{embedding::Identity, fields::Field256, linear_form::LinearForm}, - hash, - parameters::ProtocolParameters, - protocols::whir::Config as WhirConfig, - transcript::{codecs::Empty, DomainSeparator, Proof, ProverState, VerifierState}, -}; - -use whir::buffer::{ActiveBuffer, BufferOps}; - -type F = Field256; -type M = Identity; - -#[derive(Parser, Debug)] -#[command( - author, - version, - about = "Produce a WHIR proof and verify it from a serialized artifact" -)] -struct Args { - #[command(subcommand)] - command: Command, -} - -#[derive(Subcommand, Debug)] -enum Command { - Prove(RoundtripArgs), - Verify { - #[arg(long)] - input: PathBuf, - }, -} - -#[derive(Parser, Debug)] -struct RoundtripArgs { - #[arg(long)] - output: PathBuf, - - #[arg(long, default_value_t = 16)] - log_size: usize, - - #[arg(long, default_value_t = 4)] - fold: usize, - - #[arg(long, default_value_t = 1)] - rate: usize, - - #[arg(long, default_value_t = 20)] - pow_bits: usize, - - #[arg(long, default_value_t = 128)] - security_level: usize, -} - -#[derive(Serialize, Deserialize)] -struct Artifact { - log_size: usize, - fold: usize, - rate: usize, - pow_bits: usize, - security_level: usize, - proof: Proof, - proof_bytes: usize, -} - -fn main() { - let args = Args::parse(); - match args.command { - Command::Prove(args) => prove(&args), - Command::Verify { input } => verify(&input), - } -} - -fn prove(args: &RoundtripArgs) { - assert!(args.fold <= args.log_size, "fold must be <= log_size"); - let size = 1usize << args.log_size; - let whir_params = protocol_parameters(args.security_level, args.pow_bits, args.fold, args.rate); - let params = WhirConfig::::new(size, &whir_params); - let ds = DomainSeparator::protocol(¶ms) - .session(&"gpu proof cpu verify") - .instance(&Empty); - - let vector = input_vector(size); - let vector_buffer = ActiveBuffer::from_slice(&vector); - - let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &[&vector_buffer]); - let _ = params.prove( - &mut prover_state, - &[&vector_buffer], - vec![&witness], - vec![], - Cow::Owned(vec![]), - ); - let proof = prover_state.proof(); - let proof_bytes = proof.narg_string.len() + proof.hints.len(); - let artifact = Artifact { - log_size: args.log_size, - fold: args.fold, - rate: args.rate, - pow_bits: args.pow_bits, - security_level: args.security_level, - proof, - proof_bytes, - }; - let encoded = serde_json::to_vec_pretty(&artifact).expect("serialize proof artifact"); - fs::write(&args.output, encoded).expect("write proof artifact"); - println!( - "wrote proof artifact backend={} log_size={} fold={} rate={} pow_bits={} proof_bytes={}", - backend_name(), - args.log_size, - args.fold, - args.rate, - args.pow_bits, - proof_bytes - ); -} - -fn verify(input: &Path) { - let bytes = fs::read(input).expect("read proof artifact"); - let artifact: Artifact = serde_json::from_slice(&bytes).expect("decode proof artifact"); - let size = 1usize << artifact.log_size; - let whir_params = protocol_parameters( - artifact.security_level, - artifact.pow_bits, - artifact.fold, - artifact.rate, - ); - let params = WhirConfig::::new(size, &whir_params); - let ds = DomainSeparator::protocol(¶ms) - .session(&"gpu proof cpu verify") - .instance(&Empty); - - let mut verifier_state = VerifierState::new_std(&ds, &artifact.proof); - let commitment = params - .receive_commitment(&mut verifier_state) - .expect("receive commitment"); - let final_claim = params - .verify(&mut verifier_state, &[&commitment], &[]) - .expect("verify WHIR proof"); - let no_forms: Vec<&dyn LinearForm> = Vec::new(); - final_claim.verify(no_forms).expect("verify final claim"); - verifier_state.check_eof().expect("proof EOF"); - println!( - "verified proof artifact backend={} log_size={} fold={} rate={} pow_bits={} proof_bytes={}", - backend_name(), - artifact.log_size, - artifact.fold, - artifact.rate, - artifact.pow_bits, - artifact.proof_bytes - ); -} - -const fn protocol_parameters( - security_level: usize, - pow_bits: usize, - fold: usize, - rate: usize, -) -> ProtocolParameters { - ProtocolParameters { - security_level, - pow_bits, - initial_folding_factor: fold, - folding_factor: fold, - unique_decoding: false, - starting_log_inv_rate: rate, - batch_size: 1, - hash_id: hash::SHA2, - } -} - -fn input_vector(size: usize) -> Vec { - (0..size).map(|i| F::from(i as u64)).collect() -} - -#[cfg(all(feature = "metal", target_os = "macos"))] -const fn backend_name() -> &'static str { - "gpu-metal" -} - -#[cfg(not(all(feature = "metal", target_os = "macos")))] -const fn backend_name() -> &'static str { - "cpu" -} diff --git a/src/buffer/metal/buffer.rs b/src/buffer/metal/buffer.rs deleted file mode 100644 index 4a73184d..00000000 --- a/src/buffer/metal/buffer.rs +++ /dev/null @@ -1,1291 +0,0 @@ -// NOTE: 100% AI GENERATED - -use std::{ - any::type_name, - cell::OnceCell, - cmp::Ordering, - hash::{Hash, Hasher}, - marker::PhantomData, - ops::RangeBounds, -}; - -use ark_ff::{AdditiveGroup, Field}; -use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; -use metal::Buffer; -use serde::{Deserialize, Deserializer, Serialize, Serializer}; - -use crate::{ - algebra::{ - embedding::{Embedding, Identity}, - fields::Field256, - linear_form::{Covector, LinearForm, UnivariateEvaluation}, - }, - buffer::{BufferOps, BufferRead, BufferWrite}, - engines::EngineId, - hash::{self, Hash as Digest}, - protocols::{ - matrix_commit::{hash_rows, Encodable}, - merkle_tree, - }, -}; - -use super::runtime::{ - as_field256_slice, assert_bn254, copy_field_buffer, copy_field_buffer_at, download_field, - download_hash_indices, - encode_field_rows_le, field256_to_f, field256_to_target, f_to_field256, - geometric_accumulate_at_offset, geometric_accumulate_chunk_size, maybe_upload_bn254, - parallel_dot, parallel_dot_at, parallel_fold_pair_sumcheck, parallel_geometric_accumulate_point_blocks, - parallel_geometric_accumulate_point_blocks_batched, parallel_multilinear_extend_at, - parallel_sumcheck, parallel_sumcheck_at, parallel_univariate_evaluate, - parallel_univariate_evaluate_at, read_bn254_rows, run_in_place, runtime, scalar_mul_add_at, - scalar_mul_at_offset, should_use_geometric_point_blocks, should_use_geometric_point_blocks_batched, - target_to_field256, upload_field, zeroed_field_buffer, -}; -use super::sha2::MetalSha2; - -#[derive(Clone, Debug)] -pub(crate) struct MetalFieldBuffer { - pub(crate) limbs: Buffer, -} - -#[derive(Clone, Debug)] -pub(crate) struct MetalHashBuffer { - pub(crate) bytes: Buffer, -} - -#[derive(Clone, Debug)] -pub struct MetalBuffer { - len: usize, - host_cache: OnceCell>, - field: Option, - hash: Option, - _marker: PhantomData, -} - -impl PartialEq for MetalBuffer -where - T: PartialEq, -{ - fn eq(&self, other: &Self) -> bool { - self.as_slice() == other.as_slice() - } -} - -impl Eq for MetalBuffer {} - -impl PartialOrd for MetalBuffer { - fn partial_cmp(&self, other: &Self) -> Option { - self.as_slice().partial_cmp(other.as_slice()) - } -} - -impl Ord for MetalBuffer { - fn cmp(&self, other: &Self) -> Ordering { - self.as_slice().cmp(other.as_slice()) - } -} - -impl Hash for MetalBuffer { - fn hash(&self, state: &mut H) { - self.as_slice().hash(state); - } -} - -impl Default for MetalBuffer { - fn default() -> Self { - Self { - len: 0, - host_cache: OnceCell::new(), - field: None, - hash: None, - _marker: PhantomData, - } - } -} - -impl Serialize for MetalBuffer { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - self.as_slice().serialize(serializer) - } -} - -impl<'de, T> Deserialize<'de> for MetalBuffer -where - T: Clone + Deserialize<'de>, -{ - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let data = Vec::::deserialize(deserializer)?; - Ok(BufferOps::from_vec(data)) - } -} - -impl MetalBuffer { - pub fn warmup() { - super::runtime::init(); - } - - pub(crate) fn as_slice(&self) -> &[T] { - self.host_cache - .get_or_init(|| self.download_host_cache()) - .as_slice() - } - - pub(crate) fn hash_bn254_rows_sha2(&self, num_cols: usize, out: &mut [Digest]) -> bool { - if type_name::() != type_name::() || self.field.is_none() { - return false; - } - assert_eq!(self.len(), num_cols * out.len()); - let message_size = num_cols * size_of::(); - let encoded = encode_field_rows_le( - &self - .field - .as_ref() - .expect("missing Metal field buffer") - .limbs, - out.len(), - num_cols, - ); - MetalSha2::new().hash_many_buffer(message_size, &encoded, out.len(), out); - true - } - - pub(crate) fn commit_bn254_rows_sha2_merkle( - &self, - num_cols: usize, - num_rows: usize, - layers: usize, - ) -> Option> { - if type_name::() != type_name::() || self.field.is_none() { - return None; - } - if num_rows != (1usize << layers) { - return None; - } - assert_eq!(self.len(), num_cols * num_rows); - let message_size = num_cols * size_of::(); - let encoded = encode_field_rows_le( - &self - .field - .as_ref() - .expect("missing Metal field buffer") - .limbs, - num_rows, - num_cols, - ); - let sha = MetalSha2::new(); - let nodes = sha.build_merkle_tree_buffer_from_messages_buffer( - message_size, - &encoded, - num_rows, - layers, - ); - Some(MetalBuffer::::from_digest_buffer( - nodes, - (1usize << (layers + 1)) - 1, - )) - } - - pub(crate) fn read_hash_indices(&self, indices: &[usize]) -> Option> { - let buffer = self.hash.as_ref()?; - Some(download_hash_indices(&buffer.bytes, self.len, indices)) - } -} - -impl MetalBuffer { - pub(crate) fn from_digest_buffer(bytes: Buffer, len: usize) -> Self { - Self { - len, - host_cache: OnceCell::new(), - field: None, - hash: Some(MetalHashBuffer { bytes }), - _marker: PhantomData, - } - } - - pub(crate) fn read_hash_at(&self, index: usize) -> Option { - self.read_hash_indices(&[index]) - .map(|mut values| values.pop().expect("missing hash")) - } -} - -impl BufferOps for MetalBuffer { - type Nodes = MetalBuffer; - - fn as_slice(&self) -> &[T] { - self.as_slice() - } - - fn at_index(&self, index: usize) -> Option { - self.as_slice().get(index).cloned() - } - - fn gather_at_indices(&self, indices: &[usize]) -> Vec - where - T: Copy, - { - if let Some(values) = self.read_hash_indices(indices) { - return values - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - indices - .iter() - .map(|&i| self.at_index(i).expect("index out of bounds")) - .collect() - } - - fn len(&self) -> usize { - self.len - } - - fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { - if type_name::() == type_name::() && self.field.is_some() { - return read_bn254_rows( - self.field.as_ref().expect("missing Metal field buffer"), - num_cols, - indices, - ) - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - let data = self.as_slice(); - let mut result = Vec::with_capacity(indices.len() * num_cols); - for i in indices { - result.extend_from_slice(&data[i * num_cols..(i + 1) * num_cols]); - } - result - } - - fn from_vec(source: Vec) -> Self { - let len = source.len(); - let field = maybe_upload_bn254(&source); - let host_cache = if field.is_some() { - OnceCell::new() - } else { - OnceCell::from(source) - }; - Self { - len, - host_cache, - field, - hash: None, - _marker: PhantomData, - } - } - - fn from_slice(source: &[T]) -> Self { - Self::from_vec(Vec::from(source)) - } - - fn merklize( - &self, - num_cols: usize, - leaf_hash: EngineId, - merkle: &merkle_tree::Config, - ) -> (Self::Nodes, Digest) - where - T: Encodable + Send + Sync, - { - let num_rows = merkle.num_leaves; - let layers = merkle.layers.len(); - assert_eq!(self.len(), num_cols * num_rows); - - // Fast path: build the whole tree on the GPU when the leaf and every node layer is SHA2. - if leaf_hash == hash::SHA2 - && merkle - .layers - .iter() - .all(|layer| layer.hash_id == hash::SHA2) - { - if let Some(nodes) = self.commit_bn254_rows_sha2_merkle(num_cols, num_rows, layers) { - let root = nodes - .read_hash_at(merkle.num_nodes() - 1) - .expect("missing Metal Merkle root"); - return (nodes, root); - } - } - - let cpu_nodes = || { - let engine = hash::ENGINES - .retrieve(leaf_hash) - .expect("Failed to retrieve hash engine"); - let mut leaves = vec![Digest::default(); num_rows]; - hash_rows(&*engine, self.as_slice(), &mut leaves); - merkle.build_nodes(leaves) - }; - - let nodes = if leaf_hash == hash::SHA2 { - let mut leaves = vec![Digest::default(); num_rows]; - if self.hash_bn254_rows_sha2(num_cols, &mut leaves) { - merkle.build_nodes(leaves) - } else { - cpu_nodes() - } - } else { - cpu_nodes() - }; - let root = nodes[merkle.num_nodes() - 1]; - (BufferOps::from_vec(nodes), root) - } -} - -impl crate::buffer::Buffer for MetalBuffer { - fn zeros(length: usize) -> Self { - assert_bn254::(); - // Montgomery zero is all-zero bytes, so a device-side fill suffices. - Self { - len: length, - host_cache: OnceCell::new(), - field: Some(zeroed_field_buffer(length)), - hash: None, - _marker: PhantomData, - } - } - - fn geometric_sequence(base: F, length: usize) -> Self { - assert_bn254::(); - BufferOps::from_vec(crate::algebra::geometric_sequence(base, length)) - } - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution, - { - assert_bn254::(); - BufferOps::from_vec((0..length).map(|_| rng.gen()).collect()) - } - - fn zero_pad(&mut self) { - assert_bn254::(); - if !self.is_empty() { - let mut data = self.as_slice().to_vec(); - data.resize(self.len().next_power_of_two(), F::ZERO); - *self = BufferOps::from_vec(data); - } - } - - fn fold(&mut self, weight: F) { - assert_bn254::(); - if self.len() <= 1 { - return; - } - let len = self.len(); - let fold_half = len.next_power_of_two() >> 1; - let weight = upload_field(&[f_to_field256(weight)]); - let field = self.bn254_buffer(); - run_in_place( - "bn254_fold", - &[&field.limbs, &weight.limbs], - &[len as u32, fold_half as u32], - fold_half, - ); - self.len = fold_half; - self.invalidate_host_cache(); - } - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - if self.len() <= 1 { - return; - } - let len = self.len(); - let fold_half = len.next_power_of_two() >> 1; - let weight = upload_field(&[f_to_field256(weight)]); - let this = self.bn254_buffer(); - let other_buffer = other.bn254_buffer(); - run_in_place( - "bn254_fold_pair", - &[&this.limbs, &other_buffer.limbs, &weight.limbs], - &[len as u32, fold_half as u32], - fold_half, - ); - self.len = fold_half; - other.len = fold_half; - self.invalidate_host_cache(); - other.invalidate_host_cache(); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - if self.len() <= 1 { - return self.sumcheck_polynomial(other); - } - let len = self.len(); - let fold_half = len.next_power_of_two() >> 1; - if fold_half == 1 { - self.fold_pair(other, weight); - return self.sumcheck_polynomial(other); - } - let weight = upload_field(&[f_to_field256(weight)]); - let this = self.bn254_buffer(); - let other_buffer = other.bn254_buffer(); - let (c0, c2) = parallel_fold_pair_sumcheck(&this, &other_buffer, &weight, len, fold_half); - self.len = fold_half; - other.len = fold_half; - self.invalidate_host_cache(); - other.invalidate_host_cache(); - (field256_to_f::(c0), field256_to_f::(c2)) - } - - fn linear_forms_rlc( - size: usize, - linear_forms: &mut [Box>], - rlc_coeffs: &[F], - ) -> Self { - assert_bn254::(); - assert_eq!(linear_forms.len(), rlc_coeffs.len()); - let Some((first, rest)) = linear_forms.split_first_mut() else { - return Self::zeros(size); - }; - let first = (first.as_mut() as &mut dyn std::any::Any) - .downcast_mut::>() - .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); - let mut accumulator = >::from_slice(&first.vector); - for (coeff, linear_form) in rlc_coeffs[1..].iter().zip(rest) { - let covector = (linear_form.as_mut() as &mut dyn std::any::Any) - .downcast_mut::>() - .expect("MetalBuffer only supports Covector linear forms for BN254 RLC"); - let vector = >::from_slice(&covector.vector); - vector.mixed_scalar_mul_add_to(&Identity::new(), &mut accumulator, *coeff); - } - accumulator - } - - fn mixed_linear_combination>( - _embedding: &M, - vectors: &[&Self], - coeffs: &[M::Target], - ) -> Self::TargetBuffer { - assert_bn254::(); - assert_eq!(vectors.len(), coeffs.len()); - let Some((first, vectors)) = vectors.split_first() else { - return BufferOps::from_vec(Vec::new()); - }; - let mut accumulator = MetalBuffer:: { - len: first.len(), - host_cache: OnceCell::new(), - field: Some(copy_field_buffer(&first.bn254_buffer(), first.len())), - hash: None, - _marker: PhantomData, - }; - for (coeff, vector) in coeffs[1..].iter().copied().zip(vectors) { - vector.mixed_scalar_mul_add_to(_embedding, &mut accumulator, coeff); - } - accumulator - } -} - -/// Core geometric accumulate over `field[0..len]` (offset 0). Shared by the -/// owned buffer and a full-range view so the optimized chunk strategies are -/// written once. -fn geometric_accumulate_full( - field: &MetalFieldBuffer, - len: usize, - evaluators: &[UnivariateEvaluation], - scalars: &[F], -) { - let points = evaluators - .iter() - .map(|e| f_to_field256(e.point)) - .collect::>(); - let scalars = scalars - .iter() - .copied() - .map(f_to_field256) - .collect::>(); - let points = upload_field(&points); - let scalars = upload_field(&scalars); - let chunk_size = geometric_accumulate_chunk_size(len); - if std::env::var_os("WHIR_METAL_TRACE").is_some() { - eprintln!( - "metal geometric shape len={} points={} chunk={} chunks={}", - len, - evaluators.len(), - chunk_size, - len.div_ceil(chunk_size) - ); - } - if chunk_size <= 1 { - run_in_place( - "bn254_geometric_accumulate", - &[&field.limbs, &points.limbs, &scalars.limbs], - &[len as u32, evaluators.len() as u32], - len, - ); - } else { - let point_steps = evaluators - .iter() - .map(|e| f_to_field256(e.point.pow([chunk_size as u64]))) - .collect::>(); - let point_steps = upload_field(&point_steps); - if should_use_geometric_point_blocks(len, evaluators.len(), chunk_size) { - parallel_geometric_accumulate_point_blocks( - field, - &points, - &point_steps, - &scalars, - len, - evaluators.len(), - chunk_size, - ); - } else if should_use_geometric_point_blocks_batched(len, evaluators.len(), chunk_size) { - parallel_geometric_accumulate_point_blocks_batched( - field, - &points, - &point_steps, - &scalars, - len, - evaluators.len(), - chunk_size, - ); - } else { - run_in_place( - "bn254_geometric_accumulate_chunks_strided", - &[ - &field.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[len as u32, evaluators.len() as u32, chunk_size as u32], - len.div_ceil(chunk_size), - ); - } - } -} - -/// Validate that every evaluator targets `len` entries; returns `false` when -/// there is nothing to accumulate. -fn check_univariate_evaluators( - evaluators: &[UnivariateEvaluation], - scalars: &[F], - len: usize, -) -> bool { - assert_bn254::(); - assert_eq!(evaluators.len(), scalars.len()); - let Some(size) = evaluators.first().map(|e| e.size) else { - return false; - }; - assert_eq!(len, size); - for evaluator in evaluators { - assert_eq!(evaluator.size, size); - } - true -} - -impl BufferRead for MetalBuffer { - type TargetBuffer = MetalBuffer; - type Slice<'a> - = MetalSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.len() - } - - fn dot(&self, other: &Self) -> F { - assert_bn254::(); - assert_eq!(self.len(), other.len()); - let this = self.bn254_buffer(); - let other = other.bn254_buffer(); - field256_to_f::(parallel_dot(&this, &other, self.len())) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - assert_bn254::(); - let len = self.len().min(other.len()); - if len == 0 { - return (F::ZERO, F::ZERO); - } - if len == 1 { - return (self.as_slice()[0] * other.as_slice()[0], F::ZERO); - } - let fold_half = len.next_power_of_two() >> 1; - let this = self.bn254_buffer(); - let other = other.bn254_buffer(); - let (c0, c2) = parallel_sumcheck(&this, &other, len, fold_half); - (field256_to_f::(c0), field256_to_f::(c2)) - } - - fn mixed_extend, T: Field>( - &self, - _embedding: &M, - point: &[M::Target], - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let num_vars = point.len(); - let point = point - .iter() - .copied() - .map(target_to_field256) - .collect::>(); - let point = upload_field(&point); - let this = self.bn254_buffer(); - let value = parallel_multilinear_extend_at(&this, 0, self.len(), &point, num_vars); - field256_to_target::(value) - } - - fn mixed_dot, T: Field>( - &self, - _embedding: &M, - other: &MetalBuffer, - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let this = self.bn254_buffer(); - let other = other.bn254_buffer_target(); - let value = field256_to_f::(parallel_dot(&this, &other, self.len())); - field256_to_target::(f_to_field256(value)) - } - - fn mixed_univariate_evaluate>( - &self, - _embedding: &M, - point: M::Target, - ) -> M::Target { - assert_bn254::(); - let point = target_to_field256(point); - let point = upload_field(&[point]); - let this = self.bn254_buffer(); - field256_to_target::(parallel_univariate_evaluate(&this, &point, self.len())) - } - - fn mixed_scalar_mul_add_to>( - &self, - _embedding: &M, - accumulator: &mut MetalBuffer, - weight: M::Target, - ) { - assert_bn254::(); - let weight = upload_field(&[target_to_field256(weight)]); - let vector = self.bn254_buffer(); - let acc = accumulator.bn254_buffer_target(); - scalar_mul_add_at(&acc, 0, &vector, 0, &weight, self.len()); - accumulator.field = Some(acc); - accumulator.invalidate_host_cache(); - } - - fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - assert_bn254::(); - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); - MetalSlice { - field: self.bn254_buffer(), - offset: start, - len: end - start, - _parent: PhantomData, - } - } - - fn copy_to_owned(&self) -> MetalBuffer { - assert_bn254::(); - MetalBuffer { - len: self.len, - host_cache: OnceCell::new(), - field: Some(copy_field_buffer(&self.bn254_buffer(), self.len)), - hash: None, - _marker: PhantomData, - } - } -} - -impl BufferWrite for MetalBuffer { - type SliceMut<'a> - = MetalSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scalar_mul(&mut self, weight: F) { - assert_bn254::(); - if self.is_empty() { - return; - } - let weight = upload_field(&[f_to_field256(weight)]); - let field = self.bn254_buffer(); - run_in_place( - "bn254_scalar_mul", - &[&field.limbs, &weight.limbs], - &[self.len() as u32], - self.len(), - ); - self.field = Some(field); - self.invalidate_host_cache(); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - if !check_univariate_evaluators(evaluators, scalars, self.len()) { - return; - } - let field = self.bn254_buffer(); - geometric_accumulate_full(&field, self.len(), evaluators, scalars); - self.field = Some(field); - self.invalidate_host_cache(); - } - - fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { - assert_bn254::(); - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len()); - MetalSliceMut::new(self, start, end - start) - } - - fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { - assert_bn254::(); - let len = self.len(); - assert!(mid <= len, "split_at_mut mid {mid} out of bounds for {len}"); - // Materialize the parent's GPU allocation and invalidate its host - // cache once up front; the returned views share the handle and write - // disjoint ranges, so no per-op parent borrow is needed. - let field = self.bn254_buffer(); - self.field = Some(field.clone()); - self.invalidate_host_cache(); - ( - MetalSliceMut::from_field(field.clone(), 0, mid), - MetalSliceMut::from_field(field, mid, len - mid), - ) - } -} -impl MetalBuffer { - pub(crate) fn bn254_buffer(&self) -> MetalFieldBuffer { - self.field - .clone() - .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) - } - - fn invalidate_host_cache(&mut self) { - let _ = self.host_cache.take(); - } - - fn download_host_cache(&self) -> Vec { - if self.field.is_some() && type_name::() == type_name::() { - return download_field( - &self - .field - .as_ref() - .expect("missing Metal field buffer") - .limbs, - self.len, - ) - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - if self.hash.is_some() && type_name::() == type_name::() { - return download_hash_indices( - &self.hash.as_ref().expect("missing Metal hash buffer").bytes, - self.len, - &(0..self.len).collect::>(), - ) - .into_iter() - .map(|value| unsafe { std::mem::transmute_copy(&value) }) - .collect(); - } - panic!( - "MetalBuffer<{}> has no host cache and cannot be materialized", - type_name::() - ); - } -} - -impl MetalBuffer { - pub(crate) fn from_field_limb_buffer(limbs: Buffer, len: usize) -> Self { - Self { - len, - host_cache: OnceCell::new(), - field: Some(MetalFieldBuffer { limbs }), - hash: None, - _marker: PhantomData, - } - } -} - -impl MetalBuffer { - fn bn254_buffer_target(&self) -> MetalFieldBuffer { - self.field - .clone() - .unwrap_or_else(|| upload_field(as_field256_slice(self.as_slice()))) - } -} - -/// Read-only, zero-copy view into a [`MetalBuffer`]'s GPU allocation. -/// -/// Holds a clone of the parent's allocation handle plus an `(offset, len)` -/// window. Reductions bind the handle at the byte offset, so they read -/// `parent[offset + i]` without copying. -pub struct MetalSlice<'a, F> { - field: MetalFieldBuffer, - offset: usize, - len: usize, - _parent: PhantomData<&'a MetalBuffer>, -} - -/// Mutable, zero-copy view into a [`MetalBuffer`]'s GPU allocation. -/// -/// Like [`MetalSlice`] but exclusive: the parent's host cache is invalidated -/// up front when the view is created, and writes go through the shared handle -/// at the byte offset, landing in the parent's own memory. -pub struct MetalSliceMut<'a, F> { - field: MetalFieldBuffer, - offset: usize, - len: usize, - _parent: PhantomData<&'a mut MetalBuffer>, -} - -impl<'a, F: Field + Clone> MetalSliceMut<'a, F> { - /// Borrow `parent[offset..offset+len]`, materializing the parent's GPU - /// allocation and invalidating its host cache once up front. - fn new(parent: &'a mut MetalBuffer, offset: usize, len: usize) -> Self { - let field = parent.bn254_buffer(); - parent.field = Some(field.clone()); - parent.invalidate_host_cache(); - Self { - field, - offset, - len, - _parent: PhantomData, - } - } - - /// Build a view directly from an already-materialized handle (used by - /// `split_at_mut`, where the parent cache was invalidated by the caller). - fn from_field(field: MetalFieldBuffer, offset: usize, len: usize) -> Self { - Self { - field, - offset, - len, - _parent: PhantomData, - } - } - - fn as_read(&self) -> MetalSlice<'_, F> { - MetalSlice { - field: self.field.clone(), - offset: self.offset, - len: self.len, - _parent: PhantomData, - } - } -} - -impl MetalSlice<'_, F> { - pub fn len(&self) -> usize { - self.len - } - pub fn is_empty(&self) -> bool { - self.len == 0 - } -} - -impl MetalSliceMut<'_, F> { - pub fn len(&self) -> usize { - self.len - } - pub fn is_empty(&self) -> bool { - self.len == 0 - } -} - -impl BufferRead for MetalSlice<'_, F> { - type TargetBuffer = MetalBuffer; - type Slice<'a> - = MetalSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.len - } - - fn dot(&self, other: &Self) -> F { - assert_bn254::(); - assert_eq!(self.len, other.len); - field256_to_f::(parallel_dot_at( - &self.field, - self.offset, - &other.field, - other.offset, - self.len, - )) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - assert_bn254::(); - let len = self.len.min(other.len); - if len == 0 { - return (F::ZERO, F::ZERO); - } - let fold_half = len.next_power_of_two() >> 1; - let (c0, c2) = parallel_sumcheck_at( - &self.field, - self.offset, - &other.field, - other.offset, - len, - fold_half, - ); - (field256_to_f::(c0), field256_to_f::(c2)) - } - - fn mixed_extend, T: Field>( - &self, - _embedding: &M, - point: &[M::Target], - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let num_vars = point.len(); - let point = point - .iter() - .copied() - .map(target_to_field256) - .collect::>(); - let point = upload_field(&point); - let value = - parallel_multilinear_extend_at(&self.field, self.offset, self.len, &point, num_vars); - field256_to_target::(value) - } - - fn mixed_dot, T: Field>( - &self, - _embedding: &M, - other: &MetalBuffer, - ) -> M::Target { - assert_bn254::(); - assert_bn254::(); - let other = other.bn254_buffer_target(); - let value = field256_to_f::(parallel_dot_at( - &self.field, - self.offset, - &other, - 0, - self.len, - )); - field256_to_target::(f_to_field256(value)) - } - - fn mixed_univariate_evaluate>( - &self, - _embedding: &M, - point: M::Target, - ) -> M::Target { - assert_bn254::(); - let point = upload_field(&[target_to_field256(point)]); - field256_to_target::(parallel_univariate_evaluate_at( - &self.field, - self.offset, - &point, - self.len, - )) - } - - fn mixed_scalar_mul_add_to>( - &self, - _embedding: &M, - accumulator: &mut MetalBuffer, - weight: M::Target, - ) { - assert_bn254::(); - let weight = upload_field(&[target_to_field256(weight)]); - let acc = accumulator.bn254_buffer_target(); - scalar_mul_add_at(&acc, 0, &self.field, self.offset, &weight, self.len); - accumulator.field = Some(acc); - accumulator.invalidate_host_cache(); - } - - fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); - MetalSlice { - field: self.field.clone(), - offset: self.offset + start, - len: end - start, - _parent: PhantomData, - } - } - - fn copy_to_owned(&self) -> MetalBuffer { - assert_bn254::(); - MetalBuffer { - len: self.len, - host_cache: OnceCell::new(), - field: Some(copy_field_buffer_at(&self.field, self.offset, self.len)), - hash: None, - _marker: PhantomData, - } - } -} - -impl BufferRead for MetalSliceMut<'_, F> { - type TargetBuffer = MetalBuffer; - type Slice<'a> - = MetalSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.len - } - - fn dot(&self, other: &Self) -> F { - self.as_read().dot(&other.as_read()) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - self.as_read().sumcheck_polynomial(&other.as_read()) - } - - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - self.as_read().mixed_extend(embedding, point) - } - - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &MetalBuffer, - ) -> M::Target { - self.as_read().mixed_dot(embedding, other) - } - - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - self.as_read().mixed_univariate_evaluate(embedding, point) - } - - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut MetalBuffer, - weight: M::Target, - ) { - self.as_read() - .mixed_scalar_mul_add_to(embedding, accumulator, weight) - } - - fn slice(&self, range: impl RangeBounds) -> MetalSlice<'_, F> { - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); - MetalSlice { - field: self.field.clone(), - offset: self.offset + start, - len: end - start, - _parent: PhantomData, - } - } - - fn copy_to_owned(&self) -> MetalBuffer { - self.as_read().copy_to_owned() - } -} - -impl BufferWrite for MetalSliceMut<'_, F> { - type SliceMut<'a> - = MetalSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scalar_mul(&mut self, weight: F) { - assert_bn254::(); - if self.len == 0 { - return; - } - let weight = upload_field(&[f_to_field256(weight)]); - scalar_mul_at_offset(&self.field, self.offset, self.len, &weight); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ) { - if !check_univariate_evaluators(evaluators, scalars, self.len) { - return; - } - if self.offset == 0 { - // Offset 0: reuse the optimized full-buffer accumulate path. - geometric_accumulate_full(&self.field, self.len, evaluators, scalars); - } else { - // Arbitrary offset: bind the buffer at a byte offset so the - // kernel's gid-based indexing addresses field[offset + gid]. - let points = evaluators - .iter() - .map(|e| f_to_field256(e.point)) - .collect::>(); - let scalars = scalars - .iter() - .copied() - .map(f_to_field256) - .collect::>(); - let points = upload_field(&points); - let scalars = upload_field(&scalars); - geometric_accumulate_at_offset( - &self.field, - self.offset, - self.len, - &points, - &scalars, - evaluators.len(), - ); - } - } - - fn slice_mut(&mut self, range: impl RangeBounds) -> MetalSliceMut<'_, F> { - let (start, end) = crate::buffer::cpu::resolve_range(range, self.len); - MetalSliceMut::from_field(self.field.clone(), self.offset + start, end - start) - } - - fn split_at_mut(&mut self, mid: usize) -> (MetalSliceMut<'_, F>, MetalSliceMut<'_, F>) { - assert!( - mid <= self.len, - "split_at_mut mid {mid} out of bounds for {}", - self.len - ); - ( - MetalSliceMut::from_field(self.field.clone(), self.offset, mid), - MetalSliceMut::from_field(self.field.clone(), self.offset + mid, self.len - mid), - ) - } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::algebra::ntt::{Messages, NttEngine, ReedSolomon}; - use crate::algebra::univariate_evaluate; - use crate::buffer::metal::rs::rs_transpose_permute; - use crate::buffer::{Buffer, BufferOps, CpuBuffer, MetalRs}; - - fn values(len: usize, offset: u64) -> Vec { - (0..len) - .map(|i| Field256::from(i as u64 + offset)) - .collect() - } - - #[test] - fn metal_bn254_dot_matches_cpu() { - let a = values(33, 1); - let b = values(33, 9); - let cpu_a = CpuBuffer::from_slice(&a); - let cpu_b = CpuBuffer::from_slice(&b); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_b = MetalBuffer::from_slice(&b); - assert_eq!(gpu_a.dot(&gpu_b), cpu_a.dot(&cpu_b)); - } - - #[test] - fn metal_bn254_fold_matches_cpu() { - let mut cpu = CpuBuffer::from_vec(values(31, 2)); - let mut gpu = MetalBuffer::from_vec(values(31, 2)); - let weight = Field256::from(42); - cpu.fold(weight); - gpu.fold(weight); - assert_eq!(gpu.as_slice(), cpu.as_slice()); - } - - #[test] - fn metal_bn254_sumcheck_matches_cpu() { - for len in [1, 2, 27, 64, 65] { - let a = values(len, 3); - let b = values(len, 11); - let cpu_a = CpuBuffer::from_slice(&a); - let cpu_b = CpuBuffer::from_slice(&b); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_b = MetalBuffer::from_slice(&b); - assert_eq!( - gpu_a.sumcheck_polynomial(&gpu_b), - cpu_a.sumcheck_polynomial(&cpu_b) - ); - } - } - - #[test] - fn metal_bn254_fold_pair_sumcheck_matches_cpu() { - for len in [2, 3, 27, 64, 65] { - let mut cpu_a = CpuBuffer::from_vec(values(len, 3)); - let mut cpu_b = CpuBuffer::from_vec(values(len, 11)); - let mut gpu_a = MetalBuffer::from_vec(values(len, 3)); - let mut gpu_b = MetalBuffer::from_vec(values(len, 11)); - let weight = Field256::from(42); - let cpu_result = cpu_a.fold_pair_sumcheck_polynomial(&mut cpu_b, weight); - let gpu_result = gpu_a.fold_pair_sumcheck_polynomial(&mut gpu_b, weight); - assert_eq!(gpu_result, cpu_result); - assert_eq!(gpu_a.as_slice(), cpu_a.as_slice()); - assert_eq!(gpu_b.as_slice(), cpu_b.as_slice()); - } - } - - #[test] - fn metal_bn254_scalar_mul_add_matches_cpu() { - let mut cpu = CpuBuffer::from_vec(values(19, 1)); - let mut gpu = MetalBuffer::from_vec(values(19, 1)); - let vector = values(19, 5); - let cpu_vector = CpuBuffer::from_slice(&vector); - let gpu_vector = MetalBuffer::from_slice(&vector); - let weight = Field256::from(7); - cpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut cpu, weight); - gpu_vector.mixed_scalar_mul_add_to(&Identity::new(), &mut gpu, weight); - assert_eq!(gpu.as_slice(), cpu.as_slice()); - } - - #[test] - fn metal_bn254_interleaved_rs_encode_matches_cpu() { - // The GPU encoder and the CPU reference derive the same coset layout from the field. - let engine = NttEngine::::new_from_fftfield(); - let gpu_rs = MetalRs::::new_from_fftfield(); - - let a = values(8, 1); - let gpu_a = MetalBuffer::from_slice(&a); - let gpu_masks = MetalBuffer::from_slice(&[]); - - let messages = a.chunks_exact(4).collect::>(); - let generator = engine.root(8); - let mut cpu = Vec::with_capacity(16); - for index in 0..8 { - #[cfg(not(feature = "rs_in_order"))] - let index = rs_transpose_permute(index, 2, 4); - let point = generator.pow([index as u64]); - for message in &messages { - cpu.push(univariate_evaluate(message, point)); - } - } - let gpu_vectors = [&gpu_a]; - let gpu_messages = Messages::new(&gpu_vectors, 4, 2); - let gpu = gpu_rs.interleaved_encode(gpu_messages, &gpu_masks, 8); - assert_eq!(gpu.as_slice(), cpu.as_slice()); - } - - #[test] - fn metal_bn254_mixed_extend_matches_cpu() { - let values = values(8, 3); - let point = vec![Field256::from(2), Field256::from(5), Field256::from(9)]; - let cpu = CpuBuffer::from_slice(&values); - let gpu = MetalBuffer::from_slice(&values); - assert_eq!( - gpu.mixed_extend(&Identity::new(), &point), - cpu.mixed_extend(&Identity::new(), &point) - ); - } -} diff --git a/src/buffer/metal/kernels.rs b/src/buffer/metal/kernels.rs deleted file mode 100644 index d12b4a59..00000000 --- a/src/buffer/metal/kernels.rs +++ /dev/null @@ -1,943 +0,0 @@ -// NOTE: 100% AI GENERATED - -pub(crate) const METAL_SOURCE: &str = r#" -#include -using namespace metal; - -struct F { ulong v[4]; }; - -constant ulong MODULUS[4] = { - 0x43e1f593f0000001UL, - 0x2833e84879b97091UL, - 0xb85045b68181585dUL, - 0x30644e72e131a029UL, -}; - -constant ulong BN254_N0PRIME = 0xc2e1f593efffffffUL; -constant F CANONICAL_ONE = {{1UL, 0UL, 0UL, 0UL}}; -constant F MONT_ONE = {{ - 0xac96341c4ffffffbUL, - 0x36fc76959f60cd29UL, - 0x666ea36f7879462eUL, - 0xe0a77c19a07df2fUL, -}}; - -struct StageConfig { - uint row_len; - uint half_m; - uint twiddle_offset; - uint _pad0; -}; - -struct BitReverseParams { - uint row_len; - uint log_n; - uint total_elements; - uint _pad0; -}; - -struct TransposeParams { - uint rows; - uint cols; - uint total_elements; -}; - -struct ReplicateCosetsParams { - uint row_len; - uint coset_size; - uint trailing_elements; -}; - -struct PackSingleVectorParams { - uint row_count; - uint message_length; - uint codeword_length; - uint coset_size; - uint total_elements; -}; - -struct FieldBytesParams { - uint rows; - uint cols; -}; - -struct ApplyCosetTwiddlesParams { - uint row_count; - uint num_cosets; - uint coset_size; - uint codeword_length; - uint total_elements; -}; - -inline F zero_f() { - F r; - r.v[0] = 0; r.v[1] = 0; r.v[2] = 0; r.v[3] = 0; - return r; -} - -inline F one_f() { - return MONT_ONE; -} - -inline F load_f(device const ulong *data, uint idx) { - F r; - uint base = idx * 4; - r.v[0] = data[base + 0]; - r.v[1] = data[base + 1]; - r.v[2] = data[base + 2]; - r.v[3] = data[base + 3]; - return r; -} - -inline void store_f(device ulong *data, uint idx, F x) { - uint base = idx * 4; - data[base + 0] = x.v[0]; - data[base + 1] = x.v[1]; - data[base + 2] = x.v[2]; - data[base + 3] = x.v[3]; -} - -inline bool ge_mod(F a) { - for (int i = 3; i >= 0; i--) { - if (a.v[i] > MODULUS[i]) return true; - if (a.v[i] < MODULUS[i]) return false; - } - return true; -} - -inline bool ge_f(F a, F b) { - for (int i = 3; i >= 0; i--) { - if (a.v[i] > b.v[i]) return true; - if (a.v[i] < b.v[i]) return false; - } - return true; -} - -inline F sub_raw(F a, thread const ulong *b, thread bool &borrow) { - F r; - borrow = false; - for (uint i = 0; i < 4; i++) { - ulong bi = b[i] + (borrow ? 1UL : 0UL); - bool bcarry = borrow && bi == 0; - r.v[i] = a.v[i] - bi; - borrow = bcarry || a.v[i] < bi; - } - return r; -} - -inline F sub_modulus(F a) { - F r; - bool borrow = false; - for (uint i = 0; i < 4; i++) { - ulong bi = MODULUS[i] + (borrow ? 1UL : 0UL); - bool bcarry = borrow && bi == 0; - r.v[i] = a.v[i] - bi; - borrow = bcarry || a.v[i] < bi; - } - return r; -} - -inline F add_modulus(F a) { - F r; - bool carry = false; - for (uint i = 0; i < 4; i++) { - ulong mi = MODULUS[i]; - ulong sum = a.v[i] + mi; - bool c0 = sum < a.v[i]; - ulong sum2 = sum + (carry ? 1UL : 0UL); - bool c1 = carry && sum2 == 0; - r.v[i] = sum2; - carry = c0 || c1; - } - return r; -} - -inline F add_f(F a, F b) { - F r; - bool carry = false; - for (uint i = 0; i < 4; i++) { - ulong sum = a.v[i] + b.v[i]; - bool c0 = sum < a.v[i]; - ulong sum2 = sum + (carry ? 1UL : 0UL); - bool c1 = carry && sum2 == 0; - r.v[i] = sum2; - carry = c0 || c1; - } - if (carry || ge_mod(r)) { - r = sub_modulus(r); - } - return r; -} - -inline F sub_f(F a, F b) { - ulong bv[4] = { b.v[0], b.v[1], b.v[2], b.v[3] }; - bool borrow = false; - F r = sub_raw(a, bv, borrow); - if (borrow) { - r = add_modulus(r); - } - return r; -} - -inline F double_f(F a) { - return add_f(a, a); -} - -inline ulong add_with_carry(ulong a, ulong b, thread ulong &carry) { - ulong sum = a + b; - ulong c1 = sum < a ? 1UL : 0UL; - ulong sum_with_carry = sum + carry; - ulong c2 = sum_with_carry < sum ? 1UL : 0UL; - carry = c1 + c2; - return sum_with_carry; -} - -inline void add_scaled_step(thread ulong &dst, ulong s, ulong a, thread ulong &carry) { - ulong product_lo = s * a; - ulong product_hi = mulhi(s, a); - - ulong sum = dst + product_lo; - ulong carry0 = sum < dst ? 1UL : 0UL; - ulong sum_with_carry = sum + carry; - ulong carry1 = sum_with_carry < sum ? 1UL : 0UL; - - dst = sum_with_carry; - carry = product_hi + carry0 + carry1; -} - -inline void add_scaled(thread ulong *dst, ulong s, ulong a0, ulong a1, ulong a2, ulong a3) { - ulong carry = 0; - add_scaled_step(dst[0], s, a0, carry); - add_scaled_step(dst[1], s, a1, carry); - add_scaled_step(dst[2], s, a2, carry); - add_scaled_step(dst[3], s, a3, carry); - dst[4] += carry; -} - -inline F mont_mul(F lhs, F rhs) { - ulong buf[9] = {0}; - uint off = 0; - -#pragma clang loop unroll(enable) - for (uint i = 0; i < 4; i++) { - add_scaled(&buf[off], lhs.v[i], rhs.v[0], rhs.v[1], rhs.v[2], rhs.v[3]); - - ulong m = buf[off] * BN254_N0PRIME; - add_scaled( - &buf[off], - m, - MODULUS[0], - MODULUS[1], - MODULUS[2], - MODULUS[3] - ); - - off += 1; - buf[off + 4] = 0; - } - - F result; - result.v[0] = buf[off + 0]; - result.v[1] = buf[off + 1]; - result.v[2] = buf[off + 2]; - result.v[3] = buf[off + 3]; - if (ge_mod(result)) { - result = sub_modulus(result); - } - return result; -} - -inline F from_mont(F value) { - F result = mont_mul(value, CANONICAL_ONE); - if (ge_mod(result)) { - result = sub_modulus(result); - } - return result; -} - -inline F mul_f(F a, F b) { - return mont_mul(a, b); -} - -inline F pow_f(F base, uint exp) { - F acc = one_f(); - F x = base; - while (exp != 0) { - if ((exp & 1) != 0) { - acc = mul_f(acc, x); - } - exp >>= 1; - if (exp != 0) { - x = mul_f(x, x); - } - } - return acc; -} - -kernel void bn254_fold( - device ulong *values [[buffer(0)]], - device const ulong *weight_buf [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &fold_half [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= fold_half) return; - F low = gid < len ? load_f(values, gid) : zero_f(); - F high = gid + fold_half < len ? load_f(values, gid + fold_half) : zero_f(); - F weight = load_f(weight_buf, 0); - store_f(values, gid, add_f(low, mul_f(sub_f(high, low), weight))); -} - -inline F fold_value_at(device const ulong *values, uint len, uint fold_half, uint idx, F weight) { - F low = idx < len ? load_f(values, idx) : zero_f(); - F high = idx + fold_half < len ? load_f(values, idx + fold_half) : zero_f(); - return add_f(low, mul_f(sub_f(high, low), weight)); -} - -kernel void bn254_fold_pair( - device ulong *a [[buffer(0)]], - device ulong *b [[buffer(1)]], - device const ulong *weight_buf [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &fold_half [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= fold_half) return; - F weight = load_f(weight_buf, 0); - store_f(a, gid, fold_value_at(a, len, fold_half, gid, weight)); - store_f(b, gid, fold_value_at(b, len, fold_half, gid, weight)); -} - -kernel void bn254_scalar_mul_add( - device ulong *acc [[buffer(0)]], - device const ulong *vector [[buffer(1)]], - device const ulong *weight_buf [[buffer(2)]], - constant uint &len [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F weight = load_f(weight_buf, 0); - store_f(acc, gid, add_f(load_f(acc, gid), mul_f(weight, load_f(vector, gid)))); -} - -kernel void bn254_scalar_mul( - device ulong *acc [[buffer(0)]], - device const ulong *weight_buf [[buffer(1)]], - constant uint &len [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F weight = load_f(weight_buf, 0); - store_f(acc, gid, mul_f(weight, load_f(acc, gid))); -} - -kernel void bn254_dot( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - F acc = zero_f(); - for (uint i = 0; i < len; i++) { - acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); - } - store_f(out, 0, acc); -} - -kernel void bn254_sumcheck( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &fold_half [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = 0; i < fold_half; i++) { - F a0 = i < len ? load_f(a, i) : zero_f(); - F b0 = i < len ? load_f(b, i) : zero_f(); - F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); - F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); - c0 = add_f(c0, mul_f(a0, b0)); - c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); - } - store_f(out, 0, c0); - store_f(out, 1, c2); -} - -kernel void bn254_dot_chunks( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &chunk_size [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - F acc = zero_f(); - for (uint i = start; i < end; i++) { - acc = add_f(acc, mul_f(load_f(a, i), load_f(b, i))); - } - store_f(out, gid, acc); -} - -kernel void bn254_sum_chunks( - device const ulong *input [[buffer(0)]], - device ulong *out [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &offset [[buffer(3)]], - constant uint &chunk_size [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - F acc = zero_f(); - for (uint i = start; i < end; i++) { - acc = add_f(acc, load_f(input, offset + i)); - } - store_f(out, gid, acc); -} - -kernel void bn254_sumcheck_reduce_chunks( - device const ulong *input [[buffer(0)]], - device ulong *out [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &chunk_size [[buffer(3)]], - constant uint &out_len [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= out_len) return; - uint start = gid * chunk_size; - uint end = min(start + chunk_size, len); - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = start; i < end; i++) { - c0 = add_f(c0, load_f(input, i)); - c2 = add_f(c2, load_f(input, len + i)); - } - store_f(out, gid, c0); - store_f(out, out_len + gid, c2); -} - -kernel void bn254_sumcheck_chunks( - device const ulong *a [[buffer(0)]], - device const ulong *b [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &fold_half [[buffer(4)]], - constant uint &chunk_size [[buffer(5)]], - constant uint &partial_count [[buffer(6)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= fold_half) return; - uint end = min(start + chunk_size, fold_half); - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = start; i < end; i++) { - F a0 = i < len ? load_f(a, i) : zero_f(); - F b0 = i < len ? load_f(b, i) : zero_f(); - F a1 = i + fold_half < len ? load_f(a, i + fold_half) : zero_f(); - F b1 = i + fold_half < len ? load_f(b, i + fold_half) : zero_f(); - c0 = add_f(c0, mul_f(a0, b0)); - c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); - } - store_f(out, gid, c0); - store_f(out, partial_count + gid, c2); -} - -kernel void bn254_fold_pair_sumcheck_chunks( - device ulong *a [[buffer(0)]], - device ulong *b [[buffer(1)]], - device const ulong *weight_buf [[buffer(2)]], - device ulong *out [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &fold_half [[buffer(5)]], - constant uint &sum_half [[buffer(6)]], - constant uint &chunk_size [[buffer(7)]], - constant uint &partial_count [[buffer(8)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= sum_half) return; - uint end = min(start + chunk_size, sum_half); - F weight = load_f(weight_buf, 0); - F c0 = zero_f(); - F c2 = zero_f(); - for (uint i = start; i < end; i++) { - uint right = i + sum_half; - F a0 = fold_value_at(a, len, fold_half, i, weight); - F b0 = fold_value_at(b, len, fold_half, i, weight); - F a1 = right < fold_half ? fold_value_at(a, len, fold_half, right, weight) : zero_f(); - F b1 = right < fold_half ? fold_value_at(b, len, fold_half, right, weight) : zero_f(); - store_f(a, i, a0); - store_f(b, i, b0); - if (right < fold_half) { - store_f(a, right, a1); - store_f(b, right, b1); - } - c0 = add_f(c0, mul_f(a0, b0)); - c2 = add_f(c2, mul_f(sub_f(a1, a0), sub_f(b1, b0))); - } - store_f(out, gid, c0); - store_f(out, partial_count + gid, c2); -} - -kernel void bn254_geometric_accumulate( - device ulong *acc [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *scalars [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &num_points [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F value = load_f(acc, gid); - for (uint j = 0; j < num_points; j++) { - value = add_f(value, mul_f(load_f(scalars, j), pow_f(load_f(points, j), gid))); - } - store_f(acc, gid, value); -} - -kernel void bn254_geometric_accumulate_chunks( - device ulong *acc [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *scalars [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &num_points [[buffer(4)]], - constant uint &chunk_size [[buffer(5)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - for (uint j = 0; j < num_points; j++) { - F point = load_f(points, j); - F scalar = load_f(scalars, j); - F power = pow_f(point, start); - for (uint i = start; i < end; i++) { - F value = load_f(acc, i); - value = add_f(value, mul_f(scalar, power)); - store_f(acc, i, value); - power = mul_f(power, point); - } - } -} - -kernel void bn254_geometric_accumulate_chunks_strided( - device ulong *acc [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *point_steps [[buffer(2)]], - device const ulong *scalars [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &num_points [[buffer(5)]], - constant uint &chunk_size [[buffer(6)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - uint count = end - start; - - // Accumulate all points into registers so `acc` is read and written - // once per element instead of once per (element, point). - F sums[32]; - for (uint k = 0; k < count; k++) { - sums[k] = zero_f(); - } - for (uint j = 0; j < num_points; j++) { - F point = load_f(points, j); - // Fold the scalar into the running power so the inner loop needs a - // single multiplication per element. - F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), gid)); - for (uint k = 0; k < count; k++) { - sums[k] = add_f(sums[k], power); - power = mul_f(power, point); - } - } - for (uint i = start, k = 0; i < end; i++, k++) { - store_f(acc, i, add_f(load_f(acc, i), sums[k])); - } -} - -kernel void bn254_geometric_accumulate_point_blocks( - device ulong *partials [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *point_steps [[buffer(2)]], - device const ulong *scalars [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &num_points [[buffer(5)]], - constant uint &chunk_size [[buffer(6)]], - constant uint &point_block_size [[buffer(7)]], - constant uint &point_blocks [[buffer(8)]], - uint gid [[thread_position_in_grid]] -) { - uint chunk = gid / point_blocks; - uint point_block = gid - chunk * point_blocks; - uint start = chunk * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - uint point_start = point_block * point_block_size; - if (point_start >= num_points) return; - uint point_end = min(point_start + point_block_size, num_points); - - F sums[32]; - for (uint k = 0; k < chunk_size; k++) { - sums[k] = zero_f(); - } - - for (uint j = point_start; j < point_end; j++) { - F point = load_f(points, j); - F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); - for (uint i = start, k = 0; i < end; i++, k++) { - sums[k] = add_f(sums[k], power); - power = mul_f(power, point); - } - } - - uint partial_offset = point_block * len; - for (uint i = start, k = 0; i < end; i++, k++) { - store_f(partials, partial_offset + i, sums[k]); - } -} - -kernel void bn254_geometric_accumulate_point_blocks_range( - device ulong *partials [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device const ulong *point_steps [[buffer(2)]], - device const ulong *scalars [[buffer(3)]], - constant uint &len [[buffer(4)]], - constant uint &num_points [[buffer(5)]], - constant uint &chunk_size [[buffer(6)]], - constant uint &point_block_size [[buffer(7)]], - constant uint &point_block_offset [[buffer(8)]], - constant uint &batch_point_blocks [[buffer(9)]], - uint gid [[thread_position_in_grid]] -) { - uint chunk = gid / batch_point_blocks; - uint local_point_block = gid - chunk * batch_point_blocks; - uint global_point_block = point_block_offset + local_point_block; - uint start = chunk * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - uint point_start = global_point_block * point_block_size; - if (point_start >= num_points) return; - uint point_end = min(point_start + point_block_size, num_points); - - F sums[32]; - for (uint k = 0; k < chunk_size; k++) { - sums[k] = zero_f(); - } - - for (uint j = point_start; j < point_end; j++) { - F point = load_f(points, j); - F power = mul_f(load_f(scalars, j), pow_f(load_f(point_steps, j), chunk)); - for (uint i = start, k = 0; i < end; i++, k++) { - sums[k] = add_f(sums[k], power); - power = mul_f(power, point); - } - } - - uint partial_offset = local_point_block * len; - for (uint i = start, k = 0; i < end; i++, k++) { - store_f(partials, partial_offset + i, sums[k]); - } -} - -kernel void bn254_geometric_accumulate_reduce_point_blocks( - device ulong *acc [[buffer(0)]], - device const ulong *partials [[buffer(1)]], - constant uint &len [[buffer(2)]], - constant uint &point_blocks [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= len) return; - F value = load_f(acc, gid); - for (uint block = 0; block < point_blocks; block++) { - value = add_f(value, load_f(partials, block * len + gid)); - } - store_f(acc, gid, value); -} - -kernel void bn254_univariate_evaluate( - device const ulong *coeffs [[buffer(0)]], - device const ulong *point_buf [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - if (len == 0) { - store_f(out, 0, zero_f()); - return; - } - F point = load_f(point_buf, 0); - F acc = load_f(coeffs, len - 1); - for (uint i = len - 1; i > 0; i--) { - acc = add_f(mul_f(acc, point), load_f(coeffs, i - 1)); - } - store_f(out, 0, acc); -} - -kernel void bn254_univariate_eval_chunks( - device const ulong *coeffs [[buffer(0)]], - device const ulong *point_buf [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &chunk_size [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - uint start = gid * chunk_size; - if (start >= len) return; - uint end = min(start + chunk_size, len); - F point = load_f(point_buf, 0); - F power = pow_f(point, start); - F acc = zero_f(); - for (uint i = start; i < end; i++) { - acc = add_f(acc, mul_f(load_f(coeffs, i), power)); - power = mul_f(power, point); - } - store_f(out, gid, acc); -} - -kernel void bn254_interleaved_rs_encode( - device const ulong *coeffs [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &num_messages [[buffer(3)]], - constant uint &coeff_len [[buffer(4)]], - constant uint &total [[buffer(5)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= total) return; - uint row = gid / num_messages; - uint message = gid - row * num_messages; - uint base = message * coeff_len; - F point = load_f(points, row); - F acc = load_f(coeffs, base + coeff_len - 1); - for (uint i = coeff_len - 1; i > 0; i--) { - acc = add_f(mul_f(acc, point), load_f(coeffs, base + i - 1)); - } - store_f(out, gid, acc); -} - -kernel void bn254_interleaved_rs_encode_single_vector( - device const ulong *vector [[buffer(0)]], - device const ulong *points [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &interleaving_depth [[buffer(3)]], - constant uint &message_length [[buffer(4)]], - constant uint &total [[buffer(5)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= total) return; - uint row = gid / interleaving_depth; - uint message = gid - row * interleaving_depth; - uint base = message * message_length; - F point = load_f(points, row); - F acc = load_f(vector, base + message_length - 1); - for (uint i = message_length - 1; i > 0; i--) { - acc = add_f(mul_f(acc, point), load_f(vector, base + i - 1)); - } - store_f(out, gid, acc); -} - -inline uint reverse_bits_width(uint value, uint width) { - return reverse_bits(value) >> (32u - width); -} - -kernel void bn254_pack_single_vector_cosets( - device const ulong *vector [[buffer(0)]], - device ulong *out [[buffer(1)]], - constant PackSingleVectorParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint row = gid / params.codeword_length; - uint col = gid - row * params.codeword_length; - if (col < params.message_length) { - store_f(out, gid, load_f(vector, row * params.message_length + col)); - } else { - store_f(out, gid, zero_f()); - } -} - -kernel void bn254_replicate_first_coset( - device ulong *buffer [[buffer(0)]], - constant ReplicateCosetsParams ¶ms [[buffer(1)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.trailing_elements) return; - uint repeats_per_row = params.row_len - params.coset_size; - uint row = gid / repeats_per_row; - uint within = gid - row * repeats_per_row; - uint dst = row * params.row_len + params.coset_size + within; - uint src = row * params.row_len + (within % params.coset_size); - store_f(buffer, dst, load_f(buffer, src)); -} - -kernel void bn254_bit_reverse_permute_rows_in_place( - device ulong *values [[buffer(0)]], - constant BitReverseParams &config [[buffer(1)]], - uint index [[thread_position_in_grid]] -) { - if (index >= config.total_elements || config.row_len <= 1u) return; - uint row = index / config.row_len; - uint within = index - row * config.row_len; - uint reversed = reverse_bits_width(within, config.log_n); - if (reversed <= within) return; - - uint row_base = row * config.row_len; - uint mate = row_base + reversed; - uint current = row_base + within; - F tmp = load_f(values, current); - store_f(values, current, load_f(values, mate)); - store_f(values, mate, tmp); -} - -kernel void bn254_radix2_ntt_stage_rows_in_place( - device ulong *values [[buffer(0)]], - device const ulong *twiddles [[buffer(1)]], - constant StageConfig &config [[buffer(2)]], - uint index [[thread_position_in_grid]] -) { - uint butterflies_per_row = config.row_len >> 1u; - uint row = index / butterflies_per_row; - uint local = index - row * butterflies_per_row; - uint half_m = config.half_m; - uint pair_in_group = local % half_m; - uint group = local / half_m; - uint row_base = row * config.row_len; - uint base = row_base + group * (half_m << 1u) + pair_in_group; - uint mate = base + half_m; - - F even = load_f(values, base); - F odd = load_f(values, mate); - F twiddle = load_f(twiddles, config.twiddle_offset + pair_in_group); - F t = mul_f(twiddle, odd); - - store_f(values, base, add_f(even, t)); - store_f(values, mate, sub_f(even, t)); -} - -kernel void bn254_transpose_matrix_reverse_rows( - device const ulong *input [[buffer(0)]], - device ulong *output [[buffer(1)]], - constant TransposeParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint row = gid / params.cols; - uint col = gid - row * params.cols; - uint row_bits = 31u - clz(params.cols); - uint dst_row = reverse_bits_width(col, row_bits); - uint dst = dst_row * params.rows + row; - store_f(output, dst, load_f(input, gid)); -} - -kernel void bn254_transpose_matrix( - device const ulong *input [[buffer(0)]], - device ulong *output [[buffer(1)]], - constant TransposeParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint row = gid / params.cols; - uint col = gid - row * params.cols; - uint dst = col * params.rows + row; - store_f(output, dst, load_f(input, gid)); -} - -kernel void bn254_apply_coset_twiddles( - device ulong *values [[buffer(0)]], - device const ulong *root_powers [[buffer(1)]], - constant ApplyCosetTwiddlesParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= params.total_elements) return; - uint within_codeword = gid % params.codeword_length; - uint coset = within_codeword / params.coset_size; - uint col = within_codeword - coset * params.coset_size; - if (coset == 0 || col == 0) return; - uint root_index = (coset * col) % params.codeword_length; - store_f(values, gid, mul_f(load_f(values, gid), load_f(root_powers, root_index))); -} - -kernel void bn254_encode_field_rows_le( - device const ulong *input [[buffer(0)]], - device uchar *output [[buffer(1)]], - constant FieldBytesParams ¶ms [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - uint total_elements = params.rows * params.cols; - if (gid >= total_elements) return; - - F canonical = from_mont(load_f(input, gid)); - uint byte_offset = gid * 32u; - for (uint limb = 0; limb < 4; ++limb) { - ulong value = canonical.v[limb]; - for (uint byte = 0; byte < 8; ++byte) { - output[byte_offset + limb * 8u + byte] = uchar((value >> (byte * 8u)) & 0xffUL); - } - } -} - -kernel void bn254_read_rows( - device const ulong *input [[buffer(0)]], - device const uint *indices [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &num_cols [[buffer(3)]], - constant uint &total [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= total) return; - uint row = gid / num_cols; - uint col = gid - row * num_cols; - uint src = indices[row] * num_cols + col; - store_f(out, gid, load_f(input, src)); -} - -kernel void bn254_multilinear_extend( - device const ulong *values [[buffer(0)]], - device const ulong *point [[buffer(1)]], - device ulong *out [[buffer(2)]], - constant uint &len [[buffer(3)]], - constant uint &num_vars [[buffer(4)]], - uint gid [[thread_position_in_grid]] -) { - if (gid != 0) return; - F acc = zero_f(); - F one = one_f(); - for (uint i = 0; i < len; i++) { - F weight = one; - for (uint j = 0; j < num_vars; j++) { - F r = load_f(point, num_vars - 1 - j); - if (((i >> j) & 1) != 0) { - weight = mul_f(weight, r); - } else { - weight = mul_f(weight, sub_f(one, r)); - } - } - acc = add_f(acc, mul_f(load_f(values, i), weight)); - } - store_f(out, 0, acc); -} -"#; - -pub(crate) const REDUCTION_CHUNK_SIZE: usize = 64; -pub(crate) const SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 8; -pub(crate) const LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; -/// Must match the `sums` register array size in the strided kernel. -pub(crate) const MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE: usize = 32; -pub(crate) const LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD: usize = 1 << 18; -pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE: usize = 16; -pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD: usize = 1 << 17; -pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS: usize = 64; -pub(crate) const GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES: usize = 256 << 20; diff --git a/src/buffer/metal/mod.rs b/src/buffer/metal/mod.rs deleted file mode 100644 index a796bf9e..00000000 --- a/src/buffer/metal/mod.rs +++ /dev/null @@ -1,16 +0,0 @@ -// NOTE: 100% AI GENERATED - -mod buffer; -mod kernels; -mod profile; -mod rs; -mod runtime; -mod sha2; - -pub use buffer::{MetalBuffer, MetalSlice, MetalSliceMut}; -pub use profile::{ - reset_device_peak as metal_reset_device_peak, snapshot as metal_profile_snapshot, - MetalProfileSnapshot, -}; -pub use rs::MetalRs; -pub use sha2::MetalSha2; diff --git a/src/buffer/metal/profile.rs b/src/buffer/metal/profile.rs deleted file mode 100644 index cebd654d..00000000 --- a/src/buffer/metal/profile.rs +++ /dev/null @@ -1,153 +0,0 @@ -// NOTE: 100% AI GENERATED - -use std::{ - sync::atomic::{AtomicU64, Ordering}, - time::Duration, -}; - -#[derive(Clone, Copy, Debug, Default, serde::Serialize, serde::Deserialize)] -pub struct MetalProfileSnapshot { - pub upload_count: u64, - pub upload_bytes: u64, - pub upload_nanos: u64, - pub readback_count: u64, - pub readback_bytes: u64, - pub readback_nanos: u64, - pub alloc_count: u64, - pub alloc_bytes: u64, - pub command_count: u64, - pub command_wait_nanos: u64, - pub blit_count: u64, - pub blit_bytes: u64, - pub blit_wait_nanos: u64, - /// Device-wide allocated GPU memory at snapshot time. - pub device_current_bytes: u64, - /// Peak device-wide allocated GPU memory since the last `reset_device_peak`. - pub device_peak_bytes: u64, -} - -impl MetalProfileSnapshot { - pub fn delta_since(self, before: Self) -> Self { - Self { - upload_count: self.upload_count.saturating_sub(before.upload_count), - upload_bytes: self.upload_bytes.saturating_sub(before.upload_bytes), - upload_nanos: self.upload_nanos.saturating_sub(before.upload_nanos), - readback_count: self.readback_count.saturating_sub(before.readback_count), - readback_bytes: self.readback_bytes.saturating_sub(before.readback_bytes), - readback_nanos: self.readback_nanos.saturating_sub(before.readback_nanos), - alloc_count: self.alloc_count.saturating_sub(before.alloc_count), - alloc_bytes: self.alloc_bytes.saturating_sub(before.alloc_bytes), - command_count: self.command_count.saturating_sub(before.command_count), - command_wait_nanos: self - .command_wait_nanos - .saturating_sub(before.command_wait_nanos), - blit_count: self.blit_count.saturating_sub(before.blit_count), - blit_bytes: self.blit_bytes.saturating_sub(before.blit_bytes), - blit_wait_nanos: self.blit_wait_nanos.saturating_sub(before.blit_wait_nanos), - // Gauges, not counters: keep the latest values. - device_current_bytes: self.device_current_bytes, - device_peak_bytes: self.device_peak_bytes, - } - } - - pub fn upload_ms(self) -> f64 { - nanos_to_ms(self.upload_nanos) - } - - pub fn readback_ms(self) -> f64 { - nanos_to_ms(self.readback_nanos) - } - - pub fn command_wait_ms(self) -> f64 { - nanos_to_ms(self.command_wait_nanos) - } - - pub fn blit_wait_ms(self) -> f64 { - nanos_to_ms(self.blit_wait_nanos) - } -} - -static UPLOAD_COUNT: AtomicU64 = AtomicU64::new(0); -static UPLOAD_BYTES: AtomicU64 = AtomicU64::new(0); -static UPLOAD_NANOS: AtomicU64 = AtomicU64::new(0); -static READBACK_COUNT: AtomicU64 = AtomicU64::new(0); -static READBACK_BYTES: AtomicU64 = AtomicU64::new(0); -static READBACK_NANOS: AtomicU64 = AtomicU64::new(0); -static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0); -static ALLOC_BYTES: AtomicU64 = AtomicU64::new(0); -static COMMAND_COUNT: AtomicU64 = AtomicU64::new(0); -static COMMAND_WAIT_NANOS: AtomicU64 = AtomicU64::new(0); -static BLIT_COUNT: AtomicU64 = AtomicU64::new(0); -static BLIT_BYTES: AtomicU64 = AtomicU64::new(0); -static BLIT_WAIT_NANOS: AtomicU64 = AtomicU64::new(0); -static DEVICE_CURRENT_BYTES: AtomicU64 = AtomicU64::new(0); -static DEVICE_PEAK_BYTES: AtomicU64 = AtomicU64::new(0); - -pub fn snapshot() -> MetalProfileSnapshot { - MetalProfileSnapshot { - upload_count: UPLOAD_COUNT.load(Ordering::Relaxed), - upload_bytes: UPLOAD_BYTES.load(Ordering::Relaxed), - upload_nanos: UPLOAD_NANOS.load(Ordering::Relaxed), - readback_count: READBACK_COUNT.load(Ordering::Relaxed), - readback_bytes: READBACK_BYTES.load(Ordering::Relaxed), - readback_nanos: READBACK_NANOS.load(Ordering::Relaxed), - alloc_count: ALLOC_COUNT.load(Ordering::Relaxed), - alloc_bytes: ALLOC_BYTES.load(Ordering::Relaxed), - command_count: COMMAND_COUNT.load(Ordering::Relaxed), - command_wait_nanos: COMMAND_WAIT_NANOS.load(Ordering::Relaxed), - blit_count: BLIT_COUNT.load(Ordering::Relaxed), - blit_bytes: BLIT_BYTES.load(Ordering::Relaxed), - blit_wait_nanos: BLIT_WAIT_NANOS.load(Ordering::Relaxed), - device_current_bytes: DEVICE_CURRENT_BYTES.load(Ordering::Relaxed), - device_peak_bytes: DEVICE_PEAK_BYTES.load(Ordering::Relaxed), - } -} - -/// Record the device-wide allocated size (`MTLDevice.currentAllocatedSize`), -/// sampled after each buffer allocation. Peak live memory always coincides -/// with an allocation, so per-alloc sampling captures the true peak. -pub fn record_device_allocated(bytes: u64) { - DEVICE_CURRENT_BYTES.store(bytes, Ordering::Relaxed); - DEVICE_PEAK_BYTES.fetch_max(bytes, Ordering::Relaxed); -} - -/// Reset the peak gauge to the current value (e.g. at the start of a -/// measured phase). -pub fn reset_device_peak() { - DEVICE_PEAK_BYTES.store( - DEVICE_CURRENT_BYTES.load(Ordering::Relaxed), - Ordering::Relaxed, - ); -} - -pub fn record_upload(bytes: u64, duration: Duration) { - UPLOAD_COUNT.fetch_add(1, Ordering::Relaxed); - UPLOAD_BYTES.fetch_add(bytes, Ordering::Relaxed); - UPLOAD_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); -} - -pub fn record_readback(bytes: u64, duration: Duration) { - READBACK_COUNT.fetch_add(1, Ordering::Relaxed); - READBACK_BYTES.fetch_add(bytes, Ordering::Relaxed); - READBACK_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); -} - -pub fn record_alloc(bytes: u64) { - ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); - ALLOC_BYTES.fetch_add(bytes, Ordering::Relaxed); -} - -pub fn record_command_wait(duration: Duration) { - COMMAND_COUNT.fetch_add(1, Ordering::Relaxed); - COMMAND_WAIT_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); -} - -pub fn record_blit(bytes: u64, duration: Duration) { - BLIT_COUNT.fetch_add(1, Ordering::Relaxed); - BLIT_BYTES.fetch_add(bytes, Ordering::Relaxed); - BLIT_WAIT_NANOS.fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); -} - -fn nanos_to_ms(nanos: u64) -> f64 { - nanos as f64 / 1_000_000.0 -} diff --git a/src/buffer/metal/rs.rs b/src/buffer/metal/rs.rs deleted file mode 100644 index 7a1dc4d5..00000000 --- a/src/buffer/metal/rs.rs +++ /dev/null @@ -1,173 +0,0 @@ -// NOTE: 100% AI GENERATED - -use std::any::type_name; - -use ark_ff::{FftField, Field}; - -use crate::{ - algebra::{ - fields::Field256, - ntt::{Messages, NttEngine, ReedSolomon}, - }, - buffer::{BufferOps, MetalBuffer}, -}; - -use super::runtime::{assert_bn254, encode_single_vector_coset_ntt}; - -const RS_PRIMES: [usize; 2] = [2, 3]; - -fn rs_divisors(n: usize) -> Vec { - let mut result = vec![1usize]; - let mut remaining = n; - for &p in &RS_PRIMES { - let mut pk = 1usize; - let existing = result.clone(); - while remaining.is_multiple_of(p) { - pk *= p; - remaining /= p; - result.extend(existing.iter().map(|d| d * pk)); - } - } - assert_eq!(remaining, 1); - result.sort_unstable(); - result -} - -#[cfg(not(feature = "rs_in_order"))] -pub(crate) fn rs_transpose_permute(index: usize, rows: usize, cols: usize) -> usize { - debug_assert!(index < rows * cols); - let (row, col) = (index / cols, index % cols); - row + col * rows -} -/// Metal (GPU) Reed-Solomon encoder. -/// -/// Keeps the scalar Reed-Solomon parameters directly and runs encoding on Metal buffers. -#[derive(Debug, Clone)] -pub struct MetalRs { - order: usize, - divisors: Vec, - omega_order: F, -} - -impl MetalRs { - pub fn new(order: usize, omega_order: F) -> Self { - assert_eq!(omega_order.pow([order as u64]), F::ONE); - for prime in RS_PRIMES { - if order.is_multiple_of(prime) { - assert_ne!(omega_order.pow([(order / prime) as u64]), F::ONE); - } - } - Self { - order, - divisors: rs_divisors(order), - omega_order, - } - } -} - -impl MetalRs { - pub fn new_from_fftfield() -> Self { - let (mut omega, mut order) = if let (Some(mut omega), Some(b), Some(k)) = ( - F::LARGE_SUBGROUP_ROOT_OF_UNITY, - F::SMALL_SUBGROUP_BASE, - F::SMALL_SUBGROUP_BASE_ADICITY, - ) { - let mut order = 1; - let mut remaining = (b as usize).checked_pow(k).expect("Small group too large."); - for p in RS_PRIMES { - while remaining.is_multiple_of(p) { - order *= p; - remaining /= p; - } - } - omega = omega.pow([remaining as u64]); - (omega, order) - } else { - (F::TWO_ADIC_ROOT_OF_UNITY, 1) - }; - let twos = F::TWO_ADICITY.min(order.leading_zeros()) as usize; - for _ in 0..(F::TWO_ADICITY as usize - twos) { - omega.square_in_place(); - } - order <<= twos; - Self::new(order, omega) - } -} - -impl ReedSolomon for MetalRs { - fn next_order(&self, size: usize) -> Option { - match self.divisors.binary_search(&size) { - Ok(index) | Err(index) => self.divisors.get(index).copied(), - } - } - - fn generator(&self, codeword_length: usize) -> F { - self.omega_order - .pow([(self.order / codeword_length) as u64]) - } - - fn evaluation_points( - &self, - masked_message_length: usize, - codeword_length: usize, - indices: &[usize], - ) -> Vec { - assert!(masked_message_length <= codeword_length); - assert!(self.order.is_multiple_of(codeword_length)); - let mut result = Vec::new(); - let generator = self.generator(codeword_length); - - let mut coset_size = self.next_order(masked_message_length).unwrap(); - while !codeword_length.is_multiple_of(coset_size) { - coset_size = self.next_order(coset_size + 1).unwrap(); - } - let num_cosets = codeword_length / coset_size; - #[cfg(feature = "rs_in_order")] - let _ = (coset_size, num_cosets); - - for &index in indices { - assert!(index < codeword_length); - - #[cfg(not(feature = "rs_in_order"))] - let index = rs_transpose_permute(index, num_cosets, coset_size); - result.push(generator.pow([index as u64])); - } - result - } - - fn interleaved_encode( - &self, - messages: Messages<'_, F>, - masks: &MetalBuffer, - codeword_length: usize, - ) -> MetalBuffer { - let vectors = messages.vectors; - let num_messages = vectors.len() * messages.interleaving_depth; - if num_messages == 0 { - return BufferOps::from_vec(Vec::new()); - } - assert!(masks.len().is_multiple_of(num_messages)); - let mask_length = masks.len() / num_messages; - - if type_name::() == type_name::() && vectors.len() == 1 && mask_length == 0 { - assert_bn254::(); - let mut coset_size = self.next_order(messages.message_length).unwrap(); - while !codeword_length.is_multiple_of(coset_size) { - coset_size = self.next_order(coset_size + 1).unwrap(); - } - return encode_single_vector_coset_ntt( - vectors[0], - messages.message_length, - messages.interleaving_depth, - codeword_length, - coset_size, - ); - } - - NttEngine::new(self.order, self.omega_order).interleaved_encode( - messages, - masks, - codeword_length, - ) - } -} diff --git a/src/buffer/metal/runtime.rs b/src/buffer/metal/runtime.rs deleted file mode 100644 index 0fc7265a..00000000 --- a/src/buffer/metal/runtime.rs +++ /dev/null @@ -1,1218 +0,0 @@ -// NOTE: 100% AI GENERATED - -use std::{ - any::type_name, - cell::OnceCell, - collections::HashMap, - marker::PhantomData, - os::raw::c_void, - sync::{Mutex, OnceLock}, - time::Instant, -}; - -use ark_ff::{AdditiveGroup, BigInt, FftField, Field, Fp, MontBackend}; -use metal::{ - Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, MTLResourceOptions, MTLSize, -}; -use zerocopy::IntoBytes; - -use crate::{ - algebra::fields::{BN254Config, Field256}, - buffer::{BufferOps, MetalBuffer}, - hash::Hash as Digest, -}; - -use super::buffer::MetalFieldBuffer; -use super::kernels::{ - GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES, GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS, - GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE, GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD, - LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE, LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD, - MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE, METAL_SOURCE, REDUCTION_CHUNK_SIZE, - SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE, -}; -use super::profile; - -pub(crate) struct MetalRuntime { - device: Device, - queue: CommandQueue, - fold: ComputePipelineState, - fold_pair: ComputePipelineState, - scalar_mul_add: ComputePipelineState, - scalar_mul: ComputePipelineState, - dot: ComputePipelineState, - sumcheck: ComputePipelineState, - dot_chunks: ComputePipelineState, - sum_chunks: ComputePipelineState, - sumcheck_reduce_chunks: ComputePipelineState, - sumcheck_chunks: ComputePipelineState, - fold_pair_sumcheck_chunks: ComputePipelineState, - geometric_accumulate: ComputePipelineState, - geometric_accumulate_chunks: ComputePipelineState, - geometric_accumulate_chunks_strided: ComputePipelineState, - geometric_accumulate_point_blocks: ComputePipelineState, - geometric_accumulate_point_blocks_range: ComputePipelineState, - geometric_accumulate_reduce_point_blocks: ComputePipelineState, - univariate_evaluate: ComputePipelineState, - univariate_eval_chunks: ComputePipelineState, - interleaved_rs_encode: ComputePipelineState, - interleaved_rs_encode_single_vector: ComputePipelineState, - multilinear_extend: ComputePipelineState, - pack_single_vector_cosets: ComputePipelineState, - apply_coset_twiddles: ComputePipelineState, - replicate_first_coset: ComputePipelineState, - bit_reverse_rows: ComputePipelineState, - ntt_stage_rows: ComputePipelineState, - transpose: ComputePipelineState, - transpose_reverse_rows: ComputePipelineState, - encode_field_rows_le: ComputePipelineState, - read_rows: ComputePipelineState, - ntt_roots: Mutex>, - root_powers: Mutex>, -} - -pub fn init() { - let _ = runtime(); -} - -pub(crate) fn runtime() -> &'static MetalRuntime { - static RUNTIME: OnceLock = OnceLock::new(); - RUNTIME.get_or_init(|| { - let device = Device::system_default().expect("Metal device is not available"); - let library = device - .new_library_with_source(METAL_SOURCE, &CompileOptions::new()) - .expect("failed to compile Metal BN254 kernels"); - let pipeline = |name: &str| { - let function = library - .get_function(name, None) - .unwrap_or_else(|_| panic!("missing Metal kernel {name}")); - device - .new_compute_pipeline_state_with_function(&function) - .unwrap_or_else(|err| panic!("failed to compile Metal kernel {name}: {err}")) - }; - MetalRuntime { - queue: device.new_command_queue(), - fold: pipeline("bn254_fold"), - fold_pair: pipeline("bn254_fold_pair"), - scalar_mul_add: pipeline("bn254_scalar_mul_add"), - scalar_mul: pipeline("bn254_scalar_mul"), - dot: pipeline("bn254_dot"), - sumcheck: pipeline("bn254_sumcheck"), - dot_chunks: pipeline("bn254_dot_chunks"), - sum_chunks: pipeline("bn254_sum_chunks"), - sumcheck_reduce_chunks: pipeline("bn254_sumcheck_reduce_chunks"), - sumcheck_chunks: pipeline("bn254_sumcheck_chunks"), - fold_pair_sumcheck_chunks: pipeline("bn254_fold_pair_sumcheck_chunks"), - geometric_accumulate: pipeline("bn254_geometric_accumulate"), - geometric_accumulate_chunks: pipeline("bn254_geometric_accumulate_chunks"), - geometric_accumulate_chunks_strided: pipeline( - "bn254_geometric_accumulate_chunks_strided", - ), - geometric_accumulate_point_blocks: pipeline("bn254_geometric_accumulate_point_blocks"), - geometric_accumulate_point_blocks_range: pipeline( - "bn254_geometric_accumulate_point_blocks_range", - ), - geometric_accumulate_reduce_point_blocks: pipeline( - "bn254_geometric_accumulate_reduce_point_blocks", - ), - univariate_evaluate: pipeline("bn254_univariate_evaluate"), - univariate_eval_chunks: pipeline("bn254_univariate_eval_chunks"), - interleaved_rs_encode: pipeline("bn254_interleaved_rs_encode"), - interleaved_rs_encode_single_vector: pipeline( - "bn254_interleaved_rs_encode_single_vector", - ), - multilinear_extend: pipeline("bn254_multilinear_extend"), - pack_single_vector_cosets: pipeline("bn254_pack_single_vector_cosets"), - apply_coset_twiddles: pipeline("bn254_apply_coset_twiddles"), - replicate_first_coset: pipeline("bn254_replicate_first_coset"), - bit_reverse_rows: pipeline("bn254_bit_reverse_permute_rows_in_place"), - ntt_stage_rows: pipeline("bn254_radix2_ntt_stage_rows_in_place"), - transpose: pipeline("bn254_transpose_matrix"), - transpose_reverse_rows: pipeline("bn254_transpose_matrix_reverse_rows"), - encode_field_rows_le: pipeline("bn254_encode_field_rows_le"), - read_rows: pipeline("bn254_read_rows"), - ntt_roots: Mutex::new(HashMap::new()), - root_powers: Mutex::new(HashMap::new()), - device, - } - }) -} - -pub(crate) fn pipeline<'a>(rt: &'a MetalRuntime, name: &str) -> &'a ComputePipelineState { - match name { - "bn254_fold" => &rt.fold, - "bn254_fold_pair" => &rt.fold_pair, - "bn254_scalar_mul_add" => &rt.scalar_mul_add, - "bn254_scalar_mul" => &rt.scalar_mul, - "bn254_dot" => &rt.dot, - "bn254_sumcheck" => &rt.sumcheck, - "bn254_dot_chunks" => &rt.dot_chunks, - "bn254_sum_chunks" => &rt.sum_chunks, - "bn254_sumcheck_reduce_chunks" => &rt.sumcheck_reduce_chunks, - "bn254_sumcheck_chunks" => &rt.sumcheck_chunks, - "bn254_fold_pair_sumcheck_chunks" => &rt.fold_pair_sumcheck_chunks, - "bn254_geometric_accumulate" => &rt.geometric_accumulate, - "bn254_geometric_accumulate_chunks" => &rt.geometric_accumulate_chunks, - "bn254_geometric_accumulate_chunks_strided" => &rt.geometric_accumulate_chunks_strided, - "bn254_geometric_accumulate_point_blocks" => &rt.geometric_accumulate_point_blocks, - "bn254_geometric_accumulate_point_blocks_range" => { - &rt.geometric_accumulate_point_blocks_range - } - "bn254_geometric_accumulate_reduce_point_blocks" => { - &rt.geometric_accumulate_reduce_point_blocks - } - "bn254_univariate_evaluate" => &rt.univariate_evaluate, - "bn254_univariate_eval_chunks" => &rt.univariate_eval_chunks, - "bn254_interleaved_rs_encode" => &rt.interleaved_rs_encode, - "bn254_interleaved_rs_encode_single_vector" => &rt.interleaved_rs_encode_single_vector, - "bn254_multilinear_extend" => &rt.multilinear_extend, - "bn254_pack_single_vector_cosets" => &rt.pack_single_vector_cosets, - "bn254_apply_coset_twiddles" => &rt.apply_coset_twiddles, - "bn254_replicate_first_coset" => &rt.replicate_first_coset, - "bn254_bit_reverse_permute_rows_in_place" => &rt.bit_reverse_rows, - "bn254_radix2_ntt_stage_rows_in_place" => &rt.ntt_stage_rows, - "bn254_transpose_matrix" => &rt.transpose, - "bn254_transpose_matrix_reverse_rows" => &rt.transpose_reverse_rows, - "bn254_encode_field_rows_le" => &rt.encode_field_rows_le, - "bn254_read_rows" => &rt.read_rows, - _ => panic!("unknown Metal kernel {name}"), - } -} - -pub(crate) fn new_shared_buffer(rt: &MetalRuntime, bytes: u64) -> Buffer { - profile::record_alloc(bytes); - let buffer = rt - .device - .new_buffer(bytes, MTLResourceOptions::StorageModeShared); - profile::record_device_allocated(rt.device.current_allocated_size()); - buffer -} - -pub(crate) fn new_shared_buffer_with_data( - rt: &MetalRuntime, - data: *const c_void, - bytes: u64, -) -> Buffer { - let start = Instant::now(); - let buffer = rt - .device - .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); - profile::record_alloc(bytes); - profile::record_upload(bytes, start.elapsed()); - profile::record_device_allocated(rt.device.current_allocated_size()); - buffer -} - -pub(crate) fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { - command.commit(); - let start = Instant::now(); - command.wait_until_completed(); - let elapsed = start.elapsed(); - if std::env::var_os("WHIR_METAL_TRACE").is_some() { - eprintln!( - "metal command {label} {:.3} ms", - elapsed.as_secs_f64() * 1_000.0 - ); - } - profile::record_command_wait(elapsed); -} - -pub(crate) fn wait_for_blit(command: &metal::CommandBufferRef, bytes: u64) { - command.commit(); - let start = Instant::now(); - command.wait_until_completed(); - profile::record_blit(bytes, start.elapsed()); -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct BitReverseParams { - row_len: u32, - log_n: u32, - total_elements: u32, - _pad0: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct NttStageParams { - row_len: u32, - half_m: u32, - twiddle_offset: u32, - _pad0: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct TransposeParams { - rows: u32, - cols: u32, - total_elements: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct ReplicateCosetsParams { - row_len: u32, - coset_size: u32, - trailing_elements: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct PackSingleVectorParams { - row_count: u32, - message_length: u32, - codeword_length: u32, - coset_size: u32, - total_elements: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct FieldBytesParams { - rows: u32, - cols: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct ApplyCosetTwiddlesParams { - row_count: u32, - num_cosets: u32, - coset_size: u32, - codeword_length: u32, - total_elements: u32, -} - -pub(crate) fn parallel_dot(a: &MetalFieldBuffer, b: &MetalFieldBuffer, len: usize) -> Field256 { - parallel_dot_at(a, 0, b, 0, len) -} - -/// Inner product of `a[a_off..a_off+len]` and `b[b_off..b_off+len]`. -pub(crate) fn parallel_dot_at( - a: &MetalFieldBuffer, - a_off: usize, - b: &MetalFieldBuffer, - b_off: usize, - len: usize, -) -> Field256 { - if len == 0 { - return Field256::ZERO; - } - let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_dot_chunks"), - &[&a.limbs, &b.limbs, &partials.limbs], - &[field_byte_offset(a_off), field_byte_offset(b_off), 0], - &[len as u32, REDUCTION_CHUNK_SIZE as u32], - partial_count, - ); - let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); - wait_for_command_named(command, "bn254_dot_chunks"); - download_field_at(&result.limbs, offset) -} - -pub(crate) fn parallel_sumcheck( - a: &MetalFieldBuffer, - b: &MetalFieldBuffer, - len: usize, - fold_half: usize, -) -> (Field256, Field256) { - parallel_sumcheck_at(a, 0, b, 0, len, fold_half) -} - -/// Sumcheck `(c0, c2)` over `a[a_off..]` and `b[b_off..]` (logical length `len`). -pub(crate) fn parallel_sumcheck_at( - a: &MetalFieldBuffer, - a_off: usize, - b: &MetalFieldBuffer, - b_off: usize, - len: usize, - fold_half: usize, -) -> (Field256, Field256) { - let partial_count = fold_half.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_sumcheck_chunks"), - &[&a.limbs, &b.limbs, &partials.limbs], - &[field_byte_offset(a_off), field_byte_offset(b_off), 0], - &[ - len as u32, - fold_half as u32, - REDUCTION_CHUNK_SIZE as u32, - partial_count as u32, - ], - partial_count, - ); - let result = encode_sumcheck_reduction(command, partials, partial_count); - wait_for_command_named(command, "bn254_sumcheck_chunks"); - download_sumcheck_pair(&result) -} - -pub(crate) fn parallel_fold_pair_sumcheck( - a: &MetalFieldBuffer, - b: &MetalFieldBuffer, - weight: &MetalFieldBuffer, - len: usize, - fold_half: usize, -) -> (Field256, Field256) { - let sum_half = fold_half.next_power_of_two() >> 1; - debug_assert!(sum_half > 0); - let partial_count = sum_half.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * 2 * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel( - command, - pipeline(rt, "bn254_fold_pair_sumcheck_chunks"), - &[&a.limbs, &b.limbs, &weight.limbs, &partials.limbs], - &[ - len as u32, - fold_half as u32, - sum_half as u32, - REDUCTION_CHUNK_SIZE as u32, - partial_count as u32, - ], - partial_count, - ); - let result = encode_sumcheck_reduction(command, partials, partial_count); - wait_for_command_named(command, "bn254_fold_pair_sumcheck_chunks"); - download_sumcheck_pair(&result) -} - -pub(crate) fn parallel_univariate_evaluate( - coeffs: &MetalFieldBuffer, - point: &MetalFieldBuffer, - len: usize, -) -> Field256 { - parallel_univariate_evaluate_at(coeffs, 0, point, len) -} - -/// Univariate evaluation of `coeffs[coeff_off..coeff_off+len]` at `point`. -pub(crate) fn parallel_univariate_evaluate_at( - coeffs: &MetalFieldBuffer, - coeff_off: usize, - point: &MetalFieldBuffer, - len: usize, -) -> Field256 { - if len == 0 { - return Field256::ZERO; - } - let partial_count = len.div_ceil(REDUCTION_CHUNK_SIZE); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_count * size_of::()) as u64), - }; - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_univariate_eval_chunks"), - &[&coeffs.limbs, &point.limbs, &partials.limbs], - &[field_byte_offset(coeff_off), 0, 0], - &[len as u32, REDUCTION_CHUNK_SIZE as u32], - partial_count, - ); - let (result, offset) = encode_field_reduction(command, partials, partial_count, 0); - wait_for_command_named(command, "bn254_univariate_eval_chunks"); - download_field_at(&result.limbs, offset) -} - -pub(crate) fn parallel_geometric_accumulate_point_blocks( - acc: &MetalFieldBuffer, - points: &MetalFieldBuffer, - point_steps: &MetalFieldBuffer, - scalars: &MetalFieldBuffer, - len: usize, - num_points: usize, - chunk_size: usize, -) { - assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); - let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); - let chunks = len.div_ceil(chunk_size); - let partial_len = point_blocks - .checked_mul(len) - .expect("Metal geometric partial size overflow"); - let rt = runtime(); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), - }; - run_in_place( - "bn254_geometric_accumulate_point_blocks", - &[ - &partials.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[ - len as u32, - num_points as u32, - chunk_size as u32, - GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, - point_blocks as u32, - ], - chunks * point_blocks, - ); - run_in_place( - "bn254_geometric_accumulate_reduce_point_blocks", - &[&acc.limbs, &partials.limbs], - &[len as u32, point_blocks as u32], - len, - ); -} - -pub(crate) fn parallel_geometric_accumulate_point_blocks_batched( - acc: &MetalFieldBuffer, - points: &MetalFieldBuffer, - point_steps: &MetalFieldBuffer, - scalars: &MetalFieldBuffer, - len: usize, - num_points: usize, - chunk_size: usize, -) { - assert!(chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE); - let point_blocks = num_points.div_ceil(GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE); - let chunks = len.div_ceil(chunk_size); - let bytes_per_point_block = len - .checked_mul(size_of::()) - .expect("Metal geometric partial size overflow"); - let default_batch_blocks = - (GEOMETRIC_ACCUMULATE_POINT_BLOCK_BATCH_BYTES / bytes_per_point_block).max(1); - let batch_blocks = std::env::var("WHIR_METAL_GEOM_POINT_BLOCK_BATCH") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|&value| value > 0) - .unwrap_or(default_batch_blocks) - .min(point_blocks); - let rt = runtime(); - for point_block_offset in (0..point_blocks).step_by(batch_blocks) { - let current_batch = batch_blocks.min(point_blocks - point_block_offset); - let partial_len = current_batch - .checked_mul(len) - .expect("Metal geometric partial size overflow"); - let partials = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (partial_len * size_of::()) as u64), - }; - run_in_place( - "bn254_geometric_accumulate_point_blocks_range", - &[ - &partials.limbs, - &points.limbs, - &point_steps.limbs, - &scalars.limbs, - ], - &[ - len as u32, - num_points as u32, - chunk_size as u32, - GEOMETRIC_ACCUMULATE_POINT_BLOCK_SIZE as u32, - point_block_offset as u32, - current_batch as u32, - ], - chunks * current_batch, - ); - run_in_place( - "bn254_geometric_accumulate_reduce_point_blocks", - &[&acc.limbs, &partials.limbs], - &[len as u32, current_batch as u32], - len, - ); - } -} - -pub(crate) fn geometric_accumulate_chunk_size(len: usize) -> usize { - std::env::var("WHIR_METAL_GEOM_CHUNK") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|&value| value > 0) - .unwrap_or(if len >= LARGE_GEOMETRIC_ACCUMULATE_THRESHOLD { - LARGE_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - } else { - SMALL_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - }) - .min(MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE) -} - -pub(crate) fn should_use_geometric_point_blocks( - len: usize, - num_points: usize, - chunk_size: usize, -) -> bool { - chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS - && len <= GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD -} - -pub(crate) fn should_use_geometric_point_blocks_batched( - len: usize, - num_points: usize, - chunk_size: usize, -) -> bool { - chunk_size <= MAX_GEOMETRIC_ACCUMULATE_CHUNK_SIZE - && num_points >= GEOMETRIC_ACCUMULATE_POINT_BLOCK_MIN_POINTS - && len > GEOMETRIC_ACCUMULATE_POINT_BLOCK_THRESHOLD -} - -/// Encodes all tree-reduction levels into `command` and returns the buffer -/// and offset holding the final scalar. Runs with a single command wait. -pub(crate) fn encode_field_reduction( - command: &metal::CommandBufferRef, - input: MetalFieldBuffer, - len: usize, - offset: usize, -) -> (MetalFieldBuffer, usize) { - let rt = runtime(); - let mut current = input; - let mut current_len = len; - let mut current_offset = offset; - while current_len > 1 { - let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); - let next = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (next_len * size_of::()) as u64), - }; - encode_u32_kernel( - command, - pipeline(rt, "bn254_sum_chunks"), - &[¤t.limbs, &next.limbs], - &[ - current_len as u32, - current_offset as u32, - REDUCTION_CHUNK_SIZE as u32, - ], - next_len, - ); - current = next; - current_len = next_len; - current_offset = 0; - } - (current, current_offset) -} - -/// Encodes all (c0, c2) tree-reduction levels into `command` and returns the -/// buffer holding the final pair. Runs with a single command wait. -pub(crate) fn encode_sumcheck_reduction( - command: &metal::CommandBufferRef, - input: MetalFieldBuffer, - len: usize, -) -> MetalFieldBuffer { - let rt = runtime(); - let mut current = input; - let mut current_len = len; - while current_len > 1 { - let next_len = current_len.div_ceil(REDUCTION_CHUNK_SIZE); - let next = MetalFieldBuffer { - limbs: new_shared_buffer(rt, (next_len * 2 * size_of::()) as u64), - }; - encode_u32_kernel( - command, - pipeline(rt, "bn254_sumcheck_reduce_chunks"), - &[¤t.limbs, &next.limbs], - &[ - current_len as u32, - REDUCTION_CHUNK_SIZE as u32, - next_len as u32, - ], - next_len, - ); - current = next; - current_len = next_len; - } - current -} - -pub(crate) fn download_sumcheck_pair(buffer: &MetalFieldBuffer) -> (Field256, Field256) { - let values = download_field(&buffer.limbs, 2); - (values[0], values[1]) -} - -pub(crate) fn encode_single_vector_coset_ntt( - vector: &MetalBuffer, - message_length: usize, - interleaving_depth: usize, - codeword_length: usize, - coset_size: usize, -) -> MetalBuffer { - assert!(codeword_length.is_power_of_two()); - assert!(Field256::get_root_of_unity(codeword_length as u64).is_some()); - assert_eq!(vector.len(), message_length * interleaving_depth); - - assert!(codeword_length.is_multiple_of(coset_size)); - let num_cosets = codeword_length / coset_size; - let total_elements = interleaving_depth - .checked_mul(codeword_length) - .expect("Metal RS encode size overflow"); - assert!(total_elements <= u32::MAX as usize); - - let rt = runtime(); - let source = vector.bn254_buffer(); - let current = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); - let transposed = new_shared_buffer(rt, (total_elements * 4 * size_of::()) as u64); - let codeword_root_powers = root_powers_buffer(codeword_length); - let coset_roots = roots_buffer(coset_size); - - let command = rt.queue.new_command_buffer(); - - encode_kernel( - &command, - &rt.pack_single_vector_cosets, - &[&source.limbs, ¤t], - &PackSingleVectorParams { - row_count: interleaving_depth as u32, - message_length: message_length as u32, - codeword_length: codeword_length as u32, - coset_size: coset_size as u32, - total_elements: total_elements as u32, - }, - total_elements, - ); - - let trailing_elements = interleaving_depth.saturating_mul(codeword_length - coset_size); - if trailing_elements != 0 { - encode_kernel( - &command, - &rt.replicate_first_coset, - &[¤t], - &ReplicateCosetsParams { - row_len: codeword_length as u32, - coset_size: coset_size as u32, - trailing_elements: trailing_elements as u32, - }, - trailing_elements, - ); - } - - encode_kernel( - &command, - &rt.apply_coset_twiddles, - &[¤t, &codeword_root_powers.limbs], - &ApplyCosetTwiddlesParams { - row_count: interleaving_depth as u32, - num_cosets: num_cosets as u32, - coset_size: coset_size as u32, - codeword_length: codeword_length as u32, - total_elements: total_elements as u32, - }, - total_elements, - ); - - let stage_count = coset_size.trailing_zeros() as usize; - encode_kernel( - &command, - &rt.bit_reverse_rows, - &[¤t], - &BitReverseParams { - row_len: coset_size as u32, - log_n: stage_count as u32, - total_elements: total_elements as u32, - _pad0: 0, - }, - total_elements, - ); - - let total_butterflies = total_elements / 2; - let mut twiddle_offset = 0usize; - for stage in 0..stage_count { - let half_m = 1usize << stage; - encode_kernel( - &command, - &rt.ntt_stage_rows, - &[¤t, &coset_roots.limbs], - &NttStageParams { - row_len: coset_size as u32, - half_m: half_m as u32, - twiddle_offset: twiddle_offset as u32, - _pad0: 0, - }, - total_butterflies, - ); - twiddle_offset += 1usize << stage; - } - - encode_kernel( - &command, - &rt.transpose, - &[¤t, &transposed], - &TransposeParams { - rows: interleaving_depth as u32, - cols: codeword_length as u32, - total_elements: total_elements as u32, - }, - total_elements, - ); - - wait_for_command_named(&command, "bn254_rs_encode"); - - MetalBuffer::from_field_limb_buffer(transposed, total_elements) -} - -pub(crate) fn encode_kernel( - command: &metal::CommandBufferRef, - pipeline: &ComputePipelineState, - buffers: &[&Buffer], - params: &P, - threads: usize, -) { - let encoder = command.new_compute_command_encoder(); - encoder.set_compute_pipeline_state(pipeline); - for (index, buffer) in buffers.iter().enumerate() { - encoder.set_buffer(index as u64, Some(buffer), 0); - } - encoder.set_bytes( - buffers.len() as u64, - size_of::

() as u64, - (params as *const P).cast::(), - ); - dispatch(&encoder, pipeline, threads.max(1)); - encoder.end_encoding(); -} - -pub(crate) fn roots_buffer(codeword_length: usize) -> MetalFieldBuffer { - let rt = runtime(); - if let Some(buffer) = rt.ntt_roots.lock().unwrap().get(&codeword_length).cloned() { - return buffer; - } - let root = Field256::get_root_of_unity(codeword_length as u64) - .expect("BN254 root of unity unavailable for Metal NTT"); - let stage_count = codeword_length.trailing_zeros() as usize; - let mut roots = Vec::with_capacity(codeword_length.saturating_sub(1)); - for stage in 0..stage_count { - let stage_size = 1usize << (stage + 1); - let half_stage = stage_size >> 1; - let stage_root = root.pow([(codeword_length / stage_size) as u64]); - let mut current = Field256::ONE; - for _ in 0..half_stage { - roots.push(current); - current *= stage_root; - } - } - let buffer = upload_field(&roots); - rt.ntt_roots - .lock() - .unwrap() - .insert(codeword_length, buffer.clone()); - buffer -} - -pub(crate) fn root_powers_buffer(codeword_length: usize) -> MetalFieldBuffer { - let rt = runtime(); - if let Some(buffer) = rt - .root_powers - .lock() - .unwrap() - .get(&codeword_length) - .cloned() - { - return buffer; - } - let root = Field256::get_root_of_unity(codeword_length as u64) - .expect("BN254 root of unity unavailable for Metal RS twiddles"); - let mut powers = Vec::with_capacity(codeword_length); - let mut current = Field256::ONE; - for _ in 0..codeword_length { - powers.push(current); - current *= root; - } - let buffer = upload_field(&powers); - rt.root_powers - .lock() - .unwrap() - .insert(codeword_length, buffer.clone()); - buffer -} - -pub(crate) fn encode_field_rows_le(input: &Buffer, rows: usize, cols: usize) -> Buffer { - assert!(rows <= u32::MAX as usize); - assert!(cols <= u32::MAX as usize); - let rt = runtime(); - let output = new_shared_buffer(rt, (rows * cols * size_of::()) as u64); - let command = rt.queue.new_command_buffer(); - encode_kernel( - &command, - &rt.encode_field_rows_le, - &[input, &output], - &FieldBytesParams { - rows: rows as u32, - cols: cols as u32, - }, - rows * cols, - ); - wait_for_command_named(&command, "bn254_encode_field_rows_le"); - output -} - -pub(crate) fn read_bn254_rows( - source: &MetalFieldBuffer, - num_cols: usize, - indices: &[usize], -) -> Vec { - if indices.is_empty() || num_cols == 0 { - return Vec::new(); - } - assert!(num_cols <= u32::MAX as usize); - assert!(indices.iter().all(|&index| index <= u32::MAX as usize)); - let total = indices - .len() - .checked_mul(num_cols) - .expect("Metal read_rows size overflow"); - assert!(total <= u32::MAX as usize); - - let rt = runtime(); - let indices = indices - .iter() - .copied() - .map(|index| index as u32) - .collect::>(); - let indices_buffer = new_shared_buffer_with_data( - rt, - indices.as_ptr().cast(), - (indices.len() * size_of::()) as u64, - ); - let out = new_shared_buffer(rt, (total * size_of::()) as u64); - run_in_place( - "bn254_read_rows", - &[&source.limbs, &indices_buffer, &out], - &[num_cols as u32, total as u32], - total, - ); - download_field(&out, total) -} - -pub(crate) fn upload_field(values: &[Field256]) -> MetalFieldBuffer { - let rt = runtime(); - let buffer = if values.is_empty() { - new_shared_buffer(rt, 0) - } else { - // Field256 is 4 contiguous u64 limbs; upload directly without - // flattening into an intermediate Vec. - new_shared_buffer_with_data( - rt, - values.as_ptr().cast(), - std::mem::size_of_val(values) as u64, - ) - }; - MetalFieldBuffer { limbs: buffer } -} - -pub(crate) fn zeroed_field_buffer(len: usize) -> MetalFieldBuffer { - let rt = runtime(); - let bytes = (len * size_of::()) as u64; - let buffer = new_shared_buffer(rt, bytes); - if bytes > 0 { - let command = rt.queue.new_command_buffer(); - let blit = command.new_blit_command_encoder(); - blit.fill_buffer(&buffer, metal::NSRange::new(0, bytes), 0); - blit.end_encoding(); - wait_for_blit(command, bytes); - } - MetalFieldBuffer { limbs: buffer } -} - -pub(crate) fn maybe_upload_bn254(values: &[T]) -> Option { - (type_name::() == type_name::()).then(|| upload_field(as_field256_slice(values))) -} - -/// Geometric accumulate into `field[offset .. offset+len]` by binding the -/// field buffer at a byte offset; the kernel itself indexes from `gid == 0`. -pub(crate) fn geometric_accumulate_at_offset( - field: &MetalFieldBuffer, - offset: usize, - len: usize, - points: &MetalFieldBuffer, - scalars: &MetalFieldBuffer, - num_points: usize, -) { - if len == 0 || num_points == 0 { - return; - } - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - let encoder = command.new_compute_command_encoder(); - let pipe = pipeline(rt, "bn254_geometric_accumulate"); - encoder.set_compute_pipeline_state(pipe); - let byte_offset = (offset * size_of::()) as u64; - encoder.set_buffer(0, Some(&field.limbs), byte_offset); - encoder.set_buffer(1, Some(&points.limbs), 0); - encoder.set_buffer(2, Some(&scalars.limbs), 0); - let len_u32 = len as u32; - let num_u32 = num_points as u32; - encoder.set_bytes(3, size_of::() as u64, (&len_u32 as *const u32).cast()); - encoder.set_bytes(4, size_of::() as u64, (&num_u32 as *const u32).cast()); - dispatch(&encoder, pipe, len); - encoder.end_encoding(); - wait_for_command_named(command, "bn254_geometric_accumulate_window"); -} - -/// In-place scalar multiply of `field[offset .. offset+len]` by binding the -/// field buffer at a byte offset; the kernel indexes from `gid == 0`. -pub(crate) fn scalar_mul_at_offset( - field: &MetalFieldBuffer, - offset: usize, - len: usize, - weight: &MetalFieldBuffer, -) { - if len == 0 { - return; - } - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - let encoder = command.new_compute_command_encoder(); - let pipe = pipeline(rt, "bn254_scalar_mul"); - encoder.set_compute_pipeline_state(pipe); - let byte_offset = (offset * size_of::()) as u64; - encoder.set_buffer(0, Some(&field.limbs), byte_offset); - encoder.set_buffer(1, Some(&weight.limbs), 0); - let len_u32 = len as u32; - encoder.set_bytes(2, size_of::() as u64, (&len_u32 as *const u32).cast()); - dispatch(&encoder, pipe, len); - encoder.end_encoding(); - wait_for_command_named(command, "bn254_scalar_mul_window"); -} - -/// Multilinear extension of `field[off..off+len]` evaluated at `point`. -pub(crate) fn parallel_multilinear_extend_at( - field: &MetalFieldBuffer, - off: usize, - len: usize, - point: &MetalFieldBuffer, - num_vars: usize, -) -> Field256 { - let rt = runtime(); - let out = new_shared_buffer(rt, (4 * size_of::()) as u64); - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_multilinear_extend"), - &[&field.limbs, &point.limbs, &out], - &[field_byte_offset(off), 0, 0], - &[len as u32, num_vars as u32], - 1, - ); - wait_for_command_named(command, "bn254_multilinear_extend"); - download_field(&out, 1)[0] -} - -/// `acc[acc_off..] += weight * vector[vec_off..]` over `len` elements. -pub(crate) fn scalar_mul_add_at( - acc: &MetalFieldBuffer, - acc_off: usize, - vector: &MetalFieldBuffer, - vec_off: usize, - weight: &MetalFieldBuffer, - len: usize, -) { - if len == 0 { - return; - } - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - encode_u32_kernel_with_offsets( - command, - pipeline(rt, "bn254_scalar_mul_add"), - &[&acc.limbs, &vector.limbs, &weight.limbs], - &[field_byte_offset(acc_off), field_byte_offset(vec_off), 0], - &[len as u32], - len, - ); - wait_for_command_named(command, "bn254_scalar_mul_add"); -} - -pub(crate) fn copy_field_buffer(source: &MetalFieldBuffer, len: usize) -> MetalFieldBuffer { - copy_field_buffer_at(source, 0, len) -} - -pub(crate) fn copy_field_buffer_at( - source: &MetalFieldBuffer, - offset: usize, - len: usize, -) -> MetalFieldBuffer { - let rt = runtime(); - let byte_len = (len * 4 * size_of::()) as u64; - let source_offset = (offset * 4 * size_of::()) as u64; - let target = new_shared_buffer(rt, byte_len); - let command = rt.queue.new_command_buffer(); - let blit = command.new_blit_command_encoder(); - blit.copy_from_buffer(&source.limbs, source_offset, &target, 0, byte_len); - blit.end_encoding(); - wait_for_blit(&command, byte_len); - MetalFieldBuffer { limbs: target } -} - -/// Encodes a kernel dispatch with `u32` constants bound as inline bytes -/// (no per-constant buffer allocations). -pub(crate) fn encode_u32_kernel( - command: &metal::CommandBufferRef, - pipeline: &ComputePipelineState, - buffers: &[&Buffer], - constants: &[u32], - threads: usize, -) { - encode_u32_kernel_with_offsets(command, pipeline, buffers, &[], constants, threads); -} - -/// Like [`encode_u32_kernel`], but binds `buffers[i]` at byte offset -/// `offsets[i]` (defaulting to 0 when `offsets` is shorter). This is the -/// mechanism that lets a view dispatch a kernel over `parent[offset..]` -/// without copying: only the input binding shifts, the kernel still indexes -/// from `gid == 0`. -pub(crate) fn encode_u32_kernel_with_offsets( - command: &metal::CommandBufferRef, - pipeline: &ComputePipelineState, - buffers: &[&Buffer], - offsets: &[u64], - constants: &[u32], - threads: usize, -) { - let encoder = command.new_compute_command_encoder(); - encoder.set_compute_pipeline_state(pipeline); - let mut index = 0; - for (i, buffer) in buffers.iter().enumerate() { - let byte_offset = offsets.get(i).copied().unwrap_or(0); - encoder.set_buffer(index, Some(buffer), byte_offset); - index += 1; - } - for constant in constants { - encoder.set_bytes( - index, - size_of::() as u64, - (constant as *const u32).cast(), - ); - index += 1; - } - dispatch(&encoder, pipeline, threads.max(1)); - encoder.end_encoding(); -} - -/// Byte offset of element `offset` within a `Field256` buffer. -#[inline] -pub(crate) fn field_byte_offset(offset: usize) -> u64 { - (offset * size_of::()) as u64 -} - -pub(crate) fn run_in_place(name: &str, buffers: &[&Buffer], constants: &[u32], threads: usize) { - let rt = runtime(); - let command = rt.queue.new_command_buffer(); - encode_u32_kernel(command, pipeline(rt, name), buffers, constants, threads); - wait_for_command_named(command, name); -} - -pub(crate) fn dispatch( - encoder: &metal::ComputeCommandEncoderRef, - pipeline: &ComputePipelineState, - threads: usize, -) { - // Use full threadgroups (capped at 256) instead of a single execution - // width; the pipeline limit already accounts for register pressure. - let width = pipeline.max_total_threads_per_threadgroup().clamp(1, 256); - let group_width = width.min(threads as u64).max(1); - encoder.dispatch_threads( - MTLSize { - width: threads as u64, - height: 1, - depth: 1, - }, - MTLSize { - width: group_width, - height: 1, - depth: 1, - }, - ); -} - -pub(crate) fn download_field(buffer: &Buffer, len: usize) -> Vec { - if len == 0 { - return Vec::new(); - } - let start = Instant::now(); - let limbs = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 4) }; - let result = limbs - .chunks_exact(4) - .map(|chunk| { - Fp::, 4>( - BigInt([chunk[0], chunk[1], chunk[2], chunk[3]]), - PhantomData, - ) - }) - .collect(); - profile::record_readback((len * size_of::()) as u64, start.elapsed()); - result -} - -pub(crate) fn download_field_at(buffer: &Buffer, index: usize) -> Field256 { - let start = Instant::now(); - let limbs = - unsafe { std::slice::from_raw_parts(buffer.contents().cast::().add(index * 4), 4) }; - let result = Fp::, 4>( - BigInt([limbs[0], limbs[1], limbs[2], limbs[3]]), - PhantomData, - ); - profile::record_readback(size_of::() as u64, start.elapsed()); - result -} - -pub(crate) fn download_hash_indices(buffer: &Buffer, len: usize, indices: &[usize]) -> Vec { - let start = Instant::now(); - let bytes = unsafe { std::slice::from_raw_parts(buffer.contents().cast::(), len * 32) }; - let mut result = Vec::with_capacity(indices.len()); - for &index in indices { - assert!(index < len, "Metal hash index out of bounds"); - let mut hash = Digest::default(); - hash.as_mut_bytes() - .copy_from_slice(&bytes[index * 32..(index + 1) * 32]); - result.push(hash); - } - profile::record_readback( - (indices.len() * size_of::()) as u64, - start.elapsed(), - ); - result -} - -pub(crate) fn assert_bn254() { - assert_eq!( - type_name::(), - type_name::(), - "MetalBuffer only supports BN254 Field256 field operations" - ); -} - -pub(crate) fn f_to_field256(value: F) -> Field256 { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -pub(crate) fn field256_to_f(value: Field256) -> F { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -pub(crate) fn target_to_field256(value: T) -> Field256 { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -pub(crate) fn field256_to_target(value: Field256) -> T { - assert_bn254::(); - unsafe { std::mem::transmute_copy(&value) } -} - -pub(crate) fn as_field256_slice(values: &[T]) -> &[Field256] { - assert_eq!( - type_name::(), - type_name::(), - "MetalBuffer only supports BN254 Field256 buffers" - ); - unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } -} diff --git a/src/buffer/metal/sha2.rs b/src/buffer/metal/sha2.rs deleted file mode 100644 index f2862738..00000000 --- a/src/buffer/metal/sha2.rs +++ /dev/null @@ -1,722 +0,0 @@ -// NOTE: 100% AI GENERATED - -use std::{borrow::Cow, sync::OnceLock, time::Instant}; - -use const_oid::ObjectIdentifier; -use digest::const_oid::AssociatedOid; -use metal::{ - Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, MTLResourceOptions, MTLSize, -}; -use sha2::Sha256; -use zerocopy::IntoBytes; - -use super::profile; -use crate::hash::{Hash, HashEngine, HASH_COUNTER}; - -const USE_SHA256_64_SPECIALIZATION: bool = false; -const POW_MAX_WINDOW: u32 = 1 << 20; -const POW_MIN_WINDOW: u32 = 1 << 12; -const POW_THREADS_PER_GROUP: u64 = 256; - -const SHA256_METAL: &str = r#" -#include -using namespace metal; - -constant uint K[64] = { - 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, - 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, - 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, - 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, - 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, - 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, - 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, - 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, - 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, - 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, - 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, - 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, - 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, - 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, - 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, - 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2, -}; - -inline uint rotr(uint x, uint n) { - return (x >> n) | (x << (32 - n)); -} - -inline uchar padded_byte( - device const uchar *input, - uint message, - uint message_size, - uint block_offset, - uint padded_size -) { - ulong bit_len = ((ulong)message_size) * 8UL; - if (block_offset < message_size) { - return input[message * message_size + block_offset]; - } - if (block_offset == message_size) { - return 0x80; - } - uint len_start = padded_size - 8; - if (block_offset >= len_start) { - uint shift = (7 - (block_offset - len_start)) * 8; - return (uchar)((bit_len >> shift) & 0xff); - } - return 0; -} - -kernel void sha256_many( - device const uchar *input [[buffer(0)]], - device uchar *output [[buffer(1)]], - constant uint &message_size [[buffer(2)]], - constant uint &messages [[buffer(3)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= messages) return; - - uint padded_size = ((message_size + 9 + 63) / 64) * 64; - uint h0 = 0x6a09e667; - uint h1 = 0xbb67ae85; - uint h2 = 0x3c6ef372; - uint h3 = 0xa54ff53a; - uint h4 = 0x510e527f; - uint h5 = 0x9b05688c; - uint h6 = 0x1f83d9ab; - uint h7 = 0x5be0cd19; - - for (uint block = 0; block < padded_size; block += 64) { - uint w[64]; - for (uint i = 0; i < 16; i++) { - uint off = block + i * 4; - w[i] = - ((uint)padded_byte(input, gid, message_size, off + 0, padded_size) << 24) | - ((uint)padded_byte(input, gid, message_size, off + 1, padded_size) << 16) | - ((uint)padded_byte(input, gid, message_size, off + 2, padded_size) << 8) | - ((uint)padded_byte(input, gid, message_size, off + 3, padded_size)); - } - for (uint i = 16; i < 64; i++) { - uint s0 = rotr(w[i - 15], 7) ^ rotr(w[i - 15], 18) ^ (w[i - 15] >> 3); - uint s1 = rotr(w[i - 2], 17) ^ rotr(w[i - 2], 19) ^ (w[i - 2] >> 10); - w[i] = w[i - 16] + s0 + w[i - 7] + s1; - } - - uint a = h0, b = h1, c = h2, d = h3; - uint e = h4, f = h5, g = h6, h = h7; - for (uint i = 0; i < 64; i++) { - uint s1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25); - uint ch = (e & f) ^ ((~e) & g); - uint temp1 = h + s1 + ch + K[i] + w[i]; - uint s0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22); - uint maj = (a & b) ^ (a & c) ^ (b & c); - uint temp2 = s0 + maj; - h = g; - g = f; - f = e; - e = d + temp1; - d = c; - c = b; - b = a; - a = temp1 + temp2; - } - - h0 += a; h1 += b; h2 += c; h3 += d; - h4 += e; h5 += f; h6 += g; h7 += h; - } - - uint hs[8] = { h0, h1, h2, h3, h4, h5, h6, h7 }; - uint out = gid * 32; - for (uint i = 0; i < 8; i++) { - output[out + i * 4 + 0] = (uchar)((hs[i] >> 24) & 0xff); - output[out + i * 4 + 1] = (uchar)((hs[i] >> 16) & 0xff); - output[out + i * 4 + 2] = (uchar)((hs[i] >> 8) & 0xff); - output[out + i * 4 + 3] = (uchar)(hs[i] & 0xff); - } -} - -inline uint load_be32(device const uchar *input, uint off) { - return ((uint)input[off + 0] << 24) | - ((uint)input[off + 1] << 16) | - ((uint)input[off + 2] << 8) | - ((uint)input[off + 3]); -} - -inline void sha256_compress_16(thread uint h[8], thread uint w[16]) { - uint a = h[0], b = h[1], c = h[2], d = h[3]; - uint e = h[4], f = h[5], g = h[6], hh = h[7]; - - for (uint i = 0; i < 64; i++) { - uint wi; - if (i < 16) { - wi = w[i]; - } else { - uint s0 = rotr(w[(i + 1) & 15], 7) ^ rotr(w[(i + 1) & 15], 18) ^ (w[(i + 1) & 15] >> 3); - uint s1 = rotr(w[(i + 14) & 15], 17) ^ rotr(w[(i + 14) & 15], 19) ^ (w[(i + 14) & 15] >> 10); - wi = w[i & 15] + s0 + w[(i + 9) & 15] + s1; - w[i & 15] = wi; - } - uint S1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25); - uint ch = (e & f) ^ ((~e) & g); - uint temp1 = hh + S1 + ch + K[i] + wi; - uint S0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22); - uint maj = (a & b) ^ (a & c) ^ (b & c); - uint temp2 = S0 + maj; - hh = g; - g = f; - f = e; - e = d + temp1; - d = c; - c = b; - b = a; - a = temp1 + temp2; - } - - h[0] += a; h[1] += b; h[2] += c; h[3] += d; - h[4] += e; h[5] += f; h[6] += g; h[7] += hh; -} - -kernel void sha256_many_64( - device const uchar *input [[buffer(0)]], - device uchar *output [[buffer(1)]], - constant uint &messages [[buffer(2)]], - uint gid [[thread_position_in_grid]] -) { - if (gid >= messages) return; - - uint h[8] = { - 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, - 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19 - }; - - uint w[16]; - uint base = gid * 64; - for (uint i = 0; i < 16; i++) { - w[i] = load_be32(input, base + i * 4); - } - sha256_compress_16(h, w); - - w[0] = 0x80000000; - for (uint i = 1; i < 15; i++) { - w[i] = 0; - } - w[15] = 512; - sha256_compress_16(h, w); - - uint out = gid * 32; - for (uint i = 0; i < 8; i++) { - output[out + i * 4 + 0] = (uchar)((h[i] >> 24) & 0xff); - output[out + i * 4 + 1] = (uchar)((h[i] >> 16) & 0xff); - output[out + i * 4 + 2] = (uchar)((h[i] >> 8) & 0xff); - output[out + i * 4 + 3] = (uchar)(h[i] & 0xff); - } -} - -inline ulong digest_threshold_value(uint h0, uint h1) { - ulong b0 = (ulong)((h0 >> 24) & 0xff); - ulong b1 = (ulong)((h0 >> 16) & 0xff); - ulong b2 = (ulong)((h0 >> 8) & 0xff); - ulong b3 = (ulong)(h0 & 0xff); - ulong b4 = (ulong)((h1 >> 24) & 0xff); - ulong b5 = (ulong)((h1 >> 16) & 0xff); - ulong b6 = (ulong)((h1 >> 8) & 0xff); - ulong b7 = (ulong)(h1 & 0xff); - return b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) | - (b4 << 32) | (b5 << 40) | (b6 << 48) | (b7 << 56); -} - -kernel void sha256_pow_64( - device const uchar *challenge [[buffer(0)]], - device ulong *candidates [[buffer(1)]], - constant ulong &start_nonce [[buffer(2)]], - constant ulong &threshold [[buffer(3)]], - constant uint &count [[buffer(4)]], - uint gid [[thread_position_in_grid]], - uint tid [[thread_index_in_threadgroup]], - uint tgid [[threadgroup_position_in_grid]] -) { - threadgroup ulong local_best[256]; - ulong candidate = 0xffffffffffffffffUL; - - if (gid < count) { - ulong nonce = start_nonce + (ulong)gid; - - uint h[8] = { - 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, - 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19 - }; - - uint w[16]; - for (uint i = 0; i < 8; i++) { - w[i] = load_be32(challenge, i * 4); - } - w[8] = ((uint)(nonce & 0xffUL) << 24) | - ((uint)((nonce >> 8) & 0xffUL) << 16) | - ((uint)((nonce >> 16) & 0xffUL) << 8) | - ((uint)((nonce >> 24) & 0xffUL)); - w[9] = ((uint)((nonce >> 32) & 0xffUL) << 24) | - ((uint)((nonce >> 40) & 0xffUL) << 16) | - ((uint)((nonce >> 48) & 0xffUL) << 8) | - ((uint)((nonce >> 56) & 0xffUL)); - for (uint i = 10; i < 16; i++) { - w[i] = 0; - } - sha256_compress_16(h, w); - - w[0] = 0x80000000; - for (uint i = 1; i < 15; i++) { - w[i] = 0; - } - w[15] = 512; - sha256_compress_16(h, w); - - ulong value = digest_threshold_value(h[0], h[1]); - if (value <= threshold) { - candidate = nonce; - } - } - - local_best[tid] = candidate; - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint stride = 128; stride > 0; stride >>= 1) { - if (tid < stride && local_best[tid + stride] < local_best[tid]) { - local_best[tid] = local_best[tid + stride]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - if (tid == 0) { - candidates[tgid] = local_best[0]; - } -} -"#; - -#[derive(Clone, Copy, Debug, Default)] -pub struct MetalSha2; - -struct MetalSha2Runtime { - device: Device, - queue: CommandQueue, - pipeline: ComputePipelineState, - pipeline_64: ComputePipelineState, - pow_pipeline: ComputePipelineState, -} - -fn new_shared_buffer(rt: &MetalSha2Runtime, bytes: u64) -> Buffer { - profile::record_alloc(bytes); - let buffer = rt - .device - .new_buffer(bytes, MTLResourceOptions::StorageModeShared); - profile::record_device_allocated(rt.device.current_allocated_size()); - buffer -} - -fn new_shared_buffer_with_data( - rt: &MetalSha2Runtime, - data: *const std::ffi::c_void, - bytes: u64, -) -> Buffer { - let start = Instant::now(); - let buffer = rt - .device - .new_buffer_with_data(data, bytes, MTLResourceOptions::StorageModeShared); - profile::record_alloc(bytes); - profile::record_upload(bytes, start.elapsed()); - profile::record_device_allocated(rt.device.current_allocated_size()); - buffer -} - -fn upload_u32(rt: &MetalSha2Runtime, value: u32) -> Buffer { - new_shared_buffer_with_data(rt, (&value as *const u32).cast(), size_of::() as u64) -} - -fn upload_u64(rt: &MetalSha2Runtime, value: u64) -> Buffer { - new_shared_buffer_with_data(rt, (&value as *const u64).cast(), size_of::() as u64) -} - -fn wait_for_command(command: &metal::CommandBufferRef) { - wait_for_command_named(command, "sha256"); -} - -fn wait_for_command_named(command: &metal::CommandBufferRef, label: &str) { - command.commit(); - let start = Instant::now(); - command.wait_until_completed(); - let elapsed = start.elapsed(); - if std::env::var_os("WHIR_METAL_TRACE").is_some() { - eprintln!( - "metal command {label} {:.3} ms", - elapsed.as_secs_f64() * 1_000.0 - ); - } - profile::record_command_wait(elapsed); -} - -fn runtime() -> &'static MetalSha2Runtime { - static RUNTIME: OnceLock = OnceLock::new(); - RUNTIME.get_or_init(|| { - let device = Device::system_default().expect("Metal device is not available"); - let library = device - .new_library_with_source(SHA256_METAL, &CompileOptions::new()) - .expect("failed to compile Metal SHA-256 kernel"); - let function = library - .get_function("sha256_many", None) - .expect("missing Metal SHA-256 kernel"); - let function_64 = library - .get_function("sha256_many_64", None) - .expect("missing Metal SHA-256 64-byte kernel"); - let pow_function = library - .get_function("sha256_pow_64", None) - .expect("missing Metal SHA-256 PoW kernel"); - let pipeline = device - .new_compute_pipeline_state_with_function(&function) - .expect("failed to create Metal SHA-256 pipeline"); - let pipeline_64 = device - .new_compute_pipeline_state_with_function(&function_64) - .expect("failed to create Metal SHA-256 64-byte pipeline"); - let pow_pipeline = device - .new_compute_pipeline_state_with_function(&pow_function) - .expect("failed to create Metal SHA-256 PoW pipeline"); - let queue = device.new_command_queue(); - MetalSha2Runtime { - device, - queue, - pipeline, - pipeline_64, - pow_pipeline, - } - }) -} - -impl MetalSha2 { - pub const fn new() -> Self { - Self - } - - pub fn warmup() { - let _ = runtime(); - } - - pub fn prove_pow_64(challenge: &[u8; 32], threshold: u64) -> u64 { - let rt = runtime(); - let window = pow_window_for_threshold(threshold); - let challenge_buffer = - new_shared_buffer_with_data(rt, challenge.as_ptr().cast(), challenge.len() as u64); - let candidate_groups = u64::from(window).div_ceil(POW_THREADS_PER_GROUP); - let candidate_bytes = candidate_groups * size_of::() as u64; - let candidates = new_shared_buffer(rt, candidate_bytes); - let threshold = upload_u64(rt, threshold); - let count = upload_u32(rt, window); - - let mut start_nonce = 0u64; - loop { - let start = upload_u64(rt, start_nonce); - let command = rt.queue.new_command_buffer(); - let encoder = command.new_compute_command_encoder(); - encoder.set_compute_pipeline_state(&rt.pow_pipeline); - encoder.set_buffer(0, Some(&challenge_buffer), 0); - encoder.set_buffer(1, Some(&candidates), 0); - encoder.set_buffer(2, Some(&start), 0); - encoder.set_buffer(3, Some(&threshold), 0); - encoder.set_buffer(4, Some(&count), 0); - encoder.dispatch_thread_groups( - MTLSize { - width: candidate_groups, - height: 1, - depth: 1, - }, - MTLSize { - width: POW_THREADS_PER_GROUP, - height: 1, - depth: 1, - }, - ); - encoder.end_encoding(); - wait_for_command_named(&command, "sha256_pow"); - - let read_start = Instant::now(); - let candidates = unsafe { - std::slice::from_raw_parts( - candidates.contents().cast::(), - candidate_groups as usize, - ) - }; - let nonce = candidates.iter().copied().min().unwrap_or(u64::MAX); - profile::record_readback(candidate_bytes, read_start.elapsed()); - if nonce != u64::MAX { - HASH_COUNTER.add((nonce - start_nonce + 1) as usize); - return nonce; - } - HASH_COUNTER.add(window as usize); - start_nonce = start_nonce - .checked_add(window as u64) - .expect("PoW nonce range exhausted"); - } - } - - pub(crate) fn build_merkle_tree_buffer_from_messages_buffer( - &self, - message_size: usize, - messages: &Buffer, - num_leaves: usize, - layers: usize, - ) -> Buffer { - assert_eq!(num_leaves, 1usize << layers); - let rt = runtime(); - let num_nodes = (1usize << (layers + 1)) - 1; - let nodes_bytes = (num_nodes * size_of::()) as u64; - let nodes = new_shared_buffer(rt, nodes_bytes); - - let command = rt.queue.new_command_buffer(); - let mut constants = Vec::with_capacity((layers + 1) * 2); - self.encode_hash_many_into_command( - &command, - &mut constants, - message_size, - messages, - 0, - num_leaves, - &nodes, - 0, - ); - - let mut previous_offset = 0usize; - let mut previous_len = num_leaves; - let mut next_offset = num_leaves; - for _ in 0..layers { - let current_len = previous_len / 2; - self.encode_hash_many_into_command( - &command, - &mut constants, - 64, - &nodes, - (previous_offset * size_of::()) as u64, - current_len, - &nodes, - (next_offset * size_of::()) as u64, - ); - previous_offset = next_offset; - previous_len = current_len; - next_offset += current_len; - } - wait_for_command_named(&command, "sha256_merkle_fused"); - drop(constants); - nodes - } - - fn hash_many_buffer_into_buffer( - &self, - size: usize, - input_buffer: &Buffer, - input_offset: u64, - output_len: usize, - output_buffer: &Buffer, - output_offset: u64, - ) { - if output_len == 0 { - return; - } - - let command = runtime().queue.new_command_buffer(); - let mut constants = Vec::with_capacity(2); - self.encode_hash_many_into_command( - &command, - &mut constants, - size, - input_buffer, - input_offset, - output_len, - output_buffer, - output_offset, - ); - wait_for_command(&command); - drop(constants); - } - - #[allow(clippy::too_many_arguments)] - fn encode_hash_many_into_command( - &self, - command: &metal::CommandBufferRef, - constants: &mut Vec, - size: usize, - input_buffer: &Buffer, - input_offset: u64, - output_len: usize, - output_buffer: &Buffer, - output_offset: u64, - ) { - if output_len == 0 { - return; - } - let rt = runtime(); - let messages_buffer = upload_u32(rt, output_len as u32); - let encoder = command.new_compute_command_encoder(); - let pipeline = if size == 64 && USE_SHA256_64_SPECIALIZATION { - encoder.set_compute_pipeline_state(&rt.pipeline_64); - encoder.set_buffer(0, Some(input_buffer), input_offset); - encoder.set_buffer(1, Some(output_buffer), output_offset); - encoder.set_buffer(2, Some(&messages_buffer), 0); - &rt.pipeline_64 - } else { - let size_buffer = upload_u32(rt, size as u32); - encoder.set_compute_pipeline_state(&rt.pipeline); - encoder.set_buffer(0, Some(input_buffer), input_offset); - encoder.set_buffer(1, Some(output_buffer), output_offset); - encoder.set_buffer(2, Some(&size_buffer), 0); - encoder.set_buffer(3, Some(&messages_buffer), 0); - constants.push(size_buffer); - &rt.pipeline - }; - constants.push(messages_buffer); - let group_width = pipeline - .thread_execution_width() - .min(output_len as u64) - .max(1); - encoder.dispatch_threads( - MTLSize { - width: output_len as u64, - height: 1, - depth: 1, - }, - MTLSize { - width: group_width, - height: 1, - depth: 1, - }, - ); - encoder.end_encoding(); - HASH_COUNTER.add(output_len); - } - - pub(crate) fn hash_many_buffer( - &self, - size: usize, - input_buffer: &Buffer, - output_len: usize, - output: &mut [Hash], - ) { - assert_eq!(output_len, output.len()); - if output.is_empty() { - return; - } - - let rt = runtime(); - - let output_bytes = (output.len() * size_of::()) as u64; - let output_buffer = new_shared_buffer(rt, output_bytes); - self.hash_many_buffer_into_buffer(size, input_buffer, 0, output.len(), &output_buffer, 0); - - let start = Instant::now(); - let bytes = unsafe { - std::slice::from_raw_parts(output_buffer.contents().cast::(), output.len() * 32) - }; - for (out, bytes) in output.iter_mut().zip(bytes.chunks_exact(32)) { - out.as_mut_bytes().copy_from_slice(bytes); - } - profile::record_readback(output_bytes, start.elapsed()); - } -} - -fn pow_window_for_threshold(threshold: u64) -> u32 { - if threshold == 0 { - return POW_MAX_WINDOW; - } - - let expected_attempts = u64::MAX - .checked_div(threshold) - .unwrap_or(u64::MAX) - .saturating_add(1); - let target = expected_attempts.saturating_mul(4); - let window = target - .checked_next_power_of_two() - .unwrap_or(u64::from(POW_MAX_WINDOW)) - .clamp(u64::from(POW_MIN_WINDOW), u64::from(POW_MAX_WINDOW)); - window as u32 -} - -impl HashEngine for MetalSha2 { - fn name(&self) -> Cow<'_, str> { - "sha2".into() - } - - fn oid(&self) -> Option { - Some(Sha256::OID) - } - - fn supports_size(&self, _size: usize) -> bool { - Device::system_default().is_some() - } - - fn preferred_batch_size(&self) -> usize { - 256 - } - - fn hash_many(&self, size: usize, input: &[u8], output: &mut [Hash]) { - assert_eq!( - input.len(), - size * output.len(), - "Input length ({}) should be size * output.len() = {size} * {}", - input.len(), - output.len() - ); - if output.is_empty() { - return; - } - - let input_buffer = if input.is_empty() { - new_shared_buffer(runtime(), 0) - } else { - new_shared_buffer_with_data(runtime(), input.as_ptr().cast(), input.len() as u64) - }; - self.hash_many_buffer(size, &input_buffer, output.len(), output); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::hash::Sha2; - - #[test] - fn metal_sha2_matches_cpu() { - let rows = 17; - let size = 73; - let input = (0..rows * size) - .map(|i: usize| (i.wrapping_mul(31) & 0xff) as u8) - .collect::>(); - let mut cpu = vec![Hash::default(); rows]; - let mut gpu = vec![Hash::default(); rows]; - Sha2::new().hash_many(size, &input, &mut cpu); - MetalSha2::new().hash_many(size, &input, &mut gpu); - assert_eq!(gpu, cpu); - } - - #[test] - fn metal_sha2_64_matches_cpu() { - let rows = 257; - let size = 64; - let input = (0..rows * size) - .map(|i: usize| (i.wrapping_mul(17).wrapping_add(3) & 0xff) as u8) - .collect::>(); - let mut cpu = vec![Hash::default(); rows]; - let mut gpu = vec![Hash::default(); rows]; - Sha2::new().hash_many(size, &input, &mut cpu); - MetalSha2::new().hash_many(size, &input, &mut gpu); - assert_eq!(gpu, cpu); - } - - #[test] - fn metal_sha2_pow_matches_threshold() { - let challenge = [7u8; 32]; - let threshold = u64::MAX >> 8; - let nonce = MetalSha2::prove_pow_64(&challenge, threshold); - let mut input = [0u8; 64]; - input[..32].copy_from_slice(&challenge); - input[32..40].copy_from_slice(&nonce.to_le_bytes()); - let mut hash = Hash::default(); - Sha2::new().hash_many(64, &input, std::slice::from_mut(&mut hash)); - let value = u64::from_le_bytes(hash.0[..8].try_into().unwrap()); - assert!(value <= threshold); - } -} diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index daaad4ab..4d4088ef 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -13,8 +13,6 @@ //! [`DefaultRs`] selects the Reed-Solomon encoder for the active backend. pub mod cpu; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub mod metal; use std::ops::RangeBounds; use ark_ff::Field; @@ -35,25 +33,9 @@ use crate::{ pub use cpu::{CpuBuffer, CpuSlice, CpuSliceMut}; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub use metal::{MetalBuffer, MetalRs, MetalSlice, MetalSliceMut}; - -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type ActiveBuffer = MetalBuffer; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type ActiveSlice<'a, T> = MetalSlice<'a, T>; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type ActiveSliceMut<'a, T> = MetalSliceMut<'a, T>; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub type DefaultRs = MetalRs; - -#[cfg(not(all(feature = "metal", target_os = "macos")))] pub type ActiveBuffer = CpuBuffer; -#[cfg(not(all(feature = "metal", target_os = "macos")))] pub type ActiveSlice<'a, T> = CpuSlice<'a, T>; -#[cfg(not(all(feature = "metal", target_os = "macos")))] pub type ActiveSliceMut<'a, T> = CpuSliceMut<'a, T>; -#[cfg(not(all(feature = "metal", target_os = "macos")))] pub type DefaultRs = crate::algebra::ntt::NttEngine; /// Owned buffer operations over any element type. diff --git a/src/hash/mod.rs b/src/hash/mod.rs index 2618beda..53a7e84b 100644 --- a/src/hash/mod.rs +++ b/src/hash/mod.rs @@ -14,10 +14,6 @@ use serde::{Deserialize, Serialize}; use static_assertions::{assert_impl_all, assert_obj_safe}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned}; -#[cfg(all(feature = "metal", target_os = "macos"))] -pub use crate::buffer::metal::{ - metal_profile_snapshot, metal_reset_device_peak, MetalProfileSnapshot, MetalSha2, -}; pub use self::{ blake3_engine::{Blake3, BLAKE3}, copy_engine::{Copy, COPY}, @@ -36,8 +32,6 @@ pub static ENGINES: LazyLock> = LazyLock::new(|| { engines.register(Arc::new(Keccak::new())); engines.register(Arc::new(Sha3::new())); engines.register(Arc::new(Blake3::detect())); - #[cfg(all(feature = "metal", target_os = "macos"))] - engines.register(Arc::new(MetalSha2::new())); engines }); diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index ceb966fc..df7fa04e 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -511,18 +511,4 @@ pub(crate) mod tests { fn test_field256() { proptest::(); } - - #[test] - #[cfg(all(feature = "metal", target_os = "macos"))] - fn test_field256_sha2_metal_commit_open() { - test::( - StdRng::seed_from_u64(7), - hash::SHA2, - hash::SHA2, - 4, - 16, - 2, - &[0, 1, 3, 8, 8, 15], - ); - } } diff --git a/src/protocols/proof_of_work.rs b/src/protocols/proof_of_work.rs index ac6f10d7..a2a8aabc 100644 --- a/src/protocols/proof_of_work.rs +++ b/src/protocols/proof_of_work.rs @@ -20,9 +20,6 @@ use crate::{ verify, }; -#[cfg(all(feature = "metal", target_os = "macos"))] -use crate::hash::{MetalSha2, SHA2}; - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct Config { pub hash_id: EngineId, @@ -90,13 +87,6 @@ impl Config { let challenge: [u8; 32] = prover_state.verifier_message(); - #[cfg(all(feature = "metal", target_os = "macos"))] - if self.hash_id == SHA2 { - let nonce = MetalSha2::prove_pow_64(&challenge, self.threshold); - prover_state.prover_message(&U64(nonce)); - return; - } - #[cfg(not(feature = "parallel"))] let nonce = (0_u64..) .step_by(batch_size) From 72b31c0704d8cdf63a52c6a096444a3ffd8bd39d Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Fri, 19 Jun 2026 16:14:31 +0200 Subject: [PATCH 22/35] refactor: use ActiveBuffer::random --- benches/expand_from_coeff.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benches/expand_from_coeff.rs b/benches/expand_from_coeff.rs index 862da207..c1d8e334 100644 --- a/benches/expand_from_coeff.rs +++ b/benches/expand_from_coeff.rs @@ -34,7 +34,7 @@ fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) { let num_messages = 1 << coset_sz; let mut rng = ark_std::rand::thread_rng(); let coeffs: Vec> = (0..num_messages) - .map(|_| ActiveBuffer::from_vec(random_vector(&mut rng, message_length))) + .map(|_| ActiveBuffer::random(&mut rng, message_length)) .collect(); let engine = NttEngine::::new_from_fftfield(); (engine, coeffs, expansion) From d68e838227176eb596f47799d0543ffdee75d579 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Fri, 19 Jun 2026 16:24:34 +0200 Subject: [PATCH 23/35] refactor: remove visibility modifiers from CpuBuffer and CpuSlice structs --- src/buffer/cpu/buffer.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/buffer/cpu/buffer.rs b/src/buffer/cpu/buffer.rs index 1f9d145c..def98e40 100644 --- a/src/buffer/cpu/buffer.rs +++ b/src/buffer/cpu/buffer.rs @@ -31,17 +31,17 @@ use crate::{ serde::Deserialize, )] pub struct CpuBuffer { - pub(super) data: Vec, + data: Vec, } /// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. pub struct CpuSlice<'a, F> { - pub(super) data: &'a [F], + data: &'a [F], } /// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. pub struct CpuSliceMut<'a, F> { - pub(super) data: &'a mut [F], + data: &'a mut [F], } impl BufferOps for CpuBuffer { From 03d110a0531aece4cc484dfacacb235abe59b012 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 24 Jun 2026 22:06:25 +0200 Subject: [PATCH 24/35] refactor: use buffer abstraction in mask handling --- src/protocols/basecase.rs | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index e1e4adeb..7d015f87 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize}; use spongefish::{Decoding, VerificationResult}; use crate::{ - algebra::{embedding::Identity, multilinear_extend, scalar_mul_add_new, univariate_evaluate}, + algebra::{embedding::Identity, multilinear_extend, univariate_evaluate}, buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, hash::Hash, protocols::{irs_commit, sumcheck}, @@ -101,7 +101,7 @@ impl Config { let mask = ActiveBuffer::random(prover_state.rng(), vector.len()); // Commit to the masking vector. - let mask_witness = self.commit.commit(prover_state, &[&mask]); + let mut mask_witness = self.commit.commit(prover_state, &[&mask]); // Compute and send linear form of mask (μ' in paper). let mask_sum = mask.dot(&covector); @@ -110,17 +110,16 @@ impl Config { // RLC the mask with the vector let mask_rlc = prover_state.verifier_message::(); assert!(!mask_rlc.is_zero(), "Proof failed"); - let mut masked_vector = mask.clone(); + let mut masked_vector = mask; vector.mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_vector, mask_rlc); prover_state.prover_messages(masked_vector.as_slice()); // Send combined IRS randomness. (r^* in paper) - let masked_masks = scalar_mul_add_new( - mask_witness.masks.as_slice(), - mask_rlc, - witness.masks.as_slice(), - ); - prover_state.prover_messages(&masked_masks); + let mut masked_masks = std::mem::take(&mut mask_witness.masks); + witness + .masks + .mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_masks, mask_rlc); + prover_state.prover_messages(masked_masks.as_slice()); // Open the commitment and mask simultaneously. let _ = self.commit.open(prover_state, &[&mask_witness, witness]); From f7dbd7f1dfe57741bf0a360f325eb0d155ade9fd Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Wed, 24 Jun 2026 22:22:24 +0200 Subject: [PATCH 25/35] refactor: use clone instead of mem::take --- src/protocols/basecase.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 7d015f87..e2a235c8 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -101,7 +101,7 @@ impl Config { let mask = ActiveBuffer::random(prover_state.rng(), vector.len()); // Commit to the masking vector. - let mut mask_witness = self.commit.commit(prover_state, &[&mask]); + let mask_witness = self.commit.commit(prover_state, &[&mask]); // Compute and send linear form of mask (μ' in paper). let mask_sum = mask.dot(&covector); @@ -115,7 +115,7 @@ impl Config { prover_state.prover_messages(masked_vector.as_slice()); // Send combined IRS randomness. (r^* in paper) - let mut masked_masks = std::mem::take(&mut mask_witness.masks); + let mut masked_masks = mask_witness.masks.clone(); witness .masks .mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_masks, mask_rlc); From fbc63438a2f7ef39c2bad2a072eab4868547cf6b Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 2 Jul 2026 07:59:45 +0200 Subject: [PATCH 26/35] refactor: simplify the design by removing slice and mut --- benches/expand_from_coeff.rs | 11 +- src/buffer/{cpu/buffer.rs => cpu.rs} | 223 ++++++++++----------------- src/buffer/cpu/mod.rs | 5 - src/buffer/cpu/read_write.rs | 121 --------------- src/buffer/mod.rs | 189 ++++++++--------------- src/protocols/basecase.rs | 16 +- src/protocols/code_switch.rs | 8 +- src/protocols/irs_commit.rs | 2 +- src/protocols/mask_proximity.rs | 36 +++-- src/protocols/matrix_commit.rs | 7 +- src/protocols/merkle_tree.rs | 76 ++++----- src/protocols/sumcheck.rs | 2 +- src/protocols/whir/mod.rs | 2 +- src/protocols/whir/prover.rs | 2 +- src/protocols/whir_zk/mod.rs | 2 +- src/protocols/whir_zk/prover.rs | 3 +- 16 files changed, 233 insertions(+), 472 deletions(-) rename src/buffer/{cpu/buffer.rs => cpu.rs} (52%) delete mode 100644 src/buffer/cpu/mod.rs delete mode 100644 src/buffer/cpu/read_write.rs diff --git a/benches/expand_from_coeff.rs b/benches/expand_from_coeff.rs index c1d8e334..bdcfde9a 100644 --- a/benches/expand_from_coeff.rs +++ b/benches/expand_from_coeff.rs @@ -1,9 +1,10 @@ use divan::{black_box, AllocProfiler, Bencher}; -use whir::algebra::{ - buffer::{ActiveBuffer, BufferOps}, - fields::Field64, - ntt::{Messages, NttEngine, ReedSolomon}, - random_vector, +use whir::{ + algebra::{ + fields::Field64, + ntt::{Messages, NttEngine, ReedSolomon}, + }, + buffer::{ActiveBuffer, Buffer, BufferOps}, }; #[global_allocator] diff --git a/src/buffer/cpu/buffer.rs b/src/buffer/cpu.rs similarity index 52% rename from src/buffer/cpu/buffer.rs rename to src/buffer/cpu.rs index def98e40..7ab781a2 100644 --- a/src/buffer/cpu/buffer.rs +++ b/src/buffer/cpu.rs @@ -1,15 +1,16 @@ +//! In-memory CPU backend for the buffer abstraction. + use std::{any::Any, mem}; -use super::read_write::{impl_cpu_read, impl_cpu_write}; use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; use crate::{ algebra::{ embedding::Embedding, - linear_form::{Covector, LinearForm}, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, }, - buffer::{Buffer, BufferOps, BufferRead, BufferWrite}, + buffer::{Buffer, BufferOps}, engines::EngineId, hash::{self, Hash}, protocols::{ @@ -34,24 +35,16 @@ pub struct CpuBuffer { data: Vec, } -/// Read-only view into a [`CpuBuffer`], backed by a borrowed sub-slice. -pub struct CpuSlice<'a, F> { - data: &'a [F], -} - -/// Mutable view into a [`CpuBuffer`], backed by a borrowed sub-slice. -pub struct CpuSliceMut<'a, F> { - data: &'a mut [F], -} - impl BufferOps for CpuBuffer { type Nodes = CpuBuffer; - fn at_index(&self, index: usize) -> Option { - if index >= self.len() { - None - } else { - Some(self.data[index]) + fn from_vec(source: Vec) -> Self { + Self { data: source } + } + + fn from_slice(source: &[T]) -> Self { + Self { + data: source.to_vec(), } } @@ -72,20 +65,7 @@ impl BufferOps for CpuBuffer { } fn gather_at_indices(&self, indices: &[usize]) -> Vec { - indices - .iter() - .map(|&i| self.data[i]) - .collect() - } - - fn from_vec(source: Vec) -> Self { - Self { data: source } - } - - fn from_slice(source: &[T]) -> Self { - Self { - data: Vec::from(source), - } + indices.iter().map(|&i| self.data[i]).collect() } fn merklize( @@ -113,18 +93,14 @@ impl BufferOps for CpuBuffer { } impl Buffer for CpuBuffer { + type TargetBuffer = CpuBuffer; + fn zeros(length: usize) -> Self { Self { data: vec![F::ZERO; length], } } - fn geometric_sequence(base: F, length: usize) -> Self { - Self { - data: crate::algebra::geometric_sequence(base, length), - } - } - fn random(rng: &mut R, length: usize) -> Self where R: RngCore + CryptoRng, @@ -135,6 +111,58 @@ impl Buffer for CpuBuffer { } } + fn dot(&self, other: &Self) -> F { + crate::algebra::dot(&self.data, &other.data) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + crate::algebra::sumcheck::compute_sumcheck_polynomial(&self.data, &other.data) + } + + fn fold(&mut self, weight: F) { + crate::algebra::sumcheck::fold(&mut self.data, weight); + } + + fn scalar_mul(&mut self, weight: F) { + crate::algebra::scalar_mul(&mut self.data, weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + let Some(size) = evaluators.first().map(|e| e.size) else { + return; + }; + UnivariateEvaluation::accumulate_many(evaluators, &mut self.data[..size], scalars); + } + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + crate::algebra::mixed_univariate_evaluate(embedding, &self.data, point) + } + + fn mixed_dot>( + &self, + embedding: &M, + other: &CpuBuffer, + ) -> M::Target { + crate::algebra::mixed_dot(embedding, &other.data, &self.data) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut CpuBuffer, + weight: M::Target, + ) { + crate::algebra::mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, &self.data); + } + fn linear_forms_rlc( size: usize, linear_forms: &mut [Box>], @@ -158,155 +186,68 @@ impl Buffer for CpuBuffer { Self { data: covector } } - fn zero_pad(&mut self) { - if !self.is_empty() { - self.data.resize(self.len().next_power_of_two(), F::ZERO); - } - } - - fn fold(&mut self, weight: F) { - crate::algebra::sumcheck::fold(&mut self.data, weight); - } - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - self.fold(weight); - other.fold(weight); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - self.fold_pair(other, weight); - self.sumcheck_polynomial(other) - } - fn mixed_linear_combination>( embedding: &M, vectors: &[&Self], coeffs: &[M::Target], - ) -> Self::TargetBuffer { + ) -> CpuBuffer { assert_eq!(vectors.len(), coeffs.len()); let Some((first, vectors)) = vectors.split_first() else { return CpuBuffer { data: Vec::new() }; }; debug_assert_eq!(coeffs[0], M::Target::ONE); - let mut accumulator = crate::algebra::lift(embedding, first.as_slice()); + let mut accumulator = crate::algebra::lift(embedding, &first.data); for (coeff, vector) in coeffs[1..].iter().zip(vectors) { - crate::algebra::mixed_scalar_mul_add( - embedding, - &mut accumulator, - *coeff, - vector.as_slice(), - ); + crate::algebra::mixed_scalar_mul_add(embedding, &mut accumulator, *coeff, &vector.data); } CpuBuffer { data: accumulator } } } -impl_cpu_read!(CpuSlice<'_, F>); -impl_cpu_read!(CpuSliceMut<'_, F>); -impl_cpu_read!(CpuBuffer); -impl_cpu_write!(CpuSliceMut<'_, F>); -impl_cpu_write!(CpuBuffer); - -/// Resolve any range expression (`a..b`, `..b`, `a..`, `..`) against `len`, -/// returning `(start, end)` and bounds-checking like slice indexing. -pub(crate) fn resolve_range( - range: impl std::ops::RangeBounds, - len: usize, -) -> (usize, usize) { - use std::ops::Bound::{Excluded, Included, Unbounded}; - let start = match range.start_bound() { - Included(&s) => s, - Excluded(&s) => s + 1, - Unbounded => 0, - }; - let end = match range.end_bound() { - Included(&e) => e + 1, - Excluded(&e) => e, - Unbounded => len, - }; - assert!( - start <= end && end <= len, - "slice range {start}..{end} out of bounds for length {len}" - ); - (start, end) -} - #[cfg(test)] mod tests { use ark_ff::AdditiveGroup; use super::*; - use crate::algebra::{ - fields::Field64, geometric_accumulate, geometric_sequence as geometric_sequence_vec, - linear_form::UnivariateEvaluation, - }; + use crate::algebra::{fields::Field64, geometric_accumulate}; type F = Field64; #[test] - fn geometric_sequence_matches_free_function() { - let base = F::from(7u64); - for length in 0..6 { - let buffer = CpuBuffer::::geometric_sequence(base, length); - assert_eq!( - BufferOps::as_slice(&buffer), - geometric_sequence_vec(base, length).as_slice(), - "geometric_sequence mismatch at length {length}" - ); - } - } - - #[test] - fn scale_multiplies_in_place() { + fn scalar_mul_multiplies_in_place() { let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; let weight = F::from(5u64); let mut buffer = CpuBuffer::from_vec(values.clone()); buffer.scalar_mul(weight); let expected: Vec = values.iter().map(|&v| v * weight).collect(); - assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); + assert_eq!(buffer.as_slice(), expected.as_slice()); } #[test] - fn slice_accumulate_matches_restricted_geometric_accumulate() { + fn accumulate_matches_geometric_accumulate_over_prefix() { let len = 8usize; let points = vec![F::from(3u64), F::from(5u64)]; let scalars = vec![F::from(2u64), F::from(9u64)]; - // Try a prefix slice (offset 0) and an interior slice (offset > 0). - for (offset, window_len) in [(0usize, 5usize), (2usize, 4usize)] { + // Full-length and prefix accumulation. + for size in [len, 5] { let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); let evaluators: Vec<_> = points .iter() - .map(|&point| UnivariateEvaluation::new(point, window_len)) + .map(|&point| UnivariateEvaluation::new(point, size)) .collect(); - buffer - .slice_mut(offset..offset + window_len) - .accumulate_univariate_evaluations(&evaluators, &scalars); + buffer.accumulate_univariate_evaluations(&evaluators, &scalars); - // Reference: accumulate Σ_j scalars[j]·points[j]^i into the same - // sub-slice of a plain vector. + // Reference: accumulate Σ_j scalars[j]·points[j]^i into the prefix + // of a plain vector. let mut reference = vec![F::ZERO; len]; - geometric_accumulate( - &mut reference[offset..offset + window_len], - scalars.clone(), - &points, - ); + geometric_accumulate(&mut reference[..size], scalars.clone(), &points); assert_eq!( - BufferOps::as_slice(&buffer), + buffer.as_slice(), reference.as_slice(), - "slice accumulate mismatch at offset {offset}, len {window_len}" + "accumulate mismatch for evaluator size {size}" ); } } - - #[test] - fn slice_scale_multiplies_only_the_range() { - let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; - let weight = F::from(5u64); - let mut buffer = CpuBuffer::from_vec(values.clone()); - buffer.slice_mut(1..3).scalar_mul(weight); - let expected = vec![values[0], values[1] * weight, values[2] * weight, values[3]]; - assert_eq!(BufferOps::as_slice(&buffer), expected.as_slice()); - } } diff --git a/src/buffer/cpu/mod.rs b/src/buffer/cpu/mod.rs deleted file mode 100644 index d8c36c4a..00000000 --- a/src/buffer/cpu/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -mod buffer; -mod read_write; - -pub(crate) use buffer::resolve_range; -pub use buffer::{CpuBuffer, CpuSlice, CpuSliceMut}; diff --git a/src/buffer/cpu/read_write.rs b/src/buffer/cpu/read_write.rs deleted file mode 100644 index b191e86a..00000000 --- a/src/buffer/cpu/read_write.rs +++ /dev/null @@ -1,121 +0,0 @@ -macro_rules! impl_cpu_read { - ($ty:ty) => { - impl BufferRead for $ty { - type TargetBuffer = CpuBuffer; - type Slice<'a> - = CpuSlice<'a, F> - where - Self: 'a, - F: 'a; - - fn read_len(&self) -> usize { - self.data.len() - } - - fn dot(&self, other: &Self) -> F { - crate::algebra::dot(&*self.data, &*other.data) - } - - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { - crate::algebra::sumcheck::compute_sumcheck_polynomial(&*self.data, &*other.data) - } - - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target { - crate::algebra::mixed_multilinear_extend(embedding, &*self.data, point) - } - - fn mixed_dot, T: Field>( - &self, - embedding: &M, - other: &CpuBuffer, - ) -> M::Target { - crate::algebra::mixed_dot(embedding, other.as_slice(), &*self.data) - } - - fn mixed_univariate_evaluate>( - &self, - embedding: &M, - point: M::Target, - ) -> M::Target { - crate::algebra::mixed_univariate_evaluate(embedding, &*self.data, point) - } - - fn mixed_scalar_mul_add_to>( - &self, - embedding: &M, - accumulator: &mut CpuBuffer, - weight: M::Target, - ) { - crate::algebra::mixed_scalar_mul_add( - embedding, - &mut accumulator.data, - weight, - &*self.data, - ); - } - - fn slice(&self, range: impl std::ops::RangeBounds) -> CpuSlice<'_, F> { - let data = &*self.data; - let (start, end) = $crate::buffer::cpu::resolve_range(range, data.len()); - CpuSlice { - data: &data[start..end], - } - } - - fn copy_to_owned(&self) -> CpuBuffer { - CpuBuffer::from_slice(&*self.data) - } - } - }; -} - -macro_rules! impl_cpu_write { - ($ty:ty) => { - impl BufferWrite for $ty { - type SliceMut<'a> - = CpuSliceMut<'a, F> - where - Self: 'a, - F: 'a; - - fn scalar_mul(&mut self, weight: F) { - crate::algebra::scalar_mul(&mut *self.data, weight); - } - - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[crate::algebra::linear_form::UnivariateEvaluation], - scalars: &[F], - ) { - crate::algebra::linear_form::UnivariateEvaluation::accumulate_many( - evaluators, - &mut *self.data, - scalars, - ); - } - - fn slice_mut( - &mut self, - range: impl std::ops::RangeBounds, - ) -> CpuSliceMut<'_, F> { - let data = &mut *self.data; - let (start, end) = $crate::buffer::cpu::resolve_range(range, data.len()); - CpuSliceMut { - data: &mut data[start..end], - } - } - - fn split_at_mut(&mut self, mid: usize) -> (CpuSliceMut<'_, F>, CpuSliceMut<'_, F>) { - let (lo, hi) = self.data.split_at_mut(mid); - (CpuSliceMut { data: lo }, CpuSliceMut { data: hi }) - } - } - }; -} - -pub(crate) use impl_cpu_read; -pub(crate) use impl_cpu_write; diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index 4d4088ef..e1ab2cf9 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -1,19 +1,18 @@ //! Backend-agnostic buffers for protocol data. //! -//! Protocol code uses [`ActiveBuffer`], [`ActiveSlice`], and [`ActiveSliceMut`] -//! to select the CPU or GPU backend at compile time. Owned buffers manage -//! storage; slices are zero-copy views into a contiguous range. +//! Protocol code uses the [`ActiveBuffer`] alias to select the backend at +//! compile time. Buffers are owned, backend-managed storage: on the CPU +//! backend they wrap a `Vec`, on an accelerator backend they would own +//! device memory and only [`BufferOps::as_slice`] (and the other readback +//! methods) force a host copy. //! -//! The trait split follows the ownership model. [`BufferOps`] is generic over -//! any element type and is used for field elements, [`struct@Hash`] digests, and -//! Merkle nodes. [`BufferRead`] and [`BufferWrite`] provide field operations for -//! owned buffers and views. [`Buffer`] contains owned-only field operations such -//! as construction, folding, and zero-padding. +//! The trait split follows the element type. [`BufferOps`] is generic over +//! any element and also covers [`struct@Hash`] buffers for Merkle tree +//! nodes. [`Buffer`] adds field arithmetic used by the protocols. //! //! [`DefaultRs`] selects the Reed-Solomon encoder for the active backend. pub mod cpu; -use std::ops::RangeBounds; use ark_ff::Field; use ark_std::rand::{ @@ -31,39 +30,31 @@ use crate::{ protocols::{matrix_commit::Encodable, merkle_tree}, }; -pub use cpu::{CpuBuffer, CpuSlice, CpuSliceMut}; +pub use cpu::CpuBuffer; pub type ActiveBuffer = CpuBuffer; -pub type ActiveSlice<'a, T> = CpuSlice<'a, T>; -pub type ActiveSliceMut<'a, T> = CpuSliceMut<'a, T>; pub type DefaultRs = crate::algebra::ntt::NttEngine; /// Owned buffer operations over any element type. /// -/// This trait is not field-specific, so it also covers hash buffers and Merkle -/// node layers. Field arithmetic lives on [`BufferRead`] and [`BufferWrite`]. -pub trait BufferOps { +/// This trait is not field-specific, so it also covers hash buffers and +/// Merkle node layers. Field arithmetic lives on [`Buffer`]. +pub trait BufferOps { /// Same-backend buffer type used for Merkle tree nodes. type Nodes: BufferOps; fn from_vec(source: Vec) -> Self; fn from_slice(source: &[T]) -> Self; + /// Read back the buffer contents as a host slice. fn as_slice(&self) -> &[T]; - /// Number of rows when the buffer is laid out with `num_cols` columns. - fn num_rows(&self, num_cols: usize) -> usize { - self.len() / num_cols - } - /// Gather full rows `indices[i] * num_cols .. (indices[i] + 1) * num_cols`. - fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; - fn at_index(&self, index: usize) -> Option; - /// Gather elements at arbitrary indices. - fn gather_at_indices(&self, indices: &[usize]) -> Vec - where - T: Copy; fn len(&self) -> usize; fn is_empty(&self) -> bool { self.len() == 0 } + /// Gather full rows `indices[i] * num_cols .. (indices[i] + 1) * num_cols`. + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; + /// Gather elements at arbitrary indices. + fn gather_at_indices(&self, indices: &[usize]) -> Vec; /// Hash rows of width `num_cols` and build a Merkle tree. fn merklize( &self, @@ -75,55 +66,66 @@ pub trait BufferOps { T: Encodable + Send + Sync; } -/// Read-only operations over a contiguous run of field elements. -/// -/// Implemented by owned buffers and borrowed views. A view produced by -/// [`Self::slice`] aliases the original storage. -pub trait BufferRead { +/// Field operations on owned buffers. +pub trait Buffer: BufferOps + Clone { /// Same-backend owned buffer over another field. + /// + /// Used by the `mixed_*` operations that lift base-field data into an + /// extension field through an [`Embedding`]. type TargetBuffer: Buffer; - /// Read-only view type produced by slicing this buffer/view. - type Slice<'a>: BufferRead + fn zeros(length: usize) -> Self; + + fn random(rng: &mut R, length: usize) -> Self where - Self: 'a, - F: 'a; + R: RngCore + CryptoRng, + Standard: Distribution; + + /// Inner product with another buffer of the same length. + fn dot(&self, other: &Self) -> F; - /// Number of elements in this view. - fn read_len(&self) -> usize; - fn read_is_empty(&self) -> bool { - self.read_len() == 0 + /// Sumcheck round coefficients `(c0, c2)` for `dot(self, other)`. + fn sumcheck_polynomial(&self, other: &Self) -> (F, F); + + fn fold(&mut self, weight: F); + + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); } - /// Inner product with another view of the same length. - fn dot(&self, other: &Self) -> F - where - Self: Sized; + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) + } - /// Sumcheck round coefficients `(c0, c2)` for `dot(self, other)`. - fn sumcheck_polynomial(&self, other: &Self) -> (F, F) - where - Self: Sized; + /// In-place scalar multiplication: `self[i] *= weight`. + fn scalar_mul(&mut self, weight: F); - /// Multilinear extension evaluated at a target-field point. - fn mixed_extend, T: Field>( - &self, - embedding: &M, - point: &[M::Target], - ) -> M::Target; + /// Accumulate `Σ_j scalars[j] · evaluators[j].point^i` into entry `i`. + /// + /// The evaluators must share a common size `s ≤ self.len()`; only the + /// first `s` entries are updated. This allows accumulating constraints + /// that cover a prefix of the buffer (e.g. the unmasked message part of + /// a covector). + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ); - /// Inner product with a target-field buffer. - fn mixed_dot, T: Field>( + /// Univariate evaluation at a target-field point. + fn mixed_univariate_evaluate>( &self, embedding: &M, - other: &Self::TargetBuffer, + point: M::Target, ) -> M::Target; - /// Univariate evaluation at a target-field point. - fn mixed_univariate_evaluate>( + /// Inner product with a target-field buffer. + fn mixed_dot>( &self, embedding: &M, - point: M::Target, + other: &Self::TargetBuffer, ) -> M::Target; /// `accumulator += weight * self`, lifted into the target field. @@ -134,79 +136,14 @@ pub trait BufferRead { weight: M::Target, ); - /// Borrow `self[range]` as a read-only view (analogous to `&v[range]`). - fn slice(&self, range: impl RangeBounds) -> Self::Slice<'_>; - - /// Copy this view into a new owned buffer. - fn copy_to_owned(&self) -> Self::TargetBuffer; -} - -/// In-place operations over a contiguous run of field elements. -/// -/// Implemented by owned buffers and mutable views. A view produced by -/// [`Self::slice_mut`] aliases the original storage. -/// -/// Construction and length-changing operations live on [`Buffer`], because -/// borrowed views cannot construct or resize their backing storage. -pub trait BufferWrite: BufferRead { - /// Mutable view type produced by slicing this buffer/view. - type SliceMut<'a>: BufferWrite - where - Self: 'a, - F: 'a; - - /// In-place scalar multiplication: `self[i] *= weight`. - fn scalar_mul(&mut self, weight: F); - - /// Accumulate `Σ_j scalars[j] · evaluators[j].point^i` into entry `i`. - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &[F], - ); - - /// Borrow `self[range]` as a mutable, zero-copy view (analogous to - /// `&mut v[range]`). - fn slice_mut(&mut self, range: impl RangeBounds) -> Self::SliceMut<'_>; - - /// Split into two disjoint mutable views `[0, mid)` and `[mid, len)`. - fn split_at_mut(&mut self, mid: usize) -> (Self::SliceMut<'_>, Self::SliceMut<'_>); -} - -/// Owned field-buffer operations. -/// -/// This trait contains operations that construct storage or change its length, -/// so it is implemented only by owned buffers. -pub trait Buffer: BufferOps + BufferWrite + Clone { - fn zeros(length: usize) -> Self; - - /// Geometric sequence `[1, base, base², …, base^(length-1)]`. - fn geometric_sequence(base: F, length: usize) -> Self; - - fn random(rng: &mut R, length: usize) -> Self - where - R: RngCore + CryptoRng, - Standard: Distribution; - - fn zero_pad(&mut self); - fn fold(&mut self, weight: F); - - fn fold_pair(&mut self, other: &mut Self, weight: F) { - self.fold(weight); - other.fold(weight); - } - - fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { - self.fold_pair(other, weight); - self.sumcheck_polynomial(other) - } - + /// Random linear combination of linear forms into a covector buffer. fn linear_forms_rlc( size: usize, linear_forms: &mut [Box>], rlc_coeffs: &[F], ) -> Self; + /// Random linear combination of vectors, lifted into the target field. fn mixed_linear_combination>( embedding: &M, vectors: &[&Self], diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index e2a235c8..9c29cdca 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -14,7 +14,7 @@ use spongefish::{Decoding, VerificationResult}; use crate::{ algebra::{embedding::Identity, multilinear_extend, univariate_evaluate}, - buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{irs_commit, sumcheck}, transcript::{ @@ -88,12 +88,12 @@ impl Config { .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) .round_challenges; assert!( - !vector.at_index(0).unwrap_or_else(F::one).is_zero(), + !vector.as_slice().first().expect("Proof failed").is_zero(), "Proof failed" ); return Opening { evaluation_points: point, - linear_form_evaluation: covector.at_index(0).expect("Proof failed"), + linear_form_evaluation: *covector.as_slice().first().expect("Proof failed"), }; } @@ -142,13 +142,17 @@ impl Config { // no constraints on l(r) that the verifier can return. // This event is cryptographically unlikely as `F` is challenge sized. assert!( - !masked_vector.at_index(0).unwrap_or_else(F::one).is_zero(), + !masked_vector + .as_slice() + .first() + .expect("Proof failed") + .is_zero(), "Proof failed" ); Opening { evaluation_points: point, - linear_form_evaluation: covector.at_index(0).expect("Proof failed"), + linear_form_evaluation: *covector.as_slice().first().expect("Proof failed"), } } @@ -305,7 +309,7 @@ mod tests { sum, ); assert_eq!( - covector.mixed_extend(&Identity::::new(), &prover_result.evaluation_points), + multilinear_extend(covector.as_slice(), &prover_result.evaluation_points), prover_result.linear_form_evaluation ); let proof = prover_state.proof(); diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 0d7153aa..c7e7d0a6 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -19,7 +19,7 @@ use crate::{ linear_form::UnivariateEvaluation, mixed_dot, }, - buffer::{ActiveBuffer, BufferOps, BufferRead, BufferWrite}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{ geometric_challenge::geometric_challenge, @@ -246,14 +246,14 @@ impl Config { covector.accumulate_univariate_evaluations(&all_evaluators, constraint_rlc_coeffs); } else { // ZK: OOD contributes to full [f; r; s], in-domain only to [f; r]. + // The in-domain evaluators have size `masked`, so they only + // accumulate into that prefix of the covector. let ood_evaluators = univariate_evaluators(&ood_points, covector.len()); covector.accumulate_univariate_evaluations(&ood_evaluators, ood_rlc_coeffs); let masked = self.source.masked_message_length(); let in_domain_evaluators = univariate_evaluators(&eval_points, masked); - covector - .slice_mut(..masked) - .accumulate_univariate_evaluations(&in_domain_evaluators, in_domain_rlc_coeffs); + covector.accumulate_univariate_evaluations(&in_domain_evaluators, in_domain_rlc_coeffs); } Witness { diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index ad345b4b..6fe180cb 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -23,7 +23,7 @@ use std::{ ops::Neg, }; -use crate::buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}; +use crate::buffer::{ActiveBuffer, Buffer, BufferOps}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use ordered_float::OrderedFloat; diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index abce0b82..f54f291c 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -45,8 +45,8 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, R use serde::{Deserialize, Serialize}; use crate::{ - algebra::{embedding::Identity, univariate_evaluate}, - buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}, + algebra::{embedding::Identity, scalar_mul_add, univariate_evaluate}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::irs_commit::{ Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness, @@ -190,19 +190,24 @@ impl Config { .enumerate() { // ξ*_i = s_i + γ · ξ_i - let mut combined_msg = fresh_msg.copy_to_owned(); + let mut combined_msg = fresh_msg.clone(); orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); prover_state.prover_messages(combined_msg.as_slice()); // r*_i = r'_i + γ · r_i + // Combined on the host: the result is transcript data anyway, and + // the IRS randomness is only a few elements per vector. if irs_masks_per_vector > 0 { let base = i * irs_masks_per_vector; let fresh_base = (self.num_masks + i) * irs_masks_per_vector; - let orig_r = irs_masks.slice(base..base + irs_masks_per_vector); - let fresh_r = irs_masks.slice(fresh_base..fresh_base + irs_masks_per_vector); - let mut combined_r = fresh_r.copy_to_owned(); - orig_r.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_r, gamma); - prover_state.prover_messages(combined_r.as_slice()); + let masks = irs_masks.as_slice(); + let mut combined_r = masks[fresh_base..fresh_base + irs_masks_per_vector].to_vec(); + scalar_mul_add( + &mut combined_r, + gamma, + &masks[base..base + irs_masks_per_vector], + ); + prover_state.prover_messages(&combined_r); } } @@ -484,7 +489,7 @@ mod tests { .zip(witness.fresh_msgs.iter()) .enumerate() { - let mut combined_msg = fresh_msg.copy_to_owned(); + let mut combined_msg = fresh_msg.clone(); orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); if i == 0 { let mut tampered = combined_msg.as_slice().to_vec(); @@ -496,11 +501,14 @@ mod tests { if irs_masks_per_vector > 0 { let base = i * irs_masks_per_vector; let fresh_base = (config.num_masks + i) * irs_masks_per_vector; - let orig_r = irs_masks.slice(base..base + irs_masks_per_vector); - let fresh_r = irs_masks.slice(fresh_base..fresh_base + irs_masks_per_vector); - let mut combined_r = fresh_r.copy_to_owned(); - orig_r.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_r, gamma); - prover_state.prover_messages(combined_r.as_slice()); + let masks = irs_masks.as_slice(); + let mut combined_r = masks[fresh_base..fresh_base + irs_masks_per_vector].to_vec(); + scalar_mul_add( + &mut combined_r, + gamma, + &masks[base..base + irs_masks_per_vector], + ); + prover_state.prover_messages(&combined_r); } } diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index df7fa04e..02f409b9 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -114,12 +114,10 @@ where } #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] -pub struct BufferWitness { +pub struct Witness { pub nodes: ActiveBuffer, } -pub type Witness = BufferWitness; - pub type Commitment = merkle_tree::Commitment; /// Encode [`ark_ff::Field`]s using [`ArkFieldEncoder`]. @@ -224,13 +222,12 @@ impl Config { H: DuplexSpongeInterface, R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>, - ActiveBuffer: BufferOps>, { assert_eq!(matrix.len(), self.num_rows() * self.num_cols); let (nodes, root) = matrix.merklize(self.num_cols, self.leaf_hash_id, &self.merkle_tree); prover_state.prover_message(&root); - BufferWitness { nodes } + Witness { nodes } } #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index 006a2bca..f16a482f 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -274,6 +274,44 @@ fn parallel_hash(engine: &dyn HashEngine, size: usize, input: &[u8], output: &mu } } +/// Flat node indices of sibling hashes required to open at `indices`. +pub fn opening_sibling_indices( + num_leaves: usize, + layers: usize, + indices: &[usize], +) -> Vec { + debug_assert!(indices.iter().all(|&i| i < num_leaves)); + + let mut indices = indices.to_vec(); + indices.sort_unstable(); + indices.dedup(); + + let mut node_indices = Vec::new(); + let mut layer_offset = 0usize; + let mut layer_len = 1usize << layers; + while layer_len > 1 { + let mut next_indices = Vec::with_capacity(indices.len()); + let mut iter = indices.iter().copied().peekable(); + loop { + match (iter.next(), iter.peek()) { + (Some(a), Some(&b)) if b == a ^ 1 => { + next_indices.push(a >> 1); + iter.next(); + } + (Some(a), _) => { + node_indices.push(layer_offset + (a ^ 1)); + next_indices.push(a >> 1); + } + (None, _) => break, + } + } + indices = next_indices; + layer_offset += layer_len; + layer_len /= 2; + } + node_indices +} + #[cfg(test)] pub(crate) mod tests { use proptest::{collection::vec, prelude::Strategy}; @@ -343,41 +381,3 @@ pub(crate) mod tests { assert_eq!(layers_for_size(8), 3); } } - -/// Flat node indices of sibling hashes required to open at `indices`. -pub fn opening_sibling_indices( - num_leaves: usize, - layers: usize, - indices: &[usize], -) -> Vec { - debug_assert!(indices.iter().all(|&i| i < num_leaves)); - - let mut indices = indices.to_vec(); - indices.sort_unstable(); - indices.dedup(); - - let mut node_indices = Vec::new(); - let mut layer_offset = 0usize; - let mut layer_len = 1usize << layers; - while layer_len > 1 { - let mut next_indices = Vec::with_capacity(indices.len()); - let mut iter = indices.iter().copied().peekable(); - loop { - match (iter.next(), iter.peek()) { - (Some(a), Some(&b)) if b == a ^ 1 => { - next_indices.push(a >> 1); - iter.next(); - } - (Some(a), _) => { - node_indices.push(layer_offset + (a ^ 1)); - next_indices.push(a >> 1); - } - (None, _) => break, - } - } - indices = next_indices; - layer_offset += layer_len; - layer_len /= 2; - } - node_indices -} diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 7b4b36a9..eb07a3e4 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -2,7 +2,7 @@ use std::fmt; -use crate::buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead}; +use crate::buffer::{ActiveBuffer, Buffer, BufferOps}; use ark_ff::Field; use ark_std::rand::{CryptoRng, RngCore}; use serde::{Deserialize, Serialize}; diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 67497b30..11d5a253 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -77,7 +77,7 @@ impl FinalClaim { impl Config { /// Commit to one or more vectors. - #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().unwrap().len())))] + #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().map_or(0, |v| crate::buffer::BufferOps::len(*v)))))] pub fn commit( &self, prover_state: &mut ProverState, diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 5b66145f..4aba15c7 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -8,7 +8,7 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ algebra::{dot, embedding::Embedding, eq_weights, linear_form::LinearForm, tensor_product}, - buffer::{ActiveBuffer, Buffer, BufferOps, BufferRead, BufferWrite}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{geometric_challenge::geometric_challenge, irs_commit, whir::FinalClaim}, transcript::{ diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index 960f5a3c..689875bd 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -251,11 +251,11 @@ mod tests { use super::*; use crate::{ algebra::{ - buffer::ActiveBuffer, fields::Field64, linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, random_vector, }, + buffer::{ActiveBuffer, BufferOps}, hash, parameters::ProtocolParameters, transcript::{codecs::Empty, DomainSeparator, ProverState, VerifierState}, diff --git a/src/protocols/whir_zk/prover.rs b/src/protocols/whir_zk/prover.rs index a7512534..9ca26f97 100644 --- a/src/protocols/whir_zk/prover.rs +++ b/src/protocols/whir_zk/prover.rs @@ -10,10 +10,9 @@ use rayon::prelude::*; use tracing::instrument; use super::{Config, Witness}; -use crate::buffer::BufferOps; +use crate::buffer::{ActiveBuffer, BufferOps}; use crate::{ algebra::{ - buffer::ActiveBuffer, embedding::Identity, linear_form::{Covector, Evaluate, LinearForm}, mixed_dot, scalar_mul_add, From cfbcc52101c3fc377108e0325b5d97ccda6cabad Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 2 Jul 2026 08:14:35 +0200 Subject: [PATCH 27/35] docs: clarify caller obligations in build_nodes documentation --- src/protocols/merkle_tree.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index f16a482f..61f31f4a 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -81,6 +81,15 @@ impl Config { /// Build the full node array from leaf hashes, without touching the transcript. /// /// The returned vector holds the leaf layer first and the root last (at `num_nodes() - 1`). + /// + /// # Caller obligation + /// + /// This does **not** commit to anything: the caller must send the root + /// (`nodes[num_nodes() - 1]`) to the transcript before using the tree as + /// a commitment. In production this is funneled through + /// [`matrix_commit::Config::commit`](crate::protocols::matrix_commit::Config::commit), + /// which sends the root returned by `merklize`. A tree whose root was + /// never sent is not bound by the proof. pub fn build_nodes(&self, leaves: Vec) -> Vec { assert_eq!( leaves.len(), From cd657aa5ed0dbf04f8e2397bcfe30a01e37bfb94 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 2 Jul 2026 08:19:48 +0200 Subject: [PATCH 28/35] chore: run fmt --- src/protocols/merkle_tree.rs | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index 61f31f4a..5cbf6369 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -158,8 +158,7 @@ impl Config { assert_eq!(witness.nodes.len(), self.num_nodes()); assert!(indices.iter().all(|&i| i < self.num_leaves)); - let node_indices = - opening_sibling_indices(self.num_leaves, self.layers.len(), indices); + let node_indices = opening_sibling_indices(self.num_leaves, self.layers.len(), indices); for &i in &node_indices { prover_state.prover_hint(&witness.nodes[i]); } @@ -284,11 +283,7 @@ fn parallel_hash(engine: &dyn HashEngine, size: usize, input: &[u8], output: &mu } /// Flat node indices of sibling hashes required to open at `indices`. -pub fn opening_sibling_indices( - num_leaves: usize, - layers: usize, - indices: &[usize], -) -> Vec { +pub fn opening_sibling_indices(num_leaves: usize, layers: usize, indices: &[usize]) -> Vec { debug_assert!(indices.iter().all(|&i| i < num_leaves)); let mut indices = indices.to_vec(); From 676809b982aebc5e88688162e5435f3df4b03ac9 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 2 Jul 2026 08:21:52 +0200 Subject: [PATCH 29/35] refactor: update lifetime annotations in prove function --- src/protocols/whir_zk/prover.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/protocols/whir_zk/prover.rs b/src/protocols/whir_zk/prover.rs index 9ca26f97..a82a4dbf 100644 --- a/src/protocols/whir_zk/prover.rs +++ b/src/protocols/whir_zk/prover.rs @@ -209,13 +209,13 @@ impl Config { /// inner witness-side WHIR prover. #[allow(clippy::too_many_lines)] #[cfg_attr(feature = "tracing", instrument(skip_all))] - pub fn prove<'a, H, R>( + pub fn prove( &self, prover_state: &mut ProverState, vectors: &[&ActiveBuffer], witness: Witness, linear_forms: Vec>>, - evaluations: Cow<'a, [F]>, + evaluations: Cow<'_, [F]>, ) -> FinalClaim where Standard: Distribution, From 484051914b0f71a3cf0ffd94525c8bf352a90d85 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Thu, 2 Jul 2026 08:24:58 +0200 Subject: [PATCH 30/35] chore: run nightly fmt --- src/algebra/ntt/cooley_tukey.rs | 4 +--- src/bin/benchmark.rs | 3 +-- src/bin/main.rs | 3 +-- src/buffer/mod.rs | 3 +-- src/protocols/irs_commit.rs | 2 +- src/protocols/matrix_commit.rs | 3 +-- src/protocols/sumcheck.rs | 19 ++++++++++--------- src/protocols/whir/mod.rs | 2 +- src/protocols/whir_zk/committer.rs | 3 +-- src/protocols/whir_zk/prover.rs | 2 +- 10 files changed, 19 insertions(+), 25 deletions(-) diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index 4970b1a3..b713c33d 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -16,13 +16,11 @@ use super::{ utils::{lcm, sqrt_factor}, Messages, ReedSolomon, }; - -use crate::buffer::{ActiveBuffer, BufferOps}; - #[cfg(not(feature = "rs_in_order"))] use crate::algebra::ntt::transpose::transpose_permute; use crate::{ algebra::ntt::utils::divisors, + buffer::{ActiveBuffer, BufferOps}, utils::{chunks_exact_or_empty, zip_strict}, }; diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index 58f0451b..95464207 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -15,14 +15,13 @@ use whir::{ linear_form::{Evaluate, LinearForm, MultilinearExtension}, }, bits::Bits, + buffer::{ActiveBuffer, BufferOps}, cmdline_utils::{AvailableFields, AvailableHash}, hash::HASH_COUNTER, parameters::ProtocolParameters, transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; -use whir::buffer::{ActiveBuffer, BufferOps}; - #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { diff --git a/src/bin/main.rs b/src/bin/main.rs index 552ea58b..6c7de32b 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -11,14 +11,13 @@ use whir::{ linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, }, bits::Bits, + buffer::{ActiveBuffer, BufferOps}, cmdline_utils::{AvailableFields, AvailableHash}, hash::HASH_COUNTER, parameters::ProtocolParameters, transcript::{codecs::Empty, Codec, DomainSeparator, ProverState, VerifierState}, }; -use whir::buffer::{ActiveBuffer, BufferOps}; - #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index e1ab2cf9..fd58aaeb 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -19,6 +19,7 @@ use ark_std::rand::{ distributions::{Distribution, Standard}, CryptoRng, RngCore, }; +pub use cpu::CpuBuffer; use crate::{ algebra::{ @@ -30,8 +31,6 @@ use crate::{ protocols::{matrix_commit::Encodable, merkle_tree}, }; -pub use cpu::CpuBuffer; - pub type ActiveBuffer = CpuBuffer; pub type DefaultRs = crate::algebra::ntt::NttEngine; diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 6fe180cb..f3e9fc56 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -23,7 +23,6 @@ use std::{ ops::Neg, }; -use crate::buffer::{ActiveBuffer, Buffer, BufferOps}; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use ordered_float::OrderedFloat; @@ -36,6 +35,7 @@ use crate::{ dot, embedding::Embedding, fields::FieldWithSize, lift, linear_form::UnivariateEvaluation, ntt, }, + buffer::{ActiveBuffer, Buffer, BufferOps}, engines::EngineId, hash::Hash, protocols::{challenge_indices::challenge_indices, matrix_commit}, diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index 02f409b9..66fc4ebe 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -12,8 +12,7 @@ use tracing::instrument; use zerocopy::{Immutable, IntoBytes}; use crate::{ - buffer::ActiveBuffer, - buffer::BufferOps, + buffer::{ActiveBuffer, BufferOps}, engines::EngineId, hash::{self, Hash}, protocols::merkle_tree, diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index eb07a3e4..50df96fa 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -2,7 +2,6 @@ use std::fmt; -use crate::buffer::{ActiveBuffer, Buffer, BufferOps}; use ark_ff::Field; use ark_std::rand::{CryptoRng, RngCore}; use serde::{Deserialize, Serialize}; @@ -11,6 +10,7 @@ use tracing::instrument; use crate::{ algebra::univariate_evaluate, + buffer::{ActiveBuffer, Buffer, BufferOps}, protocols::proof_of_work, transcript::{ codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverState, VerificationResult, @@ -234,6 +234,15 @@ fn eval_01(coefficients: &[F]) -> F { #[cfg(test)] mod tests { + use ark_std::rand::{ + distributions::{Distribution, Standard}, + rngs::StdRng, + SeedableRng, + }; + use proptest::{prelude::Just, prop_oneof, proptest, strategy::Strategy}; + #[cfg(feature = "tracing")] + use tracing::instrument; + use super::*; use crate::{ algebra::{ @@ -244,14 +253,6 @@ mod tests { buffer::ActiveBuffer, transcript::DomainSeparator, }; - use ark_std::rand::{ - distributions::{Distribution, Standard}, - rngs::StdRng, - SeedableRng, - }; - use proptest::{prelude::Just, prop_oneof, proptest, strategy::Strategy}; - #[cfg(feature = "tracing")] - use tracing::instrument; impl Config where diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 11d5a253..057a0679 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -128,7 +128,6 @@ mod tests { use ark_std::rand::thread_rng; use super::*; - use crate::buffer::{ActiveBuffer, BufferOps}; use crate::{ algebra::{ embedding::Basefield, @@ -136,6 +135,7 @@ mod tests { linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, random_vector, }, + buffer::{ActiveBuffer, BufferOps}, hash, parameters::ProtocolParameters, transcript::{codecs::Empty, DomainSeparator, ProverState, VerifierState}, diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index 94f390c2..d44f4c1d 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -6,9 +6,8 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution}; use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; -use crate::buffer::BufferOps; use crate::{ - buffer::ActiveBuffer, + buffer::{ActiveBuffer, BufferOps}, hash::Hash, protocols::{irs_commit, whir}, transcript::{ diff --git a/src/protocols/whir_zk/prover.rs b/src/protocols/whir_zk/prover.rs index a82a4dbf..becf4b9e 100644 --- a/src/protocols/whir_zk/prover.rs +++ b/src/protocols/whir_zk/prover.rs @@ -10,13 +10,13 @@ use rayon::prelude::*; use tracing::instrument; use super::{Config, Witness}; -use crate::buffer::{ActiveBuffer, BufferOps}; use crate::{ algebra::{ embedding::Identity, linear_form::{Covector, Evaluate, LinearForm}, mixed_dot, scalar_mul_add, }, + buffer::{ActiveBuffer, BufferOps}, hash::Hash, protocols::{ whir::FinalClaim, From a3933f57f203c62d19c8e02dde3d1a1d56ac5fe7 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Mon, 6 Jul 2026 06:56:22 +0200 Subject: [PATCH 31/35] refactor: change hash_rows function visibility to public --- src/protocols/matrix_commit.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index 66fc4ebe..eaa2ce6d 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -304,7 +304,7 @@ impl fmt::Display for Config { } #[cfg(not(feature = "parallel"))] -fn hash_rows( +pub fn hash_rows( engine: &dyn hash::HashEngine, matrix: &[T], out: &mut [Hash], From 2b5bec862a163590244126fb729c0dc42dbb40b0 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Mon, 6 Jul 2026 06:56:43 +0200 Subject: [PATCH 32/35] feat: add fold_pair_sumcheck_polynomial method to CpuBuffer --- src/buffer/cpu.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/buffer/cpu.rs b/src/buffer/cpu.rs index 7ab781a2..6e789395 100644 --- a/src/buffer/cpu.rs +++ b/src/buffer/cpu.rs @@ -123,6 +123,14 @@ impl Buffer for CpuBuffer { crate::algebra::sumcheck::fold(&mut self.data, weight); } + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + crate::algebra::sumcheck::fold_and_compute_polynomial( + &mut self.data, + &mut other.data, + weight, + ) + } + fn scalar_mul(&mut self, weight: F) { crate::algebra::scalar_mul(&mut self.data, weight); } From b752e1ae9f8bfd4494c41203e5a8581ea30e8e3f Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Mon, 6 Jul 2026 07:10:48 +0200 Subject: [PATCH 33/35] refactor: replace custom Witness struct with merkle_tree::Witness and update related methods --- src/protocols/matrix_commit.rs | 16 +++------------- src/protocols/merkle_tree.rs | 16 +++++++++++----- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index eaa2ce6d..8ccef4f4 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -112,10 +112,7 @@ where pub merkle_tree: merkle_tree::Config, } -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] -pub struct Witness { - pub nodes: ActiveBuffer, -} +pub type Witness = merkle_tree::Witness; pub type Commitment = merkle_tree::Commitment; @@ -226,7 +223,7 @@ impl Config { let (nodes, root) = matrix.merklize(self.num_cols, self.leaf_hash_id, &self.merkle_tree); prover_state.prover_message(&root); - Witness { nodes } + merkle_tree::Witness::new(nodes) } #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] @@ -256,14 +253,7 @@ impl Config { R: RngCore + CryptoRng, Hash: ProverMessage<[H::U]>, { - let node_indices = merkle_tree::opening_sibling_indices( - self.merkle_tree.num_leaves, - self.merkle_tree.layers.len(), - indices, - ); - for hint in witness.nodes.gather_at_indices(&node_indices) { - prover_state.prover_hint(&hint); - } + self.merkle_tree.open(prover_state, witness, indices); } /// Verifies the commitment at the provided row indices. diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index 5cbf6369..dc18e22d 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -12,6 +12,7 @@ use tracing::{instrument, span, Level}; use zerocopy::IntoBytes; use crate::{ + buffer::{ActiveBuffer, BufferOps}, engines::EngineId, hash::{self, Hash, HashEngine, ENGINES}, transcript::{ @@ -58,7 +59,7 @@ pub struct Commitment { #[must_use] pub struct Witness { /// The nodes in the Merkle tree, starting with the leaf hash layer. - pub nodes: Vec, + nodes: ActiveBuffer, } impl Config { @@ -159,8 +160,8 @@ impl Config { assert!(indices.iter().all(|&i| i < self.num_leaves)); let node_indices = opening_sibling_indices(self.num_leaves, self.layers.len(), indices); - for &i in &node_indices { - prover_state.prover_hint(&witness.nodes[i]); + for hint in witness.nodes.gather_at_indices(&node_indices) { + prover_state.prover_hint(&hint); } } @@ -252,7 +253,12 @@ impl Config { } impl Witness { - pub const fn num_nodes(&self) -> usize { + /// Wrap a fully built node buffer (leaf layer first, root last) as a witness. + pub const fn new(nodes: ActiveBuffer) -> Self { + Self { nodes } + } + + pub fn num_nodes(&self) -> usize { self.nodes.len() } } @@ -355,7 +361,7 @@ pub(crate) mod tests { let mut prover_state = ProverState::new_std(&ds); let nodes = config.build_nodes(leaves); prover_state.prover_message(&nodes[config.num_nodes() - 1]); - let witness = Witness { nodes }; + let witness = Witness::new(ActiveBuffer::from_vec(nodes)); config.open(&mut prover_state, &witness, &[13, 42]); let proof = prover_state.proof(); From 437e2679871a1400b3a46c56706acd2a597509ee Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Mon, 6 Jul 2026 07:23:49 +0200 Subject: [PATCH 34/35] refactor: rename as_slice method to to_slice across buffer operations --- src/algebra/ntt/cooley_tukey.rs | 4 ++-- src/algebra/ntt/mod.rs | 2 +- src/buffer/cpu.rs | 6 +++--- src/buffer/mod.rs | 4 ++-- src/protocols/basecase.rs | 18 +++++++++--------- src/protocols/code_switch.rs | 4 ++-- src/protocols/mask_proximity.rs | 10 +++++----- src/protocols/sumcheck.rs | 6 +++--- src/protocols/whir/prover.rs | 2 +- 9 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index b713c33d..67e17f17 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -418,7 +418,7 @@ impl ReedSolomon for NttEngine { let vectors = messages .vectors .iter() - .map(|v| v.as_slice()) + .map(|v| v.to_slice()) .collect::>(); let messages = vectors .iter() @@ -462,7 +462,7 @@ impl ReedSolomon for NttEngine { let mut result = Vec::with_capacity(num_messages * codeword_length); for (message, mask) in zip_strict( messages, - chunks_exact_or_empty(masks.as_slice(), mask_length, num_messages), + chunks_exact_or_empty(masks.to_slice(), mask_length, num_messages), ) { // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension. for _ in 0..num_cosets { diff --git a/src/algebra/ntt/mod.rs b/src/algebra/ntt/mod.rs index 1e5d356d..9deeca44 100644 --- a/src/algebra/ntt/mod.rs +++ b/src/algebra/ntt/mod.rs @@ -244,7 +244,7 @@ mod tests { // Output values are polynomial evaluations in the evaluation points. let mut evaluation_points = ntt.evaluation_points(message_length + mask_length, codeword_length, &sampled_indices); - let codeword = codeword.as_slice(); + let codeword = codeword.to_slice(); for (&index, &evaluation_point) in zip_strict(&sampled_indices, &evaluation_points) { let evaluations = &codeword[index * num_messages.. (index + 1) * num_messages]; let masks = chunks_exact_or_empty(&masks, mask_length, num_messages); diff --git a/src/buffer/cpu.rs b/src/buffer/cpu.rs index 6e789395..85bff3f7 100644 --- a/src/buffer/cpu.rs +++ b/src/buffer/cpu.rs @@ -48,7 +48,7 @@ impl BufferOps for CpuBuffer { } } - fn as_slice(&self) -> &[T] { + fn to_slice(&self) -> &[T] { &self.data } @@ -228,7 +228,7 @@ mod tests { let mut buffer = CpuBuffer::from_vec(values.clone()); buffer.scalar_mul(weight); let expected: Vec = values.iter().map(|&v| v * weight).collect(); - assert_eq!(buffer.as_slice(), expected.as_slice()); + assert_eq!(buffer.to_slice(), expected.as_slice()); } #[test] @@ -252,7 +252,7 @@ mod tests { geometric_accumulate(&mut reference[..size], scalars.clone(), &points); assert_eq!( - buffer.as_slice(), + buffer.to_slice(), reference.as_slice(), "accumulate mismatch for evaluator size {size}" ); diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index fd58aaeb..9f320111 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -3,7 +3,7 @@ //! Protocol code uses the [`ActiveBuffer`] alias to select the backend at //! compile time. Buffers are owned, backend-managed storage: on the CPU //! backend they wrap a `Vec`, on an accelerator backend they would own -//! device memory and only [`BufferOps::as_slice`] (and the other readback +//! device memory and only [`BufferOps::to_slice`] (and the other readback //! methods) force a host copy. //! //! The trait split follows the element type. [`BufferOps`] is generic over @@ -45,7 +45,7 @@ pub trait BufferOps { fn from_vec(source: Vec) -> Self; fn from_slice(source: &[T]) -> Self; /// Read back the buffer contents as a host slice. - fn as_slice(&self) -> &[T]; + fn to_slice(&self) -> &[T]; fn len(&self) -> usize; fn is_empty(&self) -> bool { self.len() == 0 diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index 9c29cdca..b6f622b0 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -80,20 +80,20 @@ impl Config { // Even more trivial non-zk protocol: send f and r directly. if !self.masked { // TODO: avoid these big readbacks even though they are transcript-sized - prover_state.prover_messages(vector.as_slice()); - prover_state.prover_messages(witness.masks.as_slice()); + prover_state.prover_messages(vector.to_slice()); + prover_state.prover_messages(witness.masks.to_slice()); let _ = self.commit.open(prover_state, &[witness]); let point = self .sumcheck .prove(prover_state, &mut vector, &mut covector, &mut sum, &[]) .round_challenges; assert!( - !vector.as_slice().first().expect("Proof failed").is_zero(), + !vector.to_slice().first().expect("Proof failed").is_zero(), "Proof failed" ); return Opening { evaluation_points: point, - linear_form_evaluation: *covector.as_slice().first().expect("Proof failed"), + linear_form_evaluation: *covector.to_slice().first().expect("Proof failed"), }; } @@ -112,14 +112,14 @@ impl Config { assert!(!mask_rlc.is_zero(), "Proof failed"); let mut masked_vector = mask; vector.mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_vector, mask_rlc); - prover_state.prover_messages(masked_vector.as_slice()); + prover_state.prover_messages(masked_vector.to_slice()); // Send combined IRS randomness. (r^* in paper) let mut masked_masks = mask_witness.masks.clone(); witness .masks .mixed_scalar_mul_add_to(&Identity::::new(), &mut masked_masks, mask_rlc); - prover_state.prover_messages(masked_masks.as_slice()); + prover_state.prover_messages(masked_masks.to_slice()); // Open the commitment and mask simultaneously. let _ = self.commit.open(prover_state, &[&mask_witness, witness]); @@ -143,7 +143,7 @@ impl Config { // This event is cryptographically unlikely as `F` is challenge sized. assert!( !masked_vector - .as_slice() + .to_slice() .first() .expect("Proof failed") .is_zero(), @@ -152,7 +152,7 @@ impl Config { Opening { evaluation_points: point, - linear_form_evaluation: *covector.as_slice().first().expect("Proof failed"), + linear_form_evaluation: *covector.to_slice().first().expect("Proof failed"), } } @@ -309,7 +309,7 @@ mod tests { sum, ); assert_eq!( - multilinear_extend(covector.as_slice(), &prover_result.evaluation_points), + multilinear_extend(covector.to_slice(), &prover_result.evaluation_points), prover_result.linear_form_evaluation ); let proof = prover_state.proof(); diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index c7e7d0a6..d6aac94f 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -495,7 +495,7 @@ mod tests { } // Lift ι parallel masks (total length source.mask_length × ι) and fold // chunks of length source.mask_length down to a single chunk. - let raw = lift(config.source.embedding(), source_witness.masks.as_slice()); + let raw = lift(config.source.embedding(), source_witness.masks.to_slice()); let mut mask = fold_chunks(&raw, config.source.mask_length, folding_randomness); // Append fresh padding s of length message_mask_length - source.mask_length. mask.extend(random_vector::( @@ -611,7 +611,7 @@ mod tests { } else { ActiveBuffer::from_vec( folded_message - .as_slice() + .to_slice() .iter() .chain(mask_msg.iter()) .copied() diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index f54f291c..c171d5a5 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -192,7 +192,7 @@ impl Config { // ξ*_i = s_i + γ · ξ_i let mut combined_msg = fresh_msg.clone(); orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); - prover_state.prover_messages(combined_msg.as_slice()); + prover_state.prover_messages(combined_msg.to_slice()); // r*_i = r'_i + γ · r_i // Combined on the host: the result is transcript data anyway, and @@ -200,7 +200,7 @@ impl Config { if irs_masks_per_vector > 0 { let base = i * irs_masks_per_vector; let fresh_base = (self.num_masks + i) * irs_masks_per_vector; - let masks = irs_masks.as_slice(); + let masks = irs_masks.to_slice(); let mut combined_r = masks[fresh_base..fresh_base + irs_masks_per_vector].to_vec(); scalar_mul_add( &mut combined_r, @@ -492,16 +492,16 @@ mod tests { let mut combined_msg = fresh_msg.clone(); orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); if i == 0 { - let mut tampered = combined_msg.as_slice().to_vec(); + let mut tampered = combined_msg.to_slice().to_vec(); tampered[0] += F::ONE; combined_msg = ActiveBuffer::from_slice(&tampered); } - prover_state.prover_messages(combined_msg.as_slice()); + prover_state.prover_messages(combined_msg.to_slice()); if irs_masks_per_vector > 0 { let base = i * irs_masks_per_vector; let fresh_base = (config.num_masks + i) * irs_masks_per_vector; - let masks = irs_masks.as_slice(); + let masks = irs_masks.to_slice(); let mut combined_r = masks[fresh_base..fresh_base + irs_masks_per_vector].to_vec(); scalar_mul_add( &mut combined_r, diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 50df96fa..dec16236 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -316,11 +316,11 @@ mod tests { if config.final_size() == 1 { assert_eq!( multilinear_extend(&initial_vector, &point), - vector.as_slice()[0] + vector.to_slice()[0] ); assert_eq!( multilinear_extend(&initial_covector, &point), - covector.as_slice()[0] + covector.to_slice()[0] ); } else { // TODO: Check correct folding. @@ -333,7 +333,7 @@ mod tests { .sum(); assert_eq!( sum, - expected_mask_sum + mask_rlc * dot(vector.as_slice(), covector.as_slice()) + expected_mask_sum + mask_rlc * dot(vector.to_slice(), covector.to_slice()) ); let proof = prover_state.proof(); diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 4aba15c7..d3dffdc2 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -262,7 +262,7 @@ impl Config { // Directly send the vector to the verifier. assert_eq!(vector.len(), self.final_sumcheck.initial_size); - for coeff in vector.as_slice() { + for coeff in vector.to_slice() { prover_state.prover_message(coeff); } From 70fd1927e9536a80e1a2cec338bf8624d80d1bf5 Mon Sep 17 00:00:00 2001 From: zkfriendly Date: Mon, 6 Jul 2026 07:27:01 +0200 Subject: [PATCH 35/35] fix: as_slice -> to_slice in whir_zk --- src/protocols/whir_zk/committer.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index d44f4c1d..153d9dfe 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -66,7 +66,7 @@ impl Config { let num_blinding_variables = self.num_blinding_variables(); let num_witness_variables = self.num_witness_variables(); for &poly in polynomials { - let poly = poly.as_slice(); + let poly = poly.to_slice(); let blinding = BlindingPolynomials::sample( prover_state.rng(), num_blinding_variables,