poulpy/core/src/glwe/test_fft64/keyswitch.rs
core/src/glwe/test_fft64/keyswitch.rsRaw
1use backend::{
2 hal::{
3 api::{
4 MatZnxAlloc, ModuleNew, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow,
5 VecZnxAddScalarInplace, VecZnxAlloc, VecZnxAllocBytes, VecZnxFillUniform, VecZnxStd, VecZnxSwithcDegree,
6 },
7 layouts::{Backend, Module, ScratchOwned},
8 oep::{
9 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
10 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
11 },
12 },
13 implementation::cpu_spqlios::FFT64,
14};
15use sampling::source::Source;
16
17use crate::{
18 GGLWEExecLayoutFamily, GLWECiphertext, GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEPlaintext, GLWESecret, GLWESecretExec,
19 GLWESecretFamily, GLWESwitchingKey, GLWESwitchingKeyEncryptSkFamily, GLWESwitchingKeyExec, Infos,
20 noise::log2_std_noise_gglwe_product,
21};
22
23#[test]
24fn apply() {
25 let log_n: usize = 8;
26 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
27 let basek: usize = 12;
28 let k_in: usize = 45;
29 let digits: usize = k_in.div_ceil(basek);
30 (1..4).for_each(|rank_in| {
31 (1..4).for_each(|rank_out| {
32 (1..digits + 1).for_each(|di| {
33 let k_ksk: usize = k_in + basek * di;
34 let k_out: usize = k_ksk; // better capture noise
35 println!(
36 "test keyswitch digits: {} rank_in: {} rank_out: {}",
37 di, rank_in, rank_out
38 );
39 test_keyswitch(
40 &module, basek, k_out, k_in, k_ksk, di, rank_in, rank_out, 3.2,
41 );
42 })
43 });
44 });
45}
46
47#[test]
48fn apply_inplace() {
49 let log_n: usize = 8;
50 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
51 let basek: usize = 12;
52 let k_ct: usize = 45;
53 let digits: usize = k_ct.div_ceil(basek);
54 (1..4).for_each(|rank| {
55 (1..digits + 1).for_each(|di| {
56 let k_ksk: usize = k_ct + basek * di;
57 println!("test keyswitch_inplace digits: {} rank: {}", di, rank);
58 test_keyswitch_inplace(&module, basek, k_ct, k_ksk, di, rank, 3.2);
59 });
60 });
61}
62
63pub(crate) trait KeySwitchTestModuleFamily<B: Backend> = GLWESecretFamily<B>
64 + GLWESwitchingKeyEncryptSkFamily<B>
65 + GLWEKeyswitchFamily<B>
66 + GLWEDecryptFamily<B>
67 + GGLWEExecLayoutFamily<B>
68 + MatZnxAlloc
69 + VecZnxAlloc
70 + ScalarZnxAlloc
71 + ScalarZnxAllocBytes
72 + VecZnxAllocBytes
73 + VecZnxStd
74 + VecZnxSwithcDegree
75 + VecZnxAddScalarInplace;
76
77pub(crate) trait KeySwitchTestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
78 + TakeVecZnxBigImpl<B>
79 + TakeSvpPPolImpl<B>
80 + ScratchOwnedAllocImpl<B>
81 + ScratchOwnedBorrowImpl<B>
82 + ScratchAvailableImpl<B>
83 + TakeScalarZnxImpl<B>
84 + TakeVecZnxImpl<B>;
85
86fn test_keyswitch<B: Backend>(
87 module: &Module<B>,
88 basek: usize,
89 k_out: usize,
90 k_in: usize,
91 k_ksk: usize,
92 digits: usize,
93 rank_in: usize,
94 rank_out: usize,
95 sigma: f64,
96) where
97 Module<B>: KeySwitchTestModuleFamily<B>,
98 B: KeySwitchTestScratchFamily<B>,
99{
100 let rows: usize = k_in.div_ceil(basek * digits);
101
102 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(module, basek, k_ksk, rows, digits, rank_in, rank_out);
103 let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_in, rank_in);
104 let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_out, rank_out);
105 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_in);
106
107 let mut source_xs: Source = Source::new([0u8; 32]);
108 let mut source_xe: Source = Source::new([0u8; 32]);
109 let mut source_xa: Source = Source::new([0u8; 32]);
110
111 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_in, &mut source_xa);
112
113 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
114 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, ksk.k(), rank_in, rank_out)
115 | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_in.k())
116 | GLWECiphertext::keyswitch_scratch_space(
117 module,
118 basek,
119 ct_out.k(),
120 ct_in.k(),
121 ksk.k(),
122 digits,
123 rank_in,
124 rank_out,
125 ),
126 );
127
128 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_in);
129 sk_in.fill_ternary_prob(0.5, &mut source_xs);
130 let sk_in_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_in);
131
132 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out);
133 sk_out.fill_ternary_prob(0.5, &mut source_xs);
134 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
135
136 ksk.encrypt_sk(
137 module,
138 &sk_in,
139 &sk_out,
140 &mut source_xa,
141 &mut source_xe,
142 sigma,
143 scratch.borrow(),
144 );
145
146 ct_in.encrypt_sk(
147 module,
148 &pt_want,
149 &sk_in_exec,
150 &mut source_xa,
151 &mut source_xe,
152 sigma,
153 scratch.borrow(),
154 );
155
156 let ksk_exec: GLWESwitchingKeyExec<Vec<u8>, B> = GLWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
157
158 ct_out.keyswitch(module, &ct_in, &ksk_exec, scratch.borrow());
159
160 let max_noise: f64 = log2_std_noise_gglwe_product(
161 module.n() as f64,
162 basek * digits,
163 0.5,
164 0.5,
165 0f64,
166 sigma * sigma,
167 0f64,
168 rank_in as f64,
169 k_in,
170 k_ksk,
171 );
172
173 ct_out.assert_noise(module, &sk_out_exec, &pt_want, max_noise + 0.5);
174}
175
176fn test_keyswitch_inplace<B: Backend>(
177 module: &Module<B>,
178 basek: usize,
179 k_ct: usize,
180 k_ksk: usize,
181 digits: usize,
182 rank: usize,
183 sigma: f64,
184) where
185 Module<B>: KeySwitchTestModuleFamily<B>,
186 B: KeySwitchTestScratchFamily<B>,
187{
188 let rows: usize = k_ct.div_ceil(basek * digits);
189
190 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(module, basek, k_ksk, rows, digits, rank, rank);
191 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
192 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
193
194 let mut source_xs: Source = Source::new([0u8; 32]);
195 let mut source_xe: Source = Source::new([0u8; 32]);
196 let mut source_xa: Source = Source::new([0u8; 32]);
197
198 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
199
200 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
201 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, ksk.k(), rank, rank)
202 | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_glwe.k())
203 | GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, ct_glwe.k(), ksk.k(), digits, rank),
204 );
205
206 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
207 sk_in.fill_ternary_prob(0.5, &mut source_xs);
208 let sk_in_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_in);
209
210 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
211 sk_out.fill_ternary_prob(0.5, &mut source_xs);
212 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
213
214 ksk.encrypt_sk(
215 module,
216 &sk_in,
217 &sk_out,
218 &mut source_xa,
219 &mut source_xe,
220 sigma,
221 scratch.borrow(),
222 );
223
224 ct_glwe.encrypt_sk(
225 module,
226 &pt_want,
227 &sk_in_exec,
228 &mut source_xa,
229 &mut source_xe,
230 sigma,
231 scratch.borrow(),
232 );
233
234 let ksk_exec: GLWESwitchingKeyExec<Vec<u8>, B> = GLWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
235
236 ct_glwe.keyswitch_inplace(module, &ksk_exec, scratch.borrow());
237
238 let max_noise: f64 = log2_std_noise_gglwe_product(
239 module.n() as f64,
240 basek * digits,
241 0.5,
242 0.5,
243 0f64,
244 sigma * sigma,
245 0f64,
246 rank as f64,
247 k_ct,
248 k_ksk,
249 );
250
251 ct_glwe.assert_noise(module, &sk_out_exec, &pt_want, max_noise + 0.5);
252}