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

rework as discussed

Browse Source
This commit is contained in:
Jean-Philippe Bossuat
2025-05-05 17:35:35 +02:00
parent bd105497fd
commit ffa363804b
16 changed files with 1154 additions and 1153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

649
base2k/src/mat_znx_dft_ops.rs
View File

@@ -2,11 +2,11 @@ use crate::ffi::vec_znx_dft::vec_znx_dft_t;
use crate::ffi::vmp;
use crate::znx_base::{ZnxInfos, ZnxView, ZnxViewMut};
use crate::{
Backend, FFT64, MatZnxDft, MatZnxDftAllocOwned, Module, Scratch, VecZnx, VecZnxBigOps, VecZnxBigScratch, VecZnxDft,
VecZnxDftAlloc, VecZnxDftOps,
Backend, FFT64, MatZnxDft, MatZnxDftAllocOwned, MatZnxDftToMut, MatZnxDftToRef, Module, Scratch, VecZnxDft, VecZnxDftToMut,
VecZnxDftToRef,
};
pub trait MatZnxDftAlloc<B> {
pub trait MatZnxDftAlloc<B: Backend> {
/// Allocates a new [MatZnxDft] with the given number of rows and columns.
///
/// # Arguments
@@ -28,43 +28,10 @@ pub trait MatZnxDftAlloc<B> {
}
pub trait MatZnxDftScratch {
/// Returns the of bytes needed as scratch space for [MatZnxDftOps::vmp_prepare_row]
fn vmp_prepare_row_tmp_bytes(&self, cols_out: usize, size: usize) -> usize;
/// Returns the of bytes needed as scratch space for [MatZnxDftOps::vmp_extract_row]
fn vmp_extract_row_tmp_bytes(&self, cols_out: usize, size: usize) -> usize;
/// Returns the size of the stratch space necessary for [MatZnxDftOps::vmp_apply_dft].
///
/// # Arguments
///
/// * `c_size`: number of size of the output [VecZnxDft].
/// * `a_size`: number of size of the input [VecZnx].
/// * `rows`: number of rows of the input [MatZnxDft].
/// * `size`: number of size of the input [MatZnxDft].
fn vmp_apply_dft_tmp_bytes(
&self,
c_size: usize,
a_size: usize,
b_rows: usize,
b_cols_in: usize,
b_cols_out: usize,
b_size: usize,
) -> usize;
/// Returns the size of the stratch space necessary for [MatZnxDftOps::vmp_apply_dft_to_dft].
///
/// # Arguments
///
/// * `c_size`: number of size of the output [VecZnxDft].
/// * `a_size`: number of size of the input [VecZnxDft].
/// * `rows`: number of rows of the input [MatZnxDft].
/// * `size`: number of size of the input [MatZnxDft].
fn vmp_apply_dft_to_dft_tmp_bytes(
fn vmp_apply_tmp_bytes(
&self,
c_cols: usize,
c_size: usize,
a_cols: usize,
res_size: usize,
a_size: usize,
b_rows: usize,
b_cols_in: usize,
@@ -75,43 +42,7 @@ pub trait MatZnxDftScratch {
/// This trait implements methods for vector matrix product,
/// that is, multiplying a [VecZnx] with a [MatZnxDft].
pub trait MatZnxDftOps<DataMut, Data, B: Backend> {
/// Prepares the ith-row of [MatZnxDft] from a [VecZnx].
///
/// # Arguments
///
/// * `b`: [MatZnxDft] on which the values are encoded.
/// * `row_i`: the row of the [MatZnxDft] to prepare.
/// * `a`: the [VecZnx] to encode on the i-th row of the [MatZnxDft].
/// * `buf`: scratch space, the size of buf can be obtained with [MatZnxDftOps::vmp_prepare_tmp_bytes].
///
/// The size of buf can be obtained with [MatZnxDftOps::vmp_prepare_tmp_bytes].
fn vmp_prepare_row(
&self,
b: &mut MatZnxDft<DataMut, B>,
b_row: usize,
b_col_in: usize,
a: &VecZnx<Data>,
scratch: &mut Scratch,
);
/// Extracts the ith-row of [MatZnxDft] into a [VecZnxBig].
///
/// # Arguments
///
/// * `b`: the [VecZnxBig] to on which to extract the row of the [MatZnxDft].
/// * `a`: [MatZnxDft] on which the values are encoded.
/// * `row_i`: the index of the row to extract.
fn vmp_extract_row(
&self,
log_base2k: usize,
b: &mut VecZnx<DataMut>,
a: &MatZnxDft<Data, B>,
b_row: usize,
b_col_in: usize,
scratch: &mut Scratch,
);
pub trait MatZnxDftOps<BACKEND: Backend> {
/// Prepares the ith-row of [MatZnxDft] from a [VecZnxDft].
///
/// # Arguments
@@ -121,7 +52,10 @@ pub trait MatZnxDftOps<DataMut, Data, B: Backend> {
/// * `row_i`: the index of the row to prepare.
///
/// The size of buf can be obtained with [MatZnxDftOps::vmp_prepare_tmp_bytes].
fn vmp_prepare_row_dft(&self, b: &mut MatZnxDft<DataMut, B>, b_row: usize, b_col_in: usize, a: &VecZnxDft<Data, B>);
fn vmp_prepare_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
where
R: MatZnxDftToMut<BACKEND>,
A: VecZnxDftToRef<BACKEND>;
/// Extracts the ith-row of [MatZnxDft] into a [VecZnxDft].
///
@@ -130,33 +64,10 @@ pub trait MatZnxDftOps<DataMut, Data, B: Backend> {
/// * `b`: the [VecZnxDft] to on which to extract the row of the [MatZnxDft].
/// * `a`: [MatZnxDft] on which the values are encoded.
/// * `row_i`: the index of the row to extract.
fn vmp_extract_row_dft(&self, b: &mut VecZnxDft<DataMut, B>, a: &MatZnxDft<Data, B>, a_row: usize, a_col_in: usize);
/// Applies the vector matrix product [VecZnxDft] x [MatZnxDft].
///
/// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
/// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
/// and each vector a [VecZnxDft] (row) of the [MatZnxDft].
///
/// As such, given an input [VecZnx] of `i` size and a [MatZnxDft] of `i` rows and
/// `j` size, the output is a [VecZnx] of `j` size.
///
/// If there is a mismatch between the dimensions the largest valid ones are used.
///
/// ```text
/// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
/// |h i j|
/// |k l m|
/// ```
/// where each element is a [VecZnxDft].
///
/// # Arguments
///
/// * `c`: the output of the vector matrix product, as a [VecZnxDft].
/// * `a`: the left operand [VecZnx] of the vector matrix product.
/// * `b`: the right operand [MatZnxDft] of the vector matrix product.
/// * `buf`: scratch space, the size can be obtained with [MatZnxDftOps::vmp_apply_dft_tmp_bytes].
fn vmp_apply_dft(&self, c: &mut VecZnxDft<DataMut, B>, a: &VecZnx<Data>, b: &MatZnxDft<Data, B>, scratch: &mut Scratch);
fn vmp_extract_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
where
R: VecZnxDftToMut<BACKEND>,
A: MatZnxDftToRef<BACKEND>;
/// Applies the vector matrix product [VecZnxDft] x [MatZnxDft].
/// The size of `buf` is given by [MatZnxDftOps::vmp_apply_dft_to_dft_tmp_bytes].
@@ -183,13 +94,11 @@ pub trait MatZnxDftOps<DataMut, Data, B: Backend> {
/// * `a`: the left operand [VecZnxDft] of the vector matrix product.
/// * `b`: the right operand [MatZnxDft] of the vector matrix product.
/// * `buf`: scratch space, the size can be obtained with [MatZnxDftOps::vmp_apply_dft_to_dft_tmp_bytes].
fn vmp_apply_dft_to_dft(
&self,
c: &mut VecZnxDft<DataMut, B>,
a: &VecZnxDft<Data, B>,
b: &MatZnxDft<Data, B>,
scratch: &mut Scratch,
);
fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
where
R: VecZnxDftToMut<BACKEND>,
A: VecZnxDftToRef<BACKEND>,
B: MatZnxDftToRef<BACKEND>;
}
impl<B: Backend> MatZnxDftAlloc<B> for Module<B> {
@@ -213,40 +122,10 @@ impl<B: Backend> MatZnxDftAlloc<B> for Module<B> {
}
}
impl<B: Backend> MatZnxDftScratch for Module<B> {
fn vmp_prepare_row_tmp_bytes(&self, cols_out: usize, size: usize) -> usize {
<Self as VecZnxDftAlloc<_>>::bytes_of_vec_znx_dft(self, cols_out, size)
}
fn vmp_extract_row_tmp_bytes(&self, cols_out: usize, size: usize) -> usize {
<Self as VecZnxDftAlloc<_>>::bytes_of_vec_znx_dft(self, cols_out, size)
+ <Self as VecZnxBigScratch>::vec_znx_big_normalize_tmp_bytes(self)
}
fn vmp_apply_dft_tmp_bytes(
impl<BACKEND: Backend> MatZnxDftScratch for Module<BACKEND> {
fn vmp_apply_tmp_bytes(
&self,
c_size: usize,
a_size: usize,
b_rows: usize,
b_cols_in: usize,
b_cols_out: usize,
b_size: usize,
) -> usize {
unsafe {
vmp::vmp_apply_dft_tmp_bytes(
self.ptr,
c_size as u64,
a_size as u64,
(b_rows * b_cols_in) as u64,
(b_size * b_cols_out) as u64,
) as usize
}
}
fn vmp_apply_dft_to_dft_tmp_bytes(
&self,
c_cols: usize,
c_size: usize,
a_cols: usize,
res_size: usize,
a_size: usize,
b_rows: usize,
b_cols_in: usize,
@@ -256,8 +135,8 @@ impl<B: Backend> MatZnxDftScratch for Module<B> {
unsafe {
vmp::vmp_apply_dft_to_dft_tmp_bytes(
self.ptr,
(c_size * c_cols) as u64,
(a_size * a_cols) as u64,
(res_size * b_cols_out) as u64,
(a_size * b_cols_in) as u64,
(b_rows * b_cols_in) as u64,
(b_size * b_cols_out) as u64,
) as usize
@@ -265,152 +144,43 @@ impl<B: Backend> MatZnxDftScratch for Module<B> {
}
}
impl MatZnxDftOps<&mut [u8], &[u8], FFT64> for Module<FFT64> {
fn vmp_prepare_row(
&self,
b: &mut MatZnxDft<&mut [u8], FFT64>,
b_row: usize,
b_col_in: usize,
a: &VecZnx<&[u8]>,
scratch: &mut Scratch,
) {
impl MatZnxDftOps<FFT64> for Module<FFT64> {
fn vmp_prepare_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
where
R: MatZnxDftToMut<FFT64>,
A: VecZnxDftToRef<FFT64>,
{
let mut res: MatZnxDft<&mut [u8], _> = res.to_mut();
let a: VecZnxDft<&[u8], _> = a.to_ref();
#[cfg(debug_assertions)]
{
assert_eq!(b.n(), self.n());
assert_eq!(res.n(), self.n());
assert_eq!(a.n(), self.n());
assert_eq!(
a.cols(),
b.cols_out(),
"a.cols(): {} != b.cols_out(): {}",
res.cols_out(),
"a.cols(): {} != res.cols_out(): {}",
a.cols(),
b.cols_out()
res.cols_out()
);
assert!(
b_row < b.rows(),
"b_row: {} >= b.rows(): {}",
b_row,
b.rows()
res_row < res.rows(),
"res_row: {} >= res.rows(): {}",
res_row,
res.rows()
);
assert!(
b_col_in < b.cols_in(),
"b_col_in: {} >= b.cols_in(): {}",
b_col_in,
b.cols_in()
res_col_in < res.cols_in(),
"res_col_in: {} >= res.cols_in(): {}",
res_col_in,
res.cols_in()
);
assert_eq!(
b.size(),
res.size(),
a.size(),
"b.size(): {} != a.size(): {}",
b.size(),
a.size()
);
// assert!(
// tmp_bytes.len()
// >= <Self as MatZnxDftOps<DataMut, Data, FFT64>>::vmp_prepare_row_tmp_bytes(self, a.cols(), a.size())
// );
// assert!(is_aligned(tmp_bytes.as_ptr()))
}
let cols_out: usize = a.cols();
let a_size: usize = a.size();
// let (tmp_bytes_a_dft, _) = tmp_bytes.split_at_mut(self.bytes_of_vec_znx_dft(cols_out, a_size));
let (mut a_dft, _) = scratch.tmp_vec_znx_dft::<_>(self, cols_out, a_size);
(0..cols_out).for_each(|i| self.vec_znx_dft(&mut a_dft, i, &a, i));
Self::vmp_prepare_row_dft(&self, b, b_row, b_col_in, &a_dft.to_ref());
}
fn vmp_extract_row(
&self,
log_base2k: usize,
b: &mut VecZnx<&mut [u8]>,
a: &MatZnxDft<&[u8], FFT64>,
a_row: usize,
a_col_in: usize,
scratch: &mut Scratch,
) {
#[cfg(debug_assertions)]
{
assert_eq!(b.n(), self.n());
assert_eq!(a.n(), self.n());
assert_eq!(
b.cols(),
a.cols_out(),
"b.cols(): {} != a.cols_out(): {}",
b.cols(),
a.cols_out()
);
assert!(
a_row < a.rows(),
"a_row: {} >= a.rows(): {}",
a_row,
a.rows()
);
assert!(
a_col_in < a.cols_in(),
"a_col_in: {} >= a.cols_in(): {}",
a_col_in,
a.cols_in()
);
assert_eq!(
b.size(),
a.size(),
"b.size(): {} != a.size(): {}",
b.size(),
a.size()
);
// assert!(tmp_bytes.len() >= self.vmp_extract_row_tmp_bytes(a.cols(), a.size()));
// assert!(is_aligned(tmp_bytes.as_ptr()))
}
let cols_out: usize = b.cols();
let size: usize = b.size();
// let (bytes_a_dft, tmp_bytes) = tmp_bytes.split_at_mut(self.bytes_of_vec_znx_dft(cols_out, size));
let (mut b_dft, scratch) = scratch.tmp_vec_znx_dft(self, cols_out, size);
Self::vmp_extract_row_dft(&self, &mut b_dft, a, a_row, a_col_in);
let (mut b_big, scratch) = scratch.tmp_vec_znx_big(self, cols_out, size);
(0..cols_out).for_each(|i| {
<Self as VecZnxDftOps<&mut [u8], &[u8], FFT64>>::vec_znx_idft_tmp_a(self, &mut b_big, i, &mut b_dft, i);
self.vec_znx_big_normalize(log_base2k, b, i, &b_big, i, scratch);
});
}
fn vmp_prepare_row_dft(
&self,
b: &mut MatZnxDft<&mut [u8], FFT64>,
b_row: usize,
b_col_in: usize,
a: &VecZnxDft<&[u8], FFT64>,
) {
#[cfg(debug_assertions)]
{
assert_eq!(b.n(), self.n());
assert_eq!(a.n(), self.n());
assert_eq!(
a.cols(),
b.cols_out(),
"a.cols(): {} != b.cols_out(): {}",
a.cols(),
b.cols_out()
);
assert!(
b_row < b.rows(),
"b_row: {} >= b.rows(): {}",
b_row,
b.rows()
);
assert!(
b_col_in < b.cols_in(),
"b_col_in: {} >= b.cols_in(): {}",
b_col_in,
b.cols_in()
);
assert_eq!(
b.size(),
a.size(),
"b.size(): {} != a.size(): {}",
b.size(),
"res.size(): {} != a.size(): {}",
res.size(),
a.size()
);
}
@@ -418,31 +188,32 @@ impl MatZnxDftOps<&mut [u8], &[u8], FFT64> for Module<FFT64> {
unsafe {
vmp::vmp_prepare_row_dft(
self.ptr,
b.as_mut_ptr() as *mut vmp::vmp_pmat_t,
res.as_mut_ptr() as *mut vmp::vmp_pmat_t,
a.as_ptr() as *const vec_znx_dft_t,
(b_row * b.cols_in() + b_col_in) as u64,
(b.rows() * b.cols_in()) as u64,
(b.size() * b.cols_out()) as u64,
(res_row * res.cols_in() + res_col_in) as u64,
(res.rows() * res.cols_in()) as u64,
(res.size() * res.cols_out()) as u64,
);
}
}
fn vmp_extract_row_dft(
&self,
b: &mut VecZnxDft<&mut [u8], FFT64>,
a: &MatZnxDft<&[u8], FFT64>,
a_row: usize,
a_col_in: usize,
) {
fn vmp_extract_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
where
R: VecZnxDftToMut<FFT64>,
A: MatZnxDftToRef<FFT64>,
{
let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
let a: MatZnxDft<&[u8], _> = a.to_ref();
#[cfg(debug_assertions)]
{
assert_eq!(b.n(), self.n());
assert_eq!(res.n(), self.n());
assert_eq!(a.n(), self.n());
assert_eq!(
b.cols(),
res.cols(),
a.cols_out(),
"b.cols(): {} != a.cols_out(): {}",
b.cols(),
"res.cols(): {} != a.cols_out(): {}",
res.cols(),
a.cols_out()
);
assert!(
@@ -458,17 +229,17 @@ impl MatZnxDftOps<&mut [u8], &[u8], FFT64> for Module<FFT64> {
a.cols_in()
);
assert_eq!(
b.size(),
res.size(),
a.size(),
"b.size(): {} != a.size(): {}",
b.size(),
"res.size(): {} != a.size(): {}",
res.size(),
a.size()
);
}
unsafe {
vmp::vmp_extract_row_dft(
self.ptr,
b.as_mut_ptr() as *mut vec_znx_dft_t,
res.as_mut_ptr() as *mut vec_znx_dft_t,
a.as_ptr() as *const vmp::vmp_pmat_t,
(a_row * a.cols_in() + a_col_in) as u64,
(a.rows() * a.cols_in()) as u64,
@@ -477,23 +248,26 @@ impl MatZnxDftOps<&mut [u8], &[u8], FFT64> for Module<FFT64> {
}
}
fn vmp_apply_dft(
&self,
c: &mut VecZnxDft<&mut [u8], FFT64>,
a: &VecZnx<&[u8]>,
b: &MatZnxDft<&[u8], FFT64>,
scratch: &mut Scratch,
) {
fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
where
R: VecZnxDftToMut<FFT64>,
A: VecZnxDftToRef<FFT64>,
B: MatZnxDftToRef<FFT64>,
{
let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
let a: VecZnxDft<&[u8], _> = a.to_ref();
let b: MatZnxDft<&[u8], _> = b.to_ref();
#[cfg(debug_assertions)]
{
assert_eq!(c.n(), self.n());
assert_eq!(res.n(), self.n());
assert_eq!(b.n(), self.n());
assert_eq!(a.n(), self.n());
assert_eq!(
c.cols(),
res.cols(),
b.cols_out(),
"c.cols(): {} != b.cols_out: {}",
c.cols(),
"res.cols(): {} != b.cols_out: {}",
res.cols(),
b.cols_out()
);
assert_eq!(
@@ -503,37 +277,23 @@ impl MatZnxDftOps<&mut [u8], &[u8], FFT64> for Module<FFT64> {
a.cols(),
b.cols_in()
);
// assert!(
// tmp_bytes.len()
// >= self.vmp_apply_dft_tmp_bytes(
// c.size(),
// a.size(),
// b.rows(),
// b.cols_in(),
// b.cols_out(),
// b.size()
// )
// );
// assert_alignement(tmp_bytes.as_ptr());
}
let (tmp_bytes, _) = scratch.tmp_scalar_slice(<Self as MatZnxDftScratch>::vmp_apply_dft_tmp_bytes(
self,
c.size(),
let (tmp_bytes, _) = scratch.tmp_scalar_slice(self.vmp_apply_tmp_bytes(
res.size(),
a.size(),
b.rows(),
b.cols_in(),
b.cols_out(),
b.size(),
));
unsafe {
vmp::vmp_apply_dft(
vmp::vmp_apply_dft_to_dft(
self.ptr,
c.as_mut_ptr() as *mut vec_znx_dft_t,
(c.size() * c.cols()) as u64,
a.as_ptr(),
res.as_mut_ptr() as *mut vec_znx_dft_t,
(res.size() * res.cols()) as u64,
a.as_ptr() as *const vec_znx_dft_t,
(a.size() * a.cols()) as u64,
a.n() as u64,
b.as_ptr() as *const vmp::vmp_pmat_t,
(b.rows() * b.cols_in()) as u64,
(b.size() * b.cols_out()) as u64,
@@ -541,164 +301,131 @@ impl MatZnxDftOps<&mut [u8], &[u8], FFT64> for Module<FFT64> {
)
}
}
fn vmp_apply_dft_to_dft(
&self,
c: &mut VecZnxDft<&mut [u8], FFT64>,
a: &VecZnxDft<&[u8], FFT64>,
b: &MatZnxDft<&[u8], FFT64>,
scratch: &mut Scratch,
) {
{
#[cfg(debug_assertions)]
{
assert_eq!(c.n(), self.n());
assert_eq!(b.n(), self.n());
assert_eq!(a.n(), self.n());
assert_eq!(
c.cols(),
b.cols_out(),
"c.cols(): {} != b.cols_out: {}",
c.cols(),
b.cols_out()
);
assert_eq!(
a.cols(),
b.cols_in(),
"a.cols(): {} != b.cols_in: {}",
a.cols(),
b.cols_in()
);
// assert!(
// tmp_bytes.len()
// >= self.vmp_apply_dft_to_dft_tmp_bytes(
// c.cols(),
// c.size(),
// a.cols(),
// a.size(),
// b.rows(),
// b.cols_in(),
// b.cols_out(),
// b.size()
// )
// );
// assert_alignement(tmp_bytes.as_ptr());
}
let (tmp_bytes, _) = scratch.tmp_scalar_slice(self.vmp_apply_dft_to_dft_tmp_bytes(
c.cols(),
c.size(),
a.cols(),
a.size(),
b.rows(),
b.cols_in(),
b.cols_out(),
b.size(),
));
unsafe {
vmp::vmp_apply_dft_to_dft(
self.ptr,
c.as_mut_ptr() as *mut vec_znx_dft_t,
c.poly_count() as u64,
a.as_ptr() as *const vec_znx_dft_t,
a.poly_count() as u64,
b.as_ptr() as *const vmp::vmp_pmat_t,
b.rows() as u64,
(b.size() * b.cols()) as u64,
tmp_bytes.as_mut_ptr(),
)
}
}
}
}
#[cfg(test)]
mod tests {
use crate::ScratchOwned;
use crate::mat_znx_dft_ops::*;
use crate::vec_znx_big_ops::*;
use crate::vec_znx_dft_ops::*;
use crate::vec_znx_ops::*;
use crate::{
FFT64, MatZnxDft, MatZnxDftOps, Module, Sampling, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, alloc_aligned,
Encoding, FFT64, MatZnxDft, MatZnxDftOps, Module, Sampling, ScratchOwned, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc,
VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, ZnxInfos, ZnxView, ZnxViewMut,
};
use sampling::source::Source;
use super::{MatZnxDftAlloc, MatZnxDftScratch};
#[test]
fn vmp_prepare_row_dft() {
fn vmp_prepare_row() {
let module: Module<FFT64> = Module::<FFT64>::new(16);
let log_base2k: usize = 8;
let mat_rows: usize = 4;
let mat_cols_in: usize = 2;
let mat_cols_out: usize = 2;
let mat_size: usize = 5;
let mut a: VecZnx<_> = module.new_vec_znx(mat_cols_out, mat_size);
let mut b: VecZnx<_> = module.new_vec_znx(mat_cols_out, mat_size);
let mut a_dft: VecZnxDft<_, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
let mut a_big: VecZnxBig<_, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
let mut b_dft: VecZnxDft<_, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
let mut vmpmat_0: MatZnxDft<_, FFT64> = module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
let mut vmpmat_1: MatZnxDft<_, FFT64> = module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
// let mut tmp_bytes: Vec<u8> =
// alloc_aligned(module.vmp_prepare_row_tmp_bytes(mat_cols_out, mat_size) | module.vec_znx_big_normalize_tmp_bytes());
let mut scratch = ScratchOwned::new(
2 * (module.vmp_prepare_row_tmp_bytes(mat_cols_out, mat_size) + module.vec_znx_big_normalize_tmp_bytes()),
);
let mut tmp_bytes: Vec<u8> =
alloc_aligned::<u8>(<Module<FFT64> as VecZnxDftOps<Vec<u8>, Vec<u8>, _>>::vec_znx_idft_tmp_bytes(&module));
let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
let mut b_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
let mut mat: MatZnxDft<Vec<u8>, FFT64> = module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
for col_in in 0..mat_cols_in {
for row_i in 0..mat_rows {
let mut source: Source = Source::new([0u8; 32]);
(0..mat_cols_out).for_each(|col_out| {
module.fill_uniform(log_base2k, &mut a, col_out, mat_size, &mut source);
module.vec_znx_dft(&mut a_dft, col_out, &a, col_out);
});
module.vmp_prepare_row(
&mut vmpmat_0.to_mut(),
row_i,
col_in,
&a.to_ref(),
scratch.borrow(),
);
// Checks that prepare(mat_znx_dft, a) = prepare_dft(mat_znx_dft, a_dft)
module.vmp_prepare_row_dft(&mut vmpmat_1.to_mut(), row_i, col_in, &a_dft.to_ref());
assert_eq!(vmpmat_0.raw(), vmpmat_1.raw());
// Checks that a_dft = extract_dft(prepare(mat_znx_dft, a), b_dft)
module.vmp_extract_row_dft(&mut b_dft.to_mut(), &vmpmat_0.to_ref(), row_i, col_in);
module.vmp_prepare_row(&mut mat, row_i, col_in, &a_dft);
module.vmp_extract_row(&mut b_dft, &mat, row_i, col_in);
assert_eq!(a_dft.raw(), b_dft.raw());
// Checks that a_big = extract(prepare_dft(mat_znx_dft, a_dft), b_big)
module.vmp_extract_row(
log_base2k,
&mut b.to_mut(),
&vmpmat_0.to_ref(),
row_i,
col_in,
scratch.borrow(),
);
(0..mat_cols_out).for_each(|col_out| {
module.vec_znx_idft(&mut a_big, col_out, &a_dft, col_out, &mut tmp_bytes);
module.vec_znx_big_normalize(
log_base2k,
&mut a.to_mut(),
col_out,
&a_big.to_ref(),
col_out,
scratch.borrow(),
);
});
assert_eq!(a.raw(), b.raw());
}
}
module.free();
}
#[test]
fn vmp_apply() {
let log_n: i32 = 5;
let n: usize = 1 << log_n;
let module: Module<FFT64> = Module::<FFT64>::new(n);
let log_base2k: usize = 15;
let a_size: usize = 5;
let mat_size: usize = 6;
let res_size: usize = 5;
[1, 2].iter().for_each(|in_cols| {
[1, 2].iter().for_each(|out_cols| {
let a_cols: usize = *in_cols;
let res_cols: usize = *out_cols;
let mat_rows: usize = a_size;
let mat_cols_in: usize = a_cols;
let mat_cols_out: usize = res_cols;
let res_cols: usize = mat_cols_out;
let mut scratch: ScratchOwned = ScratchOwned::new(
module.vmp_apply_tmp_bytes(
res_size,
a_size,
mat_rows,
mat_cols_in,
mat_cols_out,
mat_size,
) | module.vec_znx_big_normalize_tmp_bytes(),
);
let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
(0..a_cols).for_each(|i| {
a.at_mut(i, 2)[i + 1] = 1;
});
let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
// Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
(0..a.size()).for_each(|row_i| {
(0..mat_cols_in).for_each(|col_in_i| {
(0..mat_cols_out).for_each(|col_out_i| {
let idx = 1 + col_in_i * mat_cols_out + col_out_i;
tmp.at_mut(col_out_i, row_i)[idx] = 1 as i64; // X^{idx}
module.vec_znx_dft(&mut c_dft, col_out_i, &tmp, col_out_i);
tmp.at_mut(col_out_i, row_i)[idx] = 0 as i64;
});
module.vmp_prepare_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
});
});
let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, a_size);
(0..a_cols).for_each(|i| {
module.vec_znx_dft(&mut a_dft, i, &a, i);
});
module.vmp_apply(&mut c_dft, &a_dft, &mat_znx_dft, scratch.borrow());
let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, res_size);
(0..mat_cols_out).for_each(|i| {
module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
module.vec_znx_big_normalize(log_base2k, &mut res_have, i, &c_big, i, scratch.borrow());
});
(0..mat_cols_out).for_each(|col_i| {
let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
(0..a_cols).for_each(|i| {
res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
});
res_have.decode_vec_i64(col_i, log_base2k, log_base2k * 3, &mut res_have_vi64);
assert_eq!(res_have_vi64, res_want_vi64);
});
});
});
module.free();
}
}