fixed rescaling & added all tests

This commit is contained in:
Jean-Philippe Bossuat
2025-01-08 11:45:48 +01:00
parent be87ac6ae7
commit e4c19a163e
3 changed files with 134 additions and 87 deletions

View File

@@ -33,6 +33,10 @@ impl<O: Unsigned> RingRNS<O> {
PolyRNS::<u64>::new(self.n(), self.level()) PolyRNS::<u64>::new(self.n(), self.level())
} }
pub fn new_poly(&self) -> Poly<u64> {
Poly::<u64>::new(self.n())
}
pub fn max_level(&self) -> usize { pub fn max_level(&self) -> usize {
self.0.len() - 1 self.0.len() - 1
} }

View File

@@ -1,6 +1,6 @@
use crate::modulus::barrett::Barrett; use crate::modulus::barrett::Barrett;
use crate::modulus::{BARRETT, NONE, ONCE}; use crate::modulus::{BARRETT, NONE, ONCE};
use crate::poly::PolyRNS; use crate::poly::{Poly, PolyRNS};
use crate::ring::Ring; use crate::ring::Ring;
use crate::ring::RingRNS; use crate::ring::RingRNS;
use crate::scalar::ScalarRNS; use crate::scalar::ScalarRNS;
@@ -8,18 +8,19 @@ extern crate test;
impl RingRNS<u64> { impl RingRNS<u64> {
/// Updates b to floor(a / q[b.level()]). /// Updates b to floor(a / q[b.level()]).
/// buf is unused if <ROUND=false,NTT=false>
pub fn div_by_last_modulus<const ROUND: bool, const NTT: bool>( pub fn div_by_last_modulus<const ROUND: bool, const NTT: bool>(
&self, &self,
a: &PolyRNS<u64>, a: &PolyRNS<u64>,
buf: &mut PolyRNS<u64>, buf: &mut [Poly<u64>; 2],
b: &mut PolyRNS<u64>, b: &mut PolyRNS<u64>,
) { ) {
debug_assert!(self.level() != 0, "invalid call: self.level()=0"); debug_assert!(self.level() != 0, "invalid call: self.level()=0");
debug_assert!( debug_assert!(
self.level() <= a.level(), a.level() >= self.level(),
"invalid input a: self.level()={} > a.level()={}", "invalid input a: a.level()={} < self.level()={}",
self.level(), a.level(),
a.level() self.level()
); );
debug_assert!( debug_assert!(
b.level() >= self.level() - 1, b.level() >= self.level() - 1,
@@ -35,7 +36,7 @@ impl RingRNS<u64> {
if ROUND { if ROUND {
let q_level_half: u64 = r_last.modulus.q >> 1; let q_level_half: u64 = r_last.modulus.q >> 1;
let (buf_q_scaling, buf_qi_scaling) = buf.0.split_at_mut(1); let (buf_q_scaling, buf_qi_scaling) = buf.split_at_mut(1);
if NTT { if NTT {
r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]); r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
@@ -76,7 +77,7 @@ impl RingRNS<u64> {
} }
} else { } else {
if NTT { if NTT {
let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1); let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.split_at_mut(1);
self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]); self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
for (i, r) in self.0[0..level].iter().enumerate() { for (i, r) in self.0[0..level].iter().enumerate() {
r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]); r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
@@ -104,12 +105,12 @@ impl RingRNS<u64> {
/// Expects a to be in the NTT domain. /// Expects a to be in the NTT domain.
pub fn div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>( pub fn div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(
&self, &self,
buf: &mut PolyRNS<u64>, buf: &mut [Poly<u64>; 2],
a: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>,
) { ) {
debug_assert!( debug_assert!(
self.level() <= a.level(), self.level() <= a.level(),
"invalid input a: self.level()={} > a.level()={}", "invalid input a: a.level()={} < self.level()={}",
self.level(), self.level(),
a.level() a.level()
); );
@@ -120,7 +121,7 @@ impl RingRNS<u64> {
if ROUND { if ROUND {
let q_level_half: u64 = r_last.modulus.q >> 1; let q_level_half: u64 = r_last.modulus.q >> 1;
let (buf_q_scaling, buf_qi_scaling) = buf.0.split_at_mut(1); let (buf_q_scaling, buf_qi_scaling) = buf.split_at_mut(1);
if NTT { if NTT {
r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]); r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
@@ -152,7 +153,7 @@ impl RingRNS<u64> {
} }
} else { } else {
if NTT { if NTT {
let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1); let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.split_at_mut(1);
r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]); r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
for (i, r) in self.0[0..level].iter().enumerate() { for (i, r) in self.0[0..level].iter().enumerate() {
r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]); r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
@@ -178,104 +179,106 @@ impl RingRNS<u64> {
/// Updates b to floor(a / prod_{level - nb_moduli}^{level} q[i]) /// Updates b to floor(a / prod_{level - nb_moduli}^{level} q[i])
pub fn div_by_last_moduli<const ROUND: bool, const NTT: bool>( pub fn div_by_last_moduli<const ROUND: bool, const NTT: bool>(
&self, &self,
nb_moduli: usize, nb_moduli_dropped: usize,
a: &PolyRNS<u64>, a: &PolyRNS<u64>,
buf: &mut PolyRNS<u64>, buf0: &mut [Poly<u64>; 2],
buf1: &mut PolyRNS<u64>,
c: &mut PolyRNS<u64>, c: &mut PolyRNS<u64>,
) { ) {
debug_assert!( debug_assert!(
nb_moduli <= self.level(), nb_moduli_dropped <= self.level(),
"invalid input nb_moduli: nb_moduli={} > a.level()={}", "invalid input nb_moduli_dropped: nb_moduli_dropped={} > self.level()={}",
nb_moduli, nb_moduli_dropped,
a.level() self.level()
); );
debug_assert!( debug_assert!(
a.level() <= self.level(), a.level() >= self.level(),
"invalid input a: self.level()={} > a.level()={}", "invalid input a: a.level()={} < self.level()={}",
self.level(), a.level(),
a.level() self.level()
); );
debug_assert!( debug_assert!(
buf.level() >= self.level() - 1, buf1.level() >= self.level(),
"invalid input buf: buf.level()={} < a.level()-1={}", "invalid input buf: buf.level()={} < self.level()={}",
buf.level(), buf1.level(),
a.level() - 1 self.level()
); );
debug_assert!( debug_assert!(
c.level() >= self.level() - nb_moduli, c.level() >= self.level() - nb_moduli_dropped,
"invalid input c: c.level()={} < c.level()-nb_moduli={}", "invalid input c: c.level()={} < self.level()-nb_moduli_dropped={}",
c.level(), c.level(),
a.level() - nb_moduli self.level() - nb_moduli_dropped
); );
if nb_moduli == 0 { if nb_moduli_dropped == 0 {
if a != c { if a != c {
c.copy(a); c.copy(a);
} }
} else { } else {
if NTT { if NTT {
self.intt::<false>(a, buf); self.intt::<false>(a, buf1);
(0..nb_moduli).for_each(|i| { (0..nb_moduli_dropped).for_each(|i| {
self.at_level(self.level() - i) self.at_level(self.level() - i)
.div_by_last_modulus_inplace::<ROUND, false>( .div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1)
&mut PolyRNS::<u64>::default(),
buf,
)
}); });
self.at_level(self.level() - nb_moduli).ntt::<false>(buf, c); self.at_level(self.level() - nb_moduli_dropped)
.ntt::<false>(buf1, c);
} else { } else {
println!("{} {:?}", self.level(), buf.level()); self.div_by_last_modulus::<ROUND, false>(a, buf0, buf1);
self.div_by_last_modulus::<ROUND, false>(a, buf, c);
(1..nb_moduli - 1).for_each(|i| { (1..nb_moduli_dropped - 1).for_each(|i| {
println!("{} {:?}", self.level() - i, buf.level());
self.at_level(self.level() - i) self.at_level(self.level() - i)
.div_by_last_modulus_inplace::<ROUND, false>(buf, c); .div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1);
}); });
self.at_level(self.level() - nb_moduli + 1) self.at_level(self.level() - nb_moduli_dropped + 1)
.div_by_last_modulus_inplace::<ROUND, false>(buf, c); .div_by_last_modulus::<ROUND, false>(buf1, buf0, c);
} }
} }
} }
/// Updates a to floor(a / prod_{level - nb_moduli}^{level} q[i]) /// Updates a to floor(a / prod_{level - nb_moduli_dropped}^{level} q[i])
pub fn div_by_last_moduli_inplace<const ROUND: bool, const NTT: bool>( pub fn div_by_last_moduli_inplace<const ROUND: bool, const NTT: bool>(
&self, &self,
nb_moduli: usize, nb_moduli_dropped: usize,
buf: &mut PolyRNS<u64>, buf0: &mut [Poly<u64>; 2],
buf1: &mut PolyRNS<u64>,
a: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>,
) { ) {
debug_assert!( debug_assert!(
self.level() <= a.level(), nb_moduli_dropped <= self.level(),
"invalid input a: self.level()={} > a.level()={}", "invalid input nb_moduli_dropped: nb_moduli_dropped={} > self.level()={}",
self.level(), nb_moduli_dropped,
a.level() self.level()
); );
debug_assert!( debug_assert!(
nb_moduli <= a.level(), a.level() >= self.level(),
"invalid input nb_moduli: nb_moduli={} > a.level()={}", "invalid input a: a.level()={} < self.level()={}",
nb_moduli, a.level(),
a.level() self.level()
); );
if nb_moduli == 0 { debug_assert!(
buf1.level() >= self.level(),
"invalid input buf: buf.level()={} < self.level()={}",
buf1.level(),
self.level()
);
if nb_moduli_dropped == 0 {
return; return;
} }
if NTT { if NTT {
self.intt::<false>(a, buf); self.intt::<false>(a, buf1);
(0..nb_moduli).for_each(|i| { (0..nb_moduli_dropped).for_each(|i| {
self.at_level(self.level() - i) self.at_level(self.level() - i)
.div_by_last_modulus_inplace::<ROUND, false>( .div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1)
&mut PolyRNS::<u64>::default(),
buf,
)
}); });
self.at_level(self.level() - nb_moduli).ntt::<false>(buf, a); self.at_level(self.level() - nb_moduli_dropped)
.ntt::<false>(buf1, a);
} else { } else {
(0..nb_moduli).for_each(|i| { (0..nb_moduli_dropped).for_each(|i| {
self.at_level(self.level() - i) self.at_level(self.level() - i)
.div_by_last_modulus_inplace::<ROUND, false>(buf, a) .div_by_last_modulus_inplace::<ROUND, false>(buf0, a)
}); });
} }
} }

View File

@@ -1,6 +1,6 @@
use itertools::izip; use itertools::izip;
use math::num_bigint::Div; use math::num_bigint::Div;
use math::poly::PolyRNS; use math::poly::{Poly, PolyRNS};
use math::ring::RingRNS; use math::ring::RingRNS;
use num_bigint::BigInt; use num_bigint::BigInt;
use sampling::source::Source; use sampling::source::Source;
@@ -44,8 +44,34 @@ fn rescaling_rns_u64() {
"test_div_by_last_modulus_inplace::<ROUND:true, NTT:false>", "test_div_by_last_modulus_inplace::<ROUND:true, NTT:false>",
|| test_div_by_last_modulus_inplace::<true, false>(&ring_rns), || test_div_by_last_modulus_inplace::<true, false>(&ring_rns),
); );
sub_test("test_div_by_last_moduli::<ROUND:false, NTT:false>", || {
//sub_test("test_div_by_last_moduli::<ROUND:false, NTT:false>", ||{test_div_by_last_moduli::<false, false>(&ring_rns)}); test_div_by_last_moduli::<false, false>(&ring_rns)
});
sub_test("test_div_by_last_moduli::<ROUND:false, NTT:true>", || {
test_div_by_last_moduli::<false, true>(&ring_rns)
});
sub_test("test_div_by_last_moduli::<ROUND:true, NTT:false>", || {
test_div_by_last_moduli::<true, false>(&ring_rns)
});
sub_test("test_div_by_last_moduli::<ROUND:true, NTT:true>", || {
test_div_by_last_moduli::<true, true>(&ring_rns)
});
sub_test(
"test_div_by_last_moduli_inplace::<ROUND:false, NTT:false>",
|| test_div_by_last_moduli_inplace::<false, false>(&ring_rns),
);
sub_test(
"test_div_by_last_moduli_inplace::<ROUND:false, NTT:true>",
|| test_div_by_last_moduli_inplace::<false, true>(&ring_rns),
);
sub_test(
"test_div_by_last_moduli_inplace::<ROUND:true, NTT:false>",
|| test_div_by_last_moduli_inplace::<true, false>(&ring_rns),
);
sub_test(
"test_div_by_last_moduli_inplace::<ROUND:true, NTT:true>",
|| test_div_by_last_moduli_inplace::<true, true>(&ring_rns),
);
} }
fn sub_test<F: FnOnce()>(name: &str, f: F) { fn sub_test<F: FnOnce()>(name: &str, f: F) {
@@ -58,7 +84,7 @@ fn test_div_by_last_modulus<const ROUND: bool, const NTT: bool>(ring_rns: &RingR
let mut source: Source = Source::new(seed); let mut source: Source = Source::new(seed);
let mut a: PolyRNS<u64> = ring_rns.new_polyrns(); let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
let mut b: PolyRNS<u64> = ring_rns.new_polyrns(); let mut buf: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
let mut c: PolyRNS<u64> = ring_rns.at_level(ring_rns.level() - 1).new_polyrns(); let mut c: PolyRNS<u64> = ring_rns.at_level(ring_rns.level() - 1).new_polyrns();
// Allocates a random PolyRNS // Allocates a random PolyRNS
@@ -75,7 +101,7 @@ fn test_div_by_last_modulus<const ROUND: bool, const NTT: bool>(ring_rns: &RingR
ring_rns.ntt_inplace::<false>(&mut a); ring_rns.ntt_inplace::<false>(&mut a);
} }
ring_rns.div_by_last_modulus::<ROUND, NTT>(&a, &mut b, &mut c); ring_rns.div_by_last_modulus::<ROUND, NTT>(&a, &mut buf, &mut c);
if NTT { if NTT {
ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c); ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c);
@@ -105,7 +131,7 @@ fn test_div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(ring_rns
let mut source: Source = Source::new(seed); let mut source: Source = Source::new(seed);
let mut a: PolyRNS<u64> = ring_rns.new_polyrns(); let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
let mut buf: PolyRNS<u64> = ring_rns.new_polyrns(); let mut buf: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
// Allocates a random PolyRNS // Allocates a random PolyRNS
ring_rns.fill_uniform(&mut source, &mut a); ring_rns.fill_uniform(&mut source, &mut a);
@@ -152,12 +178,13 @@ fn test_div_by_last_moduli<const ROUND: bool, const NTT: bool>(ring_rns: &RingRN
let seed: [u8; 32] = [0; 32]; let seed: [u8; 32] = [0; 32];
let mut source: Source = Source::new(seed); let mut source: Source = Source::new(seed);
let nb_moduli: usize = ring_rns.level(); let nb_moduli_dropped: usize = ring_rns.level();
let mut a: PolyRNS<u64> = ring_rns.new_polyrns(); let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
let mut buf: PolyRNS<u64> = ring_rns.new_polyrns(); let mut buf0: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
let mut buf1: PolyRNS<u64> = ring_rns.new_polyrns();
let mut c: PolyRNS<u64> = ring_rns let mut c: PolyRNS<u64> = ring_rns
.at_level(ring_rns.level() - nb_moduli) .at_level(ring_rns.level() - nb_moduli_dropped)
.new_polyrns(); .new_polyrns();
// Allocates a random PolyRNS // Allocates a random PolyRNS
@@ -174,7 +201,7 @@ fn test_div_by_last_moduli<const ROUND: bool, const NTT: bool>(ring_rns: &RingRN
ring_rns.ntt_inplace::<false>(&mut a); ring_rns.ntt_inplace::<false>(&mut a);
} }
ring_rns.div_by_last_moduli::<ROUND, NTT>(nb_moduli, &a, &mut buf, &mut c); ring_rns.div_by_last_moduli::<ROUND, NTT>(nb_moduli_dropped, &a, &mut buf0, &mut buf1, &mut c);
if NTT { if NTT {
ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c); ring_rns.at_level(c.level()).intt_inplace::<false>(&mut c);
@@ -188,7 +215,7 @@ fn test_div_by_last_moduli<const ROUND: bool, const NTT: bool>(ring_rns: &RingRN
// Performs floor division on a // Performs floor division on a
let mut scalar_big = BigInt::from(1); let mut scalar_big = BigInt::from(1);
(0..nb_moduli) (0..nb_moduli_dropped)
.for_each(|i| scalar_big *= BigInt::from(ring_rns.0[ring_rns.level() - i].modulus.q)); .for_each(|i| scalar_big *= BigInt::from(ring_rns.0[ring_rns.level() - i].modulus.q));
coeffs_a.iter_mut().for_each(|a| { coeffs_a.iter_mut().for_each(|a| {
if ROUND { if ROUND {
@@ -201,15 +228,15 @@ fn test_div_by_last_moduli<const ROUND: bool, const NTT: bool>(ring_rns: &RingRN
izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b)); izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
} }
/* fn test_div_by_last_moduli_inplace<const ROUND: bool, const NTT: bool>(ring_rns: &RingRNS<u64>) {
fn test_div_floor_by_last_moduli_inplace<const NTT: bool>(ring_rns: &RingRNS<u64>) {
let seed: [u8; 32] = [0; 32]; let seed: [u8; 32] = [0; 32];
let mut source: Source = Source::new(seed); let mut source: Source = Source::new(seed);
let nb_moduli: usize = ring_rns.level(); let nb_moduli_dropped: usize = ring_rns.level();
let mut a: PolyRNS<u64> = ring_rns.new_polyrns(); let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
let mut b: PolyRNS<u64> = ring_rns.new_polyrns(); let mut buf0: [Poly<u64>; 2] = [ring_rns.new_poly(), ring_rns.new_poly()];
let mut buf1: PolyRNS<u64> = ring_rns.new_polyrns();
// Allocates a random PolyRNS // Allocates a random PolyRNS
ring_rns.fill_uniform(&mut source, &mut a); ring_rns.fill_uniform(&mut source, &mut a);
@@ -225,23 +252,36 @@ fn test_div_floor_by_last_moduli_inplace<const NTT: bool>(ring_rns: &RingRNS<u64
ring_rns.ntt_inplace::<false>(&mut a); ring_rns.ntt_inplace::<false>(&mut a);
} }
ring_rns.div_floor_by_last_moduli_inplace::<NTT>(nb_moduli, &mut b, &mut a); ring_rns.div_by_last_moduli_inplace::<ROUND, NTT>(
nb_moduli_dropped,
&mut buf0,
&mut buf1,
&mut a,
);
if NTT { if NTT {
ring_rns.at_level(a.level()-nb_moduli).intt_inplace::<false>(&mut a); ring_rns
.at_level(a.level() - nb_moduli_dropped)
.intt_inplace::<false>(&mut a);
} }
// Exports c to coeffs_c // Exports c to coeffs_c
let mut coeffs_c = vec![BigInt::from(0); a.n()]; let mut coeffs_c = vec![BigInt::from(0); a.n()];
ring_rns ring_rns
.at_level(a.level()-nb_moduli) .at_level(a.level() - nb_moduli_dropped)
.to_bigint_inplace(&a, 1, &mut coeffs_c); .to_bigint_inplace(&a, 1, &mut coeffs_c);
// Performs floor division on a // Performs floor division on a
let mut scalar_big = BigInt::from(1); let mut scalar_big = BigInt::from(1);
(0..nb_moduli).for_each(|i|{scalar_big *= BigInt::from(ring_rns.0[ring_rns.level()-i].modulus.q)}); (0..nb_moduli_dropped)
coeffs_a.iter_mut().for_each(|a| {a.div_floor(&scalar_big)}); .for_each(|i| scalar_big *= BigInt::from(ring_rns.0[ring_rns.level() - i].modulus.q));
coeffs_a.iter_mut().for_each(|a| {
assert!(coeffs_a == coeffs_c, "test_div_floor_by_last_moduli_inplace"); if ROUND {
*a = a.div_round(&scalar_big);
} else {
*a = a.div_floor(&scalar_big);
}
});
izip!(coeffs_a, coeffs_c).for_each(|(a, b)| assert_eq!(a, b));
} }
*/