arnaucube
arnaucube/poulpyPublic Self-hosted repository
updated base2k backend
Jean-Philippe Bossuat jeanphilippe.bossuat@gmail.com
Fri, 14 Feb 2025 10:58:28 +0100
commit

68e61dc0e36a58a12d3671fb4e4c9f534c24f6d4

parent

4aeaf99fe25ee451ff51e2c54ecb782fc3e5d1f0

M base2k/examples/vector_matrix_product.rs → base2k/examples/vector_matrix_product.rs

@@ -40,8 +40,10 @@ (0..rows).for_each(|i| {

vecznx[i].data[i * n + 1] = 1 as i64; }); + let dble: Vec<&[i64]> = vecznx.iter().map(|v| v.data.as_slice()).collect(); + let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols); - module.vmp_prepare_dblptr(&mut vmp_pmat, &vecznx, &mut buf); + module.vmp_prepare_dblptr(&mut vmp_pmat, &dble, &mut buf); let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols); module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
M base2k/src/lib.rs → base2k/src/lib.rs

@@ -73,3 +73,51 @@ let ptr: *const f64 = data.as_mut_ptr() as *const f64;

let len: usize = data.len() / std::mem::size_of::<f64>(); unsafe { std::slice::from_raw_parts(ptr, len) } } + +use std::alloc::{alloc, Layout}; + +pub fn alloc_aligned_u8(size: usize, align: usize) -> Vec<u8> { + assert_eq!( + align & (align - 1), + 0, + "align={} must be a power of two", + align + ); + assert_eq!( + (size * std::mem::size_of::<u8>()) % align, + 0, + "size={} must be a multiple of align={}", + size, + align + ); + unsafe { + let layout: Layout = Layout::from_size_align(size, align).expect("Invalid alignment"); + let ptr: *mut u8 = alloc(layout); + if ptr.is_null() { + panic!("Memory allocation failed"); + } + Vec::from_raw_parts(ptr, size, size) + } +} + +pub fn alloc_aligned<T>(size: usize, align: usize) -> Vec<T> { + assert_eq!( + (size * std::mem::size_of::<T>()) % align, + 0, + "size={} must be a multiple of align={}", + size, + align + ); + let mut vec_u8: Vec<u8> = alloc_aligned_u8(std::mem::size_of::<T>() * size, align); + let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T; + let len: usize = vec_u8.len() / std::mem::size_of::<T>(); + let cap: usize = vec_u8.capacity() / std::mem::size_of::<T>(); + std::mem::forget(vec_u8); + unsafe { Vec::from_raw_parts(ptr, len, cap) } +} + +fn alias_mut_slice_to_vec<T>(slice: &mut [T]) -> Vec<T> { + let ptr = slice.as_mut_ptr(); + let len = slice.len(); + unsafe { Vec::from_raw_parts(ptr, len, len) } +}
M base2k/src/svp.rs → base2k/src/svp.rs

@@ -1,7 +1,7 @@

use crate::ffi::svp; -use crate::{Module, VecZnx, VecZnxDft}; +use crate::{alias_mut_slice_to_vec, Module, VecZnx, VecZnxDft}; -use crate::Infos; +use crate::{alloc_aligned, cast_mut_u8_to_mut_i64_slice, Infos}; use rand::seq::SliceRandom; use rand_core::RngCore; use rand_distr::{Distribution, WeightedIndex};

@@ -17,14 +17,18 @@ }

