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