poulpy/poulpy-core/src/automorphism/ggsw_ct.rs
poulpy-core/src/automorphism/ggsw_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 DFT, IDFTConsume, IDFTTmpA, ScratchAvailable, TakeVecZnxBig, TakeVecZnxDft, VecZnxAutomorphismInplace,
4 VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace,
5 VecZnxDftAllocBytes, VecZnxDftCopy, VecZnxNormalizeTmpBytes, VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
6 },
7 layouts::{Backend, DataMut, DataRef, Module, Scratch},
8};
9
10use crate::layouts::{
11 GGSWCiphertext, GLWECiphertext, Infos,
12 prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared},
13};
14
15impl GGSWCiphertext<Vec<u8>> {
16 #[allow(clippy::too_many_arguments)]
17 pub fn automorphism_scratch_space<B: Backend>(
18 module: &Module<B>,
19 basek: usize,
20 k_out: usize,
21 k_in: usize,
22 k_ksk: usize,
23 digits_ksk: usize,
24 k_tsk: usize,
25 digits_tsk: usize,
26 rank: usize,
27 ) -> usize
28 where
29 Module<B>:
30 VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
31 {
32 let out_size: usize = k_out.div_ceil(basek);
33 let ci_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, out_size);
34 let ks_internal: usize =
35 GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits_ksk, rank, rank);
36 let expand: usize = GGSWCiphertext::expand_row_scratch_space(module, basek, k_out, k_tsk, digits_tsk, rank);
37 ci_dft + (ks_internal | expand)
38 }
39
40 #[allow(clippy::too_many_arguments)]
41 pub fn automorphism_inplace_scratch_space<B: Backend>(
42 module: &Module<B>,
43 basek: usize,
44 k_out: usize,
45 k_ksk: usize,
46 digits_ksk: usize,
47 k_tsk: usize,
48 digits_tsk: usize,
49 rank: usize,
50 ) -> usize
51 where
52 Module<B>:
53 VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
54 {
55 GGSWCiphertext::automorphism_scratch_space(
56 module, basek, k_out, k_out, k_ksk, digits_ksk, k_tsk, digits_tsk, rank,
57 )
58 }
59}
60
61impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
62 pub fn automorphism<DataLhs: DataRef, DataAk: DataRef, DataTsk: DataRef, B: Backend>(
63 &mut self,
64 module: &Module<B>,
65 lhs: &GGSWCiphertext<DataLhs>,
66 auto_key: &GGLWEAutomorphismKeyPrepared<DataAk, B>,
67 tensor_key: &GGLWETensorKeyPrepared<DataTsk, B>,
68 scratch: &mut Scratch<B>,
69 ) where
70 Module<B>: VecZnxDftAllocBytes
71 + VmpApplyTmpBytes
72 + VecZnxBigNormalizeTmpBytes
73 + VmpApply<B>
74 + VmpApplyAdd<B>
75 + DFT<B>
76 + IDFTConsume<B>
77 + VecZnxBigAddSmallInplace<B>
78 + VecZnxBigNormalize<B>
79 + VecZnxAutomorphismInplace
80 + VecZnxBigAllocBytes
81 + VecZnxNormalizeTmpBytes
82 + VecZnxDftCopy<B>
83 + VecZnxDftAddInplace<B>
84 + IDFTTmpA<B>,
85 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B>,
86 {
87 #[cfg(debug_assertions)]
88 {
89 assert_eq!(self.n(), auto_key.n());
90 assert_eq!(lhs.n(), auto_key.n());
91
92 assert_eq!(
93 self.rank(),
94 lhs.rank(),
95 "ggsw_out rank: {} != ggsw_in rank: {}",
96 self.rank(),
97 lhs.rank()
98 );
99 assert_eq!(
100 self.rank(),
101 auto_key.rank(),
102 "ggsw_in rank: {} != auto_key rank: {}",
103 self.rank(),
104 auto_key.rank()
105 );
106 assert_eq!(
107 self.rank(),
108 tensor_key.rank(),
109 "ggsw_in rank: {} != tensor_key rank: {}",
110 self.rank(),
111 tensor_key.rank()
112 );
113 assert!(
114 scratch.available()
115 >= GGSWCiphertext::automorphism_scratch_space(
116 module,
117 self.basek(),
118 self.k(),
119 lhs.k(),
120 auto_key.k(),
121 auto_key.digits(),
122 tensor_key.k(),
123 tensor_key.digits(),
124 self.rank(),
125 )
126 )
127 };
128
129 self.automorphism_internal(module, lhs, auto_key, scratch);
130 self.expand_row(module, tensor_key, scratch);
131 }
132
133 pub fn automorphism_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
134 &mut self,
135 module: &Module<B>,
136 auto_key: &GGLWEAutomorphismKeyPrepared<DataKsk, B>,
137 tensor_key: &GGLWETensorKeyPrepared<DataTsk, B>,
138 scratch: &mut Scratch<B>,
139 ) where
140 Module<B>: VecZnxDftAllocBytes
141 + VmpApplyTmpBytes
142 + VecZnxBigNormalizeTmpBytes
143 + VmpApply<B>
144 + VmpApplyAdd<B>
145 + DFT<B>
146 + IDFTConsume<B>
147 + VecZnxBigAddSmallInplace<B>
148 + VecZnxBigNormalize<B>
149 + VecZnxAutomorphismInplace
150 + VecZnxBigAllocBytes
151 + VecZnxNormalizeTmpBytes
152 + VecZnxDftCopy<B>
153 + VecZnxDftAddInplace<B>
154 + IDFTTmpA<B>,
155 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B>,
156 {
157 unsafe {
158 let self_ptr: *mut GGSWCiphertext<DataSelf> = self as *mut GGSWCiphertext<DataSelf>;
159 self.automorphism(module, &*self_ptr, auto_key, tensor_key, scratch);
160 }
161 }
162
163 fn automorphism_internal<DataLhs: DataRef, DataAk: DataRef, B: Backend>(
164 &mut self,
165 module: &Module<B>,
166 lhs: &GGSWCiphertext<DataLhs>,
167 auto_key: &GGLWEAutomorphismKeyPrepared<DataAk, B>,
168 scratch: &mut Scratch<B>,
169 ) where
170 Module<B>: VecZnxDftAllocBytes
171 + VmpApplyTmpBytes
172 + VecZnxBigNormalizeTmpBytes
173 + VmpApply<B>
174 + VmpApplyAdd<B>
175 + DFT<B>
176 + IDFTConsume<B>
177 + VecZnxBigAddSmallInplace<B>
178 + VecZnxBigNormalize<B>
179 + VecZnxAutomorphismInplace,
180 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
181 {
182 // Keyswitch the j-th row of the col 0
183 (0..lhs.rows()).for_each(|row_i| {
184 // Key-switch column 0, i.e.
185 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
186 self.at_mut(row_i, 0)
187 .automorphism(module, &lhs.at(row_i, 0), auto_key, scratch);
188 });
189 }
190}