From c9fa679afdc103f0d633462a452c80d3d0f8fc2e Mon Sep 17 00:00:00 2001 From: Michael Zhu Date: Fri, 23 Feb 2024 18:17:33 -0500 Subject: [PATCH] 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.