impl Scalar { pub fn new(n: usize) -> Self { - Self(vec![i64::default(); Self::buffer_size(n)]) + Self(alloc_aligned::<i64>(n, 64)) + } + + pub fn n(&self) -> usize { + self.0.len() } pub fn buffer_size(n: usize) -> usize { n } - pub fn from_buffer(&mut self, n: usize, buf: &[i64]) { + pub fn from_buffer(&mut self, n: usize, buf: &mut [u8]) { let size: usize = Self::buffer_size(n); assert!( buf.len() >= size,

@@ -33,7 +37,7 @@ buf.len(),

n, size ); - self.0 = Vec::from(&buf[..size]) + self.0 = alias_mut_slice_to_vec(cast_mut_u8_to_mut_i64_slice(&mut buf[..size])) } pub fn as_ptr(&self) -> *const i64 {

@@ -50,6 +54,7 @@ .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);

} pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) { + assert!(hw <= self.n()); self.0[..hw] .iter_mut() .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
M base2k/src/vec_znx.rs → base2k/src/vec_znx.rs

@@ -1,6 +1,7 @@

use crate::cast_mut_u8_to_mut_i64_slice; use crate::ffi::vec_znx; use crate::ffi::znx; +use crate::{alias_mut_slice_to_vec, alloc_aligned}; use crate::{Infos, Module}; use itertools::izip; use std::cmp::min;

@@ -21,7 +22,7 @@ /// Allocates a new [VecZnx] composed of #limbs polynomials of Z\[X\].

pub fn new(n: usize, limbs: usize) -> Self { Self { n: n, - data: vec![i64::default(); n * limbs], + data: alloc_aligned::<i64>(n * limbs, 64), } }

@@ -47,7 +48,7 @@ );

VecZnx { n: n, - data: Vec::from(cast_mut_u8_to_mut_i64_slice(&mut buf[..size])), + data: alias_mut_slice_to_vec(cast_mut_u8_to_mut_i64_slice(&mut buf[..size])), } }

@@ -106,7 +107,7 @@ /// The method will panic if carry.len() < self.data.len()*8.

/// /// # Example /// ``` - /// use base2k::{VecZnx, Encoding, Infos}; + /// use base2k::{VecZnx, Encoding, Infos, alloc_aligned}; /// use itertools::izip; /// use sampling::source::Source; ///

@@ -115,8 +116,8 @@ /// let log_base2k: usize = 17; // base two logarithm of the coefficients decomposition

/// let limbs: usize = 5; // number of limbs (i.e. can store coeffs in the range +/- 2^{limbs * log_base2k - 1}) /// let log_k: usize = limbs * log_base2k - 5; /// let mut a: VecZnx = VecZnx::new(n, limbs); - /// let mut carry: Vec<u8> = vec![u8::default(); a.n()<<3]; - /// let mut have: Vec<i64> = vec![i64::default(); a.n()]; + /// let mut carry: Vec<u8> = alloc_aligned::<u8>(a.n()<<3, 64); + /// let mut have: Vec<i64> = alloc_aligned::<i64>(a.n(), 64); /// let mut source = Source::new([1; 32]); /// /// // Populates the first limb of the of polynomials with random i64 values.

@@ -135,7 +136,7 @@ /// .iter()

/// .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half)); /// /// // Ensures reconstructed normalized values are equal to non-normalized values. - /// let mut want = vec![i64::default(); n]; + /// let mut want = alloc_aligned::<i64>(a.n(), 64); /// a.decode_vec_i64(log_base2k, log_k, &mut want); /// izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b)); /// ```
M base2k/src/vec_znx_big.rs → base2k/src/vec_znx_big.rs

@@ -29,6 +29,25 @@ pub fn new_vec_znx_big(&self, limbs: usize) -> VecZnxBig {

unsafe { VecZnxBig(vec_znx_big::new_vec_znx_big(self.0, limbs as u64), limbs) } } + /// Returns a new [VecZnxBig] with the provided bytes array as backing array. + /// + /// # Arguments + /// + /// * `limbs`: the number of limbs of the [VecZnxBig]. + /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big]. + /// + /// # Panics + /// If `bytes.len()` < [Module::bytes_of_vec_znx_big]. + pub fn new_vec_znx_big_from_bytes(&self, limbs: usize, bytes: &mut [u8]) -> VecZnxBig { + assert!( + bytes.len() >= self.bytes_of_vec_znx_big(limbs), + "invalid bytes: bytes.len()={} < bytes_of_vec_znx_dft={}", + bytes.len(), + self.bytes_of_vec_znx_big(limbs) + ); + VecZnxBig::from_bytes(limbs, bytes) + } + /// Returns the minimum number of bytes necessary to allocate /// a new [VecZnxBig] through [VecZnxBig::from_bytes]. pub fn bytes_of_vec_znx_big(&self, limbs: usize) -> usize {

@@ -128,6 +147,42 @@ a.0,

