poulpy/core/src/gglwe/test/automorphism_key.rs
core/src/gglwe/test/automorphism_key.rsRaw
1use backend::{
2 hal::{
3 api::{ModuleNew, ScalarZnxAutomorphism, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxStd, VecZnxSubScalarInplace},
4 layouts::{Module, ScratchOwned},
5 },
6 implementation::cpu_spqlios::FFT64,
7};
8use sampling::source::Source;
9
10use crate::{
11 AutomorphismKey, AutomorphismKeyExec, GLWEPlaintext, GLWESecret, GLWESecretExec, Infos, noise::log2_std_noise_gglwe_product,
12};
13
14#[test]
15fn automorphism() {
16 let log_n: usize = 8;
17 let basek: usize = 12;
18 let k_in: usize = 60;
19 let k_out: usize = 40;
20 let digits: usize = k_in.div_ceil(basek);
21 let sigma: f64 = 3.2;
22 (1..4).for_each(|rank| {
23 (2..digits + 1).for_each(|di| {
24 println!("test automorphism digits: {} rank: {}", di, rank);
25 let k_apply: usize = (digits + di) * basek;
26 test_automorphism(-1, 5, log_n, basek, di, k_in, k_out, k_apply, sigma, rank);
27 });
28 });
29}
30
31#[test]
32fn automorphism_inplace() {
33 let log_n: usize = 8;
34 let basek: usize = 12;
35 let k_in: usize = 60;
36 let digits: usize = k_in.div_ceil(basek);
37 let sigma: f64 = 3.2;
38 (1..4).for_each(|rank| {
39 (2..digits + 1).for_each(|di| {
40 println!("test automorphism digits: {} rank: {}", di, rank);
41 let k_apply: usize = (digits + di) * basek;
42 test_automorphism_inplace(-1, 5, log_n, basek, di, k_in, k_apply, sigma, rank);
43 });
44 });
45}
46
47fn test_automorphism(
48 p0: i64,
49 p1: i64,
50 log_n: usize,
51 basek: usize,
52 digits: usize,
53 k_in: usize,
54 k_out: usize,
55 k_apply: usize,
56 sigma: f64,
57 rank: usize,
58) {
59 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
60
61 let digits_in: usize = 1;
62
63 let rows_in: usize = k_in / (basek * digits);
64 let rows_apply: usize = k_in.div_ceil(basek * digits);
65
66 let mut auto_key_in: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(&module, basek, k_in, rows_in, digits_in, rank);
67 let mut auto_key_out: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(&module, basek, k_out, rows_in, digits_in, rank);
68 let mut auto_key_apply: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(&module, basek, k_apply, rows_apply, digits, rank);
69
70 let mut source_xs: Source = Source::new([0u8; 32]);
71 let mut source_xe: Source = Source::new([0u8; 32]);
72 let mut source_xa: Source = Source::new([0u8; 32]);
73
74 let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
75 AutomorphismKey::encrypt_sk_scratch_space(&module, basek, k_apply, rank)
76 | AutomorphismKey::automorphism_scratch_space(&module, basek, k_out, k_in, k_apply, digits, rank),
77 );
78
79 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
80 sk.fill_ternary_prob(0.5, &mut source_xs);
81
82 // gglwe_{s1}(s0) = s0 -> s1
83 auto_key_in.encrypt_sk(
84 &module,
85 p0,
86 &sk,
87 &mut source_xa,
88 &mut source_xe,
89 sigma,
90 scratch.borrow(),
91 );
92
93 // gglwe_{s2}(s1) -> s1 -> s2
94 auto_key_apply.encrypt_sk(
95 &module,
96 p1,
97 &sk,
98 &mut source_xa,
99 &mut source_xe,
100 sigma,
101 scratch.borrow(),
102 );
103
104 let mut auto_key_apply_exec: AutomorphismKeyExec<Vec<u8>, FFT64> =
105 AutomorphismKeyExec::alloc(&module, basek, k_apply, rows_apply, digits, rank);
106
107 auto_key_apply_exec.prepare(&module, &auto_key_apply, scratch.borrow());
108
109 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
110 auto_key_out.automorphism(
111 &module,
112 &auto_key_in,
113 &auto_key_apply_exec,
114 scratch.borrow(),
115 );
116
117 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
118
119 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
120 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
121 (0..rank).for_each(|i| {
122 module.scalar_znx_automorphism(
123 module.galois_element_inv(p0 * p1),
124 &mut sk_auto.data,
125 i,
126 &sk.data,
127 i,
128 );
129 });
130
131 let sk_auto_dft: GLWESecretExec<Vec<u8>, FFT64> = GLWESecretExec::from(&module, &sk_auto);
132
133 (0..auto_key_out.rank_in()).for_each(|col_i| {
134 (0..auto_key_out.rows()).for_each(|row_i| {
135 auto_key_out
136 .at(row_i, col_i)
137 .decrypt(&module, &mut pt, &sk_auto_dft, scratch.borrow());
138
139 module.vec_znx_sub_scalar_inplace(
140 &mut pt.data,
141 0,
142 (digits_in - 1) + row_i * digits_in,
143 &sk.data,
144 col_i,
145 );
146
147 let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
148 let noise_want: f64 = log2_std_noise_gglwe_product(
149 module.n() as f64,
150 basek * digits,
151 0.5,
152 0.5,
153 0f64,
154 sigma * sigma,
155 0f64,
156 rank as f64,
157 k_out,
158 k_apply,
159 );
160
161 assert!(
162 noise_have < noise_want + 0.5,
163 "{} {}",
164 noise_have,
165 noise_want
166 );
167 });
168 });
169}
170
171fn test_automorphism_inplace(
172 p0: i64,
173 p1: i64,
174 log_n: usize,
175 basek: usize,
176 digits: usize,
177 k_in: usize,
178 k_apply: usize,
179 sigma: f64,
180 rank: usize,
181) {
182 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
183
184 let digits_in: usize = 1;
185
186 let rows_in: usize = k_in / (basek * digits);
187 let rows_apply: usize = k_in.div_ceil(basek * digits);
188
189 let mut auto_key: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(&module, basek, k_in, rows_in, digits_in, rank);
190 let mut auto_key_apply: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(&module, basek, k_apply, rows_apply, digits, rank);
191
192 let mut source_xs: Source = Source::new([0u8; 32]);
193 let mut source_xe: Source = Source::new([0u8; 32]);
194 let mut source_xa: Source = Source::new([0u8; 32]);
195
196 let mut scratch: ScratchOwned<FFT64> = ScratchOwned::alloc(
197 AutomorphismKey::encrypt_sk_scratch_space(&module, basek, k_apply, rank)
198 | AutomorphismKey::automorphism_inplace_scratch_space(&module, basek, k_in, k_apply, digits, rank),
199 );
200
201 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
202 sk.fill_ternary_prob(0.5, &mut source_xs);
203
204 // gglwe_{s1}(s0) = s0 -> s1
205 auto_key.encrypt_sk(
206 &module,
207 p0,
208 &sk,
209 &mut source_xa,
210 &mut source_xe,
211 sigma,
212 scratch.borrow(),
213 );
214
215 // gglwe_{s2}(s1) -> s1 -> s2
216 auto_key_apply.encrypt_sk(
217 &module,
218 p1,
219 &sk,
220 &mut source_xa,
221 &mut source_xe,
222 sigma,
223 scratch.borrow(),
224 );
225
226 let mut auto_key_apply_exec: AutomorphismKeyExec<Vec<u8>, FFT64> =
227 AutomorphismKeyExec::alloc(&module, basek, k_apply, rows_apply, digits, rank);
228
229 auto_key_apply_exec.prepare(&module, &auto_key_apply, scratch.borrow());
230
231 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
232 auto_key.automorphism_inplace(&module, &auto_key_apply_exec, scratch.borrow());
233
234 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_in);
235
236 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
237 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
238
239 (0..rank).for_each(|i| {
240 module.scalar_znx_automorphism(
241 module.galois_element_inv(p0 * p1),
242 &mut sk_auto.data,
243 i,
244 &sk.data,
245 i,
246 );
247 });
248
249 let sk_auto_dft: GLWESecretExec<Vec<u8>, FFT64> = GLWESecretExec::from(&module, &sk_auto);
250
251 (0..auto_key.rank_in()).for_each(|col_i| {
252 (0..auto_key.rows()).for_each(|row_i| {
253 auto_key
254 .at(row_i, col_i)
255 .decrypt(&module, &mut pt, &sk_auto_dft, scratch.borrow());
256 module.vec_znx_sub_scalar_inplace(
257 &mut pt.data,
258 0,
259 (digits_in - 1) + row_i * digits_in,
260 &sk.data,
261 col_i,
262 );
263
264 let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
265 let noise_want: f64 = log2_std_noise_gglwe_product(
266 module.n() as f64,
267 basek * digits,
268 0.5,
269 0.5,
270 0f64,
271 sigma * sigma,
272 0f64,
273 rank as f64,
274 k_in,
275 k_apply,
276 );
277
278 assert!(
279 noise_have < noise_want + 0.5,
280 "{} {}",
281 noise_have,
282 noise_want
283 );
284 });
285 });
286}