From a8bca160475a2d078110281b5ea45f7233cc955e Mon Sep 17 00:00:00 2001 From: Jean-Philippe Bossuat Date: Sat, 11 Jan 2025 15:07:22 +0100 Subject: [PATCH] streamrepacker wip --- math/src/ring/impl_u64/packing.rs | 113 +++++++++++++++++++++++++++++- math/src/ring/impl_u64/trace.rs | 23 +++--- math/tests/automorphism.rs | 50 ++++++++++++- 3 files changed, 171 insertions(+), 15 deletions(-) diff --git a/math/src/ring/impl_u64/packing.rs b/math/src/ring/impl_u64/packing.rs index f485112..8b1002f 100644 --- a/math/src/ring/impl_u64/packing.rs +++ b/math/src/ring/impl_u64/packing.rs @@ -1,8 +1,10 @@ use crate::modulus::barrett::Barrett; +use crate::modulus::montgomery::Montgomery; use crate::modulus::ONCE; use crate::poly::Poly; use crate::ring::Ring; use std::cmp::min; +use std::rc::Rc; impl Ring { @@ -53,7 +55,7 @@ impl Ring { } }); - let x_pow2: Vec> = self.gen_x_pow_2::(log_n); + let x_pow2: Vec>> = self.gen_x_pow_2::(log_n); let mut tmpa: Poly = self.new_poly(); let mut tmpb: Poly = self.new_poly(); @@ -115,3 +117,112 @@ fn max_gap(vec: &[usize]) -> usize { } gap } + + +pub struct StreamRepacker{ + accumulators: Vec, + buf0: Poly, + buf1: Poly, + buf_auto: Poly, + x_pow_2: Vec>>, + n_inv: Barrett, + counter: usize, +} + +pub struct Accumulator{ + buf: [Option>>; 2], + control: bool, +} + +impl Accumulator{ + pub fn new(r: &Ring) -> Self{ + Self { buf: [Some(Rc::new(r.new_poly())), None], control: false } + } +} + +impl StreamRepacker{ + pub fn new(r: &Ring) -> Self{ + + let mut accumulators: Vec = Vec::::new(); + + (0..r.log_n()).for_each(|_| + accumulators.push(Accumulator::new(r)) + ); + + Self{ + accumulators: accumulators, + buf0: r.new_poly(), + buf1: r.new_poly(), + buf_auto: r.new_poly(), + x_pow_2: r.gen_x_pow_2::(r.log_n()), + n_inv: r.modulus.barrett.prepare(r.modulus.inv(r.n() as u64)), + counter:0, + } + } + + fn merge_ab(&mut self, r: &Ring, a: &Poly, b: &Poly, i: usize) -> &Poly{ + + let tmp_a: &mut Poly = &mut self.buf0; + let tmp_b: &mut Poly = &mut self.buf1; + + r.a_mul_b_montgomery_into_c::(a, &self.x_pow_2[r.log_n()-i-1], tmp_a); + r.a_sub_b_into_c::<1, ONCE>(a, tmp_a, tmp_b); + r.a_add_b_into_b::(a, tmp_a); + + if i == 0{ + r.a_mul_b_scalar_barrett_into_a::(&self.n_inv, tmp_a); + r.a_mul_b_scalar_barrett_into_a::(&self.n_inv, tmp_b); + } + + let log_nth_root = r.log_n()+1; + let nth_root = 1<> 1, i == 0, log_nth_root); + + r.a_apply_automorphism_add_b_into_b::(tmp_b, gal_el, nth_root, tmp_a); + + tmp_a + } + + fn merge_a(&mut self, r: &Ring, a: &Poly, i: usize) -> &Poly{ + + let tmp_a: &mut Poly = &mut self.buf0; + + let log_nth_root = r.log_n()+1; + let nth_root = 1<> 1, i == 0, log_nth_root); + + if i == 0{ + r.a_mul_b_scalar_barrett_into_a::(&self.n_inv, tmp_a); + r.a_apply_automorphism_into_b::(tmp_a, gal_el, nth_root, &mut self.buf_auto) + r.a_add_b_into_c::(&self.buf_auto, a, tmp_a); + }else{ + r.a_apply_automorphism_into_b::(a, gal_el, nth_root, tmp_a); + r.a_add_b_into_b::(a, tmp_a); + } + + tmp_a + } + + fn merge_b(&mut self, r: &Ring, b: &Poly, i: usize) -> &Poly{ + + let tmp_a: &mut Poly = &mut self.buf0; + let tmp_b: &mut Poly = &mut self.buf1; + + let log_nth_root = r.log_n()+1; + let nth_root = 1<> 1, i == 0, log_nth_root); + + if i == 0{ + r.a_mul_b_scalar_barrett_into_c::(&self.n_inv, b, tmp_b); + r.a_mul_b_montgomery_into_a::(&self.x_pow_2[r.log_n()-i-1], tmp_b); + }else{ + r.a_mul_b_montgomery_into_c::(b, &self.x_pow_2[r.log_n()-i-1], tmp_b); + } + + r.a_apply_automorphism_into_b::(tmp_b, gal_el, nth_root, &mut self.buf_auto); + r.a_sub_b_into_a::<1, ONCE>(&self.buf_auto, tmp_b); + + tmp_b + } +} \ No newline at end of file diff --git a/math/src/ring/impl_u64/trace.rs b/math/src/ring/impl_u64/trace.rs index 6402841..78e0a56 100644 --- a/math/src/ring/impl_u64/trace.rs +++ b/math/src/ring/impl_u64/trace.rs @@ -4,29 +4,26 @@ use crate::modulus::barrett::Barrett; use crate::modulus::ONCE; impl Ring{ - pub fn trace_inplace(&self, log_start: usize, log_end: usize, a: &mut Poly){ - assert!(log_end <= self.log_n(), "invalid argument log_end: log_end={} > self.log_n()={}", log_end, self.log_n()); - assert!(log_end > log_start, "invalid argument log_start: log_start={} > log_end={}", log_start, log_end); + pub fn trace_inplace(&self, step_start: usize, a: &mut Poly){ + assert!(step_start <= self.log_n(), "invalid argument step_start: step_start={} > self.log_n()={}", step_start, self.log_n()); - let log_steps = log_end - log_start; + let log_steps: usize = self.log_n() - step_start; + let log_nth_root = self.log_n()+1; + let nth_root: usize= 1< 0 { let n_inv: Barrett = self.modulus.barrett.prepare(self.modulus.inv(1<(&n_inv, a); - if !NTT{ - self.ntt_inplace::(a); - } - let mut tmp: Poly = self.new_poly(); - (log_start..log_end).for_each(|i|{ + (step_start..self.log_n()).for_each(|i|{ + let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root); + + self.a_apply_automorphism_into_b::(a, gal_el, nth_root, &mut tmp); + self.a_add_b_into_b::(&tmp, a); }); } - - - - } } \ No newline at end of file diff --git a/math/tests/automorphism.rs b/math/tests/automorphism.rs index 3e88be4..dee2aa6 100644 --- a/math/tests/automorphism.rs +++ b/math/tests/automorphism.rs @@ -1,6 +1,5 @@ use itertools::izip; use math::poly::Poly; -use math::ring::impl_u64::ring; use math::ring::Ring; #[test] @@ -135,3 +134,52 @@ fn test_packing_sparse_u64(ring: &Ring) { }); } } + + + +#[test] +fn trace_u64() { + let n: usize = 1 << 5; + let q_base: u64 = 65537u64; + let q_power: usize = 1usize; + let ring: Ring = Ring::new(n, q_base, q_power); + + sub_test("test_trace::", || { + test_trace_u64::(&ring) + }); + sub_test("test_trace::", || { + test_trace_u64::(&ring) + }); +} + +fn test_trace_u64(ring: &Ring) { + let n: usize = ring.n(); + + let mut poly: Poly = ring.new_poly(); + + poly.0.iter_mut().enumerate().for_each(|(i, x)|{ + *x = (i+1) as u64 + }); + + if NTT{ + ring.ntt_inplace::(&mut poly); + } + + let step_start: usize = 2; + + ring.trace_inplace::(step_start, &mut poly); + + if NTT{ + ring.intt_inplace::(&mut poly); + } + + let gap: usize = 1<<(ring.log_n() - step_start); + + poly.0.iter().enumerate().for_each(|(i, x)| { + if i % gap == 0 { + assert_eq!(*x, 1 + i as u64) + } else { + assert_eq!(*x, 0u64) + } + }); +} \ No newline at end of file