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