poulpy/base2k/src/vec_znx_dft_ops.rs
base2k/src/vec_znx_dft_ops.rsRaw
1use crate::ffi::{vec_znx_big, vec_znx_dft};
2use crate::vec_znx_dft::bytes_of_vec_znx_dft;
3use crate::znx_base::ZnxInfos;
4use crate::{Backend, VecZnxDftOwned};
5use crate::{FFT64, Module, VecZnx, VecZnxBig, VecZnxDft, ZnxView, ZnxViewMut, ZnxZero, assert_alignement};
6use std::cmp::min;
7
8pub trait VecZnxDftAlloc<B> {
9 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
10 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B>;
11
12 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
13 ///
14 /// Behavior: takes ownership of the backing array.
15 ///
16 /// # Arguments
17 ///
18 /// * `cols`: the number of cols of the [VecZnxDft].
19 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
20 ///
21 /// # Panics
22 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
23 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B>;
24
25 // /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
26 // ///
27 // /// Behavior: the backing array is only borrowed.
28 // ///
29 // /// # Arguments
30 // ///
31 // /// * `cols`: the number of cols of the [VecZnxDft].
32 // /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
33 // ///
34 // /// # Panics
35 // /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
36 // fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
37
38 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
39 ///
40 /// # Arguments
41 ///
42 /// * `cols`: the number of cols of the [VecZnxDft].
43 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
44 ///
45 /// # Panics
46 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
47 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
48}
49
50pub trait VecZnxDftOps<DataMut, Data, B> {
51 /// Returns the minimum number of bytes necessary to allocate
52 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
53 fn vec_znx_idft_tmp_bytes(&self) -> usize;
54
55 /// b <- IDFT(a), uses a as scratch space.
56 fn vec_znx_idft_tmp_a(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &mut VecZnxDft<DataMut, B>, a_cols: usize);
57
58 fn vec_znx_idft(
59 &self,
60 res: &mut VecZnxBig<DataMut, B>,
61 res_col: usize,
62 a: &VecZnxDft<Data, B>,
63 a_col: usize,
64 tmp_bytes: &mut [u8],
65 );
66
67 fn vec_znx_dft(&self, res: &mut VecZnxDft<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
68}
69
70impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
71 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B> {
72 VecZnxDftOwned::new(&self, cols, size)
73 }
74
75 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B> {
76 VecZnxDftOwned::new_from_bytes(self, cols, size, bytes)
77 }
78
79 // fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxDft<FFT64> {
80 // VecZnxDft::from_bytes_borrow(self, 1, cols, size, bytes)
81 // }
82
83 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
84 bytes_of_vec_znx_dft(self, cols, size)
85 }
86}
87
88impl<DataMut, Data> VecZnxDftOps<DataMut, Data, FFT64> for Module<FFT64>
89where
90 DataMut: AsMut<[u8]> + AsRef<[u8]>,
91 Data: AsRef<[u8]>,
92{
93 fn vec_znx_idft_tmp_a(
94 &self,
95 res: &mut VecZnxBig<DataMut, FFT64>,
96 res_col: usize,
97 a: &mut VecZnxDft<DataMut, FFT64>,
98 a_col: usize,
99 ) {
100 let min_size: usize = min(res.size(), a.size());
101
102 unsafe {
103 (0..min_size).for_each(|j| {
104 vec_znx_dft::vec_znx_idft_tmp_a(
105 self.ptr,
106 res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
107 1 as u64,
108 a.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
109 1 as u64,
110 )
111 });
112 (min_size..res.size()).for_each(|j| {
113 res.zero_at(res_col, j);
114 })
115 }
116 }
117
118 fn vec_znx_idft_tmp_bytes(&self) -> usize {
119 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
120 }
121
122 /// b <- DFT(a)
123 ///
124 /// # Panics
125 /// If b.cols < a_cols
126 fn vec_znx_dft(&self, res: &mut VecZnxDft<DataMut, FFT64>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
127 let min_size: usize = min(res.size(), a.size());
128
129 unsafe {
130 (0..min_size).for_each(|j| {
131 vec_znx_dft::vec_znx_dft(
132 self.ptr,
133 res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
134 1 as u64,
135 a.at_ptr(a_col, j),
136 1 as u64,
137 a.sl() as u64,
138 )
139 });
140 (min_size..res.size()).for_each(|j| {
141 res.zero_at(res_col, j);
142 });
143 }
144 }
145
146 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
147 fn vec_znx_idft(
148 &self,
149 res: &mut VecZnxBig<DataMut, FFT64>,
150 res_col: usize,
151 a: &VecZnxDft<Data, FFT64>,
152 a_col: usize,
153 tmp_bytes: &mut [u8],
154 ) {
155 #[cfg(debug_assertions)]
156 {
157 assert!(
158 tmp_bytes.len() >= <Self as VecZnxDftOps<DataMut, DataMut, FFT64>>::vec_znx_idft_tmp_bytes(self),
159 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
160 tmp_bytes.len(),
161 <Self as VecZnxDftOps<DataMut, DataMut, FFT64>>::vec_znx_idft_tmp_bytes(self)
162 );
163 assert_alignement(tmp_bytes.as_ptr())
164 }
165
166 let min_size: usize = min(res.size(), a.size());
167
168 unsafe {
169 (0..min_size).for_each(|j| {
170 vec_znx_dft::vec_znx_idft(
171 self.ptr,
172 res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
173 1 as u64,
174 a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
175 1 as u64,
176 tmp_bytes.as_mut_ptr(),
177 )
178 });
179 (min_size..res.size()).for_each(|j| {
180 res.zero_at(res_col, j);
181 });
182 }
183 }
184}