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