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