poulpy/poulpy-core/src/glwe_trace.rs
poulpy-core/src/glwe_trace.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::{
5 ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize,
6 VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize,
7 VecZnxNormalizeTmpBytes, VecZnxRshInplace, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
8 },
9 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx},
10};
11
12use crate::{
13 TakeGLWECt,
14 layouts::{
15 Base2K, GGLWEInfos, GLWECiphertext, GLWECiphertextLayout, GLWEInfos, LWEInfos, prepared::GGLWEAutomorphismKeyPrepared,
16 },
17 operations::GLWEOperations,
18};
19
20impl GLWECiphertext<Vec<u8>> {
21 pub fn trace_galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
22 let mut gal_els: Vec<i64> = Vec::new();
23 (0..module.log_n()).for_each(|i| {
24 if i == 0 {
25 gal_els.push(-1);
26 } else {
27 gal_els.push(module.galois_element(1 << (i - 1)));
28 }
29 });
30 gal_els
31 }
32
33 pub fn trace_scratch_space<B: Backend, OUT, IN, KEY>(
34 module: &Module<B>,
35 out_infos: &OUT,
36 in_infos: &IN,
37 key_infos: &KEY,
38 ) -> usize
39 where
40 OUT: GLWEInfos,
41 IN: GLWEInfos,
42 KEY: GGLWEInfos,
43 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
44 {
45 let trace: usize = Self::automorphism_inplace_scratch_space(module, out_infos, key_infos);
46 if in_infos.base2k() != key_infos.base2k() {
47 let glwe_conv: usize = VecZnx::alloc_bytes(
48 module.n(),
49 (key_infos.rank_out() + 1).into(),
50 out_infos.k().min(in_infos.k()).div_ceil(key_infos.base2k()) as usize,
51 ) + module.vec_znx_normalize_tmp_bytes();
52 return glwe_conv + trace;
53 }
54
55 trace
56 }
57
58 pub fn trace_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
59 where
60 OUT: GLWEInfos,
61 KEY: GGLWEInfos,
62 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
63 {
64 Self::trace_scratch_space(module, out_infos, out_infos, key_infos)
65 }
66}
67
68impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
69 pub fn trace<DataLhs: DataRef, DataAK: DataRef, B: Backend>(
70 &mut self,
71 module: &Module<B>,
72 start: usize,
73 end: usize,
74 lhs: &GLWECiphertext<DataLhs>,
75 auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
76 scratch: &mut Scratch<B>,
77 ) where
78 Module<B>: VecZnxDftAllocBytes
79 + VmpApplyDftToDftTmpBytes
80 + VecZnxBigNormalizeTmpBytes
81 + VmpApplyDftToDft<B>
82 + VmpApplyDftToDftAdd<B>
83 + VecZnxDftApply<B>
84 + VecZnxIdftApplyConsume<B>
85 + VecZnxBigAddSmallInplace<B>
86 + VecZnxBigNormalize<B>
87 + VecZnxBigAutomorphismInplace<B>
88 + VecZnxRshInplace<B>
89 + VecZnxCopy
90 + VecZnxNormalizeTmpBytes
91 + VecZnxNormalize<B>,
92 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
93 {
94 self.copy(module, lhs);
95 self.trace_inplace(module, start, end, auto_keys, scratch);
96 }
97
98 pub fn trace_inplace<DataAK: DataRef, B: Backend>(
99 &mut self,
100 module: &Module<B>,
101 start: usize,
102 end: usize,
103 auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
104 scratch: &mut Scratch<B>,
105 ) where
106 Module<B>: VecZnxDftAllocBytes
107 + VmpApplyDftToDftTmpBytes
108 + VecZnxBigNormalizeTmpBytes
109 + VmpApplyDftToDft<B>
110 + VmpApplyDftToDftAdd<B>
111 + VecZnxDftApply<B>
112 + VecZnxIdftApplyConsume<B>
113 + VecZnxBigAddSmallInplace<B>
114 + VecZnxBigNormalize<B>
115 + VecZnxBigAutomorphismInplace<B>
116 + VecZnxRshInplace<B>
117 + VecZnxNormalizeTmpBytes
118 + VecZnxNormalize<B>,
119 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
120 {
121 let basek_ksk: Base2K = auto_keys
122 .get(auto_keys.keys().next().unwrap())
123 .unwrap()
124 .base2k();
125
126 #[cfg(debug_assertions)]
127 {
128 assert_eq!(self.n(), module.n() as u32);
129 assert!(start < end);
130 assert!(end <= module.log_n());
131 for key in auto_keys.values() {
132 assert_eq!(key.n(), module.n() as u32);
133 assert_eq!(key.base2k(), basek_ksk);
134 assert_eq!(key.rank_in(), self.rank());
135 assert_eq!(key.rank_out(), self.rank());
136 }
137 }
138
139 if self.base2k() != basek_ksk {
140 let (mut self_conv, scratch_1) = scratch.take_glwe_ct(&GLWECiphertextLayout {
141 n: module.n().into(),
142 base2k: basek_ksk,
143 k: self.k(),
144 rank: self.rank(),
145 });
146
147 for j in 0..(self.rank() + 1).into() {
148 module.vec_znx_normalize(
149 basek_ksk.into(),
150 &mut self_conv.data,
151 j,
152 basek_ksk.into(),
153 &self.data,
154 j,
155 scratch_1,
156 );
157 }
158
159 for i in start..end {
160 self_conv.rsh(module, 1, scratch_1);
161
162 let p: i64 = if i == 0 {
163 -1
164 } else {
165 module.galois_element(1 << (i - 1))
166 };
167
168 if let Some(key) = auto_keys.get(&p) {
169 self_conv.automorphism_add_inplace(module, key, scratch_1);
170 } else {
171 panic!("auto_keys[{p}] is empty")
172 }
173 }
174
175 for j in 0..(self.rank() + 1).into() {
176 module.vec_znx_normalize(
177 self.base2k().into(),
178 &mut self.data,
179 j,
180 basek_ksk.into(),
181 &self_conv.data,
182 j,
183 scratch_1,
184 );
185 }
186 } else {
187 for i in start..end {
188 self.rsh(module, 1, scratch);
189
190 let p: i64 = if i == 0 {
191 -1
192 } else {
193 module.galois_element(1 << (i - 1))
194 };
195
196 if let Some(key) = auto_keys.get(&p) {
197 self.automorphism_add_inplace(module, key, scratch);
198 } else {
199 panic!("auto_keys[{p}] is empty")
200 }
201 }
202 }
203 }
204}