1#include <string.h>
  2
  3#include "../commons_private.h"
  4#include "reim_fft_internal.h"
  5#include "reim_fft_private.h"
  6
  7// Computes X^i -> X^(p*i) in the fourier domain for a reim vector is size 2 * m * size
  8// This function cannot be evaluated in place.
  9EXPORT void reim_fftvec_automorphism_ref(const REIM_FFTVEC_AUTOMORPHISM_PRECOMP* tables, int64_t p, double* r,
 10                                         const double* a, uint64_t a_size) {
 11  const uint64_t m = tables->m;
 12  const uint64_t nn = 2 * m;
 13  const uint64_t* irev = tables->irev;
 14  const uint64_t mask = (4 * m - 1);
 15  const uint64_t conj = !((p & 3) == 1);
 16  p = p & mask;
 17  p *= 1 - (conj << 1);
 18  if (a_size == 1) {
 19    for (uint64_t i = 0; i < m; ++i) {
 20      uint64_t i_rev = 2 * irev[i] + 1;
 21      i_rev = (((p * i_rev) & mask) - 1) >> 1;
 22      uint64_t j = irev[i_rev];
 23      r[i] = a[j];
 24      double x = a[j + m];
 25      r[i + m] = conj ? -x : x;
 26    }
 27  } else {
 28    for (uint64_t i = 0; i < m; ++i) {
 29      uint64_t i_rev = 2 * irev[i] + 1;
 30      i_rev = (((p * i_rev) & mask) - 1) >> 1;
 31      uint64_t j = irev[i_rev];
 32      for (uint64_t k = 0; k < a_size; ++k) {
 33        uint64_t offset_re = k * nn;
 34        uint64_t offset_im = offset_re + m;
 35        r[i + offset_re] = a[j + offset_re];
 36        double x = a[j + offset_im];
 37        r[i + offset_im] = conj ? -x : x;
 38      }
 39    }
 40  }
 41}
 42
 43// Computes the permutation index for an automorphism X^{i} -> X^{i*p} in the fourier domain.
 44EXPORT void reim_fftvec_automorphism_lut_precomp_ref(const REIM_FFTVEC_AUTOMORPHISM_PRECOMP* tables, int64_t p,
 45                                                     uint64_t* precomp  // size m
 46) {
 47  const uint64_t m = tables->m;
 48  const uint64_t* irev = tables->irev;
 49  const uint64_t mask = (4 * m - 1);
 50  const uint64_t conj = !((p & 3) == 1);
 51  p = p & mask;
 52  p *= 1 - (conj << 1);
 53  for (uint64_t i = 0; i < m; ++i) {
 54    uint64_t i_rev = 2 * irev[i] + 1;
 55    i_rev = (((p * i_rev) & mask) - 1) >> 1;
 56    uint64_t j = irev[i_rev];
 57    precomp[i] = j;
 58  }
 59}
 60
 61// Computes X^{i} -> X^{i*p} in the fourier domain for a reim vector of size m using a precomputed lut permutation.
 62void reim_fftvec_automorphism_with_lut_ref(const REIM_FFTVEC_AUTOMORPHISM_PRECOMP* tables, uint64_t* precomp, double* r,
 63                                           const double* a) {
 64  const uint64_t m = tables->m;
 65  for (uint64_t i = 0; i < m; ++i) {
 66    uint64_t j = precomp[i];
 67    r[i] = a[j];
 68    r[i + m] = a[j + m];
 69  }
 70}
 71
 72// Computes X^{i} -> X^{i*-p} in the fourier domain for a reim vector of size m using a precomputed lut permutation.
 73void reim_fftvec_automorphism_conj_with_lut_ref(const REIM_FFTVEC_AUTOMORPHISM_PRECOMP* tables, uint64_t* precomp,
 74                                                double* r, const double* a) {
 75  const uint64_t m = tables->m;
 76  for (uint64_t i = 0; i < m; ++i) {
 77    uint64_t j = precomp[i];
 78    r[i] = a[j];
 79    r[i + m] = -a[j + m];
 80  }
 81}
 82
 83// Returns the minimum number of temporary bytes used by reim_fftvec_automorphism_inplace_ref.
 84EXPORT uint64_t reim_fftvec_automorphism_inplace_tmp_bytes_ref(const REIM_FFTVEC_AUTOMORPHISM_PRECOMP* tables) {
 85  const uint64_t m = tables->m;
 86  return m * (2 * sizeof(double) + sizeof(uint64_t));
 87}
 88
 89// Computes X^i -> X^(p*i) in the fourier domain for a reim vector is size 2 * m * a_size
 90// This function cannot be evaluated in place.
 91EXPORT void reim_fftvec_automorphism_inplace_ref(const REIM_FFTVEC_AUTOMORPHISM_PRECOMP* tables, int64_t p, double* a,
 92                                                 uint64_t a_size,
 93                                                 uint8_t* tmp_bytes  // m * (2*sizeof(double) + sizeof(uint64_t))
 94) {
 95  const uint64_t m = tables->m;
 96  const uint64_t nn = 2 * m;
 97  const uint64_t* irev = tables->irev;
 98  const uint64_t mask = (4 * m - 1);
 99  const uint64_t conj = !((p & 3) == 1);
100
101  double* tmp = (double*)tmp_bytes;
102  p = p & mask;
103  if (a_size == 1) {
104    p *= 1 - (conj << 1);
105    for (uint64_t i = 0; i < m; ++i) {
106      uint64_t i_rev = 2 * irev[i] + 1;
107      i_rev = (((p * i_rev) & mask) - 1) >> 1;
108      uint64_t j = irev[i_rev];
109      tmp[i] = a[j];
110      double x = a[j + m];
111      tmp[i + m] = conj ? -x : x;
112    }
113    memcpy(a, tmp, nn * sizeof(double));
114  } else {
115    uint64_t* lut = (uint64_t*)(tmp_bytes + nn * sizeof(double));
116    reim_fftvec_automorphism_lut_precomp_ref(tables, p, lut);
117    for (uint64_t i = 0; i < a_size; ++i) {
118      if (conj == 1) {
119        reim_fftvec_automorphism_conj_with_lut_ref(tables, lut, tmp, a + i * nn);
120      } else {
121        reim_fftvec_automorphism_with_lut_ref(tables, lut, tmp, a + i * nn);
122      }
123      memcpy(a + i * nn, tmp, nn * sizeof(double));
124    }
125  }
126}
127
128EXPORT void reim_fftvec_addmul_ref(const REIM_FFTVEC_ADDMUL_PRECOMP* tables, double* r, const double* a,
129                                   const double* b) {
130  const uint64_t m = tables->m;
131  for (uint64_t i = 0; i < m; ++i) {
132    double re = a[i] * b[i] - a[i + m] * b[i + m];
133    double im = a[i] * b[i + m] + a[i + m] * b[i];
134    r[i] += re;
135    r[i + m] += im;
136  }
137}
138
139EXPORT void reim_fftvec_add_ref(const REIM_FFTVEC_ADD_PRECOMP* precomp, double* r, const double* a, const double* b) {
140  const uint64_t m = precomp->m;
141  for (uint64_t i = 0; i < m; ++i) {
142    double re = a[i] + b[i];
143    double im = a[i + m] + b[i + m];
144    r[i] = re;
145    r[i + m] = im;
146  }
147}
148
149EXPORT void reim_fftvec_sub_ref(const REIM_FFTVEC_SUB_PRECOMP* precomp, double* r, const double* a, const double* b) {
150  const uint64_t m = precomp->m;
151  for (uint64_t i = 0; i < m; ++i) {
152    double re = a[i] - b[i];
153    double im = a[i + m] - b[i + m];
154    r[i] = re;
155    r[i + m] = im;
156  }
157}
158
159EXPORT void reim_fftvec_mul_ref(const REIM_FFTVEC_MUL_PRECOMP* tables, double* r, const double* a, const double* b) {
160  const uint64_t m = tables->m;
161  for (uint64_t i = 0; i < m; ++i) {
162    double re = a[i] * b[i] - a[i + m] * b[i + m];
163    double im = a[i] * b[i + m] + a[i + m] * b[i];
164    r[i] = re;
165    r[i + m] = im;
166  }
167}
168
169EXPORT REIM_FFTVEC_ADDMUL_PRECOMP* new_reim_fftvec_addmul_precomp(uint32_t m) {
170  REIM_FFTVEC_ADDMUL_PRECOMP* reps = malloc(sizeof(REIM_FFTVEC_ADDMUL_PRECOMP));
171  reps->m = m;
172  if (CPU_SUPPORTS("fma")) {
173    if (m >= 4) {
174      reps->function = reim_fftvec_addmul_fma;
175    } else {
176      reps->function = reim_fftvec_addmul_ref;
177    }
178  } else {
179    reps->function = reim_fftvec_addmul_ref;
180  }
181  return reps;
182}
183
184EXPORT REIM_FFTVEC_ADD_PRECOMP* new_reim_fftvec_add_precomp(uint32_t m) {
185  REIM_FFTVEC_ADD_PRECOMP* reps = malloc(sizeof(REIM_FFTVEC_ADD_PRECOMP));
186  reps->m = m;
187  if (CPU_SUPPORTS("fma")) {
188    if (m >= 4) {
189      reps->function = reim_fftvec_add_fma;
190    } else {
191      reps->function = reim_fftvec_add_ref;
192    }
193  } else {
194    reps->function = reim_fftvec_add_ref;
195  }
196  return reps;
197}
198
199EXPORT REIM_FFTVEC_SUB_PRECOMP* new_reim_fftvec_sub_precomp(uint32_t m) {
200  REIM_FFTVEC_SUB_PRECOMP* reps = malloc(sizeof(REIM_FFTVEC_SUB_PRECOMP));
201  reps->m = m;
202  if (CPU_SUPPORTS("fma")) {
203    if (m >= 4) {
204      reps->function = reim_fftvec_sub_fma;
205    } else {
206      reps->function = reim_fftvec_sub_ref;
207    }
208  } else {
209    reps->function = reim_fftvec_sub_ref;
210  }
211  return reps;
212}
213
214EXPORT REIM_FFTVEC_MUL_PRECOMP* new_reim_fftvec_mul_precomp(uint32_t m) {
215  REIM_FFTVEC_MUL_PRECOMP* reps = malloc(sizeof(REIM_FFTVEC_MUL_PRECOMP));
216  reps->m = m;
217  if (CPU_SUPPORTS("fma")) {
218    if (m >= 4) {
219      reps->function = reim_fftvec_mul_fma;
220    } else {
221      reps->function = reim_fftvec_mul_ref;
222    }
223  } else {
224    reps->function = reim_fftvec_mul_ref;
225  }
226  return reps;
227}
228
229EXPORT REIM_FFTVEC_AUTOMORPHISM_PRECOMP* new_reim_fftvec_automorphism_precomp(uint32_t m) {
230  REIM_FFTVEC_AUTOMORPHISM_PRECOMP* reps = malloc(sizeof(REIM_FFTVEC_AUTOMORPHISM_PRECOMP));
231  reps->m = m;
232  reps->function.apply = reim_fftvec_automorphism_ref;
233  reps->function.apply_inplace = reim_fftvec_automorphism_inplace_ref;
234  reps->function.apply_inplace_tmp_bytes = reim_fftvec_automorphism_inplace_tmp_bytes_ref;
235  const uint64_t nn = 2 * m;
236  reps->irev = malloc(sizeof(uint64_t) * nn);
237  uint32_t lognn = log2m(nn);
238  for (uint32_t i = 0; i < nn; i++) {
239    reps->irev[i] = (uint64_t)revbits(lognn, i);
240  }
241  return reps;
242}