From be6839b450620762bc14f5dbb1a4e27078c1a76e Mon Sep 17 00:00:00 2001 From: samkim-crypto Date: Mon, 10 Aug 2026 17:42:04 +0900 Subject: [PATCH 1/2] add poseidon logic --- Cargo.lock | 61 +- Cargo.toml | 2 + syscall/solana-bn254/Cargo.toml | 24 + .../solana-bn254/benches/poseidon_bench.rs | 124 ++++ .../solana-bn254/src/backend/avx512/math.rs | 153 +++++ .../solana-bn254/src/backend/avx512/mod.rs | 28 + .../solana-bn254/src/backend/avx512/pack.rs | 103 ++++ .../solana-bn254/src/backend/avx512/types.rs | 44 ++ syscall/solana-bn254/src/backend/fq.rs | 561 ++++++++++++++++++ syscall/solana-bn254/src/backend/fr.rs | 298 ++++++++++ syscall/solana-bn254/src/backend/mod.rs | 22 + syscall/solana-bn254/src/backend/portable.rs | 192 ++++++ syscall/solana-bn254/src/backend/traits.rs | 58 ++ syscall/solana-bn254/src/backend/u256.rs | 32 + syscall/solana-bn254/src/lib.rs | 14 + syscall/solana-bn254/src/poseidon/mod.rs | 247 ++++++++ 16 files changed, 1954 insertions(+), 9 deletions(-) create mode 100644 syscall/solana-bn254/Cargo.toml create mode 100644 syscall/solana-bn254/benches/poseidon_bench.rs create mode 100644 syscall/solana-bn254/src/backend/avx512/math.rs create mode 100644 syscall/solana-bn254/src/backend/avx512/mod.rs create mode 100644 syscall/solana-bn254/src/backend/avx512/pack.rs create mode 100644 syscall/solana-bn254/src/backend/avx512/types.rs create mode 100644 syscall/solana-bn254/src/backend/fq.rs create mode 100644 syscall/solana-bn254/src/backend/fr.rs create mode 100644 syscall/solana-bn254/src/backend/mod.rs create mode 100644 syscall/solana-bn254/src/backend/portable.rs create mode 100644 syscall/solana-bn254/src/backend/traits.rs create mode 100644 syscall/solana-bn254/src/backend/u256.rs create mode 100644 syscall/solana-bn254/src/lib.rs create mode 100644 syscall/solana-bn254/src/poseidon/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 829627d..5cb74fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -203,7 +203,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "246a225cc6131e9ee4f24619af0f19d67761fff15d7ccc22e42b80846e69449a" dependencies = [ "num-traits", - "rand 0.8.6", + "rand 0.8.7", ] [[package]] @@ -964,7 +964,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" dependencies = [ "ff 0.13.1", - "rand 0.8.6", + "rand 0.8.7", "rand_core 0.6.4", "rand_xorshift", "subtle", @@ -1126,6 +1126,18 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "light-poseidon" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47a1ccadd0bb5a32c196da536fd72c59183de24a055f6bf0513bf845fefab862" +dependencies = [ + "ark-bn254", + "ark-ff", + "num-bigint", + "thiserror 1.0.69", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1363,7 +1375,7 @@ dependencies = [ "rand 0.10.2", "serde", "spin", - "thiserror", + "thiserror 2.0.18", "tracing", ] @@ -1447,7 +1459,7 @@ dependencies = [ "rand 0.10.2", "serde", "spin", - "thiserror", + "thiserror 2.0.18", "tracing", ] @@ -1545,7 +1557,7 @@ dependencies = [ "p3-maybe-rayon", "p3-util", "serde", - "thiserror", + "thiserror 2.0.18", "tracing", ] @@ -1691,7 +1703,7 @@ dependencies = [ "bitflags", "lazy_static", "num-traits", - "rand 0.8.6", + "rand 0.8.7", "rand_chacha 0.3.1", "rand_xorshift", "regex-syntax", @@ -1735,9 +1747,9 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.8.6" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", "rand_chacha 0.3.1", @@ -2055,6 +2067,17 @@ dependencies = [ "pairing", ] +[[package]] +name = "solana-bn254" +version = "0.1.0" +dependencies = [ + "ark-bn254", + "ark-ff", + "criterion", + "light-poseidon", + "rand 0.10.2", +] + [[package]] name = "solana-bn254-syscall" version = "0.1.0" @@ -2177,13 +2200,33 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index fe0414b..dbcca84 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ members = [ "secp256r1", "syscall/bls12-381-syscall", "syscall/bn254-syscall", + "syscall/solana-bn254", ] resolver = "2" @@ -43,6 +44,7 @@ hashbrown = "0.15" hex = "0.4.2" hex-literal = "1.1.0" hmac = "0.12" +light-poseidon = "0.4.0" openssl = "0.10" pairing = "0.23.0" p256 = { version = "0.13", default-features = false, features = ["arithmetic", "expose-field"] } diff --git a/syscall/solana-bn254/Cargo.toml b/syscall/solana-bn254/Cargo.toml new file mode 100644 index 0000000..bcd567e --- /dev/null +++ b/syscall/solana-bn254/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "solana-bn254" +description = "BN254 and Poseidon for Solana" +version = "0.1.0" +edition = "2024" +authors = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +readme = "README.md" + +[features] + +[dependencies] + +[dev-dependencies] +criterion = { workspace = true } +rand = { workspace = true } +ark-ff = { workspace = true } +ark-bn254 = { workspace = true } +light-poseidon = { workspace = true } + +[[bench]] +name = "poseidon_bench" +harness = false diff --git a/syscall/solana-bn254/benches/poseidon_bench.rs b/syscall/solana-bn254/benches/poseidon_bench.rs new file mode 100644 index 0000000..cfd44e8 --- /dev/null +++ b/syscall/solana-bn254/benches/poseidon_bench.rs @@ -0,0 +1,124 @@ +use ark_ff::PrimeField; +use criterion::{Criterion, criterion_group, criterion_main}; +use light_poseidon::PoseidonHasher; +use rand::RngExt; // Required for the `.random()` trait method +use solana_bn254::{ + backend::U256, + poseidon::{PoseidonConstants, SparseMatrix}, +}; + +/// Generates a randomized field element safely constrained below the BN254 Fr modulus. +fn random_fr() -> U256 { + let mut rng = rand::rng(); + let mut limbs: [u64; 4] = rng.random(); + // Mask the top u64 to guarantee it stays strictly below Fr::MODULUS + limbs[3] &= 0x0FFFFFFFFFFFFFFF; + U256::new(limbs) +} + +/// Dynamically leaks a set of dummy parameters onto the heap. +/// Because Criterion runs this outside the hot-loop, the `O(1)` memory +/// leak per benchmark group is completely acceptable and ensures +/// we can provide the `'static` lifetimes required by the crate. +fn make_dummy_constants(partial_rounds: usize) -> &'static PoseidonConstants { + let num_round_constants = (9 * T) + partial_rounds; + let round_constants = (0..num_round_constants) + .map(|_| random_fr()) + .collect::>() + .leak(); + + let mut mds_matrix = [[U256::zero(); T]; T]; + for row in mds_matrix.iter_mut() { + for val in row.iter_mut() { + *val = random_fr(); + } + } + let mds_matrix = Box::leak(Box::new(mds_matrix)); + + let mut pre_sparse_matrix = [[U256::zero(); T]; T]; + for row in pre_sparse_matrix.iter_mut() { + for val in row.iter_mut() { + *val = random_fr(); + } + } + let pre_sparse_matrix = Box::leak(Box::new(pre_sparse_matrix)); + + let sparse_matrices = (0..partial_rounds) + .map(|_| { + let mut row = [U256::zero(); T]; + let mut col = [U256::zero(); T]; + for i in 0..T { + row[i] = random_fr(); + col[i] = random_fr(); + } + SparseMatrix { row, col } + }) + .collect::>() + .leak(); + + Box::leak(Box::new(PoseidonConstants { + full_rounds: 8, + partial_rounds, + round_constants, + mds_matrix, + pre_sparse_matrix, + sparse_matrices, + })) +} + +macro_rules! bench_width { + ($c:expr, $t:literal, $partial_rounds:literal) => { + let constants = make_dummy_constants::<$t>($partial_rounds); + + let mut group = $c.benchmark_group(concat!("poseidon_t", $t)); + let nr_inputs = $t - 1; + + // --- solana-bn254 --- + let state = core::array::from_fn(|_| random_fr()); + group.bench_function("solana-bn254", |b| { + b.iter(|| { + // Uses std::hint::black_box instead of deprecated criterion::black_box + solana_bn254::poseidon::poseidon(std::hint::black_box(state), constants) + }) + }); + + // --- light-poseidon --- + // Note: light-poseidon handles its own internal constant generation + // via `new_circom()`. We generate inputs from randomized byte arrays + // here, but because this happens outside the `b.iter()` block, it does + // not penalize the light-poseidon benchmark times. + let mut hasher = light_poseidon::Poseidon::::new_circom(nr_inputs).unwrap(); + let mut rng = rand::rng(); + let inputs: Vec<_> = (0..nr_inputs) + .map(|_| ark_bn254::Fr::from_be_bytes_mod_order(&rng.random::<[u8; 32]>())) + .collect(); + + group.bench_function("light-poseidon", |b| { + b.iter(|| hasher.hash(std::hint::black_box(&inputs)).unwrap()) + }); + + group.finish(); + }; +} + +fn bench_poseidon(c: &mut Criterion) { + // Solana `sol_poseidon` syscall supported parameters. + // These specific partial round counts are mandated by the SVM + // for state widths t=2..13 (which map to 1..12 inputs). + // [56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65] + bench_width!(c, 2, 56); + bench_width!(c, 3, 57); + bench_width!(c, 4, 56); + bench_width!(c, 5, 60); + bench_width!(c, 6, 60); + bench_width!(c, 7, 63); + bench_width!(c, 8, 64); + bench_width!(c, 9, 63); + bench_width!(c, 10, 60); + bench_width!(c, 11, 66); + bench_width!(c, 12, 60); + bench_width!(c, 13, 65); +} + +criterion_group!(benches, bench_poseidon); +criterion_main!(benches); diff --git a/syscall/solana-bn254/src/backend/avx512/math.rs b/syscall/solana-bn254/src/backend/avx512/math.rs new file mode 100644 index 0000000..5b64d34 --- /dev/null +++ b/syscall/solana-bn254/src/backend/avx512/math.rs @@ -0,0 +1,153 @@ +//! High-throughput Montgomery CIOS IFMA Arithmetic. + +#![allow(unused_unsafe)] +#![allow(unsafe_op_in_unsafe_fn)] + +use super::types::FieldElement8x52; +use core::arch::x86_64::*; + +// Mathematically pre-computed 52-bit modulus constants for the BN254 Fr field. +const FR_MOD_L0: i64 = 0x1f593f0000001; +const FR_MOD_L1: i64 = 0x4879b9709143e; +const FR_MOD_L2: i64 = 0x181585d2833e8; +const FR_MOD_L3: i64 = 0xa029b85045b68; +const FR_MOD_L4: i64 = 0x30644e72e131; + +// The Montgomery Inverse Multiplier for 52-bit limbs: `(-MODULUS^-1) mod 2^52` +const FR_INV_52: i64 = 0x1f593efffffff; + +/// Unreduced Parallel Addition. +/// +/// Because IFMA math logically isolates results within a 52-bit boundary inside a +/// 64-bit accumulator lane, we naturally gain 12 bits of headroom. This function +/// can be called ~4,096 times consecutively before an overflow is mathematically +/// possible, making it extremely efficient for massive MDS matrix multiplications. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn add_lazy(a: &FieldElement8x52, b: &FieldElement8x52) -> FieldElement8x52 { + FieldElement8x52 { + l0: _mm512_add_epi64(a.l0, b.l0), + l1: _mm512_add_epi64(a.l1, b.l1), + l2: _mm512_add_epi64(a.l2, b.l2), + l3: _mm512_add_epi64(a.l3, b.l3), + l4: _mm512_add_epi64(a.l4, b.l4), + } +} + +/// Computes an 8-way parallel Montgomery Multiplication for the BN254 Fr field. +/// +/// Implements a fully unrolled CIOS (Coarsely Integrated Operand Scanning) algorithm. +/// Because IFMA tracks sums in a 64-bit accumulator, cross-product carries are +/// strictly contained within the active iteration and do not require ripple logic +/// between intermediate multiplies. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn mul_8x(a: &FieldElement8x52, b: &FieldElement8x52) -> FieldElement8x52 { + // 64-bit accumulators holding the in-flight summation. + let mut t = [_mm512_setzero_si512(); 6]; + + // Broadcast the 52-bit field constants into the AVX-512 lanes. + let inv_vec = _mm512_set1_epi64(FR_INV_52); + let mod0 = _mm512_set1_epi64(FR_MOD_L0); + let mod1 = _mm512_set1_epi64(FR_MOD_L1); + let mod2 = _mm512_set1_epi64(FR_MOD_L2); + let mod3 = _mm512_set1_epi64(FR_MOD_L3); + let mod4 = _mm512_set1_epi64(FR_MOD_L4); + + let a_limbs = [a.l0, a.l1, a.l2, a.l3, a.l4]; + + // CIOS Algorithm: Loop is fully unrolled by the LLVM compiler. + for i in 0..5 { + let ai = a_limbs[i]; + + // 1. Accumulate the multiplication of `ai` against all limbs of `b`. + // `_mm512_madd52lo` adds the lower 52 bits of the product into the accumulator. + // `_mm512_madd52hi` adds the upper 52 bits of the product into the accumulator. + t[0] = _mm512_madd52lo_epu64(t[0], ai, b.l0); + t[1] = _mm512_madd52hi_epu64(t[1], ai, b.l0); + + t[1] = _mm512_madd52lo_epu64(t[1], ai, b.l1); + t[2] = _mm512_madd52hi_epu64(t[2], ai, b.l1); + + t[2] = _mm512_madd52lo_epu64(t[2], ai, b.l2); + t[3] = _mm512_madd52hi_epu64(t[3], ai, b.l2); + + t[3] = _mm512_madd52lo_epu64(t[3], ai, b.l3); + t[4] = _mm512_madd52hi_epu64(t[4], ai, b.l3); + + t[4] = _mm512_madd52lo_epu64(t[4], ai, b.l4); + t[5] = _mm512_madd52hi_epu64(t[5], ai, b.l4); + + // 2. Compute Montgomery Multiplier: `m = (t[0] * INV) mod 2^52` + // `madd52lo` automatically masks the inputs to 52 bits and ignores the high bits. + let m = _mm512_madd52lo_epu64(_mm512_setzero_si512(), t[0], inv_vec); + + // 3. Accumulate Reduction: `t += m * Modulus` + // Mathematically forces the bottom 52 bits of t[0] to exactly 0. + t[0] = _mm512_madd52lo_epu64(t[0], m, mod0); + t[1] = _mm512_madd52hi_epu64(t[1], m, mod0); + + t[1] = _mm512_madd52lo_epu64(t[1], m, mod1); + t[2] = _mm512_madd52hi_epu64(t[2], m, mod1); + + t[2] = _mm512_madd52lo_epu64(t[2], m, mod2); + t[3] = _mm512_madd52hi_epu64(t[3], m, mod2); + + t[3] = _mm512_madd52lo_epu64(t[3], m, mod3); + t[4] = _mm512_madd52hi_epu64(t[4], m, mod3); + + t[4] = _mm512_madd52lo_epu64(t[4], m, mod4); + t[5] = _mm512_madd52hi_epu64(t[5], m, mod4); + + // 4. Register shift down. Since the bottom 52 bits of t[0] are zero, we extract + // the top carry and add it into the next limb, then rotate the array. + let carry = _mm512_srli_epi64(t[0], 52); + t[1] = _mm512_add_epi64(t[1], carry); + + t[0] = t[1]; + t[1] = t[2]; + t[2] = t[3]; + t[3] = t[4]; + t[4] = t[5]; + t[5] = _mm512_setzero_si512(); + } + + // --- 5. Final Carry Propagation --- + // At the end of the CIOS loop, we strictly enforce the 52-bit boundaries + // by propagating any overflowing bits from the 64-bit accumulators upwards. + let mask_52 = _mm512_set1_epi64(0xFFFFFFFFFFFFF); + let mut out = FieldElement8x52::zero(); + + let carry0 = _mm512_srli_epi64(t[0], 52); + out.l0 = _mm512_and_si512(t[0], mask_52); + + let t1_new = _mm512_add_epi64(t[1], carry0); + let carry1 = _mm512_srli_epi64(t1_new, 52); + out.l1 = _mm512_and_si512(t1_new, mask_52); + + let t2_new = _mm512_add_epi64(t[2], carry1); + let carry2 = _mm512_srli_epi64(t2_new, 52); + out.l2 = _mm512_and_si512(t2_new, mask_52); + + let t3_new = _mm512_add_epi64(t[3], carry2); + let carry3 = _mm512_srli_epi64(t3_new, 52); + out.l3 = _mm512_and_si512(t3_new, mask_52); + + let t4_new = _mm512_add_epi64(t[4], carry3); + out.l4 = _mm512_and_si512(t4_new, mask_52); + + out +} + +/// Executes the Poseidon S-box (`x^5`) on 8 independent field elements simultaneously. +/// +/// Because IFMA CIOS handles multiplication entirely in registers without branching +/// or memory spilling, this collapses what would normally be 24 scalar multiplications +/// into just 3 massive parallel SIMD dispatches. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn sbox_8x(x: &FieldElement8x52) -> FieldElement8x52 { + let x2 = mul_8x(x, x); + let x4 = mul_8x(&x2, &x2); + mul_8x(&x4, x) +} diff --git a/syscall/solana-bn254/src/backend/avx512/mod.rs b/syscall/solana-bn254/src/backend/avx512/mod.rs new file mode 100644 index 0000000..777f142 --- /dev/null +++ b/syscall/solana-bn254/src/backend/avx512/mod.rs @@ -0,0 +1,28 @@ +//! AVX-512 IFMA batched execution engine (8-way parallel). +//! +//! This module leverages the `vpmadd52` instructions (Integer Fused Multiply-Add) +//! introduced in AVX-512 IFMA. It is strictly gated to supported hardware +//! (e.g., AMD Zen 4 EPYC processors or Intel Ice Lake). +//! +//! # 52-Bit Architecture +//! Standard 256-bit integers are represented as four 64-bit limbs. However, +//! accumulating 64-bit multiplications natively overflows the register, forcing +//! expensive software carry-propagation logic. +//! +//! By repacking the 256-bit state into five 52-bit limbs (`5 \times 52 = 260` bits), +//! we leave 12 bits of "headroom" at the top of each 64-bit SIMD lane. This allows +//! us to natively accumulate up to `2^{12} = 4096` multiplications locally within the +//! register without ever doing a Montgomery reduction or resolving carries. +//! +//! This "lazy reduction" unlocks world-class cryptographic throughput for operations +//! like Poseidon dense matrix multiplications and Batched Multi-Miller Loops. + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +pub mod math; +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +pub mod pack; +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +pub mod types; + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +pub use types::FieldElement8x52; diff --git a/syscall/solana-bn254/src/backend/avx512/pack.rs b/syscall/solana-bn254/src/backend/avx512/pack.rs new file mode 100644 index 0000000..4d3a3af --- /dev/null +++ b/syscall/solana-bn254/src/backend/avx512/pack.rs @@ -0,0 +1,103 @@ +//! Conversion bridge between public 4-limb arrays and internal 5-limb SIMD vectors. + +#![allow(unused_unsafe)] +#![allow(unsafe_op_in_unsafe_fn)] + +use super::types::FieldElement8x52; +use crate::backend::U256; +use core::arch::x86_64::{_mm512_loadu_si512, _mm512_set1_epi64, _mm512_storeu_si512}; + +/// Constant bitmask used to isolate exactly 52 bits during packing boundaries. +const MASK_52: u64 = 0xFFFFFFFFFFFFF; + +/// Packs an array of 8 scalar 4-limb (64-bit) U256s into a single 5-limb (52-bit) AVX-512 vector. +/// +/// This acts as the zero-cost API boundary, safely translating the public 64-bit +/// little-endian storage layouts into our internal 52-bit SIMD math engine. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn pack_8x(inputs: &[U256; 8]) -> FieldElement8x52 { + let mut l0 = [0u64; 8]; + let mut l1 = [0u64; 8]; + let mut l2 = [0u64; 8]; + let mut l3 = [0u64; 8]; + let mut l4 = [0u64; 8]; + + // Iteratively bit-slice the 4x64-bit boundaries into 5x52-bit boundaries. + for i in 0..8 { + let u = &inputs[i].0; + l0[i] = u[0] & MASK_52; + l1[i] = ((u[0] >> 52) | (u[1] << 12)) & MASK_52; + l2[i] = ((u[1] >> 40) | (u[2] << 24)) & MASK_52; + l3[i] = ((u[2] >> 28) | (u[3] << 36)) & MASK_52; + l4[i] = u[3] >> 16; + } + + FieldElement8x52 { + l0: _mm512_loadu_si512(l0.as_ptr() as *const _), + l1: _mm512_loadu_si512(l1.as_ptr() as *const _), + l2: _mm512_loadu_si512(l2.as_ptr() as *const _), + l3: _mm512_loadu_si512(l3.as_ptr() as *const _), + l4: _mm512_loadu_si512(l4.as_ptr() as *const _), + } +} + +/// Unpacks a 5-limb AVX-512 vector into a mutable array of 8 scalar U256s. +/// +/// Reconstructs the original base `2^{64}` boundaries by masking and shifting +/// the processed 52-bit SIMD lanes. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn unpack_8x_into(packed: &FieldElement8x52, outputs: &mut [U256; 8]) { + let mut l0 = [0u64; 8]; + let mut l1 = [0u64; 8]; + let mut l2 = [0u64; 8]; + let mut l3 = [0u64; 8]; + let mut l4 = [0u64; 8]; + + _mm512_storeu_si512(l0.as_mut_ptr() as *mut _, packed.l0); + _mm512_storeu_si512(l1.as_mut_ptr() as *mut _, packed.l1); + _mm512_storeu_si512(l2.as_mut_ptr() as *mut _, packed.l2); + _mm512_storeu_si512(l3.as_mut_ptr() as *mut _, packed.l3); + _mm512_storeu_si512(l4.as_mut_ptr() as *mut _, packed.l4); + + // Reconstruct 64-bit boundaries by reversing the shifts. + for i in 0..8 { + let u0 = l0[i] | (l1[i] << 52); + let u1 = (l1[i] >> 12) | (l2[i] << 40); + let u2 = (l2[i] >> 24) | (l3[i] << 28); + let u3 = (l3[i] >> 36) | (l4[i] << 16); + outputs[i] = U256::new([u0, u1, u2, u3]); + } +} + +/// Allocates and unpacks a 5-limb AVX-512 vector into an array of 8 scalar U256s. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn unpack_8x(packed: &FieldElement8x52) -> [U256; 8] { + let mut out = [U256::zero(); 8]; + unpack_8x_into(packed, &mut out); + out +} + +/// Broadcasts a single scalar 4-limb value identically to all 8 lanes of a SIMD register. +/// +/// Crucial for Poseidon matrix multiplication, where a single state element must +/// be multiplied against an entire SIMD chunk of the matrix column vector. +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +pub unsafe fn broadcast(u: &U256) -> FieldElement8x52 { + let l0 = u.0[0] & MASK_52; + let l1 = ((u.0[0] >> 52) | (u.0[1] << 12)) & MASK_52; + let l2 = ((u.0[1] >> 40) | (u.0[2] << 24)) & MASK_52; + let l3 = ((u.0[2] >> 28) | (u.0[3] << 36)) & MASK_52; + let l4 = u.0[3] >> 16; + + FieldElement8x52 { + l0: _mm512_set1_epi64(l0 as i64), + l1: _mm512_set1_epi64(l1 as i64), + l2: _mm512_set1_epi64(l2 as i64), + l3: _mm512_set1_epi64(l3 as i64), + l4: _mm512_set1_epi64(l4 as i64), + } +} diff --git a/syscall/solana-bn254/src/backend/avx512/types.rs b/syscall/solana-bn254/src/backend/avx512/types.rs new file mode 100644 index 0000000..dd4c0c4 --- /dev/null +++ b/syscall/solana-bn254/src/backend/avx512/types.rs @@ -0,0 +1,44 @@ +//! Fundamental data layout for AVX-512 batched cryptography. + +use core::arch::x86_64::__m512i; + +/// Holds 8 independent 256-bit field elements in parallel, packed into 52-bit limbs. +/// +/// A standard 256-bit scalar is mathematically represented as: +/// `X = l0 + (l1 << 52) + (l2 << 104) + (l3 << 156) + (l4 << 208)` +/// +/// We use a **Struct of Arrays (SoA)** layout (Vertical Vectorization). +/// Instead of one SIMD register holding the limbs of a single field element, +/// each `__m512i` register holds the *same* limb slice across 8 distinct elements. +/// This perfectly isolates carries within vertical lanes, allowing 100% horizontal +/// lane utilization. +#[derive(Copy, Clone)] +pub struct FieldElement8x52 { + /// Limb 0 (Bits 0-51) for elements 0 through 7. + pub l0: __m512i, + /// Limb 1 (Bits 52-103) for elements 0 through 7. + pub l1: __m512i, + /// Limb 2 (Bits 104-155) for elements 0 through 7. + pub l2: __m512i, + /// Limb 3 (Bits 156-207) for elements 0 through 7. + pub l3: __m512i, + /// Limb 4 (Bits 208-255) for elements 0 through 7. (Effectively 48-bits wide). + pub l4: __m512i, +} + +impl FieldElement8x52 { + /// Returns an instance where all 8 field elements are strictly zero. + #[inline(always)] + pub fn zero() -> Self { + unsafe { + use core::arch::x86_64::_mm512_setzero_si512; + Self { + l0: _mm512_setzero_si512(), + l1: _mm512_setzero_si512(), + l2: _mm512_setzero_si512(), + l3: _mm512_setzero_si512(), + l4: _mm512_setzero_si512(), + } + } + } +} diff --git a/syscall/solana-bn254/src/backend/fq.rs b/syscall/solana-bn254/src/backend/fq.rs new file mode 100644 index 0000000..f78fda2 --- /dev/null +++ b/syscall/solana-bn254/src/backend/fq.rs @@ -0,0 +1,561 @@ +//! BN254 Base Field (Fq) configuration and static addition chains. + +use super::{Backend, Field, MontgomeryBackend, U256}; + +/// BN254 Base Field parameters. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] +pub struct Fq; + +impl Field for Fq { + const MODULUS: U256 = U256::new([ + 0x3c208c16d87cfd47, + 0x97816a916871ca8d, + 0xb85045b68181585d, + 0x30644e72e131a029, + ]); + const INV: u64 = 0x87d20782e4866389; + const R2: U256 = U256::new([ + 0xf32cfc5b538afa89, + 0xb5e71911d44501fb, + 0x47ab1eff0a417ff6, + 0x06d89f71cab8351f, + ]); +} + +impl Fq { + /// Computes `a^(p-2) mod p` using a static 4-bit window addition chain. + #[inline(always)] + pub fn invert(a: &U256) -> U256 { + type B = Backend; + let a2 = B::sqr(a); + let a3 = B::mul(&a2, a); + let a4 = B::sqr(&a2); + let a5 = B::mul(&a4, a); + let a6 = B::sqr(&a3); + let a7 = B::mul(&a6, a); + let a8 = B::sqr(&a4); + let a9 = B::mul(&a8, a); + let a10 = B::sqr(&a5); + let a11 = B::mul(&a10, a); + let a12 = B::sqr(&a6); + let a13 = B::mul(&a12, a); + let a14 = B::sqr(&a7); + let a15 = B::mul(&a14, a); + + let mut t = a3; + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a13); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a13); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a13); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a13); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + t + } + + /// Computes `a^((p+1)/4) mod p` using a static 4-bit window addition chain. + #[inline(always)] + pub fn sqrt(a: &U256) -> U256 { + type B = Backend; + let a2 = B::sqr(a); + let a3 = B::mul(&a2, a); + let a4 = B::sqr(&a2); + let a5 = B::mul(&a4, a); + let a6 = B::sqr(&a3); + let a7 = B::mul(&a6, a); + let a8 = B::sqr(&a4); + let a9 = B::mul(&a8, a); + let a10 = B::sqr(&a5); + let a11 = B::mul(&a10, a); + let a12 = B::sqr(&a6); + let a13 = B::mul(&a12, a); + let a14 = B::sqr(&a7); + let a15 = B::mul(&a14, a); + + let mut t = a12; + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a13); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a12); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + t + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_fq_invert() { + type B = Backend; + let val = B::to_mont(&U256::new([5, 0, 0, 0])); + let inv = Fq::invert(&val); + let check = B::mul(&val, &inv); + assert_eq!(B::from_mont(&check), U256::one()); + + // Zero inversion maps to zero in prime fields via Fermat's Little Thm + let zero = U256::zero(); + assert_eq!(Fq::invert(&zero), U256::zero()); + } + + #[test] + fn test_fq_sqrt() { + type B = Backend; + // sqrt(25) = 5 + let sqr = B::to_mont(&U256::new([25, 0, 0, 0])); + let root = Fq::sqrt(&sqr); + + // Ensure our root squares exactly back to the original value + let check = B::sqr(&root); + assert_eq!(B::from_mont(&check), U256::new([25, 0, 0, 0])); + } +} diff --git a/syscall/solana-bn254/src/backend/fr.rs b/syscall/solana-bn254/src/backend/fr.rs new file mode 100644 index 0000000..f7303ef --- /dev/null +++ b/syscall/solana-bn254/src/backend/fr.rs @@ -0,0 +1,298 @@ +//! BN254 Scalar Field (Fr) configuration and static addition chains. + +use super::{Backend, Field, MontgomeryBackend, U256}; + +/// BN254 Scalar Field parameters. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] +pub struct Fr; + +impl Field for Fr { + const MODULUS: U256 = U256::new([ + 0x43e1f593f0000001, + 0x2833e84879b97091, + 0xb85045b68181585d, + 0x30644e72e131a029, + ]); + const INV: u64 = 0xc2e1f593efffffff; + const R2: U256 = U256::new([ + 0x1bb8e645ae216da7, + 0x53fe3ab1e35c59e3, + 0x8c49833d53bb8085, + 0x0216d0b17f4e44a5, + ]); +} + +impl Fr { + /// Computes `a^(p-2) mod p` using a static 4-bit window addition chain. + #[inline(always)] + pub fn invert(a: &U256) -> U256 { + type B = Backend; + let a2 = B::sqr(a); + let a3 = B::mul(&a2, a); + let a4 = B::sqr(&a2); + let a5 = B::mul(&a4, a); + let a6 = B::sqr(&a3); + let a7 = B::mul(&a6, a); + let a8 = B::sqr(&a4); + let a9 = B::mul(&a8, a); + let a10 = B::sqr(&a5); + let a11 = B::mul(&a10, a); + let a12 = B::sqr(&a6); + let a13 = B::mul(&a12, a); + let a14 = B::sqr(&a7); + let a15 = B::mul(&a14, a); + + let mut t = a3; + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a10); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a6); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a13); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a2); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a8); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a11); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a7); + for _ in 0..8 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a4); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, a); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a5); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a9); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a3); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a14); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + for _ in 0..4 { + t = B::sqr(&t); + } + t = B::mul(&t, &a15); + t + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_fr_invert() { + type B = Backend; + let val = B::to_mont(&U256::new([5, 0, 0, 0])); + let inv = Fr::invert(&val); + let check = B::mul(&val, &inv); + assert_eq!(B::from_mont(&check), U256::one()); + } +} diff --git a/syscall/solana-bn254/src/backend/mod.rs b/syscall/solana-bn254/src/backend/mod.rs new file mode 100644 index 0000000..f3c7d1e --- /dev/null +++ b/syscall/solana-bn254/src/backend/mod.rs @@ -0,0 +1,22 @@ +//! Unified 256-bit field arithmetic for Fq (base) and Fr (scalar). +//! +//! Handles core Montgomery arithmetic. Execution routes to either +//! pure Rust multi-limb operations or highly optimized x86_64 +//! assembly based on active feature flags. + +pub mod fq; +pub mod fr; +pub mod traits; +pub mod u256; + +pub mod portable; + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +pub mod avx512; + +pub use fq::Fq; +pub use fr::Fr; +pub use traits::{Field, MontgomeryBackend}; +pub use u256::U256; + +pub type Backend = portable::PortableBackend; diff --git a/syscall/solana-bn254/src/backend/portable.rs b/syscall/solana-bn254/src/backend/portable.rs new file mode 100644 index 0000000..98c1717 --- /dev/null +++ b/syscall/solana-bn254/src/backend/portable.rs @@ -0,0 +1,192 @@ +//! Pure Rust multi-limb arithmetic fallback. +//! +//! Provides a `MontgomeryBackend` using 64-bit limb arithmetic +//! without assuming x86_64 intrinsics. + +use super::{Field, MontgomeryBackend, U256}; +use core::marker::PhantomData; + +/// Computes `a + b + carry`, returning `(result, carry_out)`. +#[inline(always)] +const fn adc(a: u64, b: u64, carry: u64) -> (u64, u64) { + let res = (a as u128) + (b as u128) + (carry as u128); + (res as u64, (res >> 64) as u64) +} + +/// Computes `a - b - borrow`, returning `(result, borrow_out)`. +#[inline(always)] +const fn sbb(a: u64, b: u64, borrow: u64) -> (u64, u64) { + let (r1, b1) = a.overflowing_sub(b); + let (r2, b2) = r1.overflowing_sub(borrow); + (r2, (b1 as u64) | (b2 as u64)) +} + +/// Computes `a + (b * c) + carry`, returning `(result, carry_out)`. +#[inline(always)] +const fn mac(a: u64, b: u64, c: u64, carry: u64) -> (u64, u64) { + let res = (a as u128) + ((b as u128) * (c as u128)) + (carry as u128); + (res as u64, (res >> 64) as u64) +} + +/// A portable, pure-Rust backend for Montgomery arithmetic. +pub struct PortableBackend(PhantomData); + +impl MontgomeryBackend for PortableBackend { + #[inline(always)] + fn add(a: &U256, b: &U256) -> U256 { + let (r0, c) = adc(a.0[0], b.0[0], 0); + let (r1, c) = adc(a.0[1], b.0[1], c); + let (r2, c) = adc(a.0[2], b.0[2], c); + let (r3, c) = adc(a.0[3], b.0[3], c); + + let (d0, br) = sbb(r0, F::MODULUS.0[0], 0); + let (d1, br) = sbb(r1, F::MODULUS.0[1], br); + let (d2, br) = sbb(r2, F::MODULUS.0[2], br); + let (d3, br) = sbb(r3, F::MODULUS.0[3], br); + + // If no carry out and subtraction underflowed, a + b < MODULUS. + if c == 0 && br == 1 { + U256::new([r0, r1, r2, r3]) + } else { + U256::new([d0, d1, d2, d3]) + } + } + + #[inline(always)] + fn sub(a: &U256, b: &U256) -> U256 { + let (r0, br) = sbb(a.0[0], b.0[0], 0); + let (r1, br) = sbb(a.0[1], b.0[1], br); + let (r2, br) = sbb(a.0[2], b.0[2], br); + let (r3, br) = sbb(a.0[3], b.0[3], br); + + if br > 0 { + // Underflow occurred, add the modulus back to wrap around. + let (d0, c) = adc(r0, F::MODULUS.0[0], 0); + let (d1, c) = adc(r1, F::MODULUS.0[1], c); + let (d2, c) = adc(r2, F::MODULUS.0[2], c); + let (d3, _) = adc(r3, F::MODULUS.0[3], c); + U256::new([d0, d1, d2, d3]) + } else { + U256::new([r0, r1, r2, r3]) + } + } + + #[inline(always)] + fn mul(a: &U256, b: &U256) -> U256 { + let mut t = [0u64; 5]; + + for i in 0..4 { + let (r0, c) = mac(t[0], a.0[i], b.0[0], 0); + let (r1, c) = mac(t[1], a.0[i], b.0[1], c); + let (r2, c) = mac(t[2], a.0[i], b.0[2], c); + let (r3, c) = mac(t[3], a.0[i], b.0[3], c); + let (r4, r5) = adc(t[4], 0, c); + + let m = r0.wrapping_mul(F::INV); + + let (_, c2) = mac(r0, m, F::MODULUS.0[0], 0); + let (n0, c2) = mac(r1, m, F::MODULUS.0[1], c2); + let (n1, c2) = mac(r2, m, F::MODULUS.0[2], c2); + let (n2, c2) = mac(r3, m, F::MODULUS.0[3], c2); + let (n3, c2) = adc(r4, 0, c2); + + t[0] = n0; + t[1] = n1; + t[2] = n2; + t[3] = n3; + t[4] = r5 + c2; // Overflow impossible here + } + + let (d0, br) = sbb(t[0], F::MODULUS.0[0], 0); + let (d1, br) = sbb(t[1], F::MODULUS.0[1], br); + let (d2, br) = sbb(t[2], F::MODULUS.0[2], br); + let (d3, br) = sbb(t[3], F::MODULUS.0[3], br); + + if t[4] == 0 && br == 1 { + U256::new([t[0], t[1], t[2], t[3]]) + } else { + U256::new([d0, d1, d2, d3]) + } + } + + #[inline(always)] + fn sqr(a: &U256) -> U256 { + Self::mul(a, a) + } + + #[inline(always)] + fn neg(a: &U256) -> U256 { + if a.0[0] == 0 && a.0[1] == 0 && a.0[2] == 0 && a.0[3] == 0 { + *a + } else { + let (d0, br) = sbb(F::MODULUS.0[0], a.0[0], 0); + let (d1, br) = sbb(F::MODULUS.0[1], a.0[1], br); + let (d2, br) = sbb(F::MODULUS.0[2], a.0[2], br); + let (d3, _) = sbb(F::MODULUS.0[3], a.0[3], br); + U256::new([d0, d1, d2, d3]) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// BN254 scalar field (Fr) parameters for exact correctness testing. + struct TestFr; + + impl Field for TestFr { + const MODULUS: U256 = U256::new([ + 0x43e1f593f0000001, + 0x2833e84879b97091, + 0xb85045b68181585d, + 0x30644e72e131a029, + ]); + const INV: u64 = 0xc2e1f593efffffff; + const R2: U256 = U256::new([ + 0x1bb8e645ae216da7, + 0x53fe3ab1e35c59e3, + 0x8c49833d53bb8085, + 0x0216d0b17f4e44a5, + ]); + } + + type B = PortableBackend; + + #[test] + fn test_addition() { + let a = B::to_mont(&U256::new([1, 0, 0, 0])); + let b = B::to_mont(&U256::new([2, 0, 0, 0])); + let c = B::add(&a, &b); + assert_eq!(B::from_mont(&c), U256::new([3, 0, 0, 0])); + } + + #[test] + fn test_subtraction() { + let a = B::to_mont(&U256::new([5, 0, 0, 0])); + let b = B::to_mont(&U256::new([3, 0, 0, 0])); + let c = B::sub(&a, &b); + assert_eq!(B::from_mont(&c), U256::new([2, 0, 0, 0])); + + let underflow = B::sub(&b, &a); // 3 - 5 mod MODULUS + let mut expected = TestFr::MODULUS; + expected.0[0] -= 2; + assert_eq!(B::from_mont(&underflow), expected); + } + + #[test] + fn test_multiplication() { + let a = B::to_mont(&U256::new([100, 0, 0, 0])); + let b = B::to_mont(&U256::new([200, 0, 0, 0])); + let c = B::mul(&a, &b); + assert_eq!(B::from_mont(&c), U256::new([20000, 0, 0, 0])); + } + + #[test] + fn test_negation() { + let a = B::to_mont(&U256::new([12345, 0, 0, 0])); + let neg_a = B::neg(&a); + let sum = B::add(&a, &neg_a); + assert_eq!(B::from_mont(&sum), U256::zero()); + } +} diff --git a/syscall/solana-bn254/src/backend/traits.rs b/syscall/solana-bn254/src/backend/traits.rs new file mode 100644 index 0000000..bb16f18 --- /dev/null +++ b/syscall/solana-bn254/src/backend/traits.rs @@ -0,0 +1,58 @@ +//! Traits for field parameters and Montgomery arithmetic backends. + +use super::u256::U256; + +/// Compile-time configuration for a specific prime field (e.g., Fq or Fr). +pub trait Field: Send + Sync + 'static { + /// The prime modulus of the field. + const MODULUS: U256; + + /// The Montgomery inverse parameter: `-MODULUS^{-1} mod 2^64`. + /// Critical for fast, division-free Montgomery reduction. + const INV: u64; + + /// R^2 mod MODULUS. Used to move elements into Montgomery form. + const R2: U256; +} + +/// Generic trait defining a Montgomery arithmetic backend. +/// +/// # Safety and Preconditions +/// ALL functions in this trait strictly assume that inputs are fully reduced +/// field elements (i.e., `x < F::MODULUS`). Passing unreduced `U256` integers +/// will result in silently incorrect math, as the internal trial subtractions +/// and overflow logic are optimized for a maximum intermediate value of +/// `2 * MODULUS - 2`. +/// +/// Implementations process field elements in Montgomery form. Because +/// this crate operates on public data, constant-time execution is +/// explicitly NOT required. Implementations will rely on fast conditional +/// branches (e.g., simple subtraction for final reduction). +pub trait MontgomeryBackend { + /// Computes `(a + b) mod MODULUS`. + fn add(a: &U256, b: &U256) -> U256; + + /// Computes `(a - b) mod MODULUS`. + fn sub(a: &U256, b: &U256) -> U256; + + /// Computes `(a * b * R^-1) mod MODULUS` using Montgomery reduction. + fn mul(a: &U256, b: &U256) -> U256; + + /// Computes `(a * a * R^-1) mod MODULUS`. + fn sqr(a: &U256) -> U256; + + /// Computes `(-a) mod MODULUS`. + fn neg(a: &U256) -> U256; + + /// Converts a normal integer into Montgomery form: `mul(a, R^2)`. + #[inline(always)] + fn to_mont(a: &U256) -> U256 { + Self::mul(a, &F::R2) + } + + /// Converts a value out of Montgomery form: `mul(a, 1)`. + #[inline(always)] + fn from_mont(a: &U256) -> U256 { + Self::mul(a, &U256::one()) + } +} diff --git a/syscall/solana-bn254/src/backend/u256.rs b/syscall/solana-bn254/src/backend/u256.rs new file mode 100644 index 0000000..7aa3fa3 --- /dev/null +++ b/syscall/solana-bn254/src/backend/u256.rs @@ -0,0 +1,32 @@ +//! Core 256-bit integer representation. + +/// A 256-bit unsigned integer represented as four 64-bit limbs. +/// +/// Limbs are ordered little-endian: least significant at index 0, +/// most significant at index 3. +/// +/// `#[repr(C)]` guarantees a sequential memory layout, allowing the +/// x86_64 inline assembly to safely access limbs via pointer offsets. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] +#[repr(C)] +pub struct U256(pub [u64; 4]); + +impl U256 { + /// Creates a new `U256` from a fixed-size array of limbs. + #[inline(always)] + pub const fn new(limbs: [u64; 4]) -> Self { + Self(limbs) + } + + /// Creates a `U256` initialized to zero. + #[inline(always)] + pub const fn zero() -> Self { + Self([0, 0, 0, 0]) + } + + /// Creates a `U256` initialized to one. + #[inline(always)] + pub const fn one() -> Self { + Self([1, 0, 0, 0]) + } +} diff --git a/syscall/solana-bn254/src/lib.rs b/syscall/solana-bn254/src/lib.rs new file mode 100644 index 0000000..89dd127 --- /dev/null +++ b/syscall/solana-bn254/src/lib.rs @@ -0,0 +1,14 @@ +#![no_std] + +//! `solana-bn254` +//! +//! BN254 elliptic curve operations and Poseidon hashing, optimized +//! for the Solana Agave validator. +//! +//! # Security Warning +//! This crate is designed exclusively for PUBLIC DATA CONTEXTS. +//! It intentionally bypasses constant-time execution guarantees +//! to prioritize cycle efficiency and lowest possible Compute Units. + +pub mod backend; +pub mod poseidon; diff --git a/syscall/solana-bn254/src/poseidon/mod.rs b/syscall/solana-bn254/src/poseidon/mod.rs new file mode 100644 index 0000000..0ca7571 --- /dev/null +++ b/syscall/solana-bn254/src/poseidon/mod.rs @@ -0,0 +1,247 @@ +//! Poseidon hash implementation using the Fr scalar field. +//! +//! Includes optimized scalar execution routes using sparse matrices +//! for partial rounds, alongside an AVX-512 IFMA batched execution route +//! for maximum throughput on supporting hardware. +//! +//! # Hybrid Execution Architecture +//! - Partial Rounds: Uses strictly scalar logic. Because only the first element +//! `state[0]` receives an S-box, parallelization provides no benefit here. +//! - Full Rounds: When compiled for native targets supporting `AVX-512 IFMA`, +//! the full rounds dynamically route to an 8-way batched SIMD engine. + +use crate::backend::{Backend, Fr, MontgomeryBackend, U256}; + +/// Sparse matrix representation for Poseidon partial rounds. +/// +/// Precomputing the MDS matrix transitions into a sparse vector form reduces +/// the `O(T^2)` dense matrix multiplication down to an `O(T)` sparse matrix +/// computation, saving massive Compute Units on-chain. +#[derive(Clone, Debug)] +pub struct SparseMatrix { + pub row: [U256; T], + pub col: [U256; T], +} + +/// Constants required for Poseidon execution over a state of width `T`. +pub struct PoseidonConstants { + pub full_rounds: usize, + pub partial_rounds: usize, + pub round_constants: &'static [U256], + pub mds_matrix: &'static [[U256; T]; T], + pub pre_sparse_matrix: &'static [[U256; T]; T], + pub sparse_matrices: &'static [SparseMatrix], +} + +/// Computes the scalar Poseidon S-box (`x^5`) using exactly 3 inlined multiplications. +#[inline(always)] +pub fn sbox(x: &U256) -> U256 { + type B = Backend; + let x2 = B::sqr(x); + let x4 = B::sqr(&x2); + B::mul(&x4, x) +} + +/// Computes 8 Poseidon S-boxes simultaneously utilizing AVX-512 vectorization. +/// +/// The state is dynamically chunked in groups of 8. This guarantees full utilization +/// of the SIMD registers for large state widths. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +unsafe fn apply_sbox_simd(state: &mut [U256; T]) { + use crate::backend::avx512::{ + math::sbox_8x, + pack::{pack_8x, unpack_8x_into}, + }; + let mut i = 0; + + // Process state in chunks of 8 to saturate the SIMD lanes + while i < T { + let chunk_size = core::cmp::min(8, T - i); + let mut chunk = [U256::zero(); 8]; + chunk[..chunk_size].copy_from_slice(&state[i..i + chunk_size]); + + unsafe { + let packed = pack_8x(&chunk); + let sboxed = sbox_8x(&packed); + unpack_8x_into(&sboxed, &mut chunk); + } + + state[i..i + chunk_size].copy_from_slice(&chunk[..chunk_size]); + i += chunk_size; + } +} + +/// Computes a dense Matrix-Vector multiplication using AVX-512 Column-Accumulation. +/// +/// Instead of computing rows sequentially (`T^2` multiplications), we broadcast +/// the scalar `state[j]`, pack the Matrix Column `j` into a SIMD vector, and execute +/// parallel multiplications. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] +#[inline] +#[target_feature(enable = "avx512f,avx512ifma,avx512dq")] +unsafe fn apply_dense_matrix_simd(state: &mut [U256; T], mds: &[[U256; T]; T]) { + use crate::backend::avx512::{ + math::mul_8x, + pack::{broadcast, pack_8x, unpack_8x_into}, + }; + + let mut new_state = [U256::zero(); T]; + + for j in 0..T { + // Broadcast state[j] horizontally across all 8 SIMD lanes + let sj_broadcast = unsafe { broadcast(&state[j]) }; + + // Extract column j from the dense matrix + let mut col = [U256::zero(); T]; + for i in 0..T { + col[i] = mds[i][j]; + } + + let mut i = 0; + while i < T { + let chunk_size = core::cmp::min(8, T - i); + let mut col_chunk = [U256::zero(); 8]; + col_chunk[..chunk_size].copy_from_slice(&col[i..i + chunk_size]); + + let mut terms = [U256::zero(); 8]; + unsafe { + let col_packed = pack_8x(&col_chunk); + // Compute (Column * state[j]) simultaneously across up to 8 elements + let term_packed = mul_8x(&sj_broadcast, &col_packed); + unpack_8x_into(&term_packed, &mut terms); + } + + // Accumulate safely using the scalar backend to automatically handle Modulus bounds + for k in 0..chunk_size { + new_state[i + k] = Backend::::add(&new_state[i + k], &terms[k]); + } + i += chunk_size; + } + } + *state = new_state; +} + +/// Branchlessly executes a dense matrix multiplication on the scalar state. +#[inline(always)] +fn apply_dense_matrix(state: &mut [U256; T], m: &[[U256; T]; T]) { + type B = Backend; + let mut new_state = [U256::zero(); T]; + for i in 0..T { + let mut sum = U256::zero(); + for (j, state_val) in state.iter().enumerate() { + let term = B::mul(&m[i][j], state_val); + sum = B::add(&sum, &term); + } + new_state[i] = sum; + } + *state = new_state; +} + +/// Executes an `O(T)` sparse matrix multiplication on the scalar state. +#[inline(always)] +fn apply_sparse_matrix(state: &mut [U256; T], m: &SparseMatrix) { + type B = Backend; + let mut first_word = U256::zero(); + + // Row vector dot product for the new state[0] + for (j, state_val) in state.iter().enumerate() { + let term = B::mul(&m.row[j], state_val); + first_word = B::add(&first_word, &term); + } + + let prev_first = state[0]; + state[0] = first_word; + + // Identity operations scaled by the sparse column vector + for (i, state_val) in state.iter_mut().enumerate().skip(1) { + let term = B::mul(&m.col[i], &prev_first); + *state_val = B::add(state_val, &term); + } +} + +/// Executes the highly optimized Poseidon hash function. +pub fn poseidon( + mut state: [U256; T], + constants: &PoseidonConstants, +) -> [U256; T] { + type B = Backend; + let half_full = constants.full_rounds / 2; + let mut rc_idx = 0; + + // --- First Half: Full Rounds --- + for _ in 0..half_full { + for state_val in state.iter_mut() { + *state_val = B::add(state_val, &constants.round_constants[rc_idx]); + rc_idx += 1; + } + + // Execute S-boxes based on hardware capabilities + #[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] + unsafe { + apply_sbox_simd(&mut state); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512ifma")))] + for state_val in state.iter_mut() { + *state_val = sbox(state_val); + } + + // Execute MDS Matrix based on hardware capabilities + #[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] + unsafe { + apply_dense_matrix_simd(&mut state, constants.mds_matrix); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512ifma")))] + apply_dense_matrix(&mut state, constants.mds_matrix); + } + + // --- Transition to Partial Rounds --- + for state_val in state.iter_mut() { + *state_val = B::add(state_val, &constants.round_constants[rc_idx]); + rc_idx += 1; + } + + #[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] + unsafe { + apply_dense_matrix_simd(&mut state, constants.pre_sparse_matrix); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512ifma")))] + apply_dense_matrix(&mut state, constants.pre_sparse_matrix); + + // --- Middle: Partial Rounds --- + // Kept strictly scalar because only state[0] receives the S-box computation, + // rendering SIMD parallelization overhead highly inefficient here. + for sparse_idx in 0..constants.partial_rounds { + state[0] = sbox(&state[0]); + state[0] = B::add(&state[0], &constants.round_constants[rc_idx]); + rc_idx += 1; + apply_sparse_matrix(&mut state, &constants.sparse_matrices[sparse_idx]); + } + + // --- Second Half: Full Rounds --- + for _ in 0..half_full { + for state_val in state.iter_mut() { + *state_val = B::add(state_val, &constants.round_constants[rc_idx]); + rc_idx += 1; + } + + #[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] + unsafe { + apply_sbox_simd(&mut state); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512ifma")))] + for state_val in state.iter_mut() { + *state_val = sbox(state_val); + } + + #[cfg(all(target_arch = "x86_64", target_feature = "avx512ifma"))] + unsafe { + apply_dense_matrix_simd(&mut state, constants.mds_matrix); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512ifma")))] + apply_dense_matrix(&mut state, constants.mds_matrix); + } + + state +} From 5262b83199c89953af7d1141441d691ec3c3d547 Mon Sep 17 00:00:00 2001 From: samkim-crypto Date: Mon, 10 Aug 2026 17:51:06 +0900 Subject: [PATCH 2/2] add additional clarifying comments --- syscall/solana-bn254/src/backend/avx512/pack.rs | 2 +- syscall/solana-bn254/src/backend/fq.rs | 17 ++++++++++++++++- syscall/solana-bn254/src/backend/fr.rs | 8 +++++++- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/syscall/solana-bn254/src/backend/avx512/pack.rs b/syscall/solana-bn254/src/backend/avx512/pack.rs index 4d3a3af..18acd15 100644 --- a/syscall/solana-bn254/src/backend/avx512/pack.rs +++ b/syscall/solana-bn254/src/backend/avx512/pack.rs @@ -82,7 +82,7 @@ pub unsafe fn unpack_8x(packed: &FieldElement8x52) -> [U256; 8] { /// Broadcasts a single scalar 4-limb value identically to all 8 lanes of a SIMD register. /// -/// Crucial for Poseidon matrix multiplication, where a single state element must +/// Important for Poseidon matrix multiplication, where a single state element must /// be multiplied against an entire SIMD chunk of the matrix column vector. #[inline] #[target_feature(enable = "avx512f,avx512ifma,avx512dq")] diff --git a/syscall/solana-bn254/src/backend/fq.rs b/syscall/solana-bn254/src/backend/fq.rs index f78fda2..c8c868b 100644 --- a/syscall/solana-bn254/src/backend/fq.rs +++ b/syscall/solana-bn254/src/backend/fq.rs @@ -23,7 +23,22 @@ impl Field for Fq { } impl Fq { - /// Computes `a^(p-2) mod p` using a static 4-bit window addition chain. + /// Computes the modular inverse of `a` in the base field Fq. + /// + /// This relies on Fermat's Little Theorem, computing `a^(p-2) mod p`. + /// + /// Rather than using the Extended Euclidean Algorithm (which relies on + /// unpredictable loops and branching), this function uses a fully + /// unrolled, static 4-bit window addition chain. This minimizes + /// Solana Compute Units by leveraging our highly optimized, branchless + /// Montgomery multiplications and squarings. + /// + /// # Algorithm Breakdown: + /// 1. Precomputation: Calculates `a^2` through `a^15` (`a2`..`a15`). + /// 2. Chain Execution: Traverses the bits of the exponent `p-2`. + /// The accumulator `t` is squared 4 times (effectively shifting left + /// by 4 bits), followed by a multiplication with the precomputed + /// window value that matches the next 4 bits of the exponent. #[inline(always)] pub fn invert(a: &U256) -> U256 { type B = Backend; diff --git a/syscall/solana-bn254/src/backend/fr.rs b/syscall/solana-bn254/src/backend/fr.rs index f7303ef..afa077c 100644 --- a/syscall/solana-bn254/src/backend/fr.rs +++ b/syscall/solana-bn254/src/backend/fr.rs @@ -23,7 +23,13 @@ impl Field for Fr { } impl Fr { - /// Computes `a^(p-2) mod p` using a static 4-bit window addition chain. + /// Computes the modular inverse of `a` in the scalar field Fr. + /// + /// Computes `a^(r-2) mod r` via Fermat's Little Theorem, where `r` is + /// the Fr modulus. + /// + /// Like the Fq implementation, this strictly avoids dynamic branching or + /// loops in favor of a static 4-bit window addition chain. #[inline(always)] pub fn invert(a: &U256) -> U256 { type B = Backend;