mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 13:16:44 +01:00
Added size and memory layout to VecZnxBig, VecZnxDft and VmpPmat
This commit is contained in:
@@ -35,7 +35,7 @@ fn main() {
|
||||
module.fill_uniform(log_base2k, &mut a, cols, &mut source);
|
||||
|
||||
// Scratch space for DFT values
|
||||
let mut buf_dft: VecZnxDft = module.new_vec_znx_dft(a.cols());
|
||||
let mut buf_dft: VecZnxDft = module.new_vec_znx_dft(1, a.cols());
|
||||
|
||||
// Applies buf_dft <- s * a
|
||||
module.svp_apply_dft(&mut buf_dft, &s_ppol, &a);
|
||||
|
||||
@@ -30,7 +30,7 @@ fn main() {
|
||||
a.print(0, a.cols(), n);
|
||||
println!();
|
||||
|
||||
let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
|
||||
let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(1, rows, cols);
|
||||
|
||||
(0..a.cols()).for_each(|row_i| {
|
||||
let mut tmp: VecZnx = module.new_vec_znx(1, cols);
|
||||
@@ -38,7 +38,7 @@ fn main() {
|
||||
module.vmp_prepare_row(&mut vmp_pmat, tmp.raw(), row_i, &mut buf);
|
||||
});
|
||||
|
||||
let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
|
||||
let mut c_dft: VecZnxDft = module.new_vec_znx_dft(1, cols);
|
||||
module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
|
||||
|
||||
let mut c_big: VecZnxBig = c_dft.as_vec_znx_big();
|
||||
|
||||
@@ -261,7 +261,7 @@ fn decode_coeff_i64(a: &VecZnx, poly_idx: usize, log_base2k: usize, log_k: usize
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::{Encoding, VecZnx};
|
||||
use crate::{Encoding, Infos, VecZnx};
|
||||
use itertools::izip;
|
||||
use sampling::source::Source;
|
||||
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use crate::LAYOUT;
|
||||
|
||||
pub trait Infos {
|
||||
/// Returns the ring degree of the receiver.
|
||||
fn n(&self) -> usize;
|
||||
@@ -5,6 +7,12 @@ pub trait Infos {
|
||||
/// Returns the base two logarithm of the ring dimension of the receiver.
|
||||
fn log_n(&self) -> usize;
|
||||
|
||||
/// Returns the number of stacked polynomials.
|
||||
fn size(&self) -> usize;
|
||||
|
||||
/// Returns the memory layout of the stacked polynomials.
|
||||
fn layout(&self) -> LAYOUT;
|
||||
|
||||
/// Returns the number of columns of the receiver.
|
||||
/// This method is equivalent to [Infos::cols].
|
||||
fn cols(&self) -> usize;
|
||||
|
||||
@@ -34,7 +34,6 @@ pub enum LAYOUT {
|
||||
COL,
|
||||
}
|
||||
|
||||
|
||||
pub fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
|
||||
(ptr as usize) % align == 0
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
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, LAYOUT};
|
||||
use crate::{BACKEND, LAYOUT, Module, VecZnx, VecZnxDft, assert_alignement};
|
||||
|
||||
use crate::{Infos, alloc_aligned, cast_mut};
|
||||
use rand::seq::SliceRandom;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::LAYOUT;
|
||||
use crate::cast_mut;
|
||||
use crate::ffi::vec_znx;
|
||||
use crate::ffi::znx;
|
||||
use crate::LAYOUT;
|
||||
use crate::{Infos, Module};
|
||||
use crate::{alloc_aligned, assert_alignement};
|
||||
use itertools::izip;
|
||||
@@ -99,11 +99,6 @@ impl VecZnx {
|
||||
self.data.len() == 0
|
||||
}
|
||||
|
||||
/// TODO: when SML refactoring is done, move this to the [Infos] trait.
|
||||
pub fn size(&self) -> usize {
|
||||
self.size
|
||||
}
|
||||
|
||||
/// Total size is [VecZnx::n()] * [VecZnx::size()] * [VecZnx::cols()].
|
||||
pub fn raw(&self) -> &[i64] {
|
||||
unsafe { std::slice::from_raw_parts(self.ptr, self.n * self.size * self.cols) }
|
||||
@@ -225,6 +220,14 @@ impl Infos for VecZnx {
|
||||
self.n
|
||||
}
|
||||
|
||||
fn size(&self) -> usize {
|
||||
self.size
|
||||
}
|
||||
|
||||
fn layout(&self) -> LAYOUT {
|
||||
self.layout
|
||||
}
|
||||
|
||||
/// Returns the number of cols of the [VecZnx].
|
||||
fn cols(&self) -> usize {
|
||||
self.cols
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
use crate::ffi::vec_znx_big::{self, vec_znx_big_t};
|
||||
use crate::{BACKEND, Infos, Module, VecZnx, VecZnxDft, alloc_aligned, assert_alignement};
|
||||
use crate::{BACKEND, Infos, LAYOUT, Module, VecZnx, VecZnxDft, alloc_aligned, assert_alignement};
|
||||
|
||||
pub struct VecZnxBig {
|
||||
pub data: Vec<u8>,
|
||||
pub ptr: *mut u8,
|
||||
pub n: usize,
|
||||
pub size: usize,
|
||||
pub cols: usize,
|
||||
pub layout: LAYOUT,
|
||||
pub backend: BACKEND,
|
||||
}
|
||||
|
||||
@@ -13,10 +15,10 @@ impl VecZnxBig {
|
||||
/// Returns a new [VecZnxBig] with the provided data as backing array.
|
||||
/// User must ensure that data is properly alligned and that
|
||||
/// the size of data is at least equal to [Module::bytes_of_vec_znx_big].
|
||||
pub fn from_bytes(module: &Module, cols: usize, bytes: &mut [u8]) -> Self {
|
||||
pub fn from_bytes(module: &Module, size: usize, cols: usize, bytes: &mut [u8]) -> Self {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_big(cols));
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_big(size, cols));
|
||||
assert_alignement(bytes.as_ptr())
|
||||
};
|
||||
unsafe {
|
||||
@@ -24,22 +26,26 @@ impl VecZnxBig {
|
||||
data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
|
||||
ptr: bytes.as_mut_ptr(),
|
||||
n: module.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
backend: module.backend,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_bytes_borrow(module: &Module, cols: usize, bytes: &mut [u8]) -> Self {
|
||||
pub fn from_bytes_borrow(module: &Module, size: usize, cols: usize, bytes: &mut [u8]) -> Self {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_big(cols));
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_big(size, cols));
|
||||
assert_alignement(bytes.as_ptr());
|
||||
}
|
||||
Self {
|
||||
data: Vec::new(),
|
||||
ptr: bytes.as_mut_ptr(),
|
||||
n: module.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
backend: module.backend,
|
||||
}
|
||||
@@ -50,6 +56,8 @@ impl VecZnxBig {
|
||||
data: Vec::new(),
|
||||
ptr: self.ptr,
|
||||
n: self.n,
|
||||
size: self.size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: self.cols,
|
||||
backend: self.backend,
|
||||
}
|
||||
@@ -81,6 +89,14 @@ impl Infos for VecZnxBig {
|
||||
self.n
|
||||
}
|
||||
|
||||
fn size(&self) -> usize {
|
||||
self.size
|
||||
}
|
||||
|
||||
fn layout(&self) -> LAYOUT {
|
||||
self.layout
|
||||
}
|
||||
|
||||
/// Returns the number of cols of the [VecZnx].
|
||||
fn cols(&self) -> usize {
|
||||
self.cols
|
||||
@@ -94,7 +110,7 @@ impl Infos for VecZnxBig {
|
||||
|
||||
pub trait VecZnxBigOps {
|
||||
/// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
|
||||
fn new_vec_znx_big(&self, cols: usize) -> VecZnxBig;
|
||||
fn new_vec_znx_big(&self, size: usize, cols: usize) -> VecZnxBig;
|
||||
|
||||
/// Returns a new [VecZnxBig] with the provided bytes array as backing array.
|
||||
///
|
||||
@@ -107,7 +123,7 @@ pub trait VecZnxBigOps {
|
||||
///
|
||||
/// # Panics
|
||||
/// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
|
||||
fn new_vec_znx_big_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxBig;
|
||||
fn new_vec_znx_big_from_bytes(&self, size: usize, cols: usize, bytes: &mut [u8]) -> VecZnxBig;
|
||||
|
||||
/// Returns a new [VecZnxBig] with the provided bytes array as backing array.
|
||||
///
|
||||
@@ -120,11 +136,11 @@ pub trait VecZnxBigOps {
|
||||
///
|
||||
/// # Panics
|
||||
/// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
|
||||
fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxBig;
|
||||
fn new_vec_znx_big_from_bytes_borrow(&self, size: usize, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxBig;
|
||||
|
||||
/// Returns the minimum number of bytes necessary to allocate
|
||||
/// a new [VecZnxBig] through [VecZnxBig::from_bytes].
|
||||
fn bytes_of_vec_znx_big(&self, cols: usize) -> usize;
|
||||
fn bytes_of_vec_znx_big(&self, size: usize, cols: usize) -> usize;
|
||||
|
||||
/// b <- b - a
|
||||
fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig, a: &VecZnx);
|
||||
@@ -162,28 +178,30 @@ pub trait VecZnxBigOps {
|
||||
}
|
||||
|
||||
impl VecZnxBigOps for Module {
|
||||
fn new_vec_znx_big(&self, cols: usize) -> VecZnxBig {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_big(cols));
|
||||
fn new_vec_znx_big(&self, size: usize, cols: usize) -> VecZnxBig {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_big(size, cols));
|
||||
let ptr: *mut u8 = data.as_mut_ptr();
|
||||
VecZnxBig {
|
||||
data: data,
|
||||
ptr: ptr,
|
||||
n: self.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
backend: self.backend(),
|
||||
}
|
||||
}
|
||||
|
||||
fn new_vec_znx_big_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxBig {
|
||||
VecZnxBig::from_bytes(self, cols, bytes)
|
||||
fn new_vec_znx_big_from_bytes(&self, size: usize, cols: usize, bytes: &mut [u8]) -> VecZnxBig {
|
||||
VecZnxBig::from_bytes(self, size, cols, bytes)
|
||||
}
|
||||
|
||||
fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxBig {
|
||||
VecZnxBig::from_bytes_borrow(self, cols, tmp_bytes)
|
||||
fn new_vec_znx_big_from_bytes_borrow(&self, size: usize, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxBig {
|
||||
VecZnxBig::from_bytes_borrow(self, size, cols, tmp_bytes)
|
||||
}
|
||||
|
||||
fn bytes_of_vec_znx_big(&self, cols: usize) -> usize {
|
||||
unsafe { vec_znx_big::bytes_of_vec_znx_big(self.ptr, cols as u64) as usize }
|
||||
fn bytes_of_vec_znx_big(&self, size: usize, cols: usize) -> usize {
|
||||
unsafe { vec_znx_big::bytes_of_vec_znx_big(self.ptr, cols as u64) as usize * size }
|
||||
}
|
||||
|
||||
fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig, a: &VecZnx) {
|
||||
|
||||
@@ -1,13 +1,15 @@
|
||||
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::{BACKEND, Infos, Module, VecZnxBig, assert_alignement};
|
||||
use crate::{BACKEND, Infos, LAYOUT, Module, VecZnxBig, assert_alignement};
|
||||
use crate::{DEFAULTALIGN, VecZnx, alloc_aligned};
|
||||
|
||||
pub struct VecZnxDft {
|
||||
pub data: Vec<u8>,
|
||||
pub ptr: *mut u8,
|
||||
pub n: usize,
|
||||
pub size: usize,
|
||||
pub layout: LAYOUT,
|
||||
pub cols: usize,
|
||||
pub backend: BACKEND,
|
||||
}
|
||||
@@ -16,10 +18,10 @@ impl VecZnxDft {
|
||||
/// Returns a new [VecZnxDft] with the provided data as backing array.
|
||||
/// User must ensure that data is properly alligned and that
|
||||
/// the size of data is at least equal to [Module::bytes_of_vec_znx_dft].
|
||||
pub fn from_bytes(module: &Module, cols: usize, bytes: &mut [u8]) -> VecZnxDft {
|
||||
pub fn from_bytes(module: &Module, size: usize, cols: usize, bytes: &mut [u8]) -> VecZnxDft {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(cols));
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(size, cols));
|
||||
assert_alignement(bytes.as_ptr())
|
||||
}
|
||||
unsafe {
|
||||
@@ -27,22 +29,26 @@ impl VecZnxDft {
|
||||
data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
|
||||
ptr: bytes.as_mut_ptr(),
|
||||
n: module.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
backend: module.backend,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_bytes_borrow(module: &Module, cols: usize, bytes: &mut [u8]) -> VecZnxDft {
|
||||
pub fn from_bytes_borrow(module: &Module, size: usize, cols: usize, bytes: &mut [u8]) -> VecZnxDft {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(cols));
|
||||
assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(size, cols));
|
||||
assert_alignement(bytes.as_ptr());
|
||||
}
|
||||
VecZnxDft {
|
||||
data: Vec::new(),
|
||||
ptr: bytes.as_mut_ptr(),
|
||||
n: module.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
backend: module.backend,
|
||||
}
|
||||
@@ -56,6 +62,8 @@ impl VecZnxDft {
|
||||
data: Vec::new(),
|
||||
ptr: self.ptr,
|
||||
n: self.n,
|
||||
layout: LAYOUT::COL,
|
||||
size: self.size,
|
||||
cols: self.cols,
|
||||
backend: self.backend,
|
||||
}
|
||||
@@ -105,6 +113,14 @@ impl Infos for VecZnxDft {
|
||||
self.n
|
||||
}
|
||||
|
||||
fn size(&self) -> usize {
|
||||
self.size
|
||||
}
|
||||
|
||||
fn layout(&self) -> LAYOUT {
|
||||
self.layout
|
||||
}
|
||||
|
||||
/// Returns the number of cols of the [VecZnx].
|
||||
fn cols(&self) -> usize {
|
||||
self.cols
|
||||
@@ -118,7 +134,7 @@ impl Infos for VecZnxDft {
|
||||
|
||||
pub trait VecZnxDftOps {
|
||||
/// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
|
||||
fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft;
|
||||
fn new_vec_znx_dft(&self, size: usize, cols: usize) -> VecZnxDft;
|
||||
|
||||
/// Returns a new [VecZnxDft] with the provided bytes array as backing array.
|
||||
///
|
||||
@@ -131,7 +147,7 @@ pub trait VecZnxDftOps {
|
||||
///
|
||||
/// # Panics
|
||||
/// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
|
||||
fn new_vec_znx_dft_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
|
||||
fn new_vec_znx_dft_from_bytes(&self, size: usize, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
|
||||
|
||||
/// Returns a new [VecZnxDft] with the provided bytes array as backing array.
|
||||
///
|
||||
@@ -144,7 +160,7 @@ pub trait VecZnxDftOps {
|
||||
///
|
||||
/// # Panics
|
||||
/// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
|
||||
fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
|
||||
fn new_vec_znx_dft_from_bytes_borrow(&self, size: usize, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
|
||||
|
||||
/// Returns a new [VecZnxDft] with the provided bytes array as backing array.
|
||||
///
|
||||
@@ -155,7 +171,7 @@ pub trait VecZnxDftOps {
|
||||
///
|
||||
/// # Panics
|
||||
/// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
|
||||
fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize;
|
||||
fn bytes_of_vec_znx_dft(&self, size: usize, cols: usize) -> usize;
|
||||
|
||||
/// Returns the minimum number of bytes necessary to allocate
|
||||
/// a new [VecZnxDft] through [VecZnxDft::from_bytes].
|
||||
@@ -176,28 +192,30 @@ pub trait VecZnxDftOps {
|
||||
}
|
||||
|
||||
impl VecZnxDftOps for Module {
|
||||
fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_dft(cols));
|
||||
fn new_vec_znx_dft(&self, size: usize, cols: usize) -> VecZnxDft {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_dft(size, cols));
|
||||
let ptr: *mut u8 = data.as_mut_ptr();
|
||||
VecZnxDft {
|
||||
data: data,
|
||||
ptr: ptr,
|
||||
n: self.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
backend: self.backend(),
|
||||
}
|
||||
}
|
||||
|
||||
fn new_vec_znx_dft_from_bytes(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
|
||||
VecZnxDft::from_bytes(self, cols, tmp_bytes)
|
||||
fn new_vec_znx_dft_from_bytes(&self, size: usize, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
|
||||
VecZnxDft::from_bytes(self, size, cols, tmp_bytes)
|
||||
}
|
||||
|
||||
fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
|
||||
VecZnxDft::from_bytes_borrow(self, cols, tmp_bytes)
|
||||
fn new_vec_znx_dft_from_bytes_borrow(&self, size: usize, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
|
||||
VecZnxDft::from_bytes_borrow(self, size, cols, tmp_bytes)
|
||||
}
|
||||
|
||||
fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize {
|
||||
unsafe { bytes_of_vec_znx_dft(self.ptr, cols as u64) as usize }
|
||||
fn bytes_of_vec_znx_dft(&self, size: usize, cols: usize) -> usize {
|
||||
unsafe { bytes_of_vec_znx_dft(self.ptr, cols as u64) as usize * size }
|
||||
}
|
||||
|
||||
fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft) {
|
||||
@@ -318,8 +336,8 @@ mod tests {
|
||||
let cols: usize = 2;
|
||||
let log_base2k: usize = 17;
|
||||
let mut a: VecZnx = module.new_vec_znx(1, cols);
|
||||
let mut a_dft: VecZnxDft = module.new_vec_znx_dft(cols);
|
||||
let mut b_dft: VecZnxDft = module.new_vec_znx_dft(cols);
|
||||
let mut a_dft: VecZnxDft = module.new_vec_znx_dft(1, cols);
|
||||
let mut b_dft: VecZnxDft = module.new_vec_znx_dft(1, cols);
|
||||
|
||||
let mut source: Source = Source::new(new_seed());
|
||||
module.fill_uniform(log_base2k, &mut a, cols, &mut source);
|
||||
|
||||
@@ -1,7 +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::{BACKEND, Infos, Module, VecZnx, VecZnxBig, VecZnxDft, alloc_aligned, assert_alignement};
|
||||
use crate::{BACKEND, Infos, LAYOUT, 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.
|
||||
@@ -23,8 +23,11 @@ pub struct VmpPMat {
|
||||
cols: usize,
|
||||
/// The ring degree of each [VecZnxDft].
|
||||
n: usize,
|
||||
|
||||
#[warn(dead_code)]
|
||||
/// The number of stacked [VmpPMat], must be a square.
|
||||
size: usize,
|
||||
/// The memory layout of the stacked [VmpPMat].
|
||||
layout: LAYOUT,
|
||||
/// The backend fft or ntt.
|
||||
backend: BACKEND,
|
||||
}
|
||||
|
||||
@@ -38,6 +41,14 @@ impl Infos for VmpPMat {
|
||||
(usize::BITS - (self.n() - 1).leading_zeros()) as _
|
||||
}
|
||||
|
||||
fn size(&self) -> usize {
|
||||
self.size
|
||||
}
|
||||
|
||||
fn layout(&self) -> LAYOUT {
|
||||
self.layout
|
||||
}
|
||||
|
||||
/// Returns the number of rows (i.e. of [VecZnxDft]) of the [VmpPMat]
|
||||
fn rows(&self) -> usize {
|
||||
self.rows
|
||||
@@ -120,12 +131,16 @@ impl VmpPMat {
|
||||
&self.raw::<T>()[blk * nrows * ncols * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
|
||||
}
|
||||
}
|
||||
|
||||
fn backend(&self) -> BACKEND {
|
||||
self.backend
|
||||
}
|
||||
}
|
||||
|
||||
/// This trait implements methods for vector matrix product,
|
||||
/// that is, multiplying a [VecZnx] with a [VmpPMat].
|
||||
pub trait VmpPMatOps {
|
||||
fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize) -> usize;
|
||||
fn bytes_of_vmp_pmat(&self, size: usize, rows: usize, cols: usize) -> usize;
|
||||
|
||||
/// Allocates a new [VmpPMat] with the given number of rows and columns.
|
||||
///
|
||||
@@ -133,7 +148,7 @@ pub trait VmpPMatOps {
|
||||
///
|
||||
/// * `rows`: number of rows (number of [VecZnxDft]).
|
||||
/// * `cols`: number of cols (number of cols of each [VecZnxDft]).
|
||||
fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat;
|
||||
fn new_vmp_pmat(&self, size: usize, rows: usize, cols: usize) -> VmpPMat;
|
||||
|
||||
/// Returns the number of bytes needed as scratch space for [VmpPMatOps::vmp_prepare_contiguous].
|
||||
///
|
||||
@@ -360,17 +375,19 @@ pub trait VmpPMatOps {
|
||||
}
|
||||
|
||||
impl VmpPMatOps for Module {
|
||||
fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize) -> usize {
|
||||
unsafe { vmp::bytes_of_vmp_pmat(self.ptr, rows as u64, cols as u64) as usize }
|
||||
fn bytes_of_vmp_pmat(&self, size: usize, rows: usize, cols: usize) -> usize {
|
||||
unsafe { vmp::bytes_of_vmp_pmat(self.ptr, rows as u64, cols as u64) as usize * size }
|
||||
}
|
||||
|
||||
fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vmp_pmat(rows, cols));
|
||||
fn new_vmp_pmat(&self, size: usize, rows: usize, cols: usize) -> VmpPMat {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vmp_pmat(size, rows, cols));
|
||||
let ptr: *mut u8 = data.as_mut_ptr();
|
||||
VmpPMat {
|
||||
data: data,
|
||||
ptr: ptr,
|
||||
n: self.n(),
|
||||
size: size,
|
||||
layout: LAYOUT::COL,
|
||||
cols: cols,
|
||||
rows: rows,
|
||||
backend: self.backend(),
|
||||
@@ -643,12 +660,12 @@ mod tests {
|
||||
let vpmat_cols: usize = 5;
|
||||
let log_base2k: usize = 8;
|
||||
let mut a: VecZnx = module.new_vec_znx(1, vpmat_cols);
|
||||
let mut a_dft: VecZnxDft = module.new_vec_znx_dft(vpmat_cols);
|
||||
let mut a_big: VecZnxBig = module.new_vec_znx_big(vpmat_cols);
|
||||
let mut b_big: VecZnxBig = module.new_vec_znx_big(vpmat_cols);
|
||||
let mut b_dft: VecZnxDft = module.new_vec_znx_dft(vpmat_cols);
|
||||
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 a_dft: VecZnxDft = module.new_vec_znx_dft(1, vpmat_cols);
|
||||
let mut a_big: VecZnxBig = module.new_vec_znx_big(1, vpmat_cols);
|
||||
let mut b_big: VecZnxBig = module.new_vec_znx_big(1, vpmat_cols);
|
||||
let mut b_dft: VecZnxDft = module.new_vec_znx_dft(1, vpmat_cols);
|
||||
let mut vmpmat_0: VmpPMat = module.new_vmp_pmat(1, vpmat_rows, vpmat_cols);
|
||||
let mut vmpmat_1: VmpPMat = module.new_vmp_pmat(1, vpmat_rows, vpmat_cols);
|
||||
|
||||
let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vmp_prepare_tmp_bytes(vpmat_rows, vpmat_cols));
|
||||
|
||||
|
||||
Reference in New Issue
Block a user