poulpy/core/src/elem.rs
core/src/elem.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftToRef, Module, Scratch, VecZnx, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps,
3 VecZnxDftToMut, VecZnxDftToRef, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxZero,
4};
5
6use crate::{
7 grlwe::GRLWECt,
8 rgsw::RGSWCt,
9 rlwe::{RLWECt, RLWECtDft},
10 utils::derive_size,
11};
12
13pub trait Infos {
14 type Inner: ZnxInfos;
15
16 fn inner(&self) -> &Self::Inner;
17
18 /// Returns the ring degree of the polynomials.
19 fn n(&self) -> usize {
20 self.inner().n()
21 }
22
23 /// Returns the base two logarithm of the ring dimension of the polynomials.
24 fn log_n(&self) -> usize {
25 self.inner().log_n()
26 }
27
28 /// Returns the number of rows.
29 fn rows(&self) -> usize {
30 self.inner().rows()
31 }
32
33 /// Returns the number of polynomials in each row.
34 fn cols(&self) -> usize {
35 self.inner().cols()
36 }
37
38 /// Returns the number of size per polynomial.
39 fn size(&self) -> usize {
40 let size: usize = self.inner().size();
41 debug_assert_eq!(size, derive_size(self.log_base2k(), self.log_k()));
42 size
43 }
44
45 /// Returns the total number of small polynomials.
46 fn poly_count(&self) -> usize {
47 self.rows() * self.cols() * self.size()
48 }
49
50 /// Returns the base 2 logarithm of the ciphertext base.
51 fn log_base2k(&self) -> usize;
52
53 /// Returns the bit precision of the ciphertext.
54 fn log_k(&self) -> usize;
55}
56
57pub trait GetRow<B: Backend> {
58 fn get_row<R>(&self, module: &Module<B>, row_i: usize, col_j: usize, res: &mut RLWECtDft<R, B>)
59 where
60 VecZnxDft<R, B>: VecZnxDftToMut<B>;
61}
62
63pub trait SetRow<B: Backend> {
64 fn set_row<A>(&mut self, module: &Module<B>, row_i: usize, col_j: usize, a: &RLWECtDft<A, B>)
65 where
66 VecZnxDft<A, B>: VecZnxDftToRef<B>;
67}
68
69pub trait ProdInplaceScratchSpace {
70 fn prod_by_grlwe_inplace_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize;
71 fn prod_by_rgsw_inplace_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize;
72}
73
74pub trait ProdInplace<MUT, REF>
75where
76 MatZnxDft<REF, FFT64>: MatZnxDftToRef<FFT64>,
77{
78 fn prod_by_grlwe_inplace(&mut self, module: &Module<FFT64>, rhs: &GRLWECt<REF, FFT64>, scratch: &mut Scratch);
79 fn prod_by_rgsw_inplace(&mut self, module: &Module<FFT64>, rhs: &RGSWCt<REF, FFT64>, scratch: &mut Scratch);
80}
81
82pub trait ProdScratchSpace {
83 fn prod_by_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize;
84 fn prod_by_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize;
85}
86
87pub trait Product<MUT, REF>
88where
89 MatZnxDft<REF, FFT64>: MatZnxDftToRef<FFT64>,
90{
91 type Lhs;
92
93 fn prod_by_grlwe(&mut self, module: &Module<FFT64>, lhs: &Self::Lhs, rhs: &GRLWECt<REF, FFT64>, scratch: &mut Scratch);
94 fn prod_by_rgsw(&mut self, module: &Module<FFT64>, lhs: &Self::Lhs, rhs: &RGSWCt<REF, FFT64>, scratch: &mut Scratch);
95}
96
97pub(crate) trait MatRLWEProductScratchSpace {
98 fn prod_with_rlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize;
99
100 fn prod_with_rlwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, mat_size: usize) -> usize {
101 Self::prod_with_rlwe_scratch_space(module, res_size, res_size, mat_size)
102 }
103
104 fn prod_with_rlwe_dft_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, mat_size: usize) -> usize {
105 (Self::prod_with_rlwe_scratch_space(module, res_size, a_size, mat_size) | module.vec_znx_idft_tmp_bytes())
106 + module.bytes_of_vec_znx(2, a_size)
107 + module.bytes_of_vec_znx(2, res_size)
108 }
109
110 fn prod_with_rlwe_dft_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, mat_size: usize) -> usize {
111 (Self::prod_with_rlwe_inplace_scratch_space(module, res_size, mat_size) | module.vec_znx_idft_tmp_bytes())
112 + module.bytes_of_vec_znx(2, res_size)
113 }
114
115 fn prod_with_mat_rlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, mat_size: usize) -> usize {
116 Self::prod_with_rlwe_dft_scratch_space(module, res_size, a_size, mat_size)
117 + module.bytes_of_vec_znx_dft(2, a_size)
118 + module.bytes_of_vec_znx_dft(2, res_size)
119 }
120
121 fn prod_with_mat_rlwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, mat_size: usize) -> usize {
122 Self::prod_with_rlwe_dft_inplace_scratch_space(module, res_size, mat_size) + module.bytes_of_vec_znx_dft(2, res_size)
123 }
124}
125
126pub(crate) trait MatRLWEProduct: Infos {
127 fn prod_with_rlwe<MUT, REF>(&self, module: &Module<FFT64>, res: &mut RLWECt<MUT>, a: &RLWECt<REF>, scratch: &mut Scratch)
128 where
129 VecZnx<MUT>: VecZnxToMut,
130 VecZnx<REF>: VecZnxToRef;
131
132 fn prod_with_rlwe_inplace<MUT>(&self, module: &Module<FFT64>, res: &mut RLWECt<MUT>, scratch: &mut Scratch)
133 where
134 VecZnx<MUT>: VecZnxToMut + VecZnxToRef,
135 {
136 unsafe {
137 let res_ptr: *mut RLWECt<MUT> = res as *mut RLWECt<MUT>; // This is ok because [Self::mul_rlwe] only updates res at the end.
138 self.prod_with_rlwe(&module, &mut *res_ptr, &*res_ptr, scratch);
139 }
140 }
141
142 fn prod_with_rlwe_dft<MUT, REF>(
143 &self,
144 module: &Module<FFT64>,
145 res: &mut RLWECtDft<MUT, FFT64>,
146 a: &RLWECtDft<REF, FFT64>,
147 scratch: &mut Scratch,
148 ) where
149 VecZnxDft<MUT, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
150 VecZnxDft<REF, FFT64>: VecZnxDftToRef<FFT64> + ZnxInfos,
151 {
152 let log_base2k: usize = self.log_base2k();
153
154 #[cfg(debug_assertions)]
155 {
156 assert_eq!(res.log_base2k(), log_base2k);
157 assert_eq!(self.n(), module.n());
158 assert_eq!(res.n(), module.n());
159 }
160
161 let (a_data, scratch_1) = scratch.tmp_vec_znx(module, 2, a.size());
162
163 let mut a_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
164 data: a_data,
165 log_base2k: a.log_base2k(),
166 log_k: a.log_k(),
167 };
168
169 a.idft(module, &mut a_idft, scratch_1);
170
171 let (res_data, scratch_2) = scratch_1.tmp_vec_znx(module, 2, res.size());
172
173 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
174 data: res_data,
175 log_base2k: res.log_base2k(),
176 log_k: res.log_k(),
177 };
178
179 self.prod_with_rlwe(module, &mut res_idft, &a_idft, scratch_2);
180
181 module.vec_znx_dft(res, 0, &res_idft, 0);
182 module.vec_znx_dft(res, 1, &res_idft, 1);
183 }
184
185 fn prod_with_rlwe_dft_inplace<MUT>(&self, module: &Module<FFT64>, res: &mut RLWECtDft<MUT, FFT64>, scratch: &mut Scratch)
186 where
187 VecZnxDft<MUT, FFT64>: VecZnxDftToRef<FFT64> + VecZnxDftToMut<FFT64>,
188 {
189 let log_base2k: usize = self.log_base2k();
190
191 #[cfg(debug_assertions)]
192 {
193 assert_eq!(res.log_base2k(), log_base2k);
194 assert_eq!(self.n(), module.n());
195 assert_eq!(res.n(), module.n());
196 }
197
198 let (res_data, scratch_1) = scratch.tmp_vec_znx(module, 2, res.size());
199
200 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
201 data: res_data,
202 log_base2k: res.log_base2k(),
203 log_k: res.log_k(),
204 };
205
206 res.idft(module, &mut res_idft, scratch_1);
207
208 self.prod_with_rlwe_inplace(module, &mut res_idft, scratch_1);
209
210 module.vec_znx_dft(res, 0, &res_idft, 0);
211 module.vec_znx_dft(res, 1, &res_idft, 1);
212 }
213
214 fn prod_with_mat_rlwe<RES, LHS>(&self, module: &Module<FFT64>, res: &mut RES, a: &LHS, scratch: &mut Scratch)
215 where
216 LHS: GetRow<FFT64> + Infos,
217 RES: SetRow<FFT64> + Infos,
218 {
219 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
220
221 let mut tmp_a_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
222 data: tmp_row_data,
223 log_base2k: a.log_base2k(),
224 log_k: a.log_k(),
225 };
226
227 let (tmp_res_data, scratch2) = scratch1.tmp_vec_znx_dft(module, 2, res.size());
228
229 let mut tmp_res_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
230 data: tmp_res_data,
231 log_base2k: res.log_base2k(),
232 log_k: res.log_k(),
233 };
234
235 let min_rows: usize = res.rows().min(a.rows());
236
237 (0..res.rows()).for_each(|row_i| {
238 (0..res.cols()).for_each(|col_j| {
239 a.get_row(module, row_i, col_j, &mut tmp_a_row);
240 self.prod_with_rlwe_dft(module, &mut tmp_res_row, &tmp_a_row, scratch2);
241 res.set_row(module, row_i, col_j, &tmp_res_row);
242 });
243 });
244
245 tmp_res_row.data.zero();
246
247 (min_rows..res.rows()).for_each(|row_i| {
248 (0..self.cols()).for_each(|col_j| {
249 res.set_row(module, row_i, col_j, &tmp_res_row);
250 });
251 });
252 }
253
254 fn prod_with_mat_rlwe_inplace<RES>(&self, module: &Module<FFT64>, res: &mut RES, scratch: &mut Scratch)
255 where
256 RES: GetRow<FFT64> + SetRow<FFT64> + Infos,
257 {
258 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
259
260 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
261 data: tmp_row_data,
262 log_base2k: res.log_base2k(),
263 log_k: res.log_k(),
264 };
265
266 (0..res.rows()).for_each(|row_i| {
267 (0..res.cols()).for_each(|col_j| {
268 res.get_row(module, row_i, col_j, &mut tmp_row);
269 self.prod_with_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
270 res.set_row(module, row_i, col_j, &tmp_row);
271 });
272 });
273 }
274}