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
 40#ifdef __x86_64__
 41// FIXME: currently, it only works on avx
 42static void test_ntt120_vec_znx_dft(VEC_ZNX_DFT_F dft) {
 43  for (uint64_t n : {2, 4, 128}) {
 44    MODULE* module = new_module_info(n, NTT120);
 45    for (uint64_t sa : {3, 5, 8}) {
 46      for (uint64_t sr : {3, 5, 8}) {
 47        uint64_t a_sl = n + uniform_u64_bits(2);
 48        znx_vec_i64_layout a(n, sa, a_sl);
 49        ntt120_vec_znx_dft_layout res(n, sr);
 50        a.fill_random(42);
 51        std::vector<q120_nttvec> expect(sr);
 52        for (uint64_t i = 0; i < sr; ++i) {
 53          expect[i] = simple_ntt120(a.get_copy_zext(i));
 54        }
 55        // test the function
 56        thash hash_before = a.content_hash();
 57        dft(module, res.data, sr, a.data(), sa, a_sl);
 58        ASSERT_EQ(a.content_hash(), hash_before);
 59        for (uint64_t i = 0; i < sr; ++i) {
 60          q120_nttvec actual = res.get_copy_zext(i);
 61          if (!(actual == expect[i])) {
 62            for (uint64_t j = 0; j < n; ++j) {
 63              std::cerr << actual.v[j] << " vs " << expect[i].v[j] << std::endl;
 64            }
 65          }
 66          ASSERT_EQ(actual, expect[i]);
 67        }
 68      }
 69    }
 70    delete_module_info(module);
 71  }
 72}
 73#endif
 74
 75TEST(vec_znx_dft, fft64_vec_znx_dft) { test_fft64_vec_znx_dft(fft64_vec_znx_dft); }
 76#ifdef __x86_64__
 77// FIXME: currently, it only works on avx
 78TEST(vec_znx_dft, ntt120_vec_znx_dft) { test_ntt120_vec_znx_dft(ntt120_vec_znx_dft_avx); }
 79#endif
 80TEST(vec_znx_dft, vec_znx_dft) {
 81  test_fft64_vec_znx_dft(vec_znx_dft);
 82#ifdef __x86_64__
 83  // FIXME: currently, it only works on avx
 84  test_ntt120_vec_znx_dft(ntt120_vec_znx_dft_avx);
 85#endif
 86}
 87
 88static void test_fft64_vec_znx_idft(VEC_ZNX_IDFT_F idft, VEC_ZNX_IDFT_TMP_A_F idft_tmp_a,
 89                                    VEC_ZNX_IDFT_TMP_BYTES_F idft_tmp_bytes) {
 90  for (uint64_t n : {2, 4, 64, 128}) {
 91    MODULE* module = new_module_info(n, FFT64);
 92    uint64_t tmp_size = idft_tmp_bytes ? idft_tmp_bytes(module) : 0;
 93    std::vector<uint8_t> tmp(tmp_size);
 94    for (uint64_t sa : {3, 5, 8}) {
 95      for (uint64_t sr : {3, 5, 8}) {
 96        fft64_vec_znx_dft_layout a(n, sa);
 97        fft64_vec_znx_big_layout res(n, sr);
 98        a.fill_dft_random_log2bound(22);
 99        std::vector<znx_i64> expect(sr);
100        for (uint64_t i = 0; i < sr; ++i) {
101          expect[i] = simple_rint_ifft64(a.get_copy_zext(i));
102        }
103        // test the function
104        if (idft_tmp_bytes) {
105          thash hash_before = a.content_hash();
106          idft(module, res.data, sr, a.data, sa, tmp.data());
107          ASSERT_EQ(a.content_hash(), hash_before);
108        } else {
109          idft_tmp_a(module, res.data, sr, a.data, sa);
110        }
111        for (uint64_t i = 0; i < sr; ++i) {
112          znx_i64 actual = res.get_copy_zext(i);
113          // ASSERT_EQ(res.get_copy_zext(i), expect[i]);
114          if (!(actual == expect[i])) {
115            for (uint64_t j = 0; j < n; ++j) {
116              std::cerr << actual.get_coeff(j) << " dft vs. " << expect[i].get_coeff(j) << std::endl;
117            }
118            FAIL();
119          }
120        }
121      }
122    }
123    delete_module_info(module);
124  }
125}
126
127TEST(vec_znx_dft, fft64_vec_znx_idft) {
128  test_fft64_vec_znx_idft(fft64_vec_znx_idft, nullptr, fft64_vec_znx_idft_tmp_bytes);
129}
130TEST(vec_znx_dft, fft64_vec_znx_idft_tmp_a) { test_fft64_vec_znx_idft(nullptr, fft64_vec_znx_idft_tmp_a, nullptr); }
131
132#ifdef __x86_64__
133// FIXME: currently, it only works on avx
134static void test_ntt120_vec_znx_idft(VEC_ZNX_IDFT_F idft, VEC_ZNX_IDFT_TMP_A_F idft_tmp_a,
135                                     VEC_ZNX_IDFT_TMP_BYTES_F idft_tmp_bytes) {
136  for (uint64_t n : {2, 4, 64, 128}) {
137    MODULE* module = new_module_info(n, NTT120);
138    uint64_t tmp_size = idft_tmp_bytes ? idft_tmp_bytes(module) : 0;
139    std::vector<uint8_t> tmp(tmp_size);
140    for (uint64_t sa : {3, 5, 8}) {
141      for (uint64_t sr : {3, 5, 8}) {
142        ntt120_vec_znx_dft_layout a(n, sa);
143        ntt120_vec_znx_big_layout res(n, sr);
144        a.fill_random();
145        std::vector<znx_i128> expect(sr);
146        for (uint64_t i = 0; i < sr; ++i) {
147          expect[i] = simple_intt120(a.get_copy_zext(i));
148        }
149        // test the function
150        if (idft_tmp_bytes) {
151          thash hash_before = a.content_hash();
152          idft(module, res.data, sr, a.data, sa, tmp.data());
153          ASSERT_EQ(a.content_hash(), hash_before);
154        } else {
155          idft_tmp_a(module, res.data, sr, a.data, sa);
156        }
157        for (uint64_t i = 0; i < sr; ++i) {
158          znx_i128 actual = res.get_copy_zext(i);
159          ASSERT_EQ(res.get_copy_zext(i), expect[i]);
160          // if (!(actual == expect[i])) {
161          //   for (uint64_t j = 0; j < n; ++j) {
162          //     std::cerr << actual.get_coeff(j) << " dft vs. " << expect[i].get_coeff(j) << std::endl;
163          //   }
164          //   FAIL();
165          // }
166        }
167      }
168    }
169    delete_module_info(module);
170  }
171}
172
173TEST(vec_znx_dft, ntt120_vec_znx_idft) {
174  test_ntt120_vec_znx_idft(ntt120_vec_znx_idft_avx, nullptr, ntt120_vec_znx_idft_tmp_bytes_avx);
175}
176TEST(vec_znx_dft, ntt120_vec_znx_idft_tmp_a) {
177  test_ntt120_vec_znx_idft(nullptr, ntt120_vec_znx_idft_tmp_a_avx, nullptr);
178}
179#endif
180TEST(vec_znx_dft, vec_znx_idft) {
181  test_fft64_vec_znx_idft(vec_znx_idft, nullptr, vec_znx_idft_tmp_bytes);
182#ifdef __x86_64__
183  // FIXME: currently, only supported on avx
184  test_ntt120_vec_znx_idft(vec_znx_idft, nullptr, vec_znx_idft_tmp_bytes);
185#endif
186}
187TEST(vec_znx_dft, vec_znx_idft_tmp_a) {
188  test_fft64_vec_znx_idft(nullptr, vec_znx_idft_tmp_a, nullptr);
189#ifdef __x86_64__
190  // FIXME: currently, only supported on avx
191  test_ntt120_vec_znx_idft(nullptr, vec_znx_idft_tmp_a, nullptr);
192#endif
193}