poulpy/base2k/src/vec_znx_dft.rs
base2k/src/vec_znx_dft.rsRaw
1use crate::ffi::vec_znx_big::vec_znx_big_t;
2use crate::ffi::vec_znx_dft;
3use crate::ffi::vec_znx_dft::{bytes_of_vec_znx_dft, vec_znx_dft_t};
4use crate::{Backend, FFT64, Infos, Module, VecZnxBig, VecZnxLayout, assert_alignement};
5use crate::{VecZnx, alloc_aligned};
6use std::marker::PhantomData;
7
8pub struct VecZnxDft<B: Backend> {
9 pub data: Vec<u8>,
10 pub ptr: *mut u8,
11 pub n: usize,
12 pub cols: usize,
13 pub limbs: usize,
14 pub _marker: PhantomData<B>,
15}
16
17impl VecZnxDft<FFT64> {
18 pub fn new(module: &Module<FFT64>, cols: usize, limbs: usize) -> Self {
19 #[cfg(debug_assertions)]
20 {
21 assert!(cols > 0);
22 assert!(limbs > 0);
23 }
24 let mut data: Vec<u8> = alloc_aligned::<u8>(module.bytes_of_vec_znx_dft(cols, limbs));
25 let ptr: *mut u8 = data.as_mut_ptr();
26 Self {
27 data: data,
28 ptr: ptr,
29 n: module.n(),
30 limbs: limbs,
31 cols: cols,
32 _marker: PhantomData,
33 }
34 }
35
36 fn bytes_of(module: &Module<FFT64>, cols: usize, limbs: usize) -> usize {
37 unsafe { bytes_of_vec_znx_dft(module.ptr, limbs as u64) as usize * cols }
38 }
39
40 /// Returns a new [VecZnxDft] with the provided data as backing array.
41 /// User must ensure that data is properly alligned and that
42 /// the size of data is at least equal to [Module::bytes_of_vec_znx_dft].
43 pub fn from_bytes(module: &Module<FFT64>, cols: usize, limbs: usize, bytes: &mut [u8]) -> Self {
44 #[cfg(debug_assertions)]
45 {
46 assert!(cols > 0);
47 assert!(limbs > 0);
48 assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(cols, limbs));
49 assert_alignement(bytes.as_ptr())
50 }
51 unsafe {
52 Self {
53 data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
54 ptr: bytes.as_mut_ptr(),
55 n: module.n(),
56 cols: cols,
57 limbs: limbs,
58 _marker: PhantomData,
59 }
60 }
61 }
62
63 pub fn from_bytes_borrow(module: &Module<FFT64>, cols: usize, limbs: usize, bytes: &mut [u8]) -> Self {
64 #[cfg(debug_assertions)]
65 {
66 assert!(cols > 0);
67 assert!(limbs > 0);
68 assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(cols, limbs));
69 assert_alignement(bytes.as_ptr());
70 }
71 Self {
72 data: Vec::new(),
73 ptr: bytes.as_mut_ptr(),
74 n: module.n(),
75 cols: cols,
76 limbs: limbs,
77 _marker: PhantomData,
78 }
79 }
80
81 /// Cast a [VecZnxDft] into a [VecZnxBig].
82 /// The returned [VecZnxBig] shares the backing array
83 /// with the original [VecZnxDft].
84 pub fn as_vec_znx_big(&mut self) -> VecZnxBig<FFT64> {
85 VecZnxBig::<FFT64> {
86 data: Vec::new(),
87 ptr: self.ptr,
88 n: self.n,
89 cols: self.cols,
90 limbs: self.limbs,
91 _marker: PhantomData,
92 }
93 }
94
95 pub fn print(&self, n: usize) {
96 (0..self.limbs()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]));
97 }
98}
99
100impl<B: Backend> Infos for VecZnxDft<B> {
101 fn n(&self) -> usize {
102 self.n
103 }
104
105 fn log_n(&self) -> usize {
106 (usize::BITS - (self.n() - 1).leading_zeros()) as _
107 }
108
109 fn rows(&self) -> usize {
110 1
111 }
112
113 fn cols(&self) -> usize {
114 self.cols
115 }
116
117 fn limbs(&self) -> usize {
118 self.limbs
119 }
120
121 fn poly_count(&self) -> usize {
122 self.cols * self.limbs
123 }
124}
125
126impl VecZnxLayout for VecZnxDft<FFT64> {
127 type Scalar = f64;
128
129 fn as_ptr(&self) -> *const Self::Scalar {
130 self.ptr as *const Self::Scalar
131 }
132
133 fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
134 self.ptr as *mut Self::Scalar
135 }
136}
137
138pub trait VecZnxDftOps<B: Backend> {
139 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
140 fn new_vec_znx_dft(&self, cols: usize, limbs: usize) -> VecZnxDft<B>;
141
142 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
143 ///
144 /// Behavior: takes ownership of the backing array.
145 ///
146 /// # Arguments
147 ///
148 /// * `cols`: the number of cols of the [VecZnxDft].
149 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
150 ///
151 /// # Panics
152 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
153 fn new_vec_znx_dft_from_bytes(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
154
155 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
156 ///
157 /// Behavior: the backing array is only borrowed.
158 ///
159 /// # Arguments
160 ///
161 /// * `cols`: the number of cols of the [VecZnxDft].
162 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
163 ///
164 /// # Panics
165 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
166 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
167
168 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
169 ///
170 /// # Arguments
171 ///
172 /// * `cols`: the number of cols of the [VecZnxDft].
173 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
174 ///
175 /// # Panics
176 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
177 fn bytes_of_vec_znx_dft(&self, cols: usize, limbs: usize) -> usize;
178
179 /// Returns the minimum number of bytes necessary to allocate
180 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
181 fn vec_znx_idft_tmp_bytes(&self) -> usize;
182
183 /// b <- IDFT(a), uses a as scratch space.
184 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig<B>, a: &mut VecZnxDft<B>);
185
186 fn vec_znx_idft(&self, b: &mut VecZnxBig<B>, a: &VecZnxDft<B>, tmp_bytes: &mut [u8]);
187
188 fn vec_znx_dft(&self, b: &mut VecZnxDft<B>, a: &VecZnx);
189
190 fn vec_znx_dft_automorphism(&self, k: i64, b: &mut VecZnxDft<B>, a: &VecZnxDft<B>);
191
192 fn vec_znx_dft_automorphism_inplace(&self, k: i64, a: &mut VecZnxDft<B>, tmp_bytes: &mut [u8]);
193
194 fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize;
195}
196
197impl VecZnxDftOps<FFT64> for Module<FFT64> {
198 fn new_vec_znx_dft(&self, cols: usize, limbs: usize) -> VecZnxDft<FFT64> {
199 VecZnxDft::<FFT64>::new(&self, cols, limbs)
200 }
201
202 fn new_vec_znx_dft_from_bytes(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxDft<FFT64> {
203 VecZnxDft::from_bytes(self, cols, limbs, tmp_bytes)
204 }
205
206 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxDft<FFT64> {
207 VecZnxDft::from_bytes_borrow(self, cols, limbs, tmp_bytes)
208 }
209
210 fn bytes_of_vec_znx_dft(&self, cols: usize, limbs: usize) -> usize {
211 VecZnxDft::bytes_of(&self, cols, limbs)
212 }
213
214 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig<FFT64>, a: &mut VecZnxDft<FFT64>) {
215 unsafe {
216 vec_znx_dft::vec_znx_idft_tmp_a(
217 self.ptr,
218 b.ptr as *mut vec_znx_big_t,
219 b.poly_count() as u64,
220 a.ptr as *mut vec_znx_dft_t,
221 a.poly_count() as u64,
222 )
223 }
224 }
225
226 fn vec_znx_idft_tmp_bytes(&self) -> usize {
227 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
228 }
229
230 /// b <- DFT(a)
231 ///
232 /// # Panics
233 /// If b.cols < a_cols
234 fn vec_znx_dft(&self, b: &mut VecZnxDft<FFT64>, a: &VecZnx) {
235 unsafe {
236 vec_znx_dft::vec_znx_dft(
237 self.ptr,
238 b.ptr as *mut vec_znx_dft_t,
239 b.limbs() as u64,
240 a.as_ptr(),
241 a.limbs() as u64,
242 (a.n() * a.cols()) as u64,
243 )
244 }
245 }
246
247 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
248 fn vec_znx_idft(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxDft<FFT64>, tmp_bytes: &mut [u8]) {
249 #[cfg(debug_assertions)]
250 {
251 assert!(
252 tmp_bytes.len() >= Self::vec_znx_idft_tmp_bytes(self),
253 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
254 tmp_bytes.len(),
255 Self::vec_znx_idft_tmp_bytes(self)
256 );
257 assert_alignement(tmp_bytes.as_ptr())
258 }
259 unsafe {
260 vec_znx_dft::vec_znx_idft(
261 self.ptr,
262 b.ptr as *mut vec_znx_big_t,
263 b.poly_count() as u64,
264 a.ptr as *const vec_znx_dft_t,
265 a.poly_count() as u64,
266 tmp_bytes.as_mut_ptr(),
267 )
268 }
269 }
270
271 fn vec_znx_dft_automorphism(&self, k: i64, b: &mut VecZnxDft<FFT64>, a: &VecZnxDft<FFT64>) {
272 unsafe {
273 vec_znx_dft::vec_znx_dft_automorphism(
274 self.ptr,
275 k,
276 b.ptr as *mut vec_znx_dft_t,
277 b.poly_count() as u64,
278 a.ptr as *const vec_znx_dft_t,
279 a.poly_count() as u64,
280 [0u8; 0].as_mut_ptr(),
281 );
282 }
283 }
284
285 fn vec_znx_dft_automorphism_inplace(&self, k: i64, a: &mut VecZnxDft<FFT64>, tmp_bytes: &mut [u8]) {
286 #[cfg(debug_assertions)]
287 {
288 assert!(
289 tmp_bytes.len() >= Self::vec_znx_dft_automorphism_tmp_bytes(self),
290 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_dft_automorphism_tmp_bytes()={}",
291 tmp_bytes.len(),
292 Self::vec_znx_dft_automorphism_tmp_bytes(self)
293 );
294 assert_alignement(tmp_bytes.as_ptr())
295 }
296 println!("{}", a.poly_count());
297 unsafe {
298 vec_znx_dft::vec_znx_dft_automorphism(
299 self.ptr,
300 k,
301 a.ptr as *mut vec_znx_dft_t,
302 a.poly_count() as u64,
303 a.ptr as *const vec_znx_dft_t,
304 a.poly_count() as u64,
305 tmp_bytes.as_mut_ptr(),
306 );
307 }
308 }
309
310 fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize {
311 unsafe { vec_znx_dft::vec_znx_dft_automorphism_tmp_bytes(self.ptr) as usize }
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use crate::{FFT64, Module, Sampling, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxLayout, VecZnxOps, alloc_aligned};
318 use itertools::izip;
319 use sampling::source::Source;
320
321 #[test]
322 fn test_automorphism_dft() {
323 let n: usize = 8;
324 let module: Module<FFT64> = Module::<FFT64>::new(n);
325
326 let limbs: usize = 2;
327 let log_base2k: usize = 17;
328 let mut a: VecZnx = module.new_vec_znx(1, limbs);
329 let mut a_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs);
330 let mut b_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs);
331
332 let mut source: Source = Source::new([0u8; 32]);
333 module.fill_uniform(log_base2k, &mut a, 0, limbs, &mut source);
334
335 let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vec_znx_dft_automorphism_tmp_bytes());
336
337 let p: i64 = -5;
338
339 // a_dft <- DFT(a)
340 module.vec_znx_dft(&mut a_dft, &a);
341
342 // a_dft <- AUTO(a_dft)
343 module.vec_znx_dft_automorphism_inplace(p, &mut a_dft, &mut tmp_bytes);
344
345 // a <- AUTO(a)
346 module.vec_znx_automorphism_inplace(p, &mut a);
347
348 // b_dft <- DFT(AUTO(a))
349 module.vec_znx_dft(&mut b_dft, &a);
350
351 let a_f64: &[f64] = a_dft.raw();
352 let b_f64: &[f64] = b_dft.raw();
353 izip!(a_f64.iter(), b_f64.iter()).for_each(|(ai, bi)| {
354 assert!((ai - bi).abs() <= 1e-9, "{:+e} > 1e-9", (ai - bi).abs());
355 });
356
357 module.free()
358 }
359}