1use crate::znx_base::{GetZnxBase, ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize};
  2use crate::{Backend, FFT64, Module, alloc_aligned};
  3use std::marker::PhantomData;
  4
  5/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
  6/// stored as a 3D matrix in the DFT domain in a single contiguous array.
  7/// Each col of the [MatZnxDft] can be seen as a collection of [VecZnxDft].
  8///
  9/// [MatZnxDft] is used to permform a vector matrix product between a [VecZnx]/[VecZnxDft] and a [MatZnxDft].
 10/// See the trait [MatZnxDftOps] for additional information.
 11pub struct MatZnxDft<B: Backend> {
 12    pub inner: ZnxBase,
 13    pub cols_in: usize,
 14    pub cols_out: usize,
 15    _marker: PhantomData<B>,
 16}
 17
 18impl<B: Backend> GetZnxBase for MatZnxDft<B> {
 19    fn znx(&self) -> &ZnxBase {
 20        &self.inner
 21    }
 22
 23    fn znx_mut(&mut self) -> &mut ZnxBase {
 24        &mut self.inner
 25    }
 26}
 27
 28impl<B: Backend> ZnxInfos for MatZnxDft<B> {}
 29
 30impl ZnxSliceSize for MatZnxDft<FFT64> {
 31    fn sl(&self) -> usize {
 32        self.n()
 33    }
 34}
 35
 36impl ZnxLayout for MatZnxDft<FFT64> {
 37    type Scalar = f64;
 38}
 39
 40impl<B: Backend> MatZnxDft<B> {
 41    pub fn new(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
 42        let bytes: Vec<u8> = alloc_aligned(Self::bytes_of(module, rows, cols_in, cols_out, size));
 43        Self::from_bytes(module, rows, cols_in, cols_out, size, bytes)
 44    }
 45
 46    pub fn from_bytes(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize, mut bytes: Vec<u8>) -> Self {
 47        let mut mat: MatZnxDft<B> = Self::from_bytes_borrow(module, rows, cols_in, cols_out, size, &mut bytes);
 48        mat.znx_mut().data = bytes;
 49        mat
 50    }
 51
 52    pub fn from_bytes_borrow(
 53        module: &Module<B>,
 54        rows: usize,
 55        cols_in: usize,
 56        cols_out: usize,
 57        size: usize,
 58        bytes: &mut [u8],
 59    ) -> Self {
 60        debug_assert_eq!(
 61            bytes.len(),
 62            Self::bytes_of(module, rows, cols_in, cols_out, size)
 63        );
 64        Self {
 65            inner: ZnxBase::from_bytes_borrow(module.n(), rows, cols_out, size, bytes),
 66            cols_in: cols_in,
 67            cols_out: cols_out,
 68            _marker: PhantomData,
 69        }
 70    }
 71
 72    pub fn bytes_of(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 73        unsafe {
 74            crate::ffi::vmp::bytes_of_vmp_pmat(
 75                module.ptr,
 76                (rows * cols_in) as u64,
 77                (size * cols_out) as u64,
 78            ) as usize
 79        }
 80    }
 81
 82    pub fn cols_in(&self) -> usize {
 83        self.cols_in
 84    }
 85
 86    pub fn cols_out(&self) -> usize {
 87        self.cols_out
 88    }
 89}
 90
 91impl MatZnxDft<FFT64> {
 92    /// Returns a copy of the backend array at index (i, j) of the [MatZnxDft].
 93    ///
 94    /// # Arguments
 95    ///
 96    /// * `row`: row index (i).
 97    /// * `col`: col index (j).
 98    #[allow(dead_code)]
 99    fn at(&self, row: usize, col: usize) -> Vec<f64> {
100        let n: usize = self.n();
101
102        let mut res: Vec<f64> = alloc_aligned(n);
103
104        if n < 8 {
105            res.copy_from_slice(&self.raw()[(row + col * self.rows()) * n..(row + col * self.rows()) * (n + 1)]);
106        } else {
107            (0..n >> 3).for_each(|blk| {
108                res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
109            });
110        }
111
112        res
113    }
114
115    #[allow(dead_code)]
116    fn at_block(&self, row: usize, col: usize, blk: usize) -> &[f64] {
117        let nrows: usize = self.rows();
118        let nsize: usize = self.size();
119        if col == (nsize - 1) && (nsize & 1 == 1) {
120            &self.raw()[blk * nrows * nsize * 8 + col * nrows * 8 + row * 8..]
121        } else {
122            &self.raw()[blk * nrows * nsize * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
123        }
124    }
125}