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