1#include <immintrin.h>
  2#include <string.h>
  3
  4#include "vec_rnx_arithmetic_private.h"
  5
  6void rnx_add_avx(uint64_t nn, double* res, const double* a, const double* b) {
  7  if (nn < 8) {
  8    if (nn == 4) {
  9      _mm256_storeu_pd(res, _mm256_add_pd(_mm256_loadu_pd(a), _mm256_loadu_pd(b)));
 10    } else if (nn == 2) {
 11      _mm_storeu_pd(res, _mm_add_pd(_mm_loadu_pd(a), _mm_loadu_pd(b)));
 12    } else if (nn == 1) {
 13      *res = *a + *b;
 14    } else {
 15      NOT_SUPPORTED();  // not a power of 2
 16    }
 17    return;
 18  }
 19  // general case: nn >= 8
 20  __m256d x0, x1, x2, x3, x4, x5;
 21  const double* aa = a;
 22  const double* bb = b;
 23  double* rr = res;
 24  double* const rrend = res + nn;
 25  do {
 26    x0 = _mm256_loadu_pd(aa);
 27    x1 = _mm256_loadu_pd(aa + 4);
 28    x2 = _mm256_loadu_pd(bb);
 29    x3 = _mm256_loadu_pd(bb + 4);
 30    x4 = _mm256_add_pd(x0, x2);
 31    x5 = _mm256_add_pd(x1, x3);
 32    _mm256_storeu_pd(rr, x4);
 33    _mm256_storeu_pd(rr + 4, x5);
 34    aa += 8;
 35    bb += 8;
 36    rr += 8;
 37  } while (rr < rrend);
 38}
 39
 40void rnx_sub_avx(uint64_t nn, double* res, const double* a, const double* b) {
 41  if (nn < 8) {
 42    if (nn == 4) {
 43      _mm256_storeu_pd(res, _mm256_sub_pd(_mm256_loadu_pd(a), _mm256_loadu_pd(b)));
 44    } else if (nn == 2) {
 45      _mm_storeu_pd(res, _mm_sub_pd(_mm_loadu_pd(a), _mm_loadu_pd(b)));
 46    } else if (nn == 1) {
 47      *res = *a - *b;
 48    } else {
 49      NOT_SUPPORTED();  // not a power of 2
 50    }
 51    return;
 52  }
 53  // general case: nn >= 8
 54  __m256d x0, x1, x2, x3, x4, x5;
 55  const double* aa = a;
 56  const double* bb = b;
 57  double* rr = res;
 58  double* const rrend = res + nn;
 59  do {
 60    x0 = _mm256_loadu_pd(aa);
 61    x1 = _mm256_loadu_pd(aa + 4);
 62    x2 = _mm256_loadu_pd(bb);
 63    x3 = _mm256_loadu_pd(bb + 4);
 64    x4 = _mm256_sub_pd(x0, x2);
 65    x5 = _mm256_sub_pd(x1, x3);
 66    _mm256_storeu_pd(rr, x4);
 67    _mm256_storeu_pd(rr + 4, x5);
 68    aa += 8;
 69    bb += 8;
 70    rr += 8;
 71  } while (rr < rrend);
 72}
 73
 74void rnx_negate_avx(uint64_t nn, double* res, const double* b) {
 75  if (nn < 8) {
 76    if (nn == 4) {
 77      _mm256_storeu_pd(res, _mm256_sub_pd(_mm256_set1_pd(0), _mm256_loadu_pd(b)));
 78    } else if (nn == 2) {
 79      _mm_storeu_pd(res, _mm_sub_pd(_mm_set1_pd(0), _mm_loadu_pd(b)));
 80    } else if (nn == 1) {
 81      *res = -*b;
 82    } else {
 83      NOT_SUPPORTED();  // not a power of 2
 84    }
 85    return;
 86  }
 87  // general case: nn >= 8
 88  __m256d x2, x3, x4, x5;
 89  const __m256d ZERO = _mm256_set1_pd(0);
 90  const double* bb = b;
 91  double* rr = res;
 92  double* const rrend = res + nn;
 93  do {
 94    x2 = _mm256_loadu_pd(bb);
 95    x3 = _mm256_loadu_pd(bb + 4);
 96    x4 = _mm256_sub_pd(ZERO, x2);
 97    x5 = _mm256_sub_pd(ZERO, x3);
 98    _mm256_storeu_pd(rr, x4);
 99    _mm256_storeu_pd(rr + 4, x5);
100    bb += 8;
101    rr += 8;
102  } while (rr < rrend);
103}
104
105/** @brief sets res = a + b */
106EXPORT void vec_rnx_add_avx(                          //
107    const MOD_RNX* module,                            // N
108    double* res, uint64_t res_size, uint64_t res_sl,  // res
109    const double* a, uint64_t a_size, uint64_t a_sl,  // a
110    const double* b, uint64_t b_size, uint64_t b_sl   // b
111) {
112  const uint64_t nn = module->n;
113  if (a_size < b_size) {
114    const uint64_t msize = res_size < a_size ? res_size : a_size;
115    const uint64_t nsize = res_size < b_size ? res_size : b_size;
116    for (uint64_t i = 0; i < msize; ++i) {
117      rnx_add_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
118    }
119    for (uint64_t i = msize; i < nsize; ++i) {
120      memcpy(res + i * res_sl, b + i * b_sl, nn * sizeof(double));
121    }
122    for (uint64_t i = nsize; i < res_size; ++i) {
123      memset(res + i * res_sl, 0, nn * sizeof(double));
124    }
125  } else {
126    const uint64_t msize = res_size < b_size ? res_size : b_size;
127    const uint64_t nsize = res_size < a_size ? res_size : a_size;
128    for (uint64_t i = 0; i < msize; ++i) {
129      rnx_add_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
130    }
131    for (uint64_t i = msize; i < nsize; ++i) {
132      memcpy(res + i * res_sl, a + i * a_sl, nn * sizeof(double));
133    }
134    for (uint64_t i = nsize; i < res_size; ++i) {
135      memset(res + i * res_sl, 0, nn * sizeof(double));
136    }
137  }
138}
139
140/** @brief sets res = -a */
141EXPORT void vec_rnx_negate_avx(                       //
142    const MOD_RNX* module,                            // N
143    double* res, uint64_t res_size, uint64_t res_sl,  // res
144    const double* a, uint64_t a_size, uint64_t a_sl   // a
145) {
146  const uint64_t nn = module->n;
147  const uint64_t msize = res_size < a_size ? res_size : a_size;
148  for (uint64_t i = 0; i < msize; ++i) {
149    rnx_negate_avx(nn, res + i * res_sl, a + i * a_sl);
150  }
151  for (uint64_t i = msize; i < res_size; ++i) {
152    memset(res + i * res_sl, 0, nn * sizeof(double));
153  }
154}
155
156/** @brief sets res = a - b */
157EXPORT void vec_rnx_sub_avx(                          //
158    const MOD_RNX* module,                            // N
159    double* res, uint64_t res_size, uint64_t res_sl,  // res
160    const double* a, uint64_t a_size, uint64_t a_sl,  // a
161    const double* b, uint64_t b_size, uint64_t b_sl   // b
162) {
163  const uint64_t nn = module->n;
164  if (a_size < b_size) {
165    const uint64_t msize = res_size < a_size ? res_size : a_size;
166    const uint64_t nsize = res_size < b_size ? res_size : b_size;
167    for (uint64_t i = 0; i < msize; ++i) {
168      rnx_sub_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
169    }
170    for (uint64_t i = msize; i < nsize; ++i) {
171      rnx_negate_avx(nn, res + i * res_sl, b + i * b_sl);
172    }
173    for (uint64_t i = nsize; i < res_size; ++i) {
174      memset(res + i * res_sl, 0, nn * sizeof(double));
175    }
176  } else {
177    const uint64_t msize = res_size < b_size ? res_size : b_size;
178    const uint64_t nsize = res_size < a_size ? res_size : a_size;
179    for (uint64_t i = 0; i < msize; ++i) {
180      rnx_sub_avx(nn, res + i * res_sl, a + i * a_sl, b + i * b_sl);
181    }
182    for (uint64_t i = msize; i < nsize; ++i) {
183      memcpy(res + i * res_sl, a + i * a_sl, nn * sizeof(double));
184    }
185    for (uint64_t i = nsize; i < res_size; ++i) {
186      memset(res + i * res_sl, 0, nn * sizeof(double));
187    }
188  }
189}