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