poulpy/poulpy-core/src/tests/test_suite/keyswitch/ggsw_ct.rs
poulpy-core/src/tests/test_suite/keyswitch/ggsw_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes,
4 SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
5 VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAlloc,
6 VecZnxDftAllocBytes, VecZnxDftApply, VecZnxDftCopy, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA,
7 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwitchRing,
8 VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
9 },
10 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
11 oep::{
12 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
13 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
14 },
15 source::Source,
16};
17
18use crate::{
19 encryption::SIGMA,
20 layouts::{
21 GGLWESwitchingKey, GGLWETensorKey, GGSWCiphertext, GLWESecret,
22 prepared::{GGLWESwitchingKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
23 },
24 noise::noise_ggsw_keyswitch,
25};
26
27#[allow(clippy::too_many_arguments)]
28pub fn test_ggsw_keyswitch<B>(module: &Module<B>)
29where
30 Module<B>: VecZnxDftAllocBytes
31 + VecZnxBigNormalize<B>
32 + VecZnxDftApply<B>
33 + SvpApplyDftToDftInplace<B>
34 + VecZnxIdftApplyConsume<B>
35 + VecZnxFillUniform
36 + VecZnxSubABInplace
37 + VecZnxAddInplace
38 + VecZnxNormalizeInplace<B>
39 + VecZnxAddNormal
40 + VecZnxNormalize<B>
41 + VecZnxSub
42 + SvpPrepare<B>
43 + SvpPPolAllocBytes
44 + SvpPPolAlloc<B>
45 + VecZnxBigAllocBytes
46 + VecZnxBigAddInplace<B>
47 + VecZnxBigAddSmallInplace<B>
48 + VecZnxNormalizeTmpBytes
49 + VecZnxAddScalarInplace
50 + VmpPMatAlloc<B>
51 + VmpPrepare<B>
52 + VmpApplyDftToDftTmpBytes
53 + VmpApplyDftToDft<B>
54 + VmpApplyDftToDftAdd<B>
55 + VecZnxBigNormalizeTmpBytes
56 + VecZnxSwitchRing
57 + SvpApplyDftToDft<B>
58 + VecZnxIdftApplyTmpA<B>
59 + VecZnxDftCopy<B>
60 + VecZnxDftAddInplace<B>
61 + VecZnxBigAlloc<B>
62 + VecZnxDftAlloc<B>,
63 B: Backend
64 + TakeVecZnxDftImpl<B>
65 + TakeVecZnxBigImpl<B>
66 + TakeSvpPPolImpl<B>
67 + ScratchOwnedAllocImpl<B>
68 + ScratchOwnedBorrowImpl<B>
69 + ScratchAvailableImpl<B>
70 + TakeScalarZnxImpl<B>
71 + TakeVecZnxImpl<B>,
72{
73 let basek: usize = 12;
74 let k_in: usize = 54;
75 let digits: usize = k_in.div_ceil(basek);
76 (1..4).for_each(|rank| {
77 (1..digits + 1).for_each(|di| {
78 let k_ksk: usize = k_in + basek * di;
79 let k_tsk: usize = k_ksk;
80 let k_out: usize = k_ksk; // Better capture noise.
81
82 let n: usize = module.n();
83 let rows: usize = k_in.div_ceil(di * basek);
84
85 let digits_in: usize = 1;
86
87 let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows, digits_in, rank);
88 let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows, digits_in, rank);
89 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_ksk, rows, di, rank);
90 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, di, rank, rank);
91 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
92
93 let mut source_xs: Source = Source::new([0u8; 32]);
94 let mut source_xe: Source = Source::new([0u8; 32]);
95 let mut source_xa: Source = Source::new([0u8; 32]);
96
97 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
98 GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k_in, rank)
99 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank, rank)
100 | GGLWETensorKey::encrypt_sk_scratch_space(module, basek, k_tsk, rank)
101 | GGSWCiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, di, k_tsk, di, rank),
102 );
103
104 let var_xs: f64 = 0.5;
105
106 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
107 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
108 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
109
110 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
111 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
112 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
113
114 ksk.encrypt_sk(
115 module,
116 &sk_in,
117 &sk_out,
118 &mut source_xa,
119 &mut source_xe,
120 scratch.borrow(),
121 );
122 tsk.encrypt_sk(
123 module,
124 &sk_out,
125 &mut source_xa,
126 &mut source_xe,
127 scratch.borrow(),
128 );
129
130 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
131
132 ct_in.encrypt_sk(
133 module,
134 &pt_scalar,
135 &sk_in_dft,
136 &mut source_xa,
137 &mut source_xe,
138 scratch.borrow(),
139 );
140
141 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
142 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
143
144 ct_out.keyswitch(
145 module,
146 &ct_in,
147 &ksk_prepared,
148 &tsk_prepared,
149 scratch.borrow(),
150 );
151
152 let max_noise = |col_j: usize| -> f64 {
153 noise_ggsw_keyswitch(
154 n as f64,
155 basek * di,
156 col_j,
157 var_xs,
158 0f64,
159 SIGMA * SIGMA,
160 0f64,
161 rank as f64,
162 k_in,
163 k_ksk,
164 k_tsk,
165 ) + 0.5
166 };
167
168 ct_out.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
169 });
170 });
171}
172
173#[allow(clippy::too_many_arguments)]
174pub fn test_ggsw_keyswitch_inplace<B>(module: &Module<B>)
175where
176 Module<B>: VecZnxDftAllocBytes
177 + VecZnxBigNormalize<B>
178 + VecZnxDftApply<B>
179 + SvpApplyDftToDftInplace<B>
180 + VecZnxIdftApplyConsume<B>
181 + VecZnxFillUniform
182 + VecZnxSubABInplace
183 + VecZnxAddInplace
184 + VecZnxNormalizeInplace<B>
185 + VecZnxAddNormal
186 + VecZnxNormalize<B>
187 + VecZnxSub
188 + SvpPrepare<B>
189 + SvpPPolAllocBytes
190 + SvpPPolAlloc<B>
191 + VecZnxBigAllocBytes
192 + VecZnxBigAddInplace<B>
193 + VecZnxBigAddSmallInplace<B>
194 + VecZnxNormalizeTmpBytes
195 + VecZnxAddScalarInplace
196 + VmpPMatAlloc<B>
197 + VmpPrepare<B>
198 + VmpApplyDftToDftTmpBytes
199 + VmpApplyDftToDft<B>
200 + VmpApplyDftToDftAdd<B>
201 + VecZnxBigNormalizeTmpBytes
202 + VecZnxSwitchRing
203 + SvpApplyDftToDft<B>
204 + VecZnxIdftApplyTmpA<B>
205 + VecZnxDftCopy<B>
206 + VecZnxDftAddInplace<B>
207 + VecZnxBigAlloc<B>
208 + VecZnxDftAlloc<B>,
209 B: Backend
210 + TakeVecZnxDftImpl<B>
211 + TakeVecZnxBigImpl<B>
212 + TakeSvpPPolImpl<B>
213 + ScratchOwnedAllocImpl<B>
214 + ScratchOwnedBorrowImpl<B>
215 + ScratchAvailableImpl<B>
216 + TakeScalarZnxImpl<B>
217 + TakeVecZnxImpl<B>,
218{
219 let basek: usize = 12;
220 let k_ct: usize = 54;
221 let digits: usize = k_ct.div_ceil(basek);
222 (1..4).for_each(|rank| {
223 (1..digits + 1).for_each(|di| {
224 let k_ksk: usize = k_ct + basek * di;
225 let k_tsk: usize = k_ksk;
226
227 let n: usize = module.n();
228 let rows: usize = k_ct.div_ceil(di * basek);
229
230 let digits_in: usize = 1;
231
232 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows, digits_in, rank);
233 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, di, rank);
234 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, di, rank, rank);
235 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
236
237 let mut source_xs: Source = Source::new([0u8; 32]);
238 let mut source_xe: Source = Source::new([0u8; 32]);
239 let mut source_xa: Source = Source::new([0u8; 32]);
240
241 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
242 GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k_ct, rank)
243 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank, rank)
244 | GGLWETensorKey::encrypt_sk_scratch_space(module, basek, k_tsk, rank)
245 | GGSWCiphertext::keyswitch_inplace_scratch_space(module, basek, k_ct, k_ksk, di, k_tsk, di, rank),
246 );
247
248 let var_xs: f64 = 0.5;
249
250 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
251 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
252 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
253
254 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
255 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
256 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
257
258 ksk.encrypt_sk(
259 module,
260 &sk_in,
261 &sk_out,
262 &mut source_xa,
263 &mut source_xe,
264 scratch.borrow(),
265 );
266 tsk.encrypt_sk(
267 module,
268 &sk_out,
269 &mut source_xa,
270 &mut source_xe,
271 scratch.borrow(),
272 );
273
274 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
275
276 ct.encrypt_sk(
277 module,
278 &pt_scalar,
279 &sk_in_dft,
280 &mut source_xa,
281 &mut source_xe,
282 scratch.borrow(),
283 );
284
285 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
286 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
287
288 ct.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
289
290 let max_noise = |col_j: usize| -> f64 {
291 noise_ggsw_keyswitch(
292 n as f64,
293 basek * di,
294 col_j,
295 var_xs,
296 0f64,
297 SIGMA * SIGMA,
298 0f64,
299 rank as f64,
300 k_ct,
301 k_ksk,
302 k_tsk,
303 ) + 0.5
304 };
305
306 ct.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
307 });
308 });
309}