1#include <gtest/gtest.h>
  2
  3#include <cassert>
  4#include <cstdint>
  5#include <cstdio>
  6#include <random>
  7#include <vector>
  8
  9#include "spqlios/q120/q120_common.h"
 10#include "spqlios/q120/q120_ntt.h"
 11#include "testlib/mod_q120.h"
 12
 13std::vector<mod_q120> q120_ntt(const std::vector<mod_q120>& x) {
 14  const uint64_t n = x.size();
 15
 16  mod_q120 omega_2pow17{OMEGA1, OMEGA2, OMEGA3, OMEGA4};
 17  mod_q120 omega = pow(omega_2pow17, (1 << 16) / n);
 18
 19  std::vector<mod_q120> res(n);
 20  for (uint64_t i = 0; i < n; ++i) {
 21    res[i] = x[i];
 22  }
 23
 24  for (uint64_t i = 0; i < n; ++i) {
 25    res[i] = res[i] * pow(omega, i);
 26  }
 27
 28  for (uint64_t nn = n; nn > 1; nn /= 2) {
 29    const uint64_t halfnn = nn / 2;
 30    const uint64_t m = n / halfnn;
 31
 32    for (uint64_t j = 0; j < n; j += nn) {
 33      for (uint64_t k = 0; k < halfnn; ++k) {
 34        mod_q120 a = res[j + k];
 35        mod_q120 b = res[j + halfnn + k];
 36
 37        res[j + k] = a + b;
 38        res[j + halfnn + k] = (a - b) * pow(omega, k * m);
 39      }
 40    }
 41  }
 42
 43  return res;
 44}
 45
 46std::vector<mod_q120> q120_intt(const std::vector<mod_q120>& x) {
 47  const uint64_t n = x.size();
 48
 49  mod_q120 omega_2pow17{OMEGA1, OMEGA2, OMEGA3, OMEGA4};
 50  mod_q120 omega = pow(omega_2pow17, (1 << 16) / n);
 51
 52  std::vector<mod_q120> res(n);
 53  for (uint64_t i = 0; i < n; ++i) {
 54    res[i] = x[i];
 55  }
 56
 57  for (uint64_t nn = 2; nn <= n; nn *= 2) {
 58    const uint64_t halfnn = nn / 2;
 59    const uint64_t m = n / halfnn;
 60
 61    for (uint64_t j = 0; j < n; j += nn) {
 62      for (uint64_t k = 0; k < halfnn; ++k) {
 63        mod_q120 a = res[j + k];
 64        mod_q120 b = res[j + halfnn + k];
 65
 66        mod_q120 bo = b * pow(omega, -k * m);
 67        res[j + k] = a + bo;
 68        res[j + halfnn + k] = a - bo;
 69      }
 70    }
 71  }
 72
 73  mod_q120 n_q120{(int64_t)n, (int64_t)n, (int64_t)n, (int64_t)n};
 74  mod_q120 n_inv_q120 = pow(n_q120, -1);
 75
 76  for (uint64_t i = 0; i < n; ++i) {
 77    mod_q120 po = pow(omega, -i) * n_inv_q120;
 78    res[i] = res[i] * po;
 79  }
 80
 81  return res;
 82}
 83
 84class ntt : public testing::TestWithParam<uint64_t> {};
 85
 86#ifdef __x86_64__
 87
 88TEST_P(ntt, q120_ntt_bb_avx2) {
 89  const uint64_t n = GetParam();
 90  q120_ntt_precomp* precomp = q120_new_ntt_bb_precomp(n);
 91
 92  std::vector<uint64_t> x(n * 4);
 93  uint64_t* px = x.data();
 94  for (uint64_t i = 0; i < 4 * n; i += 4) {
 95    uniform_q120b(px + i);
 96  }
 97
 98  std::vector<mod_q120> x_modq(n);
 99  for (uint64_t i = 0; i < n; ++i) {
100    x_modq[i] = mod_q120::from_q120b(px + 4 * i);
101  }
102
103  std::vector<mod_q120> y_exp = q120_ntt(x_modq);
104
105  q120_ntt_bb_avx2(precomp, (q120b*)px);
106
107  for (uint64_t i = 0; i < n; ++i) {
108    mod_q120 comp_r = mod_q120::from_q120b(px + 4 * i);
109    ASSERT_EQ(comp_r, y_exp[i]) << i;
110  }
111
112  q120_del_ntt_bb_precomp(precomp);
113}
114
115TEST_P(ntt, q120_intt_bb_avx2) {
116  const uint64_t n = GetParam();
117  q120_ntt_precomp* precomp = q120_new_intt_bb_precomp(n);
118
119  std::vector<uint64_t> x(n * 4);
120  uint64_t* px = x.data();
121  for (uint64_t i = 0; i < 4 * n; i += 4) {
122    uniform_q120b(px + i);
123  }
124
125  std::vector<mod_q120> x_modq(n);
126  for (uint64_t i = 0; i < n; ++i) {
127    x_modq[i] = mod_q120::from_q120b(px + 4 * i);
128  }
129
130  q120_intt_bb_avx2(precomp, (q120b*)px);
131
132  std::vector<mod_q120> y_exp = q120_intt(x_modq);
133  for (uint64_t i = 0; i < n; ++i) {
134    mod_q120 comp_r = mod_q120::from_q120b(px + 4 * i);
135    ASSERT_EQ(comp_r, y_exp[i]) << i;
136  }
137
138  q120_del_intt_bb_precomp(precomp);
139}
140
141TEST_P(ntt, q120_ntt_intt_bb_avx2) {
142  const uint64_t n = GetParam();
143  q120_ntt_precomp* precomp_ntt = q120_new_ntt_bb_precomp(n);
144  q120_ntt_precomp* precomp_intt = q120_new_intt_bb_precomp(n);
145
146  std::vector<uint64_t> x(n * 4);
147  uint64_t* px = x.data();
148  for (uint64_t i = 0; i < 4 * n; i += 4) {
149    uniform_q120b(px + i);
150  }
151
152  std::vector<mod_q120> x_modq(n);
153  for (uint64_t i = 0; i < n; ++i) {
154    x_modq[i] = mod_q120::from_q120b(px + 4 * i);
155  }
156
157  q120_ntt_bb_avx2(precomp_ntt, (q120b*)px);
158  q120_intt_bb_avx2(precomp_intt, (q120b*)px);
159
160  for (uint64_t i = 0; i < n; ++i) {
161    mod_q120 comp_r = mod_q120::from_q120b(px + 4 * i);
162    ASSERT_EQ(comp_r, x_modq[i]) << i;
163  }
164
165  q120_del_intt_bb_precomp(precomp_intt);
166  q120_del_ntt_bb_precomp(precomp_ntt);
167}
168
169INSTANTIATE_TEST_SUITE_P(q120, ntt,
170                         testing::Values(1, 2, 4, 16, 256, UINT64_C(1) << 10, UINT64_C(1) << 11, UINT64_C(1) << 12,
171                                         UINT64_C(1) << 13, UINT64_C(1) << 14, UINT64_C(1) << 15, UINT64_C(1) << 16),
172                         testing::PrintToStringParamName());
173
174#endif