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