poulpy/backend/src/vec_znx_dft_ops.rs
backend/src/vec_znx_dft_ops.rsRaw
1use crate::ffi::{vec_znx_big, vec_znx_dft};
2use crate::vec_znx_dft::bytes_of_vec_znx_dft;
3use crate::znx_base::ZnxInfos;
4use crate::{
5 Backend, Scratch, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef,
6 ZnxSliceSize,
7};
8use crate::{FFT64, Module, ZnxView, ZnxViewMut, ZnxZero};
9use std::cmp::min;
10
11pub trait VecZnxDftAlloc<B: Backend> {
12 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
13 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B>;
14
15 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
16 ///
17 /// Behavior: takes ownership of the backing array.
18 ///
19 /// # Arguments
20 ///
21 /// * `cols`: the number of cols of the [VecZnxDft].
22 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
23 ///
24 /// # Panics
25 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
26 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B>;
27
28 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
29 ///
30 /// # Arguments
31 ///
32 /// * `cols`: the number of cols of the [VecZnxDft].
33 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
34 ///
35 /// # Panics
36 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
37 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
38}
39
40pub trait VecZnxDftOps<B: Backend> {
41 /// Returns the minimum number of bytes necessary to allocate
42 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
43 fn vec_znx_idft_tmp_bytes(&self) -> usize;
44
45 fn vec_znx_dft_add<R, A, D>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
46 where
47 R: VecZnxDftToMut<B>,
48 A: VecZnxDftToRef<B>,
49 D: VecZnxDftToRef<B>;
50
51 fn vec_znx_dft_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
52 where
53 R: VecZnxDftToMut<B>,
54 A: VecZnxDftToRef<B>;
55
56 fn vec_znx_dft_copy<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
57 where
58 R: VecZnxDftToMut<B>,
59 A: VecZnxDftToRef<B>;
60
61 /// b <- IDFT(a), uses a as scratch space.
62 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
63 where
64 R: VecZnxBigToMut<B>,
65 A: VecZnxDftToMut<B>;
66
67 /// Consumes a to return IDFT(a) in big coeff space.
68 fn vec_znx_idft_consume<D>(&self, a: VecZnxDft<D, B>) -> VecZnxBig<D, FFT64>
69 where
70 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>;
71
72 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
73 where
74 R: VecZnxBigToMut<B>,
75 A: VecZnxDftToRef<B>;
76
77 fn vec_znx_dft<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
78 where
79 R: VecZnxDftToMut<FFT64>,
80 A: VecZnxToRef;
81}
82
83impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
84 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B> {
85 VecZnxDftOwned::new(&self, cols, size)
86 }
87
88 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B> {
89 VecZnxDftOwned::new_from_bytes(self, cols, size, bytes)
90 }
91
92 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
93 bytes_of_vec_znx_dft(self, cols, size)
94 }
95}
96
97impl VecZnxDftOps<FFT64> for Module<FFT64> {
98 fn vec_znx_dft_add<R, A, D>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
99 where
100 R: VecZnxDftToMut<FFT64>,
101 A: VecZnxDftToRef<FFT64>,
102 D: VecZnxDftToRef<FFT64>,
103 {
104 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
105 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
106 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
107
108 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
109
110 unsafe {
111 (0..min_size).for_each(|j| {
112 vec_znx_dft::vec_dft_add(
113 self.ptr,
114 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
115 1,
116 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
117 1,
118 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
119 1,
120 );
121 });
122 }
123 (min_size..res_mut.size()).for_each(|j| {
124 res_mut.zero_at(res_col, j);
125 })
126 }
127
128 fn vec_znx_dft_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
129 where
130 R: VecZnxDftToMut<FFT64>,
131 A: VecZnxDftToRef<FFT64>,
132 {
133 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
134 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
135
136 let min_size: usize = res_mut.size().min(a_ref.size());
137
138 unsafe {
139 (0..min_size).for_each(|j| {
140 vec_znx_dft::vec_dft_add(
141 self.ptr,
142 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
143 1,
144 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
145 1,
146 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
147 1,
148 );
149 });
150 }
151 }
152
153 fn vec_znx_dft_copy<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
154 where
155 R: VecZnxDftToMut<FFT64>,
156 A: VecZnxDftToRef<FFT64>,
157 {
158 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
159 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
160
161 let steps: usize = (a_ref.size() + step - 1) / step;
162 let min_steps: usize = min(res_mut.size(), steps);
163
164 (0..min_steps).for_each(|j| {
165 let limb: usize = offset + j * step;
166 if limb < a_ref.size() {
167 res_mut
168 .at_mut(res_col, j)
169 .copy_from_slice(a_ref.at(a_col, limb));
170 }
171 });
172 (min_steps..res_mut.size()).for_each(|j| {
173 res_mut.zero_at(res_col, j);
174 })
175 }
176
177 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
178 where
179 R: VecZnxBigToMut<FFT64>,
180 A: VecZnxDftToMut<FFT64>,
181 {
182 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
183 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
184
185 let min_size: usize = min(res_mut.size(), a_mut.size());
186
187 unsafe {
188 (0..min_size).for_each(|j| {
189 vec_znx_dft::vec_znx_idft_tmp_a(
190 self.ptr,
191 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
192 1 as u64,
193 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
194 1 as u64,
195 )
196 });
197 (min_size..res_mut.size()).for_each(|j| {
198 res_mut.zero_at(res_col, j);
199 })
200 }
201 }
202
203 fn vec_znx_idft_consume<D>(&self, mut a: VecZnxDft<D, FFT64>) -> VecZnxBig<D, FFT64>
204 where
205 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>,
206 {
207 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
208
209 unsafe {
210 // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
211 (0..a_mut.size()).for_each(|j| {
212 (0..a_mut.cols()).for_each(|i| {
213 vec_znx_dft::vec_znx_idft_tmp_a(
214 self.ptr,
215 a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
216 1 as u64,
217 a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
218 1 as u64,
219 )
220 });
221 });
222 }
223
224 a.into_big()
225 }
226
227 fn vec_znx_idft_tmp_bytes(&self) -> usize {
228 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
229 }
230
231 fn vec_znx_dft<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
232 where
233 R: VecZnxDftToMut<FFT64>,
234 A: VecZnxToRef,
235 {
236 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
237 let a_ref: crate::VecZnx<&[u8]> = a.to_ref();
238 let steps: usize = (a_ref.size() + step - 1) / step;
239 let min_steps: usize = min(res_mut.size(), steps);
240 unsafe {
241 (0..min_steps).for_each(|j| {
242 let limb: usize = offset + j * step;
243 if limb < a_ref.size() {
244 vec_znx_dft::vec_znx_dft(
245 self.ptr,
246 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
247 1 as u64,
248 a_ref.at_ptr(a_col, limb),
249 1 as u64,
250 a_ref.sl() as u64,
251 )
252 }
253 });
254 (min_steps..res_mut.size()).for_each(|j| {
255 res_mut.zero_at(res_col, j);
256 });
257 }
258 }
259
260 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
261 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
262 where
263 R: VecZnxBigToMut<FFT64>,
264 A: VecZnxDftToRef<FFT64>,
265 {
266 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
267 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
268
269 let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_idft_tmp_bytes());
270
271 let min_size: usize = min(res_mut.size(), a_ref.size());
272
273 unsafe {
274 (0..min_size).for_each(|j| {
275 vec_znx_dft::vec_znx_idft(
276 self.ptr,
277 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
278 1 as u64,
279 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
280 1 as u64,
281 tmp_bytes.as_mut_ptr(),
282 )
283 });
284 (min_size..res_mut.size()).for_each(|j| {
285 res_mut.zero_at(res_col, j);
286 });
287 }
288 }
289}