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