poulpy/poulpy-core/src/keyswitching/gglwe_ct.rs
poulpy-core/src/keyswitching/gglwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
4 VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
5 VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
6 },
7 layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero},
8};
9
10use crate::layouts::{
11 GGLWEAutomorphismKey, GGLWEInfos, GGLWESwitchingKey, GLWECiphertext, GLWEInfos,
12 prepared::{GGLWEAutomorphismKeyPrepared, GGLWESwitchingKeyPrepared},
13};
14
15impl GGLWEAutomorphismKey<Vec<u8>> {
16 pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY>(
17 module: &Module<B>,
18 out_infos: &OUT,
19 in_infos: &IN,
20 key_infos: &KEY,
21 ) -> usize
22 where
23 OUT: GGLWEInfos,
24 IN: GGLWEInfos,
25 KEY: GGLWEInfos,
26 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
27 {
28 GGLWESwitchingKey::keyswitch_scratch_space(module, out_infos, in_infos, key_infos)
29 }
30
31 pub fn keyswitch_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
32 where
33 OUT: GGLWEInfos,
34 KEY: GGLWEInfos,
35 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
36 {
37 GGLWESwitchingKey::keyswitch_inplace_scratch_space(module, out_infos, key_infos)
38 }
39}
40
41impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
42 pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
43 &mut self,
44 module: &Module<B>,
45 lhs: &GGLWEAutomorphismKey<DataLhs>,
46 rhs: &GGLWESwitchingKeyPrepared<DataRhs, B>,
47 scratch: &mut Scratch<B>,
48 ) where
49 Module<B>: VecZnxDftAllocBytes
50 + VmpApplyDftToDftTmpBytes
51 + VecZnxBigNormalizeTmpBytes
52 + VmpApplyDftToDft<B>
53 + VmpApplyDftToDftAdd<B>
54 + VecZnxDftApply<B>
55 + VecZnxIdftApplyConsume<B>
56 + VecZnxBigAddSmallInplace<B>
57 + VecZnxBigNormalize<B>
58 + VecZnxNormalize<B>
59 + VecZnxNormalizeTmpBytes,
60 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
61 {
62 self.key.keyswitch(module, &lhs.key, rhs, scratch);
63 }
64
65 pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
66 &mut self,
67 module: &Module<B>,
68 rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
69 scratch: &mut Scratch<B>,
70 ) where
71 Module<B>: VecZnxDftAllocBytes
72 + VmpApplyDftToDftTmpBytes
73 + VecZnxBigNormalizeTmpBytes
74 + VmpApplyDftToDft<B>
75 + VmpApplyDftToDftAdd<B>
76 + VecZnxDftApply<B>
77 + VecZnxIdftApplyConsume<B>
78 + VecZnxBigAddSmallInplace<B>
79 + VecZnxBigNormalize<B>
80 + VecZnxNormalize<B>
81 + VecZnxNormalizeTmpBytes,
82 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
83 {
84 self.key.keyswitch_inplace(module, &rhs.key, scratch);
85 }
86}
87
88impl GGLWESwitchingKey<Vec<u8>> {
89 pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY>(
90 module: &Module<B>,
91 out_infos: &OUT,
92 in_infos: &IN,
93 key_apply: &KEY,
94 ) -> usize
95 where
96 OUT: GGLWEInfos,
97 IN: GGLWEInfos,
98 KEY: GGLWEInfos,
99 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
100 {
101 GLWECiphertext::keyswitch_scratch_space(module, out_infos, in_infos, key_apply)
102 }
103
104 pub fn keyswitch_inplace_scratch_space<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_apply: &KEY) -> usize
105 where
106 OUT: GGLWEInfos + GLWEInfos,
107 KEY: GGLWEInfos + GLWEInfos,
108 Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
109 {
110 GLWECiphertext::keyswitch_inplace_scratch_space(module, out_infos, key_apply)
111 }
112}
113
114impl<DataSelf: DataMut> GGLWESwitchingKey<DataSelf> {
115 pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
116 &mut self,
117 module: &Module<B>,
118 lhs: &GGLWESwitchingKey<DataLhs>,
119 rhs: &GGLWESwitchingKeyPrepared<DataRhs, B>,
120 scratch: &mut Scratch<B>,
121 ) where
122 Module<B>: VecZnxDftAllocBytes
123 + VmpApplyDftToDftTmpBytes
124 + VecZnxBigNormalizeTmpBytes
125 + VmpApplyDftToDft<B>
126 + VmpApplyDftToDftAdd<B>
127 + VecZnxDftApply<B>
128 + VecZnxIdftApplyConsume<B>
129 + VecZnxBigAddSmallInplace<B>
130 + VecZnxBigNormalize<B>
131 + VecZnxNormalize<B>
132 + VecZnxNormalizeTmpBytes,
133 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
134 {
135 #[cfg(debug_assertions)]
136 {
137 assert_eq!(
138 self.rank_in(),
139 lhs.rank_in(),
140 "ksk_out input rank: {} != ksk_in input rank: {}",
141 self.rank_in(),
142 lhs.rank_in()
143 );
144 assert_eq!(
145 lhs.rank_out(),
146 rhs.rank_in(),
147 "ksk_in output rank: {} != ksk_apply input rank: {}",
148 self.rank_out(),
149 rhs.rank_in()
150 );
151 assert_eq!(
152 self.rank_out(),
153 rhs.rank_out(),
154 "ksk_out output rank: {} != ksk_apply output rank: {}",
155 self.rank_out(),
156 rhs.rank_out()
157 );
158 assert!(
159 self.dnum() <= lhs.dnum(),
160 "self.dnum()={} > lhs.dnum()={}",
161 self.dnum(),
162 lhs.dnum()
163 );
164 assert_eq!(
165 self.dsize(),
166 lhs.dsize(),
167 "ksk_out dsize: {} != ksk_in dsize: {}",
168 self.dsize(),
169 lhs.dsize()
170 )
171 }
172
173 (0..self.rank_in().into()).for_each(|col_i| {
174 (0..self.dnum().into()).for_each(|row_j| {
175 self.at_mut(row_j, col_i)
176 .keyswitch(module, &lhs.at(row_j, col_i), rhs, scratch);
177 });
178 });
179
180 (self.dnum().min(lhs.dnum()).into()..self.dnum().into()).for_each(|row_i| {
181 (0..self.rank_in().into()).for_each(|col_j| {
182 self.at_mut(row_i, col_j).data.zero();
183 });
184 });
185 }
186
187 pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
188 &mut self,
189 module: &Module<B>,
190 rhs: &GGLWESwitchingKeyPrepared<DataRhs, B>,
191 scratch: &mut Scratch<B>,
192 ) where
193 Module<B>: VecZnxDftAllocBytes
194 + VmpApplyDftToDftTmpBytes
195 + VecZnxBigNormalizeTmpBytes
196 + VmpApplyDftToDft<B>
197 + VmpApplyDftToDftAdd<B>
198 + VecZnxDftApply<B>
199 + VecZnxIdftApplyConsume<B>
200 + VecZnxBigAddSmallInplace<B>
201 + VecZnxBigNormalize<B>
202 + VecZnxNormalize<B>
203 + VecZnxNormalizeTmpBytes,
204 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnx,
205 {
206 #[cfg(debug_assertions)]
207 {
208 assert_eq!(
209 self.rank_out(),
210 rhs.rank_out(),
211 "ksk_out output rank: {} != ksk_apply output rank: {}",
212 self.rank_out(),
213 rhs.rank_out()
214 );
215 }
216
217 (0..self.rank_in().into()).for_each(|col_i| {
218 (0..self.dnum().into()).for_each(|row_j| {
219 self.at_mut(row_j, col_i)
220 .keyswitch_inplace(module, rhs, scratch)
221 });
222 });
223 }
224}