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

@@ -1,6 +1,6 @@
use crate::modulus::barrett::Barrett;
use crate::modulus::{BARRETT, NONE, ONCE};
use crate::poly::PolyRNS;
use crate::poly::{Poly, PolyRNS};
use crate::ring::Ring;
use crate::ring::RingRNS;
use crate::scalar::ScalarRNS;
@@ -8,18 +8,19 @@ extern crate test;
impl RingRNS<u64> {
/// 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>(
&self,
a: &PolyRNS<u64>,
buf: &mut PolyRNS<u64>,
buf: &mut [Poly<u64>; 2],
b: &mut PolyRNS<u64>,
) {
debug_assert!(self.level() != 0, "invalid call: self.level()=0");
debug_assert!(
self.level() <= a.level(),
"invalid input a: self.level()={} > a.level()={}",
self.level(),
a.level()
a.level() >= self.level(),
"invalid input a: a.level()={} < self.level()={}",
a.level(),
self.level()
);
debug_assert!(
b.level() >= self.level() - 1,
@@ -35,7 +36,7 @@ impl RingRNS<u64> {
if ROUND {
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 {
r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
@@ -76,7 +77,7 @@ impl RingRNS<u64> {
}
} else {
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]);
for (i, r) in self.0[0..level].iter().enumerate() {
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.
pub fn div_by_last_modulus_inplace<const ROUND: bool, const NTT: bool>(
&self,
buf: &mut PolyRNS<u64>,
buf: &mut [Poly<u64>; 2],
a: &mut PolyRNS<u64>,
) {
debug_assert!(
self.level() <= a.level(),
"invalid input a: self.level()={} > a.level()={}",
"invalid input a: a.level()={} < self.level()={}",
self.level(),
a.level()
);
@@ -120,7 +121,7 @@ impl RingRNS<u64> {
if ROUND {
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 {
r_last.intt::<false>(a.at(level), &mut buf_q_scaling[0]);
@@ -152,7 +153,7 @@ impl RingRNS<u64> {
}
} else {
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]);
for (i, r) in self.0[0..level].iter().enumerate() {
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])
pub fn div_by_last_moduli<const ROUND: bool, const NTT: bool>(
&self,
nb_moduli: usize,
nb_moduli_dropped: usize,
a: &PolyRNS<u64>,
buf: &mut PolyRNS<u64>,
buf0: &mut [Poly<u64>; 2],
buf1: &mut PolyRNS<u64>,
c: &mut PolyRNS<u64>,
) {
debug_assert!(
nb_moduli <= self.level(),
"invalid input nb_moduli: nb_moduli={} > a.level()={}",
nb_moduli,
a.level()
nb_moduli_dropped <= self.level(),
"invalid input nb_moduli_dropped: nb_moduli_dropped={} > self.level()={}",
nb_moduli_dropped,
self.level()
);
debug_assert!(
a.level() <= self.level(),
"invalid input a: self.level()={} > a.level()={}",
self.level(),
a.level()
a.level() >= self.level(),
"invalid input a: a.level()={} < self.level()={}",
a.level(),
self.level()
);
debug_assert!(
buf.level() >= self.level() - 1,
"invalid input buf: buf.level()={} < a.level()-1={}",
buf.level(),
a.level() - 1
buf1.level() >= self.level(),
"invalid input buf: buf.level()={} < self.level()={}",
buf1.level(),
self.level()
);
debug_assert!(
c.level() >= self.level() - nb_moduli,
"invalid input c: c.level()={} < c.level()-nb_moduli={}",
c.level() >= self.level() - nb_moduli_dropped,
"invalid input c: c.level()={} < self.level()-nb_moduli_dropped={}",
c.level(),
a.level() - nb_moduli
self.level() - nb_moduli_dropped
);
if nb_moduli == 0 {
if nb_moduli_dropped == 0 {
if a != c {
c.copy(a);
}
} else {
if NTT {
self.intt::<false>(a, buf);
(0..nb_moduli).for_each(|i| {
self.intt::<false>(a, buf1);
(0..nb_moduli_dropped).for_each(|i| {
self.at_level(self.level() - i)
.div_by_last_modulus_inplace::<ROUND, false>(
&mut PolyRNS::<u64>::default(),
buf,
)
.div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1)
});
self.at_level(self.level() - nb_moduli).ntt::<false>(buf, c);
self.at_level(self.level() - nb_moduli_dropped)
.ntt::<false>(buf1, c);
} else {
println!("{} {:?}", self.level(), buf.level());
self.div_by_last_modulus::<ROUND, false>(a, buf, c);
self.div_by_last_modulus::<ROUND, false>(a, buf0, buf1);
(1..nb_moduli - 1).for_each(|i| {
println!("{} {:?}", self.level() - i, buf.level());
(1..nb_moduli_dropped - 1).for_each(|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)
.div_by_last_modulus_inplace::<ROUND, false>(buf, c);
self.at_level(self.level() - nb_moduli_dropped + 1)
.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>(
&self,
nb_moduli: usize,
buf: &mut PolyRNS<u64>,
nb_moduli_dropped: usize,
buf0: &mut [Poly<u64>; 2],
buf1: &mut PolyRNS<u64>,
a: &mut PolyRNS<u64>,
) {
debug_assert!(
self.level() <= a.level(),
"invalid input a: self.level()={} > a.level()={}",
self.level(),
a.level()
nb_moduli_dropped <= self.level(),
"invalid input nb_moduli_dropped: nb_moduli_dropped={} > self.level()={}",
nb_moduli_dropped,
self.level()
);
debug_assert!(
nb_moduli <= a.level(),
"invalid input nb_moduli: nb_moduli={} > a.level()={}",
nb_moduli,
a.level()
a.level() >= self.level(),
"invalid input a: a.level()={} < self.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;
}
if NTT {
self.intt::<false>(a, buf);
(0..nb_moduli).for_each(|i| {
self.intt::<false>(a, buf1);
(0..nb_moduli_dropped).for_each(|i| {
self.at_level(self.level() - i)
.div_by_last_modulus_inplace::<ROUND, false>(
&mut PolyRNS::<u64>::default(),
buf,
)
.div_by_last_modulus_inplace::<ROUND, false>(buf0, buf1)
});
self.at_level(self.level() - nb_moduli).ntt::<false>(buf, a);
self.at_level(self.level() - nb_moduli_dropped)
.ntt::<false>(buf1, a);
} else {
(0..nb_moduli).for_each(|i| {
(0..nb_moduli_dropped).for_each(|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)
});
}
}