poulpy/poulpy-core/src/glwe_trace.rs
poulpy-core/src/glwe_trace.rsRaw
1use poulpy_hal::{
2 api::{ModuleLogN, VecZnxNormalize, VecZnxNormalizeTmpBytes},
3 layouts::{Backend, CyclotomicOrder, DataMut, GaloisElement, Module, Scratch, VecZnx, galois_element},
4};
5
6use crate::{
7 GLWEAutomorphism, GLWECopy, GLWEShift, ScratchTakeCore,
8 layouts::{
9 GGLWEInfos, GGLWELayout, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWELayout, GLWEToMut,
10 GLWEToRef, GetGaloisElement, LWEInfos,
11 },
12};
13
14impl GLWE<Vec<u8>> {
15 pub fn trace_galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
16 where
17 M: GLWETrace<BE>,
18 {
19 module.glwe_trace_galois_elements()
20 }
21
22 pub fn trace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
23 where
24 R: GLWEInfos,
25 A: GLWEInfos,
26 K: GGLWEInfos,
27 M: GLWETrace<BE>,
28 {
29 module.glwe_trace_tmp_bytes(res_infos, a_infos, key_infos)
30 }
31}
32
33impl<D: DataMut> GLWE<D> {
34 pub fn trace<A, H, K, M, BE: Backend>(&mut self, module: &M, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
35 where
36 A: GLWEToRef,
37 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
38 H: GLWEAutomorphismKeyHelper<K, BE>,
39 Scratch<BE>: ScratchTakeCore<BE>,
40 M: GLWETrace<BE>,
41 {
42 module.glwe_trace(self, skip, a, keys, scratch);
43 }
44
45 pub fn trace_inplace<H, K, M, BE: Backend>(&mut self, module: &M, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
46 where
47 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
48 H: GLWEAutomorphismKeyHelper<K, BE>,
49 Scratch<BE>: ScratchTakeCore<BE>,
50 M: GLWETrace<BE>,
51 {
52 module.glwe_trace_inplace(self, skip, keys, scratch);
53 }
54}
55
56#[inline(always)]
57pub fn trace_galois_elements(log_n: usize, cyclotomic_order: i64) -> Vec<i64> {
58 (0..log_n)
59 .map(|i| {
60 if i == 0 {
61 -1
62 } else {
63 galois_element(1 << (i - 1), cyclotomic_order)
64 }
65 })
66 .collect()
67}
68
69impl<BE: Backend> GLWETrace<BE> for Module<BE>
70where
71 Self: ModuleLogN
72 + GaloisElement
73 + GLWEAutomorphism<BE>
74 + GLWEShift<BE>
75 + GLWECopy
76 + CyclotomicOrder
77 + VecZnxNormalizeTmpBytes
78 + VecZnxNormalize<BE>,
79 Scratch<BE>: ScratchTakeCore<BE>,
80{
81 fn glwe_trace_galois_elements(&self) -> Vec<i64> {
82 trace_galois_elements(self.log_n(), self.cyclotomic_order())
83 }
84
85 fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
86 where
87 R: GLWEInfos,
88 A: GLWEInfos,
89 K: GGLWEInfos,
90 {
91 let trace: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
92 if a_infos.base2k() != key_infos.base2k() {
93 let glwe_conv: usize = VecZnx::bytes_of(
94 self.n(),
95 (key_infos.rank_out() + 1).into(),
96 res_infos.k().min(a_infos.k()).div_ceil(key_infos.base2k()) as usize,
97 ) + self.vec_znx_normalize_tmp_bytes();
98 return glwe_conv + trace;
99 }
100
101 trace
102 }
103
104 fn glwe_trace<R, A, K, H>(&self, res: &mut R, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
105 where
106 R: GLWEToMut,
107 A: GLWEToRef,
108 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
109 H: GLWEAutomorphismKeyHelper<K, BE>,
110 {
111 self.glwe_copy(res, a);
112 self.glwe_trace_inplace(res, skip, keys, scratch);
113 }
114
115 fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
116 where
117 R: GLWEToMut,
118 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
119 H: GLWEAutomorphismKeyHelper<K, BE>,
120 {
121 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
122
123 let ksk_infos: &GGLWELayout = &keys.automorphism_key_infos();
124 let log_n: usize = self.log_n();
125
126 assert_eq!(res.n(), self.n() as u32);
127 assert_eq!(ksk_infos.n(), self.n() as u32);
128 assert!(skip <= log_n);
129 assert_eq!(ksk_infos.rank_in(), res.rank());
130 assert_eq!(ksk_infos.rank_out(), res.rank());
131
132 if res.base2k() != ksk_infos.base2k() {
133 let (mut self_conv, scratch_1) = scratch.take_glwe(&GLWELayout {
134 n: self.n().into(),
135 base2k: ksk_infos.base2k(),
136 k: res.k(),
137 rank: res.rank(),
138 });
139
140 for j in 0..(res.rank() + 1).into() {
141 self.vec_znx_normalize(
142 ksk_infos.base2k().into(),
143 &mut self_conv.data,
144 j,
145 res.base2k().into(),
146 res.data(),
147 j,
148 scratch_1,
149 );
150 }
151
152 for i in skip..log_n {
153 self.glwe_rsh(1, &mut self_conv, scratch_1);
154
155 let p: i64 = if i == 0 {
156 -1
157 } else {
158 self.galois_element(1 << (i - 1))
159 };
160
161 if let Some(key) = keys.get_automorphism_key(p) {
162 self.glwe_automorphism_add_inplace(&mut self_conv, key, scratch_1);
163 } else {
164 panic!("keys[{p}] is empty")
165 }
166 }
167
168 for j in 0..(res.rank() + 1).into() {
169 self.vec_znx_normalize(
170 res.base2k().into(),
171 res.data_mut(),
172 j,
173 ksk_infos.base2k().into(),
174 &self_conv.data,
175 j,
176 scratch_1,
177 );
178 }
179 } else {
180 // println!("res: {}", res);
181
182 for i in skip..log_n {
183 self.glwe_rsh(1, res, scratch);
184
185 let p: i64 = if i == 0 {
186 -1
187 } else {
188 self.galois_element(1 << (i - 1))
189 };
190
191 if let Some(key) = keys.get_automorphism_key(p) {
192 self.glwe_automorphism_add_inplace(res, key, scratch);
193 } else {
194 panic!("keys[{p}] is empty")
195 }
196 }
197 }
198 }
199}
200
201pub trait GLWETrace<BE: Backend> {
202 fn glwe_trace_galois_elements(&self) -> Vec<i64>;
203
204 fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
205 where
206 R: GLWEInfos,
207 A: GLWEInfos,
208 K: GGLWEInfos;
209
210 fn glwe_trace<R, A, K, H>(&self, res: &mut R, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
211 where
212 R: GLWEToMut,
213 A: GLWEToRef,
214 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
215 H: GLWEAutomorphismKeyHelper<K, BE>;
216
217 fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
218 where
219 R: GLWEToMut,
220 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
221 H: GLWEAutomorphismKeyHelper<K, BE>;
222}