arnaucube/phantom-zone
1
0
Fork 0
mirror of https://github.com/arnaucube/phantom-zone.git synced 2026-01-10 16:11:30 +01:00
Code Issues Projects Releases Wiki Activity

add more ntt benches

Browse Source
This commit is contained in:
Janmajaya Mall
2024-06-10 13:59:05 +05:30
parent e9bdaaea69
commit 0f496a1032
5 changed files with 134 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
benches/ntt.rs
View File

@@ -13,6 +13,15 @@ fn forward_lazy_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
.for_each(|r| nttop.forward_lazy(r.as_mut_slice()));
}
fn backward_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
a.iter_mut().for_each(|r| nttop.backward(r.as_mut_slice()));
}
fn backward_lazy_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
a.iter_mut()
.for_each(|r| nttop.backward_lazy(r.as_mut_slice()));
}
fn benchmark(c: &mut Criterion) {
let mut group = c.benchmark_group("ntt");
// 55
@@ -85,6 +94,17 @@ fn benchmark(c: &mut Criterion) {
}
{
group.bench_function(
BenchmarkId::new("backward", format!("q={prime}/N={ring_size}")),
|b| {
b.iter_batched_ref(
|| a.clone(),
|mut a| black_box(ntt.backward(&mut a)),
criterion::BatchSize::PerIteration,
)
},
);
group.bench_function(
BenchmarkId::new("backward_lazy", format!("q={prime}/N={ring_size}")),
|b| {
@@ -95,6 +115,31 @@ fn benchmark(c: &mut Criterion) {
)
},
);
group.bench_function(
BenchmarkId::new("backward_matrix", format!("q={prime}/N={ring_size}")),
|b| {
b.iter_batched_ref(
|| a_matrix.clone(),
|a_matrix| black_box(backward_matrix(a_matrix, &ntt)),
criterion::BatchSize::PerIteration,
)
},
);
group.bench_function(
BenchmarkId::new(
"backward_lazy_matrix",
format!("q={prime}/N={ring_size}/d={d}"),
),
|b| {
b.iter_batched_ref(
|| a_matrix.clone(),
|a_matrix| black_box(backward_lazy_matrix(a_matrix, &ntt)),
criterion::BatchSize::PerIteration,
)
},
);
}
}
}

120
src/ntt.rs
View File

