poulpy/base2k/src/scalar_znx.rs

use crate::znx_base::ZnxInfos;
use crate::{Backend, DataView, DataViewMut, Module, ZnxView, ZnxViewMut, alloc_aligned};
use rand::seq::SliceRandom;
use rand_core::RngCore;
use rand_distr::{Distribution, weighted::WeightedIndex};
use sampling::source::Source;

// pub const SCALAR_ZNX_ROWS: usize = 1;
// pub const SCALAR_ZNX_SIZE: usize = 1;

pub struct Scalar<D> {
    data: D,
    n: usize,
    cols: usize,
}

impl<D> ZnxInfos for Scalar<D> {
    fn cols(&self) -> usize {
        self.cols
    }

    fn rows(&self) -> usize {
        1
    }

    fn n(&self) -> usize {
        self.n
    }

    fn size(&self) -> usize {
        1
    }

    fn sl(&self) -> usize {
        self.n()
    }
}

impl<D> DataView for Scalar<D> {
    type D = D;
    fn data(&self) -> &Self::D {
        &self.data
    }
}

impl<D> DataViewMut for Scalar<D> {
    fn data_mut(&mut self) -> &mut Self::D {
        &mut self.data
    }
}

impl<D: AsRef<[u8]>> ZnxView for Scalar<D> {
    type Scalar = i64;
}

impl<D: AsMut<[u8]> + AsRef<[u8]>> Scalar<D> {
    pub fn fill_ternary_prob(&mut self, col: usize, 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.at_mut(col, 0)
            .iter_mut()
            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
    }

    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
        assert!(hw <= self.n());
        self.at_mut(col, 0)[..hw]
            .iter_mut()
            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
        self.at_mut(col, 0).shuffle(source);
    }

    // pub fn alias_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,
    //         },
    //     }
    // }
}

impl<D: From<Vec<u8>>> Scalar<D> {
    pub(crate) fn bytes_of<S: Sized>(n: usize, cols: usize) -> usize {
        n * cols * size_of::<S>()
    }

    pub(crate) fn new<S: Sized>(n: usize, cols: usize) -> Self {
        let data = alloc_aligned::<u8>(Self::bytes_of::<S>(n, cols));
        Self {
            data: data.into(),
            n,
            cols,
        }
    }

    pub(crate) fn new_from_bytes<S: Sized>(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
        let data: Vec<u8> = bytes.into();
        assert!(data.len() == Self::bytes_of::<S>(n, cols));
        Self {
            data: data.into(),
            n,
            cols,
        }
    }
}

pub type ScalarOwned = Scalar<Vec<u8>>;

pub trait ScalarAlloc {
    fn bytes_of_scalar(&self, cols: usize) -> usize;
    fn new_scalar(&self, cols: usize) -> ScalarOwned;
    fn new_scalar_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarOwned;
    // fn new_scalar_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> Scalar;
}

impl<B: Backend> ScalarAlloc for Module<B> {
    fn bytes_of_scalar(&self, cols: usize) -> usize {
        ScalarOwned::bytes_of::<i64>(self.n(), cols)
    }
    fn new_scalar(&self, cols: usize) -> ScalarOwned {
        ScalarOwned::new::<i64>(self.n(), cols)
    }
    fn new_scalar_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarOwned {
        ScalarOwned::new_from_bytes::<i64>(self.n(), cols, bytes)
    }
    // fn new_scalar_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> Scalar {
    //     Scalar::from_bytes_borrow(self, SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes)
    // }
}

// impl<B: Backend> ZnxAlloc<B> for Scalar {
//     type Scalar = i64;

//     fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, _size: usize, bytes: &mut [u8]) -> Self {
//         Self {
//             inner: ZnxBase::from_bytes_borrow(module.n(), SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes),
//         }
//     }

//     fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, _size: usize) -> usize {
//         debug_assert_eq!(
//             _rows, SCALAR_ZNX_ROWS,
//             "rows != {} not supported for Scalar",
//             SCALAR_ZNX_ROWS
//         );
//         debug_assert_eq!(
//             _size, SCALAR_ZNX_SIZE,
//             "rows != {} not supported for Scalar",
//             SCALAR_ZNX_SIZE
//         );
//         module.n() * cols * std::mem::size_of::<self::Scalar>()
//     }
// }