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};
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 n(&self) -> usize {
65 self.n
66 }
67
68 pub fn cols(&self) -> usize {
69 self.cols
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
105pub trait VecZnxDftOps {
106 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
107 fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft;
108
109 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
110 ///
111 /// # Arguments
112 ///
113 /// * `cols`: the number of cols of the [VecZnxDft].
114 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
115 ///
116 /// # Panics
117 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
118 fn new_vec_znx_dft_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
119
120 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
121 ///
122 /// # Arguments
123 ///
124 /// * `cols`: the number of cols of the [VecZnxDft].
125 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
126 ///
127 /// # Panics
128 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
129 fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize;
130
131 /// Returns the minimum number of bytes necessary to allocate
132 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
133 fn vec_znx_idft_tmp_bytes(&self) -> usize;
134
135 /// b <- IDFT(a), uses a as scratch space.
136 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_cols: usize);
137
138 fn vec_znx_idft(
139 &self,
140 b: &mut VecZnxBig,
141 a: &VecZnxDft,
142 a_cols: usize,
143 tmp_bytes: &mut [u8],
144 );
145
146 fn vec_znx_dft(&self, b: &mut VecZnxDft, a: &VecZnx, a_cols: usize);
147}
148
149impl VecZnxDftOps for Module {
150 fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft {
151 let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_dft(cols));
152 let ptr: *mut u8 = data.as_mut_ptr();
153 VecZnxDft {
154 data: data,
155 ptr: ptr,
156 n: self.n(),
157 cols: cols,
158 backend: self.backend(),
159 }
160 }
161
162 fn new_vec_znx_dft_from_bytes(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
163 debug_assert!(
164 tmp_bytes.len() >= <Module as VecZnxDftOps>::bytes_of_vec_znx_dft(self, cols),
165 "invalid bytes: bytes.len()={} < bytes_of_vec_znx_dft={}",
166 tmp_bytes.len(),
167 <Module as VecZnxDftOps>::bytes_of_vec_znx_dft(self, cols)
168 );
169 #[cfg(debug_assertions)]
170 {
171 assert_alignement(tmp_bytes.as_ptr())
172 }
173 VecZnxDft::from_bytes(self, cols, tmp_bytes)
174 }
175
176 fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize {
177 unsafe { bytes_of_vec_znx_dft(self.ptr, cols as u64) as usize }
178 }
179
180 fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_cols: usize) {
181 debug_assert!(
182 b.cols() >= a_cols,
183 "invalid c_vector: b_vector.cols()={} < a_cols={}",
184 b.cols(),
185 a_cols
186 );
187 unsafe {
188 vec_znx_dft::vec_znx_idft_tmp_a(
189 self.ptr,
190 b.ptr as *mut vec_znx_big_t,
191 b.cols() as u64,
192 a.ptr as *mut vec_znx_dft_t,
193 a_cols as u64,
194 )
195 }
196 }
197
198 fn vec_znx_idft_tmp_bytes(&self) -> usize {
199 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
200 }
201
202 /// b <- DFT(a)
203 ///
204 /// # Panics
205 /// If b.cols < a_cols
206 fn vec_znx_dft(&self, b: &mut VecZnxDft, a: &VecZnx, a_cols: usize) {
207 debug_assert!(
208 b.cols() >= a_cols,
209 "invalid a_cols: b.cols()={} < a_cols={}",
210 b.cols(),
211 a_cols
212 );
213 unsafe {
214 vec_znx_dft::vec_znx_dft(
215 self.ptr,
216 b.ptr as *mut vec_znx_dft_t,
217 b.cols() as u64,
218 a.as_ptr(),
219 a_cols as u64,
220 a.n() as u64,
221 )
222 }
223 }
224
225 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
226 fn vec_znx_idft(
227 &self,
228 b: &mut VecZnxBig,
229 a: &VecZnxDft,
230 a_cols: usize,
231 tmp_bytes: &mut [u8],
232 ) {
233 debug_assert!(
234 b.cols() >= a_cols,
235 "invalid c_vector: b.cols()={} < a_cols={}",
236 b.cols(),
237 a_cols
238 );
239 debug_assert!(
240 a.cols() >= a_cols,
241 "invalid c_vector: a.cols()={} < a_cols={}",
242 a.cols(),
243 a_cols
244 );
245 debug_assert!(
246 tmp_bytes.len() >= <Module as VecZnxDftOps>::vec_znx_idft_tmp_bytes(self),
247 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
248 tmp_bytes.len(),
249 <Module as VecZnxDftOps>::vec_znx_idft_tmp_bytes(self)
250 );
251 #[cfg(debug_assertions)]
252 {
253 assert_alignement(tmp_bytes.as_ptr())
254 }
255 unsafe {
256 vec_znx_dft::vec_znx_idft(
257 self.ptr,
258 b.ptr as *mut vec_znx_big_t,
259 a.cols() as u64,
260 a.ptr as *const vec_znx_dft_t,
261 a_cols as u64,
262 tmp_bytes.as_mut_ptr(),
263 )
264 }
265 }
266}