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

Added basic key-switching + file formatting

Browse Source
This commit is contained in:
Jean-Philippe Bossuat
2025-04-24 10:43:51 +02:00
parent 4196477300
commit ad6e8169e5
33 changed files with 319 additions and 715 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
base2k/benches/fft.rs
View File

@@ -1,14 +1,9 @@
use base2k::ffi::reim::*;
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
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> {
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());
@@ -16,12 +11,7 @@ fn fft(c: &mut Criterion) {
}
}
fn backward<'a>(
m: u32,
log_bound: u32,
reim_ifft_precomp: *mut reim_ifft_precomp,
a: &'a [i64],
) -> Box<dyn FnMut() + '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());
@@ -29,8 +19,7 @@ fn fft(c: &mut Criterion) {
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
c.benchmark_group("fft");
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("fft");
for log_n in 10..17 {
let n: usize = 1 << log_n;

6
base2k/examples/rlwe_encrypt.rs
View File

@@ -1,6 +1,6 @@
use base2k::{
alloc_aligned, Encoding, Infos, Module, Sampling, Scalar, SvpPPol, SvpPPolOps, VecZnx,
VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, BACKEND,
BACKEND, Encoding, Infos, Module, Sampling, Scalar, SvpPPol, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft,
VecZnxDftOps, VecZnxOps, alloc_aligned,
};
use itertools::izip;
use sampling::source::Source;
@@ -71,7 +71,7 @@ fn main() {
19.0,
);
//Decrypt
// Decrypt
// buf_big <- a * s
module.svp_apply_dft(&mut buf_dft, &s_ppol, &a);

7
base2k/examples/vector_matrix_product.rs
View File

@@ -1,6 +1,6 @@
use base2k::{
alloc_aligned, Encoding, Infos, Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft,
VecZnxDftOps, VecZnxOps, VecZnxVec, VmpPMat, VmpPMatOps, BACKEND,
BACKEND, Encoding, Infos, Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VecZnxVec, VmpPMat,
VmpPMatOps, alloc_aligned,
};
fn main() {
@@ -16,8 +16,7 @@ fn main() {
let cols: usize = cols + 1;
// Maximum size of the byte scratch needed
let tmp_bytes: usize = module.vmp_prepare_tmp_bytes(rows, cols)
| module.vmp_apply_dft_tmp_bytes(cols, cols, rows, cols);
let tmp_bytes: usize = module.vmp_prepare_tmp_bytes(rows, cols) | module.vmp_apply_dft_tmp_bytes(cols, cols, rows, cols);
let mut buf: Vec<u8> = alloc_aligned(tmp_bytes);

30
base2k/src/encoding.rs
View File

@@ -40,14 +40,7 @@ pub trait Encoding {
/// * `i`: index of the coefficient on which to encode the data.
/// * `data`: data to encode on the receiver.
/// * `log_max`: base two logarithm of the infinity norm of the input data.
fn encode_coeff_i64(
&mut self,
log_base2k: usize,
log_k: usize,
i: usize,
data: i64,
log_max: usize,
);
fn encode_coeff_i64(&mut self, log_base2k: usize, log_k: usize, i: usize, data: i64, log_max: usize);
/// decode a single of i64 from the receiver at the given index.
///
@@ -73,14 +66,7 @@ impl Encoding for VecZnx {
decode_vec_float(self, log_base2k, data)
}
fn encode_coeff_i64(
&mut self,
log_base2k: usize,
log_k: usize,
i: usize,
value: i64,
log_max: usize,
) {
fn encode_coeff_i64(&mut self, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
encode_coeff_i64(self, log_base2k, log_k, i, value, log_max)
}
@@ -119,8 +105,7 @@ fn encode_vec_i64(a: &mut VecZnx, log_base2k: usize, log_k: usize, data: &[i64],
.enumerate()
.for_each(|(i, i_rev)| {
let shift: usize = i * log_base2k;
izip!(a.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
.for_each(|(y, x)| *y = (x >> shift) & mask);
izip!(a.at_mut(i_rev)[..size].iter_mut(), data[..size].iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
})
}
@@ -189,14 +174,7 @@ fn decode_vec_float(a: &VecZnx, log_base2k: usize, data: &mut [Float]) {
});
}
fn encode_coeff_i64(
a: &mut VecZnx,
log_base2k: usize,
log_k: usize,
i: usize,
value: i64,
log_max: usize,
) {
fn encode_coeff_i64(a: &mut VecZnx, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
debug_assert!(i < a.n());
let cols: usize = (log_k + log_base2k - 1) / log_base2k;
debug_assert!(

115
base2k/src/ffi/reim.rs
View File

@@ -62,10 +62,7 @@ unsafe extern "C" {
pub unsafe fn new_reim_fft_precomp(m: u32, num_buffers: u32) -> *mut REIM_FFT_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_fft_precomp_get_buffer(
tables: *const REIM_FFT_PRECOMP,
buffer_index: u32,
) -> *mut f64;
pub unsafe fn reim_fft_precomp_get_buffer(tables: *const REIM_FFT_PRECOMP, buffer_index: u32) -> *mut f64;
}
unsafe extern "C" {
pub unsafe fn new_reim_fft_buffer(m: u32) -> *mut f64;
@@ -80,10 +77,7 @@ unsafe extern "C" {
pub unsafe fn new_reim_ifft_precomp(m: u32, num_buffers: u32) -> *mut REIM_IFFT_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_ifft_precomp_get_buffer(
tables: *const REIM_IFFT_PRECOMP,
buffer_index: u32,
) -> *mut f64;
pub unsafe fn reim_ifft_precomp_get_buffer(tables: *const REIM_IFFT_PRECOMP, buffer_index: u32) -> *mut f64;
}
unsafe extern "C" {
pub unsafe fn reim_ifft(tables: *const REIM_IFFT_PRECOMP, data: *mut f64);
@@ -92,120 +86,58 @@ unsafe extern "C" {
pub unsafe fn new_reim_fftvec_mul_precomp(m: u32) -> *mut REIM_FFTVEC_MUL_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_fftvec_mul(
tables: *const REIM_FFTVEC_MUL_PRECOMP,
r: *mut f64,
a: *const f64,
b: *const f64,
);
pub unsafe fn reim_fftvec_mul(tables: *const REIM_FFTVEC_MUL_PRECOMP, r: *mut f64, a: *const f64, b: *const f64);
}
unsafe extern "C" {
pub unsafe fn new_reim_fftvec_addmul_precomp(m: u32) -> *mut REIM_FFTVEC_ADDMUL_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_fftvec_addmul(
tables: *const REIM_FFTVEC_ADDMUL_PRECOMP,
r: *mut f64,
a: *const f64,
b: *const f64,
);
pub unsafe fn reim_fftvec_addmul(tables: *const REIM_FFTVEC_ADDMUL_PRECOMP, r: *mut f64, a: *const f64, b: *const f64);
}
unsafe extern "C" {
pub unsafe fn new_reim_from_znx32_precomp(
m: u32,
log2bound: u32,
) -> *mut REIM_FROM_ZNX32_PRECOMP;
pub unsafe fn new_reim_from_znx32_precomp(m: u32, log2bound: u32) -> *mut REIM_FROM_ZNX32_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_from_znx32(
tables: *const REIM_FROM_ZNX32_PRECOMP,
r: *mut ::std::os::raw::c_void,
a: *const i32,
);
pub unsafe fn reim_from_znx32(tables: *const REIM_FROM_ZNX32_PRECOMP, r: *mut ::std::os::raw::c_void, a: *const i32);
}
unsafe extern "C" {
pub unsafe fn reim_from_znx64(
tables: *const REIM_FROM_ZNX64_PRECOMP,
r: *mut ::std::os::raw::c_void,
a: *const i64,
);
pub unsafe fn reim_from_znx64(tables: *const REIM_FROM_ZNX64_PRECOMP, r: *mut ::std::os::raw::c_void, a: *const i64);
}
unsafe extern "C" {
pub unsafe fn new_reim_from_znx64_precomp(m: u32, maxbnd: u32) -> *mut REIM_FROM_ZNX64_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_from_znx64_simple(
m: u32,
log2bound: u32,
r: *mut ::std::os::raw::c_void,
a: *const i64,
);
pub unsafe fn reim_from_znx64_simple(m: u32, log2bound: u32, r: *mut ::std::os::raw::c_void, a: *const i64);
}
unsafe extern "C" {
pub unsafe fn new_reim_from_tnx32_precomp(m: u32) -> *mut REIM_FROM_TNX32_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_from_tnx32(
tables: *const REIM_FROM_TNX32_PRECOMP,
r: *mut ::std::os::raw::c_void,
a: *const i32,
);
pub unsafe fn reim_from_tnx32(tables: *const REIM_FROM_TNX32_PRECOMP, r: *mut ::std::os::raw::c_void, a: *const i32);
}
unsafe extern "C" {
pub unsafe fn new_reim_to_tnx32_precomp(
m: u32,
divisor: f64,
log2overhead: u32,
) -> *mut REIM_TO_TNX32_PRECOMP;
pub unsafe fn new_reim_to_tnx32_precomp(m: u32, divisor: f64, log2overhead: u32) -> *mut REIM_TO_TNX32_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_to_tnx32(
tables: *const REIM_TO_TNX32_PRECOMP,
r: *mut i32,
a: *const ::std::os::raw::c_void,
);
pub unsafe fn reim_to_tnx32(tables: *const REIM_TO_TNX32_PRECOMP, r: *mut i32, a: *const ::std::os::raw::c_void);
}
unsafe extern "C" {
pub unsafe fn new_reim_to_tnx_precomp(
m: u32,
divisor: f64,
log2overhead: u32,
) -> *mut REIM_TO_TNX_PRECOMP;
pub unsafe fn new_reim_to_tnx_precomp(m: u32, divisor: f64, log2overhead: u32) -> *mut REIM_TO_TNX_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_to_tnx(tables: *const REIM_TO_TNX_PRECOMP, r: *mut f64, a: *const f64);
}
unsafe extern "C" {
pub unsafe fn reim_to_tnx_simple(
m: u32,
divisor: f64,
log2overhead: u32,
r: *mut f64,
a: *const f64,
);
pub unsafe fn reim_to_tnx_simple(m: u32, divisor: f64, log2overhead: u32, r: *mut f64, a: *const f64);
}
unsafe extern "C" {
pub unsafe fn new_reim_to_znx64_precomp(
m: u32,
divisor: f64,
log2bound: u32,
) -> *mut REIM_TO_ZNX64_PRECOMP;
pub unsafe fn new_reim_to_znx64_precomp(m: u32, divisor: f64, log2bound: u32) -> *mut REIM_TO_ZNX64_PRECOMP;
}
unsafe extern "C" {
pub unsafe fn reim_to_znx64(
precomp: *const REIM_TO_ZNX64_PRECOMP,
r: *mut i64,
a: *const ::std::os::raw::c_void,
);
pub unsafe fn reim_to_znx64(precomp: *const REIM_TO_ZNX64_PRECOMP, r: *mut i64, a: *const ::std::os::raw::c_void);
}
unsafe extern "C" {
pub unsafe fn reim_to_znx64_simple(
m: u32,
divisor: f64,
log2bound: u32,
r: *mut i64,
a: *const ::std::os::raw::c_void,
);
pub unsafe fn reim_to_znx64_simple(m: u32, divisor: f64, log2bound: u32, r: *mut i64, a: *const ::std::os::raw::c_void);
}
unsafe extern "C" {
pub unsafe fn reim_fft_simple(m: u32, data: *mut ::std::os::raw::c_void);
@@ -230,22 +162,11 @@ unsafe extern "C" {
);
}
unsafe extern "C" {
pub unsafe fn reim_from_znx32_simple(
m: u32,
log2bound: u32,
r: *mut ::std::os::raw::c_void,
x: *const i32,
);
pub unsafe fn reim_from_znx32_simple(m: u32, log2bound: u32, r: *mut ::std::os::raw::c_void, x: *const i32);
}
unsafe extern "C" {
pub unsafe fn reim_from_tnx32_simple(m: u32, r: *mut ::std::os::raw::c_void, x: *const i32);
}
unsafe extern "C" {
pub unsafe fn reim_to_tnx32_simple(
m: u32,
divisor: f64,
log2overhead: u32,
r: *mut i32,
x: *const ::std::os::raw::c_void,
);
pub unsafe fn reim_to_tnx32_simple(m: u32, divisor: f64, log2overhead: u32, r: *mut i32, x: *const ::std::os::raw::c_void);
}

9
base2k/src/ffi/vec_znx_dft.rs
View File

@@ -44,14 +44,7 @@ unsafe extern "C" {
);
}
unsafe extern "C" {
pub unsafe fn vec_znx_dft(
module: *const MODULE,
res: *mut VEC_ZNX_DFT,
res_size: u64,
a: *const i64,
a_size: u64,
a_sl: u64,
);
pub unsafe fn vec_znx_dft(module: *const MODULE, res: *mut VEC_ZNX_DFT, res_size: u64, a: *const i64, a_size: u64, a_sl: u64);
}
unsafe extern "C" {
pub unsafe fn vec_znx_idft(

8
base2k/src/ffi/vmp.rs
View File

@@ -52,13 +52,7 @@ unsafe extern "C" {
}
unsafe extern "C" {
pub unsafe fn vmp_apply_dft_tmp_bytes(
module: *const MODULE,
res_size: u64,
a_size: u64,
nrows: u64,
ncols: u64,
) -> u64;
pub unsafe fn vmp_apply_dft_tmp_bytes(module: *const MODULE, res_size: u64, a_size: u64, nrows: u64, ncols: u64) -> u64;
}
unsafe extern "C" {

17
base2k/src/ffi/znx.rs
View File

@@ -64,24 +64,11 @@ unsafe extern "C" {
pub unsafe fn rnx_mul_xp_minus_one_inplace(nn: u64, p: i64, res: *mut f64);
}
unsafe extern "C" {
pub unsafe fn znx_normalize(
nn: u64,
base_k: u64,
out: *mut i64,
carry_out: *mut i64,
in_: *const i64,
carry_in: *const i64,
);
pub unsafe fn znx_normalize(nn: u64, base_k: u64, out: *mut i64, carry_out: *mut i64, in_: *const i64, carry_in: *const i64);
}
unsafe extern "C" {
pub unsafe fn znx_small_single_product(
module: *const MODULE,
res: *mut i64,
a: *const i64,
b: *const i64,
tmp: *mut u8,
);
pub unsafe fn znx_small_single_product(module: *const MODULE, res: *mut i64, a: *const i64, b: *const i64, tmp: *mut u8);
}
unsafe extern "C" {

15
base2k/src/lib.rs
View File

@@ -1,11 +1,5 @@
pub mod encoding;
#[allow(
non_camel_case_types,
non_snake_case,
non_upper_case_globals,
dead_code,
improper_ctypes
)]
#[allow(non_camel_case_types, non_snake_case, non_upper_case_globals, dead_code, improper_ctypes)]
// Other modules and exports
pub mod ffi;
pub mod infos;
@@ -42,7 +36,10 @@ pub fn is_aligned<T>(ptr: *const T) -> bool {
}
pub fn assert_alignement<T>(ptr: *const T) {
assert!(is_aligned(ptr), "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]")
assert!(
is_aligned(ptr),
"invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
)
}
pub fn cast<T, V>(data: &[T]) -> &[V] {
@@ -57,7 +54,7 @@ pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
unsafe { std::slice::from_raw_parts_mut(ptr, len) }
}
use std::alloc::{alloc, Layout};
use std::alloc::{Layout, alloc};
use std::ptr;
/// Allocates a block of bytes with a custom alignement.

8
base2k/src/module.rs
View File

@@ -1,5 +1,5 @@
use crate::ffi::module::{delete_module_info, module_info_t, new_module_info, MODULE};
use crate::GALOISGENERATOR;
use crate::ffi::module::{MODULE, delete_module_info, module_info_t, new_module_info};
#[derive(Copy, Clone)]
#[repr(u8)]
@@ -56,8 +56,7 @@ impl Module {
if gen == 0 {
return 1;
}
((mod_exp_u64(GALOISGENERATOR, gen.abs() as usize) & (self.cyclotomic_order() - 1)) as i64)
* gen.signum()
((mod_exp_u64(GALOISGENERATOR, gen.abs() as usize) & (self.cyclotomic_order() - 1)) as i64) * gen.signum()
}
// Returns gen^-1
@@ -65,8 +64,7 @@ impl Module {
if gen == 0 {
panic!("cannot invert 0")
}
((mod_exp_u64(gen.abs() as u64, (self.cyclotomic_order() - 1) as usize)
& (self.cyclotomic_order() - 1)) as i64)
((mod_exp_u64(gen.abs() as u64, (self.cyclotomic_order() - 1) as usize) & (self.cyclotomic_order() - 1)) as i64)
* gen.signum()
}

20
base2k/src/sampling.rs
View File

@@ -18,15 +18,7 @@ pub trait Sampling {
);
/// Adds a discrete normal vector scaled by 2^{-log_k} with the provided standard deviation and bounded to \[-bound, bound\].
fn add_normal(
&self,
log_base2k: usize,
a: &mut VecZnx,
log_k: usize,
source: &mut Source,
sigma: f64,
bound: f64,
);
fn add_normal(&self, log_base2k: usize, a: &mut VecZnx, log_k: usize, source: &mut Source, sigma: f64, bound: f64);
}
impl Sampling for Module {
@@ -76,15 +68,7 @@ impl Sampling for Module {
}
}
fn add_normal(
&self,
log_base2k: usize,
a: &mut VecZnx,
log_k: usize,
source: &mut Source,
sigma: f64,
bound: f64,
) {
fn add_normal(&self, log_base2k: usize, a: &mut VecZnx, log_k: usize, source: &mut Source, sigma: f64, bound: f64) {
self.add_dist_f64(
log_base2k,
a,

2
base2k/src/stats.rs
View File

@@ -1,7 +1,7 @@
use crate::{Encoding, Infos, VecZnx};
use rug::Float;
use rug::float::Round;
use rug::ops::{AddAssignRound, DivAssignRound, SubAssignRound};
use rug::Float;
impl VecZnx {
pub fn std(&self, log_base2k: usize) -> f64 {

4
base2k/src/svp.rs
View File

@@ -1,8 +1,8 @@
use crate::ffi::svp::{self, svp_ppol_t};
use crate::ffi::vec_znx_dft::vec_znx_dft_t;
use crate::{assert_alignement, Module, VecZnx, VecZnxDft, BACKEND};
use crate::{BACKEND, Module, VecZnx, VecZnxDft, assert_alignement};
use crate::{alloc_aligned, cast_mut, Infos};
use crate::{Infos, alloc_aligned, cast_mut};
use rand::seq::SliceRandom;
use rand_core::RngCore;
use rand_distr::{Distribution, WeightedIndex};

2
base2k/src/vec_znx.rs
View File

@@ -1,8 +1,8 @@
use crate::cast_mut;
use crate::ffi::vec_znx;
use crate::ffi::znx;
use crate::{alloc_aligned, assert_alignement};
use crate::{Infos, Module};
use crate::{alloc_aligned, assert_alignement};
use itertools::izip;
use std::cmp::min;

18
base2k/src/vec_znx_big.rs
View File

@@ -1,5 +1,5 @@
use crate::ffi::vec_znx_big::{self, vec_znx_big_t};
use crate::{alloc_aligned, assert_alignement, Infos, Module, VecZnx, VecZnxDft, BACKEND};
use crate::{BACKEND, Infos, Module, VecZnx, VecZnxDft, alloc_aligned, assert_alignement};
pub struct VecZnxBig {
pub data: Vec<u8>,
@@ -141,13 +141,7 @@ pub trait VecZnxBigOps {
fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
/// b <- normalize(a)
fn vec_znx_big_normalize(
&self,
log_base2k: usize,
b: &mut VecZnx,
a: &VecZnxBig,
tmp_bytes: &mut [u8],
);
fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig, tmp_bytes: &mut [u8]);
fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize;
@@ -256,13 +250,7 @@ impl VecZnxBigOps for Module {
unsafe { vec_znx_big::vec_znx_big_normalize_base2k_tmp_bytes(self.ptr) as usize }
}
fn vec_znx_big_normalize(
&self,
log_base2k: usize,
b: &mut VecZnx,
a: &VecZnxBig,
tmp_bytes: &mut [u8],
) {
fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig, tmp_bytes: &mut [u8]) {
debug_assert!(
tmp_bytes.len() >= <Module as VecZnxBigOps>::vec_znx_big_normalize_tmp_bytes(self),
"invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",

10
base2k/src/vec_znx_dft.rs
View File

@@ -1,8 +1,8 @@
use crate::ffi::vec_znx_big::vec_znx_big_t;
use crate::ffi::vec_znx_dft;
use crate::ffi::vec_znx_dft::{bytes_of_vec_znx_dft, vec_znx_dft_t};
use crate::{alloc_aligned, VecZnx, DEFAULTALIGN};
use crate::{assert_alignement, Infos, Module, VecZnxBig, BACKEND};
use crate::{BACKEND, Infos, Module, VecZnxBig, assert_alignement};
use crate::{DEFAULTALIGN, VecZnx, alloc_aligned};
pub struct VecZnxDft {
pub data: Vec<u8>,
@@ -307,11 +307,9 @@ impl VecZnxDftOps for Module {
#[cfg(test)]
mod tests {
use crate::{
alloc_aligned, Module, Sampling, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxOps, BACKEND,
};
use crate::{BACKEND, Module, Sampling, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxOps, alloc_aligned};
use itertools::izip;
use sampling::source::{new_seed, Source};
use sampling::source::{Source, new_seed};
#[test]
fn test_automorphism_dft() {

97
base2k/src/vmp.rs
View File

@@ -1,9 +1,7 @@
use crate::ffi::vec_znx_big::vec_znx_big_t;
use crate::ffi::vec_znx_dft::vec_znx_dft_t;
use crate::ffi::vmp::{self, vmp_pmat_t};
use crate::{
alloc_aligned, assert_alignement, Infos, Module, VecZnx, VecZnxBig, VecZnxDft, BACKEND,
};
use crate::{BACKEND, Infos, Module, VecZnx, VecZnxBig, VecZnxDft, alloc_aligned, assert_alignement};
/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
/// stored as a 3D matrix in the DFT domain in a single contiguous array.
@@ -100,8 +98,7 @@ impl VmpPMat {
if self.n < 8 {
res.copy_from_slice(
&self.raw::<T>()[(row + col * self.rows()) * self.n()
..(row + col * self.rows()) * (self.n() + 1)],
&self.raw::<T>()[(row + col * self.rows()) * self.n()..(row + col * self.rows()) * (self.n() + 1)],
);
} else {
(0..self.n >> 3).for_each(|blk| {
@@ -120,10 +117,7 @@ impl VmpPMat {
if col == (ncols - 1) && (ncols & 1 == 1) {
&self.raw::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
} else {
&self.raw::<T>()[blk * nrows * ncols * 8
+ (col / 2) * (2 * nrows) * 8
+ row * 2 * 8
+ (col % 2) * 8..]
&self.raw::<T>()[blk * nrows * ncols * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
}
}
}
@@ -220,13 +214,7 @@ pub trait VmpPMatOps {
/// * `a_cols`: number of cols of the input [VecZnx].
/// * `rows`: number of rows of the input [VmpPMat].
/// * `cols`: number of cols of the input [VmpPMat].
fn vmp_apply_dft_tmp_bytes(
&self,
c_cols: usize,
a_cols: usize,
rows: usize,
cols: usize,
) -> usize;
fn vmp_apply_dft_tmp_bytes(&self, c_cols: usize, a_cols: usize, rows: usize, cols: usize) -> usize;
/// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
///
@@ -288,13 +276,7 @@ pub trait VmpPMatOps {
/// * `a_cols`: number of cols of the input [VecZnxDft].
/// * `rows`: number of rows of the input [VmpPMat].
/// * `cols`: number of cols of the input [VmpPMat].
fn vmp_apply_dft_to_dft_tmp_bytes(
&self,
c_cols: usize,
a_cols: usize,
rows: usize,
cols: usize,
) -> usize;
fn vmp_apply_dft_to_dft_tmp_bytes(&self, c_cols: usize, a_cols: usize, rows: usize, cols: usize) -> usize;
/// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
/// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
@@ -348,13 +330,7 @@ pub trait VmpPMatOps {
/// * `a`: the left operand [VecZnxDft] of the vector matrix product.
/// * `b`: the right operand [VmpPMat] of the vector matrix product.
/// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
fn vmp_apply_dft_to_dft_add(
&self,
c: &mut VecZnxDft,
a: &VecZnxDft,
b: &VmpPMat,
buf: &mut [u8],
);
fn vmp_apply_dft_to_dft_add(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
/// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
/// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
@@ -521,13 +497,7 @@ impl VmpPMatOps for Module {
}
}
fn vmp_apply_dft_tmp_bytes(
&self,
res_cols: usize,
a_cols: usize,
gct_rows: usize,
gct_cols: usize,
) -> usize {
fn vmp_apply_dft_tmp_bytes(&self, res_cols: usize, a_cols: usize, gct_rows: usize, gct_cols: usize) -> usize {
unsafe {
vmp::vmp_apply_dft_tmp_bytes(
self.ptr,
@@ -540,9 +510,7 @@ impl VmpPMatOps for Module {
}
fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
debug_assert!(
tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
);
debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols()));
#[cfg(debug_assertions)]
{
assert_alignement(tmp_bytes.as_ptr());
@@ -564,9 +532,7 @@ impl VmpPMatOps for Module {
}
fn vmp_apply_dft_add(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
debug_assert!(
tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
);
debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols()));
#[cfg(debug_assertions)]
{
assert_alignement(tmp_bytes.as_ptr());
@@ -587,13 +553,7 @@ impl VmpPMatOps for Module {
}
}
fn vmp_apply_dft_to_dft_tmp_bytes(
&self,
res_cols: usize,
a_cols: usize,
gct_rows: usize,
gct_cols: usize,
) -> usize {
fn vmp_apply_dft_to_dft_tmp_bytes(&self, res_cols: usize, a_cols: usize, gct_rows: usize, gct_cols: usize) -> usize {
unsafe {
vmp::vmp_apply_dft_to_dft_tmp_bytes(
self.ptr,
@@ -605,17 +565,8 @@ impl VmpPMatOps for Module {
}
}
fn vmp_apply_dft_to_dft(
&self,
c: &mut VecZnxDft,
a: &VecZnxDft,
b: &VmpPMat,
tmp_bytes: &mut [u8],
) {
debug_assert!(
tmp_bytes.len()
>= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
);
fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, tmp_bytes: &mut [u8]) {
debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols()));
#[cfg(debug_assertions)]
{
assert_alignement(tmp_bytes.as_ptr());
@@ -635,17 +586,8 @@ impl VmpPMatOps for Module {
}
}
fn vmp_apply_dft_to_dft_add(
&self,
c: &mut VecZnxDft,
a: &VecZnxDft,
b: &VmpPMat,
tmp_bytes: &mut [u8],
) {
debug_assert!(
tmp_bytes.len()
>= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
);
fn vmp_apply_dft_to_dft_add(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, tmp_bytes: &mut [u8]) {
debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols()));
#[cfg(debug_assertions)]
{
assert_alignement(tmp_bytes.as_ptr());
@@ -666,10 +608,7 @@ impl VmpPMatOps for Module {
}
fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, tmp_bytes: &mut [u8]) {
debug_assert!(
tmp_bytes.len()
>= self.vmp_apply_dft_to_dft_tmp_bytes(b.cols(), b.cols(), a.rows(), a.cols())
);
debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_to_dft_tmp_bytes(b.cols(), b.cols(), a.rows(), a.cols()));
#[cfg(debug_assertions)]
{
assert_alignement(tmp_bytes.as_ptr());
@@ -693,8 +632,7 @@ impl VmpPMatOps for Module {
#[cfg(test)]
mod tests {
use crate::{
alloc_aligned, Module, Sampling, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps,
VecZnxOps, VmpPMat, VmpPMatOps,
Module, Sampling, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat, VmpPMatOps, alloc_aligned,
};
use sampling::source::Source;
@@ -712,8 +650,7 @@ mod tests {
let mut vmpmat_0: VmpPMat = module.new_vmp_pmat(vpmat_rows, vpmat_cols);
let mut vmpmat_1: VmpPMat = module.new_vmp_pmat(vpmat_rows, vpmat_cols);
let mut tmp_bytes: Vec<u8> =
alloc_aligned(module.vmp_prepare_tmp_bytes(vpmat_rows, vpmat_cols));
let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vmp_prepare_tmp_bytes(vpmat_rows, vpmat_cols));
for row_i in 0..vpmat_rows {
let mut source: Source = Source::new([0u8; 32]);