1#include <gtest/gtest.h>
  2
  3#include "../spqlios/arithmetic/vec_znx_arithmetic_private.h"
  4#include "testlib/fft64_layouts.h"
  5#include "testlib/polynomial_vector.h"
  6
  7static void test_vmp_prepare_contiguous(VMP_PREPARE_CONTIGUOUS_F* prepare_contiguous,
  8                                        VMP_PREPARE_CONTIGUOUS_TMP_BYTES_F* tmp_bytes) {
  9  // tests when n < 8
 10  for (uint64_t nn : {2, 4}) {
 11    MODULE* module = new_module_info(nn, FFT64);
 12    for (uint64_t nrows : {1, 2, 5}) {
 13      for (uint64_t ncols : {2, 6, 7}) {
 14        znx_vec_i64_layout mat(nn, nrows * ncols, nn);
 15        fft64_vmp_pmat_layout pmat(nn, nrows, ncols);
 16        mat.fill_random(30);
 17        std::vector<uint8_t> tmp_space(fft64_vmp_prepare_contiguous_tmp_bytes(module, nrows, ncols));
 18        thash hash_before = mat.content_hash();
 19        prepare_contiguous(module, pmat.data, mat.data(), nrows, ncols, tmp_space.data());
 20        ASSERT_EQ(mat.content_hash(), hash_before);
 21        for (uint64_t row = 0; row < nrows; ++row) {
 22          for (uint64_t col = 0; col < ncols; ++col) {
 23            const double* pmatv = (double*)pmat.data + (col * nrows + row) * nn;
 24            reim_fft64vec tmp = simple_fft64(mat.get_copy(row * ncols + col));
 25            const double* tmpv = tmp.data();
 26            for (uint64_t i = 0; i < nn; ++i) {
 27              ASSERT_LE(abs(pmatv[i] - tmpv[i]), 1e-10);
 28            }
 29          }
 30        }
 31      }
 32    }
 33    delete_module_info(module);
 34  }
 35  // tests when n >= 8
 36  for (uint64_t nn : {8, 32}) {
 37    MODULE* module = new_module_info(nn, FFT64);
 38    uint64_t nblk = nn / 8;
 39    for (uint64_t nrows : {1, 2, 5}) {
 40      for (uint64_t ncols : {2, 6, 7}) {
 41        znx_vec_i64_layout mat(nn, nrows * ncols, nn);
 42        fft64_vmp_pmat_layout pmat(nn, nrows, ncols);
 43        mat.fill_random(30);
 44        std::vector<uint8_t> tmp_space(tmp_bytes(module, nrows, ncols));
 45        thash hash_before = mat.content_hash();
 46        prepare_contiguous(module, pmat.data, mat.data(), nrows, ncols, tmp_space.data());
 47        ASSERT_EQ(mat.content_hash(), hash_before);
 48        for (uint64_t row = 0; row < nrows; ++row) {
 49          for (uint64_t col = 0; col < ncols; ++col) {
 50            reim_fft64vec tmp = simple_fft64(mat.get_copy(row * ncols + col));
 51            for (uint64_t blk = 0; blk < nblk; ++blk) {
 52              reim4_elem expect = tmp.get_blk(blk);
 53              reim4_elem actual = pmat.get(row, col, blk);
 54              ASSERT_LE(infty_dist(actual, expect), 1e-10);
 55            }
 56          }
 57        }
 58      }
 59    }
 60    delete_module_info(module);
 61  }
 62}
 63
 64TEST(vec_znx, vmp_prepare_contiguous) {
 65  test_vmp_prepare_contiguous(vmp_prepare_contiguous, vmp_prepare_contiguous_tmp_bytes);
 66}
 67TEST(vec_znx, fft64_vmp_prepare_contiguous_ref) {
 68  test_vmp_prepare_contiguous(fft64_vmp_prepare_contiguous_ref, fft64_vmp_prepare_contiguous_tmp_bytes);
 69}
 70#ifdef __x86_64__
 71TEST(vec_znx, fft64_vmp_prepare_contiguous_avx) {
 72  test_vmp_prepare_contiguous(fft64_vmp_prepare_contiguous_avx, fft64_vmp_prepare_contiguous_tmp_bytes);
 73}
 74#endif
 75
 76static void test_vmp_apply(VMP_APPLY_DFT_TO_DFT_F* apply, VMP_APPLY_DFT_TO_DFT_TMP_BYTES_F* tmp_bytes) {
 77  for (uint64_t nn : {2, 4, 8, 64}) {
 78    MODULE* module = new_module_info(nn, FFT64);
 79    for (uint64_t mat_nrows : {1, 4, 7}) {
 80      for (uint64_t mat_ncols : {1, 2, 5}) {
 81        for (uint64_t in_size : {1, 4, 7}) {
 82          for (uint64_t out_size : {1, 2, 5}) {
 83            fft64_vec_znx_dft_layout in(nn, in_size);
 84            fft64_vmp_pmat_layout pmat(nn, mat_nrows, mat_ncols);
 85            fft64_vec_znx_dft_layout out(nn, out_size);
 86            in.fill_random(0);
 87            pmat.fill_random(0);
 88            // naive computation of the product
 89            std::vector<reim_fft64vec> expect(out_size, reim_fft64vec(nn));
 90            for (uint64_t col = 0; col < out_size; ++col) {
 91              reim_fft64vec ex = reim_fft64vec::zero(nn);
 92              for (uint64_t row = 0; row < std::min(mat_nrows, in_size); ++row) {
 93                ex += pmat.get_zext(row, col) * in.get_copy_zext(row);
 94              }
 95              expect[col] = ex;
 96            }
 97            // apply the product
 98            std::vector<uint8_t> tmp(tmp_bytes(module, out_size, in_size, mat_nrows, mat_ncols));
 99            apply(module, out.data, out_size, in.data, in_size, pmat.data, mat_nrows, mat_ncols, tmp.data());
100            // check that the output is close from the expectation
101            for (uint64_t col = 0; col < out_size; ++col) {
102              reim_fft64vec actual = out.get_copy_zext(col);
103              ASSERT_LE(infty_dist(actual, expect[col]), 1e-10);
104            }
105          }
106        }
107      }
108    }
109    delete_module_info(module);
110  }
111}
112
113TEST(vec_znx, vmp_apply_to_dft) { test_vmp_apply(vmp_apply_dft_to_dft, vmp_apply_dft_to_dft_tmp_bytes); }
114TEST(vec_znx, fft64_vmp_apply_dft_to_dft_ref) {
115  test_vmp_apply(fft64_vmp_apply_dft_to_dft_ref, fft64_vmp_apply_dft_to_dft_tmp_bytes);
116}
117#ifdef __x86_64__
118TEST(vec_znx, fft64_vmp_apply_dft_to_dft_avx) {
119  test_vmp_apply(fft64_vmp_apply_dft_to_dft_avx, fft64_vmp_apply_dft_to_dft_tmp_bytes);
120}
121#endif