1use std::marker::PhantomData;
  2
  3use crate::{
  4    alloc_aligned,
  5    hal::{
  6        api::{DataView, DataViewMut, ZnxInfos},
  7        layouts::{Backend, Data, DataMut, DataRef},
  8    },
  9};
 10
 11#[derive(PartialEq, Eq)]
 12pub struct VmpPMat<D: Data, B: Backend> {
 13    data: D,
 14    n: usize,
 15    size: usize,
 16    rows: usize,
 17    cols_in: usize,
 18    cols_out: usize,
 19    _phantom: PhantomData<B>,
 20}
 21
 22impl<D: Data, B: Backend> ZnxInfos for VmpPMat<D, B> {
 23    fn cols(&self) -> usize {
 24        self.cols_in
 25    }
 26
 27    fn rows(&self) -> usize {
 28        self.rows
 29    }
 30
 31    fn n(&self) -> usize {
 32        self.n
 33    }
 34
 35    fn size(&self) -> usize {
 36        self.size
 37    }
 38}
 39
 40impl<D: Data, B: Backend> DataView for VmpPMat<D, B> {
 41    type D = D;
 42    fn data(&self) -> &Self::D {
 43        &self.data
 44    }
 45}
 46
 47impl<D: Data, B: Backend> DataViewMut for VmpPMat<D, B> {
 48    fn data_mut(&mut self) -> &mut Self::D {
 49        &mut self.data
 50    }
 51}
 52
 53impl<D: Data, B: Backend> VmpPMat<D, B> {
 54    pub fn cols_in(&self) -> usize {
 55        self.cols_in
 56    }
 57
 58    pub fn cols_out(&self) -> usize {
 59        self.cols_out
 60    }
 61}
 62
 63pub trait VmpPMatBytesOf {
 64    fn vmp_pmat_bytes_of(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize;
 65}
 66
 67impl<D: DataRef + From<Vec<u8>>, B: Backend> VmpPMat<D, B>
 68where
 69    B: VmpPMatBytesOf,
 70{
 71    pub(crate) fn alloc(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
 72        let data: Vec<u8> = alloc_aligned(B::vmp_pmat_bytes_of(n, rows, cols_in, cols_out, size));
 73        Self {
 74            data: data.into(),
 75            n,
 76            size,
 77            rows,
 78            cols_in,
 79            cols_out,
 80            _phantom: PhantomData,
 81        }
 82    }
 83
 84    pub(crate) fn from_bytes(
 85        n: usize,
 86        rows: usize,
 87        cols_in: usize,
 88        cols_out: usize,
 89        size: usize,
 90        bytes: impl Into<Vec<u8>>,
 91    ) -> Self {
 92        let data: Vec<u8> = bytes.into();
 93        assert!(data.len() == B::vmp_pmat_bytes_of(n, rows, cols_in, cols_out, size));
 94        Self {
 95            data: data.into(),
 96            n,
 97            size,
 98            rows,
 99            cols_in,
100            cols_out,
101            _phantom: PhantomData,
102        }
103    }
104}
105
106pub type VmpPMatOwned<B> = VmpPMat<Vec<u8>, B>;
107pub type VmpPMatRef<'a, B> = VmpPMat<&'a [u8], B>;
108
109pub trait VmpPMatToRef<B: Backend> {
110    fn to_ref(&self) -> VmpPMat<&[u8], B>;
111}
112
113impl<D: DataRef, B: Backend> VmpPMatToRef<B> for VmpPMat<D, B> {
114    fn to_ref(&self) -> VmpPMat<&[u8], B> {
115        VmpPMat {
116            data: self.data.as_ref(),
117            n: self.n,
118            rows: self.rows,
119            cols_in: self.cols_in,
120            cols_out: self.cols_out,
121            size: self.size,
122            _phantom: std::marker::PhantomData,
123        }
124    }
125}
126
127pub trait VmpPMatToMut<B: Backend> {
128    fn to_mut(&mut self) -> VmpPMat<&mut [u8], B>;
129}
130
131impl<D: DataMut, B: Backend> VmpPMatToMut<B> for VmpPMat<D, B> {
132    fn to_mut(&mut self) -> VmpPMat<&mut [u8], B> {
133        VmpPMat {
134            data: self.data.as_mut(),
135            n: self.n,
136            rows: self.rows,
137            cols_in: self.cols_in,
138            cols_out: self.cols_out,
139            size: self.size,
140            _phantom: std::marker::PhantomData,
141        }
142    }
143}
144
145impl<D: Data, B: Backend> VmpPMat<D, B> {
146    pub(crate) fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
147        Self {
148            data,
149            n,
150            rows,
151            cols_in,
152            cols_out,
153            size,
154            _phantom: PhantomData,
155        }
156    }
157}