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 ProdByScratchSpace {
70 fn prod_by_grlwe_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize;
71 fn prod_by_rgsw_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize;
72}
73
74pub trait ProdBy<D> {
75 fn prod_by_grlwe<R>(&mut self, module: &Module<FFT64>, rhs: &GRLWECt<R, FFT64>, scratch: &mut Scratch)
76 where
77 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>;
78
79 fn prod_by_rgsw<R>(&mut self, module: &Module<FFT64>, rhs: &RGSWCt<R, FFT64>, scratch: &mut Scratch)
80 where
81 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>;
82}
83
84pub trait FromProdByScratchSpace {
85 fn from_prod_by_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize;
86 fn from_prod_by_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize;
87}
88
89pub trait FromProdBy<D, L> {
90 fn from_prod_by_grlwe<R>(&mut self, module: &Module<FFT64>, lhs: &L, rhs: &GRLWECt<R, FFT64>, scratch: &mut Scratch)
91 where
92 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>;
93
94 fn from_prod_by_rgsw<R>(&mut self, module: &Module<FFT64>, lhs: &L, rhs: &RGSWCt<R, FFT64>, scratch: &mut Scratch)
95 where
96 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>;
97}
98
99pub(crate) trait MatZnxDftProducts<D, C>: Infos
100where
101 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
102{
103 fn mul_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, a: &RLWECt<A>, scratch: &mut Scratch)
104 where
105 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
106 VecZnx<R>: VecZnxToMut,
107 VecZnx<A>: VecZnxToRef;
108
109 fn mul_rlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize;
110
111 fn mul_rlwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, mat_size: usize) -> usize {
112 Self::mul_rlwe_scratch_space(module, res_size, res_size, mat_size)
113 }
114
115 fn mul_rlwe_dft_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, mat_size: usize) -> usize {
116 (Self::mul_rlwe_scratch_space(module, res_size, a_size, mat_size) | module.vec_znx_idft_tmp_bytes())
117 + module.bytes_of_vec_znx(2, a_size)
118 + module.bytes_of_vec_znx(2, res_size)
119 }
120
121 fn mul_rlwe_dft_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, mat_size: usize) -> usize {
122 (Self::mul_rlwe_inplace_scratch_space(module, res_size, mat_size) | module.vec_znx_idft_tmp_bytes())
123 + module.bytes_of_vec_znx(2, res_size)
124 }
125
126 fn mul_mat_rlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, mat_size: usize) -> usize {
127 Self::mul_rlwe_dft_inplace_scratch_space(module, res_size, mat_size) + module.bytes_of_vec_znx_dft(2, a_size)
128 }
129
130 fn mul_mat_rlwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, mat_size: usize) -> usize {
131 Self::mul_rlwe_dft_inplace_scratch_space(module, res_size, mat_size) + module.bytes_of_vec_znx_dft(2, res_size)
132 }
133
134 fn mul_rlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, scratch: &mut Scratch)
135 where
136 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
137 VecZnx<R>: VecZnxToMut + VecZnxToRef,
138 {
139 unsafe {
140 let res_ptr: *mut RLWECt<R> = res as *mut RLWECt<R>; // This is ok because [Self::mul_rlwe] only updates res at the end.
141 self.mul_rlwe(&module, &mut *res_ptr, &*res_ptr, scratch);
142 }
143 }
144
145 fn mul_rlwe_dft<R, A>(
146 &self,
147 module: &Module<FFT64>,
148 res: &mut RLWECtDft<R, FFT64>,
149 a: &RLWECtDft<A, FFT64>,
150 scratch: &mut Scratch,
151 ) where
152 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
153 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
154 VecZnxDft<A, FFT64>: VecZnxDftToRef<FFT64> + ZnxInfos,
155 {
156 let log_base2k: usize = self.log_base2k();
157
158 #[cfg(debug_assertions)]
159 {
160 assert_eq!(res.log_base2k(), log_base2k);
161 assert_eq!(self.n(), module.n());
162 assert_eq!(res.n(), module.n());
163 }
164
165 let (a_data, scratch_1) = scratch.tmp_vec_znx(module, 2, a.size());
166
167 let mut a_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
168 data: a_data,
169 log_base2k: a.log_base2k(),
170 log_k: a.log_k(),
171 };
172
173 a.idft(module, &mut a_idft, scratch_1);
174
175 let (res_data, scratch_2) = scratch_1.tmp_vec_znx(module, 2, res.size());
176
177 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
178 data: res_data,
179 log_base2k: res.log_base2k(),
180 log_k: res.log_k(),
181 };
182
183 self.mul_rlwe(module, &mut res_idft, &a_idft, scratch_2);
184
185 module.vec_znx_dft(res, 0, &res_idft, 0);
186 module.vec_znx_dft(res, 1, &res_idft, 1);
187 }
188
189 fn mul_rlwe_dft_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECtDft<R, FFT64>, scratch: &mut Scratch)
190 where
191 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64> + VecZnxDftToMut<FFT64>,
192 {
193 let log_base2k: usize = self.log_base2k();
194
195 #[cfg(debug_assertions)]
196 {
197 assert_eq!(res.log_base2k(), log_base2k);
198 assert_eq!(self.n(), module.n());
199 assert_eq!(res.n(), module.n());
200 }
201
202 let (res_data, scratch_1) = scratch.tmp_vec_znx(module, 2, res.size());
203
204 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
205 data: res_data,
206 log_base2k: res.log_base2k(),
207 log_k: res.log_k(),
208 };
209
210 res.idft(module, &mut res_idft, scratch_1);
211
212 self.mul_rlwe_inplace(module, &mut res_idft, scratch_1);
213
214 module.vec_znx_dft(res, 0, &res_idft, 0);
215 module.vec_znx_dft(res, 1, &res_idft, 1);
216 }
217
218 fn mul_mat_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut R, a: &A, scratch: &mut Scratch)
219 where
220 A: GetRow<FFT64> + Infos,
221 R: SetRow<FFT64> + Infos,
222 {
223 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
224
225 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
226 data: tmp_row_data,
227 log_base2k: a.log_base2k(),
228 log_k: a.log_k(),
229 };
230
231 let min_rows: usize = res.rows().min(a.rows());
232
233 (0..res.rows()).for_each(|row_i| {
234 (0..self.cols()).for_each(|col_j| {
235 a.get_row(module, row_i, col_j, &mut tmp_row);
236 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
237 res.set_row(module, row_i, col_j, &tmp_row);
238 });
239 });
240
241 tmp_row.data.zero();
242
243 (min_rows..res.rows()).for_each(|row_i| {
244 (0..self.cols()).for_each(|col_j| {
245 res.set_row(module, row_i, col_j, &tmp_row);
246 });
247 });
248 }
249
250 fn mul_mat_rlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut R, scratch: &mut Scratch)
251 where
252 R: GetRow<FFT64> + SetRow<FFT64> + Infos,
253 {
254 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
255
256 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
257 data: tmp_row_data,
258 log_base2k: res.log_base2k(),
259 log_k: res.log_k(),
260 };
261
262 (0..res.rows()).for_each(|row_i| {
263 (0..self.cols()).for_each(|col_j| {
264 res.get_row(module, row_i, col_j, &mut tmp_row);
265 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
266 res.set_row(module, row_i, col_j, &tmp_row);
267 });
268 });
269 }
270}