poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/q120/q120_ntt_avx2.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/q120/q120_ntt_avx2.cRaw
1#include <assert.h>
2#include <immintrin.h>
3#include <inttypes.h>
4#include <math.h>
5#include <stdint.h>
6#include <stdio.h>
7
8#include "q120_common.h"
9#include "q120_ntt_private.h"
10
11// at which level to switch from computations by level to computations by block
12#define CHANGE_MODE_N 1024
13
14__always_inline __m256i split_precompmul_si256(__m256i inp, __m256i powomega, const uint64_t h, const __m256i mask) {
15 const __m256i inp_low = _mm256_and_si256(inp, mask);
16 const __m256i t1 = _mm256_mul_epu32(inp_low, powomega);
17
18 const __m256i inp_high = _mm256_srli_epi64(inp, h);
19 const __m256i powomega_high = _mm256_srli_epi64(powomega, 32);
20 const __m256i t2 = _mm256_mul_epu32(inp_high, powomega_high);
21
22 return _mm256_add_epi64(t1, t2);
23}
24
25__always_inline __m256i modq_red(const __m256i x, const uint64_t h, const __m256i mask, const __m256i _2_pow_h_modq) {
26 __m256i xh = _mm256_srli_epi64(x, h);
27 __m256i xl = _mm256_and_si256(x, mask);
28 __m256i xh_1 = _mm256_mul_epu32(xh, _2_pow_h_modq);
29 return _mm256_add_epi64(xl, xh_1);
30}
31
32void print_data(const uint64_t n, const uint64_t* const data, const uint64_t q) {
33 for (uint64_t i = 0; i < n; i++) {
34 printf("%" PRIu64 " ", *(data + i) % q);
35 }
36 printf("\n");
37}
38
39double max_bit_size(const void* const begin, const void* const end) {
40 double bs = 0;
41 const uint64_t* data = (uint64_t*)begin;
42 while (data != end) {
43 double t = log2(*(data++));
44 if (bs < t) {
45 bs = t;
46 }
47 }
48 return bs;
49}
50
51void ntt_iter_first(__m256i* const begin, const __m256i* const end, const q120_ntt_step_precomp* const itData,
52 const __m256i* powomega) {
53 const uint64_t h = itData->half_bs;
54 const __m256i vmask = _mm256_set1_epi64x(itData->mask);
55
56 __m256i* data = begin;
57 while (data < end) {
58 __m256i x = _mm256_loadu_si256(data);
59 __m256i po = _mm256_loadu_si256(powomega);
60 __m256i r = split_precompmul_si256(x, po, h, vmask);
61 _mm256_storeu_si256(data, r);
62
63 data++;
64 powomega++;
65 }
66}
67
68void ntt_iter(const uint64_t nn, __m256i* const begin, const __m256i* const end,
69 const q120_ntt_step_precomp* const itData, const __m256i* const powomega) {
70 assert(nn % 2 == 0);
71 const uint64_t halfnn = nn / 2;
72
73 const __m256i vq2bs = _mm256_loadu_si256((__m256i*)itData->q2bs);
74 const __m256i vmask = _mm256_set1_epi64x(itData->mask);
75
76 __m256i* data = begin;
77 while (data < end) {
78 __m256i* ptr1 = data;
79 __m256i* ptr2 = data + halfnn;
80
81 const __m256i a = _mm256_loadu_si256(ptr1);
82 const __m256i b = _mm256_loadu_si256(ptr2);
83
84 const __m256i ap = _mm256_add_epi64(a, b);
85 _mm256_storeu_si256(ptr1, ap);
86
87 const __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), b);
88 _mm256_storeu_si256(ptr2, bp);
89
90 ptr1++;
91 ptr2++;
92
93 const __m256i* po_ptr = powomega;
94 for (uint64_t i = 1; i < halfnn; ++i) {
95 __m256i a = _mm256_loadu_si256(ptr1);
96 __m256i b = _mm256_loadu_si256(ptr2);
97
98 __m256i ap = _mm256_add_epi64(a, b);
99
100 _mm256_storeu_si256(ptr1, ap);
101
102 __m256i b1 = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), b);
103 __m256i po = _mm256_loadu_si256(po_ptr);
104
105 __m256i bp = split_precompmul_si256(b1, po, itData->half_bs, vmask);
106
107 _mm256_storeu_si256(ptr2, bp);
108
109 ptr1++;
110 ptr2++;
111 po_ptr++;
112 }
113 data += nn;
114 }
115}
116
117void ntt_iter_red(const uint64_t nn, __m256i* const begin, const __m256i* const end,
118 const q120_ntt_step_precomp* const itData, const __m256i* const powomega,
119 const q120_ntt_reduc_step_precomp* const reduc_precomp) {
120 assert(nn % 2 == 0);
121 const uint64_t halfnn = nn / 2;
122
123 const __m256i vq2bs = _mm256_loadu_si256((__m256i*)itData->q2bs);
124 const __m256i vmask = _mm256_set1_epi64x(itData->mask);
125
126 const __m256i reduc_mask = _mm256_set1_epi64x(reduc_precomp->mask);
127 const __m256i reduc_cst = _mm256_loadu_si256((__m256i*)reduc_precomp->modulo_red_cst);
128
129 __m256i* data = begin;
130 while (data < end) {
131 __m256i* ptr1 = data;
132 __m256i* ptr2 = data + halfnn;
133
134 __m256i a = _mm256_loadu_si256(ptr1);
135 __m256i b = _mm256_loadu_si256(ptr2);
136
137 a = modq_red(a, reduc_precomp->h, reduc_mask, reduc_cst);
138 b = modq_red(b, reduc_precomp->h, reduc_mask, reduc_cst);
139
140 const __m256i ap = _mm256_add_epi64(a, b);
141 _mm256_storeu_si256(ptr1, ap);
142
143 const __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), b);
144 _mm256_storeu_si256(ptr2, bp);
145
146 ptr1++;
147 ptr2++;
148
149 const __m256i* po_ptr = powomega;
150 for (uint64_t i = 1; i < halfnn; ++i) {
151 __m256i a = _mm256_loadu_si256(ptr1);
152 __m256i b = _mm256_loadu_si256(ptr2);
153
154 a = modq_red(a, reduc_precomp->h, reduc_mask, reduc_cst);
155 b = modq_red(b, reduc_precomp->h, reduc_mask, reduc_cst);
156
157 __m256i ap = _mm256_add_epi64(a, b);
158
159 _mm256_storeu_si256(ptr1, ap);
160
161 __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), b);
162 __m256i po = _mm256_loadu_si256(po_ptr);
163 bp = split_precompmul_si256(bp, po, itData->half_bs, vmask);
164
165 _mm256_storeu_si256(ptr2, bp);
166
167 ptr1++;
168 ptr2++;
169 po_ptr++;
170 }
171 data += nn;
172 }
173}
174
175EXPORT void q120_ntt_bb_avx2(const q120_ntt_precomp* const precomp, q120b* const data_ptr) {
176 // assert((size_t)data_ptr % 32 == 0); // alignment check
177
178 const uint64_t n = precomp->n;
179 if (n == 1) return;
180
181 const q120_ntt_step_precomp* itData = precomp->level_metadata;
182 const __m256i* powomega = (__m256i*)precomp->powomega;
183
184 __m256i* const begin = (__m256i*)data_ptr;
185 const __m256i* const end = ((__m256i*)data_ptr) + n;
186
187 if (CHECK_BOUNDS) {
188 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
189 LOG("Input %lf %" PRIu64 "\n", bs, precomp->input_bit_size);
190 assert(bs <= precomp->input_bit_size);
191 }
192
193 // first iteration a_k.omega^k
194 ntt_iter_first(begin, end, itData, powomega);
195
196 if (CHECK_BOUNDS) {
197 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
198 LOG("Iter %3" PRIu64 " - %lf %" PRIu64 "\n", n, bs, itData->bs);
199 assert(bs < itData->bs);
200 }
201
202 powomega += n;
203 itData++;
204
205 const uint64_t split_nn = (CHANGE_MODE_N > n) ? n : CHANGE_MODE_N;
206 // const uint64_t split_nn = 2;
207
208 // computations by level
209 for (uint64_t nn = n; nn > split_nn; nn /= 2) {
210 const uint64_t halfnn = nn / 2;
211
212 if (itData->reduce) {
213 ntt_iter_red(nn, begin, end, itData, powomega, &precomp->reduc_metadata);
214 } else {
215 ntt_iter(nn, begin, end, itData, powomega);
216 }
217
218 if (CHECK_BOUNDS) {
219 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
220 LOG("Iter %3" PRIu64 " - %lf %" PRIu64 " %c\n", nn / 2, bs, itData->bs, itData->reduce ? '*' : ' ');
221 assert(bs < itData->bs);
222 }
223
224 powomega += halfnn - 1;
225 itData++;
226 }
227
228 // computations by memory block
229 if (split_nn >= 2) {
230 const q120_ntt_step_precomp* itData1 = itData;
231 const __m256i* powomega1 = powomega;
232 for (__m256i* it = begin; it < end; it += split_nn) {
233 __m256i* const begin1 = it;
234 const __m256i* const end1 = it + split_nn;
235
236 itData = itData1;
237 powomega = powomega1;
238 for (uint64_t nn = split_nn; nn >= 2; nn /= 2) {
239 const uint64_t halfnn = nn / 2;
240
241 if (itData->reduce) {
242 ntt_iter_red(nn, begin1, end1, itData, powomega, &precomp->reduc_metadata);
243 } else {
244 ntt_iter(nn, begin1, end1, itData, powomega);
245 }
246
247 if (CHECK_BOUNDS) {
248 double bs __attribute__((unused)) = max_bit_size((uint64_t*)begin1, (uint64_t*)end1);
249 // LOG("Iter %3lu - %lf %lu\n", nn / 2, bs, itData->bs);
250 assert(bs < itData->bs);
251 }
252
253 powomega += halfnn - 1;
254 itData++;
255 }
256 }
257 }
258
259 if (CHECK_BOUNDS) {
260 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
261 LOG("Iter %3" PRIu64 " - %lf %" PRIu64 "\n", UINT64_C(1), bs, precomp->output_bit_size);
262 assert(bs < precomp->output_bit_size);
263 }
264}
265
266void intt_iter(const uint64_t nn, __m256i* const begin, const __m256i* const end,
267 const q120_ntt_step_precomp* const itData, const __m256i* const powomega) {
268 assert(nn % 2 == 0);
269 const uint64_t halfnn = nn / 2;
270
271 const __m256i vq2bs = _mm256_loadu_si256((__m256i*)itData->q2bs);
272 const __m256i vmask = _mm256_set1_epi64x(itData->mask);
273
274 __m256i* data = begin;
275 while (data < end) {
276 __m256i* ptr1 = data;
277 __m256i* ptr2 = data + halfnn;
278
279 const __m256i a = _mm256_loadu_si256(ptr1);
280 const __m256i b = _mm256_loadu_si256(ptr2);
281
282 const __m256i ap = _mm256_add_epi64(a, b);
283 _mm256_storeu_si256(ptr1, ap);
284
285 const __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), b);
286 _mm256_storeu_si256(ptr2, bp);
287
288 ptr1++;
289 ptr2++;
290
291 const __m256i* po_ptr = powomega;
292 for (uint64_t i = 1; i < halfnn; ++i) {
293 __m256i a = _mm256_loadu_si256(ptr1);
294 __m256i b = _mm256_loadu_si256(ptr2);
295
296 __m256i po = _mm256_loadu_si256(po_ptr);
297 __m256i bo = split_precompmul_si256(b, po, itData->half_bs, vmask);
298
299 __m256i ap = _mm256_add_epi64(a, bo);
300
301 _mm256_storeu_si256(ptr1, ap);
302
303 __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), bo);
304
305 _mm256_storeu_si256(ptr2, bp);
306
307 ptr1++;
308 ptr2++;
309 po_ptr++;
310 }
311 data += nn;
312 }
313}
314
315void intt_iter_red(const uint64_t nn, __m256i* const begin, const __m256i* const end,
316 const q120_ntt_step_precomp* const itData, const __m256i* const powomega,
317 const q120_ntt_reduc_step_precomp* const reduc_precomp) {
318 assert(nn % 2 == 0);
319 const uint64_t halfnn = nn / 2;
320
321 const __m256i vq2bs = _mm256_loadu_si256((__m256i*)itData->q2bs);
322 const __m256i vmask = _mm256_set1_epi64x(itData->mask);
323
324 const __m256i reduc_mask = _mm256_set1_epi64x(reduc_precomp->mask);
325 const __m256i reduc_cst = _mm256_loadu_si256((__m256i*)reduc_precomp->modulo_red_cst);
326
327 __m256i* data = begin;
328 while (data < end) {
329 __m256i* ptr1 = data;
330 __m256i* ptr2 = data + halfnn;
331
332 __m256i a = _mm256_loadu_si256(ptr1);
333 __m256i b = _mm256_loadu_si256(ptr2);
334
335 a = modq_red(a, reduc_precomp->h, reduc_mask, reduc_cst);
336 b = modq_red(b, reduc_precomp->h, reduc_mask, reduc_cst);
337
338 const __m256i ap = _mm256_add_epi64(a, b);
339 _mm256_storeu_si256(ptr1, ap);
340
341 const __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), b);
342 _mm256_storeu_si256(ptr2, bp);
343
344 ptr1++;
345 ptr2++;
346
347 const __m256i* po_ptr = powomega;
348 for (uint64_t i = 1; i < halfnn; ++i) {
349 __m256i a = _mm256_loadu_si256(ptr1);
350 __m256i b = _mm256_loadu_si256(ptr2);
351
352 a = modq_red(a, reduc_precomp->h, reduc_mask, reduc_cst);
353 b = modq_red(b, reduc_precomp->h, reduc_mask, reduc_cst);
354
355 __m256i po = _mm256_loadu_si256(po_ptr);
356 __m256i bo = split_precompmul_si256(b, po, itData->half_bs, vmask);
357
358 __m256i ap = _mm256_add_epi64(a, bo);
359
360 _mm256_storeu_si256(ptr1, ap);
361
362 __m256i bp = _mm256_sub_epi64(_mm256_add_epi64(a, vq2bs), bo);
363
364 _mm256_storeu_si256(ptr2, bp);
365
366 ptr1++;
367 ptr2++;
368 po_ptr++;
369 }
370 data += nn;
371 }
372}
373
374void ntt_iter_first_red(__m256i* const begin, const __m256i* const end, const q120_ntt_step_precomp* const itData,
375 const __m256i* powomega, const q120_ntt_reduc_step_precomp* const reduc_precomp) {
376 const uint64_t h = itData->half_bs;
377 const __m256i vmask = _mm256_set1_epi64x(itData->mask);
378
379 const __m256i reduc_mask = _mm256_set1_epi64x(reduc_precomp->mask);
380 const __m256i reduc_cst = _mm256_loadu_si256((__m256i*)reduc_precomp->modulo_red_cst);
381
382 __m256i* data = begin;
383 while (data < end) {
384 __m256i x = _mm256_loadu_si256(data);
385 x = modq_red(x, reduc_precomp->h, reduc_mask, reduc_cst);
386 __m256i po = _mm256_loadu_si256(powomega);
387 __m256i r = split_precompmul_si256(x, po, h, vmask);
388 _mm256_storeu_si256(data, r);
389
390 data++;
391 powomega++;
392 }
393}
394
395EXPORT void q120_intt_bb_avx2(const q120_ntt_precomp* const precomp, q120b* const data_ptr) {
396 // assert((size_t)data_ptr % 32 == 0); // alignment check
397
398 const uint64_t n = precomp->n;
399 if (n == 1) return;
400
401 const q120_ntt_step_precomp* itData = precomp->level_metadata;
402 const __m256i* powomega = (__m256i*)precomp->powomega;
403
404 __m256i* const begin = (__m256i*)data_ptr;
405 const __m256i* const end = ((__m256i*)data_ptr) + n;
406
407 if (CHECK_BOUNDS) {
408 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
409 LOG("Input %lf %" PRIu64 "\n", bs, precomp->input_bit_size);
410 assert(bs <= precomp->input_bit_size);
411 }
412
413 const uint64_t split_nn = (CHANGE_MODE_N > n) ? n : CHANGE_MODE_N;
414
415 // computations by memory block
416 if (split_nn >= 2) {
417 const q120_ntt_step_precomp* itData1 = itData;
418 const __m256i* powomega1 = powomega;
419 for (__m256i* it = begin; it < end; it += split_nn) {
420 __m256i* const begin1 = it;
421 const __m256i* const end1 = it + split_nn;
422
423 itData = itData1;
424 powomega = powomega1;
425 for (uint64_t nn = 2; nn <= split_nn; nn *= 2) {
426 const uint64_t halfnn = nn / 2;
427
428 if (itData->reduce) {
429 intt_iter_red(nn, begin1, end1, itData, powomega, &precomp->reduc_metadata);
430 } else {
431 intt_iter(nn, begin1, end1, itData, powomega);
432 }
433
434 if (CHECK_BOUNDS) {
435 double bs __attribute__((unused)) = max_bit_size((uint64_t*)begin1, (uint64_t*)end1);
436 // LOG("Iter %3lu - %lf %lu\n", nn / 2, bs, itData->bs);
437 assert(bs < itData->bs);
438 }
439
440 powomega += halfnn - 1;
441 itData++;
442 }
443 }
444 }
445
446 // computations by level
447 // for (uint64_t nn = 2; nn <= n; nn *= 2) {
448 for (uint64_t nn = 2 * split_nn; nn <= n; nn *= 2) {
449 const uint64_t halfnn = nn / 2;
450
451 if (itData->reduce) {
452 intt_iter_red(nn, begin, end, itData, powomega, &precomp->reduc_metadata);
453 } else {
454 intt_iter(nn, begin, end, itData, powomega);
455 }
456
457 if (CHECK_BOUNDS) {
458 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
459 LOG("Iter %3" PRIu64 " - %lf %" PRIu64 " %c\n", nn / 2, bs, itData->bs, itData->reduce ? '*' : ' ');
460 assert(bs < itData->bs);
461 }
462
463 powomega += halfnn - 1;
464 itData++;
465 }
466
467 // last iteration a_k . omega^k . n^-1
468 if (itData->reduce) {
469 ntt_iter_first_red(begin, end, itData, powomega, &precomp->reduc_metadata);
470 } else {
471 ntt_iter_first(begin, end, itData, powomega);
472 }
473
474 if (CHECK_BOUNDS) {
475 double bs __attribute__((unused)) = max_bit_size((void*)begin, (void*)end);
476 LOG("Iter %3" PRIu64 " - %lf %" PRIu64 "\n", n, bs, itData->bs);
477 assert(bs < itData->bs);
478 }
479}