more additions to automorphisms

This commit is contained in:
Jean-Philippe Bossuat
2025-01-10 21:52:40 +01:00
parent a6c1aa1c09
commit 74bfb52ada
8 changed files with 177 additions and 114 deletions

View File

@@ -44,7 +44,7 @@ fn main() {
p0.0[i] = i as u64
}
r.automorphism::<false>(&p0, 2 * r.n - 1, nth_root, &mut p1);
r.a_apply_automorphism_into_b::<false>(&p0, 2 * r.n - 1, nth_root, &mut p1);
println!("{:?}", p1);
}

View File

@@ -1,6 +1,7 @@
use crate::modulus::WordOps;
use crate::modulus::{REDUCEMOD, WordOps};
use crate::poly::Poly;
use crate::ring::Ring;
use crate::modulus::{ONCE, ScalarOperations};
/// Returns a lookup table for the automorphism X^{i} -> X^{i * k mod nth_root}.
/// Method will panic if n or nth_root are not power-of-two.
@@ -44,7 +45,41 @@ pub fn automorphism_index<const NTT: bool>(n: usize, nth_root: usize, gal_el: us
}
impl Ring<u64> {
pub fn automorphism<const NTT: bool>(
// b <- auto(a)
pub fn a_apply_automorphism_into_b<const NTT: bool>(
&self,
a: &Poly<u64>,
gal_el: usize,
nth_root: usize,
b: &mut Poly<u64>,
){
self.apply_automorphism_core::<0, ONCE, NTT>(a, gal_el, nth_root, b)
}
// b <- REDUCEMOD(b + auto(a))
pub fn a_apply_automorphism_add_b_into_b<const REDUCE:REDUCEMOD, const NTT: bool>(
&self,
a: &Poly<u64>,
gal_el: usize,
nth_root: usize,
b: &mut Poly<u64>,
){
self.apply_automorphism_core::<1, REDUCE, NTT>(a, gal_el, nth_root, b)
}
// b <- REDUCEMOD(b - auto(a))
pub fn a_apply_automorphism_sub_b_into_b<const REDUCE:REDUCEMOD, const NTT:bool>(
&self,
a: &Poly<u64>,
gal_el: usize,
nth_root: usize,
b: &mut Poly<u64>,
){
self.apply_automorphism_core::<2, REDUCE, NTT>(a, gal_el, nth_root, b)
}
fn apply_automorphism_core<const MOD:u8, const REDUCE:REDUCEMOD, const NTT: bool>(
&self,
a: &Poly<u64>,
gal_el: usize,
@@ -81,7 +116,12 @@ impl Ring<u64> {
let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
let gal_el_i: usize = (((gal_el * i_rev) & mask) - 1) >> 1;
let idx: usize = gal_el_i.reverse_bits_msb(log_nth_root_half);
b_vec[idx] = *ai;
match MOD{
0 =>{b_vec[idx] = *ai}
1=>{self.modulus.sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx])}
2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx])}
_=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
}
});
} else {
let n: usize = a.n();
@@ -92,12 +132,49 @@ impl Ring<u64> {
let gal_el_i: usize = i * gal_el;
let sign: u64 = ((gal_el_i >> log_n) & 1) as u64;
let i_out: usize = gal_el_i & mask;
b_vec[i_out] = ai * (sign ^ 1) | (q - ai) * sign
let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
match MOD{
0 =>{b_vec[i_out] = v}
1=>{self.modulus.sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[i_out])}
2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[i_out])}
_=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
}
});
}
}
pub fn automorphism_from_index<const NTT: bool>(
// b <- auto(a)
pub fn a_apply_automorphism_from_index_into_b<const NTT: bool>(
&self,
a: &Poly<u64>,
idx: &[usize],
b: &mut Poly<u64>,
){
self.automorphism_from_index_core::<0, ONCE, NTT>(a, idx, b)
}
// b <- REDUCEMOD(b + auto(a))
pub fn a_apply_automorphism_from_index_add_b_into_b<const REDUCE:REDUCEMOD, const NTT: bool>(
&self,
a: &Poly<u64>,
idx: &[usize],
b: &mut Poly<u64>,
){
self.automorphism_from_index_core::<1, REDUCE, NTT>(a, idx, b)
}
// b <- REDUCEMOD(b - auto(a))
pub fn a_apply_automorphism_from_index_sub_b_into_b<const REDUCE:REDUCEMOD, const NTT:bool>(
&self,
a: &Poly<u64>,
idx: &[usize],
b: &mut Poly<u64>,
){
self.automorphism_from_index_core::<2, REDUCE, NTT>(a, idx, b)
}
// b <- auto(a) if OVERWRITE else b <- REDUCEMOD(b + auto(a))
fn automorphism_from_index_core<const MOD:u8, const REDUCE:REDUCEMOD, const NTT: bool>(
&self,
a: &Poly<u64>,
idx: &[usize],
@@ -115,7 +192,12 @@ impl Ring<u64> {
if NTT {
a_vec.iter().enumerate().for_each(|(i, ai)| {
b_vec[idx[i]] = *ai;
match MOD{
0 =>{ b_vec[idx[i]] = *ai}
1 =>{self.modulus.sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx[i]])}
2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx[i]])}
_=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
}
});
} else {
let n: usize = a.n();
@@ -123,7 +205,13 @@ impl Ring<u64> {
let q: u64 = self.modulus.q();
a_vec.iter().enumerate().for_each(|(i, ai)| {
let sign: u64 = (idx[i] >> usize::BITS - 1) as u64;
b_vec[idx[i] & mask] = ai * (sign ^ 1) | (q - ai) * sign;
let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
match MOD{
0 =>{b_vec[idx[i] & mask] = v}
1 =>{self.modulus.sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[idx[i] & mask])}
2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[idx[i] & mask])}
_=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
}
});
}
}

