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