poulpy/core/src/glwe/automorphism.rs
core/src/glwe/automorphism.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAutomorphismInplace, VecZnxBigSubSmallAInplace,
4 VecZnxBigSubSmallBInplace,
5 },
6 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnxBig},
7};
8
9use crate::{AutomorphismKeyExec, GLWECiphertext, GLWEKeyswitchFamily, Infos, glwe::keyswitch::keyswitch};
10
11impl GLWECiphertext<Vec<u8>> {
12 pub fn automorphism_scratch_space<B: Backend>(
13 module: &Module<B>,
14 n: usize,
15 basek: usize,
16 k_out: usize,
17 k_in: usize,
18 k_ksk: usize,
19 digits: usize,
20 rank: usize,
21 ) -> usize
22 where
23 Module<B>: GLWEKeyswitchFamily<B>,
24 {
25 Self::keyswitch_scratch_space(module, n, basek, k_out, k_in, k_ksk, digits, rank, rank)
26 }
27
28 pub fn automorphism_inplace_scratch_space<B: Backend>(
29 module: &Module<B>,
30 n: usize,
31 basek: usize,
32 k_out: usize,
33 k_ksk: usize,
34 digits: usize,
35 rank: usize,
36 ) -> usize
37 where
38 Module<B>: GLWEKeyswitchFamily<B>,
39 {
40 Self::keyswitch_inplace_scratch_space(module, n, basek, k_out, k_ksk, digits, rank)
41 }
42}
43
44impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
45 pub fn automorphism<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
46 &mut self,
47 module: &Module<B>,
48 lhs: &GLWECiphertext<DataLhs>,
49 rhs: &AutomorphismKeyExec<DataRhs, B>,
50 scratch: &mut Scratch<B>,
51 ) where
52 Module<B>: GLWEKeyswitchFamily<B> + VecZnxAutomorphismInplace,
53 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
54 {
55 self.keyswitch(module, lhs, &rhs.key, scratch);
56 (0..self.rank() + 1).for_each(|i| {
57 module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
58 })
59 }
60
61 pub fn automorphism_inplace<DataRhs: DataRef, B: Backend>(
62 &mut self,
63 module: &Module<B>,
64 rhs: &AutomorphismKeyExec<DataRhs, B>,
65 scratch: &mut Scratch<B>,
66 ) where
67 Module<B>: GLWEKeyswitchFamily<B> + VecZnxAutomorphismInplace,
68 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
69 {
70 self.keyswitch_inplace(module, &rhs.key, scratch);
71 (0..self.rank() + 1).for_each(|i| {
72 module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
73 })
74 }
75
76 pub fn automorphism_add<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
77 &mut self,
78 module: &Module<B>,
79 lhs: &GLWECiphertext<DataLhs>,
80 rhs: &AutomorphismKeyExec<DataRhs, B>,
81 scratch: &mut Scratch<B>,
82 ) where
83 Module<B>: GLWEKeyswitchFamily<B> + VecZnxBigAutomorphismInplace<B>,
84 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
85 {
86 #[cfg(debug_assertions)]
87 {
88 self.assert_keyswitch(module, lhs, &rhs.key, scratch);
89 }
90 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
91 let mut res_big: VecZnxBig<_, B> = keyswitch(module, res_dft, lhs, &rhs.key, scratch1);
92 (0..self.cols()).for_each(|i| {
93 module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
94 module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i);
95 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
96 })
97 }
98
99 pub fn automorphism_add_inplace<DataRhs: DataRef, B: Backend>(
100 &mut self,
101 module: &Module<B>,
102 rhs: &AutomorphismKeyExec<DataRhs, B>,
103 scratch: &mut Scratch<B>,
104 ) where
105 Module<B>: GLWEKeyswitchFamily<B> + VecZnxBigAutomorphismInplace<B>,
106 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
107 {
108 unsafe {
109 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
110 self.automorphism_add(module, &*self_ptr, rhs, scratch);
111 }
112 }
113
114 pub fn automorphism_sub_ab<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
115 &mut self,
116 module: &Module<B>,
117 lhs: &GLWECiphertext<DataLhs>,
118 rhs: &AutomorphismKeyExec<DataRhs, B>,
119 scratch: &mut Scratch<B>,
120 ) where
121 Module<B>: GLWEKeyswitchFamily<B> + VecZnxBigAutomorphismInplace<B> + VecZnxBigSubSmallAInplace<B>,
122 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
123 {
124 #[cfg(debug_assertions)]
125 {
126 self.assert_keyswitch(module, lhs, &rhs.key, scratch);
127 }
128 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
129 let mut res_big: VecZnxBig<_, B> = keyswitch(module, res_dft, lhs, &rhs.key, scratch1);
130 (0..self.cols()).for_each(|i| {
131 module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
132 module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i);
133 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
134 })
135 }
136
137 pub fn automorphism_sub_ab_inplace<DataRhs: DataRef, B: Backend>(
138 &mut self,
139 module: &Module<B>,
140 rhs: &AutomorphismKeyExec<DataRhs, B>,
141 scratch: &mut Scratch<B>,
142 ) where
143 Module<B>: GLWEKeyswitchFamily<B> + VecZnxBigAutomorphismInplace<B> + VecZnxBigSubSmallAInplace<B>,
144 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
145 {
146 unsafe {
147 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
148 self.automorphism_sub_ab(module, &*self_ptr, rhs, scratch);
149 }
150 }
151
152 pub fn automorphism_sub_ba<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
153 &mut self,
154 module: &Module<B>,
155 lhs: &GLWECiphertext<DataLhs>,
156 rhs: &AutomorphismKeyExec<DataRhs, B>,
157 scratch: &mut Scratch<B>,
158 ) where
159 Module<B>: GLWEKeyswitchFamily<B> + VecZnxBigAutomorphismInplace<B> + VecZnxBigSubSmallBInplace<B>,
160 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
161 {
162 #[cfg(debug_assertions)]
163 {
164 self.assert_keyswitch(module, lhs, &rhs.key, scratch);
165 }
166 let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
167 let mut res_big: VecZnxBig<_, B> = keyswitch(module, res_dft, lhs, &rhs.key, scratch1);
168 (0..self.cols()).for_each(|i| {
169 module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
170 module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i);
171 module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
172 })
173 }
174
175 pub fn automorphism_sub_ba_inplace<DataRhs: DataRef, B: Backend>(
176 &mut self,
177 module: &Module<B>,
178 rhs: &AutomorphismKeyExec<DataRhs, B>,
179 scratch: &mut Scratch<B>,
180 ) where
181 Module<B>: GLWEKeyswitchFamily<B> + VecZnxBigAutomorphismInplace<B> + VecZnxBigSubSmallBInplace<B>,
182 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
183 {
184 unsafe {
185 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
186 self.automorphism_sub_ba(module, &*self_ptr, rhs, scratch);
187 }
188 }
189}