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