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