1use crate::znx_base::ZnxInfos;
  2use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxSliceSize, 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: Backend> {
 12    data: D,
 13    n: usize,
 14    size: usize,
 15    rows: usize,
 16    cols_in: usize,
 17    cols_out: usize,
 18    _phantom: PhantomData<B>,
 19}
 20
 21impl<D, B: Backend> 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
 39impl<D> ZnxSliceSize for MatZnxDft<D, FFT64> {
 40    fn sl(&self) -> usize {
 41        self.n() * self.cols_out()
 42    }
 43}
 44
 45impl<D, B: Backend> DataView for MatZnxDft<D, B> {
 46    type D = D;
 47    fn data(&self) -> &Self::D {
 48        &self.data
 49    }
 50}
 51
 52impl<D, B: Backend> DataViewMut for MatZnxDft<D, B> {
 53    fn data_mut(&mut self) -> &mut Self::D {
 54        &mut self.data
 55    }
 56}
 57
 58impl<D: AsRef<[u8]>> ZnxView for MatZnxDft<D, FFT64> {
 59    type Scalar = f64;
 60}
 61
 62impl<D, B: Backend> MatZnxDft<D, B> {
 63    pub fn cols_in(&self) -> usize {
 64        self.cols_in
 65    }
 66
 67    pub fn cols_out(&self) -> usize {
 68        self.cols_out
 69    }
 70}
 71
 72impl<D: From<Vec<u8>>, B: Backend> MatZnxDft<D, B> {
 73    pub(crate) fn bytes_of(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 74        unsafe {
 75            crate::ffi::vmp::bytes_of_vmp_pmat(
 76                module.ptr,
 77                (rows * cols_in) as u64,
 78                (size * cols_out) as u64,
 79            ) as usize
 80        }
 81    }
 82
 83    pub(crate) fn new(module: &Module<B>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
 84        let data: Vec<u8> = alloc_aligned(Self::bytes_of(module, rows, cols_in, cols_out, size));
 85        Self {
 86            data: data.into(),
 87            n: module.n(),
 88            size,
 89            rows,
 90            cols_in,
 91            cols_out,
 92            _phantom: PhantomData,
 93        }
 94    }
 95
 96    pub(crate) fn new_from_bytes(
 97        module: &Module<B>,
 98        rows: usize,
 99        cols_in: usize,
100        cols_out: usize,
101        size: usize,
102        bytes: impl Into<Vec<u8>>,
103    ) -> Self {
104        let data: Vec<u8> = bytes.into();
105        assert!(data.len() == Self::bytes_of(module, rows, cols_in, cols_out, size));
106        Self {
107            data: data.into(),
108            n: module.n(),
109            size,
110            rows,
111            cols_in,
112            cols_out,
113            _phantom: PhantomData,
114        }
115    }
116}
117
118impl<D: AsRef<[u8]>> MatZnxDft<D, FFT64> {
119    /// Returns a copy of the backend array at index (i, j) of the [MatZnxDft].
120    ///
121    /// # Arguments
122    ///
123    /// * `row`: row index (i).
124    /// * `col`: col index (j).
125    #[allow(dead_code)]
126    fn at(&self, row: usize, col: usize) -> Vec<f64> {
127        let n: usize = self.n();
128
129        let mut res: Vec<f64> = alloc_aligned(n);
130
131        if n < 8 {
132            res.copy_from_slice(&self.raw()[(row + col * self.rows()) * n..(row + col * self.rows()) * (n + 1)]);
133        } else {
134            (0..n >> 3).for_each(|blk| {
135                res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
136            });
137        }
138
139        res
140    }
141
142    #[allow(dead_code)]
143    fn at_block(&self, row: usize, col: usize, blk: usize) -> &[f64] {
144        let nrows: usize = self.rows();
145        let nsize: usize = self.size();
146        if col == (nsize - 1) && (nsize & 1 == 1) {
147            &self.raw()[blk * nrows * nsize * 8 + col * nrows * 8 + row * 8..]
148        } else {
149            &self.raw()[blk * nrows * nsize * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
150        }
151    }
152}
153
154pub type MatZnxDftOwned<B> = MatZnxDft<Vec<u8>, B>;
155
156pub trait MatZnxDftToRef<B: Backend> {
157    fn to_ref(&self) -> MatZnxDft<&[u8], B>;
158}
159
160pub trait MatZnxDftToMut<B: Backend> {
161    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B>;
162}
163
164impl<B: Backend> MatZnxDftToMut<B> for MatZnxDft<Vec<u8>, B> {
165    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
166        MatZnxDft {
167            data: self.data.as_mut_slice(),
168            n: self.n,
169            rows: self.rows,
170            cols_in: self.cols_in,
171            cols_out: self.cols_out,
172            size: self.size,
173            _phantom: PhantomData,
174        }
175    }
176}
177
178impl<B: Backend> MatZnxDftToRef<B> for MatZnxDft<Vec<u8>, B> {
179    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
180        MatZnxDft {
181            data: self.data.as_slice(),
182            n: self.n,
183            rows: self.rows,
184            cols_in: self.cols_in,
185            cols_out: self.cols_out,
186            size: self.size,
187            _phantom: PhantomData,
188        }
189    }
190}
191
192impl<B: Backend> MatZnxDftToMut<B> for MatZnxDft<&mut [u8], B> {
193    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
194        MatZnxDft {
195            data: self.data,
196            n: self.n,
197            rows: self.rows,
198            cols_in: self.cols_in,
199            cols_out: self.cols_out,
200            size: self.size,
201            _phantom: PhantomData,
202        }
203    }
204}
205
206impl<B: Backend> MatZnxDftToRef<B> for MatZnxDft<&mut [u8], B> {
207    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
208        MatZnxDft {
209            data: self.data,
210            n: self.n,
211            rows: self.rows,
212            cols_in: self.cols_in,
213            cols_out: self.cols_out,
214            size: self.size,
215            _phantom: PhantomData,
216        }
217    }
218}
219
220impl<B: Backend> MatZnxDftToRef<B> for MatZnxDft<&[u8], B> {
221    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
222        MatZnxDft {
223            data: self.data,
224            n: self.n,
225            rows: self.rows,
226            cols_in: self.cols_in,
227            cols_out: self.cols_out,
228            size: self.size,
229            _phantom: PhantomData,
230        }
231    }
232}