[pack]: replaced HashMap by [Option<Poly<u64>]
Jean-Philippe Bossuat jeanphilippe.bossuat@gmail.com
Fri, 10 Jan 2025 11:48:51 +0100
8 files changed,
232 insertions(+),
168 deletions(-)
M
Cargo.lock
→
Cargo.lock
@@ -12,6 +12,17 @@ "memchr",
] [[package]] +name = "alga" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f823d037a7ec6ea2197046bafd4ae150e6bc36f9ca347404f46a46823fa84f2" +dependencies = [ + "approx", + "num-complex 0.2.4", + "num-traits", +] + +[[package]] name = "anes" version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index"@@ -22,6 +33,15 @@ name = "anstyle"
version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" + +[[package]] +name = "approx" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0e60b75072ecd4168020818c0107f2857bb6c4e64252d8d3983f6263b40a5c3" +dependencies = [ + "num-traits", +] [[package]] name = "autocfg"@@ -179,12 +199,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" [[package]] -name = "equivalent" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" - -[[package]] name = "getrandom" version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index"@@ -206,10 +220,10 @@ "crunchy",
] [[package]] -name = "hashbrown" -version = "0.15.2" +name = "hermit-abi" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" [[package]] name = "hermit-abi"@@ -218,22 +232,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" [[package]] -name = "indexmap" -version = "2.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f" -dependencies = [ - "equivalent", - "hashbrown", -] - -[[package]] name = "is-terminal" version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "261f68e344040fbd0edea105bef17c66edf46f984ddb1115b775ce31be948f4b" dependencies = [ - "hermit-abi", + "hermit-abi 0.4.0", "libc", "windows-sys 0.52.0", ]@@ -295,7 +299,6 @@ name = "math"
version = "0.1.0" dependencies = [ "criterion", - "indexmap", "itertools 0.14.0", "num", "num-bigint",@@ -305,6 +308,17 @@ "primality-test",
"prime_factorization", "rand_distr", "sampling", + "sprs", +] + +[[package]] +name = "matrixmultiply" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9380b911e3e96d10c1f415da0876389aaf1b56759054eeb0de7df940c456ba1a" +dependencies = [ + "autocfg", + "rawpointer", ] [[package]]@@ -314,13 +328,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex 0.4.6", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] name = "num" version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" dependencies = [ "num-bigint", - "num-complex", + "num-complex 0.4.6", "num-integer", "num-iter", "num-rational",@@ -334,6 +363,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" dependencies = [ "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6b19411a9719e753aff12e5187b74d60d3dc449ec3f4dc21e3989c3f554bc95" +dependencies = [ + "autocfg", "num-traits", ]@@ -388,6 +427,16 @@ "libm",
] [[package]] +name = "num_cpus" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +dependencies = [ + "hermit-abi 0.3.9", + "libc", +] + +[[package]] name = "once_cell" version = "1.20.2" source = "registry+https://github.com/rust-lang/crates.io-index"@@ -428,6 +477,21 @@ "plotters-backend",
] [[package]] +name = "portable-atomic" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6" + +[[package]] +name = "portable-atomic-util" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507" +dependencies = [ + "portable-atomic", +] + +[[package]] name = "ppv-lite86" version = "0.2.20" source = "registry+https://github.com/rust-lang/crates.io-index"@@ -510,6 +574,12 @@ dependencies = [
"num-traits", "rand", ] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" [[package]] name = "rayon"@@ -613,6 +683,27 @@ "itoa",
"memchr", "ryu", "serde", +] + +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + +[[package]] +name = "sprs" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704ef26d974e8a452313ed629828cd9d4e4fa34667ca1ad9d6b1fffa43c6e166" +dependencies = [ + "alga", + "ndarray", + "num-complex 0.4.6", + "num-traits", + "num_cpus", + "rayon", + "smallvec", ] [[package]]
M
math/Cargo.toml
→
math/Cargo.toml
@@ -13,7 +13,8 @@ prime_factorization = "1.0.5"
itertools = "0.14.0" criterion = "0.5.1" rand_distr = "0.4.3" -indexmap = "2.7.0" +sprs = "0.11.2" + sampling = { path = "../sampling" } [[bench]]
M
math/src/modulus/impl_u64/operations.rs
→
math/src/modulus/impl_u64/operations.rs
@@ -100,7 +100,11 @@ *a = self.montgomery.prepare::<REDUCE>(*a);
} #[inline(always)] - fn sa_prepare_montgomery_into_sb<const REDUCE: REDUCEMOD>(&self, a: &u64, b: &mut Montgomery<u64>) { + fn sa_prepare_montgomery_into_sb<const REDUCE: REDUCEMOD>( + &self, + a: &u64, + b: &mut Montgomery<u64>, + ) { self.montgomery.prepare_assign::<REDUCE>(*a, b); }@@ -330,7 +334,13 @@ &self,
a: &[u64], b: &mut [Montgomery<u64>], ) { - apply_vv!(self, Self::sa_prepare_montgomery_into_sb::<REDUCE>, a, b, CHUNK); + apply_vv!( + self, + Self::sa_prepare_montgomery_into_sb::<REDUCE>, + a, + b, + CHUNK + ); } #[inline(always)]@@ -338,9 +348,14 @@ fn va_prepare_montgomery_into_va<const CHUNK: usize, const REDUCE: REDUCEMOD>(
&self, a: &mut [Montgomery<u64>], ) { - apply_v!(self, Self::sa_prepare_montgomery_into_sa::<REDUCE>, a, CHUNK); + apply_v!( + self, + Self::sa_prepare_montgomery_into_sa::<REDUCE>, + a, + CHUNK + ); } - + #[inline(always)] fn va_mont_mul_vb_into_vc<const CHUNK: usize, const REDUCE: REDUCEMOD>( &self,
M
math/src/ring.rs
→
math/src/ring.rs
@@ -2,8 +2,8 @@ pub mod impl_u64;
use crate::dft::DFT; use crate::modulus::prime::Prime; -use crate::poly::{Poly, PolyRNS}; use crate::modulus::WordOps; +use crate::poly::{Poly, PolyRNS}; use num::traits::Unsigned; use std::rc::Rc;@@ -14,8 +14,7 @@ pub dft: Box<dyn DFT<O>>,
} impl<O: Unsigned> Ring<O> { - - pub fn log_n(&self) -> usize{ + pub fn log_n(&self) -> usize { return self.n().log2(); }@@ -31,8 +30,7 @@
pub struct RingRNS<O: Unsigned>(pub Vec<Rc<Ring<O>>>); impl<O: Unsigned> RingRNS<O> { - - pub fn log_n(&self) -> usize{ + pub fn log_n(&self) -> usize { return self.n().log2(); }
M
math/src/ring/impl_u64/mod.rs
→
math/src/ring/impl_u64/mod.rs
@@ -1,6 +1,6 @@
pub mod automorphism; +pub mod packing; pub mod rescaling_rns; pub mod ring; pub mod ring_rns; pub mod sampling; -pub mod packing;
M
math/src/ring/impl_u64/packing.rs
→
math/src/ring/impl_u64/packing.rs
@@ -1,188 +1,143 @@
-use std::collections::HashMap; +use crate::modulus::barrett::Barrett; +use crate::modulus::{WordOps, ONCE}; use crate::poly::Poly; use crate::ring::Ring; -use crate::modulus::{ONCE, WordOps}; -use crate::modulus::barrett::Barrett; use std::cmp::min; +use std::collections::HashSet; use std::mem::transmute; - -impl Ring<u64>{ - +impl Ring<u64> { // Generates a vector storing {X^{2^0}, X^{2^1}, .., X^{2^log_n}}. - pub fn gen_x_pow_2<const NTT: bool, const INV: bool>(&self, log_n: usize) -> Vec<Poly<u64>>{ + pub fn gen_x_pow_2<const NTT: bool, const INV: bool>(&self, log_n: usize) -> Vec<Poly<u64>> { let mut x_pow: Vec<Poly<u64>> = Vec::<Poly<u64>>::with_capacity(log_n); - (0..log_n).for_each(|i|{ - let mut idx: usize = 1<<i; + (0..log_n).for_each(|i| { + let mut idx: usize = 1 << i; - if INV{ + if INV { idx = self.n() - idx; } x_pow.push(self.new_poly()); - if i == 0{ + if i == 0 { x_pow[i].0[idx] = self.modulus.montgomery.one(); self.ntt_inplace::<false>(&mut x_pow[i]); - }else{ + } else { let (left, right) = x_pow.split_at_mut(i); - self.a_mul_b_montgomery_into_c::<ONCE>(&left[i-1], &left[i-1], &mut right[0]); + self.a_mul_b_montgomery_into_c::<ONCE>(&left[i - 1], &left[i - 1], &mut right[0]); } }); - if INV{ + if INV { self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]); } - if !NTT{ + if !NTT { x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x)); } x_pow } - pub fn pack<'a, const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: &'a mut HashMap<usize, &'a mut Poly<u64>>, log_gap: usize) -> &'a Poly<u64>{ - + pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>( + &self, + polys: &mut [Option<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_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 keys: Vec<usize> = polys.keys().copied().collect(); - keys.sort(); + let mut indices: Vec<usize> = Vec::<usize>::new(); - let mut gap = 0usize; + // Retrives non-empty indexes + polys.iter().enumerate().for_each(|(i, poly)| { + if Some(poly) != None { + indices.push(i); + } + }); - if keys.len() > 1{ - gap = max_pow2_gap(&keys); - }else{ - gap = 1<<log_n; - } + let gap: usize = max_gap(&indices); + + let set: HashSet<_> = indices.into_iter().collect(); - let log_gap: usize = gap.log2(); + let max_pow2_gap_divisor: usize = 1 << gap.trailing_zeros(); - if !ZEROGARBAGE{ + if !ZEROGARBAGE { if gap > 0 { - log_end -= log_gap; + log_end -= max_pow2_gap_divisor; } } - - let n_inv: Barrett<u64> = self.modulus.barrett.prepare(self.modulus.inv(1<<(log_end - log_start))); - for (_, poly) in polys.iter_mut() { - if !NTT{ - self.ntt_inplace::<true>(poly); - } + let n_inv: Barrett<u64> = self + .modulus + .barrett + .prepare(self.modulus.inv(1 << (log_end - log_start))); - self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly); - } + set.iter().for_each(|i| { + if let Some(poly) = polys[*i].as_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<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{ + for i in log_start..log_end { + let t: usize = 1 << (log_n - 1 - i); - let t: usize = 1<<(log_n-1-i); + let (polys_lo, polys_hi) = polys.split_at_mut(t); - for j in 0..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); - let option_lo: Option<&&mut Poly<u64>> = polys.get(&i); - let option_hi: Option<&&mut Poly<u64>> = polys.get(&(i+t)); - let mut hi_exists: bool = false; - - match option_hi{ - Some(hi) =>{ - - // Unsafe code is necessary because two mutables references are - // accessed from the map. - unsafe{ - self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n-i-1], transmute(*hi as *const Poly<u64> as *mut Poly<u64>)); - } - - hi_exists = true; - - match option_lo{ - Some(lo) =>{ - - self.a_sub_b_into_c::<1, ONCE>(lo, hi, &mut tmpa); - - // Ensures unsafe blocks are "safe". - let ptr_hi: *mut Poly<u64> = *hi as *const Poly<u64> as *mut Poly<u64>; - let ptr_lo: *mut Poly<u64> = *lo as *const Poly<u64> as *mut Poly<u64>; - assert!(ptr_hi != ptr_lo, "something went seriously wrong"); - - unsafe{ - self.a_add_b_into_b::<ONCE>(hi, transmute(ptr_lo)); - } - } - None =>{ - unsafe{ - polys.insert(j, transmute(*hi as *const Poly<u64> as *mut Poly<u64>)); - } - }, - } - - polys.remove(&(j+t)); - } - - None =>{}, + 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); + } else { + std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]); + } } - let option_lo: Option<&&mut Poly<u64>> = polys.get(&i); - let option_hi: Option<&&mut Poly<u64>> = polys.get(&(i+t)); + 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); - match option_lo{ - Some(lo) =>{ - - let gal_el: usize = self.galois_element(1<<(i-1), i == 0, log_nth_root); - - if hi_exists{ - self.automorphism::<true>(&tmpa, gal_el, 2<<self.log_n(), &mut tmpb); - }else{ - self.automorphism::<true>(*lo, gal_el, nth_root, &mut tmpa); - } - unsafe{ - self.a_add_b_into_b::<ONCE>(&tmpa, transmute(*lo as *const Poly<u64> as *mut Poly<u64>)); - } + if !polys_hi[j].is_none() { + self.automorphism::<true>(&tmpa, gal_el, 2 << self.log_n(), &mut tmpb); + } else { + self.automorphism::<true>(poly_lo, gal_el, nth_root, &mut tmpa); } - None =>{ - match option_hi{ - Some(hi) =>{ - let gal_el: usize = self.galois_element(1<<(i-1), i == 0, log_nth_root); + 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.automorphism::<true>(*hi, gal_el, nth_root, &mut tmpa); - - unsafe{ - self.a_sub_b_into_a::<1, ONCE>(&tmpa, transmute(*hi as *const Poly<u64> as *mut Poly<u64>)) - } - } - - None =>{} - } - } + self.automorphism::<true>(poly_hi, gal_el, nth_root, &mut tmpa); + self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi) } } } - - *polys.get(&0).unwrap() } } - -// Returns the largest -fn max_pow2_gap(vec: &[usize]) -> usize{ +// Returns the largest gap. +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]); + for i in 1..vec.len() { + let (l, r) = (vec[i - 1], vec[i]); assert!(l > r, "invalid input vec: not sorted"); - gap = min(gap, r-l); - if gap == 1{ + gap = min(gap, r - l); + if gap == 1 { break; } - }; - 1 << gap.trailing_zeros() -}+ } + gap +}
M
math/src/ring/impl_u64/ring.rs
→
math/src/ring/impl_u64/ring.rs
@@ -42,24 +42,24 @@ .for_each(|(i, v)| *v = (&coeffs[i] % &q_big).to_u64().unwrap());
} // Returns GALOISGENERATOR^gen_1 * (-1)^gen_2 mod 2^log_nth_root. - pub fn galois_element(&self, gen_1: usize, gen_2: bool, log_nth_root: usize) -> usize{ + pub fn galois_element(&self, gen_1: usize, gen_2: bool, log_nth_root: usize) -> usize { let mut gal_el: usize = 1; let mut gen_1_pow: usize = GALOISGENERATOR; let mut e: usize = gen_1; - while e > 0{ - if e & 1 == 1{ + while e > 0 { + if e & 1 == 1 { gal_el = gal_el.wrapping_mul(gen_1_pow); } gen_1_pow *= gen_1_pow; - e>>=1; + e >>= 1; } - let nth_root = 1<<log_nth_root; - gal_el &= (nth_root-1); + let nth_root = 1 << log_nth_root; + gal_el &= (nth_root - 1); - if gen_2{ - return nth_root - gal_el + if gen_2 { + return nth_root - gal_el; } gal_el }@@ -223,9 +223,13 @@ .va_neg_into_va::<CHUNK, ARANGE, REDUCE>(&mut a.0);
} #[inline(always)] - pub fn a_prepare_montgomery_into_a<const REDUCE: REDUCEMOD>(&self, a: &mut Poly<Montgomery<u64>>){ + pub fn a_prepare_montgomery_into_a<const REDUCE: REDUCEMOD>( + &self, + a: &mut Poly<Montgomery<u64>>, + ) { debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n()); - self.modulus.va_prepare_montgomery_into_va::<CHUNK, REDUCE>(&mut a.0); + self.modulus + .va_prepare_montgomery_into_va::<CHUNK, REDUCE>(&mut a.0); } #[inline(always)]
M
sampling/src/source.rs
→
sampling/src/source.rs
@@ -36,7 +36,7 @@ min + ((self.next_u64() << 11 >> 11) as f64) / MAXF64 * (max - min)
} } -impl RngCore for Source{ +impl RngCore for Source { #[inline(always)] fn next_u32(&mut self) -> u32 { self.source.next_u32()@@ -53,7 +53,7 @@ self.source.fill_bytes(bytes)
} #[inline(always)] - fn try_fill_bytes(&mut self, bytes: &mut [u8]) -> Result<(), rand_core::Error>{ + fn try_fill_bytes(&mut self, bytes: &mut [u8]) -> Result<(), rand_core::Error> { self.source.try_fill_bytes(bytes) } -}+}