poulpy/core/src/tests/generics/keyswitch/ggsw_ct.rs
core/src/tests/generics/keyswitch/ggsw_ct.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxCopy, VecZnxStd, VecZnxSubABInplace,
4 VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
5 },
6 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
7 oep::{
8 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
9 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
10 },
11};
12use sampling::source::Source;
13
14use crate::{
15 layouts::{
16 GGLWETensorKey, GGSWCiphertext, GLWESecret, GGLWESwitchingKey,
17 prepared::{GLWESecretExec, GGLWESwitchingKeyExec, GGLWETensorKeyExec},
18 },
19 noise::noise_ggsw_keyswitch,
20 trait_families::GGSWAssertNoiseFamily,
21};
22
23use crate::trait_families::{
24 GGLWESwitchingKeyEncryptSkFamily, GGLWETensorKeyEncryptSkFamily, GGSWEncryptSkFamily, GGSWKeySwitchFamily,
25 GLWESecretExecModuleFamily,
26};
27
28pub fn test_ggsw_keyswitch<B: Backend>(
29 module: &Module<B>,
30 basek: usize,
31 k_out: usize,
32 k_in: usize,
33 k_ksk: usize,
34 k_tsk: usize,
35 digits: usize,
36 rank: usize,
37 sigma: f64,
38) where
39 Module<B>: GLWESecretExecModuleFamily<B>
40 + GGSWEncryptSkFamily<B>
41 + GGSWAssertNoiseFamily<B>
42 + VecZnxAddScalarInplace
43 + VecZnxSubABInplace
44 + VecZnxStd
45 + VecZnxCopy
46 + VmpPMatAlloc<B>
47 + VmpPMatPrepare<B>
48 + GGSWAssertNoiseFamily<B>
49 + GGSWKeySwitchFamily<B>
50 + GGLWESwitchingKeyEncryptSkFamily<B>
51 + GGLWETensorKeyEncryptSkFamily<B>
52 + VecZnxSwithcDegree,
53 B: TakeVecZnxDftImpl<B>
54 + TakeVecZnxBigImpl<B>
55 + TakeSvpPPolImpl<B>
56 + ScratchOwnedAllocImpl<B>
57 + ScratchOwnedBorrowImpl<B>
58 + ScratchAvailableImpl<B>
59 + TakeScalarZnxImpl<B>
60 + TakeVecZnxImpl<B>
61 + VecZnxDftAllocBytesImpl<B>
62 + VecZnxBigAllocBytesImpl<B>
63 + TakeSvpPPolImpl<B>
64 + VecZnxDftAllocBytesImpl<B>
65 + VecZnxBigAllocBytesImpl<B>
66 + TakeSvpPPolImpl<B>,
67{
68 let n: usize = module.n();
69 let rows: usize = k_in.div_ceil(digits * basek);
70
71 let digits_in: usize = 1;
72
73 let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows, digits_in, rank);
74 let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows, digits_in, rank);
75 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_ksk, rows, digits, rank);
76 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
77 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
78
79 let mut source_xs: Source = Source::new([0u8; 32]);
80 let mut source_xe: Source = Source::new([0u8; 32]);
81 let mut source_xa: Source = Source::new([0u8; 32]);
82
83 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
84 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_in, rank)
85 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank, rank)
86 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
87 | GGSWCiphertext::keyswitch_scratch_space(
88 module, n, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
89 ),
90 );
91
92 let var_xs: f64 = 0.5;
93
94 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
95 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
96 let sk_in_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_in);
97
98 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
99 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
100 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
101
102 ksk.encrypt_sk(
103 module,
104 &sk_in,
105 &sk_out,
106 &mut source_xa,
107 &mut source_xe,
108 sigma,
109 scratch.borrow(),
110 );
111 tsk.encrypt_sk(
112 module,
113 &sk_out,
114 &mut source_xa,
115 &mut source_xe,
116 sigma,
117 scratch.borrow(),
118 );
119
120 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
121
122 ct_in.encrypt_sk(
123 module,
124 &pt_scalar,
125 &sk_in_dft,
126 &mut source_xa,
127 &mut source_xe,
128 sigma,
129 scratch.borrow(),
130 );
131
132 let mut ksk_exec: GGLWESwitchingKeyExec<Vec<u8>, B> =
133 GGLWESwitchingKeyExec::alloc(module, n, basek, k_ksk, rows, digits, rank, rank);
134 let mut tsk_exec: GGLWETensorKeyExec<Vec<u8>, B> = GGLWETensorKeyExec::alloc(module, n, basek, k_ksk, rows, digits, rank);
135
136 ksk_exec.prepare(module, &ksk, scratch.borrow());
137 tsk_exec.prepare(module, &tsk, scratch.borrow());
138
139 ct_out.keyswitch(module, &ct_in, &ksk_exec, &tsk_exec, scratch.borrow());
140
141 let max_noise = |col_j: usize| -> f64 {
142 noise_ggsw_keyswitch(
143 n as f64,
144 basek * digits,
145 col_j,
146 var_xs,
147 0f64,
148 sigma * sigma,
149 0f64,
150 rank as f64,
151 k_in,
152 k_ksk,
153 k_tsk,
154 ) + 0.5
155 };
156
157 ct_out.assert_noise(module, &sk_out_exec, &pt_scalar, &max_noise);
158}
159
160pub fn test_ggsw_keyswitch_inplace<B: Backend>(
161 module: &Module<B>,
162 basek: usize,
163 k_ct: usize,
164 k_ksk: usize,
165 k_tsk: usize,
166 digits: usize,
167 rank: usize,
168 sigma: f64,
169) where
170 Module<B>: GLWESecretExecModuleFamily<B>
171 + GGSWEncryptSkFamily<B>
172 + GGSWAssertNoiseFamily<B>
173 + VecZnxAddScalarInplace
174 + VecZnxSubABInplace
175 + VecZnxStd
176 + VecZnxCopy
177 + VmpPMatAlloc<B>
178 + VmpPMatPrepare<B>
179 + GGSWAssertNoiseFamily<B>
180 + GGSWKeySwitchFamily<B>
181 + GGLWESwitchingKeyEncryptSkFamily<B>
182 + GGLWETensorKeyEncryptSkFamily<B>
183 + VecZnxSwithcDegree,
184 B: TakeVecZnxDftImpl<B>
185 + TakeVecZnxBigImpl<B>
186 + TakeSvpPPolImpl<B>
187 + ScratchOwnedAllocImpl<B>
188 + ScratchOwnedBorrowImpl<B>
189 + ScratchAvailableImpl<B>
190 + TakeScalarZnxImpl<B>
191 + TakeVecZnxImpl<B>
192 + VecZnxDftAllocBytesImpl<B>
193 + VecZnxBigAllocBytesImpl<B>
194 + TakeSvpPPolImpl<B>,
195{
196 let n: usize = module.n();
197 let rows: usize = k_ct.div_ceil(digits * basek);
198
199 let digits_in: usize = 1;
200
201 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows, digits_in, rank);
202 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
203 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
204 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
205
206 let mut source_xs: Source = Source::new([0u8; 32]);
207 let mut source_xe: Source = Source::new([0u8; 32]);
208 let mut source_xa: Source = Source::new([0u8; 32]);
209
210 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
211 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_ct, rank)
212 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank, rank)
213 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
214 | GGSWCiphertext::keyswitch_inplace_scratch_space(module, n, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
215 );
216
217 let var_xs: f64 = 0.5;
218
219 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
220 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
221 let sk_in_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_in);
222
223 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
224 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
225 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
226
227 ksk.encrypt_sk(
228 module,
229 &sk_in,
230 &sk_out,
231 &mut source_xa,
232 &mut source_xe,
233 sigma,
234 scratch.borrow(),
235 );
236 tsk.encrypt_sk(
237 module,
238 &sk_out,
239 &mut source_xa,
240 &mut source_xe,
241 sigma,
242 scratch.borrow(),
243 );
244
245 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
246
247 ct.encrypt_sk(
248 module,
249 &pt_scalar,
250 &sk_in_dft,
251 &mut source_xa,
252 &mut source_xe,
253 sigma,
254 scratch.borrow(),
255 );
256
257 let mut ksk_exec: GGLWESwitchingKeyExec<Vec<u8>, B> =
258 GGLWESwitchingKeyExec::alloc(module, n, basek, k_ksk, rows, digits, rank, rank);
259 let mut tsk_exec: GGLWETensorKeyExec<Vec<u8>, B> = GGLWETensorKeyExec::alloc(module, n, basek, k_ksk, rows, digits, rank);
260
261 ksk_exec.prepare(module, &ksk, scratch.borrow());
262 tsk_exec.prepare(module, &tsk, scratch.borrow());
263
264 ct.keyswitch_inplace(module, &ksk_exec, &tsk_exec, scratch.borrow());
265
266 let max_noise = |col_j: usize| -> f64 {
267 noise_ggsw_keyswitch(
268 n as f64,
269 basek * digits,
270 col_j,
271 var_xs,
272 0f64,
273 sigma * sigma,
274 0f64,
275 rank as f64,
276 k_ct,
277 k_ksk,
278 k_tsk,
279 ) + 0.5
280 };
281
282 ct.assert_noise(module, &sk_out_exec, &pt_scalar, &max_noise);
283}