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 pub fn raw(&self) -> &[f64] {
86 let ptr: *mut f64 = self.ptr as *mut f64;
87 let size: usize = self.n() * self.poly_count();
88 unsafe { &std::slice::from_raw_parts(ptr, size) }
89 }
90
91 pub fn at(&self, col_i: usize) -> &[f64] {
92 &self.raw()[col_i * self.n() * self.limbs()..(col_i + 1) * self.n() * self.limbs()]
93 }
94
95 pub fn raw_mut(&mut self) -> &mut [f64] {
96 let ptr: *mut f64 = self.ptr as *mut f64;
97 let size: usize = self.n() * self.poly_count();
98 unsafe { std::slice::from_raw_parts_mut(ptr, size) }
99 }
100
101 pub fn at_mut(&mut self, col_i: usize) -> &mut [f64] {
102 let n: usize = self.n();
103 let limbs:usize = self.limbs();
104 &mut self.raw_mut()[col_i * n * limbs..(col_i + 1) * n * limbs]
105 }
106}
107
108impl<B: Backend> Infos for VecZnxDft<B> {
109 fn n(&self) -> usize {
110 self.n
111 }
112
113 fn log_n(&self) -> usize {
114 (usize::BITS - (self.n() - 1).leading_zeros()) as _
115 }
116
117 fn rows(&self) -> usize {
118 1
119 }
120
121 fn cols(&self) -> usize {
122 self.cols
123 }
124
125 fn limbs(&self) -> usize {
126 self.limbs
127 }
128
129 fn poly_count(&self) -> usize {
130 self.cols * self.limbs
131 }
132}
133
134pub trait VecZnxDftOps<B: Backend> {
135 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
136 fn new_vec_znx_dft(&self, cols: usize, limbs: usize) -> VecZnxDft<B>;
137
138 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
139 ///
140 /// Behavior: takes ownership of the backing array.
141 ///
142 /// # Arguments
143 ///
144 /// * `cols`: the number of cols of the [VecZnxDft].
145 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
146 ///
147 /// # Panics
148 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
149 fn new_vec_znx_dft_from_bytes(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
150
151 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
152 ///
153 /// Behavior: the backing array is only borrowed.
154 ///
155 /// # Arguments
156 ///
157 /// * `cols`: the number of cols of the [VecZnxDft].
158 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
159 ///
160 /// # Panics
161 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
162 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
163
164 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
165 ///
166 /// # Arguments
167 ///
168 /// * `cols`: the number of cols of the [VecZnxDft].
169 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
170 ///
171 /// # Panics
172 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
173 fn bytes_of_vec_znx_dft(&self, cols: usize, limbs: usize) -> usize;
174
175 /// Returns the minimum number of bytes necessary to allocate
176 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
177 fn vec_znx_idft_tmp_bytes(&self) -> usize;
178
179 /// b <- IDFT(a), uses a as scratch space.
180 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig<B>, a: &mut VecZnxDft<B>);
181
182 fn vec_znx_idft(&self, b: &mut VecZnxBig<B>, a: &VecZnxDft<B>, tmp_bytes: &mut [u8]);
183
184 fn vec_znx_dft(&self, b: &mut VecZnxDft<B>, a: &VecZnx);
185
186 fn vec_znx_dft_automorphism(&self, k: i64, b: &mut VecZnxDft<B>, a: &VecZnxDft<B>);
187
188 fn vec_znx_dft_automorphism_inplace(&self, k: i64, a: &mut VecZnxDft<B>, tmp_bytes: &mut [u8]);
189
190 fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize;
191}
192
193impl VecZnxDftOps<FFT64> for Module<FFT64> {
194 fn new_vec_znx_dft(&self, cols: usize, limbs: usize) -> VecZnxDft<FFT64> {
195 VecZnxDft::<FFT64>::new(&self, cols, limbs)
196 }
197
198 fn new_vec_znx_dft_from_bytes(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxDft<FFT64> {
199 VecZnxDft::from_bytes(self, cols, limbs, tmp_bytes)
200 }
201
202 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxDft<FFT64> {
203 VecZnxDft::from_bytes_borrow(self, cols, limbs, tmp_bytes)
204 }
205
206 fn bytes_of_vec_znx_dft(&self, cols: usize, limbs: usize) -> usize {
207 unsafe { bytes_of_vec_znx_dft(self.ptr, limbs as u64) as usize * cols }
208 }
209
210 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig<FFT64>, a: &mut VecZnxDft<FFT64>) {
211 unsafe {
212 vec_znx_dft::vec_znx_idft_tmp_a(
213 self.ptr,
214 b.ptr as *mut vec_znx_big_t,
215 b.poly_count() as u64,
216 a.ptr as *mut vec_znx_dft_t,
217 a.poly_count() as u64,
218 )
219 }
220 }
221
222 fn vec_znx_idft_tmp_bytes(&self) -> usize {
223 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
224 }
225
226 /// b <- DFT(a)
227 ///
228 /// # Panics
229 /// If b.cols < a_cols
230 fn vec_znx_dft(&self, b: &mut VecZnxDft<FFT64>, a: &VecZnx) {
231 unsafe {
232 vec_znx_dft::vec_znx_dft(
233 self.ptr,
234 b.ptr as *mut vec_znx_dft_t,
235 b.limbs() as u64,
236 a.as_ptr(),
237 a.limbs() as u64,
238 (a.n() * a.cols()) as u64,
239 )
240 }
241 }
242
243 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
244 fn vec_znx_idft(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxDft<FFT64>, tmp_bytes: &mut [u8]) {
245 #[cfg(debug_assertions)]
246 {
247 assert!(
248 tmp_bytes.len() >= Self::vec_znx_idft_tmp_bytes(self),
249 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
250 tmp_bytes.len(),
251 Self::vec_znx_idft_tmp_bytes(self)
252 );
253 assert_alignement(tmp_bytes.as_ptr())
254 }
255 unsafe {
256 vec_znx_dft::vec_znx_idft(
257 self.ptr,
258 b.ptr as *mut vec_znx_big_t,
259 b.poly_count() as u64,
260 a.ptr as *const vec_znx_dft_t,
261 a.poly_count() as u64,
262 tmp_bytes.as_mut_ptr(),
263 )
264 }
265 }
266
267 fn vec_znx_dft_automorphism(&self, k: i64, b: &mut VecZnxDft<FFT64>, a: &VecZnxDft<FFT64>) {
268 unsafe {
269 vec_znx_dft::vec_znx_dft_automorphism(
270 self.ptr,
271 k,
272 b.ptr as *mut vec_znx_dft_t,
273 b.poly_count() as u64,
274 a.ptr as *const vec_znx_dft_t,
275 a.poly_count() as u64,
276 [0u8; 0].as_mut_ptr(),
277 );
278 }
279 }
280
281 fn vec_znx_dft_automorphism_inplace(&self, k: i64, a: &mut VecZnxDft<FFT64>, tmp_bytes: &mut [u8]) {
282 #[cfg(debug_assertions)]
283 {
284 assert!(
285 tmp_bytes.len() >= Self::vec_znx_dft_automorphism_tmp_bytes(self),
286 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_dft_automorphism_tmp_bytes()={}",
287 tmp_bytes.len(),
288 Self::vec_znx_dft_automorphism_tmp_bytes(self)
289 );
290 assert_alignement(tmp_bytes.as_ptr())
291 }
292 unsafe {
293 vec_znx_dft::vec_znx_dft_automorphism(
294 self.ptr,
295 k,
296 a.ptr as *mut vec_znx_dft_t,
297 a.poly_count() as u64,
298 a.ptr as *const vec_znx_dft_t,
299 a.poly_count() as u64,
300 tmp_bytes.as_mut_ptr(),
301 );
302 }
303 }
304
305 fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize {
306 unsafe {
307 std::cmp::max(
308 vec_znx_dft::vec_znx_dft_automorphism_tmp_bytes(self.ptr) as usize,
309 DEFAULTALIGN,
310 )
311 }
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use crate::{FFT64, Module, Sampling, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxOps, alloc_aligned};
318 use itertools::izip;
319 use sampling::source::{Source, new_seed};
320
321 #[test]
322 fn test_automorphism_dft() {
323 let module: Module<FFT64> = Module::<FFT64>::new(128);
324
325 let limbs: usize = 2;
326 let log_base2k: usize = 17;
327 let mut a: VecZnx = module.new_vec_znx(1, limbs);
328 let mut a_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs);
329 let mut b_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs);
330
331 let mut source: Source = Source::new(new_seed());
332 module.fill_uniform(log_base2k, &mut a, 0, limbs, &mut source);
333
334 let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vec_znx_dft_automorphism_tmp_bytes());
335
336 let p: i64 = -5;
337
338
339
340 // a_dft <- DFT(a)
341 module.vec_znx_dft(&mut a_dft, &a);
342
343
344
345 // a_dft <- AUTO(a_dft)
346 module.vec_znx_dft_automorphism_inplace(p, &mut a_dft, &mut tmp_bytes);
347
348 println!("123");
349
350 // a <- AUTO(a)
351 module.vec_znx_automorphism_inplace(p, &mut a);
352
353 // b_dft <- DFT(AUTO(a))
354 module.vec_znx_dft(&mut b_dft, &a);
355
356 let a_f64: &[f64] = a_dft.raw();
357 let b_f64: &[f64] = b_dft.raw();
358 izip!(a_f64.iter(), b_f64.iter()).for_each(|(ai, bi)| {
359 assert!((ai - bi).abs() <= 1e-9, "{:+e} > 1e-9", (ai - bi).abs());
360 });
361
362 module.free()
363 }
364}