moved plaintext ops related to fhevm to fhevm repo
Jean-Philippe Bossuat jeanphilippe.bossuat@gmail.com
Wed, 15 Jan 2025 16:11:15 +0100
5 files changed,
6 insertions(+),
514 deletions(-)
M
math/src/poly.rs
→
math/src/poly.rs
@@ -1,4 +1,5 @@
pub mod poly; +use std::cmp::min; use std::cmp::PartialEq; #[derive(Clone, Debug, Eq)]@@ -36,6 +37,11 @@ }
pub fn resize(&mut self, n: usize) { self.0.resize(n, O::default()); + } + + pub fn set(&mut self, v: &[O]) { + let size: usize = min(v.len(), self.n()); + self.0[..size].copy_from_slice(&v[..size]); } pub fn fill(&mut self, v: &O) {
M
math/src/ring/impl_u64/mod.rs
→
math/src/ring/impl_u64/mod.rs
@@ -1,8 +1,6 @@
pub mod automorphism; -pub mod packing; pub mod rescaling_rns; pub mod ring; pub mod ring_rns; pub mod sampling; -pub mod trace; pub mod utils;
D
math/src/ring/impl_u64/packing.rs
@@ -1,318 +0,0 @@
-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; - -impl Ring<u64> { - pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>( - &self, - polys: &mut Vec<Option<&mut Poly<u64>>>, - log_gap: usize, - ) { - let log_n: usize = self.log_n(); - let log_nth_root: usize = log_n + 1; - let nth_root: usize = 1 << log_nth_root; - let log_start: usize = log_n - log_gap; - let mut log_end: usize = log_n; - - let mut indices: Vec<usize> = Vec::<usize>::new(); - - // Retrives non-empty indexes - polys.iter().enumerate().for_each(|(i, poly)| { - if !poly.is_none() { - indices.push(i); - } - }); - - let gap: usize = max_gap(&indices); - - if !ZEROGARBAGE { - if gap > 0 { - log_end -= gap.trailing_zeros() as usize; - } - } - - assert!( - log_start < log_end, - "invalid input polys: gap between non None value is smaller than 2^log_gap" - ); - - let n_inv: Barrett<u64> = self - .modulus - .barrett - .prepare(self.modulus.inv(1 << (log_end - log_start))); - - indices.iter().for_each(|i| { - if let Some(poly) = polys[*i].as_deref_mut() { - if !NTT { - self.ntt_inplace::<true>(poly); - } - self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly); - } - }); - - let x_pow2: Vec<Poly<Montgomery<u64>>> = self.gen_x_pow_2::<true, false>(log_n); - let mut tmpa: Poly<u64> = self.new_poly(); - let mut tmpb: Poly<u64> = self.new_poly(); - - for i in log_start..log_end { - let t: usize = 1 << (log_n - 1 - i); - - let (polys_lo, polys_hi) = polys.split_at_mut(t); - - for j in 0..t { - if let Some(poly_hi) = polys_hi[j].as_mut() { - self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi); - - if let Some(poly_lo) = polys_lo[j].as_mut() { - self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa); - self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo); - } - } - - if let Some(poly_lo) = polys_lo[j].as_mut() { - let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root); - - if !polys_hi[j].is_none() { - self.a_apply_automorphism_add_b_into_b::<ONCE, true>( - &tmpa, - gal_el, - 2 << self.log_n(), - poly_lo, - ); - } else { - self.a_apply_automorphism_into_b::<true>( - poly_lo, gal_el, nth_root, &mut tmpa, - ); - self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo); - } - } else if let Some(poly_hi) = polys_hi[j].as_mut() { - let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root); - self.a_apply_automorphism_into_b::<true>(poly_hi, gal_el, nth_root, &mut tmpa); - self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi); - std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]); - } - } - - polys.truncate(t); - } - - if !NTT { - if let Some(poly) = polys[0].as_mut() { - self.intt_inplace::<false>(poly); - } - } - } -} - -// Returns the largest gap between two values in an ordered array of distinct values. -// Panics if the array is not ordered or values are not distincts. -fn max_gap(vec: &[usize]) -> usize { - let mut gap: usize = usize::MAX; - for i in 1..vec.len() { - let (l, r) = (vec[i - 1], vec[i]); - assert!( - r > l, - "invalid input vec: not sorted or collision between indices" - ); - gap = min(gap, r - l); - if gap == 1 { - break; - } - } - gap -} - - -pub struct StreamRepacker { - accumulators: Vec<Accumulator>, - tmp_a: Poly<u64>, - tmp_b: Poly<u64>, - x_pow_2: Vec<Poly<Montgomery<u64>>>, - n_inv: Barrett<u64>, - pub results: Vec<Poly<u64>>, - counter: usize, -} - -pub struct Accumulator { - buf: Poly<u64>, - value: bool, - control: bool, -} - -impl Accumulator { - pub fn new(r: &Ring<u64>) -> Self { - Self { - buf: r.new_poly(), - value: false, - control: false, - } - } -} - -impl StreamRepacker { - pub fn new(r: &Ring<u64>) -> Self { - let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new(); - - (0..r.log_n()).for_each(|_| accumulators.push(Accumulator::new(r))); - - Self { - accumulators: accumulators, - tmp_a: r.new_poly(), - tmp_b: r.new_poly(), - x_pow_2: r.gen_x_pow_2::<true, false>(r.log_n()), - n_inv: r.modulus.barrett.prepare(r.modulus.inv(r.n() as u64)), - results: Vec::<Poly<u64>>::new(), - counter: 0, - } - } - - pub fn reset(&mut self) { - for i in 0..self.accumulators.len() { - self.accumulators[i].value = false; - self.accumulators[i].control = false; - } - self.counter = 0; - } - - pub fn add<const NTT: bool>(&mut self, r: &Ring<u64>, a: Option<&Poly<u64>>) { - assert!(NTT, "invalid parameterization: const NTT must be true"); - pack_core::<NTT>( - r, - a, - &mut self.accumulators, - &self.n_inv, - &self.x_pow_2, - &mut self.tmp_a, - &mut self.tmp_b, - 0, - ); - self.counter += 1; - if self.counter == r.n() { - self.results - .push(self.accumulators[r.log_n() - 1].buf.clone()); - self.reset(); - } - } - - pub fn flush<const NTT: bool>(&mut self, r: &Ring<u64>) { - assert!(NTT, "invalid parameterization: const NTT must be true"); - if self.counter != 0 { - while self.counter != r.n() - 1 { - self.add::<NTT>(r, None); - } - } - } -} - -fn pack_core<const NTT: bool>( - r: &Ring<u64>, - a: Option<&Poly<u64>>, - accumulators: &mut [Accumulator], - n_inv: &Barrett<u64>, - x_pow_2: &[Poly<u64>], - tmp_a: &mut Poly<u64>, - tmp_b: &mut Poly<u64>, - i: usize, -) { - if i == r.log_n() { - return; - } - - let (acc_prev, acc_next) = accumulators.split_at_mut(1); - - if !acc_prev[0].control { - let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut - - if let Some(a_ref) = a { - acc_mut_ref.buf.copy_from(a_ref); - acc_mut_ref.value = true - } else { - acc_mut_ref.value = false - } - acc_mut_ref.control = true; - } else { - combine::<true>(r, &mut acc_prev[0], a, n_inv, x_pow_2, tmp_a, tmp_b, i); - acc_prev[0].control = false; - - if acc_prev[0].value { - pack_core::<NTT>( - r, - Some(&acc_prev[0].buf), - acc_next, - n_inv, - x_pow_2, - tmp_a, - tmp_b, - i + 1, - ); - } else { - pack_core::<NTT>(r, None, acc_next, n_inv, x_pow_2, tmp_a, tmp_b, i + 1); - } - } -} - -fn combine<const NTT: bool>( - r: &Ring<u64>, - acc: &mut Accumulator, - b: Option<&Poly<u64>>, - n_inv: &Barrett<u64>, - x_pow_2: &[Poly<u64>], - tmp_a: &mut Poly<u64>, - tmp_b: &mut Poly<u64>, - i: usize, -) { - let log_n = r.log_n(); - let log_nth_root = log_n + 1; - let nth_root = 1 << log_nth_root; - let gal_el: usize = r.galois_element((1 << i) >> 1, i == 0, log_nth_root); - - let a: &mut Poly<u64> = &mut acc.buf; - - if acc.value { - if i == 0 { - r.a_mul_b_scalar_barrett_into_a::<ONCE>(n_inv, a); - } - - if let Some(b) = b { - // tmp_a = b * X^t - r.a_mul_b_montgomery_into_c::<ONCE>(b, &x_pow_2[log_n - i - 1], tmp_a); - - if i == 0 { - r.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, tmp_a); - } - - // tmp_b = a - b*X^t - r.a_sub_b_into_c::<1, ONCE>(a, tmp_a, tmp_b); - - // a = a + b * X^t - r.a_add_b_into_b::<ONCE>(tmp_a, a); - - // a = a + b * X^t + phi(a - b * X^t) - r.a_apply_automorphism_add_b_into_b::<ONCE, NTT>(tmp_b, gal_el, nth_root, a); - } else { - // tmp_a = phi(a) - r.a_apply_automorphism_into_b::<NTT>(a, gal_el, nth_root, tmp_a); - // a = a + phi(a) - r.a_add_b_into_b::<ONCE>(tmp_a, a); - } - } else { - if let Some(b) = b { - // tmp_b = b * X^t - r.a_mul_b_montgomery_into_c::<ONCE>(b, &x_pow_2[log_n - i - 1], tmp_b); - - if i == 0 { - r.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, tmp_b); - } - - // tmp_a = phi(b * X^t) - r.a_apply_automorphism_into_b::<NTT>(tmp_b, gal_el, nth_root, tmp_a); - - // a = (b* X^t - phi(b* X^t)) - r.a_sub_b_into_c::<1, ONCE>(tmp_b, tmp_a, a); - acc.value = true - } - } -}
D
math/src/ring/impl_u64/trace.rs
@@ -1,36 +0,0 @@
-use crate::modulus::barrett::Barrett; -use crate::modulus::ONCE; -use crate::poly::Poly; -use crate::ring::Ring; - -impl Ring<u64> { - pub fn trace_inplace<const NTT: bool>(&self, step_start: usize, a: &mut Poly<u64>) { - assert!( - step_start <= self.log_n(), - "invalid argument step_start: step_start={} > self.log_n()={}", - step_start, - self.log_n() - ); - - let log_steps: usize = self.log_n() - step_start; - let log_nth_root = self.log_n() + 1; - let nth_root: usize = 1 << log_nth_root; - - if log_steps > 0 { - let n_inv: Barrett<u64> = self - .modulus - .barrett - .prepare(self.modulus.inv(1 << log_steps)); - self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, a); - - let mut tmp: Poly<u64> = self.new_poly(); - - (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::<NTT>(a, gal_el, nth_root, &mut tmp); - self.a_add_b_into_b::<ONCE>(&tmp, a); - }); - } - } -}
M
math/tests/automorphism.rs
→
math/tests/automorphism.rs
@@ -1,7 +1,5 @@
use itertools::izip; -use math::modulus::WordOps; use math::poly::Poly; -use math::ring::impl_u64::packing::StreamRepacker; use math::ring::Ring; #[test]@@ -53,159 +51,3 @@ }
izip!(p0.0, p1.0).for_each(|(a, b)| assert_eq!(a, b)); } - -#[test] -fn packing_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_u64::<NTT:false>", || { - test_packing_sparse_u64::<false>(&ring, 1) - }); - sub_test("test_packing_u64::<NTT:true>", || { - test_packing_sparse_u64::<true>(&ring, 1) - }); - sub_test("test_packing_sparse_u64::<NTT:false>", || { - test_packing_sparse_u64::<false>(&ring, 3) - }); - sub_test("test_packing_sparse_u64::<NTT:true>", || { - test_packing_sparse_u64::<true>(&ring, 3) - }); -} - -fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>, gap: usize) { - let n: usize = ring.n(); - - let mut result: Vec<Option<&mut Poly<u64>>> = Vec::with_capacity(n); - result.resize_with(n, || None); - - let mut polys: Vec<Poly<u64>> = vec![ring.new_poly(); (n+gap-1)/gap]; - - polys.iter_mut().enumerate().for_each(|(i , poly)|{ - poly.fill(&((1 + i*gap) as u64)); - if NTT { - ring.ntt_inplace::<false>(poly); - } - result[i*gap] = Some(poly); - }); - - ring.pack::<true, NTT>(&mut result, ring.log_n()); - - if let Some(poly) = result[0].as_mut() { - - if NTT { - ring.intt_inplace::<false>(poly); - } - - poly.0.iter().enumerate().for_each(|(i, x)| { - if i % gap == 0 { - assert_eq!(*x, (1+i) as u64) - } else { - assert_eq!(*x, 0u64) - } - }); - } -} - -#[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() { - 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_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 { - ring.ntt_inplace::<false>(&mut poly); - } - - let step_start: usize = 2; - - ring.trace_inplace::<NTT>(step_start, &mut poly); - - if NTT { - ring.intt_inplace::<false>(&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) - } - }); -}