added spqlios as submodule

This commit is contained in:
Jean-Philippe Bossuat
2025-01-27 14:10:59 +01:00
parent 250d1a4942
commit c30f598776
244 changed files with 51 additions and 29899 deletions

19
base2k/Cargo.toml Normal file
View File

@@ -0,0 +1,19 @@
[package]
name = "spqlios"
version = "0.1.0"
edition = "2021"
[dependencies]
rand = "0.8.5"
rand_core = "0.6.4"
itertools = "0.14.0"
criterion = "0.5.1"
rand_distr = "0.4.3"
sampling = { path = "../sampling" }
[build-dependencies]
bindgen = "0.71.1"
[[bench]]
name = "fft"
harness = false

71
base2k/benches/fft.rs Normal file
View File

@@ -0,0 +1,71 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use spqlios::bindings::{
new_reim_fft_precomp, new_reim_ifft_precomp, reim_fft, reim_fft_precomp,
reim_fft_precomp_get_buffer, reim_from_znx64_simple, reim_ifft, reim_ifft_precomp,
reim_ifft_precomp_get_buffer,
};
use std::ffi::c_void;
fn fft(c: &mut Criterion) {
fn forward<'a>(
m: u32,
log_bound: u32,
reim_fft_precomp: *mut reim_fft_precomp,
a: &'a [i64],
) -> Box<dyn FnMut() + 'a> {
unsafe {
let buf_a: *mut f64 = reim_fft_precomp_get_buffer(reim_fft_precomp, 0);
reim_from_znx64_simple(m as u32, log_bound as u32, buf_a as *mut c_void, a.as_ptr());
Box::new(move || reim_fft(reim_fft_precomp, buf_a))
}
}
fn backward<'a>(
m: u32,
log_bound: u32,
reim_ifft_precomp: *mut reim_ifft_precomp,
a: &'a [i64],
) -> Box<dyn FnMut() + 'a> {
Box::new(move || unsafe {
let buf_a: *mut f64 = reim_ifft_precomp_get_buffer(reim_ifft_precomp, 0);
reim_from_znx64_simple(m as u32, log_bound as u32, buf_a as *mut c_void, a.as_ptr());
reim_ifft(reim_ifft_precomp, buf_a);
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
c.benchmark_group("fft");
for log_n in 10..17 {
let n: usize = 1 << log_n;
let m: usize = n >> 1;
let log_bound: u32 = 19;
let mut a: Vec<i64> = vec![i64::default(); n];
a.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
unsafe {
let reim_fft_precomp: *mut reim_fft_precomp = new_reim_fft_precomp(m as u32, 1);
let reim_ifft_precomp: *mut reim_ifft_precomp = new_reim_ifft_precomp(m as u32, 1);
let runners: [(String, Box<dyn FnMut()>); 2] = [
(format!("forward"), {
forward(m as u32, log_bound, reim_fft_precomp, &a)
}),
(format!("backward"), {
backward(m as u32, log_bound, reim_ifft_precomp, &a)
}),
];
for (name, mut runner) in runners {
let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
b.iter(&mut runner)
});
}
}
}
}
criterion_group!(benches, fft,);
criterion_main!(benches);

52
base2k/build.rs Normal file
View File

