1use crate::znx_base::{GetZnxBase, ZnxAlloc, 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    _marker: PhantomData<B>,
14}
15
16impl<B: Backend> GetZnxBase for MatZnxDft<B> {
17    fn znx(&self) -> &ZnxBase {
18        &self.inner
19    }
20
21    fn znx_mut(&mut self) -> &mut ZnxBase {
22        &mut self.inner
23    }
24}
25
26impl<B: Backend> ZnxInfos for MatZnxDft<B> {}
27
28impl ZnxSliceSize for MatZnxDft<FFT64> {
29    fn sl(&self) -> usize {
30        self.n()
31    }
32}
33
34impl ZnxLayout for MatZnxDft<FFT64> {
35    type Scalar = f64;
36}
37
38impl<B: Backend> ZnxAlloc<B> for MatZnxDft<B> {
39    type Scalar = u8;
40
41    fn from_bytes_borrow(module: &Module<B>, rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
42        Self {
43            inner: ZnxBase::from_bytes_borrow(module.n(), rows, cols, size, bytes),
44            _marker: PhantomData,
45        }
46    }
47
48    fn bytes_of(module: &Module<B>, rows: usize, cols: usize, size: usize) -> usize {
49        unsafe { crate::ffi::vmp::bytes_of_vmp_pmat(module.ptr, rows as u64, size as u64) as usize * cols }
50    }
51}
52
53impl MatZnxDft<FFT64> {
54    /// Returns a copy of the backend array at index (i, j) of the [MatZnxDft].
55    ///
56    /// # Arguments
57    ///
58    /// * `row`: row index (i).
59    /// * `col`: col index (j).
60    #[allow(dead_code)]
61    fn at(&self, row: usize, col: usize) -> Vec<f64> {
62        let n: usize = self.n();
63
64        let mut res: Vec<f64> = alloc_aligned(n);
65
66        if n < 8 {
67            res.copy_from_slice(&self.raw()[(row + col * self.rows()) * n..(row + col * self.rows()) * (n + 1)]);
68        } else {
69            (0..n >> 3).for_each(|blk| {
70                res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
71            });
72        }
73
74        res
75    }
76
77    #[allow(dead_code)]
78    fn at_block(&self, row: usize, col: usize, blk: usize) -> &[f64] {
79        let nrows: usize = self.rows();
80        let nsize: usize = self.size();
81        if col == (nsize - 1) && (nsize & 1 == 1) {
82            &self.raw()[blk * nrows * nsize * 8 + col * nrows * 8 + row * 8..]
83        } else {
84            &self.raw()[blk * nrows * nsize * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
85        }
86    }
87}