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