poulpy/base2k/src/scalar_znx_dft.rs
base2k/src/scalar_znx_dft.rsRaw
1use std::marker::PhantomData;
2
3use crate::ffi::svp;
4use crate::znx_base::ZnxInfos;
5use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxView, alloc_aligned};
6
7pub const SCALAR_ZNX_DFT_ROWS: usize = 1;
8pub const SCALAR_ZNX_DFT_SIZE: usize = 1;
9
10pub struct ScalarZnxDft<D, B> {
11 data: D,
12 n: usize,
13 cols: usize,
14 _phantom: PhantomData<B>,
15}
16
17impl<D, B> ZnxInfos for ScalarZnxDft<D, B> {
18 fn cols(&self) -> usize {
19 self.cols
20 }
21
22 fn rows(&self) -> usize {
23 1
24 }
25
26 fn n(&self) -> usize {
27 self.n
28 }
29
30 fn size(&self) -> usize {
31 1
32 }
33
34 fn sl(&self) -> usize {
35 self.n()
36 }
37}
38
39impl<D, B> DataView for ScalarZnxDft<D, B> {
40 type D = D;
41 fn data(&self) -> &Self::D {
42 &self.data
43 }
44}
45
46impl<D, B> DataViewMut for ScalarZnxDft<D, B> {
47 fn data_mut(&mut self) -> &mut Self::D {
48 &mut self.data
49 }
50}
51
52impl<D: AsRef<[u8]>> ZnxView for ScalarZnxDft<D, FFT64> {
53 type Scalar = f64;
54}
55
56impl<D: From<Vec<u8>>, B: Backend> ScalarZnxDft<D, B> {
57 pub(crate) fn bytes_of(module: &Module<B>, cols: usize) -> usize {
58 unsafe { svp::bytes_of_svp_ppol(module.ptr) as usize * cols }
59 }
60
61 pub(crate) fn new(module: &Module<B>, cols: usize) -> Self {
62 let data = alloc_aligned::<u8>(Self::bytes_of(module, cols));
63 Self {
64 data: data.into(),
65 n: module.n(),
66 cols,
67 _phantom: PhantomData,
68 }
69 }
70
71 pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
72 let data: Vec<u8> = bytes.into();
73 assert!(data.len() == Self::bytes_of(module, cols));
74 Self {
75 data: data.into(),
76 n: module.n(),
77 cols,
78 _phantom: PhantomData,
79 }
80 }
81
82 // fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, _size: usize, bytes: &mut [u8]) -> Self {
83 // debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, _size));
84 // Self {
85 // inner: ZnxBase::from_bytes_borrow(
86 // module.n(),
87 // SCALAR_ZNX_DFT_ROWS,
88 // cols,
89 // SCALAR_ZNX_DFT_SIZE,
90 // bytes,
91 // ),
92 // _phantom: PhantomData,
93 // }
94 // }
95}
96
97pub type ScalarZnxDftOwned<B> = ScalarZnxDft<Vec<u8>, B>;