arnaucube/fhe-study
1
1
Fork 0
mirror of https://github.com/arnaucube/fhe-study.git synced 2026-01-24 04:33:52 +01:00
Code Issues Packages Projects Releases Wiki Activity

ntt: get rid of Zq and use u64 instead (>2x speed improvement)

Browse Source
This commit is contained in:
arnaucube
2025-08-18 14:14:19 +00:00
parent fb1fb6b4e9
commit 4dca2c6ff5
3 changed files with 86 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
arith/src/ntt.rs
View File

@@ -6,7 +6,6 @@
//! generics; but once using real-world parameters, the stack could not handle
//! it, so moved to use Vec instead of fixed-sized arrays, and adapted the NTT
//! implementation to that too.
use crate::{ring::RingParam, ring_nq::Rq, zq::Zq};
use std::collections::HashMap;
@@ -15,22 +14,19 @@ pub struct NTT {}
use std::sync::{Mutex, OnceLock};
static CACHE: OnceLock<Mutex<HashMap<(u64, usize), (Vec<Zq>, Vec<Zq>, Zq)>>> = OnceLock::new();
static CACHE: OnceLock<Mutex<HashMap<(u64, usize), (Vec<u64>, Vec<u64>, u64)>>> = OnceLock::new();
fn roots(q: u64, n: usize) -> (Vec<Zq>, Vec<Zq>, Zq) {
fn roots(q: u64, n: usize) -> (Vec<u64>, Vec<u64>, u64) {
let cache_lock = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
let mut cache = cache_lock.lock().unwrap();
if let Some(value) = cache.get(&(q, n)) {
return value.clone();
}
let n_inv: Zq = Zq {
q,
v: const_inv_mod(q, n as u64),
};
let n_inv: u64 = const_inv_mod(q, n as u64);
let root_of_unity: u64 = primitive_root_of_unity(q, 2 * n);
let roots_of_unity: Vec<Zq> = roots_of_unity(q, n, root_of_unity);
let roots_of_unity_inv: Vec<Zq> = roots_of_unity_inv(q, n, roots_of_unity.clone());
let roots_of_unity: Vec<u64> = roots_of_unity(q, n, root_of_unity);
let roots_of_unity_inv: Vec<u64> = roots_of_unity_inv(q, n, roots_of_unity.clone());
let value = (roots_of_unity, roots_of_unity_inv, n_inv);
cache.insert((q, n), value.clone());
@@ -41,56 +37,70 @@ impl NTT {
/// implements the Cooley-Tukey (CT) algorithm. Details at
/// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.1 of
/// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
pub fn ntt(a: &Rq) -> Rq {
let (q, n) = (a.param.q, a.param.n);
pub fn ntt(q: u64, n: usize, a: &Vec<u64>) -> Vec<u64> {
debug_assert_eq!(n, a.len());
let (roots_of_unity, _, _) = roots(q, n);
let mut t = n / 2;
let mut m = 1;
let mut r: Vec<Zq> = a.coeffs.clone();
let mut r: Vec<u64> = a.clone();
while m < n {
let mut k = 0;
for i in 0..m {
let S: Zq = roots_of_unity[m + i];
let S: u64 = roots_of_unity[m + i];
for j in k..k + t {
let U: Zq = r[j];
let V: Zq = r[j + t] * S;
let U: u64 = r[j];
let V: u64 = (r[j + t] * S) % q;
// compute r[j] = (U + V) % q:
r[j] = U + V;
if r[j] >= q {
r[j] -= q;
}
// compute r[j + t] = (U - V) % q:
if U >= V {
r[j + t] = U - V;
} else {
r[j + t] = (q + U) - V;
}
}
k = k + 2 * t;
}
t /= 2;
m *= 2;
}
// TODO think if maybe not return a Rq type, or if returned Rq, maybe
// fill the `evals` field, which is what we're actually returning here
Rq {
param: RingParam { q, n },
coeffs: r,
evals: None,
}
r
}
/// implements the Cooley-Tukey (CT) algorithm. Details at
/// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.2 of
/// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
pub fn intt(a: &Rq) -> Rq {
let (q, n) = (a.param.q, a.param.n);
pub fn intt(q: u64, n: usize, a: &Vec<u64>) -> Vec<u64> {
debug_assert_eq!(n, a.len());
let (_, roots_of_unity_inv, n_inv) = roots(q, n);
let mut t = 1;
let mut m = n / 2;
let mut r: Vec<Zq> = a.coeffs.clone();
let mut r: Vec<u64> = a.clone();
while m > 0 {
let mut k = 0;
for i in 0..m {
let S: Zq = roots_of_unity_inv[m + i];
let S: u64 = roots_of_unity_inv[m + i];
for j in k..k + t {
let U: Zq = r[j];
let V: Zq = r[j + t];
let U: u64 = r[j];
let V: u64 = r[j + t];
// compute r[j] = (U + V) % q:
r[j] = U + V;
r[j + t] = (U - V) * S;
if r[j] >= q {
r[j] -= q;
}
// compute r[j + t] = ((U - V) * S) % q;
if U >= V {
r[j + t] = ((U - V) * S) % q;
} else {
r[j + t] = ((q + U - V) * S) % q;
}
}
k += 2 * t;
}
@@ -98,15 +108,9 @@ impl NTT {
m /= 2;
}
for i in 0..n {
r[i] = r[i] * n_inv;
}
Rq {
param: RingParam { q, n },
coeffs: r,
// TODO maybe at `evals` place the inputed `a` which is the evals
// format
evals: None,
r[i] = (r[i] * n_inv) % q;
}
r
}
}
@@ -130,31 +134,25 @@ const fn primitive_root_of_unity(q: u64, n: usize) -> u64 {
panic!("No primitive root of unity");
}
fn roots_of_unity(q: u64, n: usize, w: u64) -> Vec<Zq> {
let mut r: Vec<Zq> = vec![Zq { q, v: 0 }; n];
fn roots_of_unity(q: u64, n: usize, w: u64) -> Vec<u64> {
let mut r: Vec<u64> = vec![0; n];
let mut i = 0;
let log_n = n.ilog2();
while i < n {
// (return the roots in bit-reverset order)
let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
r[i] = Zq {
q,
v: const_exp_mod(q, w, j as u64),
};
r[i] = const_exp_mod(q, w, j as u64);
i += 1;
}
r
}
fn roots_of_unity_inv(q: u64, n: usize, v: Vec<Zq>) -> Vec<Zq> {
fn roots_of_unity_inv(q: u64, n: usize, v: Vec<u64>) -> Vec<u64> {
// assumes that the inputted roots are already in bit-reverset order
let mut r: Vec<Zq> = vec![Zq { q, v: 0 }; n];
let mut r: Vec<u64> = vec![0; n];
let mut i = 0;
while i < n {
r[i] = Zq {
q,
v: const_inv_mod(q, v[i].v),
};
r[i] = const_inv_mod(q, v[i]);
i += 1;
}
r
@@ -187,7 +185,7 @@ const fn const_inv_mod(q: u64, x: u64) -> u64 {
#[cfg(test)]
mod tests {
use super::*;
use crate::Ring;
use rand_distr::Distribution;
use anyhow::Result;
@@ -195,14 +193,12 @@ mod tests {
fn test_ntt() -> Result<()> {
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 4;
let param = RingParam { q, n };
let a: Vec<u64> = vec![1u64, 2, 3, 4];
let a: Rq = Rq::from_vec_u64(&param, a);
let a_ntt = NTT::ntt(&a);
let a_ntt = NTT::ntt(q, n, &a);
let a_intt = NTT::intt(&a_ntt);
let a_intt = NTT::intt(q, n, &a_ntt);
dbg!(&a);
dbg!(&a_ntt);
@@ -218,16 +214,17 @@ mod tests {
fn test_ntt_loop() -> Result<()> {
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 512;
let param = RingParam { q, n };
use rand::distributions::Uniform;
let mut rng = rand::thread_rng();
let dist = Uniform::new(0_f64, q as f64);
let dist = Uniform::new(0_u64, q as u64);
for _ in 0..1000 {
let a: Rq = Rq::rand(&mut rng, dist, &param);
let a_ntt = NTT::ntt(&a);
let a_intt = NTT::intt(&a_ntt);
let a: Vec<u64> = std::iter::repeat_with(|| dist.sample(&mut rng))
.take(n)
.collect();
let a_ntt = NTT::ntt(q, n, &a);
let a_intt = NTT::intt(q, n, &a_ntt);
assert_eq!(a, a_intt);
}
Ok(())

36
arith/src/ring_nq.rs
View File

@@ -113,6 +113,24 @@ impl Ring for Rq {
}
}
impl Rq {
fn coeffs_u64(&self) -> Vec<u64> {
self.coeffs.iter().map(|c_i| c_i.v).collect()
}
fn ntt(&self) -> Vec<Zq> {
NTT::ntt(self.param.q, self.param.n, &self.coeffs_u64())
.iter()
.map(|c_i| Zq::from_u64(self.param.q, *c_i))
.collect()
}
fn intt(&self) -> Vec<Zq> {
NTT::intt(self.param.q, self.param.n, &self.coeffs_u64())
.iter()
.map(|c_i| Zq::from_u64(self.param.q, *c_i))
.collect()
}
}
impl From<(u64, crate::ring_n::R)> for Rq {
fn from(qr: (u64, crate::ring_n::R)) -> Self {
let (q, r) = qr;
@@ -145,7 +163,7 @@ impl Rq {
self.coeffs.clone()
}
pub fn compute_evals(&mut self) {
self.evals = Some(NTT::ntt(self).coeffs);
self.evals = Some(self.ntt());
// TODO improve, ntt returns Rq but here just needs Vec<Zq>
}
pub fn to_r(self) -> crate::R {
@@ -566,10 +584,10 @@ fn mul_mut(lhs: &mut Rq, rhs: &mut Rq) -> Rq {
// reuse evaluations if already computed
if !lhs.evals.is_some() {
lhs.evals = Some(NTT::ntt(lhs).coeffs);
lhs.evals = Some(lhs.ntt());
};
if !rhs.evals.is_some() {
rhs.evals = Some(NTT::ntt(rhs).coeffs);
rhs.evals = Some(rhs.ntt());
};
let lhs_evals = lhs.evals.clone().unwrap();
let rhs_evals = rhs.evals.clone().unwrap();
@@ -578,8 +596,8 @@ fn mul_mut(lhs: &mut Rq, rhs: &mut Rq) -> Rq {
&lhs.param,
zip_eq(lhs_evals, rhs_evals).map(|(l, r)| l * r).collect(),
);
let c = NTT::intt(&c_ntt);
Rq::new(&lhs.param, c.coeffs, Some(c_ntt.coeffs))
let c: Vec<Zq> = c_ntt.intt();
Rq::new(&lhs.param, c, Some(c_ntt.coeffs))
}
// note: this assumes that Q is prime
// TODO impl karatsuba for non-prime Q. Alternatively check NTT with RNS trick.
@@ -590,20 +608,20 @@ fn mul(lhs: &Rq, rhs: &Rq) -> Rq {
let lhs_evals: Vec<Zq> = if lhs.evals.is_some() {
lhs.evals.clone().unwrap()
} else {
NTT::ntt(lhs).coeffs
lhs.ntt()
};
let rhs_evals: Vec<Zq> = if rhs.evals.is_some() {
rhs.evals.clone().unwrap()
} else {
NTT::ntt(rhs).coeffs
rhs.ntt()
};
let c_ntt: Rq = Rq::from_vec(
&lhs.param,
zip_eq(lhs_evals, rhs_evals).map(|(l, r)| l * r).collect(),
);
let c = NTT::intt(&c_ntt);
Rq::new(&lhs.param, c.coeffs, Some(c_ntt.coeffs))
let c = c_ntt.intt();
Rq::new(&lhs.param, c, Some(c_ntt.coeffs))
}
impl fmt::Display for Rq {

2
arith/src/ring_torus.rs
View File

@@ -252,6 +252,7 @@ impl Mul<Tn> for Tn {
type Output = Self;
fn mul(self, rhs: Self) -> Self {
// TODO NTT/FFT
naive_poly_mul(&self, &rhs)
}
}
@@ -259,6 +260,7 @@ impl Mul<&Tn> for &Tn {
type Output = Tn;
fn mul(self, rhs: &Tn) -> Self::Output {
// TODO NTT/FFT
naive_poly_mul(self, rhs)
}
}