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 limbs: 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, limbs: usize) -> Self {
21 #[cfg(debug_assertions)]
22 {
23 assert!(cols > 0);
24 assert!(limbs > 0);
25 }
26 let mut data: Vec<Self::Scalar> = alloc_aligned(Self::bytes_of(module, cols, limbs));
27 let ptr: *mut Self::Scalar = data.as_mut_ptr();
28 Self {
29 data: data,
30 ptr: ptr,
31 n: module.n(),
32 limbs: limbs,
33 cols: cols,
34 _marker: PhantomData,
35 }
36 }
37
38 fn bytes_of(module: &Module<B>, cols: usize, limbs: usize) -> usize {
39 unsafe { bytes_of_vec_znx_dft(module.ptr, limbs 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, limbs: usize, bytes: &mut [Self::Scalar]) -> Self {
46 #[cfg(debug_assertions)]
47 {
48 assert!(cols > 0);
49 assert!(limbs > 0);
50 assert_eq!(bytes.len(), Self::bytes_of(module, cols, limbs));
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 limbs: limbs,
60 _marker: PhantomData,
61 }
62 }
63 }
64
65 fn from_bytes_borrow(module: &Module<B>, cols: usize, limbs: usize, bytes: &mut [Self::Scalar]) -> Self {
66 #[cfg(debug_assertions)]
67 {
68 assert!(cols > 0);
69 assert!(limbs > 0);
70 assert_eq!(bytes.len(), Self::bytes_of(module, cols, limbs));
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 limbs: limbs,
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 limbs: self.limbs,
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 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 ZnxLayout 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
138impl VecZnxDft<FFT64> {
139 pub fn print(&self, n: usize) {
140 (0..self.limbs()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]));
141 }
142}
143
144pub trait VecZnxDftOps<B: Backend> {
145 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
146 fn new_vec_znx_dft(&self, cols: usize, limbs: usize) -> VecZnxDft<B>;
147
148 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
149 ///
150 /// Behavior: takes ownership of the backing array.
151 ///
152 /// # Arguments
153 ///
154 /// * `cols`: the number of cols of the [VecZnxDft].
155 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
156 ///
157 /// # Panics
158 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
159 fn new_vec_znx_dft_from_bytes(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
160
161 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
162 ///
163 /// Behavior: the backing array is only borrowed.
164 ///
165 /// # Arguments
166 ///
167 /// * `cols`: the number of cols of the [VecZnxDft].
168 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
169 ///
170 /// # Panics
171 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
172 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
173
174 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
175 ///
176 /// # Arguments
177 ///
178 /// * `cols`: the number of cols of the [VecZnxDft].
179 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
180 ///
181 /// # Panics
182 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
183 fn bytes_of_vec_znx_dft(&self, cols: usize, limbs: usize) -> usize;
184
185 /// Returns the minimum number of bytes necessary to allocate
186 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
187 fn vec_znx_idft_tmp_bytes(&self) -> usize;
188
189 /// b <- IDFT(a), uses a as scratch space.
190 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig<B>, a: &mut VecZnxDft<B>);
191
192 fn vec_znx_idft(&self, b: &mut VecZnxBig<B>, a: &VecZnxDft<B>, tmp_bytes: &mut [u8]);
193
194 fn vec_znx_dft(&self, b: &mut VecZnxDft<B>, a: &VecZnx);
195
196 fn vec_znx_dft_automorphism(&self, k: i64, b: &mut VecZnxDft<B>, a: &VecZnxDft<B>);
197
198 fn vec_znx_dft_automorphism_inplace(&self, k: i64, a: &mut VecZnxDft<B>, tmp_bytes: &mut [u8]);
199
200 fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize;
201}
202
203impl VecZnxDftOps<FFT64> for Module<FFT64> {
204 fn new_vec_znx_dft(&self, cols: usize, limbs: usize) -> VecZnxDft<FFT64> {
205 VecZnxDft::<FFT64>::new(&self, cols, limbs)
206 }
207
208 fn new_vec_znx_dft_from_bytes(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxDft<FFT64> {
209 VecZnxDft::from_bytes(self, cols, limbs, tmp_bytes)
210 }
211
212 fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxDft<FFT64> {
213 VecZnxDft::from_bytes_borrow(self, cols, limbs, tmp_bytes)
214 }
215
216 fn bytes_of_vec_znx_dft(&self, cols: usize, limbs: usize) -> usize {
217 VecZnxDft::bytes_of(&self, cols, limbs)
218 }
219
220 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig<FFT64>, a: &mut VecZnxDft<FFT64>) {
221 unsafe {
222 vec_znx_dft::vec_znx_idft_tmp_a(
223 self.ptr,
224 b.ptr as *mut vec_znx_big_t,
225 b.poly_count() as u64,
226 a.ptr as *mut vec_znx_dft_t,
227 a.poly_count() as u64,
228 )
229 }
230 }
231
232 fn vec_znx_idft_tmp_bytes(&self) -> usize {
233 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
234 }
235
236 /// b <- DFT(a)
237 ///
238 /// # Panics
239 /// If b.cols < a_cols
240 fn vec_znx_dft(&self, b: &mut VecZnxDft<FFT64>, a: &VecZnx) {
241 unsafe {
242 vec_znx_dft::vec_znx_dft(
243 self.ptr,
244 b.ptr as *mut vec_znx_dft_t,
245 b.limbs() as u64,
246 a.as_ptr(),
247 a.limbs() as u64,
248 (a.n() * a.cols()) as u64,
249 )
250 }
251 }
252
253 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
254 fn vec_znx_idft(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxDft<FFT64>, tmp_bytes: &mut [u8]) {
255 #[cfg(debug_assertions)]
256 {
257 assert!(
258 tmp_bytes.len() >= Self::vec_znx_idft_tmp_bytes(self),
259 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
260 tmp_bytes.len(),
261 Self::vec_znx_idft_tmp_bytes(self)
262 );
263 assert_alignement(tmp_bytes.as_ptr())
264 }
265 unsafe {
266 vec_znx_dft::vec_znx_idft(
267 self.ptr,
268 b.ptr as *mut vec_znx_big_t,
269 b.poly_count() as u64,
270 a.ptr as *const vec_znx_dft_t,
271 a.poly_count() as u64,
272 tmp_bytes.as_mut_ptr(),
273 )
274 }
275 }
276
277 fn vec_znx_dft_automorphism(&self, k: i64, b: &mut VecZnxDft<FFT64>, a: &VecZnxDft<FFT64>) {
278 unsafe {
279 vec_znx_dft::vec_znx_dft_automorphism(
280 self.ptr,
281 k,
282 b.ptr as *mut vec_znx_dft_t,
283 b.poly_count() as u64,
284 a.ptr as *const vec_znx_dft_t,
285 a.poly_count() as u64,
286 [0u8; 0].as_mut_ptr(),
287 );
288 }
289 }
290
291 fn vec_znx_dft_automorphism_inplace(&self, k: i64, a: &mut VecZnxDft<FFT64>, tmp_bytes: &mut [u8]) {
292 #[cfg(debug_assertions)]
293 {
294 assert!(
295 tmp_bytes.len() >= Self::vec_znx_dft_automorphism_tmp_bytes(self),
296 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_dft_automorphism_tmp_bytes()={}",
297 tmp_bytes.len(),
298 Self::vec_znx_dft_automorphism_tmp_bytes(self)
299 );
300 assert_alignement(tmp_bytes.as_ptr())
301 }
302 println!("{}", a.poly_count());
303 unsafe {
304 vec_znx_dft::vec_znx_dft_automorphism(
305 self.ptr,
306 k,
307 a.ptr as *mut vec_znx_dft_t,
308 a.poly_count() as u64,
309 a.ptr as *const vec_znx_dft_t,
310 a.poly_count() as u64,
311 tmp_bytes.as_mut_ptr(),
312 );
313 }
314 }
315
316 fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize {
317 unsafe { vec_znx_dft::vec_znx_dft_automorphism_tmp_bytes(self.ptr) as usize }
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 use crate::{FFT64, Module, Sampling, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxOps, ZnxLayout, alloc_aligned};
324 use itertools::izip;
325 use sampling::source::Source;
326
327 #[test]
328 fn test_automorphism_dft() {
329 let n: usize = 8;
330 let module: Module<FFT64> = Module::<FFT64>::new(n);
331
332 let limbs: usize = 2;
333 let log_base2k: usize = 17;
334 let mut a: VecZnx = module.new_vec_znx(1, limbs);
335 let mut a_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs);
336 let mut b_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, limbs);
337
338 let mut source: Source = Source::new([0u8; 32]);
339 module.fill_uniform(log_base2k, &mut a, 0, limbs, &mut source);
340
341 let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vec_znx_dft_automorphism_tmp_bytes());
342
343 let p: i64 = -5;
344
345 // a_dft <- DFT(a)
346 module.vec_znx_dft(&mut a_dft, &a);
347
348 // a_dft <- AUTO(a_dft)
349 module.vec_znx_dft_automorphism_inplace(p, &mut a_dft, &mut tmp_bytes);
350
351 // a <- AUTO(a)
352 module.vec_znx_automorphism_inplace(p, &mut a);
353
354 // b_dft <- DFT(AUTO(a))
355 module.vec_znx_dft(&mut b_dft, &a);
356
357 let a_f64: &[f64] = a_dft.raw();
358 let b_f64: &[f64] = b_dft.raw();
359 izip!(a_f64.iter(), b_f64.iter()).for_each(|(ai, bi)| {
360 assert!((ai - bi).abs() <= 1e-9, "{:+e} > 1e-9", (ai - bi).abs());
361 });
362
363 module.free()
364 }
365}