poulpy/poulpy-hal/src/layouts/encoding.rs
poulpy-hal/src/layouts/encoding.rsRaw
1use itertools::izip;
2use rug::{Assign, Float};
3
4use crate::{
5 layouts::{DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef, Zn, ZnToMut, ZnToRef, ZnxInfos, ZnxView, ZnxViewMut},
6 reference::znx::znx_zero_ref,
7};
8
9impl<D: DataMut> VecZnx<D> {
10 pub fn encode_vec_i64(&mut self, basek: usize, col: usize, k: usize, data: &[i64], log_max: usize) {
11 let size: usize = k.div_ceil(basek);
12
13 #[cfg(debug_assertions)]
14 {
15 let a: VecZnx<&mut [u8]> = self.to_mut();
16 assert!(
17 size <= a.size(),
18 "invalid argument k.div_ceil(basek)={} > a.size()={}",
19 size,
20 a.size()
21 );
22 assert!(col < a.cols());
23 assert!(data.len() <= a.n())
24 }
25
26 let data_len: usize = data.len();
27 let mut a: VecZnx<&mut [u8]> = self.to_mut();
28 let k_rem: usize = basek - (k % basek);
29
30 // Zeroes coefficients of the i-th column
31 (0..a.size()).for_each(|i| {
32 znx_zero_ref(a.at_mut(col, i));
33 });
34
35 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
36 // values on the last limb.
37 // Else we decompose values base2k.
38 if log_max + k_rem < 63 || k_rem == basek {
39 a.at_mut(col, size - 1)[..data_len].copy_from_slice(&data[..data_len]);
40 } else {
41 let mask: i64 = (1 << basek) - 1;
42 let steps: usize = size.min(log_max.div_ceil(basek));
43 (size - steps..size)
44 .rev()
45 .enumerate()
46 .for_each(|(i, i_rev)| {
47 let shift: usize = i * basek;
48 izip!(a.at_mut(col, i_rev).iter_mut(), data.iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
49 })
50 }
51
52 // Case where self.prec % self.k != 0.
53 if k_rem != basek {
54 let steps: usize = size.min(log_max.div_ceil(basek));
55 (size - steps..size).rev().for_each(|i| {
56 a.at_mut(col, i)[..data_len]
57 .iter_mut()
58 .for_each(|x| *x <<= k_rem);
59 })
60 }
61 }
62
63 pub fn encode_coeff_i64(&mut self, basek: usize, col: usize, k: usize, idx: usize, data: i64, log_max: usize) {
64 let size: usize = k.div_ceil(basek);
65
66 #[cfg(debug_assertions)]
67 {
68 let a: VecZnx<&mut [u8]> = self.to_mut();
69 assert!(idx < a.n());
70 assert!(
71 size <= a.size(),
72 "invalid argument k.div_ceil(basek)={} > a.size()={}",
73 size,
74 a.size()
75 );
76 assert!(col < a.cols());
77 }
78
79 let k_rem: usize = basek - (k % basek);
80 let mut a: VecZnx<&mut [u8]> = self.to_mut();
81 (0..a.size()).for_each(|j| a.at_mut(col, j)[idx] = 0);
82
83 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
84 // values on the last limb.
85 // Else we decompose values base2k.
86 if log_max + k_rem < 63 || k_rem == basek {
87 a.at_mut(col, size - 1)[idx] = data;
88 } else {
89 let mask: i64 = (1 << basek) - 1;
90 let steps: usize = size.min(log_max.div_ceil(basek));
91 (size - steps..size)
92 .rev()
93 .enumerate()
94 .for_each(|(j, j_rev)| {
95 a.at_mut(col, j_rev)[idx] = (data >> (j * basek)) & mask;
96 })
97 }
98
99 // Case where prec % k != 0.
100 if k_rem != basek {
101 let steps: usize = size.min(log_max.div_ceil(basek));
102 (size - steps..size).rev().for_each(|j| {
103 a.at_mut(col, j)[idx] <<= k_rem;
104 })
105 }
106 }
107}
108
109impl<D: DataRef> VecZnx<D> {
110 pub fn decode_vec_i64(&self, basek: usize, col: usize, k: usize, data: &mut [i64]) {
111 let size: usize = k.div_ceil(basek);
112 #[cfg(debug_assertions)]
113 {
114 let a: VecZnx<&[u8]> = self.to_ref();
115 assert!(
116 data.len() >= a.n(),
117 "invalid data: data.len()={} < a.n()={}",
118 data.len(),
119 a.n()
120 );
121 assert!(col < a.cols());
122 }
123
124 let a: VecZnx<&[u8]> = self.to_ref();
125 data.copy_from_slice(a.at(col, 0));
126 let rem: usize = basek - (k % basek);
127 if k < basek {
128 data.iter_mut().for_each(|x| *x >>= rem);
129 } else {
130 (1..size).for_each(|i| {
131 if i == size - 1 && rem != basek {
132 let k_rem: usize = basek - rem;
133 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
134 *y = (*y << k_rem) + (x >> rem);
135 });
136 } else {
137 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
138 *y = (*y << basek) + x;
139 });
140 }
141 })
142 }
143 }
144
145 pub fn decode_coeff_i64(&self, basek: usize, col: usize, k: usize, idx: usize) -> i64 {
146 #[cfg(debug_assertions)]
147 {
148 let a: VecZnx<&[u8]> = self.to_ref();
149 assert!(idx < a.n());
150 assert!(col < a.cols())
151 }
152
153 let a: VecZnx<&[u8]> = self.to_ref();
154 let size: usize = k.div_ceil(basek);
155 let mut res: i64 = 0;
156 let rem: usize = basek - (k % basek);
157 (0..size).for_each(|j| {
158 let x: i64 = a.at(col, j)[idx];
159 if j == size - 1 && rem != basek {
160 let k_rem: usize = basek - rem;
161 res = (res << k_rem) + (x >> rem);
162 } else {
163 res = (res << basek) + x;
164 }
165 });
166 res
167 }
168
169 pub fn decode_vec_float(&self, basek: usize, col: usize, data: &mut [Float]) {
170 #[cfg(debug_assertions)]
171 {
172 let a: VecZnx<&[u8]> = self.to_ref();
173 assert!(
174 data.len() >= a.n(),
175 "invalid data: data.len()={} < a.n()={}",
176 data.len(),
177 a.n()
178 );
179 assert!(col < a.cols());
180 }
181
182 let a: VecZnx<&[u8]> = self.to_ref();
183 let size: usize = a.size();
184 let prec: u32 = (basek * size) as u32;
185
186 // 2^{basek}
187 let base = Float::with_val(prec, (1u64 << basek) as f64);
188
189 // y[i] = sum x[j][i] * 2^{-basek*j}
190 (0..size).for_each(|i| {
191 if i == 0 {
192 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
193 y.assign(*x);
194 *y /= &base;
195 });
196 } else {
197 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
198 *y += Float::with_val(prec, *x);
199 *y /= &base;
200 });
201 }
202 });
203 }
204}
205
206impl<D: DataMut> Zn<D> {
207 pub fn encode_i64(&mut self, basek: usize, k: usize, data: i64, log_max: usize) {
208 let size: usize = k.div_ceil(basek);
209
210 #[cfg(debug_assertions)]
211 {
212 let a: Zn<&mut [u8]> = self.to_mut();
213 assert!(
214 size <= a.size(),
215 "invalid argument k.div_ceil(basek)={} > a.size()={}",
216 size,
217 a.size()
218 );
219 }
220
221 let k_rem: usize = basek - (k % basek);
222 let mut a: Zn<&mut [u8]> = self.to_mut();
223 (0..a.size()).for_each(|j| a.at_mut(0, j)[0] = 0);
224
225 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
226 // values on the last limb.
227 // Else we decompose values base2k.
228 if log_max + k_rem < 63 || k_rem == basek {
229 a.at_mut(0, size - 1)[0] = data;
230 } else {
231 let mask: i64 = (1 << basek) - 1;
232 let steps: usize = size.min(log_max.div_ceil(basek));
233 (size - steps..size)
234 .rev()
235 .enumerate()
236 .for_each(|(j, j_rev)| {
237 a.at_mut(0, j_rev)[0] = (data >> (j * basek)) & mask;
238 })
239 }
240
241 // Case where prec % k != 0.
242 if k_rem != basek {
243 let steps: usize = size.min(log_max.div_ceil(basek));
244 (size - steps..size).rev().for_each(|j| {
245 a.at_mut(0, j)[0] <<= k_rem;
246 })
247 }
248 }
249}
250
251impl<D: DataRef> Zn<D> {
252 pub fn decode_i64(&self, basek: usize, k: usize) -> i64 {
253 let a: Zn<&[u8]> = self.to_ref();
254 let size: usize = k.div_ceil(basek);
255 let mut res: i64 = 0;
256 let rem: usize = basek - (k % basek);
257 (0..size).for_each(|j| {
258 let x: i64 = a.at(0, j)[0];
259 if j == size - 1 && rem != basek {
260 let k_rem: usize = basek - rem;
261 res = (res << k_rem) + (x >> rem);
262 } else {
263 res = (res << basek) + x;
264 }
265 });
266 res
267 }
268
269 pub fn decode_float(&self, basek: usize) -> Float {
270 let a: Zn<&[u8]> = self.to_ref();
271 let size: usize = a.size();
272 let prec: u32 = (basek * size) as u32;
273
274 // 2^{basek}
275 let base: Float = Float::with_val(prec, (1 << basek) as f64);
276 let mut res: Float = Float::with_val(prec, (1 << basek) as f64);
277
278 // y[i] = sum x[j][i] * 2^{-basek*j}
279 (0..size).for_each(|i| {
280 if i == 0 {
281 res.assign(a.at(0, size - i - 1)[0]);
282 res /= &base;
283 } else {
284 res += Float::with_val(prec, a.at(0, size - i - 1)[0]);
285 res /= &base;
286 }
287 });
288
289 res
290 }
291}