@@ -0,0 +1,52 @@
use bindgen;
use std::env;
use std::fs;
use std::path::absolute;
use std::path::PathBuf;
use std::time::SystemTime;
fn main() {
// Path to the C header file
let header_paths = [
"lib/spqlios/coeffs/coeffs_arithmetic.h",
"lib/spqlios/arithmetic/vec_znx_arithmetic.h",
];
let out_path: PathBuf = PathBuf::from(env::var("OUT_DIR").unwrap());
let bindings_file = out_path.join("bindings.rs");
let regenerate: bool = header_paths.iter().any(|header| {
let header_metadata: SystemTime = fs::metadata(header)
.and_then(|m| m.modified())
.unwrap_or(SystemTime::UNIX_EPOCH);
let bindings_metadata: SystemTime = fs::metadata(&bindings_file)
.and_then(|m| m.modified())
.unwrap_or(SystemTime::UNIX_EPOCH);
header_metadata > bindings_metadata
});
if regenerate {
// Generate the Rust bindings
let mut builder: bindgen::Builder = bindgen::Builder::default();
for header in header_paths {
builder = builder.header(header);
}
let bindings = builder
.generate_comments(false) // Optional: includes comments in bindings
.generate_inline_functions(true) // Optional: includes inline functions
.generate()
.expect("Unable to generate bindings");
// Write the bindings to the OUT_DIR
bindings
.write_to_file(&bindings_file)
.expect("Couldn't write bindings!");
}
println!(
"cargo:rustc-link-search=native={}",
absolute("./lib/build/spqlios").unwrap().to_str().unwrap()
);
println!("cargo:rustc-link-lib=static=spqlios"); //"cargo:rustc-link-lib=dylib=spqlios"
}

57
base2k/examples/fft.rs Normal file
View File

@@ -0,0 +1,57 @@
use std::ffi::c_void;
use std::time::Instant;
use spqlios::bindings::*;
fn main() {
let log_bound: usize = 19;
let n: usize = 2048;
let m: usize = n >> 1;
let mut a: Vec<i64> = vec![i64::default(); n];
let mut b: Vec<i64> = vec![i64::default(); n];
let mut c: Vec<i64> = vec![i64::default(); n];
a.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
b[1] = 1;
println!("{:?}", b);
unsafe {
let reim_fft_precomp = new_reim_fft_precomp(m as u32, 2);
let reim_ifft_precomp = new_reim_ifft_precomp(m as u32, 1);
let buf_a = reim_fft_precomp_get_buffer(reim_fft_precomp, 0);
let buf_b = reim_fft_precomp_get_buffer(reim_fft_precomp, 1);
let buf_c = reim_ifft_precomp_get_buffer(reim_ifft_precomp, 0);
let now = Instant::now();
(0..1024).for_each(|_| {
reim_from_znx64_simple(m as u32, log_bound as u32, buf_a as *mut c_void, a.as_ptr());
reim_fft(reim_fft_precomp, buf_a);
reim_from_znx64_simple(m as u32, log_bound as u32, buf_b as *mut c_void, b.as_ptr());
reim_fft(reim_fft_precomp, buf_b);
reim_fftvec_mul_simple(
m as u32,
buf_c as *mut c_void,
buf_a as *mut c_void,
buf_b as *mut c_void,
);
reim_ifft(reim_ifft_precomp, buf_c);
reim_to_znx64_simple(
m as u32,
m as f64,
log_bound as u32,
c.as_mut_ptr(),
buf_c as *mut c_void,
)
});
println!("time: {}us", now.elapsed().as_micros());
println!("{:?}", &c[..16]);
}
}

View File

