mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 13:16:44 +01:00
vmp & svp doc
This commit is contained in:
Jean-Philippe Bossuat
592
base2k/src/vmp.rs
Normal file
592
base2k/src/vmp.rs
Normal file
@@ -0,0 +1,592 @@
|
||||
use crate::ffi::vmp::{
|
||||
delete_vmp_pmat, new_vmp_pmat, vmp_apply_dft, vmp_apply_dft_tmp_bytes, vmp_apply_dft_to_dft,
|
||||
vmp_apply_dft_to_dft_tmp_bytes, vmp_pmat_t, vmp_prepare_contiguous,
|
||||
vmp_prepare_contiguous_tmp_bytes,
|
||||
};
|
||||
use crate::Free;
|
||||
use crate::{Module, VecZnx, VecZnxDft};
|
||||
use std::cmp::min;
|
||||
|
||||
/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
|
||||
/// stored as a 3D matrix in the DFT domain in a single contiguous array.
|
||||
/// Each row of the [VmpPMat] can be seen as a [VecZnxDft].
|
||||
///
|
||||
/// The backend array of [VmpPMat] is allocate in C,
|
||||
/// and thus must be manually freed.
|
||||
///
|
||||
/// [VmpPMat] is used to permform a vector matrix product between a [VecZnx] and a [VmpPMat].
|
||||
/// See the trait [VectorMatrixProduct] for additional information.
|
||||
pub struct VmpPMat {
|
||||
/// The pointer to the C memory.
|
||||
pub data: *mut vmp_pmat_t,
|
||||
/// The number of [VecZnxDft].
|
||||
pub rows: usize,
|
||||
/// The number of limbs in each [VecZnxDft].
|
||||
pub cols: usize,
|
||||
/// The ring degree of each [VecZnxDft].
|
||||
pub n: usize,
|
||||
}
|
||||
|
||||
impl VmpPMat {
|
||||
/// Returns the pointer to the [vmp_pmat_t].
|
||||
pub fn data(&self) -> *mut vmp_pmat_t {
|
||||
self.data
|
||||
}
|
||||
|
||||
/// Returns the number of rows of the [VmpPMat].
|
||||
/// The number of rows (i.e. of [VecZnx]) of the [VmpPMat].
|
||||
pub fn rows(&self) -> usize {
|
||||
self.rows
|
||||
}
|
||||
|
||||
/// Returns the number of cols of the [VmpPMat].
|
||||
/// The number of cols refers to the number of limbs
|
||||
/// of the prepared [VecZnx].
|
||||
pub fn cols(&self) -> usize {
|
||||
self.cols
|
||||
}
|
||||
|
||||
/// Returns the ring dimension of the [VmpPMat].
|
||||
pub fn n(&self) -> usize {
|
||||
self.n
|
||||
}
|
||||
|
||||
/// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
|
||||
/// When using [`crate::FFT64`] as backend, `T` should be [f64].
|
||||
/// When using [`crate::NTT120`] as backend, `T` should be [i64].
|
||||
pub fn at<T: Default + Copy>(&self, row: usize, col: usize) -> Vec<T> {
|
||||
let mut res: Vec<T> = vec![T::default(); self.n];
|
||||
|
||||
if self.n < 8 {
|
||||
res.copy_from_slice(
|
||||
&self.get_backend_array::<T>()[(row + col * self.rows()) * self.n()
|
||||
..(row + col * self.rows()) * (self.n() + 1)],
|
||||
);
|
||||
} else {
|
||||
(0..self.n >> 3).for_each(|blk| {
|
||||
res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.get_array(row, col, blk)[..8]);
|
||||
});
|
||||
}
|
||||
|
||||
res
|
||||
}
|
||||
|
||||
/// When using [`crate::FFT64`] as backend, `T` should be [f64].
|
||||
/// When using [`crate::NTT120`] as backend, `T` should be [i64].
|
||||
fn get_array<T>(&self, row: usize, col: usize, blk: usize) -> &[T] {
|
||||
let nrows: usize = self.rows();
|
||||
let ncols: usize = self.cols();
|
||||
if col == (ncols - 1) && (ncols & 1 == 1) {
|
||||
&self.get_backend_array::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
|
||||
} else {
|
||||
&self.get_backend_array::<T>()[blk * nrows * ncols * 8
|
||||
+ (col / 2) * (2 * nrows) * 8
|
||||
+ row * 2 * 8
|
||||
+ (col % 2) * 8..]
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a non-mutable reference of [T] of the entire contiguous array of the [VmpPMat].
|
||||
/// When using [`crate::FFT64`] as backend, `T` should be [f64].
|
||||
/// When using [`crate::NTT120`] as backend, `T` should be [i64].
|
||||
/// The length of the returned array is rows * cols * n.
|
||||
pub fn get_backend_array<T>(&self) -> &[T] {
|
||||
let ptr: *const T = self.data as *const T;
|
||||
let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
|
||||
unsafe { &std::slice::from_raw_parts(ptr, len) }
|
||||
}
|
||||
}
|
||||
|
||||
impl Free for VmpPMat {
|
||||
fn free(self) {
|
||||
unsafe { delete_vmp_pmat(self.data) };
|
||||
drop(self);
|
||||
}
|
||||
}
|
||||
|
||||
/// This trait implements methods for vector matrix product,
|
||||
/// that is, multiplying a [VecZnx] with a [VmpPMat].
|
||||
pub trait VectorMatrixProduct {
|
||||
/// Allocates a new [VmpPMat] with the given number of rows and columns.
|
||||
fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat;
|
||||
|
||||
/// Returns the number of bytes needed as scratch space for [VectorMatrixProduct::vmp_prepare_contiguous].
|
||||
fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize;
|
||||
|
||||
/// Prepares a [VmpPMat] from a contiguous array of [i64].
|
||||
/// The helper struct [Matrix3D] can be used to contruct and populate
|
||||
/// the appropriate contiguous array.
|
||||
///
|
||||
/// The size of buf can be obtained with [VectorMatrixProduct::vmp_prepare_contiguous_tmp_bytes].
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::{Module, Matrix3D, VmpPMat, FFT64, Free};
|
||||
/// use base2k::vmp::VectorMatrixProduct;
|
||||
/// use std::cmp::min;
|
||||
///
|
||||
/// let n: usize = 1024;
|
||||
/// let module = Module::new::<FFT64>(n);
|
||||
/// let rows = 5;
|
||||
/// let cols = 6;
|
||||
///
|
||||
/// let mut b_mat: Matrix3D<i64> = Matrix3D::new(rows, cols, n);
|
||||
///
|
||||
/// // Populates the i-th row of b_math with X^1 * 2^(i * log_w) (here log_w is undefined)
|
||||
/// (0..min(rows, cols)).for_each(|i| {
|
||||
/// b_mat.at_mut(i, i)[1] = 1 as i64;
|
||||
/// });
|
||||
///
|
||||
/// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_contiguous_tmp_bytes(rows, cols)];
|
||||
///
|
||||
/// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
|
||||
/// module.vmp_prepare_contiguous(&mut vmp_pmat, &b_mat.data, &mut buf);
|
||||
///
|
||||
/// vmp_pmat.free() // don't forget to free the memory once vmp_pmat is not needed anymore.
|
||||
/// ```
|
||||
fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]);
|
||||
|
||||
/// Prepares a [VmpPMat] from a vector of [VecZnx].
|
||||
///
|
||||
/// The size of buf can be obtained with [VectorMatrixProduct::vmp_prepare_contiguous_tmp_bytes].
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::{Module, FFT64, Matrix3D, VmpPMat, VecZnx, Free};
|
||||
/// use base2k::vmp::VectorMatrixProduct;
|
||||
/// use std::cmp::min;
|
||||
///
|
||||
/// let n: usize = 1024;
|
||||
/// let module: Module = Module::new::<FFT64>(n);
|
||||
/// let rows: usize = 5;
|
||||
/// let cols: usize = 6;
|
||||
///
|
||||
/// let mut vecznx: Vec<VecZnx>= Vec::new();
|
||||
/// (0..rows).for_each(|_|{
|
||||
/// vecznx.push(module.new_vec_znx(cols));
|
||||
/// });
|
||||
///
|
||||
/// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_contiguous_tmp_bytes(rows, cols)];
|
||||
///
|
||||
/// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
|
||||
/// module.vmp_prepare_dblptr(&mut vmp_pmat, &vecznx, &mut buf);
|
||||
///
|
||||
/// vmp_pmat.free();
|
||||
/// module.free();
|
||||
/// ```
|
||||
fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]);
|
||||
|
||||
/// Returns the size of the stratch space necessary for [VectorMatrixProduct::vmp_apply_dft].
|
||||
fn vmp_apply_dft_tmp_bytes(
|
||||
&self,
|
||||
c_limbs: usize,
|
||||
a_limbs: usize,
|
||||
rows: usize,
|
||||
cols: usize,
|
||||
) -> usize;
|
||||
|
||||
/// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
|
||||
/// The size of `buf` is given by [VectorMatrixProduct::vmp_apply_dft_to_dft_tmp_bytes].
|
||||
///
|
||||
/// A vector matrix product is equivalent to a sum of [ScalarVectorProduct::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 [VmpPMat].
|
||||
///
|
||||
/// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
|
||||
/// `j` cols, the output is a [VecZnx] of `j` limbs.
|
||||
///
|
||||
/// 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].
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, FFT64, Free};
|
||||
/// use base2k::vmp::VectorMatrixProduct;
|
||||
///
|
||||
/// let n = 1024;
|
||||
///
|
||||
/// let module: Module = Module::new::<FFT64>(n);
|
||||
/// let limbs: usize = 5;
|
||||
///
|
||||
/// let rows: usize = limbs;
|
||||
/// let cols: usize = limbs + 1;
|
||||
/// let c_limbs: usize = cols;
|
||||
/// let a_limbs: usize = limbs;
|
||||
/// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
|
||||
///
|
||||
/// let mut buf: Vec<u8> = vec![0; tmp_bytes];
|
||||
/// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
|
||||
///
|
||||
/// let a: VecZnx = module.new_vec_znx(limbs);
|
||||
/// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
|
||||
/// module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
|
||||
///
|
||||
/// c_dft.free();
|
||||
/// vmp_pmat.free();
|
||||
/// module.free();
|
||||
/// ```
|
||||
fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
|
||||
|
||||
/// Returns the size of the stratch space necessary for [VectorMatrixProduct::vmp_apply_dft_to_dft].
|
||||
fn vmp_apply_dft_to_dft_tmp_bytes(
|
||||
&self,
|
||||
c_limbs: usize,
|
||||
a_limbs: usize,
|
||||
rows: usize,
|
||||
cols: usize,
|
||||
) -> usize;
|
||||
|
||||
/// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
|
||||
/// The size of `buf` is given by [VectorMatrixProduct::vmp_apply_dft_to_dft_tmp_bytes].
|
||||
///
|
||||
/// A vector matrix product is equivalent to a sum of [ScalarVectorProduct::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 [VmpPMat].
|
||||
///
|
||||
/// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
|
||||
/// `j` cols, the output is a [VecZnx] of `j` limbs.
|
||||
///
|
||||
/// 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].
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, FFT64, Free};
|
||||
/// use base2k::vmp::VectorMatrixProduct;
|
||||
///
|
||||
/// let n = 1024;
|
||||
///
|
||||
/// let module: Module = Module::new::<FFT64>(n);
|
||||
/// let limbs: usize = 5;
|
||||
///
|
||||
/// let rows: usize = limbs;
|
||||
/// let cols: usize = limbs + 1;
|
||||
/// let c_limbs: usize = cols;
|
||||
/// let a_limbs: usize = limbs;
|
||||
/// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
|
||||
///
|
||||
/// let mut buf: Vec<u8> = vec![0; tmp_bytes];
|
||||
/// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
|
||||
///
|
||||
/// let a_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
|
||||
/// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
|
||||
/// module.vmp_apply_dft_to_dft(&mut c_dft, &a_dft, &vmp_pmat, &mut buf);
|
||||
///
|
||||
/// a_dft.free();
|
||||
/// c_dft.free();
|
||||
/// vmp_pmat.free();
|
||||
/// module.free();
|
||||
/// ```
|
||||
fn vmp_apply_dft_to_dft(&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 [VectorMatrixProduct::vmp_apply_dft_to_dft_tmp_bytes].
|
||||
///
|
||||
/// A vector matrix product is equivalent to a sum of [ScalarVectorProduct::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 [VmpPMat].
|
||||
///
|
||||
/// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
|
||||
/// `j` cols, the output is a [VecZnx] of `j` limbs.
|
||||
///
|
||||
/// 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].
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, FFT64, Free};
|
||||
/// use base2k::vmp::VectorMatrixProduct;
|
||||
///
|
||||
/// let n = 1024;
|
||||
///
|
||||
/// let module: Module = Module::new::<FFT64>(n);
|
||||
/// let limbs: usize = 5;
|
||||
///
|
||||
/// let rows: usize = limbs;
|
||||
/// let cols: usize = limbs + 1;
|
||||
/// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(limbs, limbs, rows, cols);
|
||||
///
|
||||
/// let mut buf: Vec<u8> = vec![0; tmp_bytes];
|
||||
/// let a: VecZnx = module.new_vec_znx(limbs);
|
||||
/// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
|
||||
///
|
||||
/// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
|
||||
/// module.vmp_apply_dft_to_dft_inplace(&mut c_dft, &vmp_pmat, &mut buf);
|
||||
///
|
||||
/// c_dft.free();
|
||||
/// vmp_pmat.free();
|
||||
/// module.free();
|
||||
/// ```
|
||||
fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
|
||||
}
|
||||
|
||||
impl VectorMatrixProduct for Module {
|
||||
fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
|
||||
unsafe {
|
||||
VmpPMat {
|
||||
data: new_vmp_pmat(self.0, rows as u64, cols as u64),
|
||||
rows,
|
||||
cols,
|
||||
n: self.n(),
|
||||
}
|
||||
}
|
||||
}
|
||||
fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
|
||||
unsafe { vmp_prepare_contiguous_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
|
||||
}
|
||||
|
||||
fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
|
||||
unsafe {
|
||||
vmp_prepare_contiguous(
|
||||
self.0,
|
||||
b.data(),
|
||||
a.as_ptr(),
|
||||
b.rows() as u64,
|
||||
b.cols() as u64,
|
||||
buf.as_mut_ptr(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]) {
|
||||
let rows: usize = b.rows();
|
||||
let cols: usize = b.cols();
|
||||
|
||||
let mut mat: Matrix3D<i64> = Matrix3D::<i64>::new(rows, cols, self.n());
|
||||
|
||||
(0..min(rows, a.len())).for_each(|i| {
|
||||
mat.set_row(i, &a[i].data);
|
||||
});
|
||||
|
||||
self.vmp_prepare_contiguous(b, &mat.data, buf);
|
||||
|
||||
/*
|
||||
NOT IMPLEMENTED IN SPQLIOS
|
||||
let mut ptrs: Vec<*const i64> = a.iter().map(|v| v.data.as_ptr()).collect();
|
||||
unsafe {
|
||||
vmp_prepare_dblptr(
|
||||
self.0,
|
||||
b.data(),
|
||||
ptrs.as_mut_ptr(),
|
||||
b.rows() as u64,
|
||||
b.cols() as u64,
|
||||
buf.as_mut_ptr(),
|
||||
);
|
||||
}
|
||||
*/
|
||||
}
|
||||
|
||||
fn vmp_apply_dft_tmp_bytes(
|
||||
&self,
|
||||
c_limbs: usize,
|
||||
a_limbs: usize,
|
||||
rows: usize,
|
||||
cols: usize,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
vmp_apply_dft_tmp_bytes(
|
||||
self.0,
|
||||
c_limbs as u64,
|
||||
a_limbs as u64,
|
||||
rows as u64,
|
||||
cols as u64,
|
||||
) as usize
|
||||
}
|
||||
}
|
||||
|
||||
fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
|
||||
unsafe {
|
||||
vmp_apply_dft(
|
||||
self.0,
|
||||
c.0,
|
||||
c.limbs() as u64,
|
||||
a.as_ptr(),
|
||||
a.limbs() as u64,
|
||||
a.n() as u64,
|
||||
b.data(),
|
||||
b.rows() as u64,
|
||||
b.cols() as u64,
|
||||
buf.as_mut_ptr(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn vmp_apply_dft_to_dft_tmp_bytes(
|
||||
&self,
|
||||
c_limbs: usize,
|
||||
a_limbs: usize,
|
||||
rows: usize,
|
||||
cols: usize,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
vmp_apply_dft_to_dft_tmp_bytes(
|
||||
self.0,
|
||||
c_limbs as u64,
|
||||
a_limbs as u64,
|
||||
rows as u64,
|
||||
cols as u64,
|
||||
) as usize
|
||||
}
|
||||
}
|
||||
|
||||
fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]) {
|
||||
unsafe {
|
||||
vmp_apply_dft_to_dft(
|
||||
self.0,
|
||||
c.0,
|
||||
c.limbs() as u64,
|
||||
a.0,
|
||||
a.limbs() as u64,
|
||||
b.data(),
|
||||
b.rows() as u64,
|
||||
b.cols() as u64,
|
||||
buf.as_mut_ptr(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
|
||||
unsafe {
|
||||
vmp_apply_dft_to_dft(
|
||||
self.0,
|
||||
b.0,
|
||||
b.limbs() as u64,
|
||||
b.0,
|
||||
b.limbs() as u64,
|
||||
a.data(),
|
||||
a.rows() as u64,
|
||||
a.cols() as u64,
|
||||
buf.as_mut_ptr(),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A helper struture that stores a 3D matrix as a contiguous array.
|
||||
/// To be passed to [VectorMatrixProduct::vmp_prepare_contiguous].
|
||||
///
|
||||
/// rows: index of the i-th base2K power.
|
||||
/// cols: index of the j-th limb of the i-th row.
|
||||
/// n : polynomial degree.
|
||||
///
|
||||
/// A [Matrix3D] can be seen as a vector of [VecZnx].
|
||||
pub struct Matrix3D<T> {
|
||||
pub data: Vec<T>,
|
||||
pub rows: usize,
|
||||
pub cols: usize,
|
||||
pub n: usize,
|
||||
}
|
||||
|
||||
impl<T: Default + Clone + std::marker::Copy> Matrix3D<T> {
|
||||
/// Allocates a new [Matrix3D] with the respective dimensions.
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::Matrix3D;
|
||||
///
|
||||
/// let rows = 5; // #decomp
|
||||
/// let cols = 5; // #limbs
|
||||
/// let n = 1024; // #coeffs
|
||||
///
|
||||
/// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
|
||||
/// ```
|
||||
pub fn new(rows: usize, cols: usize, n: usize) -> Self {
|
||||
let size = rows * cols * n;
|
||||
Self {
|
||||
data: vec![T::default(); size],
|
||||
rows,
|
||||
cols,
|
||||
n,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a non-mutable reference to the entry (row, col) of the [Matrix3D].
|
||||
/// The returned array is of size n.
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::Matrix3D;
|
||||
///
|
||||
/// let rows = 5; // #decomp
|
||||
/// let cols = 5; // #limbs
|
||||
/// let n = 1024; // #coeffs
|
||||
///
|
||||
/// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
|
||||
///
|
||||
/// let elem: &[i64] = mat.at(4, 4); // size n
|
||||
/// ```
|
||||
pub fn at(&self, row: usize, col: usize) -> &[T] {
|
||||
assert!(row < self.rows && col < self.cols);
|
||||
let idx: usize = row * (self.n * self.cols) + col * self.n;
|
||||
&self.data[idx..idx + self.n]
|
||||
}
|
||||
|
||||
/// Returns a mutable reference of the array at the (row, col) entry of the [Matrix3D].
|
||||
/// The returned array is of size n.
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::Matrix3D;
|
||||
///
|
||||
/// let rows = 5; // #decomp
|
||||
/// let cols = 5; // #limbs
|
||||
/// let n = 1024; // #coeffs
|
||||
///
|
||||
/// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
|
||||
///
|
||||
/// let elem: &mut [i64] = mat.at_mut(4, 4); // size n
|
||||
/// ```
|
||||
pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
|
||||
assert!(row < self.rows && col < self.cols);
|
||||
let idx: usize = row * (self.n * self.cols) + col * self.n;
|
||||
&mut self.data[idx..idx + self.n]
|
||||
}
|
||||
|
||||
/// Sets the entry \[row\] of the [Matrix3D].
|
||||
/// Typicall this is used to assign a [VecZnx] to the i-th row
|
||||
/// of the [Matrix3D].
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use base2k::{Matrix3D, VecZnx};
|
||||
///
|
||||
/// let rows = 5; // #decomp
|
||||
/// let cols = 5; // #limbs
|
||||
/// let n = 1024; // #coeffs
|
||||
///
|
||||
/// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
|
||||
///
|
||||
/// let a: VecZnx = VecZnx::new(n, cols);
|
||||
///
|
||||
/// mat.set_row(1, &a.data);
|
||||
/// ```
|
||||
pub fn set_row(&mut self, row: usize, a: &[T]) {
|
||||
assert!(
|
||||
row < self.rows,
|
||||
"invalid argument row: row={} > self.rows={}",
|
||||
row,
|
||||
self.rows
|
||||
);
|
||||
let idx: usize = row * (self.n * self.cols);
|
||||
let size: usize = min(a.len(), self.cols * self.n);
|
||||
self.data[idx..idx + size].copy_from_slice(&a[..size]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user