arnaucube/poulpy
1
0
Fork 0
mirror of https://github.com/arnaucube/poulpy.git synced 2026-02-10 05:06:44 +01:00
Code Issues Packages Projects Releases Wiki Activity

[pack]: replaced HashMap by [Option<Poly<u64>]

Browse Source
This commit is contained in:
Jean-Philippe Bossuat
2025-01-10 11:48:51 +01:00
parent 0886e439e0
commit 332e07dbbf
8 changed files with 232 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
math/Cargo.toml
View File

@@ -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]]

23
math/src/modulus/impl_u64/operations.rs
View File

@@ -100,7 +100,11 @@ impl ScalarOperations<u64> for Prime<u64> {
}
#[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 @@ impl VectorOperations<u64> for Prime<u64> {
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 @@ impl VectorOperations<u64> for Prime<u64> {
&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,

8
math/src/ring.rs
View File

@@ -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 struct Ring<O: Unsigned> {
}
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 @@ impl<O: Unsigned> Ring<O> {
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();
}

2
math/src/ring/impl_u64/mod.rs
View File

@@ -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;

199
math/src/ring/impl_u64/packing.rs
View File

@@ -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 log_gap: usize = gap.log2();
let set: HashSet<_> = indices.into_iter().collect();
if !ZEROGARBAGE{
let max_pow2_gap_divisor: usize = 1 << gap.trailing_zeros();
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)));
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);
}
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
}

24
math/src/ring/impl_u64/ring.rs
View File

@@ -42,24 +42,24 @@ impl Ring<u64> {
}
// 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 @@ impl Ring<u64> {
}
#[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)]