@@ -0,0 +1,86 @@
use itertools::izip;
use sampling::source::Source;
use spqlios::module::{Module, FFT64};
use spqlios::scalar::Scalar;
use spqlios::vector::Vector;
fn main() {
let n: usize = 16;
let log_base2k: usize = 40;
let prec: usize = 54;
let log_scale: usize = 18;
let module: Module = Module::new::<FFT64>(n);
let mut carry: Vec<u8> = vec![0; module.vec_znx_big_normalize_tmp_bytes()];
let seed: [u8; 32] = [0; 32];
let mut source: Source = Source::new(seed);
let mut res: Vector = Vector::new(n, log_base2k, prec);
// s <- Z_{-1, 0, 1}[X]/(X^{N}+1)
let mut s: Scalar = Scalar::new(n);
s.fill_ternary_prob(0.5, &mut source);
// Buffer to store s in the DFT domain
let mut s_ppol: spqlios::module::SVPPOL = module.svp_new_ppol();
// s_ppol <- DFT(s)
module.svp_prepare(&mut s_ppol, &s);
// a <- Z_{2^prec}[X]/(X^{N}+1)
let mut a: Vector = Vector::new(n, log_base2k, prec);
a.fill_uniform(&mut source);
// Scratch space for DFT values
let mut buf_dft: spqlios::module::VECZNXDFT = module.new_vec_znx_dft(a.limbs());
// Applies buf_dft <- s * a
module.svp_apply_dft(&mut buf_dft, &s_ppol, &a);
// Alias scratch space
let mut buf_big: spqlios::module::VECZNXBIG = buf_dft.as_vec_znx_big();
// buf_big <- IDFT(buf_dft) (not normalized)
module.vec_znx_idft_tmp_a(&mut buf_big, &mut buf_dft, a.limbs());
let mut m: Vector = Vector::new(n, log_base2k, prec - log_scale);
let mut want: Vec<i64> = vec![0; n];
want.iter_mut()
.for_each(|x| *x = source.next_u64n(16, 15) as i64);
// m
m.set_i64(&want, 4);
m.normalize(&mut carry);
// buf_big <- m - buf_big
module.vec_znx_big_sub_small_a_inplace(&mut buf_big, &m);
// b <- normalize(buf_big) + e
let mut b: Vector = Vector::new(n, log_base2k, prec);
module.vec_znx_big_normalize(&mut b, &buf_big, &mut carry);
b.add_normal(&mut source, 3.2, 19.0);
//Decrypt
// buf_big <- a * s
module.svp_apply_dft(&mut buf_dft, &s_ppol, &a);
module.vec_znx_idft_tmp_a(&mut buf_big, &mut buf_dft, b.limbs());
// buf_big <- a * s + b
module.vec_znx_big_add_small_inplace(&mut buf_big, &b);
// res <- normalize(buf_big)
module.vec_znx_big_normalize(&mut res, &buf_big, &mut carry);
// have = m * 2^{log_scale} + e
let mut have: Vec<i64> = vec![i64::default(); n];
res.get_i64(&mut have);
let scale: f64 = (1 << log_scale) as f64;
izip!(want.iter(), have.iter())
.enumerate()
.for_each(|(i, (a, b))| {
println!("{}: {} {}", i, a, (*b as f64) / scale);
})
}

35
base2k/src/lib.rs Normal file
View File

@@ -0,0 +1,35 @@
pub mod module;
pub mod scalar;
pub mod vector;
#[allow(
non_camel_case_types,
non_snake_case,
non_upper_case_globals,
dead_code,
improper_ctypes
)]
pub mod bindings {
include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
}
pub mod vec_znx_arithmetic;
pub use vec_znx_arithmetic::*;
pub mod vec_znx_big_arithmetic;
pub use vec_znx_big_arithmetic::*;
pub mod vec_znx_dft;
pub use vec_znx_dft::*;
pub mod scalar_vector_product;
pub use scalar_vector_product::*;
fn cast_mut_u64_to_mut_u8_slice(data: &mut [u64]) -> &mut [u8] {
let ptr: *mut u8 = data.as_mut_ptr() as *mut u8;
let len: usize = data.len() * std::mem::size_of::<u64>();
unsafe { std::slice::from_raw_parts_mut(ptr, len) }
}
fn cast_mut_u8_to_mut_i64_slice(data: &mut [u8]) -> &mut [i64] {
let ptr: *mut i64 = data.as_mut_ptr() as *mut i64;
let len: usize = data.len() / std::mem::size_of::<i64>();
unsafe { std::slice::from_raw_parts_mut(ptr, len) }
}

47
base2k/src/module.rs Normal file
View File

