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