1#include <string.h>
  2
  3#include "../q120/q120_arithmetic.h"
  4#include "vec_znx_arithmetic_private.h"
  5
  6EXPORT void vec_znx_dft(const MODULE* module,                             // N
  7                        VEC_ZNX_DFT* res, uint64_t res_size,              // res
  8                        const int64_t* a, uint64_t a_size, uint64_t a_sl  // a
  9) {
 10  return module->func.vec_znx_dft(module, res, res_size, a, a_size, a_sl);
 11}
 12
 13EXPORT void vec_dft_add(const MODULE* module,                   // N
 14                        VEC_ZNX_DFT* res, uint64_t res_size,    // res
 15                        const VEC_ZNX_DFT* a, uint64_t a_size,  // a
 16                        const VEC_ZNX_DFT* b, uint64_t b_size   // b
 17) {
 18  return module->func.vec_dft_add(module, res, res_size, a, a_size, b, b_size);
 19}
 20
 21EXPORT void vec_dft_sub(const MODULE* module,                   // N
 22                        VEC_ZNX_DFT* res, uint64_t res_size,    // res
 23                        const VEC_ZNX_DFT* a, uint64_t a_size,  // a
 24                        const VEC_ZNX_DFT* b, uint64_t b_size   // b
 25) {
 26  return module->func.vec_dft_sub(module, res, res_size, a, a_size, b, b_size);
 27}
 28
 29EXPORT void vec_znx_idft(const MODULE* module,                       // N
 30                         VEC_ZNX_BIG* res, uint64_t res_size,        // res
 31                         const VEC_ZNX_DFT* a_dft, uint64_t a_size,  // a
 32                         uint8_t* tmp                                // scratch space
 33) {
 34  return module->func.vec_znx_idft(module, res, res_size, a_dft, a_size, tmp);
 35}
 36
 37EXPORT uint64_t vec_znx_idft_tmp_bytes(const MODULE* module, uint64_t nn) { return module->func.vec_znx_idft_tmp_bytes(module, nn); }
 38
 39EXPORT void vec_znx_idft_tmp_a(const MODULE* module,                 // N
 40                               VEC_ZNX_BIG* res, uint64_t res_size,  // res
 41                               VEC_ZNX_DFT* a_dft, uint64_t a_size   // a is overwritten
 42) {
 43  return module->func.vec_znx_idft_tmp_a(module, res, res_size, a_dft, a_size);
 44}
 45
 46EXPORT uint64_t bytes_of_vec_znx_dft(const MODULE* module,  // N
 47                                     uint64_t size) {
 48  return module->func.bytes_of_vec_znx_dft(module, size);
 49}
 50
 51// fft64 backend
 52EXPORT uint64_t fft64_bytes_of_vec_znx_dft(const MODULE* module,  // N
 53                                           uint64_t size) {
 54  return module->nn * size * sizeof(double);
 55}
 56
 57EXPORT VEC_ZNX_DFT* new_vec_znx_dft(const MODULE* module,  // N
 58                                    uint64_t size) {
 59  return spqlios_alloc(bytes_of_vec_znx_dft(module, size));
 60}
 61
 62EXPORT void delete_vec_znx_dft(VEC_ZNX_DFT* res) { spqlios_free(res); }
 63
 64EXPORT void fft64_vec_znx_dft(const MODULE* module,                             // N
 65                              VEC_ZNX_DFT* res, uint64_t res_size,              // res
 66                              const int64_t* a, uint64_t a_size, uint64_t a_sl  // a
 67) {
 68  const uint64_t smin = res_size < a_size ? res_size : a_size;
 69  const uint64_t nn = module->nn;
 70
 71  for (uint64_t i = 0; i < smin; i++) {
 72    reim_from_znx64(module->mod.fft64.p_conv, ((double*)res) + i * nn, a + i * a_sl);
 73    reim_fft(module->mod.fft64.p_fft, ((double*)res) + i * nn);
 74  }
 75
 76  // fill up remaining part with 0's
 77  double* const dres = (double*)res;
 78  memset(dres + smin * nn, 0, (res_size - smin) * nn * sizeof(double));
 79}
 80
 81EXPORT void fft64_vec_dft_add(const MODULE* module,                   // N
 82                              VEC_ZNX_DFT* res, uint64_t res_size,    // res
 83                              const VEC_ZNX_DFT* a, uint64_t a_size,  // a
 84                              const VEC_ZNX_DFT* b, uint64_t b_size   // b
 85) {
 86  const uint64_t smin0 = a_size < b_size ? a_size : b_size;
 87  const uint64_t smin = res_size < smin0 ? res_size : smin0;
 88  const uint64_t nn = module->nn;
 89
 90  for (uint64_t i = 0; i < smin; i++) {
 91    reim_fftvec_add(module->mod.fft64.add_fft, ((double*)res) + i * nn, ((double*)a) + i * nn, ((double*)b) + i * nn);
 92  }
 93
 94  // fill remain `res` part with 0's
 95  double* const dres = (double*)res;
 96  memset(dres + smin * nn, 0, (res_size - smin) * nn * sizeof(double));
 97}
 98
 99EXPORT void fft64_vec_dft_sub(const MODULE* module,                   // N
100                              VEC_ZNX_DFT* res, uint64_t res_size,    // res
101                              const VEC_ZNX_DFT* a, uint64_t a_size,  // a
102                              const VEC_ZNX_DFT* b, uint64_t b_size   // b
103) {
104  const uint64_t smin0 = a_size < b_size ? a_size : b_size;
105  const uint64_t smin = res_size < smin0 ? res_size : smin0;
106  const uint64_t nn = module->nn;
107
108  for (uint64_t i = 0; i < smin; i++) {
109    reim_fftvec_sub(module->mod.fft64.sub_fft, ((double*)res) + i * nn, ((double*)a) + i * nn, ((double*)b) + i * nn);
110  }
111
112  // fill remain `res` part with 0's
113  double* const dres = (double*)res;
114  memset(dres + smin * nn, 0, (res_size - smin) * nn * sizeof(double));
115}
116
117EXPORT void fft64_vec_znx_idft(const MODULE* module,                       // N
118                               VEC_ZNX_BIG* res, uint64_t res_size,        // res
119                               const VEC_ZNX_DFT* a_dft, uint64_t a_size,  // a
120                               uint8_t* tmp                                // unused
121) {
122  const uint64_t nn = module->nn;
123  const uint64_t smin = res_size < a_size ? res_size : a_size;
124  if ((double*)res != (double*)a_dft) {
125    memcpy(res, a_dft, smin * nn * sizeof(double));
126  }
127
128  for (uint64_t i = 0; i < smin; i++) {
129    reim_ifft(module->mod.fft64.p_ifft, ((double*)res) + i * nn);
130    reim_to_znx64(module->mod.fft64.p_reim_to_znx, ((int64_t*)res) + i * nn, ((int64_t*)res) + i * nn);
131  }
132
133  // fill up remaining part with 0's
134  int64_t* const dres = (int64_t*)res;
135  memset(dres + smin * nn, 0, (res_size - smin) * nn * sizeof(double));
136}
137
138EXPORT uint64_t fft64_vec_znx_idft_tmp_bytes(const MODULE* module, uint64_t nn) { return 0; }
139
140EXPORT void fft64_vec_znx_idft_tmp_a(const MODULE* module,                 // N
141                                     VEC_ZNX_BIG* res, uint64_t res_size,  // res
142                                     VEC_ZNX_DFT* a_dft, uint64_t a_size   // a is overwritten
143) {
144  const uint64_t nn = module->nn;
145  const uint64_t smin = res_size < a_size ? res_size : a_size;
146
147  int64_t* const tres = (int64_t*)res;
148  double* const ta = (double*)a_dft;
149  for (uint64_t i = 0; i < smin; i++) {
150    reim_ifft(module->mod.fft64.p_ifft, ta + i * nn);
151    reim_to_znx64(module->mod.fft64.p_reim_to_znx, tres + i * nn, ta + i * nn);
152  }
153
154  // fill up remaining part with 0's
155  memset(tres + smin * nn, 0, (res_size - smin) * nn * sizeof(double));
156}
157
158// ntt120 backend
159
160EXPORT void ntt120_vec_znx_dft_avx(const MODULE* module,                             // N
161                                   VEC_ZNX_DFT* res, uint64_t res_size,              // res
162                                   const int64_t* a, uint64_t a_size, uint64_t a_sl  // a
163) {
164  const uint64_t nn = module->nn;
165  const uint64_t smin = res_size < a_size ? res_size : a_size;
166
167  int64_t* tres = (int64_t*)res;
168  for (uint64_t i = 0; i < smin; i++) {
169    q120_b_from_znx64_simple(nn, (q120b*)(tres + i * nn * 4), a + i * a_sl);
170    q120_ntt_bb_avx2(module->mod.q120.p_ntt, (q120b*)(tres + i * nn * 4));
171  }
172
173  // fill up remaining part with 0's
174  memset(tres + smin * nn * 4, 0, (res_size - smin) * nn * 4 * sizeof(int64_t));
175}
176
177EXPORT void ntt120_vec_znx_idft_avx(const MODULE* module,                       // N
178                                    VEC_ZNX_BIG* res, uint64_t res_size,        // res
179                                    const VEC_ZNX_DFT* a_dft, uint64_t a_size,  // a
180                                    uint8_t* tmp) {
181  const uint64_t nn = module->nn;
182  const uint64_t smin = res_size < a_size ? res_size : a_size;
183
184  __int128_t* const tres = (__int128_t*)res;
185  const int64_t* const ta = (int64_t*)a_dft;
186  for (uint64_t i = 0; i < smin; i++) {
187    memcpy(tmp, ta + i * nn * 4, nn * 4 * sizeof(uint64_t));
188    q120_intt_bb_avx2(module->mod.q120.p_intt, (q120b*)tmp);
189    q120_b_to_znx128_simple(nn, tres + i * nn, (q120b*)tmp);
190  }
191
192  // fill up remaining part with 0's
193  memset(tres + smin * nn, 0, (res_size - smin) * nn * sizeof(*tres));
194}
195
196EXPORT uint64_t ntt120_vec_znx_idft_tmp_bytes_avx(const MODULE* module, uint64_t nn) { return nn * 4 * sizeof(uint64_t); }
197
198EXPORT void ntt120_vec_znx_idft_tmp_a_avx(const MODULE* module,                 // N
199                                          VEC_ZNX_BIG* res, uint64_t res_size,  // res
200                                          VEC_ZNX_DFT* a_dft, uint64_t a_size   // a is overwritten
201) {
202  const uint64_t nn = module->nn;
203  const uint64_t smin = res_size < a_size ? res_size : a_size;
204
205  __int128_t* const tres = (__int128_t*)res;
206  int64_t* const ta = (int64_t*)a_dft;
207  for (uint64_t i = 0; i < smin; i++) {
208    q120_intt_bb_avx2(module->mod.q120.p_intt, (q120b*)(ta + i * nn * 4));
209    q120_b_to_znx128_simple(nn, tres + i * nn, (q120b*)(ta + i * nn * 4));
210  }
211
212  // fill up remaining part with 0's
213  memset(tres + smin * nn, 0, (res_size - smin) * nn * sizeof(*tres));
214}