From c9fa679afdc103f0d633462a452c80d3d0f8fc2e Mon Sep 17 00:00:00 2001 From: Michael Zhu Date: Fri, 23 Feb 2024 18:17:33 -0500 Subject: [PATCH 1/4] Introduce optimized from_u64 function for conversion to Montgomery form --- bench-templates/src/macros/field.rs | 11 ++++++ ff/src/fields/models/fp/mod.rs | 14 ++++++++ ff/src/fields/models/fp/montgomery_backend.rs | 34 +++++++++++++++++++ ff/src/fields/prime.rs | 3 ++ 4 files changed, 62 insertions(+) diff --git a/bench-templates/src/macros/field.rs b/bench-templates/src/macros/field.rs index 585b5000a..81874b40c 100644 --- a/bench-templates/src/macros/field.rs +++ b/bench-templates/src/macros/field.rs @@ -6,6 +6,7 @@ macro_rules! f_bench { mod [<$F:lower>] { use super::*; use ark_ff::{Field, PrimeField, UniformRand}; + use ark_std::rand::RngCore; field_common!($bench_group_name, $F); sqrt!($bench_group_name, $F); prime_field!($bench_group_name, $F); @@ -402,6 +403,16 @@ macro_rules! prime_field { f[i].into_bigint() }) }); + let u64s = (0..SAMPLES) + .map(|_| rng.next_u64() % 65536) + .collect::>(); + conversions.bench_function("From u64", |b| { + let mut i = 0; + b.iter(|| { + i = (i + 1) % SAMPLES; + <$F>::from_u64(u64s[i]) + }) + }); conversions.finish() } }; diff --git a/ff/src/fields/models/fp/mod.rs b/ff/src/fields/models/fp/mod.rs index 7f417ed9a..337a21b0f 100644 --- a/ff/src/fields/models/fp/mod.rs +++ b/ff/src/fields/models/fp/mod.rs @@ -64,6 +64,10 @@ pub trait FpConfig: Send + Sync + 'static + Sized { /// which works for every modulus. const SQRT_PRECOMP: Option>>; + /// Precomputed lookup table for values 0..2^16 in Montgomery form. + /// Otherwise, conversion to Montgomery form requires a multiplication by R^2. + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; 65536]; + /// Set a += b. fn add_assign(a: &mut Fp, b: &Fp); @@ -96,6 +100,11 @@ pub trait FpConfig: Send + Sync + 'static + Sized { /// Convert a field element to an integer in the range `0..(Self::MODULUS - /// 1)`. fn into_bigint(other: Fp) -> BigInt; + + /// Construct a field element from a u64 in the range + /// `0..(Self::MODULUS - 1)`. Returns `None` if the integer is outside + /// this range. + fn from_u64(other: u64) -> Option>; } /// Represents an element of the prime field F_p, where `p == P::MODULUS`. @@ -350,6 +359,11 @@ impl, const N: usize> PrimeField for Fp { fn into_bigint(self) -> BigInt { P::into_bigint(self) } + + #[inline] + fn from_u64(r: u64) -> Option { + P::from_u64(r) + } } impl, const N: usize> FftField for Fp { diff --git a/ff/src/fields/models/fp/montgomery_backend.rs b/ff/src/fields/models/fp/montgomery_backend.rs index 76344227c..78fc9493e 100644 --- a/ff/src/fields/models/fp/montgomery_backend.rs +++ b/ff/src/fields/models/fp/montgomery_backend.rs @@ -2,6 +2,7 @@ use ark_std::{marker::PhantomData, Zero}; use super::{Fp, FpConfig}; use crate::{biginteger::arithmetic as fa, BigInt, BigInteger, PrimeField, SqrtPrecomputation}; +use crate::fields::Field; use ark_ff_macros::unroll_for_loops; /// A trait that specifies the constants and arithmetic procedures @@ -76,6 +77,10 @@ pub trait MontConfig: 'static + Sync + Send + Sized { const SQRT_PRECOMP: Option, N>>> = sqrt_precomputation::(); + #[allow(long_running_const_eval)] + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp, N>; 65536] = + small_element_montgomery_precomputation::(); + /// (MODULUS + 1) / 4 when MODULUS % 4 == 3. Used for square root precomputations. #[doc(hidden)] const MODULUS_PLUS_ONE_DIV_FOUR: Option> = { @@ -354,6 +359,14 @@ pub trait MontConfig: 'static + Sync + Send + Sized { } } + fn from_u64(r: u64) -> Option, N>> { + if r < 65536 { + Some(Self::SMALL_ELEMENT_MONTGOMERY_PRECOMP[r as usize]) + } else { + Self::from_bigint(r.into()) + } + } + fn from_bigint(r: BigInt) -> Option, N>> { let mut r = Fp::new_unchecked(r); if r.is_zero() { @@ -559,6 +572,21 @@ pub const fn sqrt_precomputation>( } } +pub const fn small_element_montgomery_precomputation>( +) -> [Fp, N>; 65536] { + let mut lookup_table: [Fp, N>; 65536] = + [, N>>::ZERO; 65536]; + + let mut i: usize = 1; + while i < 65536 { + let mut limbs = [0u64; N]; + limbs[0] = i as u64; + lookup_table[i] = , N>>::new(BigInt::new(limbs)); + i += 1; + } + lookup_table +} + /// Construct a [`Fp, N>`] element from a literal string. This /// should be used primarily for constructing constant field elements; in a /// non-const context, [`Fp::from_str`](`ark_std::str::FromStr::from_str`) is @@ -624,6 +652,8 @@ impl, const N: usize> FpConfig for MontBackend { const SMALL_SUBGROUP_BASE_ADICITY: Option = T::SMALL_SUBGROUP_BASE_ADICITY; const LARGE_SUBGROUP_ROOT_OF_UNITY: Option> = T::LARGE_SUBGROUP_ROOT_OF_UNITY; const SQRT_PRECOMP: Option>> = T::SQRT_PRECOMP; + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; 65536] = + T::SMALL_ELEMENT_MONTGOMERY_PRECOMP; fn add_assign(a: &mut Fp, b: &Fp) { T::add_assign(a, b) @@ -675,6 +705,10 @@ impl, const N: usize> FpConfig for MontBackend { fn into_bigint(a: Fp) -> BigInt { T::into_bigint(a) } + + fn from_u64(r: u64) -> Option> { + T::from_u64(r) + } } impl, const N: usize> Fp, N> { diff --git a/ff/src/fields/prime.rs b/ff/src/fields/prime.rs index 8a8fe89b4..7c2e50e0d 100644 --- a/ff/src/fields/prime.rs +++ b/ff/src/fields/prime.rs @@ -57,6 +57,9 @@ pub trait PrimeField: /// Converts an element of the prime field into an integer in the range 0..(p - 1). fn into_bigint(self) -> Self::BigInt; + /// Construct a prime field element from a u64 in the range 0..(p - 1). + fn from_u64(repr: u64) -> Option; + /// Reads bytes in big-endian, and converts them to a field element. /// If the integer represented by `bytes` is larger than the modulus `p`, this method /// performs the appropriate reduction. From 72fc43e860f09c90c94b14ffe0c2bc297f10435e Mon Sep 17 00:00:00 2001 From: Michael Zhu Date: Mon, 26 Feb 2024 12:35:41 -0500 Subject: [PATCH 2/4] Use optimized mul_u64 for conversion to Montgomery form --- ff/src/fields/models/fp/montgomery_backend.rs | 57 +++++++++++++++++-- 1 file changed, 52 insertions(+), 5 deletions(-) diff --git a/ff/src/fields/models/fp/montgomery_backend.rs b/ff/src/fields/models/fp/montgomery_backend.rs index 78fc9493e..7cd1db03f 100644 --- a/ff/src/fields/models/fp/montgomery_backend.rs +++ b/ff/src/fields/models/fp/montgomery_backend.rs @@ -1,8 +1,11 @@ use ark_std::{marker::PhantomData, Zero}; use super::{Fp, FpConfig}; -use crate::{biginteger::arithmetic as fa, BigInt, BigInteger, PrimeField, SqrtPrecomputation}; -use crate::fields::Field; +use crate::{ + biginteger::arithmetic::{self as fa}, + fields::Field, + BigInt, BigInteger, PrimeField, SqrtPrecomputation, +}; use ark_ff_macros::unroll_for_loops; /// A trait that specifies the constants and arithmetic procedures @@ -362,8 +365,10 @@ pub trait MontConfig: 'static + Sync + Send + Sized { fn from_u64(r: u64) -> Option, N>> { if r < 65536 { Some(Self::SMALL_ELEMENT_MONTGOMERY_PRECOMP[r as usize]) + } else if BigInt::from(r) >= >::MODULUS { + None } else { - Self::from_bigint(r.into()) + Some(Fp::new_unchecked(Self::R2).mul_u64(r)) } } @@ -708,7 +713,7 @@ impl, const N: usize> FpConfig for MontBackend { fn from_u64(r: u64) -> Option> { T::from_u64(r) - } + } } impl, const N: usize> Fp, N> { @@ -821,6 +826,37 @@ impl, const N: usize> Fp, N> { } } + #[unroll_for_loops(12)] + #[inline(always)] + pub fn mul_u64(mut self, other: u64) -> Self { + let (mut lo, mut hi) = ([0u64; N], [0u64; N]); + + for i in 0..N - 1 { + lo[i] = mac_with_carry!(lo[i], (self.0).0[i], other, &mut lo[i + 1]); + } + lo[N - 1] = mac_with_carry!(lo[N - 1], (self.0).0[N - 1], other, &mut hi[0]); + + // Montgomery reduction + let mut carry2 = 0; + for i in 0..N { + let tmp = lo[i].wrapping_mul(T::INV); + let mut carry = 0u64; + mac!(lo[i], tmp, T::MODULUS.0[0], &mut carry); + for j in 1..N { + let k: usize = i + j; + if k >= N { + hi[k - N] = mac_with_carry!(hi[k - N], tmp, T::MODULUS.0[j], &mut carry); + } else { + lo[k] = mac_with_carry!(lo[k], tmp, T::MODULUS.0[j], &mut carry); + } + } + hi[i] = adc!(hi[i], carry, &mut carry2); + } + +(self.0).0 = hi; + self.const_subtract_modulus_with_carry(carry2 != 0) + } + const fn const_is_valid(&self) -> bool { crate::const_for!((i in 0..N) { if (self.0).0[N - i - 1] < T::MODULUS.0[N - i - 1] { @@ -855,10 +891,21 @@ impl, const N: usize> Fp, N> { #[cfg(test)] mod test { - use ark_std::{str::FromStr, vec::Vec}; + use ark_std::{rand::RngCore, str::FromStr}; use ark_test_curves::secp256k1::Fr; use num_bigint::{BigInt, BigUint, Sign}; + #[test] + fn test_mul_u64() { + let r2 = Fr::new_unchecked(Fr::R2); + let mut rng = ark_std::test_rng(); + let value = rng.next_u64(); + assert_eq!( + r2.mul_u64(value), + r2 * Fr::new_unchecked(value.into()) + ); + } + #[test] fn test_mont_macro_correctness() { let (is_positive, limbs) = str_to_limbs_u64( From 5339f7c5a9b4603bb8fd0b6a3a8bef4c21546248 Mon Sep 17 00:00:00 2001 From: Michael Zhu Date: Mon, 26 Feb 2024 14:22:47 -0500 Subject: [PATCH 3/4] Add attribution comment + put 2^16 in a constant --- bench-templates/src/macros/field.rs | 2 +- ff/src/fields/models/fp/mod.rs | 2 +- ff/src/fields/models/fp/montgomery_backend.rs | 18 +++++++++++------- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/bench-templates/src/macros/field.rs b/bench-templates/src/macros/field.rs index 81874b40c..62063205b 100644 --- a/bench-templates/src/macros/field.rs +++ b/bench-templates/src/macros/field.rs @@ -404,7 +404,7 @@ macro_rules! prime_field { }) }); let u64s = (0..SAMPLES) - .map(|_| rng.next_u64() % 65536) + .map(|_| rng.next_u64()) .collect::>(); conversions.bench_function("From u64", |b| { let mut i = 0; diff --git a/ff/src/fields/models/fp/mod.rs b/ff/src/fields/models/fp/mod.rs index 337a21b0f..be8b3c923 100644 --- a/ff/src/fields/models/fp/mod.rs +++ b/ff/src/fields/models/fp/mod.rs @@ -66,7 +66,7 @@ pub trait FpConfig: Send + Sync + 'static + Sized { /// Precomputed lookup table for values 0..2^16 in Montgomery form. /// Otherwise, conversion to Montgomery form requires a multiplication by R^2. - const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; 65536]; + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; PRECOMP_TABLE_SIZE]; /// Set a += b. fn add_assign(a: &mut Fp, b: &Fp); diff --git a/ff/src/fields/models/fp/montgomery_backend.rs b/ff/src/fields/models/fp/montgomery_backend.rs index 7cd1db03f..f6b2d8985 100644 --- a/ff/src/fields/models/fp/montgomery_backend.rs +++ b/ff/src/fields/models/fp/montgomery_backend.rs @@ -8,6 +8,8 @@ use crate::{ }; use ark_ff_macros::unroll_for_loops; +pub const PRECOMP_TABLE_SIZE: usize = 65536; + /// A trait that specifies the constants and arithmetic procedures /// for Montgomery arithmetic over the prime field defined by `MODULUS`. /// @@ -81,7 +83,7 @@ pub trait MontConfig: 'static + Sync + Send + Sized { sqrt_precomputation::(); #[allow(long_running_const_eval)] - const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp, N>; 65536] = + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp, N>; PRECOMP_TABLE_SIZE] = small_element_montgomery_precomputation::(); /// (MODULUS + 1) / 4 when MODULUS % 4 == 3. Used for square root precomputations. @@ -363,7 +365,7 @@ pub trait MontConfig: 'static + Sync + Send + Sized { } fn from_u64(r: u64) -> Option, N>> { - if r < 65536 { + if r < PRECOMP_TABLE_SIZE as u64 { Some(Self::SMALL_ELEMENT_MONTGOMERY_PRECOMP[r as usize]) } else if BigInt::from(r) >= >::MODULUS { None @@ -577,13 +579,15 @@ pub const fn sqrt_precomputation>( } } +/// Adapted the `bn256-table` feature from `halo2curves`: +/// https://github.com/privacy-scaling-explorations/halo2curves/blob/main/script/bn256.py pub const fn small_element_montgomery_precomputation>( -) -> [Fp, N>; 65536] { - let mut lookup_table: [Fp, N>; 65536] = - [, N>>::ZERO; 65536]; +) -> [Fp, N>; PRECOMP_TABLE_SIZE] { + let mut lookup_table: [Fp, N>; PRECOMP_TABLE_SIZE] = + [, N>>::ZERO; PRECOMP_TABLE_SIZE]; let mut i: usize = 1; - while i < 65536 { + while i < PRECOMP_TABLE_SIZE { let mut limbs = [0u64; N]; limbs[0] = i as u64; lookup_table[i] = , N>>::new(BigInt::new(limbs)); @@ -657,7 +661,7 @@ impl, const N: usize> FpConfig for MontBackend { const SMALL_SUBGROUP_BASE_ADICITY: Option = T::SMALL_SUBGROUP_BASE_ADICITY; const LARGE_SUBGROUP_ROOT_OF_UNITY: Option> = T::LARGE_SUBGROUP_ROOT_OF_UNITY; const SQRT_PRECOMP: Option>> = T::SQRT_PRECOMP; - const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; 65536] = + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; PRECOMP_TABLE_SIZE] = T::SMALL_ELEMENT_MONTGOMERY_PRECOMP; fn add_assign(a: &mut Fp, b: &Fp) { From ccd64e513d16e627c782a4f4519475d691977fc5 Mon Sep 17 00:00:00 2001 From: Michael Zhu Date: Wed, 28 Feb 2024 11:32:06 -0500 Subject: [PATCH 4/4] Reduce PRECOMP_TABLE_SIZE to avoid stack overflow --- ff/src/fields/models/fp/montgomery_backend.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ff/src/fields/models/fp/montgomery_backend.rs b/ff/src/fields/models/fp/montgomery_backend.rs index f6b2d8985..ac0b52e90 100644 --- a/ff/src/fields/models/fp/montgomery_backend.rs +++ b/ff/src/fields/models/fp/montgomery_backend.rs @@ -8,7 +8,7 @@ use crate::{ }; use ark_ff_macros::unroll_for_loops; -pub const PRECOMP_TABLE_SIZE: usize = 65536; +pub const PRECOMP_TABLE_SIZE: usize = 1 << 14; /// A trait that specifies the constants and arithmetic procedures /// for Montgomery arithmetic over the prime field defined by `MODULUS`.