@@ -0,0 +1,47 @@
use crate::bindings::{
module_info_t, new_module_info, svp_ppol_t, vec_znx_bigcoeff_t, vec_znx_dft_t, MODULE,
};
pub type MODULETYPE = u8;
pub const FFT64: u8 = 0;
pub const NTT120: u8 = 1;
pub struct Module(pub *mut MODULE);
impl Module {
// Instantiates a new module.
pub fn new<const MODULETYPE: MODULETYPE>(n: usize) -> Self {
unsafe {
let m: *mut module_info_t = new_module_info(n as u64, MODULETYPE as u32);
if m.is_null() {
panic!("Failed to create module.");
}
Self(m)
}
}
}
pub struct SVPPOL(pub *mut svp_ppol_t);
pub struct VECZNXBIG(pub *mut vec_znx_bigcoeff_t, pub usize);
// Stores a vector of
impl VECZNXBIG {
pub fn as_vec_znx_dft(&mut self) -> VECZNXDFT {
VECZNXDFT(self.0 as *mut vec_znx_dft_t, self.1)
}
pub fn limbs(&self) -> usize {
self.1
}
}
pub struct VECZNXDFT(pub *mut vec_znx_dft_t, pub usize);
impl VECZNXDFT {
pub fn as_vec_znx_big(&mut self) -> VECZNXBIG {
VECZNXBIG(self.0 as *mut vec_znx_bigcoeff_t, self.1)
}
pub fn limbs(&self) -> usize {
self.1
}
}

48
base2k/src/scalar.rs Normal file
View File

@@ -0,0 +1,48 @@
use rand::distributions::{Distribution, WeightedIndex};
use rand::seq::SliceRandom;
use rand_core::RngCore;
use sampling::source::Source;
pub struct Scalar(pub Vec<i64>);
impl Scalar {
pub fn new(n: usize) -> Self {
Self(vec![i64::default(); Self::buffer_size(n)])
}
pub fn buffer_size(n: usize) -> usize {
n
}
pub fn from_buffer(&mut self, n: usize, buf: &[i64]) {
let size = Self::buffer_size(n);
assert!(
buf.len() >= size,
"invalid buffer: buf.len()={} < self.buffer_size(n={})={}",
buf.len(),
n,
size
);
self.0 = Vec::from(&buf[..size])
}
pub fn as_ptr(&self) -> *const i64 {
self.0.as_ptr()
}
pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
let choices: [i64; 3] = [-1, 0, 1];
let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
self.0
.iter_mut()
.for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
}
pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
self.0[..hw]
.iter_mut()
.for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
self.0.shuffle(source);
}
}

View File

@@ -0,0 +1,29 @@
use crate::bindings::{new_svp_ppol, svp_apply_dft, svp_prepare};
use crate::module::{Module, SVPPOL, VECZNXDFT};
use crate::scalar::Scalar;
use crate::vector::Vector;
impl Module {
// Prepares a scalar polynomial (1 limb) for a scalar x vector product.
// Method will panic if a.limbs() != 1.
pub fn svp_prepare(&self, svp_ppol: &mut SVPPOL, a: &Scalar) {
unsafe { svp_prepare(self.0, svp_ppol.0, a.as_ptr()) }
}
// Allocates a scalar-vector-product prepared-poly (SVPPOL).
pub fn svp_new_ppol(&self) -> SVPPOL {
unsafe { SVPPOL(new_svp_ppol(self.0)) }
}
// Applies a scalar x vector product: res <- a (ppol) x b
pub fn svp_apply_dft(&self, c: &mut VECZNXDFT, a: &SVPPOL, b: &Vector) {
let limbs: u64 = b.limbs() as u64;
assert!(
c.limbs() as u64 >= limbs,
"invalid c_vector: c_vector.limbs()={} < b.limbs()={}",
c.limbs(),
limbs
);
unsafe { svp_apply_dft(self.0, c.0, limbs, a.0, b.as_ptr(), limbs, b.n() as u64) }
}
}

View File

@@ -0,0 +1,35 @@
use crate::bindings::vec_znx_automorphism;
use crate::module::Module;
use crate::vector::Vector;
impl Module {
pub fn vec_znx_automorphism(&self, gal_el: i64, b: &mut Vector, a: &Vector) {
unsafe {
vec_znx_automorphism(
self.0,
gal_el,
b.as_mut_ptr(),
b.limbs() as u64,
b.n() as u64,
a.as_ptr(),
a.limbs() as u64,
a.n() as u64,
);
}
}
pub fn vec_znx_automorphism_inplace(&self, gal_el: i64, a: &mut Vector) {
unsafe {
vec_znx_automorphism(
self.0,
gal_el,
a.as_mut_ptr(),
a.limbs() as u64,
a.n() as u64,
a.as_ptr(),
a.limbs() as u64,
a.n() as u64,
);
}
}
}

