poulpy/core/src/glwe_ops.rs
core/src/glwe_ops.rsRaw
1use backend::{Backend, Module, Scratch, VecZnx, VecZnxOps, VecZnxToMut, VecZnxToRef, ZnxZero};
2
3use crate::elem::{Infos, SetMetaData};
4
5pub trait GLWEOps<BACKEND: Backend>
6where
7 Self: Sized + VecZnxToMut + SetMetaData + Infos,
8{
9 fn add<A, B>(&mut self, module: &Module<BACKEND>, a: &A, b: &B)
10 where
11 A: VecZnxToRef + Infos,
12 B: VecZnxToRef + Infos,
13 {
14 #[cfg(debug_assertions)]
15 {
16 assert_eq!(a.n(), module.n());
17 assert_eq!(b.n(), module.n());
18 assert_eq!(self.n(), module.n());
19 assert_eq!(a.basek(), b.basek());
20 assert!(self.rank() >= a.rank().max(b.rank()));
21 }
22
23 let min_col: usize = a.rank().min(b.rank()) + 1;
24 let max_col: usize = a.rank().max(b.rank() + 1);
25 let self_col: usize = self.rank() + 1;
26
27 (0..min_col).for_each(|i| {
28 module.vec_znx_add(self, i, a, i, b, i);
29 });
30
31 if a.rank() > b.rank() {
32 (min_col..max_col).for_each(|i| {
33 module.vec_znx_copy(self, i, a, i);
34 });
35 } else {
36 (min_col..max_col).for_each(|i| {
37 module.vec_znx_copy(self, i, b, i);
38 });
39 }
40
41 let size: usize = self.size();
42 let mut self_mut: VecZnx<&mut [u8]> = self.to_mut();
43 (max_col..self_col).for_each(|i| {
44 (0..size).for_each(|j| {
45 self_mut.zero_at(i, j);
46 });
47 });
48
49 self.set_basek(a.basek());
50 self.set_k(a.k().max(b.k()));
51 }
52
53 fn add_inplace<A>(&mut self, module: &Module<BACKEND>, a: &A)
54 where
55 A: VecZnxToRef + Infos,
56 {
57 #[cfg(debug_assertions)]
58 {
59 assert_eq!(a.n(), module.n());
60 assert_eq!(self.n(), module.n());
61 assert_eq!(self.basek(), a.basek());
62 assert!(self.rank() >= a.rank())
63 }
64
65 (0..a.rank() + 1).for_each(|i| {
66 module.vec_znx_add_inplace(self, i, a, i);
67 });
68
69 self.set_k(a.k().max(self.k()));
70 }
71
72 fn sub<A, B>(&mut self, module: &Module<BACKEND>, a: &A, b: &B)
73 where
74 A: VecZnxToRef + Infos,
75 B: VecZnxToRef + Infos,
76 {
77 #[cfg(debug_assertions)]
78 {
79 assert_eq!(a.n(), module.n());
80 assert_eq!(b.n(), module.n());
81 assert_eq!(self.n(), module.n());
82 assert_eq!(a.basek(), b.basek());
83 assert!(self.rank() >= a.rank().max(b.rank()));
84 }
85
86 let min_col: usize = a.rank().min(b.rank()) + 1;
87 let max_col: usize = a.rank().max(b.rank() + 1);
88 let self_col: usize = self.rank() + 1;
89
90 (0..min_col).for_each(|i| {
91 module.vec_znx_sub(self, i, a, i, b, i);
92 });
93
94 if a.rank() > b.rank() {
95 (min_col..max_col).for_each(|i| {
96 module.vec_znx_copy(self, i, a, i);
97 });
98 } else {
99 (min_col..max_col).for_each(|i| {
100 module.vec_znx_copy(self, i, b, i);
101 module.vec_znx_negate_inplace(self, i);
102 });
103 }
104
105 let size: usize = self.size();
106 let mut self_mut: VecZnx<&mut [u8]> = self.to_mut();
107 (max_col..self_col).for_each(|i| {
108 (0..size).for_each(|j| {
109 self_mut.zero_at(i, j);
110 });
111 });
112
113 self.set_basek(a.basek());
114 self.set_k(a.k().max(b.k()));
115 }
116
117 fn sub_inplace_ab<A>(&mut self, module: &Module<BACKEND>, a: &A)
118 where
119 A: VecZnxToRef + Infos,
120 {
121 #[cfg(debug_assertions)]
122 {
123 assert_eq!(a.n(), module.n());
124 assert_eq!(self.n(), module.n());
125 assert_eq!(self.basek(), a.basek());
126 assert!(self.rank() >= a.rank())
127 }
128
129 (0..a.rank() + 1).for_each(|i| {
130 module.vec_znx_sub_ab_inplace(self, i, a, i);
131 });
132
133 self.set_k(a.k().max(self.k()));
134 }
135
136 fn sub_inplace_ba<A>(&mut self, module: &Module<BACKEND>, a: &A)
137 where
138 A: VecZnxToRef + Infos,
139 {
140 #[cfg(debug_assertions)]
141 {
142 assert_eq!(a.n(), module.n());
143 assert_eq!(self.n(), module.n());
144 assert_eq!(self.basek(), a.basek());
145 assert!(self.rank() >= a.rank())
146 }
147
148 (0..a.rank() + 1).for_each(|i| {
149 module.vec_znx_sub_ba_inplace(self, i, a, i);
150 });
151
152 self.set_k(a.k().max(self.k()));
153 }
154
155 fn rotate<A>(&mut self, module: &Module<BACKEND>, k: i64, a: &A)
156 where
157 A: VecZnxToRef + Infos,
158 {
159 #[cfg(debug_assertions)]
160 {
161 assert_eq!(a.n(), module.n());
162 assert_eq!(self.n(), module.n());
163 assert_eq!(self.basek(), a.basek());
164 assert_eq!(self.rank(), a.rank())
165 }
166
167 (0..a.rank() + 1).for_each(|i| {
168 module.vec_znx_rotate(k, self, i, a, i);
169 });
170
171 self.set_k(a.k());
172 }
173
174 fn rotate_inplace<A>(&mut self, module: &Module<BACKEND>, k: i64)
175 where
176 A: VecZnxToRef + Infos,
177 {
178 #[cfg(debug_assertions)]
179 {
180 assert_eq!(self.n(), module.n());
181 }
182
183 (0..self.rank() + 1).for_each(|i| {
184 module.vec_znx_rotate_inplace(k, self, i);
185 });
186 }
187
188 fn copy<A>(&mut self, module: &Module<BACKEND>, a: &A)
189 where
190 A: VecZnxToRef + Infos,
191 {
192 #[cfg(debug_assertions)]
193 {
194 assert_eq!(self.n(), module.n());
195 assert_eq!(a.n(), module.n());
196 }
197
198 let cols: usize = self.rank().min(a.rank()) + 1;
199
200 (0..cols).for_each(|i| {
201 module.vec_znx_copy(self, i, a, i);
202 });
203
204 self.set_k(a.k());
205 self.set_basek(a.basek());
206 }
207
208 fn rsh(&mut self, k: usize, scratch: &mut Scratch) {
209 let basek: usize = self.basek();
210 let mut self_mut: VecZnx<&mut [u8]> = self.to_mut();
211 self_mut.rsh(basek, k, scratch);
212 }
213}