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