View File

@@ -0,0 +1,162 @@
use crate::bindings::{
new_vec_znx_big, vec_znx_big_add_small, vec_znx_big_automorphism, vec_znx_big_normalize_base2k,
vec_znx_big_normalize_base2k_tmp_bytes, vec_znx_big_sub_small_a,
};
use crate::module::{Module, VECZNXBIG};
use crate::vector::Vector;
impl Module {
// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
pub fn new_vec_znx_big(&self, limbs: usize) -> VECZNXBIG {
unsafe { VECZNXBIG(new_vec_znx_big(self.0, limbs as u64), limbs) }
}
// b <- b - a
pub fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VECZNXBIG, a: &Vector) {
let limbs: usize = a.limbs();
assert!(
b.limbs() >= limbs,
"invalid c_vector: b.limbs()={} < a.limbs()={}",
b.limbs(),
limbs
);
unsafe {
vec_znx_big_sub_small_a(
self.0,
b.0,
b.limbs() as u64,
a.as_ptr(),
limbs as u64,
a.n() as u64,
b.0,
b.limbs() as u64,
)
}
}
// c <- b - a
pub fn big_sub_small_a(&self, c: &mut VECZNXBIG, a: &Vector, b: &VECZNXBIG) {
let limbs: usize = a.limbs();
assert!(
b.limbs() >= limbs,
"invalid c: b.limbs()={} < a.limbs()={}",
b.limbs(),
limbs
);
assert!(
c.limbs() >= limbs,
"invalid c: c.limbs()={} < a.limbs()={}",
c.limbs(),
limbs
);
unsafe {
vec_znx_big_sub_small_a(
self.0,
c.0,
c.limbs() as u64,
a.as_ptr(),
limbs as u64,
a.n() as u64,
b.0,
b.limbs() as u64,
)
}
}
// c <- b + a
pub fn vec_znx_big_add_small(&self, c: &mut VECZNXBIG, a: &Vector, b: &VECZNXBIG) {
let limbs: usize = a.limbs();
assert!(
b.limbs() >= limbs,
"invalid c: b.limbs()={} < a.limbs()={}",
b.limbs(),
limbs
);
assert!(
c.limbs() >= limbs,
"invalid c: c.limbs()={} < a.limbs()={}",
c.limbs(),
limbs
);
unsafe {
vec_znx_big_add_small(
self.0,
c.0,
limbs as u64,
b.0,
limbs as u64,
a.as_ptr(),
limbs as u64,
a.n() as u64,
)
}
}
// b <- b + a
pub fn vec_znx_big_add_small_inplace(&self, b: &mut VECZNXBIG, a: &Vector) {
let limbs: usize = a.limbs();
assert!(
b.limbs() >= limbs,
"invalid c_vector: b.limbs()={} < a.limbs()={}",
b.limbs(),
limbs
);
unsafe {
vec_znx_big_add_small(
self.0,
b.0,
limbs as u64,
b.0,
limbs as u64,
a.as_ptr(),
limbs as u64,
a.n() as u64,
)
}
}
pub fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
unsafe { vec_znx_big_normalize_base2k_tmp_bytes(self.0) as usize }
}
// b <- normalize(a)
pub fn vec_znx_big_normalize(&self, b: &mut Vector, a: &VECZNXBIG, tmp_bytes: &mut [u8]) {
let limbs: usize = b.limbs();
assert!(
b.limbs() >= limbs,
"invalid c_vector: b.limbs()={} < a.limbs()={}",
b.limbs(),
limbs
);
assert!(
tmp_bytes.len() <= self.vec_znx_big_normalize_tmp_bytes(),
"invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",
tmp_bytes.len(),
self.vec_znx_big_normalize_tmp_bytes()
);
unsafe {
vec_znx_big_normalize_base2k(
self.0,
b.log_base2k as u64,
b.as_mut_ptr(),
limbs as u64,
b.n() as u64,
a.0,
limbs as u64,
tmp_bytes.as_mut_ptr(),
)
}
}
pub fn big_automorphism(&self, gal_el: i64, b: &mut VECZNXBIG, a: &VECZNXBIG) {
unsafe {
vec_znx_big_automorphism(self.0, gal_el, b.0, b.limbs() as u64, a.0, a.limbs() as u64);
}
}
pub fn big_automorphism_inplace(&self, gal_el: i64, a: &mut VECZNXBIG) {
unsafe {
vec_znx_big_automorphism(self.0, gal_el, a.0, a.limbs() as u64, a.0, a.limbs() as u64);
}
}
}

