[ring]: added ring degree switching
Jean-Philippe Bossuat jeanphilippe.bossuat@gmail.com
Mon, 20 Jan 2025 11:44:27 +0100
5 files changed,
153 insertions(+),
4 deletions(-)
M
math/src/dft/ntt.rs
→
math/src/dft/ntt.rs
@@ -100,9 +100,17 @@ ) {
let n: usize = a.len(); assert!( n & n - 1 == 0, - "invalid x.len()= {} must be a power of two", + "invalid a.len()={} must be a power of two", n ); + + assert!( + n <= self.psi_forward_rev.len(), + "invalid a.len()={} > psi_forward_rev.len()={}", + n, + self.psi_forward_rev.len() + ); + let log_n: u32 = usize::BITS - ((n as usize) - 1).leading_zeros(); let start: u32 = SKIPSTART as u32;@@ -204,6 +212,14 @@ n & n - 1 == 0,
"invalid x.len()= {} must be a power of two", n ); + + assert!( + n <= self.psi_backward_rev.len(), + "invalid a.len()={} > psi_backward_rev.len()={}", + n, + self.psi_backward_rev.len() + ); + let log_n = usize::BITS - ((n as usize) - 1).leading_zeros(); let start: u32 = SKIPEND as u32;
M
math/src/ring/impl_u64/mod.rs
→
math/src/ring/impl_u64/mod.rs
@@ -2,5 +2,6 @@ pub mod automorphism;
pub mod rescaling_rns; pub mod ring; pub mod ring_rns; +pub mod ring_switch; pub mod sampling; pub mod utils;
A
math/src/ring/impl_u64/ring_switch.rs
@@ -0,0 +1,43 @@
+use crate::poly::Poly; +use crate::ring::Ring; + +impl Ring<u64> { + pub fn switch_degree<const NTT: bool>( + &self, + a: &Poly<u64>, + buf: &mut Poly<u64>, + b: &mut Poly<u64>, + ) { + let (n_in, n_out) = (a.n(), b.n()); + + if n_in > n_out { + let (gap_in, gap_out) = (1, n_in / n_out); + if NTT { + self.intt::<false>(&a, buf); + b.0.iter_mut() + .step_by(gap_in) + .zip(buf.0.iter().step_by(gap_out)) + .for_each(|(x_out, x_in)| *x_out = *x_in); + self.ntt_inplace::<false>(b); + } else { + b.0.iter_mut() + .step_by(gap_in) + .zip(a.0.iter().step_by(gap_out)) + .for_each(|(x_out, x_in)| *x_out = *x_in); + } + } else { + let gap: usize = n_out / n_in; + + if NTT { + a.0.iter() + .enumerate() + .for_each(|(i, &c)| (0..gap).for_each(|j| b.0[i * gap + j] = c)); + } else { + b.0.iter_mut() + .step_by(gap) + .zip(a.0.iter()) + .for_each(|(x_out, x_in)| *x_out = *x_in); + } + } + } +}
M
math/tests/automorphism.rs
→
math/tests/automorphism.rs
@@ -19,10 +19,10 @@ test_automorphism_native_u64::<true>(&ring, nth_root)
}); sub_test("test_automorphism_from_perm_u64::<NTT:false>", || { - test_automorphism_from_perm_u64::<false>(&ring, nth_root) + test_automorphism_from_perm_u64::<false>(&ring) }); sub_test("test_automorphism_from_perm_u64::<NTT:true>", || { - test_automorphism_from_perm_u64::<true>(&ring, nth_root) + test_automorphism_from_perm_u64::<true>(&ring) }); }@@ -62,7 +62,7 @@
izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b)); } -fn test_automorphism_from_perm_u64<const NTT: bool>(ring: &Ring<u64>, nth_root: usize) { +fn test_automorphism_from_perm_u64<const NTT: bool>(ring: &Ring<u64>) { let n: usize = ring.n(); let q: u64 = ring.modulus.q;
A
math/tests/ring_switch.rs
@@ -0,0 +1,89 @@
+use itertools::izip; +use math::automorphism::AutoPerm; +use math::poly::Poly; +use math::ring::Ring; + +#[test] +fn ring_switch_u64() { + let n: usize = 1 << 4; + let q_base: u64 = 65537u64; + let q_power: usize = 1usize; + let ring_small: Ring<u64> = Ring::new(n, q_base, q_power); + let ring_large = Ring::new(2 * n, q_base, q_power); + + sub_test("test_ring_switch_small_to_large_u64::<NTT:false>", || { + test_ring_switch_small_to_large_u64::<false>(&ring_small, &ring_large) + }); + sub_test("test_ring_switch_small_to_large_u64::<NTT:true>", || { + test_ring_switch_small_to_large_u64::<true>(&ring_small, &ring_large) + }); + sub_test("test_ring_switch_large_to_small_u64::<NTT:false>", || { + test_ring_switch_large_to_small_u64::<false>(&ring_small, &ring_large) + }); + sub_test("test_ring_switch_large_to_small_u64::<NTT:true>", || { + test_ring_switch_large_to_small_u64::<true>(&ring_small, &ring_large) + }); +} + +fn sub_test<F: FnOnce()>(name: &str, f: F) { + println!("Running {}", name); + f(); +} + +fn test_ring_switch_small_to_large_u64<const NTT: bool>( + ring_small: &Ring<u64>, + ring_large: &Ring<u64>, +) { + let mut a: Poly<u64> = ring_small.new_poly(); + let mut buf: Poly<u64> = ring_small.new_poly(); + let mut b: Poly<u64> = ring_large.new_poly(); + + a.0.iter_mut().enumerate().for_each(|(i, x)| *x = i as u64); + + if NTT { + ring_small.ntt_inplace::<false>(&mut a); + } + + ring_large.switch_degree::<NTT>(&a, &mut buf, &mut b); + + if NTT { + ring_small.intt_inplace::<false>(&mut a); + ring_large.intt_inplace::<false>(&mut b); + } + + let gap: usize = ring_large.n() / ring_small.n(); + + b.0.iter() + .step_by(gap) + .zip(a.0.iter()) + .for_each(|(x_out, x_in)| assert_eq!(x_out, x_in)); +} + +fn test_ring_switch_large_to_small_u64<const NTT: bool>( + ring_small: &Ring<u64>, + ring_large: &Ring<u64>, +) { + let mut a: Poly<u64> = ring_large.new_poly(); + let mut buf: Poly<u64> = ring_large.new_poly(); + let mut b: Poly<u64> = ring_small.new_poly(); + + a.0.iter_mut().enumerate().for_each(|(i, x)| *x = i as u64); + + if NTT { + ring_large.ntt_inplace::<false>(&mut a); + } + + ring_large.switch_degree::<NTT>(&a, &mut buf, &mut b); + + if NTT { + ring_large.intt_inplace::<false>(&mut a); + ring_small.intt_inplace::<false>(&mut b); + } + + let gap: usize = ring_large.n() / ring_small.n(); + + a.0.iter() + .step_by(gap) + .zip(b.0.iter()) + .for_each(|(x_out, x_in)| assert_eq!(x_out, x_in)); +}