mirror of
https://github.com/arnaucube/poulpy.git
synced 2026-02-10 21:26:41 +01:00
reworked scalar
This commit is contained in:
Jean-Philippe Bossuat
@@ -1,279 +1,66 @@
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use crate::ffi::svp::{self, svp_ppol_t};
|
||||
use crate::ffi::vec_znx_dft::vec_znx_dft_t;
|
||||
use crate::znx_base::{ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize};
|
||||
use crate::{Backend, FFT64, Module, VecZnx, VecZnxDft, alloc_aligned, assert_alignement, cast_mut};
|
||||
use rand::seq::SliceRandom;
|
||||
use rand_core::RngCore;
|
||||
use rand_distr::{Distribution, weighted::WeightedIndex};
|
||||
use sampling::source::Source;
|
||||
use crate::ffi::svp;
|
||||
use crate::znx_base::{ZnxAlloc, ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize};
|
||||
use crate::{Backend, FFT64, GetZnxBase, Module};
|
||||
|
||||
pub struct Scalar {
|
||||
pub n: usize,
|
||||
pub data: Vec<i64>,
|
||||
pub ptr: *mut i64,
|
||||
}
|
||||
|
||||
impl<B: Backend> Module<B> {
|
||||
pub fn new_scalar(&self) -> Scalar {
|
||||
Scalar::new(self.n())
|
||||
}
|
||||
}
|
||||
|
||||
impl Scalar {
|
||||
pub fn new(n: usize) -> Self {
|
||||
let mut data: Vec<i64> = alloc_aligned::<i64>(n);
|
||||
let ptr: *mut i64 = data.as_mut_ptr();
|
||||
Self {
|
||||
n: n,
|
||||
data: data,
|
||||
ptr: ptr,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn n(&self) -> usize {
|
||||
self.n
|
||||
}
|
||||
|
||||
pub fn bytes_of(n: usize) -> usize {
|
||||
n * std::mem::size_of::<i64>()
|
||||
}
|
||||
|
||||
pub fn from_bytes(n: usize, bytes: &mut [u8]) -> Self {
|
||||
let size: usize = Self::bytes_of(n);
|
||||
debug_assert!(
|
||||
bytes.len() == size,
|
||||
"invalid buffer: bytes.len()={} < self.bytes_of(n={})={}",
|
||||
bytes.len(),
|
||||
n,
|
||||
size
|
||||
);
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_alignement(bytes.as_ptr())
|
||||
}
|
||||
unsafe {
|
||||
let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
|
||||
let ptr: *mut i64 = bytes_i64.as_mut_ptr();
|
||||
Self {
|
||||
n: n,
|
||||
data: Vec::from_raw_parts(bytes_i64.as_mut_ptr(), bytes.len(), bytes.len()),
|
||||
ptr: ptr,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_bytes_borrow(n: usize, bytes: &mut [u8]) -> Self {
|
||||
let size: usize = Self::bytes_of(n);
|
||||
debug_assert!(
|
||||
bytes.len() == size,
|
||||
"invalid buffer: bytes.len()={} < self.bytes_of(n={})={}",
|
||||
bytes.len(),
|
||||
n,
|
||||
size
|
||||
);
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_alignement(bytes.as_ptr())
|
||||
}
|
||||
let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
|
||||
let ptr: *mut i64 = bytes_i64.as_mut_ptr();
|
||||
Self {
|
||||
n: n,
|
||||
data: Vec::new(),
|
||||
ptr: ptr,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_ptr(&self) -> *const i64 {
|
||||
self.ptr
|
||||
}
|
||||
|
||||
pub fn raw(&self) -> &[i64] {
|
||||
unsafe { std::slice::from_raw_parts(self.ptr, self.n) }
|
||||
}
|
||||
|
||||
pub fn raw_mut(&self) -> &mut [i64] {
|
||||
unsafe { std::slice::from_raw_parts_mut(self.ptr, self.n) }
|
||||
}
|
||||
|
||||
pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
|
||||
let choices: [i64; 3] = [-1, 0, 1];
|
||||
let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
|
||||
let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
|
||||
self.data
|
||||
.iter_mut()
|
||||
.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.data[..hw]
|
||||
.iter_mut()
|
||||
.for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
|
||||
self.data.shuffle(source);
|
||||
}
|
||||
|
||||
pub fn as_vec_znx(&self) -> VecZnx {
|
||||
VecZnx {
|
||||
inner: ZnxBase {
|
||||
n: self.n,
|
||||
rows: 1,
|
||||
cols: 1,
|
||||
size: 1,
|
||||
data: Vec::new(),
|
||||
ptr: self.ptr as *mut u8,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ScalarOps {
|
||||
fn bytes_of_scalar(&self) -> usize;
|
||||
fn new_scalar(&self) -> Scalar;
|
||||
fn new_scalar_from_bytes(&self, bytes: &mut [u8]) -> Scalar;
|
||||
fn new_scalar_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> Scalar;
|
||||
}
|
||||
impl<B: Backend> ScalarOps for Module<B> {
|
||||
fn bytes_of_scalar(&self) -> usize {
|
||||
Scalar::bytes_of(self.n())
|
||||
}
|
||||
fn new_scalar(&self) -> Scalar {
|
||||
Scalar::new(self.n())
|
||||
}
|
||||
fn new_scalar_from_bytes(&self, bytes: &mut [u8]) -> Scalar {
|
||||
Scalar::from_bytes(self.n(), bytes)
|
||||
}
|
||||
fn new_scalar_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> Scalar {
|
||||
Scalar::from_bytes_borrow(self.n(), tmp_bytes)
|
||||
}
|
||||
}
|
||||
pub const SCALAR_ZNX_DFT_ROWS: usize = 1;
|
||||
pub const SCALAR_ZNX_DFT_SIZE: usize = 1;
|
||||
|
||||
pub struct ScalarZnxDft<B: Backend> {
|
||||
pub n: usize,
|
||||
pub data: Vec<u8>,
|
||||
pub ptr: *mut u8,
|
||||
pub inner: ZnxBase,
|
||||
_marker: PhantomData<B>,
|
||||
}
|
||||
|
||||
/// A prepared [crate::Scalar] for [SvpPPolOps::svp_apply_dft].
|
||||
/// An [SvpPPol] an be seen as a [VecZnxDft] of one limb.
|
||||
impl ScalarZnxDft<FFT64> {
|
||||
pub fn new(module: &Module<FFT64>) -> Self {
|
||||
module.new_scalar_znx_dft()
|
||||
impl<B: Backend> GetZnxBase for ScalarZnxDft<B> {
|
||||
fn znx(&self) -> &ZnxBase {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
/// Returns the ring degree of the [SvpPPol].
|
||||
pub fn n(&self) -> usize {
|
||||
self.n
|
||||
fn znx_mut(&mut self) -> &mut ZnxBase {
|
||||
&mut self.inner
|
||||
}
|
||||
}
|
||||
|
||||
pub fn bytes_of(module: &Module<FFT64>) -> usize {
|
||||
module.bytes_of_scalar_znx_dft()
|
||||
}
|
||||
impl<B: Backend> ZnxInfos for ScalarZnxDft<B> {}
|
||||
|
||||
pub fn from_bytes(module: &Module<FFT64>, bytes: &mut [u8]) -> Self {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_alignement(bytes.as_ptr());
|
||||
assert_eq!(bytes.len(), module.bytes_of_scalar_znx_dft());
|
||||
}
|
||||
unsafe {
|
||||
Self {
|
||||
n: module.n(),
|
||||
data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
|
||||
ptr: bytes.as_mut_ptr(),
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl<B: Backend> ZnxAlloc<B> for ScalarZnxDft<B> {
|
||||
type Scalar = u8;
|
||||
|
||||
pub fn from_bytes_borrow(module: &Module<FFT64>, tmp_bytes: &mut [u8]) -> Self {
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
assert_alignement(tmp_bytes.as_ptr());
|
||||
assert_eq!(tmp_bytes.len(), module.bytes_of_scalar_znx_dft());
|
||||
}
|
||||
fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, _size: usize, bytes: &mut [u8]) -> Self {
|
||||
Self {
|
||||
n: module.n(),
|
||||
data: Vec::new(),
|
||||
ptr: tmp_bytes.as_mut_ptr(),
|
||||
inner: ZnxBase::from_bytes_borrow(
|
||||
module.n(),
|
||||
SCALAR_ZNX_DFT_ROWS,
|
||||
cols,
|
||||
SCALAR_ZNX_DFT_SIZE,
|
||||
bytes,
|
||||
),
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the number of cols of the [SvpPPol], which is always 1.
|
||||
pub fn cols(&self) -> usize {
|
||||
1
|
||||
fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, _size: usize) -> usize {
|
||||
debug_assert_eq!(
|
||||
_rows, SCALAR_ZNX_DFT_ROWS,
|
||||
"rows != {} not supported for ScalarZnxDft",
|
||||
SCALAR_ZNX_DFT_ROWS
|
||||
);
|
||||
debug_assert_eq!(
|
||||
_size, SCALAR_ZNX_DFT_SIZE,
|
||||
"rows != {} not supported for ScalarZnxDft",
|
||||
SCALAR_ZNX_DFT_SIZE
|
||||
);
|
||||
unsafe { svp::bytes_of_svp_ppol(module.ptr) as usize * cols }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ScalarZnxDftOps<B: Backend> {
|
||||
/// Allocates a new [SvpPPol].
|
||||
fn new_scalar_znx_dft(&self) -> ScalarZnxDft<B>;
|
||||
|
||||
/// Returns the minimum number of bytes necessary to allocate
|
||||
/// a new [SvpPPol] through [SvpPPol::from_bytes] ro.
|
||||
fn bytes_of_scalar_znx_dft(&self) -> usize;
|
||||
|
||||
/// Allocates a new [SvpPPol] from an array of bytes.
|
||||
/// The array of bytes is owned by the [SvpPPol].
|
||||
/// The method will panic if bytes.len() < [SvpPPolOps::bytes_of_svp_ppol]
|
||||
fn new_scalar_znx_dft_from_bytes(&self, bytes: &mut [u8]) -> ScalarZnxDft<B>;
|
||||
|
||||
/// Allocates a new [SvpPPol] from an array of bytes.
|
||||
/// The array of bytes is borrowed by the [SvpPPol].
|
||||
/// The method will panic if bytes.len() < [SvpPPolOps::bytes_of_svp_ppol]
|
||||
fn new_scalar_znx_dft_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> ScalarZnxDft<B>;
|
||||
|
||||
/// Prepares a [crate::Scalar] for a [SvpPPolOps::svp_apply_dft].
|
||||
fn svp_prepare(&self, svp_ppol: &mut ScalarZnxDft<B>, a: &Scalar);
|
||||
|
||||
/// Applies the [SvpPPol] x [VecZnxDft] product, where each limb of
|
||||
/// the [VecZnxDft] is multiplied with [SvpPPol].
|
||||
fn svp_apply_dft(&self, res: &mut VecZnxDft<B>, res_col: usize, a: &ScalarZnxDft<B>, b: &VecZnx, b_col: usize);
|
||||
impl ZnxLayout for ScalarZnxDft<FFT64> {
|
||||
type Scalar = f64;
|
||||
}
|
||||
|
||||
impl ScalarZnxDftOps<FFT64> for Module<FFT64> {
|
||||
fn new_scalar_znx_dft(&self) -> ScalarZnxDft<FFT64> {
|
||||
let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_scalar_znx_dft());
|
||||
let ptr: *mut u8 = data.as_mut_ptr();
|
||||
ScalarZnxDft::<FFT64> {
|
||||
data: data,
|
||||
ptr: ptr,
|
||||
n: self.n(),
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
fn bytes_of_scalar_znx_dft(&self) -> usize {
|
||||
unsafe { svp::bytes_of_svp_ppol(self.ptr) as usize }
|
||||
}
|
||||
|
||||
fn new_scalar_znx_dft_from_bytes(&self, bytes: &mut [u8]) -> ScalarZnxDft<FFT64> {
|
||||
ScalarZnxDft::from_bytes(self, bytes)
|
||||
}
|
||||
|
||||
fn new_scalar_znx_dft_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> ScalarZnxDft<FFT64> {
|
||||
ScalarZnxDft::from_bytes_borrow(self, tmp_bytes)
|
||||
}
|
||||
|
||||
fn svp_prepare(&self, res: &mut ScalarZnxDft<FFT64>, a: &Scalar) {
|
||||
unsafe { svp::svp_prepare(self.ptr, res.ptr as *mut svp_ppol_t, a.as_ptr()) }
|
||||
}
|
||||
|
||||
fn svp_apply_dft(&self, res: &mut VecZnxDft<FFT64>, res_col: usize, a: &ScalarZnxDft<FFT64>, b: &VecZnx, b_col: usize) {
|
||||
unsafe {
|
||||
svp::svp_apply_dft(
|
||||
self.ptr,
|
||||
res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
|
||||
res.size() as u64,
|
||||
a.ptr as *const svp_ppol_t,
|
||||
b.at_ptr(b_col, 0),
|
||||
b.size() as u64,
|
||||
b.sl() as u64,
|
||||
)
|
||||
}
|
||||
impl ZnxSliceSize for ScalarZnxDft<FFT64> {
|
||||
fn sl(&self) -> usize {
|
||||
self.n()
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user