Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Monty-related performance improvements #777

Merged
merged 12 commits into from
Mar 5, 2025
11 changes: 11 additions & 0 deletions benches/monty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,17 @@ fn bench_montgomery_ops<M: Measurement>(group: &mut BenchmarkGroup<'_, M>) {
)
});

group.bench_function("div_by_2, U256", |b| {
b.iter_batched(
|| {
let x = U256::random_mod(&mut rng, params.modulus().as_nz_ref());
MontyForm::new(&x, params)
},
|x| black_box(x.div_by_2()),
BatchSize::SmallInput,
)
});

#[cfg(feature = "alloc")]
for i in [1, 2, 3, 4, 10, 100] {
group.bench_function(
Expand Down
92 changes: 57 additions & 35 deletions src/modular/boxed_monty_form.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,9 @@ mod neg;
mod pow;
mod sub;

use super::{
ConstMontyParams, Retrieve, div_by_2,
reduction::{montgomery_reduction_boxed, montgomery_reduction_boxed_mut},
};
use super::{ConstMontyParams, Retrieve, div_by_2};
use mul::BoxedMontyMultiplier;

use crate::{BoxedUint, Limb, Monty, Odd, Word};
use alloc::sync::Arc;
use subtle::Choice;
Expand Down Expand Up @@ -59,7 +58,27 @@ impl BoxedMontyParams {
.rem(&modulus.as_nz_ref().widen(bits_precision * 2))
.shorten(bits_precision);

Self::new_inner(modulus, one, r2)
// The modular inverse should always exist, because it was ensured odd above, which also ensures it's non-zero
let (inv_mod_limb, inv_mod_limb_exists) = modulus.inv_mod2k_vartime(Word::BITS);
debug_assert!(bool::from(inv_mod_limb_exists));

let mod_neg_inv = Limb(Word::MIN.wrapping_sub(inv_mod_limb.limbs[0].0));

let mod_leading_zeros = modulus.as_ref().leading_zeros().min(Word::BITS - 1);

let r3 = {
let mut mm = BoxedMontyMultiplier::new(&modulus, mod_neg_inv);
mm.square(&r2)
};

Self {
modulus,
one,
r2,
r3,
mod_neg_inv,
mod_leading_zeros,
}
}

/// Instantiates a new set of [`BoxedMontyParams`] representing the given `modulus`, which
Expand All @@ -82,23 +101,18 @@ impl BoxedMontyParams {
.rem_vartime(&modulus.as_nz_ref().widen(bits_precision * 2))
.shorten(bits_precision);

Self::new_inner(modulus, one, r2)
}

/// Common functionality of `new` and `new_vartime`.
fn new_inner(modulus: Odd<BoxedUint>, one: BoxedUint, r2: BoxedUint) -> Self {
debug_assert_eq!(one.bits_precision(), modulus.bits_precision());
debug_assert_eq!(r2.bits_precision(), modulus.bits_precision());

// If the inverse exists, it means the modulus is odd.
let (inv_mod_limb, modulus_is_odd) = modulus.inv_mod2k(Word::BITS);
debug_assert!(bool::from(modulus_is_odd));
// The modular inverse should always exist, because it was ensured odd above, which also ensures it's non-zero
let (inv_mod_limb, inv_mod_limb_exists) = modulus.inv_mod2k_full_vartime(Word::BITS);
debug_assert!(bool::from(inv_mod_limb_exists));

let mod_neg_inv = Limb(Word::MIN.wrapping_sub(inv_mod_limb.limbs[0].0));

let mod_leading_zeros = modulus.as_ref().leading_zeros().min(Word::BITS - 1);

let r3 = montgomery_reduction_boxed(&mut r2.square(), &modulus, mod_neg_inv);
let r3 = {
let mut mm = BoxedMontyMultiplier::new(&modulus, mod_neg_inv);
mm.square(&r2)
};

Self {
modulus,
Expand Down Expand Up @@ -173,19 +187,8 @@ impl BoxedMontyForm {

/// Retrieves the integer currently encoded in this [`BoxedMontyForm`], guaranteed to be reduced.
pub fn retrieve(&self) -> BoxedUint {
let mut montgomery_form = self.montgomery_form.widen(self.bits_precision() * 2);

let ret = montgomery_reduction_boxed(
&mut montgomery_form,
&self.params.modulus,
self.params.mod_neg_inv,
);

#[cfg(feature = "zeroize")]
montgomery_form.zeroize();

debug_assert!(ret < self.params.modulus);
ret
let mut mm = BoxedMontyMultiplier::from(self.params.as_ref());
mm.mul_by_one(&self.montgomery_form)
}

/// Instantiates a new `ConstMontyForm` that represents zero.
Expand Down Expand Up @@ -256,6 +259,12 @@ impl BoxedMontyForm {
params: self.params.clone(),
}
}

/// Performs division by 2 inplace, that is finds `x` such that `x + x = self`
/// and writes it into `self`.
pub fn div_by_2_assign(&mut self) {
div_by_2::div_by_2_boxed_assign(&mut self.montgomery_form, &self.params.modulus)
}
}

impl Retrieve for BoxedMontyForm {
Expand All @@ -268,6 +277,7 @@ impl Retrieve for BoxedMontyForm {
impl Monty for BoxedMontyForm {
type Integer = BoxedUint;
type Params = BoxedMontyParams;
type Multiplier<'a> = BoxedMontyMultiplier<'a>;

fn new_params_vartime(modulus: Odd<Self::Integer>) -> Self::Params {
BoxedMontyParams::new_vartime(modulus)
Expand All @@ -293,6 +303,17 @@ impl Monty for BoxedMontyForm {
&self.montgomery_form
}

fn copy_montgomery_from(&mut self, other: &Self) {
debug_assert_eq!(
self.montgomery_form.bits_precision(),
other.montgomery_form.bits_precision()
);
debug_assert_eq!(self.params, other.params);
self.montgomery_form
.limbs
.copy_from_slice(&other.montgomery_form.limbs);
}

fn double(&self) -> Self {
BoxedMontyForm::double(self)
}
Expand All @@ -301,6 +322,10 @@ impl Monty for BoxedMontyForm {
BoxedMontyForm::div_by_2(self)
}

fn div_by_2_assign(&mut self) {
BoxedMontyForm::div_by_2_assign(self)
}

fn lincomb_vartime(products: &[(&Self, &Self)]) -> Self {
BoxedMontyForm::lincomb_vartime(products)
}
Expand All @@ -317,11 +342,8 @@ impl Zeroize for BoxedMontyForm {
/// Convert the given integer into the Montgomery domain.
#[inline]
fn convert_to_montgomery(integer: &mut BoxedUint, params: &BoxedMontyParams) {
let mut product = integer.mul(&params.r2);
montgomery_reduction_boxed_mut(&mut product, &params.modulus, params.mod_neg_inv, integer);

#[cfg(feature = "zeroize")]
product.zeroize();
let mut mm = BoxedMontyMultiplier::from(params);
mm.mul_assign(integer, &params.r2);
}

#[cfg(test)]
Expand Down
5 changes: 2 additions & 3 deletions src/modular/boxed_monty_form/add.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,8 @@ impl Add<BoxedMontyForm> for BoxedMontyForm {
impl AddAssign<&BoxedMontyForm> for BoxedMontyForm {
fn add_assign(&mut self, rhs: &BoxedMontyForm) {
debug_assert_eq!(self.params, rhs.params);
self.montgomery_form = self
.montgomery_form
.add_mod(&rhs.montgomery_form, &self.params.modulus)
self.montgomery_form
.add_mod_assign(&rhs.montgomery_form, &self.params.modulus);
}
}

Expand Down
Loading