poulpy/spqlios/lib/spqlios/q120/q120_ntt.c
spqlios/lib/spqlios/q120/q120_ntt.cRaw
1#include <assert.h>
2#include <inttypes.h>
3#include <math.h>
4#include <stdint.h>
5
6#include "q120_ntt_private.h"
7
8q120_ntt_precomp* new_precomp(const uint64_t n) {
9 q120_ntt_precomp* precomp = malloc(sizeof(*precomp));
10 precomp->n = n;
11
12 assert(n && !(n & (n - 1)) && n <= (1 << 16)); // n is a power of 2 smaller than 2^16
13 const uint64_t logN = ceil(log2(n));
14 precomp->level_metadata = malloc((logN + 2) * sizeof(*precomp->level_metadata));
15
16 precomp->powomega = spqlios_alloc_custom_align(32, 4 * 2 * n * sizeof(*(precomp->powomega)));
17
18 return precomp;
19}
20
21uint32_t modq_pow(const uint32_t x, const int64_t n, const uint32_t q) {
22 uint64_t np = (n % (q - 1) + q - 1) % (q - 1);
23
24 uint64_t val_pow = x;
25 uint64_t res = 1;
26 while (np != 0) {
27 if (np & 1) res = (res * val_pow) % q;
28 val_pow = (val_pow * val_pow) % q;
29 np >>= 1;
30 }
31 return res;
32}
33
34void fill_omegas(const uint64_t n, uint32_t omegas[4]) {
35 for (uint64_t k = 0; k < 4; ++k) {
36 omegas[k] = modq_pow(OMEGAS_VEC[k], (1 << 16) / n, PRIMES_VEC[k]);
37 }
38
39#ifndef NDEBUG
40
41 const uint64_t logQ = ceil(log2(PRIMES_VEC[0]));
42 for (int k = 1; k < 4; ++k) {
43 if (logQ != ceil(log2(PRIMES_VEC[k]))) {
44 fprintf(stderr, "The 4 primes must have the same bit-size\n");
45 exit(-1);
46 }
47 }
48
49 // check if each omega is a 2.n primitive root of unity
50 for (uint64_t k = 0; k < 4; ++k) {
51 assert(modq_pow(omegas[k], 2 * n, PRIMES_VEC[k]) == 1);
52 for (uint64_t i = 1; i < 2 * n; ++i) {
53 assert(modq_pow(omegas[k], i, PRIMES_VEC[k]) != 1);
54 }
55 }
56
57 if (logQ > 31) {
58 fprintf(stderr, "Modulus q bit-size is larger than 30 bit\n");
59 exit(-1);
60 }
61#endif
62}
63
64uint64_t fill_reduction_meta(const uint64_t bs_start, q120_ntt_reduc_step_precomp* reduc_metadata) {
65 // fill reduction metadata
66 uint64_t bs_after_reduc = -1;
67 {
68 uint64_t min_h = -1;
69
70 for (uint64_t h = bs_start / 2; h < bs_start; ++h) {
71 uint64_t t = 0;
72 for (uint64_t k = 0; k < 4; ++k) {
73 const uint64_t t1 = bs_start - h + (uint64_t)ceil(log2((UINT64_C(1) << h) % PRIMES_VEC[k]));
74 const uint64_t t2 = UINT64_C(1) + ((t1 > h) ? t1 : h);
75 if (t < t2) t = t2;
76 }
77 if (t < bs_after_reduc) {
78 min_h = h;
79 bs_after_reduc = t;
80 }
81 }
82
83 reduc_metadata->h = min_h;
84 reduc_metadata->mask = (UINT64_C(1) << min_h) - 1;
85 for (uint64_t k = 0; k < 4; ++k) {
86 reduc_metadata->modulo_red_cst[k] = (UINT64_C(1) << min_h) % PRIMES_VEC[k];
87 }
88
89 assert(bs_after_reduc < 64);
90 }
91
92 return bs_after_reduc;
93}
94
95uint64_t round_up_half_n(const uint64_t n) { return (n + 1) / 2; }
96
97EXPORT q120_ntt_precomp* q120_new_ntt_bb_precomp(const uint64_t n) {
98 uint32_t omega_vec[4];
99 fill_omegas(n, omega_vec);
100
101 const uint64_t logQ = ceil(log2(PRIMES_VEC[0]));
102
103 q120_ntt_precomp* precomp = new_precomp(n);
104
105 uint64_t bs = precomp->input_bit_size = 64;
106
107 LOG("NTT parameters:\n");
108 LOG("\tsize = %" PRIu64 "\n", n)
109 LOG("\tlogQ = %" PRIu64 "\n", logQ);
110 LOG("\tinput bit-size = %" PRIu64 "\n", bs);
111
112 if (n == 1) return precomp;
113
114 // fill reduction metadata
115 uint64_t bs_after_reduc = fill_reduction_meta(bs, &(precomp->reduc_metadata));
116
117 // forward metadata
118 q120_ntt_step_precomp* level_metadata_ptr = precomp->level_metadata;
119
120 // first level a_k.omega^k
121 {
122 const uint64_t half_bs = (bs + 1) / 2;
123 level_metadata_ptr->half_bs = half_bs;
124 level_metadata_ptr->mask = (UINT64_C(1) << half_bs) - UINT64_C(1);
125 level_metadata_ptr->bs = bs = half_bs + logQ + 1;
126 LOG("\tlevel %6" PRIu64 " output bit-size = %" PRIu64 " (a_k.omega^k) \n", n, bs);
127 level_metadata_ptr++;
128 }
129
130 for (uint64_t nn = n; nn >= 4; nn /= 2) {
131 level_metadata_ptr->reduce = (bs == 64);
132 if (level_metadata_ptr->reduce) {
133 bs = bs_after_reduc;
134 LOG("\treduce output bit-size = %" PRIu64 "\n", bs);
135 }
136
137 for (int k = 0; k < 4; ++k) level_metadata_ptr->q2bs[k] = (uint64_t)PRIMES_VEC[k] << (bs - logQ);
138
139 double bs_1 = bs + 1; // bit-size of term a+b or a-b
140
141 const uint64_t half_bs = round_up_half_n(bs_1);
142 uint64_t bs_2 = half_bs + logQ + 1; // bit-size of term (a-b).omega^k
143 bs = (bs_1 > bs_2) ? bs_1 : bs_2;
144 assert(bs <= 64);
145
146 level_metadata_ptr->bs = bs;
147 level_metadata_ptr->half_bs = half_bs;
148 level_metadata_ptr->mask = (UINT64_C(1) << half_bs) - UINT64_C(1);
149 level_metadata_ptr++;
150
151 LOG("\tlevel %6" PRIu64 " output bit-size = %" PRIu64 "\n", nn / 2, bs);
152 }
153
154 // last level (a-b, a+b)
155 {
156 level_metadata_ptr->reduce = (bs == 64);
157 if (level_metadata_ptr->reduce) {
158 bs = bs_after_reduc;
159 LOG("\treduce output bit-size = %" PRIu64 "\n", bs);
160 }
161 for (int k = 0; k < 4; ++k) level_metadata_ptr->q2bs[k] = ((uint64_t)PRIMES_VEC[k] << (bs - logQ));
162 level_metadata_ptr->bs = ++bs;
163 level_metadata_ptr->half_bs = level_metadata_ptr->mask = UINT64_C(0); // not used
164
165 LOG("\tlevel %6" PRIu64 " output bit-size = %" PRIu64 "\n", UINT64_C(1), bs);
166 }
167 precomp->output_bit_size = bs;
168
169 // omega powers
170 uint64_t* powomega = malloc(sizeof(*powomega) * 2 * n);
171 for (uint64_t k = 0; k < 4; ++k) {
172 const uint64_t q = PRIMES_VEC[k];
173
174 for (uint64_t i = 0; i < 2 * n; ++i) {
175 powomega[i] = modq_pow(omega_vec[k], i, q);
176 }
177
178 uint64_t* powomega_ptr = precomp->powomega + k;
179 level_metadata_ptr = precomp->level_metadata;
180
181 {
182 // const uint64_t hpow = UINT64_C(1) << level_metadata_ptr->half_bs;
183 for (uint64_t i = 0; i < n; ++i) {
184 uint64_t t = powomega[i];
185 uint64_t t1 = (t << level_metadata_ptr->half_bs) % q;
186 powomega_ptr[4 * i] = (t1 << 32) + t;
187 }
188 powomega_ptr += 4 * n;
189 level_metadata_ptr++;
190 }
191
192 for (uint64_t nn = n; nn >= 4; nn /= 2) {
193 const uint64_t halfnn = nn / 2;
194 const uint64_t m = n / halfnn;
195
196 // const uint64_t hpow = UINT64_C(1) << level_metadata_ptr->half_bs;
197 for (uint64_t i = 1; i < halfnn; ++i) {
198 uint64_t t = powomega[i * m];
199 uint64_t t1 = (t << level_metadata_ptr->half_bs) % q;
200 powomega_ptr[4 * (i - 1)] = (t1 << 32) + t;
201 }
202 powomega_ptr += 4 * (halfnn - 1);
203 level_metadata_ptr++;
204 }
205 }
206 free(powomega);
207
208 return precomp;
209}
210
211EXPORT q120_ntt_precomp* q120_new_intt_bb_precomp(const uint64_t n) {
212 uint32_t omega_vec[4];
213 fill_omegas(n, omega_vec);
214
215 const uint64_t logQ = ceil(log2(PRIMES_VEC[0]));
216
217 q120_ntt_precomp* precomp = new_precomp(n);
218
219 uint64_t bs = precomp->input_bit_size = 64;
220
221 LOG("iNTT parameters:\n");
222 LOG("\tsize = %" PRIu64 "\n", n)
223 LOG("\tlogQ = %" PRIu64 "\n", logQ);
224 LOG("\tinput bit-size = %" PRIu64 "\n", bs);
225
226 if (n == 1) return precomp;
227
228 // fill reduction metadata
229 uint64_t bs_after_reduc = fill_reduction_meta(bs, &(precomp->reduc_metadata));
230
231 // backward metadata
232 q120_ntt_step_precomp* level_metadata_ptr = precomp->level_metadata;
233
234 // first level (a+b, a-b) adds 1-bit
235 {
236 level_metadata_ptr->reduce = (bs == 64);
237 if (level_metadata_ptr->reduce) {
238 bs = bs_after_reduc;
239 LOG("\treduce output bit-size = %" PRIu64 "\n", bs);
240 }
241
242 for (int k = 0; k < 4; ++k) level_metadata_ptr->q2bs[k] = (uint64_t)PRIMES_VEC[k] << (bs - logQ);
243
244 level_metadata_ptr->bs = ++bs;
245 level_metadata_ptr->half_bs = level_metadata_ptr->mask = UINT64_C(0); // not used
246 LOG("\tlevel %6" PRIu64 " output bit-size = %" PRIu64 "\n", UINT64_C(1), bs);
247 level_metadata_ptr++;
248 }
249
250 for (uint64_t nn = 4; nn <= n; nn *= 2) {
251 level_metadata_ptr->reduce = (bs == 64);
252 if (level_metadata_ptr->reduce) {
253 bs = bs_after_reduc;
254 LOG("\treduce output bit-size = %" PRIu64 "\n", bs);
255 }
256
257 const uint64_t half_bs = round_up_half_n(bs);
258 const uint64_t bs_mult = half_bs + logQ + 1; // bit-size of term b.omega^k
259 bs = 1 + ((bs > bs_mult) ? bs : bs_mult); // bit-size of a+b.omega^k or a-b.omega^k
260 assert(bs <= 64);
261
262 level_metadata_ptr->bs = bs;
263 level_metadata_ptr->half_bs = half_bs;
264 level_metadata_ptr->mask = (UINT64_C(1) << half_bs) - UINT64_C(1);
265 for (int k = 0; k < 4; ++k) level_metadata_ptr->q2bs[k] = (uint64_t)PRIMES_VEC[k] << (bs_mult - logQ);
266 level_metadata_ptr++;
267
268 LOG("\tlevel %6" PRIu64 " output bit-size = %" PRIu64 "\n", nn / 2, bs);
269 }
270
271 // last level a_k.omega^k
272 {
273 level_metadata_ptr->reduce = (bs == 64);
274 if (level_metadata_ptr->reduce) {
275 bs = bs_after_reduc;
276 LOG("\treduce output bit-size = %" PRIu64 "\n", bs);
277 }
278
279 const uint64_t half_bs = round_up_half_n(bs);
280
281 bs = half_bs + logQ + 1; // bit-size of term a.omega^k
282 assert(bs <= 64);
283
284 level_metadata_ptr->bs = bs;
285 level_metadata_ptr->half_bs = half_bs;
286 level_metadata_ptr->mask = (UINT64_C(1) << half_bs) - UINT64_C(1);
287 for (int k = 0; k < 4; ++k) level_metadata_ptr->q2bs[k] = (uint64_t)PRIMES_VEC[k] << (bs - logQ);
288
289 LOG("\tlevel %6" PRIu64 " output bit-size = %" PRIu64 "\n", n, bs);
290 }
291
292 // omega powers
293 uint32_t* powomegabar = malloc(sizeof(*powomegabar) * 2 * n);
294 for (int k = 0; k < 4; ++k) {
295 const uint64_t q = PRIMES_VEC[k];
296
297 for (uint64_t i = 0; i < 2 * n; ++i) {
298 powomegabar[i] = modq_pow(omega_vec[k], -i, q);
299 }
300
301 uint64_t* powomega_ptr = precomp->powomega + k;
302 level_metadata_ptr = precomp->level_metadata + 1;
303
304 for (uint64_t nn = 4; nn <= n; nn *= 2) {
305 const uint64_t halfnn = nn / 2;
306 const uint64_t m = n / halfnn;
307
308 for (uint64_t i = 1; i < halfnn; ++i) {
309 uint64_t t = powomegabar[i * m];
310 uint64_t t1 = (t << level_metadata_ptr->half_bs) % q;
311 powomega_ptr[4 * (i - 1)] = (t1 << 32) + t;
312 }
313 powomega_ptr += 4 * (halfnn - 1);
314 level_metadata_ptr++;
315 }
316
317 {
318 const uint64_t invNmod = modq_pow(n, -1, q);
319 for (uint64_t i = 0; i < n; ++i) {
320 uint64_t t = (powomegabar[i] * invNmod) % q;
321 uint64_t t1 = (t << level_metadata_ptr->half_bs) % q;
322 powomega_ptr[4 * i] = (t1 << 32) + t;
323 }
324 }
325 }
326
327 free(powomegabar);
328
329 return precomp;
330}
331
332void del_precomp(q120_ntt_precomp* precomp) {
333 spqlios_free(precomp->powomega);
334 free(precomp->level_metadata);
335 free(precomp);
336}
337
338EXPORT void q120_del_ntt_bb_precomp(q120_ntt_precomp* precomp) { del_precomp(precomp); }
339
340EXPORT void q120_del_intt_bb_precomp(q120_ntt_precomp* precomp) { del_precomp(precomp); }