poulpy/poulpy-core/src/external_product/glwe_ct.rs
poulpy-core/src/external_product/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, TakeVecZnxDft, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume,
4 VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
5 },
6 layouts::{Backend, DataMut, DataRef, DataViewMut, Module, Scratch, VecZnxBig},
7};
8
9use crate::layouts::{GLWECiphertext, Infos, prepared::GGSWCiphertextPrepared};
10
11impl GLWECiphertext<Vec<u8>> {
12 #[allow(clippy::too_many_arguments)]
13 pub fn external_product_scratch_space<B: Backend>(
14 module: &Module<B>,
15 basek: usize,
16 k_out: usize,
17 k_in: usize,
18 k_ggsw: usize,
19 digits: usize,
20 rank: usize,
21 ) -> usize
22 where
23 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxNormalizeTmpBytes,
24 {
25 let in_size: usize = k_in.div_ceil(basek).div_ceil(digits);
26 let out_size: usize = k_out.div_ceil(basek);
27 let ggsw_size: usize = k_ggsw.div_ceil(basek);
28 let res_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, ggsw_size);
29 let a_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, in_size);
30 let vmp: usize = module.vmp_apply_dft_to_dft_tmp_bytes(
31 out_size,
32 in_size,
33 in_size, // rows
34 rank + 1, // cols in
35 rank + 1, // cols out
36 ggsw_size,
37 );
38 let normalize: usize = module.vec_znx_normalize_tmp_bytes();
39 res_dft + a_dft + (vmp | normalize)
40 }
41
42 pub fn external_product_inplace_scratch_space<B: Backend>(
43 module: &Module<B>,
44 basek: usize,
45 k_out: usize,
46 k_ggsw: usize,
47 digits: usize,
48 rank: usize,
49 ) -> usize
50 where
51 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxNormalizeTmpBytes,
52 {
53 Self::external_product_scratch_space(module, basek, k_out, k_out, k_ggsw, digits, rank)
54 }
55}
56
57impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
58 pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
59 &mut self,
60 module: &Module<B>,
61 lhs: &GLWECiphertext<DataLhs>,
62 rhs: &GGSWCiphertextPrepared<DataRhs, B>,
63 scratch: &mut Scratch<B>,
64 ) where
65 Module<B>: VecZnxDftAllocBytes
66 + VmpApplyDftToDftTmpBytes
67 + VecZnxNormalizeTmpBytes
68 + VecZnxDftApply<B>
69 + VmpApplyDftToDft<B>
70 + VmpApplyDftToDftAdd<B>
71 + VecZnxIdftApplyConsume<B>
72 + VecZnxBigNormalize<B>,
73 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
74 {
75 let basek: usize = self.basek();
76
77 #[cfg(debug_assertions)]
78 {
79 use poulpy_hal::api::ScratchAvailable;
80
81 assert_eq!(rhs.rank(), lhs.rank());
82 assert_eq!(rhs.rank(), self.rank());
83 assert_eq!(self.basek(), basek);
84 assert_eq!(lhs.basek(), basek);
85 assert_eq!(rhs.n(), self.n());
86 assert_eq!(lhs.n(), self.n());
87 assert!(
88 scratch.available()
89 >= GLWECiphertext::external_product_scratch_space(
90 module,
91 self.basek(),
92 self.k(),
93 lhs.k(),
94 rhs.k(),
95 rhs.digits(),
96 rhs.rank(),
97 )
98 );
99 }
100
101 let cols: usize = rhs.rank() + 1;
102 let digits: usize = rhs.digits();
103
104 let (mut res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), cols, rhs.size()); // Todo optimise
105 let (mut a_dft, scratch_2) = scratch_1.take_vec_znx_dft(self.n(), cols, lhs.size().div_ceil(digits));
106
107 a_dft.data_mut().fill(0);
108
109 {
110 (0..digits).for_each(|di| {
111 // (lhs.size() + di) / digits = (a - (digit - di - 1)).div_ceil(digits)
112 a_dft.set_size((lhs.size() + di) / digits);
113
114 // Small optimization for digits > 2
115 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
116 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
117 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
118 // It is possible to further ignore the last digits-1 limbs, but this introduce
119 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
120 // noise is kept with respect to the ideal functionality.
121 res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
122
123 (0..cols).for_each(|col_i| {
124 module.vec_znx_dft_apply(digits, digits - 1 - di, &mut a_dft, col_i, &lhs.data, col_i);
125 });
126
127 if di == 0 {
128 module.vmp_apply_dft_to_dft(&mut res_dft, &a_dft, &rhs.data, scratch_2);
129 } else {
130 module.vmp_apply_dft_to_dft_add(&mut res_dft, &a_dft, &rhs.data, di, scratch_2);
131 }
132 });
133 }
134
135 let res_big: VecZnxBig<&mut [u8], B> = module.vec_znx_idft_apply_consume(res_dft);
136
137 (0..cols).for_each(|i| {
138 module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch_1);
139 });
140 }
141
142 pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
143 &mut self,
144 module: &Module<B>,
145 rhs: &GGSWCiphertextPrepared<DataRhs, B>,
146 scratch: &mut Scratch<B>,
147 ) where
148 Module<B>: VecZnxDftAllocBytes
149 + VmpApplyDftToDftTmpBytes
150 + VecZnxNormalizeTmpBytes
151 + VecZnxDftApply<B>
152 + VmpApplyDftToDft<B>
153 + VmpApplyDftToDftAdd<B>
154 + VecZnxIdftApplyConsume<B>
155 + VecZnxBigNormalize<B>,
156 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
157 {
158 let basek: usize = self.basek();
159
160 #[cfg(debug_assertions)]
161 {
162 use poulpy_hal::api::ScratchAvailable;
163
164 assert_eq!(rhs.rank(), self.rank());
165 assert_eq!(self.basek(), basek);
166 assert_eq!(rhs.n(), self.n());
167 assert!(
168 scratch.available()
169 >= GLWECiphertext::external_product_scratch_space(
170 module,
171 self.basek(),
172 self.k(),
173 self.k(),
174 rhs.k(),
175 rhs.digits(),
176 rhs.rank(),
177 )
178 );
179 }
180
181 let cols: usize = rhs.rank() + 1;
182 let digits: usize = rhs.digits();
183
184 let (mut res_dft, scratch_1) = scratch.take_vec_znx_dft(self.n(), cols, rhs.size()); // Todo optimise
185 let (mut a_dft, scratch_2) = scratch_1.take_vec_znx_dft(self.n(), cols, self.size().div_ceil(digits));
186
187 a_dft.data_mut().fill(0);
188
189 {
190 (0..digits).for_each(|di| {
191 // (lhs.size() + di) / digits = (a - (digit - di - 1)).div_ceil(digits)
192 a_dft.set_size((self.size() + di) / digits);
193
194 // Small optimization for digits > 2
195 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
196 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
197 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
198 // It is possible to further ignore the last digits-1 limbs, but this introduce
199 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
200 // noise is kept with respect to the ideal functionality.
201 res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
202
203 (0..cols).for_each(|col_i| {
204 module.vec_znx_dft_apply(
205 digits,
206 digits - 1 - di,
207 &mut a_dft,
208 col_i,
209 &self.data,
210 col_i,
211 );
212 });
213
214 if di == 0 {
215 module.vmp_apply_dft_to_dft(&mut res_dft, &a_dft, &rhs.data, scratch_2);
216 } else {
217 module.vmp_apply_dft_to_dft_add(&mut res_dft, &a_dft, &rhs.data, di, scratch_2);
218 }
219 });
220 }
221
222 let res_big: VecZnxBig<&mut [u8], B> = module.vec_znx_idft_apply_consume(res_dft);
223
224 (0..cols).for_each(|i| {
225 module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch_1);
226 });
227 }
228}