1#include <gtest/gtest.h>
  2
  3#include <cstdint>
  4#include <vector>
  5
  6#include "spqlios/q120/q120_arithmetic.h"
  7#include "test/testlib/negacyclic_polynomial.h"
  8#include "test/testlib/ntt120_layouts.h"
  9#include "testlib/mod_q120.h"
 10
 11typedef typeof(q120_vec_mat1col_product_baa_ref) vec_mat1col_product_baa_f;
 12
 13void test_vec_mat1col_product_baa(vec_mat1col_product_baa_f vec_mat1col_product_baa) {
 14  q120_mat1col_product_baa_precomp* precomp = q120_new_vec_mat1col_product_baa_precomp();
 15  for (uint64_t ell : {1, 2, 100, 10000}) {
 16    std::vector<uint64_t> a(ell * 4);
 17    std::vector<uint64_t> b(ell * 4);
 18    std::vector<uint64_t> res(4);
 19    uint64_t* pa = a.data();
 20    uint64_t* pb = b.data();
 21    uint64_t* pr = res.data();
 22    // generate some random data
 23    uniform_q120b(pr);
 24    for (uint64_t i = 0; i < ell; ++i) {
 25      uniform_q120a(pa + 4 * i);
 26      uniform_q120a(pb + 4 * i);
 27    }
 28    // compute the expected result
 29    mod_q120 expect_r;
 30    for (uint64_t i = 0; i < ell; ++i) {
 31      expect_r += mod_q120::from_q120a(pa + 4 * i) * mod_q120::from_q120a(pb + 4 * i);
 32    }
 33    // compute the function
 34    vec_mat1col_product_baa(precomp, ell, (q120b*)pr, (q120a*)pa, (q120a*)pb);
 35    mod_q120 comp_r = mod_q120::from_q120b(pr);
 36    // check for equality
 37    ASSERT_EQ(comp_r, expect_r) << ell;
 38  }
 39  q120_delete_vec_mat1col_product_baa_precomp(precomp);
 40}
 41
 42TEST(q120_arithmetic, q120_vec_mat1col_product_baa_ref) {
 43  test_vec_mat1col_product_baa(q120_vec_mat1col_product_baa_ref);
 44}
 45#ifdef __x86_64__
 46TEST(q120_arithmetic, q120_vec_mat1col_product_baa_avx2) {
 47  test_vec_mat1col_product_baa(q120_vec_mat1col_product_baa_avx2);
 48}
 49#endif
 50
 51typedef typeof(q120_vec_mat1col_product_bbb_ref) vec_mat1col_product_bbb_f;
 52
 53void test_vec_mat1col_product_bbb(vec_mat1col_product_bbb_f vec_mat1col_product_bbb) {
 54  q120_mat1col_product_bbb_precomp* precomp = q120_new_vec_mat1col_product_bbb_precomp();
 55  for (uint64_t ell : {1, 2, 100, 10000}) {
 56    std::vector<uint64_t> a(ell * 4);
 57    std::vector<uint64_t> b(ell * 4);
 58    std::vector<uint64_t> res(4);
 59    uint64_t* pa = a.data();
 60    uint64_t* pb = b.data();
 61    uint64_t* pr = res.data();
 62    // generate some random data
 63    uniform_q120b(pr);
 64    for (uint64_t i = 0; i < ell; ++i) {
 65      uniform_q120b(pa + 4 * i);
 66      uniform_q120b(pb + 4 * i);
 67    }
 68    // compute the expected result
 69    mod_q120 expect_r;
 70    for (uint64_t i = 0; i < ell; ++i) {
 71      expect_r += mod_q120::from_q120b(pa + 4 * i) * mod_q120::from_q120b(pb + 4 * i);
 72    }
 73    // compute the function
 74    vec_mat1col_product_bbb(precomp, ell, (q120b*)pr, (q120b*)pa, (q120b*)pb);
 75    mod_q120 comp_r = mod_q120::from_q120b(pr);
 76    // check for equality
 77    ASSERT_EQ(comp_r, expect_r);
 78  }
 79  q120_delete_vec_mat1col_product_bbb_precomp(precomp);
 80}
 81
 82TEST(q120_arithmetic, q120_vec_mat1col_product_bbb_ref) {
 83  test_vec_mat1col_product_bbb(q120_vec_mat1col_product_bbb_ref);
 84}
 85#ifdef __x86_64__
 86TEST(q120_arithmetic, q120_vec_mat1col_product_bbb_avx2) {
 87  test_vec_mat1col_product_bbb(q120_vec_mat1col_product_bbb_avx2);
 88}
 89#endif
 90
 91typedef typeof(q120_vec_mat1col_product_bbc_ref) vec_mat1col_product_bbc_f;
 92
 93void test_vec_mat1col_product_bbc(vec_mat1col_product_bbc_f vec_mat1col_product_bbc) {
 94  q120_mat1col_product_bbc_precomp* precomp = q120_new_vec_mat1col_product_bbc_precomp();
 95  for (uint64_t ell : {1, 2, 100, 10000}) {
 96    std::vector<uint64_t> a(ell * 4);
 97    std::vector<uint64_t> b(ell * 4);
 98    std::vector<uint64_t> res(4);
 99    uint64_t* pa = a.data();
100    uint64_t* pb = b.data();
101    uint64_t* pr = res.data();
102    // generate some random data
103    uniform_q120b(pr);
104    for (uint64_t i = 0; i < ell; ++i) {
105      uniform_q120b(pa + 4 * i);
106      uniform_q120c(pb + 4 * i);
107    }
108    // compute the expected result
109    mod_q120 expect_r;
110    for (uint64_t i = 0; i < ell; ++i) {
111      expect_r += mod_q120::from_q120b(pa + 4 * i) * mod_q120::from_q120c(pb + 4 * i);
112    }
113    // compute the function
114    vec_mat1col_product_bbc(precomp, ell, (q120b*)pr, (q120b*)pa, (q120c*)pb);
115    mod_q120 comp_r = mod_q120::from_q120b(pr);
116    // check for equality
117    ASSERT_EQ(comp_r, expect_r);
118  }
119  q120_delete_vec_mat1col_product_bbc_precomp(precomp);
120}
121
122TEST(q120_arithmetic, q120_vec_mat1col_product_bbc_ref) {
123  test_vec_mat1col_product_bbc(q120_vec_mat1col_product_bbc_ref);
124}
125#ifdef __x86_64__
126TEST(q120_arithmetic, q120_vec_mat1col_product_bbc_avx2) {
127  test_vec_mat1col_product_bbc(q120_vec_mat1col_product_bbc_avx2);
128}
129#endif
130
131EXPORT void q120x2_vec_mat2cols_product_bbc_avx2(q120_mat1col_product_bbc_precomp* precomp, const uint64_t ell,
132                                                 q120b* const res, const q120b* const x, const q120c* const y);
133EXPORT void q120x2_vec_mat1col_product_bbc_avx2(q120_mat1col_product_bbc_precomp* precomp, const uint64_t ell,
134                                                q120b* const res, const q120b* const x, const q120c* const y);
135
136typedef typeof(q120x2_vec_mat2cols_product_bbc_avx2) q120x2_prod_bbc_f;
137
138void test_q120x2_vec_mat2cols_product_bbc(q120x2_prod_bbc_f q120x2_prod_bbc) {
139  q120_mat1col_product_bbc_precomp* precomp = q120_new_vec_mat1col_product_bbc_precomp();
140  for (uint64_t ell : {1, 2, 100, 10000}) {
141    std::vector<uint64_t> a(ell * 8);
142    std::vector<uint64_t> b(ell * 16);
143    std::vector<uint64_t> res(16);
144    uint64_t* pa = a.data();
145    uint64_t* pb = b.data();
146    uint64_t* pr = res.data();
147    // generate some random data
148    uniform_q120b(pr);
149    for (uint64_t i = 0; i < 2 * ell; ++i) {
150      uniform_q120b(pa + 4 * i);
151    }
152    for (uint64_t i = 0; i < 4 * ell; ++i) {
153      uniform_q120c(pb + 4 * i);
154    }
155    // compute the expected result
156    mod_q120 expect_r[4];
157    for (uint64_t i = 0; i < ell; ++i) {
158      mod_q120 va = mod_q120::from_q120b(pa + 8 * i);
159      mod_q120 vb = mod_q120::from_q120b(pa + 8 * i + 4);
160      mod_q120 m1a = mod_q120::from_q120c(pb + 16 * i);
161      mod_q120 m1b = mod_q120::from_q120c(pb + 16 * i + 4);
162      mod_q120 m2a = mod_q120::from_q120c(pb + 16 * i + 8);
163      mod_q120 m2b = mod_q120::from_q120c(pb + 16 * i + 12);
164      expect_r[0] += va * m1a;
165      expect_r[1] += vb * m1b;
166      expect_r[2] += va * m2a;
167      expect_r[3] += vb * m2b;
168    }
169    // compute the function
170    q120x2_prod_bbc(precomp, ell, (q120b*)pr, (q120b*)pa, (q120c*)pb);
171    // check for equality
172    ASSERT_EQ(mod_q120::from_q120b(pr), expect_r[0]);
173    ASSERT_EQ(mod_q120::from_q120b(pr + 4), expect_r[1]);
174    ASSERT_EQ(mod_q120::from_q120b(pr + 8), expect_r[2]);
175    ASSERT_EQ(mod_q120::from_q120b(pr + 12), expect_r[3]);
176  }
177  q120_delete_vec_mat1col_product_bbc_precomp(precomp);
178}
179
180TEST(q120_arithmetic, q120x2_vec_mat2cols_product_bbc_ref) {
181  test_q120x2_vec_mat2cols_product_bbc(q120x2_vec_mat2cols_product_bbc_ref);
182}
183#ifdef __x86_64__
184TEST(q120_arithmetic, q120x2_vec_mat2cols_product_bbc_avx2) {
185  test_q120x2_vec_mat2cols_product_bbc(q120x2_vec_mat2cols_product_bbc_avx2);
186}
187#endif
188
189typedef typeof(q120x2_vec_mat1col_product_bbc_avx2) q120x2_c1_prod_bbc_f;
190
191void test_q120x2_vec_mat1col_product_bbc(q120x2_c1_prod_bbc_f q120x2_c1_prod_bbc) {
192  q120_mat1col_product_bbc_precomp* precomp = q120_new_vec_mat1col_product_bbc_precomp();
193  for (uint64_t ell : {1, 2, 100, 10000}) {
194    std::vector<uint64_t> a(ell * 8);
195    std::vector<uint64_t> b(ell * 8);
196    std::vector<uint64_t> res(8);
197    uint64_t* pa = a.data();
198    uint64_t* pb = b.data();
199    uint64_t* pr = res.data();
200    // generate some random data
201    uniform_q120b(pr);
202    for (uint64_t i = 0; i < 2 * ell; ++i) {
203      uniform_q120b(pa + 4 * i);
204    }
205    for (uint64_t i = 0; i < 2 * ell; ++i) {
206      uniform_q120c(pb + 4 * i);
207    }
208    // compute the expected result
209    mod_q120 expect_r[2];
210    for (uint64_t i = 0; i < ell; ++i) {
211      mod_q120 va = mod_q120::from_q120b(pa + 8 * i);
212      mod_q120 vb = mod_q120::from_q120b(pa + 8 * i + 4);
213      mod_q120 m1a = mod_q120::from_q120c(pb + 8 * i);
214      mod_q120 m1b = mod_q120::from_q120c(pb + 8 * i + 4);
215      expect_r[0] += va * m1a;
216      expect_r[1] += vb * m1b;
217    }
218    // compute the function
219    q120x2_c1_prod_bbc(precomp, ell, (q120b*)pr, (q120b*)pa, (q120c*)pb);
220    // check for equality
221    ASSERT_EQ(mod_q120::from_q120b(pr), expect_r[0]);
222    ASSERT_EQ(mod_q120::from_q120b(pr + 4), expect_r[1]);
223  }
224  q120_delete_vec_mat1col_product_bbc_precomp(precomp);
225}
226
227TEST(q120_arithmetic, q120x2_vec_mat1col_product_bbc_ref) {
228  test_q120x2_vec_mat1col_product_bbc(q120x2_vec_mat1col_product_bbc_ref);
229}
230#ifdef __x86_64__
231TEST(q120_arithmetic, q120x2_vec_mat1col_product_bbc_avx2) {
232  test_q120x2_vec_mat1col_product_bbc(q120x2_vec_mat1col_product_bbc_avx2);
233}
234#endif
235
236typedef typeof(q120x2_extract_1blk_from_q120b_ref) q120x2_extract_f;
237void test_q120x2_extract_1blk(q120x2_extract_f q120x2_extract) {
238  for (uint64_t n : {2, 4, 64}) {
239    ntt120_vec_znx_dft_layout v(n, 1);
240    std::vector<uint64_t> r(8);
241    std::vector<uint64_t> expect(8);
242    for (uint64_t blk = 0; blk < n / 2; ++blk) {
243      for (uint64_t i = 0; i < 8; ++i) {
244        expect[i] = uniform_u64();
245      }
246      memcpy(v.get_blk(0, blk), expect.data(), 8 * sizeof(uint64_t));
247      q120x2_extract_1blk_from_q120b_ref(n, blk, (q120x2b*)r.data(), (q120b*)v.data);
248      ASSERT_EQ(r, expect);
249    }
250  }
251}
252
253TEST(q120_arithmetic, q120x2_extract_1blk_from_q120b_ref) {
254  test_q120x2_extract_1blk(q120x2_extract_1blk_from_q120b_ref);
255}
256
257typedef typeof(q120x2_extract_1blk_from_contiguous_q120b_ref) q120x2_extract_vec_f;
258void test_q120x2_extract_1blk_vec(q120x2_extract_vec_f q120x2_extract) {
259  for (uint64_t n : {2, 4, 32}) {
260    for (uint64_t size : {1, 2, 7}) {
261      ntt120_vec_znx_dft_layout v(n, size);
262      std::vector<uint64_t> r(8 * size);
263      std::vector<uint64_t> expect(8 * size);
264      for (uint64_t blk = 0; blk < n / 2; ++blk) {
265        for (uint64_t i = 0; i < 8 * size; ++i) {
266          expect[i] = uniform_u64();
267        }
268        for (uint64_t i = 0; i < size; ++i) {
269          memcpy(v.get_blk(i, blk), expect.data() + 8 * i, 8 * sizeof(uint64_t));
270        }
271        q120x2_extract(n, size, blk, (q120x2b*)r.data(), (q120b*)v.data);
272        ASSERT_EQ(r, expect);
273      }
274    }
275  }
276}
277
278TEST(q120_arithmetic, q120x2_extract_1blk_from_contiguous_q120b_ref) {
279  test_q120x2_extract_1blk_vec(q120x2_extract_1blk_from_contiguous_q120b_ref);
280}
281
282typedef typeof(q120x2b_save_1blk_to_q120b_ref) q120x2_save_f;
283void test_q120x2_save_1blk(q120x2_save_f q120x2_save) {
284  for (uint64_t n : {2, 4, 64}) {
285    ntt120_vec_znx_dft_layout v(n, 1);
286    std::vector<uint64_t> r(8);
287    std::vector<uint64_t> expect(8);
288    for (uint64_t blk = 0; blk < n / 2; ++blk) {
289      for (uint64_t i = 0; i < 8; ++i) {
290        expect[i] = uniform_u64();
291      }
292      q120x2_save(n, blk, (q120b*)v.data, (q120x2b*)expect.data());
293      memcpy(r.data(), v.get_blk(0, blk), 8 * sizeof(uint64_t));
294      ASSERT_EQ(r, expect);
295    }
296  }
297}
298
299TEST(q120_arithmetic, q120x2b_save_1blk_to_q120b_ref) { test_q120x2_save_1blk(q120x2b_save_1blk_to_q120b_ref); }
300
301TEST(q120_arithmetic, q120_add_bbb_simple) {
302  for (const uint64_t n : {2, 4, 1024}) {
303    std::vector<uint64_t> a(n * 4);
304    std::vector<uint64_t> b(n * 4);
305    std::vector<uint64_t> r(n * 4);
306    uint64_t* pa = a.data();
307    uint64_t* pb = b.data();
308    uint64_t* pr = r.data();
309
310    // generate some random data
311    for (uint64_t i = 0; i < n; ++i) {
312      uniform_q120b(pa + 4 * i);
313      uniform_q120b(pb + 4 * i);
314    }
315
316    // compute the function
317    q120_add_bbb_simple(n, (q120b*)pr, (q120b*)pa, (q120b*)pb);
318
319    for (uint64_t i = 0; i < n; ++i) {
320      mod_q120 ae = mod_q120::from_q120b(pa + 4 * i);
321      mod_q120 be = mod_q120::from_q120b(pb + 4 * i);
322      mod_q120 re = mod_q120::from_q120b(pr + 4 * i);
323
324      ASSERT_EQ(ae + be, re);
325    }
326  }
327}
328
329TEST(q120_arithmetic, q120_add_ccc_simple) {
330  for (const uint64_t n : {2, 4, 1024}) {
331    std::vector<uint64_t> a(n * 4);
332    std::vector<uint64_t> b(n * 4);
333    std::vector<uint64_t> r(n * 4);
334    uint64_t* pa = a.data();
335    uint64_t* pb = b.data();
336    uint64_t* pr = r.data();
337
338    // generate some random data
339    for (uint64_t i = 0; i < n; ++i) {
340      uniform_q120c(pa + 4 * i);
341      uniform_q120c(pb + 4 * i);
342    }
343
344    // compute the function
345    q120_add_ccc_simple(n, (q120c*)pr, (q120c*)pa, (q120c*)pb);
346
347    for (uint64_t i = 0; i < n; ++i) {
348      mod_q120 ae = mod_q120::from_q120c(pa + 4 * i);
349      mod_q120 be = mod_q120::from_q120c(pb + 4 * i);
350      mod_q120 re = mod_q120::from_q120c(pr + 4 * i);
351
352      ASSERT_EQ(ae + be, re);
353    }
354  }
355}
356
357TEST(q120_arithmetic, q120_c_from_b_simple) {
358  for (const uint64_t n : {2, 4, 1024}) {
359    std::vector<uint64_t> a(n * 4);
360    std::vector<uint64_t> r(n * 4);
361    uint64_t* pa = a.data();
362    uint64_t* pr = r.data();
363
364    // generate some random data
365    for (uint64_t i = 0; i < n; ++i) {
366      uniform_q120b(pa + 4 * i);
367    }
368
369    // compute the function
370    q120_c_from_b_simple(n, (q120c*)pr, (q120b*)pa);
371
372    for (uint64_t i = 0; i < n; ++i) {
373      mod_q120 ae = mod_q120::from_q120b(pa + 4 * i);
374      mod_q120 re = mod_q120::from_q120c(pr + 4 * i);
375
376      ASSERT_EQ(ae, re);
377    }
378  }
379}
380
381TEST(q120_arithmetic, q120_b_from_znx64_simple) {
382  for (const uint64_t n : {2, 4, 1024}) {
383    znx_i64 x = znx_i64::random_log2bound(n, 62);
384    std::vector<uint64_t> r(n * 4);
385    uint64_t* pr = r.data();
386
387    q120_b_from_znx64_simple(n, (q120b*)pr, x.data());
388
389    for (uint64_t i = 0; i < n; ++i) {
390      mod_q120 re = mod_q120::from_q120b(pr + 4 * i);
391
392      for (uint64_t k = 0; k < 4; ++k) {
393        ASSERT_EQ(centermod(x.get_coeff(i), mod_q120::Qi[k]), re.a[k]);
394      }
395    }
396  }
397}
398
399TEST(q120_arithmetic, q120_c_from_znx64_simple) {
400  for (const uint64_t n : {2, 4, 1024}) {
401    znx_i64 x = znx_i64::random(n);
402    std::vector<uint64_t> r(n * 4);
403    uint64_t* pr = r.data();
404
405    q120_c_from_znx64_simple(n, (q120c*)pr, x.data());
406
407    for (uint64_t i = 0; i < n; ++i) {
408      mod_q120 re = mod_q120::from_q120c(pr + 4 * i);
409
410      for (uint64_t k = 0; k < 4; ++k) {
411        ASSERT_EQ(centermod(x.get_coeff(i), mod_q120::Qi[k]), re.a[k]);
412      }
413    }
414  }
415}
416
417TEST(q120_arithmetic, q120_b_to_znx128_simple) {
418  for (const uint64_t n : {2, 4, 1024}) {
419    std::vector<uint64_t> x(n * 4);
420    uint64_t* px = x.data();
421
422    // generate some random data
423    for (uint64_t i = 0; i < n; ++i) {
424      uniform_q120b(px + 4 * i);
425    }
426
427    znx_i128 r(n);
428    q120_b_to_znx128_simple(n, r.data(), (q120b*)px);
429
430    for (uint64_t i = 0; i < n; ++i) {
431      mod_q120 xe = mod_q120::from_q120b(px + 4 * i);
432      for (uint64_t k = 0; k < 4; ++k) {
433        ASSERT_EQ(centermod((int64_t)(r.get_coeff(i) % mod_q120::Qi[k]), mod_q120::Qi[k]), xe.a[k]);
434      }
435    }
436  }
437}