1use crate::znx_base::ZnxInfos;
  2use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxView, 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<D, B> {
 12    data: D,
 13    n: usize,
 14    size: usize,
 15    rows: usize,
 16    cols_in: usize,
 17    cols_out: usize,
 18    _marker: PhantomData<B>,
 19}
 20
 21impl<D, B> ZnxInfos for MatZnxDft<D, B> {
 22    fn cols(&self) -> usize {
 23        self.cols_in
 24    }
 25
 26    fn rows(&self) -> usize {
 27        self.rows
 28    }
 29
 30    fn n(&self) -> usize {
 31        self.n
 32    }
 33
 34    fn size(&self) -> usize {
 35        self.size
 36    }
 37
 38    fn sl(&self) -> usize {
 39        self.n()
 40    }
 41}
 42
 43impl<D, B> DataView for MatZnxDft<D, B> {
 44    type D = D;
 45    fn data(&self) -> &Self::D {
 46        &self.data
 47    }
 48}
 49
 50impl<D, B> DataViewMut for MatZnxDft<D, B> {
 51    fn data_mut(&mut self) -> &mut Self::D {
 52        &mut self.data
 53    }
 54}
 55
 56impl<D: AsRef<[u8]>> ZnxView for MatZnxDft<D, FFT64> {
 57    type Scalar = f64;
 58}
 59
 60impl<D, B> MatZnxDft<D, B> {
 61    pub(crate) fn cols_in(&self) -> usize {
 62        self.cols_in
 63    }
 64
 65    pub(crate) fn cols_out(&self) -> usize {
 66        self.cols_out
 67    }
 68}
 69
 70impl<D: From<Vec<u8>>, B: Backend> MatZnxDft<D, B> {
 71    pub(crate) fn bytes_of(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 72        unsafe {
 73            crate::ffi::vmp::bytes_of_vmp_pmat(
 74                module.ptr,
 75                (rows * cols_in) as u64,
 76                (size * cols_out) as u64,
 77            ) as usize
 78        }
 79    }
 80
 81    pub(crate) fn new(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
 82        let data: Vec<u8> = alloc_aligned(Self::bytes_of(module, rows, cols_in, cols_out, size));
 83        Self {
 84            data: data.into(),
 85            n: module.n(),
 86            size,
 87            rows,
 88            cols_in,
 89            cols_out,
 90            _marker: PhantomData,
 91        }
 92    }
 93
 94    pub(crate) fn new_from_bytes(
 95        module: &Module<B>,
 96        rows: usize,
 97        cols_in: usize,
 98        cols_out: usize,
 99        size: usize,
100        bytes: impl Into<Vec<u8>>,
101    ) -> Self {
102        let data: Vec<u8> = bytes.into();
103        assert!(data.len() == Self::bytes_of(module, rows, cols_in, cols_out, size));
104        Self {
105            data: data.into(),
106            n: module.n(),
107            size,
108            rows,
109            cols_in,
110            cols_out,
111            _marker: PhantomData,
112        }
113    }
114}
115
116impl<D: AsRef<[u8]>> MatZnxDft<D, FFT64> {
117    /// Returns a copy of the backend array at index (i, j) of the [MatZnxDft].
118    ///
119    /// # Arguments
120    ///
121    /// * `row`: row index (i).
122    /// * `col`: col index (j).
123    #[allow(dead_code)]
124    fn at(&self, row: usize, col: usize) -> Vec<f64> {
125        let n: usize = self.n();
126
127        let mut res: Vec<f64> = alloc_aligned(n);
128
129        if n < 8 {
130            res.copy_from_slice(&self.raw()[(row + col * self.rows()) * n..(row + col * self.rows()) * (n + 1)]);
131        } else {
132            (0..n >> 3).for_each(|blk| {
133                res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
134            });
135        }
136
137        res
138    }
139
140    #[allow(dead_code)]
141    fn at_block(&self, row: usize, col: usize, blk: usize) -> &[f64] {
142        let nrows: usize = self.rows();
143        let nsize: usize = self.size();
144        if col == (nsize - 1) && (nsize & 1 == 1) {
145            &self.raw()[blk * nrows * nsize * 8 + col * nrows * 8 + row * 8..]
146        } else {
147            &self.raw()[blk * nrows * nsize * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
148        }
149    }
150}
151
152pub type MatZnxDftAllocOwned<B> = MatZnxDft<Vec<u8>, B>;
153
154impl<B> MatZnxDft<Vec<u8>, B> {
155    pub fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
156        MatZnxDft {
157            data: self.data.as_mut_slice(),
158            n: self.n,
159            size: self.size,
160            rows: self.rows,
161            cols_in: self.cols_in,
162            cols_out: self.cols_out,
163            _marker: PhantomData,
164        }
165    }
166
167    pub fn to_ref(&self) -> MatZnxDft<&[u8], B> {
168        MatZnxDft {
169            data: self.data.as_slice(),
170            n: self.n,
171            size: self.size,
172            rows: self.rows,
173            cols_in: self.cols_in,
174            cols_out: self.cols_out,
175            _marker: PhantomData,
176        }
177    }
178}