arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_conversions.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_conversions.cRaw
  1#include <errno.h>
  2#include <string.h>
  3
  4#include "../commons_private.h"
  5#include "reim_fft_internal.h"
  6#include "reim_fft_private.h"
  7
  8void reim_from_znx64_ref(const REIM_FROM_ZNX64_PRECOMP* precomp, void* r, const int64_t* x) {
  9  // naive version of the function (just cast)
 10  const uint64_t nn = precomp->m << 1;
 11  double* res = (double*)r;
 12  for (uint64_t i = 0; i < nn; ++i) {
 13    res[i] = (double)x[i];
 14  }
 15}
 16
 17void* init_reim_from_znx64_precomp(REIM_FROM_ZNX64_PRECOMP* const res, uint32_t m, uint32_t log2bound) {
 18  if (m & (m - 1)) return spqlios_error("m must be a power of 2");
 19  // currently, we are going to use the trick add 3.2^51, mask the exponent, reinterpret bits as double.
 20  // therefore we need the input values to be < 2^50.
 21  if (log2bound > 50) return spqlios_error("Invalid log2bound error: must be in [0,50]");
 22  res->m = m;
 23  FROM_ZNX64_FUNC resf = reim_from_znx64_ref;
 24  if (m >= 8) {
 25    if (CPU_SUPPORTS("avx2")) {
 26      resf = reim_from_znx64_bnd50_fma;
 27    }
 28  }
 29  res->function = resf;
 30  return res;
 31}
 32
 33EXPORT REIM_FROM_ZNX64_PRECOMP* new_reim_from_znx64_precomp(uint32_t m, uint32_t log2bound) {
 34  REIM_FROM_ZNX64_PRECOMP* res = malloc(sizeof(*res));
 35  if (!res) return spqlios_error(strerror(errno));
 36  return spqlios_keep_or_free(res, init_reim_from_znx64_precomp(res, m, log2bound));
 37}
 38
 39EXPORT void reim_from_znx64(const REIM_FROM_ZNX64_PRECOMP* tables, void* r, const int64_t* a) {
 40  tables->function(tables, r, a);
 41}
 42
 43/**
 44 * @brief Simpler API for the znx64 to reim conversion.
 45 * For each dimension, the precomputed tables for this dimension are generated automatically the first time.
 46 * It is advised to do one dry-run call per desired dimension before using in a multithread environment
 47 */
 48EXPORT void reim_from_znx64_simple(uint32_t m, uint32_t log2bound, void* r, const int64_t* a) {
 49  // not checking for log2bound which is not relevant here
 50  static REIM_FROM_ZNX64_PRECOMP precomp[32];
 51  REIM_FROM_ZNX64_PRECOMP* p = precomp + log2m(m);
 52  if (!p->function) {
 53    if (!init_reim_from_znx64_precomp(p, m, log2bound)) abort();
 54  }
 55  p->function(p, r, a);
 56}
 57
 58void reim_from_znx32_ref(const REIM_FROM_ZNX32_PRECOMP* precomp, void* r, const int32_t* x) { NOT_IMPLEMENTED(); }
 59void reim_from_znx32_avx2_fma(const REIM_FROM_ZNX32_PRECOMP* precomp, void* r, const int32_t* x) { NOT_IMPLEMENTED(); }
 60
 61void* init_reim_from_znx32_precomp(REIM_FROM_ZNX32_PRECOMP* const res, uint32_t m, uint32_t log2bound) {
 62  if (m & (m - 1)) return spqlios_error("m must be a power of 2");
 63  if (log2bound > 32) return spqlios_error("Invalid log2bound error: must be in [0,32]");
 64  res->m = m;
 65  // TODO: check selection logic
 66  if (CPU_SUPPORTS("avx2")) {
 67    if (m >= 8) {
 68      res->function = reim_from_znx32_avx2_fma;
 69    } else {
 70      res->function = reim_from_znx32_ref;
 71    }
 72  } else {
 73    res->function = reim_from_znx32_ref;
 74  }
 75  return res;
 76}
 77
 78EXPORT REIM_FROM_ZNX32_PRECOMP* new_reim_from_znx32_precomp(uint32_t m, uint32_t log2bound) {
 79  REIM_FROM_ZNX32_PRECOMP* res = malloc(sizeof(*res));
 80  if (!res) return spqlios_error(strerror(errno));
 81  return spqlios_keep_or_free(res, init_reim_from_znx32_precomp(res, m, log2bound));
 82}
 83
 84void reim_from_tnx32_ref(const REIM_FROM_TNX32_PRECOMP* precomp, void* r, const int32_t* x) { NOT_IMPLEMENTED(); }
 85void reim_from_tnx32_avx2_fma(const REIM_FROM_TNX32_PRECOMP* precomp, void* r, const int32_t* x) { NOT_IMPLEMENTED(); }
 86
 87void* init_reim_from_tnx32_precomp(REIM_FROM_TNX32_PRECOMP* const res, uint32_t m) {
 88  if (m & (m - 1)) return spqlios_error("m must be a power of 2");
 89  res->m = m;
 90  // TODO: check selection logic
 91  if (CPU_SUPPORTS("avx2")) {
 92    if (m >= 8) {
 93      res->function = reim_from_tnx32_avx2_fma;
 94    } else {
 95      res->function = reim_from_tnx32_ref;
 96    }
 97  } else {
 98    res->function = reim_from_tnx32_ref;
 99  }
100  return res;
101}
102
103EXPORT REIM_FROM_TNX32_PRECOMP* new_reim_from_tnx32_precomp(uint32_t m) {
104  REIM_FROM_TNX32_PRECOMP* res = malloc(sizeof(*res));
105  if (!res) return spqlios_error(strerror(errno));
106  return spqlios_keep_or_free(res, init_reim_from_tnx32_precomp(res, m));
107}
108
109void reim_to_tnx32_ref(const REIM_TO_TNX32_PRECOMP* precomp, int32_t* r, const void* x) { NOT_IMPLEMENTED(); }
110void reim_to_tnx32_avx2_fma(const REIM_TO_TNX32_PRECOMP* precomp, int32_t* r, const void* x) { NOT_IMPLEMENTED(); }
111
112void* init_reim_to_tnx32_precomp(REIM_TO_TNX32_PRECOMP* const res, uint32_t m, double divisor, uint32_t log2overhead) {
113  if (m & (m - 1)) return spqlios_error("m must be a power of 2");
114  if (is_not_pow2_double(&divisor)) return spqlios_error("divisor must be a power of 2");
115  if (log2overhead > 52) return spqlios_error("log2overhead is too large");
116  res->m = m;
117  res->divisor = divisor;
118  // TODO: check selection logic
119  if (CPU_SUPPORTS("avx2")) {
120    if (log2overhead <= 18) {
121      if (m >= 8) {
122        res->function = reim_to_tnx32_avx2_fma;
123      } else {
124        res->function = reim_to_tnx32_ref;
125      }
126    } else {
127      res->function = reim_to_tnx32_ref;
128    }
129  } else {
130    res->function = reim_to_tnx32_ref;
131  }
132  return res;
133}
134
135EXPORT REIM_TO_TNX32_PRECOMP* new_reim_to_tnx32_precomp(uint32_t m, double divisor, uint32_t log2overhead) {
136  REIM_TO_TNX32_PRECOMP* res = malloc(sizeof(*res));
137  if (!res) return spqlios_error(strerror(errno));
138  return spqlios_keep_or_free(res, init_reim_to_tnx32_precomp(res, m, divisor, log2overhead));
139}
140
141EXPORT void reim_from_znx32_simple(uint32_t m, uint32_t log2bound, void* r, const int32_t* x) {
142  static REIM_FROM_ZNX32_PRECOMP* p[31] = {0};
143  REIM_FROM_ZNX32_PRECOMP** f = p + log2m(m);
144  if (!*f) *f = new_reim_from_znx32_precomp(m, log2bound);
145  (*f)->function(*f, r, x);
146}
147
148EXPORT void reim_from_tnx32_simple(uint32_t m, void* r, const int32_t* x) {
149  static REIM_FROM_TNX32_PRECOMP* p[31] = {0};
150  REIM_FROM_TNX32_PRECOMP** f = p + log2m(m);
151  if (!*f) *f = new_reim_from_tnx32_precomp(m);
152  (*f)->function(*f, r, x);
153}
154
155EXPORT void reim_to_tnx32_simple(uint32_t m, double divisor, uint32_t log2overhead, int32_t* r, const void* x) {
156  static REIM_TO_TNX32_PRECOMP* p[31] = {0};
157  REIM_TO_TNX32_PRECOMP** f = p + log2m(m);
158  if (!*f) *f = new_reim_to_tnx32_precomp(m, divisor, log2overhead);
159  (*f)->function(*f, r, x);
160}
161
162void reim_to_znx64_ref(const REIM_TO_ZNX64_PRECOMP* precomp, int64_t* r, const void* x) {
163  // for now, we stick to a slow implem
164  uint64_t nn = precomp->m << 1;
165  const double* v = (double*)x;
166  double invdiv = 1. / precomp->divisor;
167  for (uint64_t i = 0; i < nn; ++i) {
168    r[i] = (int64_t)rint(v[i] * invdiv);
169  }
170}
171
172void* init_reim_to_znx64_precomp(REIM_TO_ZNX64_PRECOMP* const res, uint32_t m, double divisor, uint32_t log2bound) {
173  if (m & (m - 1)) return spqlios_error("m must be a power of 2");
174  if (is_not_pow2_double(&divisor)) return spqlios_error("divisor must be a power of 2");
175  if (log2bound > 64) return spqlios_error("log2bound is too large");
176  res->m = m;
177  res->divisor = divisor;
178  TO_ZNX64_FUNC resf = reim_to_znx64_ref;
179  if (CPU_SUPPORTS("avx2") && m >= 8) {
180    if (log2bound <= 50) {
181      resf = reim_to_znx64_avx2_bnd50_fma;
182    } else {
183      resf = reim_to_znx64_avx2_bnd63_fma;
184    }
185  }
186  res->function = resf;  // must be the last one set
187  return res;
188}
189
190EXPORT REIM_TO_ZNX64_PRECOMP* new_reim_to_znx64_precomp(uint32_t m, double divisor, uint32_t log2bound) {
191  REIM_TO_ZNX64_PRECOMP* res = malloc(sizeof(*res));
192  if (!res) return spqlios_error(strerror(errno));
193  return spqlios_keep_or_free(res, init_reim_to_znx64_precomp(res, m, divisor, log2bound));
194}
195
196EXPORT void reim_to_znx64(const REIM_TO_ZNX64_PRECOMP* precomp, int64_t* r, const void* a) {
197  precomp->function(precomp, r, a);
198}
199
200/**
201 * @brief Simpler API for the znx64 to reim conversion.
202 */
203EXPORT void reim_to_znx64_simple(uint32_t m, double divisor, uint32_t log2bound, int64_t* r, const void* a) {
204  // not checking distinguishing <=50 or not
205  static __thread REIM_TO_ZNX64_PRECOMP p;
206  static __thread uint32_t prev_log2bound;
207  if (!p.function || p.m != m || p.divisor != divisor || prev_log2bound != log2bound) {
208    if (!init_reim_to_znx64_precomp(&p, m, divisor, log2bound)) abort();
209    prev_log2bound = log2bound;
210  }
211  p.function(&p, r, a);
212}