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