View File

@@ -1,6 +1,8 @@
pub mod automorphism;
pub mod packing;
pub mod trace;
pub mod rescaling_rns;
pub mod ring;
pub mod ring_rns;
pub mod sampling;
pub mod utils;

View File

@@ -3,41 +3,8 @@ use crate::modulus::ONCE;
use crate::poly::Poly;
use crate::ring::Ring;
use std::cmp::min;
use std::collections::HashSet;
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>> {
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;
if INV {
idx = self.n() - idx;
}
x_pow.push(self.new_poly());
if i == 0 {
x_pow[i].0[idx] = self.modulus.montgomery.one();
self.ntt_inplace::<false>(&mut x_pow[i]);
} 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]);
}
});
if INV {
self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
}
if !NTT {
x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
}
x_pow
}
pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
&self,
@@ -61,8 +28,6 @@ impl Ring<u64> {
let gap: usize = max_gap(&indices);
let set: HashSet<_> = indices.into_iter().collect();
if !ZEROGARBAGE {
if gap > 0 {
log_end -= gap.trailing_zeros() as usize;
@@ -79,7 +44,7 @@ impl Ring<u64> {
.barrett
.prepare(self.modulus.inv(1 << (log_end - log_start)));
set.iter().for_each(|i| {
indices.iter().for_each(|i| {
if let Some(poly) = polys[*i].as_mut() {
if !NTT {
self.ntt_inplace::<true>(poly);
@@ -111,15 +76,14 @@ impl Ring<u64> {
let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
if !polys_hi[j].is_none() {
self.automorphism::<true>(&tmpa, gal_el, 2 << self.log_n(), &mut tmpb);
self.a_add_b_into_b::<ONCE>(&tmpb, poly_lo);
self.a_apply_automorphism_add_b_into_b::<ONCE, true>(&tmpa, gal_el, 2 << self.log_n(), poly_lo);
} else {
self.automorphism::<true>(poly_lo, gal_el, nth_root, &mut tmpa);
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.automorphism::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
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]);
}
@@ -136,12 +100,14 @@ impl Ring<u64> {
}
}
// Returns the largest gap.
// 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");
assert!(r > l, "invalid input vec: not sorted or collision between indices");
gap = min(gap, r - l);
if gap == 1 {
break;

View File

@@ -0,0 +1,32 @@
use crate::ring::Ring;
use crate::poly::Poly;
use crate::modulus::barrett::Barrett;
use crate::modulus::ONCE;
impl Ring<u64>{
pub fn trace_inplace<const NTT:bool>(&self, log_start: usize, log_end: usize, a: &mut Poly<u64>){
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);
let log_steps = log_end - log_start;
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);
if !NTT{
self.ntt_inplace::<false>(a);
}
let mut tmp: Poly<u64> = self.new_poly();
(log_start..log_end).for_each(|i|{
});
}
}
}

View File

@@ -0,0 +1,38 @@
use crate::ring::Ring;
use crate::poly::Poly;
use crate::modulus::ONCE;
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>> {
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;
if INV {
idx = self.n() - idx;
}
x_pow.push(self.new_poly());
if i == 0 {
x_pow[i].0[idx] = self.modulus.montgomery.one();
self.ntt_inplace::<false>(&mut x_pow[i]);
} 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]);
}
});
if INV {
self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
}
if !NTT {
x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
}
x_pow
}
}

View File

@@ -39,7 +39,7 @@ fn test_automorphism_u64<const NTT: bool>(ring: &Ring<u64>, nth_root: usize) {
ring.ntt_inplace::<false>(&mut p0);
}
ring.automorphism::<NTT>(&p0, 2 * n - 1, nth_root, &mut p1);
ring.a_apply_automorphism_into_b::<NTT>(&p0, 2 * n - 1, nth_root, &mut p1);
if NTT {
ring.intt_inplace::<false>(&mut p1);