1#include <gtest/gtest.h>
  2
  3#include <cstdint>
  4
  5#include "../spqlios/arithmetic/vec_znx_arithmetic_private.h"
  6#include "spqlios/arithmetic/vec_znx_arithmetic.h"
  7#include "test/testlib/ntt120_dft.h"
  8#include "test/testlib/ntt120_layouts.h"
  9#include "testlib/fft64_dft.h"
 10#include "testlib/fft64_layouts.h"
 11#include "testlib/polynomial_vector.h"
 12
 13static void test_fft64_vec_znx_dft(VEC_ZNX_DFT_F dft) {
 14  for (uint64_t n : {2, 4, 128}) {
 15    MODULE* module = new_module_info(n, FFT64);
 16    for (uint64_t sa : {3, 5, 8}) {
 17      for (uint64_t sr : {3, 5, 8}) {
 18        uint64_t a_sl = n + uniform_u64_bits(2);
 19        znx_vec_i64_layout a(n, sa, a_sl);
 20        fft64_vec_znx_dft_layout res(n, sr);
 21        a.fill_random(42);
 22        std::vector<reim_fft64vec> expect(sr);
 23        for (uint64_t i = 0; i < sr; ++i) {
 24          expect[i] = simple_fft64(a.get_copy_zext(i));
 25        }
 26        // test the function
 27        thash hash_before = a.content_hash();
 28        dft(module, res.data, sr, a.data(), sa, a_sl);
 29        ASSERT_EQ(a.content_hash(), hash_before);
 30        for (uint64_t i = 0; i < sr; ++i) {
 31          reim_fft64vec actual = res.get_copy_zext(i);
 32          ASSERT_LE(infty_dist(actual, expect[i]), 1e-10);
 33        }
 34      }
 35    }
 36    delete_module_info(module);
 37  }
 38}
 39
 40static void test_fft64_vec_dft_add(VEC_DFT_ADD_F dft_add) {
 41  for (uint64_t n : {2, 4, 128}) {
 42    MODULE* module = new_module_info(n, FFT64);
 43    for (uint64_t sa : {3, 5, 8}) {
 44      for (uint64_t sr : {3, 5, 8}) {
 45        fft64_vec_znx_dft_layout a(n, sa);
 46        fft64_vec_znx_dft_layout b(n, sa);
 47        fft64_vec_znx_dft_layout res(n, sr);
 48        a.fill_dft_random_log2bound(42);
 49        b.fill_dft_random_log2bound(42);
 50        std::vector<reim_fft64vec> expect(sr);
 51        for (uint64_t i = 0; i < sr; ++i) {
 52          std::vector<double> v(n);
 53          if (i < sa) {
 54            reim_fftvec_add(module->mod.fft64.add_fft, v.data(), ((double*)a.data) + i * n, ((double*)b.data) + i * n);
 55          } else {
 56            std::fill(v.begin(), v.end(), 0.0);
 57          }
 58          expect[i] = reim_fft64vec(n, v.data());
 59        }
 60        // test the function
 61        thash a_hash_before = a.content_hash();
 62        thash b_hash_before = b.content_hash();
 63        dft_add(module, res.data, sr, a.data, sa, b.data, sa);
 64        ASSERT_EQ(a.content_hash(), a_hash_before);
 65        ASSERT_EQ(b.content_hash(), b_hash_before);
 66        for (uint64_t i = 0; i < sr; ++i) {
 67          reim_fft64vec actual = res.get_copy_zext(i);
 68          ASSERT_LE(infty_dist(actual, expect[i]), 1e-10);
 69        }
 70      }
 71    }
 72    delete_module_info(module);
 73  }
 74}
 75
 76#ifdef __x86_64__
 77// FIXME: currently, it only works on avx
 78static void test_ntt120_vec_znx_dft(VEC_ZNX_DFT_F dft) {
 79  for (uint64_t n : {2, 4, 128}) {
 80    MODULE* module = new_module_info(n, NTT120);
 81    for (uint64_t sa : {3, 5, 8}) {
 82      for (uint64_t sr : {3, 5, 8}) {
 83        uint64_t a_sl = n + uniform_u64_bits(2);
 84        znx_vec_i64_layout a(n, sa, a_sl);
 85        ntt120_vec_znx_dft_layout res(n, sr);
 86        a.fill_random(42);
 87        std::vector<q120_nttvec> expect(sr);
 88        for (uint64_t i = 0; i < sr; ++i) {
 89          expect[i] = simple_ntt120(a.get_copy_zext(i));
 90        }
 91        // test the function
 92        thash hash_before = a.content_hash();
 93        dft(module, res.data, sr, a.data(), sa, a_sl);
 94        ASSERT_EQ(a.content_hash(), hash_before);
 95        for (uint64_t i = 0; i < sr; ++i) {
 96          q120_nttvec actual = res.get_copy_zext(i);
 97          if (!(actual == expect[i])) {
 98            for (uint64_t j = 0; j < n; ++j) {
 99              std::cerr << actual.v[j] << " vs " << expect[i].v[j] << std::endl;
100            }
101          }
102          ASSERT_EQ(actual, expect[i]);
103        }
104      }
105    }
106    delete_module_info(module);
107  }
108}
109#endif
110
111TEST(vec_znx_dft, fft64_vec_znx_dft) { test_fft64_vec_znx_dft(fft64_vec_znx_dft); }
112#ifdef __x86_64__
113// FIXME: currently, it only works on avx
114TEST(vec_znx_dft, ntt120_vec_znx_dft) { test_ntt120_vec_znx_dft(ntt120_vec_znx_dft_avx); }
115#endif
116TEST(vec_znx_dft, vec_znx_dft) {
117  test_fft64_vec_znx_dft(vec_znx_dft);
118#ifdef __x86_64__
119  // FIXME: currently, it only works on avx
120  test_ntt120_vec_znx_dft(ntt120_vec_znx_dft_avx);
121#endif
122}
123
124TEST(vec_dft_add, fft64_vec_dft_add) { test_fft64_vec_dft_add(fft64_vec_dft_add); }
125
126static void test_fft64_vec_znx_idft(VEC_ZNX_IDFT_F idft, VEC_ZNX_IDFT_TMP_A_F idft_tmp_a,
127                                    VEC_ZNX_IDFT_TMP_BYTES_F idft_tmp_bytes) {
128  for (uint64_t n : {2, 4, 64, 128}) {
129    MODULE* module = new_module_info(n, FFT64);
130    uint64_t tmp_size = idft_tmp_bytes ? idft_tmp_bytes(module, n) : 0;
131    std::vector<uint8_t> tmp(tmp_size);
132    for (uint64_t sa : {3, 5, 8}) {
133      for (uint64_t sr : {3, 5, 8}) {
134        fft64_vec_znx_dft_layout a(n, sa);
135        fft64_vec_znx_big_layout res(n, sr);
136        a.fill_dft_random_log2bound(22);
137        std::vector<znx_i64> expect(sr);
138        for (uint64_t i = 0; i < sr; ++i) {
139          expect[i] = simple_rint_ifft64(a.get_copy_zext(i));
140        }
141        // test the function
142        if (idft_tmp_bytes) {
143          thash hash_before = a.content_hash();
144          idft(module, res.data, sr, a.data, sa, tmp.data());
145          ASSERT_EQ(a.content_hash(), hash_before);
146        } else {
147          idft_tmp_a(module, res.data, sr, a.data, sa);
148        }
149        for (uint64_t i = 0; i < sr; ++i) {
150          znx_i64 actual = res.get_copy_zext(i);
151          // ASSERT_EQ(res.get_copy_zext(i), expect[i]);
152          if (!(actual == expect[i])) {
153            for (uint64_t j = 0; j < n; ++j) {
154              std::cerr << actual.get_coeff(j) << " dft vs. " << expect[i].get_coeff(j) << std::endl;
155            }
156            FAIL();
157          }
158        }
159      }
160    }
161    delete_module_info(module);
162  }
163}
164
165TEST(vec_znx_dft, fft64_vec_znx_idft) {
166  test_fft64_vec_znx_idft(fft64_vec_znx_idft, nullptr, fft64_vec_znx_idft_tmp_bytes);
167}
168TEST(vec_znx_dft, fft64_vec_znx_idft_tmp_a) { test_fft64_vec_znx_idft(nullptr, fft64_vec_znx_idft_tmp_a, nullptr); }
169
170#ifdef __x86_64__
171// FIXME: currently, it only works on avx
172static void test_ntt120_vec_znx_idft(VEC_ZNX_IDFT_F idft, VEC_ZNX_IDFT_TMP_A_F idft_tmp_a,
173                                     VEC_ZNX_IDFT_TMP_BYTES_F idft_tmp_bytes) {
174  for (uint64_t n : {2, 4, 64, 128}) {
175    MODULE* module = new_module_info(n, NTT120);
176    uint64_t tmp_size = idft_tmp_bytes ? idft_tmp_bytes(module, n) : 0;
177    std::vector<uint8_t> tmp(tmp_size);
178    for (uint64_t sa : {3, 5, 8}) {
179      for (uint64_t sr : {3, 5, 8}) {
180        ntt120_vec_znx_dft_layout a(n, sa);
181        ntt120_vec_znx_big_layout res(n, sr);
182        a.fill_random();
183        std::vector<znx_i128> expect(sr);
184        for (uint64_t i = 0; i < sr; ++i) {
185          expect[i] = simple_intt120(a.get_copy_zext(i));
186        }
187        // test the function
188        if (idft_tmp_bytes) {
189          thash hash_before = a.content_hash();
190          idft(module, res.data, sr, a.data, sa, tmp.data());
191          ASSERT_EQ(a.content_hash(), hash_before);
192        } else {
193          idft_tmp_a(module, res.data, sr, a.data, sa);
194        }
195        for (uint64_t i = 0; i < sr; ++i) {
196          znx_i128 actual = res.get_copy_zext(i);
197          ASSERT_EQ(res.get_copy_zext(i), expect[i]);
198          // if (!(actual == expect[i])) {
199          //   for (uint64_t j = 0; j < n; ++j) {
200          //     std::cerr << actual.get_coeff(j) << " dft vs. " << expect[i].get_coeff(j) << std::endl;
201          //   }
202          //   FAIL();
203          // }
204        }
205      }
206    }
207    delete_module_info(module);
208  }
209}
210
211TEST(vec_znx_dft, ntt120_vec_znx_idft) {
212  test_ntt120_vec_znx_idft(ntt120_vec_znx_idft_avx, nullptr, ntt120_vec_znx_idft_tmp_bytes_avx);
213}
214TEST(vec_znx_dft, ntt120_vec_znx_idft_tmp_a) {
215  test_ntt120_vec_znx_idft(nullptr, ntt120_vec_znx_idft_tmp_a_avx, nullptr);
216}
217#endif
218TEST(vec_znx_dft, vec_znx_idft) {
219  test_fft64_vec_znx_idft(vec_znx_idft, nullptr, vec_znx_idft_tmp_bytes);
220#ifdef __x86_64__
221  // FIXME: currently, only supported on avx
222  test_ntt120_vec_znx_idft(vec_znx_idft, nullptr, vec_znx_idft_tmp_bytes);
223#endif
224}
225TEST(vec_znx_dft, vec_znx_idft_tmp_a) {
226  test_fft64_vec_znx_idft(nullptr, vec_znx_idft_tmp_a, nullptr);
227#ifdef __x86_64__
228  // FIXME: currently, only supported on avx
229  test_ntt120_vec_znx_idft(nullptr, vec_znx_idft_tmp_a, nullptr);
230#endif
231}