added stream repacking & associated tests

This commit is contained in:
Jean-Philippe Bossuat
2025-01-13 17:20:30 +01:00
parent a8bca16047
commit 47a0894414
6 changed files with 333 additions and 155 deletions

View File

@@ -1,5 +1,7 @@
use itertools::izip;
use math::modulus::WordOps;
use math::poly::Poly;
use math::ring::impl_u64::packing::StreamRepacker;
use math::ring::Ring;
#[test]
@@ -135,7 +137,62 @@ fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>) {
}
}
#[test]
fn packing_streaming_u64() {
let n: usize = 1 << 5;
let q_base: u64 = 65537u64;
let q_power: usize = 1usize;
let ring: Ring<u64> = Ring::new(n, q_base, q_power);
sub_test("test_packing_streaming_dense_u64::<NTT:true>", || {
test_packing_streaming_dense_u64::<true>(&ring)
});
}
fn test_packing_streaming_dense_u64<const NTT: bool>(ring: &Ring<u64>) {
let n: usize = ring.n();
let mut values: Vec<u64> = vec![0; n];
values
.iter_mut()
.enumerate()
.for_each(|(i, x)| *x = (i + 1) as u64);
let gap: usize = 3;
let mut packer = StreamRepacker::new(ring);
let mut poly: Poly<u64> = ring.new_poly();
for i in 0..n {
let i_rev: usize = i.reverse_bits_msb(ring.log_n() as u32);
if i_rev % gap == 0 {
poly.fill(&values[i_rev]);
if NTT {
ring.ntt_inplace::<false>(&mut poly);
}
packer.add::<NTT>(ring, Some(&poly))
} else {
packer.add::<NTT>(ring, None)
}
}
packer.flush::<NTT>(ring);
let result: &mut Poly<u64> = &mut packer.results[0];
if NTT {
ring.intt_inplace::<false>(result);
}
result.0.iter().enumerate().for_each(|(i, x)| {
if i % gap == 0 {
assert_eq!(*x, values[i] as u64)
} else {
assert_eq!(*x, 0u64)
}
});
}
#[test]
fn trace_u64() {
@@ -144,24 +201,21 @@ fn trace_u64() {
let q_power: usize = 1usize;
let ring: Ring<u64> = Ring::new(n, q_base, q_power);
sub_test("test_trace::<NTT:false>", || {
test_trace_u64::<false>(&ring)
});
sub_test("test_trace::<NTT:true>", || {
test_trace_u64::<true>(&ring)
});
sub_test("test_trace::<NTT:false>", || test_trace_u64::<false>(&ring));
sub_test("test_trace::<NTT:true>", || test_trace_u64::<true>(&ring));
}
fn test_trace_u64<const NTT: bool>(ring: &Ring<u64>) {
let n: usize = ring.n();
let mut poly: Poly<u64> = ring.new_poly();
poly.0.iter_mut().enumerate().for_each(|(i, x)|{
*x = (i+1) as u64
});
if NTT{
poly.0
.iter_mut()
.enumerate()
.for_each(|(i, x)| *x = (i + 1) as u64);
if NTT {
ring.ntt_inplace::<false>(&mut poly);
}
@@ -169,11 +223,11 @@ fn test_trace_u64<const NTT: bool>(ring: &Ring<u64>) {
ring.trace_inplace::<NTT>(step_start, &mut poly);
if NTT{
if NTT {
ring.intt_inplace::<false>(&mut poly);
}
let gap: usize = 1<<(ring.log_n() - step_start);
let gap: usize = 1 << (ring.log_n() - step_start);
poly.0.iter().enumerate().for_each(|(i, x)| {
if i % gap == 0 {
@@ -182,4 +236,4 @@ fn test_trace_u64<const NTT: bool>(ring: &Ring<u64>) {
assert_eq!(*x, 0u64)
}
});
}
}