@@ -4,7 +4,7 @@ use rand_chacha::{rand_core::le, ChaCha8Rng};
use crate::{
backend::{ArithmeticOps, ModInit, ModularOpsU64, Modulus},
utils::{mod_exponent, mod_inverse, shoup_representation_fq},
utils::{mod_exponent, mod_inverse, ShoupMul},
};
pub trait NttInit<M> {
@@ -43,9 +43,7 @@ pub fn forward_butterly_0_to_4q(
x = x - q_twice;
}
// TODO (Jay): Hot path expected. How expensive is it?
let k = ((w_shoup as u128 * y as u128) >> 64) as u64;
let t = w.wrapping_mul(y).wrapping_sub(k.wrapping_mul(q));
let t = ShoupMul::mul(y, w, w_shoup, q);
(x + t, x + q_twice - t)
}
@@ -65,8 +63,7 @@ pub fn forward_butterly_0_to_2q(
x = x - q_twice;
}
let k = ((w_shoup as u128 * y as u128) >> 64) as u64;
let t = w.wrapping_mul(y).wrapping_sub(k.wrapping_mul(q));
let t = ShoupMul::mul(y, w, w_shoup, q);
let ox = x.wrapping_add(t);
let oy = x.wrapping_sub(t);
@@ -84,7 +81,7 @@ pub fn forward_butterly_0_to_2q(
/// and both x' and y' are \in [0, 2q)
///
/// Implements Algorithm 3 of [FASTER ARITHMETIC FOR NUMBER-THEORETIC TRANSFORMS](https://arxiv.org/pdf/1205.2926.pdf)
pub fn inverse_butterfly(
pub fn inverse_butterfly_0_to_2q(
x: u64,
y: u64,
w_inv: u64,
@@ -101,8 +98,7 @@ pub fn inverse_butterfly(
}
let t = x + q_twice - y;
let k = ((w_inv_shoup as u128 * t as u128) >> 64) as u64;
let y = w_inv.wrapping_mul(t).wrapping_sub(k.wrapping_mul(q));
let y = ShoupMul::mul(t, w_inv, w_inv_shoup, q);
(x_dash, y)
}
@@ -202,68 +198,82 @@ pub fn ntt_inv_lazy(
let mut t = 1;
while m > 0 {
let w_inv = &psi_inv[m..];
let w_inv_shoup = &psi_inv_shoup[m..];
if m == 1 {
let (left, right) = a.split_at_mut(t);
for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
let (ox, oy) = inverse_butterfly(*x, *y, w_inv[0], w_inv_shoup[0], q, q_twice);
*x = (n_inv.wrapping_mul(ox)).wrapping_sub(
q.wrapping_mul(((ox as u128 * n_inv_shoup as u128) >> 64) as u64),
);
*y = (n_inv.wrapping_mul(oy)).wrapping_sub(
q.wrapping_mul(((oy as u128 * n_inv_shoup as u128) >> 64) as u64),
);
let (ox, oy) =
inverse_butterfly_0_to_2q(*x, *y, psi_inv[1], psi_inv_shoup[1], q, q_twice);
*x = ShoupMul::mul(ox, n_inv, n_inv_shoup, q);
*y = ShoupMul::mul(oy, n_inv, n_inv_shoup, q);
}
} else {
let w_inv = &psi_inv[m..];
let w_inv_shoup = &psi_inv_shoup[m..];
for i in 0..m {
let a = &mut a[2 * i * t..2 * (i + 1) * t];
let (left, right) = a.split_at_mut(t);
for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
let (ox, oy) = inverse_butterfly(*x, *y, w_inv[i], w_inv_shoup[i], q, q_twice);
let (ox, oy) =
inverse_butterfly_0_to_2q(*x, *y, w_inv[i], w_inv_shoup[i], q, q_twice);
*x = ox;
*y = oy;
}
}
}
// for i in 0..m {
// let a = &mut a[2 * i * t..2 * (i + 1) * t];
// let (left, right) = a.split_at_mut(t);
// for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
// let (ox, oy) = inverse_butterfly(*x, *y, w_inv[i], w_inv_shoup[i], q,
// q_twice); *x = ox;
// *y = oy;
// }
// }
// for i in 0..h {
// let j_2 = j_1 + t;
// unsafe {
// let w_inv = psi_inv.get_unchecked(h + i);
// let w_inv_shoup = psi_inv_shoup.get_unchecked(h + i);
// for j in j_1..j_2 {
// let x = a.get_unchecked_mut(j) as *mut u64;
// let y = a.get_unchecked_mut(j + t) as *mut u64;
// inverse_butterfly(x, y, w_inv, w_inv_shoup, &q, &q_twice);
// }
// }
// j_1 = j_1 + 2 * t;
// }
t *= 2;
m >>= 1;
}
}
// a.iter_mut().for_each(|a0| {
// *a0 = (n_inv.wrapping_mul(*a0))
// .wrapping_sub(((*a0 as u128 * n_inv_shoup as u128) >> 64) as u64)
// });
/// Same as `ntt_inv_lazy` with output in range [0, q)
pub fn ntt_inv(
a: &mut [u64],
psi_inv: &[u64],
psi_inv_shoup: &[u64],
n_inv: u64,
n_inv_shoup: u64,
q: u64,
q_twice: u64,
) {
assert!(a.len() == psi_inv.len());
let mut m = a.len() >> 1;
let mut t = 1;
while m > 0 {
if m == 1 {
let (left, right) = a.split_at_mut(t);
for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
let (ox, oy) =
inverse_butterfly_0_to_2q(*x, *y, psi_inv[1], psi_inv_shoup[1], q, q_twice);
let ox = ShoupMul::mul(ox, n_inv, n_inv_shoup, q);
let oy = ShoupMul::mul(oy, n_inv, n_inv_shoup, q);
*x = ox.min(ox.wrapping_sub(q));
*y = oy.min(oy.wrapping_sub(q));
}
} else {
let w_inv = &psi_inv[m..];
let w_inv_shoup = &psi_inv_shoup[m..];
for i in 0..m {
let a = &mut a[2 * i * t..2 * (i + 1) * t];
let (left, right) = a.split_at_mut(t);
for (x, y) in izip!(left.iter_mut(), right.iter_mut()) {
let (ox, oy) =
inverse_butterfly_0_to_2q(*x, *y, w_inv[i], w_inv_shoup[i], q, q_twice);
*x = ox;
*y = oy;
}
}
}
t *= 2;
m >>= 1;
}
}
/// Find n^{th} root of unity in field F_q, if one exists
@@ -350,11 +360,11 @@ impl NttBackendU64 {
// shoup representation
let psi_powers_bo_shoup = psi_powers_bo
.iter()
.map(|v| shoup_representation_fq(*v, q))
.map(|v| ShoupMul::representation(*v, q))
.collect_vec();
let psi_inv_powers_bo_shoup = psi_inv_powers_bo
.iter()
.map(|v| shoup_representation_fq(*v, q))
.map(|v| ShoupMul::representation(*v, q))
.collect_vec();
// n^{-1} \mod{q}
@@ -365,7 +375,7 @@ impl NttBackendU64 {
q_twice: 2 * q,
n: n as u64,
n_inv,
n_inv_shoup: shoup_representation_fq(n_inv, q),
n_inv_shoup: ShoupMul::representation(n_inv, q),
psi_powers_bo: psi_powers_bo.into_boxed_slice(),
psi_inv_powers_bo: psi_inv_powers_bo.into_boxed_slice(),
psi_powers_bo_shoup: psi_powers_bo_shoup.into_boxed_slice(),
@@ -429,7 +439,7 @@ impl Ntt for NttBackendU64 {
}
fn backward(&self, v: &mut [Self::Element]) {
ntt_inv_lazy(
ntt_inv(
v,
&self.psi_inv_powers_bo,
&self.psi_inv_powers_bo_shoup,
@@ -438,10 +448,10 @@ impl Ntt for NttBackendU64 {
self.q,
self.q_twice,
);
self.reduce_from_lazy(v);
}
}
#[cfg(test)]
mod tests {
use itertools::Itertools;
use rand::{thread_rng, Rng};

2
src/pbs.rs
View File

@@ -342,7 +342,7 @@ fn blind_rotation<
mod_op,
);
});
println!("Auto count: {count}");
// println!("Auto count: {count}");
}
fn mod_switch_odd(v: f64, from_q: f64, to_q: f64) -> usize {

6
src/shortint/mod.rs
View File

@@ -306,7 +306,7 @@ mod tests {
bool::{
aggregate_public_key_shares, aggregate_server_key_shares, gen_client_key, gen_keys,
gen_mp_keys_phase1, gen_mp_keys_phase2,
parameters::{MP_BOOL_PARAMS, SP_BOOL_PARAMS},
parameters::{MP_BOOL_PARAMS, SMALL_MP_BOOL_PARAMS, SP_BOOL_PARAMS},
set_mp_seed, set_parameter_set,
},
shortint::types::FheUint8,
@@ -463,7 +463,7 @@ mod tests {
#[test]
fn fheuint8_test_multi_party() {
set_parameter_set(&MP_BOOL_PARAMS);
set_parameter_set(&SMALL_MP_BOOL_PARAMS);
set_mp_seed([0; 32]);
let parties = 8;
@@ -497,10 +497,12 @@ mod tests {
let ct_b = public_key.encrypt(&b);
let ct_c = public_key.encrypt(&c);
let now = std::time::Instant::now();
// server computes
// a*b + c
let mut ct_ab = &ct_a * &ct_b;
ct_ab += &ct_c;
println!("Circuit time: {:?}", now.elapsed());
// decrypt ab and check
// generate decryption shares

23
src/utils.rs
View File

@@ -25,6 +25,25 @@ pub trait Global {
fn global() -> &'static Self;
}
pub trait ShoupMul {
fn representation(value: Self, q: Self) -> Self;
fn mul(a: Self, b: Self, b_shoup: Self, q: Self) -> Self;
}
impl ShoupMul for u64 {
#[inline]
fn representation(value: Self, q: Self) -> Self {
((value as u128 * (1u128 << 64)) / q as u128) as u64
}
#[inline]
/// Returns a * b % q
fn mul(a: Self, b: Self, b_shoup: Self, q: Self) -> Self {
(b.wrapping_mul(a))
.wrapping_sub(q.wrapping_mul(((b_shoup as u128 * a as u128) >> 64) as u64))
}
}
pub fn fill_random_ternary_secret_with_hamming_weight<
T: Signed,
R: RandomFill<[u8]> + RandomElementInModulus<usize, usize>,
@@ -121,10 +140,6 @@ pub fn mod_inverse(a: u64, q: u64) -> u64 {
mod_exponent(a, q - 2, q)
}
pub fn shoup_representation_fq(v: u64, q: u64) -> u64 {
((v as u128 * (1u128 << 64)) / q as u128) as u64
}
pub fn negacyclic_mul<T: PrimInt, F: Fn(&T, &T) -> T>(
a: &[T],
b: &[T],