poulpy/rlwe/src/elem.rs
rlwe/src/elem.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftToMut, MatZnxDftToRef, Module, Scratch, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxDftToMut,
3 VecZnxDftToRef, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxZero,
4};
5
6use crate::{
7 elem_grlwe::GRLWECt,
8 elem_rlwe::{RLWECt, RLWECtDft},
9 utils::derive_size,
10};
11
12pub trait Infos {
13 type Inner: ZnxInfos;
14
15 fn inner(&self) -> &Self::Inner;
16
17 /// Returns the ring degree of the polynomials.
18 fn n(&self) -> usize {
19 self.inner().n()
20 }
21
22 /// Returns the base two logarithm of the ring dimension of the polynomials.
23 fn log_n(&self) -> usize {
24 self.inner().log_n()
25 }
26
27 /// Returns the number of rows.
28 fn rows(&self) -> usize {
29 self.inner().rows()
30 }
31
32 /// Returns the number of polynomials in each row.
33 fn cols(&self) -> usize {
34 self.inner().cols()
35 }
36
37 /// Returns the number of size per polynomial.
38 fn size(&self) -> usize {
39 let size: usize = self.inner().size();
40 debug_assert_eq!(size, derive_size(self.log_base2k(), self.log_k()));
41 size
42 }
43
44 /// Returns the total number of small polynomials.
45 fn poly_count(&self) -> usize {
46 self.rows() * self.cols() * self.size()
47 }
48
49 /// Returns the base 2 logarithm of the ciphertext base.
50 fn log_base2k(&self) -> usize;
51
52 /// Returns the bit precision of the ciphertext.
53 fn log_k(&self) -> usize;
54}
55
56pub trait GetRow<B: Backend> {
57 fn get_row<R>(&self, module: &Module<B>, row_i: usize, col_j: usize, res: &mut RLWECtDft<R, B>)
58 where
59 VecZnxDft<R, B>: VecZnxDftToMut<B>;
60}
61
62pub trait SetRow<B: Backend> {
63 fn set_row<A>(&mut self, module: &Module<B>, row_i: usize, col_j: usize, a: &RLWECtDft<A, B>)
64 where
65 VecZnxDft<A, B>: VecZnxDftToRef<B>;
66}
67
68pub(crate) trait MatZnxDftProducts<D, C>: Infos
69where
70 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
71{
72 fn mul_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, a: &RLWECt<A>, scratch: &mut Scratch)
73 where
74 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
75 VecZnx<R>: VecZnxToMut,
76 VecZnx<A>: VecZnxToRef;
77
78 fn mul_rlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, scratch: &mut Scratch)
79 where
80 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
81 VecZnx<R>: VecZnxToMut + VecZnxToRef,
82 {
83 unsafe {
84 let res_ptr: *mut RLWECt<R> = res as *mut RLWECt<R>; // This is ok because [Self::mul_rlwe] only updates res at the end.
85 self.mul_rlwe(&module, &mut *res_ptr, &*res_ptr, scratch);
86 }
87 }
88
89 fn mul_rlwe_dft<R, A>(
90 &self,
91 module: &Module<FFT64>,
92 res: &mut RLWECtDft<R, FFT64>,
93 a: &RLWECtDft<A, FFT64>,
94 scratch: &mut Scratch,
95 ) where
96 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
97 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
98 VecZnxDft<A, FFT64>: VecZnxDftToRef<FFT64> + ZnxInfos,
99 {
100 let log_base2k: usize = self.log_base2k();
101
102 #[cfg(debug_assertions)]
103 {
104 assert_eq!(res.log_base2k(), log_base2k);
105 assert_eq!(self.n(), module.n());
106 assert_eq!(res.n(), module.n());
107 }
108
109 let (a_data, scratch_1) = scratch.tmp_vec_znx(module, 2, a.size());
110
111 let mut a_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
112 data: a_data,
113 log_base2k: a.log_base2k(),
114 log_k: a.log_k(),
115 };
116
117 a.idft(module, &mut a_idft, scratch_1);
118
119 let (res_data, scratch_2) = scratch_1.tmp_vec_znx(module, 2, res.size());
120
121 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
122 data: res_data,
123 log_base2k: res.log_base2k(),
124 log_k: res.log_k(),
125 };
126
127 self.mul_rlwe(module, &mut res_idft, &a_idft, scratch_2);
128
129 module.vec_znx_dft(res, 0, &res_idft, 0);
130 module.vec_znx_dft(res, 1, &res_idft, 1);
131 }
132
133 fn mul_rlwe_dft_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECtDft<R, FFT64>, scratch: &mut Scratch)
134 where
135 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
136 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64> + VecZnxDftToMut<FFT64>,
137 {
138 let log_base2k: usize = self.log_base2k();
139
140 #[cfg(debug_assertions)]
141 {
142 assert_eq!(res.log_base2k(), log_base2k);
143 assert_eq!(self.n(), module.n());
144 assert_eq!(res.n(), module.n());
145 }
146
147 let (res_data, scratch_1) = scratch.tmp_vec_znx(module, 2, res.size());
148
149 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
150 data: res_data,
151 log_base2k: res.log_base2k(),
152 log_k: res.log_k(),
153 };
154
155 res.idft(module, &mut res_idft, scratch_1);
156
157 self.mul_rlwe_inplace(module, &mut res_idft, scratch_1);
158
159 module.vec_znx_dft(res, 0, &res_idft, 0);
160 module.vec_znx_dft(res, 1, &res_idft, 1);
161 }
162
163 fn mul_grlwe<R, A>(&self, module: &Module<FFT64>, res: &mut GRLWECt<R, FFT64>, a: &GRLWECt<A, FFT64>, scratch: &mut Scratch)
164 where
165 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
166 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
167 MatZnxDft<A, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
168 {
169 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
170
171 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
172 data: tmp_row_data,
173 log_base2k: a.log_base2k(),
174 log_k: a.log_k(),
175 };
176
177 let min_rows: usize = res.rows().min(a.rows());
178
179 (0..min_rows).for_each(|row_i| {
180 a.get_row(module, row_i, &mut tmp_row);
181 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
182 res.set_row(module, row_i, &tmp_row);
183 });
184
185 tmp_row.data.zero();
186
187 (min_rows..res.rows()).for_each(|row_i| {
188 res.set_row(module, row_i, &tmp_row);
189 })
190 }
191
192 fn mul_grlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut R, scratch: &mut Scratch)
193 where
194 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
195 R: GetRow<FFT64> + SetRow<FFT64> + Infos,
196 {
197 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
198
199 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
200 data: tmp_row_data,
201 log_base2k: res.log_base2k(),
202 log_k: res.log_k(),
203 };
204
205 (0..self.cols()).for_each(|col_j| {
206 (0..res.rows()).for_each(|row_i| {
207 res.get_row(module, row_i, col_j, &mut tmp_row);
208 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
209 res.set_row(module, row_i, col_j, &tmp_row);
210 });
211 })
212 }
213}