poulpy/poulpy-core/src/tests/test_suite/encryption/glwe_tsk.rs
poulpy-core/src/tests/test_suite/encryption/glwe_tsk.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes,
4 SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
5 VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxCopy, VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftApply,
6 VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace,
7 VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSubScalarInplace, VecZnxSwitchRing,
8 },
9 layouts::{Backend, Module, ScratchOwned, VecZnxDft},
10 oep::{
11 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
12 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
13 },
14 source::Source,
15};
16
17use crate::{
18 encryption::SIGMA,
19 layouts::{
20 GGLWETensorKey, GLWEPlaintext, GLWESecret, Infos,
21 compressed::{Decompress, GGLWETensorKeyCompressed},
22 prepared::{GLWESecretPrepared, PrepareAlloc},
23 },
24};
25
26pub fn test_gglwe_tensor_key_encrypt_sk<B>(module: &Module<B>)
27where
28 Module<B>: VecZnxDftAllocBytes
29 + VecZnxBigNormalize<B>
30 + VecZnxDftApply<B>
31 + SvpApplyDftToDftInplace<B>
32 + VecZnxIdftApplyConsume<B>
33 + VecZnxNormalizeTmpBytes
34 + VecZnxFillUniform
35 + VecZnxSubABInplace
36 + VecZnxAddInplace
37 + VecZnxNormalizeInplace<B>
38 + VecZnxAddNormal
39 + VecZnxNormalize<B>
40 + VecZnxSub
41 + SvpPrepare<B>
42 + SvpPPolAllocBytes
43 + SvpPPolAlloc<B>
44 + VecZnxBigAddSmallInplace<B>
45 + VecZnxBigAllocBytes
46 + VecZnxBigAddInplace<B>
47 + VecZnxCopy
48 + VecZnxDftAlloc<B>
49 + SvpApplyDftToDft<B>
50 + VecZnxBigAlloc<B>
51 + VecZnxIdftApplyTmpA<B>
52 + VecZnxAddScalarInplace
53 + VecZnxSwitchRing
54 + VecZnxSubScalarInplace,
55 B: Backend
56 + TakeVecZnxDftImpl<B>
57 + TakeVecZnxBigImpl<B>
58 + ScratchOwnedAllocImpl<B>
59 + ScratchOwnedBorrowImpl<B>
60 + ScratchAvailableImpl<B>
61 + TakeScalarZnxImpl<B>
62 + TakeVecZnxImpl<B>
63 + VecZnxDftAllocBytesImpl<B>
64 + VecZnxBigAllocBytesImpl<B>
65 + TakeSvpPPolImpl<B>,
66{
67 let basek: usize = 8;
68 let k: usize = 54;
69
70 (1..3).for_each(|rank| {
71 let n: usize = module.n();
72 let rows: usize = k / basek;
73
74 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k, rows, 1, rank);
75
76 let mut source_xs: Source = Source::new([0u8; 32]);
77 let mut source_xe: Source = Source::new([0u8; 32]);
78 let mut source_xa: Source = Source::new([0u8; 32]);
79
80 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWETensorKey::encrypt_sk_scratch_space(
81 module,
82 basek,
83 tensor_key.k(),
84 rank,
85 ));
86
87 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
88 sk.fill_ternary_prob(0.5, &mut source_xs);
89 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
90
91 tensor_key.encrypt_sk(
92 module,
93 &sk,
94 &mut source_xa,
95 &mut source_xe,
96 scratch.borrow(),
97 );
98
99 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
100
101 let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
102 let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
103 let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, 1);
104 let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(rank, 1);
105
106 (0..rank).for_each(|i| {
107 module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
108 });
109
110 (0..rank).for_each(|i| {
111 (0..rank).for_each(|j| {
112 module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
113 module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
114 module.vec_znx_big_normalize(
115 basek,
116 &mut sk_ij.data.as_vec_znx_mut(),
117 0,
118 &sk_ij_big,
119 0,
120 scratch.borrow(),
121 );
122 (0..tensor_key.rank_in()).for_each(|col_i| {
123 (0..tensor_key.rows()).for_each(|row_i| {
124 tensor_key
125 .at(i, j)
126 .at(row_i, col_i)
127 .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
128
129 module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
130
131 let std_pt: f64 = pt.data.std(basek, 0) * (k as f64).exp2();
132 assert!((SIGMA - std_pt).abs() <= 0.5, "{} {}", SIGMA, std_pt);
133 });
134 });
135 });
136 });
137 });
138}
139
140pub fn test_gglwe_tensor_key_compressed_encrypt_sk<B>(module: &Module<B>)
141where
142 Module<B>: VecZnxDftAllocBytes
143 + VecZnxBigNormalize<B>
144 + VecZnxDftApply<B>
145 + SvpApplyDftToDftInplace<B>
146 + VecZnxIdftApplyConsume<B>
147 + VecZnxNormalizeTmpBytes
148 + VecZnxFillUniform
149 + VecZnxSubABInplace
150 + VecZnxAddInplace
151 + VecZnxNormalizeInplace<B>
152 + VecZnxAddNormal
153 + VecZnxNormalize<B>
154 + VecZnxSub
155 + SvpPrepare<B>
156 + SvpPPolAllocBytes
157 + SvpPPolAlloc<B>
158 + VecZnxBigAddSmallInplace<B>
159 + VecZnxBigAllocBytes
160 + VecZnxBigAddInplace<B>
161 + VecZnxCopy
162 + VecZnxDftAlloc<B>
163 + SvpApplyDftToDft<B>
164 + VecZnxBigAlloc<B>
165 + VecZnxIdftApplyTmpA<B>
166 + VecZnxAddScalarInplace
167 + VecZnxSwitchRing
168 + VecZnxSubScalarInplace,
169 B: Backend
170 + TakeVecZnxDftImpl<B>
171 + TakeVecZnxBigImpl<B>
172 + ScratchOwnedAllocImpl<B>
173 + ScratchOwnedBorrowImpl<B>
174 + ScratchAvailableImpl<B>
175 + TakeScalarZnxImpl<B>
176 + TakeVecZnxImpl<B>
177 + VecZnxDftAllocBytesImpl<B>
178 + VecZnxBigAllocBytesImpl<B>
179 + TakeSvpPPolImpl<B>,
180{
181 let basek = 8;
182 let k = 54;
183 (1..3).for_each(|rank| {
184 let n: usize = module.n();
185 let rows: usize = k / basek;
186
187 let mut tensor_key_compressed: GGLWETensorKeyCompressed<Vec<u8>> =
188 GGLWETensorKeyCompressed::alloc(n, basek, k, rows, 1, rank);
189
190 let mut source_xs: Source = Source::new([0u8; 32]);
191 let mut source_xe: Source = Source::new([0u8; 32]);
192
193 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWETensorKeyCompressed::encrypt_sk_scratch_space(
194 module,
195 basek,
196 tensor_key_compressed.k(),
197 rank,
198 ));
199
200 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
201 sk.fill_ternary_prob(0.5, &mut source_xs);
202 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
203
204 let seed_xa: [u8; 32] = [1u8; 32];
205
206 tensor_key_compressed.encrypt_sk(module, &sk, seed_xa, &mut source_xe, scratch.borrow());
207
208 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k, rows, 1, rank);
209 tensor_key.decompress(module, &tensor_key_compressed);
210
211 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
212
213 let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
214 let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
215 let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, 1);
216 let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(rank, 1);
217
218 (0..rank).for_each(|i| {
219 module.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
220 });
221
222 (0..rank).for_each(|i| {
223 (0..rank).for_each(|j| {
224 module.svp_apply_dft_to_dft(&mut sk_ij_dft, 0, &sk_prepared.data, j, &sk_dft, i);
225 module.vec_znx_idft_apply_tmpa(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
226 module.vec_znx_big_normalize(
227 basek,
228 &mut sk_ij.data.as_vec_znx_mut(),
229 0,
230 &sk_ij_big,
231 0,
232 scratch.borrow(),
233 );
234 (0..tensor_key.rank_in()).for_each(|col_i| {
235 (0..tensor_key.rows()).for_each(|row_i| {
236 tensor_key
237 .at(i, j)
238 .at(row_i, col_i)
239 .decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
240
241 module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
242
243 let std_pt: f64 = pt.data.std(basek, 0) * (k as f64).exp2();
244 assert!((SIGMA - std_pt).abs() <= 0.5, "{} {}", SIGMA, std_pt);
245 });
246 });
247 });
248 });
249 });
250}