63
base2k/src/vec_znx_dft.rs Normal file
View File

@@ -0,0 +1,63 @@
use crate::bindings::{new_vec_znx_dft, vec_znx_idft, vec_znx_idft_tmp_a, vec_znx_idft_tmp_bytes};
use crate::module::{Module, VECZNXBIG, VECZNXDFT};
impl Module {
// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
pub fn new_vec_znx_dft(&self, limbs: usize) -> VECZNXDFT {
unsafe { VECZNXDFT(new_vec_znx_dft(self.0, limbs as u64), limbs) }
}
// b <- IDFT(a), uses a as scratch space.
pub fn vec_znx_idft_tmp_a(&self, b: &mut VECZNXBIG, a: &mut VECZNXDFT, a_limbs: usize) {
assert!(
b.limbs() >= a_limbs,
"invalid c_vector: b_vector.limbs()={} < a_limbs={}",
b.limbs(),
a_limbs
);
unsafe { vec_znx_idft_tmp_a(self.0, b.0, a_limbs as u64, a.0, a_limbs as u64) }
}
// Returns the size of the scratch space for [vec_znx_idft].
pub fn vec_znx_idft_tmp_bytes(&self) -> usize {
unsafe { vec_znx_idft_tmp_bytes(self.0) as usize }
}
// b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
pub fn vec_znx_idft(
&self,
b_vector: &mut VECZNXBIG,
a_vector: &mut VECZNXDFT,
a_limbs: usize,
tmp_bytes: &mut [u8],
) {
assert!(
b_vector.limbs() >= a_limbs,
"invalid c_vector: b_vector.limbs()={} < a_limbs={}",
b_vector.limbs(),
a_limbs
);
assert!(
a_vector.limbs() >= a_limbs,
"invalid c_vector: c_vector.limbs()={} < a_limbs={}",
a_vector.limbs(),
a_limbs
);
assert!(
tmp_bytes.len() <= self.vec_znx_idft_tmp_bytes(),
"invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
tmp_bytes.len(),
self.vec_znx_idft_tmp_bytes()
);
unsafe {
vec_znx_idft(
self.0,
b_vector.0,
a_limbs as u64,
a_vector.0,
a_limbs as u64,
tmp_bytes.as_mut_ptr(),
)
}
}
}

328
base2k/src/vector.rs Normal file
View File