a.limbs() as u64, tmp_bytes.as_mut_ptr(), ) + } + } + + pub fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize { + unsafe { vec_znx_big::vec_znx_big_range_normalize_base2k_tmp_bytes(self.0) as usize } + } + + pub fn vec_znx_big_range_normalize_base2k( + &self, + log_base2k: usize, + res: &mut VecZnx, + a: &VecZnxBig, + a_range_begin: usize, + a_range_xend: usize, + a_range_step: usize, + tmp_bytes: &mut [u8], + ) { + assert!( + tmp_bytes.len() >= self.vec_znx_big_range_normalize_base2k_tmp_bytes(), + "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_range_normalize_base2k_tmp_bytes()={}", + tmp_bytes.len(), + self.vec_znx_big_range_normalize_base2k_tmp_bytes() + ); + unsafe { + vec_znx_big::vec_znx_big_range_normalize_base2k( + self.0, + log_base2k as u64, + res.as_mut_ptr(), + res.limbs() as u64, + res.n() as u64, + a.0, + a_range_begin as u64, + a_range_xend as u64, + a_range_step as u64, + tmp_bytes.as_mut_ptr(), + ); } }
M base2k/src/vec_znx_dft.rs → base2k/src/vec_znx_dft.rs

@@ -39,7 +39,7 @@ /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].

/// /// # Panics /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft]. - pub fn new_vec_znx_from_bytes(&self, limbs: usize, bytes: &mut [u8]) -> VecZnxDft { + pub fn new_vec_znx_dft_from_bytes(&self, limbs: usize, bytes: &mut [u8]) -> VecZnxDft { assert!( bytes.len() >= self.bytes_of_vec_znx_dft(limbs), "invalid bytes: bytes.len()={} < bytes_of_vec_znx_dft={}",

@@ -63,7 +63,9 @@ "invalid c_vector: b_vector.limbs()={} < a_limbs={}",

b.limbs(), a_limbs ); - unsafe { vec_znx_dft::vec_znx_idft_tmp_a(self.0, b.0, b.limbs() as u64, a.0, a_limbs as u64) } + unsafe { + vec_znx_dft::vec_znx_idft_tmp_a(self.0, b.0, b.limbs() as u64, a.0, a_limbs as u64) + } } // Returns the size of the scratch space for [vec_znx_idft].
M base2k/src/vmp.rs → base2k/src/vmp.rs

@@ -159,15 +159,17 @@ /// (0..rows).for_each(|_|{

/// vecznx.push(module.new_vec_znx(cols)); /// }); /// + /// let dble: Vec<&[i64]> = vecznx.iter().map(|v| v.data.as_slice()).collect(); + /// /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_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); + /// module.vmp_prepare_dblptr(&mut vmp_pmat, &dble, &mut buf); /// /// vmp_pmat.free(); /// module.free(); /// ``` - fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]); + fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], buf: &mut [u8]); /// Prepares the ith-row of [VmpPMat] from a vector of [VecZnx]. ///

@@ -189,7 +191,7 @@ /// let module: Module = Module::new::<FFT64>(n);

/// let rows: usize = 5; /// let cols: usize = 6; /// - /// let vecznx = module.new_vec_znx(cols); + /// let vecznx = vec![0i64; cols*n]; /// /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_tmp_bytes(rows, cols)]; ///

@@ -199,7 +201,7 @@ ///

/// vmp_pmat.free(); /// module.free(); /// ``` - fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &VecZnx, row_i: usize, tmp_bytes: &mut [u8]); + fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]); /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft]. ///

@@ -422,8 +424,8 @@ );

} } - fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]) { - let ptrs: Vec<*const i64> = a.iter().map(|v| v.data.as_ptr()).collect(); + fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], buf: &mut [u8]) { + let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect(); unsafe { vmp::vmp_prepare_dblptr( self.0,

@@ -436,12 +438,12 @@ );

} } - fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &VecZnx, row_i: usize, buf: &mut [u8]) { + fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, buf: &mut [u8]) { unsafe { vmp::vmp_prepare_row( self.0, b.data(), - a.data.as_ptr(), + a.as_ptr(), row_i as u64, b.rows() as u64, b.cols() as u64,