poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_znx.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/arithmetic/vec_znx.cRaw
1#include <assert.h>
2#include <math.h>
3#include <stdint.h>
4#include <stdlib.h>
5#include <string.h>
6
7#include "../coeffs/coeffs_arithmetic.h"
8#include "../q120/q120_arithmetic.h"
9#include "../q120/q120_ntt.h"
10#include "../reim/reim_fft_internal.h"
11#include "../reim4/reim4_arithmetic.h"
12#include "vec_znx_arithmetic.h"
13#include "vec_znx_arithmetic_private.h"
14
15// general function (virtual dispatch)
16
17EXPORT void vec_znx_add(const MODULE* module, // N
18 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
19 const int64_t* a, uint64_t a_size, uint64_t a_sl, // a
20 const int64_t* b, uint64_t b_size, uint64_t b_sl // b
21) {
22 module->func.vec_znx_add(module, // N
23 res, res_size, res_sl, // res
24 a, a_size, a_sl, // a
25 b, b_size, b_sl // b
26 );
27}
28
29EXPORT void vec_znx_sub(const MODULE* module, // N
30 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
31 const int64_t* a, uint64_t a_size, uint64_t a_sl, // a
32 const int64_t* b, uint64_t b_size, uint64_t b_sl // b
33) {
34 module->func.vec_znx_sub(module, // N
35 res, res_size, res_sl, // res
36 a, a_size, a_sl, // a
37 b, b_size, b_sl // b
38 );
39}
40
41EXPORT void vec_znx_rotate(const MODULE* module, // N
42 const int64_t p, // rotation value
43 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
44 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
45) {
46 module->func.vec_znx_rotate(module, // N
47 p, // p
48 res, res_size, res_sl, // res
49 a, a_size, a_sl // a
50 );
51}
52
53EXPORT void vec_znx_mul_xp_minus_one(const MODULE* module, // N
54 const int64_t p, // p
55 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
56 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
57) {
58 module->func.vec_znx_mul_xp_minus_one(module, // N
59 p, // p
60 res, res_size, res_sl, // res
61 a, a_size, a_sl // a
62 );
63}
64
65EXPORT void vec_znx_automorphism(const MODULE* module, // N
66 const int64_t p, // X->X^p
67 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
68 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
69) {
70 module->func.vec_znx_automorphism(module, // N
71 p, // p
72 res, res_size, res_sl, // res
73 a, a_size, a_sl // a
74 );
75}
76
77EXPORT void vec_znx_normalize_base2k(const MODULE* module, uint64_t nn, // N
78 uint64_t log2_base2k, // output base 2^K
79 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
80 const int64_t* a, uint64_t a_size, uint64_t a_sl, // a
81 uint8_t* tmp_space // scratch space of size >= N
82) {
83 module->func.vec_znx_normalize_base2k(module, nn, // N
84 log2_base2k, // log2_base2k
85 res, res_size, res_sl, // res
86 a, a_size, a_sl, // a
87 tmp_space);
88}
89
90EXPORT uint64_t vec_znx_normalize_base2k_tmp_bytes(const MODULE* module, uint64_t nn // N
91) {
92 return module->func.vec_znx_normalize_base2k_tmp_bytes(module, nn // N
93 );
94}
95
96// specialized function (ref)
97
98EXPORT void vec_znx_add_ref(const MODULE* module, // N
99 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
100 const int64_t* a, uint64_t a_size, uint64_t a_sl, // a
101 const int64_t* b, uint64_t b_size, uint64_t b_sl // b
102) {
103 const uint64_t nn = module->nn;
104 if (a_size <= b_size) {
105 const uint64_t sum_idx = res_size < a_size ? res_size : a_size;
106 const uint64_t copy_idx = res_size < b_size ? res_size : b_size;
107 // add up to the smallest dimension
108 for (uint64_t i = 0; i < sum_idx; ++i) {
109 znx_add_i64_ref(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
110 }
111 // then copy to the largest dimension
112 for (uint64_t i = sum_idx; i < copy_idx; ++i) {
113 znx_copy_i64_ref(nn, res + i * res_sl, b + i * b_sl);
114 }
115 // then extend with zeros
116 for (uint64_t i = copy_idx; i < res_size; ++i) {
117 znx_zero_i64_ref(nn, res + i * res_sl);
118 }
119 } else {
120 const uint64_t sum_idx = res_size < b_size ? res_size : b_size;
121 const uint64_t copy_idx = res_size < a_size ? res_size : a_size;
122 // add up to the smallest dimension
123 for (uint64_t i = 0; i < sum_idx; ++i) {
124 znx_add_i64_ref(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
125 }
126 // then copy to the largest dimension
127 for (uint64_t i = sum_idx; i < copy_idx; ++i) {
128 znx_copy_i64_ref(nn, res + i * res_sl, a + i * a_sl);
129 }
130 // then extend with zeros
131 for (uint64_t i = copy_idx; i < res_size; ++i) {
132 znx_zero_i64_ref(nn, res + i * res_sl);
133 }
134 }
135}
136
137EXPORT void vec_znx_sub_ref(const MODULE* module, // N
138 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
139 const int64_t* a, uint64_t a_size, uint64_t a_sl, // a
140 const int64_t* b, uint64_t b_size, uint64_t b_sl // b
141) {
142 const uint64_t nn = module->nn;
143 if (a_size <= b_size) {
144 const uint64_t sub_idx = res_size < a_size ? res_size : a_size;
145 const uint64_t copy_idx = res_size < b_size ? res_size : b_size;
146 // subtract up to the smallest dimension
147 for (uint64_t i = 0; i < sub_idx; ++i) {
148 znx_sub_i64_ref(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
149 }
150 // then negate to the largest dimension
151 for (uint64_t i = sub_idx; i < copy_idx; ++i) {
152 znx_negate_i64_ref(nn, res + i * res_sl, b + i * b_sl);
153 }
154 // then extend with zeros
155 for (uint64_t i = copy_idx; i < res_size; ++i) {
156 znx_zero_i64_ref(nn, res + i * res_sl);
157 }
158 } else {
159 const uint64_t sub_idx = res_size < b_size ? res_size : b_size;
160 const uint64_t copy_idx = res_size < a_size ? res_size : a_size;
161 // subtract up to the smallest dimension
162 for (uint64_t i = 0; i < sub_idx; ++i) {
163 znx_sub_i64_ref(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
164 }
165 // then copy to the largest dimension
166 for (uint64_t i = sub_idx; i < copy_idx; ++i) {
167 znx_copy_i64_ref(nn, res + i * res_sl, a + i * a_sl);
168 }
169 // then extend with zeros
170 for (uint64_t i = copy_idx; i < res_size; ++i) {
171 znx_zero_i64_ref(nn, res + i * res_sl);
172 }
173 }
174}
175
176EXPORT void vec_znx_rotate_ref(const MODULE* module, // N
177 const int64_t p, // rotation value
178 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
179 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
180) {
181 const uint64_t nn = module->nn;
182
183 const uint64_t rot_end_idx = res_size < a_size ? res_size : a_size;
184 // rotate up to the smallest dimension
185 for (uint64_t i = 0; i < rot_end_idx; ++i) {
186 int64_t* res_ptr = res + i * res_sl;
187 const int64_t* a_ptr = a + i * a_sl;
188 if (res_ptr == a_ptr) {
189 znx_rotate_inplace_i64(nn, p, res_ptr);
190 } else {
191 znx_rotate_i64(nn, p, res_ptr, a_ptr);
192 }
193 }
194 // then extend with zeros
195 for (uint64_t i = rot_end_idx; i < res_size; ++i) {
196 znx_zero_i64_ref(nn, res + i * res_sl);
197 }
198}
199
200EXPORT void vec_znx_mul_xp_minus_one_ref(const MODULE* module, // N
201 const int64_t p, // p
202 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
203 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
204) {
205 const uint64_t nn = module->nn;
206
207 const uint64_t rot_end_idx = res_size < a_size ? res_size : a_size;
208 for (uint64_t i = 0; i < rot_end_idx; ++i) {
209 int64_t* res_ptr = res + i * res_sl;
210 const int64_t* a_ptr = a + i * a_sl;
211 if (res_ptr == a_ptr) {
212 znx_mul_xp_minus_one_inplace_i64(nn, p, res_ptr);
213 } else {
214 znx_mul_xp_minus_one_i64(nn, p, res_ptr, a_ptr);
215 }
216 }
217 // then extend with zeros
218 for (uint64_t i = rot_end_idx; i < res_size; ++i) {
219 znx_zero_i64_ref(nn, res + i * res_sl);
220 }
221}
222
223EXPORT void vec_znx_automorphism_ref(const MODULE* module, // N
224 const int64_t p, // X->X^p
225 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
226 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
227) {
228 const uint64_t nn = module->nn;
229
230 const uint64_t auto_end_idx = res_size < a_size ? res_size : a_size;
231
232 for (uint64_t i = 0; i < auto_end_idx; ++i) {
233 int64_t* res_ptr = res + i * res_sl;
234 const int64_t* a_ptr = a + i * a_sl;
235 if (res_ptr == a_ptr) {
236 znx_automorphism_inplace_i64(nn, p, res_ptr);
237 } else {
238 znx_automorphism_i64(nn, p, res_ptr, a_ptr);
239 }
240 }
241 // then extend with zeros
242 for (uint64_t i = auto_end_idx; i < res_size; ++i) {
243 znx_zero_i64_ref(nn, res + i * res_sl);
244 }
245}
246
247EXPORT void vec_znx_normalize_base2k_ref(const MODULE* module, uint64_t nn, // N
248 uint64_t log2_base2k, // output base 2^K
249 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
250 const int64_t* a, uint64_t a_size, uint64_t a_sl, // a
251 uint8_t* tmp_space // scratch space of size >= N
252) {
253
254 // use MSB limb of res for carry propagation
255 int64_t* cout = (int64_t*)tmp_space;
256 int64_t* cin = 0x0;
257
258 // propagate carry until first limb of res
259 int64_t i = a_size - 1;
260 for (; i >= res_size; --i) {
261 znx_normalize(nn, log2_base2k, 0x0, cout, a + i * a_sl, cin);
262 cin = cout;
263 }
264
265 // propagate carry and normalize
266 for (; i >= 1; --i) {
267 znx_normalize(nn, log2_base2k, res + i * res_sl, cout, a + i * a_sl, cin);
268 cin = cout;
269 }
270
271 // normalize last limb
272 znx_normalize(nn, log2_base2k, res, 0x0, a, cin);
273
274 // extend result with zeros
275 for (uint64_t i = a_size; i < res_size; ++i) {
276 znx_zero_i64_ref(nn, res + i * res_sl);
277 }
278}
279
280EXPORT uint64_t vec_znx_normalize_base2k_tmp_bytes_ref(const MODULE* module, uint64_t nn // N
281) {
282 return nn * sizeof(int64_t);
283}
284
285// alias have to be defined in this unit: do not move
286#ifdef __APPLE__
287EXPORT uint64_t fft64_vec_znx_big_range_normalize_base2k_tmp_bytes( //
288 const MODULE* module, // N
289 uint64_t nn
290) {
291 return vec_znx_normalize_base2k_tmp_bytes_ref(module, nn);
292}
293EXPORT uint64_t fft64_vec_znx_big_normalize_base2k_tmp_bytes( //
294 const MODULE* module, // N
295 uint64_t nn
296) {
297 return vec_znx_normalize_base2k_tmp_bytes_ref(module, nn);
298}
299#else
300EXPORT uint64_t fft64_vec_znx_big_normalize_base2k_tmp_bytes( //
301 const MODULE* module, // N
302 uint64_t nn
303 ) __attribute((alias("vec_znx_normalize_base2k_tmp_bytes_ref")));
304
305EXPORT uint64_t fft64_vec_znx_big_range_normalize_base2k_tmp_bytes( //
306 const MODULE* module, // N
307 uint64_t nn
308 ) __attribute((alias("vec_znx_normalize_base2k_tmp_bytes_ref")));
309#endif
310
311/** @brief sets res = 0 */
312EXPORT void vec_znx_zero(const MODULE* module, // N
313 int64_t* res, uint64_t res_size, uint64_t res_sl // res
314) {
315 module->func.vec_znx_zero(module, res, res_size, res_sl);
316}
317
318/** @brief sets res = a */
319EXPORT void vec_znx_copy(const MODULE* module, // N
320 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
321 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
322) {
323 module->func.vec_znx_copy(module, res, res_size, res_sl, a, a_size, a_sl);
324}
325
326/** @brief sets res = a */
327EXPORT void vec_znx_negate(const MODULE* module, // N
328 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
329 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
330) {
331 module->func.vec_znx_negate(module, res, res_size, res_sl, a, a_size, a_sl);
332}
333
334EXPORT void vec_znx_zero_ref(const MODULE* module, // N
335 int64_t* res, uint64_t res_size, uint64_t res_sl // res
336) {
337 uint64_t nn = module->nn;
338 for (uint64_t i = 0; i < res_size; ++i) {
339 znx_zero_i64_ref(nn, res + i * res_sl);
340 }
341}
342
343EXPORT void vec_znx_copy_ref(const MODULE* module, // N
344 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
345 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
346) {
347 uint64_t nn = module->nn;
348 uint64_t smin = res_size < a_size ? res_size : a_size;
349 for (uint64_t i = 0; i < smin; ++i) {
350 znx_copy_i64_ref(nn, res + i * res_sl, a + i * a_sl);
351 }
352 for (uint64_t i = smin; i < res_size; ++i) {
353 znx_zero_i64_ref(nn, res + i * res_sl);
354 }
355}
356
357EXPORT void vec_znx_negate_ref(const MODULE* module, // N
358 int64_t* res, uint64_t res_size, uint64_t res_sl, // res
359 const int64_t* a, uint64_t a_size, uint64_t a_sl // a
360) {
361 uint64_t nn = module->nn;
362 uint64_t smin = res_size < a_size ? res_size : a_size;
363 for (uint64_t i = 0; i < smin; ++i) {
364 znx_negate_i64_ref(nn, res + i * res_sl, a + i * a_sl);
365 }
366 for (uint64_t i = smin; i < res_size; ++i) {
367 znx_zero_i64_ref(nn, res + i * res_sl);
368 }
369}