@@ -0,0 +1,328 @@
use crate::bindings::{
znx_automorphism_i64, znx_automorphism_inplace_i64, znx_normalize, znx_zero_i64_ref,
};
use crate::cast_mut_u8_to_mut_i64_slice;
use itertools::izip;
use rand_distr::{Distribution, Normal};
use sampling::source::Source;
use std::cmp::min;
pub struct Vector {
pub n: usize,
pub log_base2k: usize,
pub prec: usize,
pub data: Vec<i64>,
}
impl Vector {
pub fn new(n: usize, log_base2k: usize, prec: usize) -> Self {
Self {
n: n,
log_base2k: log_base2k,
prec: prec,
data: vec![i64::default(); Self::buffer_size(n, log_base2k, prec)],
}
}
pub fn buffer_size(n: usize, log_base2k: usize, prec: usize) -> usize {
n * ((prec + log_base2k - 1) / log_base2k)
}
pub fn from_buffer(&mut self, n: usize, log_base2k: usize, prec: usize, buf: &[i64]) {
let size = Self::buffer_size(n, log_base2k, prec);
assert!(
buf.len() >= size,
"invalid buffer: buf.len()={} < self.buffer_size(n={}, k={}, prec={})={}",
buf.len(),
n,
log_base2k,
prec,
size
);
self.n = n;
self.log_base2k = log_base2k;
self.prec = prec;
self.data = Vec::from(&buf[..size])
}
pub fn log_n(&self) -> u64 {
(u64::BITS - (self.n - 1).leading_zeros()) as _
}
pub fn n(&self) -> usize {
self.n
}
pub fn prec(&self) -> usize {
self.prec
}
pub fn limbs(&self) -> usize {
self.data.len() / self.n
}
pub fn as_ptr(&self) -> *const i64 {
self.data.as_ptr()
}
pub fn as_mut_ptr(&mut self) -> *mut i64 {
self.data.as_mut_ptr()
}
pub fn at(&self, i: usize) -> &[i64] {
&self.data[i * self.n..(i + 1) * self.n]
}
pub fn at_ptr(&self, i: usize) -> *const i64 {
&self.data[i * self.n] as *const i64
}
pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
&mut self.data[i * self.n] as *mut i64
}
pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
&mut self.data[i * self.n..(i + 1) * self.n]
}
pub fn set_i64(&mut self, data: &[i64], log_max: usize) {
let size: usize = min(data.len(), self.n());
let k_rem: usize = self.log_base2k - (self.prec % self.log_base2k);
// If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
// values on the last limb.
// Else we decompose values base2k.
if log_max + k_rem < 63 || k_rem == self.log_base2k {
self.at_mut(self.limbs() - 1).copy_from_slice(&data[..size]);
} else {
let mask: i64 = (1 << self.log_base2k) - 1;
let limbs = self.limbs();
let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
(limbs - steps..limbs)
.rev()
.enumerate()
.for_each(|(i, i_rev)| {
let shift: usize = i * self.log_base2k;
izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
.for_each(|(y, x)| *y = (x >> shift) & mask);
})
}
// Case where self.prec % self.k != 0.
if k_rem != self.log_base2k {
let limbs = self.limbs();
let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
(limbs - steps..limbs).rev().for_each(|i| {
self.at_mut(i)[..size].iter_mut().for_each(|x| *x <<= k_rem);
})
}
}
pub fn normalize(&mut self, carry: &mut [u8]) {
assert!(
carry.len() >= self.n * 8,
"invalid carry: carry.len()={} < self.n()={}",
carry.len(),
self.n()
);
let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
unsafe {
znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
(0..self.limbs()).rev().for_each(|i| {
znx_normalize(
self.n as u64,
self.log_base2k as u64,
self.at_mut_ptr(i),
carry_i64.as_mut_ptr(),
self.at_mut_ptr(i),
carry_i64.as_mut_ptr(),
)
});
}
}
pub fn get_i64(&self, data: &mut [i64]) {
assert!(
data.len() >= self.n,
"invalid data: data.len()={} < self.n()={}",
data.len(),
self.n
);
data.copy_from_slice(self.at(0));
let rem: usize = self.log_base2k - (self.prec % self.log_base2k);
(1..self.limbs()).for_each(|i| {
if i == self.limbs() - 1 && rem != self.log_base2k {
let k_rem: usize = self.log_base2k - rem;
izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
*y = (*y << k_rem) + (x >> rem);
});
} else {
izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
*y = (*y << self.log_base2k) + x;
});
}
})
}
pub fn automorphism_inplace(&mut self, gal_el: i64) {
unsafe {
(0..self.limbs()).for_each(|i| {
znx_automorphism_inplace_i64(self.n as u64, gal_el, self.at_mut_ptr(i))
})
}
}
pub fn automorphism(&mut self, gal_el: i64, a: &mut Vector) {
unsafe {
(0..self.limbs()).for_each(|i| {
znx_automorphism_i64(self.n as u64, gal_el, a.at_mut_ptr(i), self.at_ptr(i))
})
}
}
pub fn fill_uniform(&mut self, source: &mut Source) {
let mut base2k: u64 = 1 << self.log_base2k;
let mut mask: u64 = base2k - 1;
let mut base2k_half: i64 = (base2k >> 1) as i64;
let size: usize = self.n() * (self.limbs() - 1);
self.data[..size]
.iter_mut()
.for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
let log_base2k_rem: usize = self.prec % self.log_base2k;
if log_base2k_rem != 0 {
base2k = 1 << log_base2k_rem;
mask = (base2k - 1) << (self.log_base2k - log_base2k_rem);
base2k_half = ((mask >> 1) + 1) as i64;
}
self.data[size..]
.iter_mut()
.for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
}
pub fn add_dist_f64<T: Distribution<f64>>(&mut self, source: &mut Source, dist: T, bound: f64) {
let log_base2k_rem: usize = self.prec % self.log_base2k;
if log_base2k_rem != 0 {
self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
let mut dist_f64: f64 = dist.sample(source);
while dist_f64.abs() > bound {
dist_f64 = dist.sample(source)
}
*a += (dist_f64.round() as i64) << log_base2k_rem
});
} else {
self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
let mut dist_f64: f64 = dist.sample(source);
while dist_f64.abs() > bound {
dist_f64 = dist.sample(source)
}
*a += dist_f64.round() as i64
});
}
}
pub fn add_normal(&mut self, source: &mut Source, sigma: f64, bound: f64) {
self.add_dist_f64(source, Normal::new(0.0, sigma).unwrap(), bound);
}
pub fn trunc_pow2(&mut self, k: usize) {
if k == 0 {
return;
}
assert!(
k <= self.prec,
"invalid argument k: k={} > self.prec()={}",
k,
self.prec()
);
self.prec -= k;
self.data
.truncate((self.limbs() - k / self.log_base2k) * self.n());
let k_rem: usize = k % self.log_base2k;
if k_rem != 0 {
let mask: i64 = ((1 << (self.log_base2k - k_rem - 1)) - 1) << k_rem;
self.at_mut(self.limbs() - 1)
.iter_mut()
.for_each(|x: &mut i64| *x &= mask)
}
}
}
#[cfg(test)]
mod tests {
use crate::vector::Vector;
use itertools::izip;
use sampling::source::Source;
#[test]
fn test_set_get_i64_lo_norm() {
let n: usize = 32;
let k: usize = 19;
let prec: usize = 128;
let mut a: Vector = Vector::new(n, k, prec);
let mut have: Vec<i64> = vec![i64::default(); n];
have.iter_mut()
.enumerate()
.for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
a.set_i64(&have, 10);
let mut want = vec![i64::default(); n];
a.get_i64(&mut want);
izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
}
#[test]
fn test_set_get_i64_hi_norm() {
let n: usize = 8;
let k: usize = 17;
let prec: usize = 84;
let mut a: Vector = Vector::new(n, k, prec);
let mut have: Vec<i64> = vec![i64::default(); n];
let mut source = Source::new([1; 32]);
have.iter_mut().for_each(|x| {
*x = source
.next_u64n(u64::MAX, u64::MAX)
.wrapping_sub(u64::MAX / 2 + 1) as i64;
});
a.set_i64(&have, 63);
//(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
let mut want = vec![i64::default(); n];
//(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
a.get_i64(&mut want);
izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
}
#[test]
fn test_normalize() {
let n: usize = 8;
let k: usize = 17;
let prec: usize = 84;
let mut a: Vector = Vector::new(n, k, prec);
let mut have: Vec<i64> = vec![i64::default(); n];
let mut source = Source::new([1; 32]);
have.iter_mut().for_each(|x| {
*x = source
.next_u64n(u64::MAX, u64::MAX)
.wrapping_sub(u64::MAX / 2 + 1) as i64;
});
a.set_i64(&have, 63);
let mut carry: Vec<u8> = vec![u8::default(); n * 8];
a.normalize(&mut carry);
let base_half = 1 << (k - 1);
a.data
.iter()
.for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
let mut want = vec![i64::default(); n];
a.get_i64(&mut want);
izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
}
}