1#include <immintrin.h>
  2
  3#include "../commons_private.h"
  4#include "coeffs_arithmetic.h"
  5
  6// res = a + b. dimension n must be a power of 2
  7EXPORT void znx_add_i64_avx(uint64_t nn, int64_t* res, const int64_t* a, const int64_t* b) {
  8  if (nn <= 2) {
  9    if (nn == 1) {
 10      res[0] = a[0] + b[0];
 11    } else {
 12      _mm_storeu_si128((__m128i*)res,                     //
 13                       _mm_add_epi64(                     //
 14                           _mm_loadu_si128((__m128i*)a),  //
 15                           _mm_loadu_si128((__m128i*)b)));
 16    }
 17  } else {
 18    const __m256i* aa = (__m256i*)a;
 19    const __m256i* bb = (__m256i*)b;
 20    __m256i* rr = (__m256i*)res;
 21    __m256i* const rrend = (__m256i*)(res + nn);
 22    do {
 23      _mm256_storeu_si256(rr,                          //
 24                          _mm256_add_epi64(            //
 25                              _mm256_loadu_si256(aa),  //
 26                              _mm256_loadu_si256(bb)));
 27      ++rr;
 28      ++aa;
 29      ++bb;
 30    } while (rr < rrend);
 31  }
 32}
 33
 34// res = a - b. dimension n must be a power of 2
 35EXPORT void znx_sub_i64_avx(uint64_t nn, int64_t* res, const int64_t* a, const int64_t* b) {
 36  if (nn <= 2) {
 37    if (nn == 1) {
 38      res[0] = a[0] - b[0];
 39    } else {
 40      _mm_storeu_si128((__m128i*)res,                     //
 41                       _mm_sub_epi64(                     //
 42                           _mm_loadu_si128((__m128i*)a),  //
 43                           _mm_loadu_si128((__m128i*)b)));
 44    }
 45  } else {
 46    const __m256i* aa = (__m256i*)a;
 47    const __m256i* bb = (__m256i*)b;
 48    __m256i* rr = (__m256i*)res;
 49    __m256i* const rrend = (__m256i*)(res + nn);
 50    do {
 51      _mm256_storeu_si256(rr,                          //
 52                          _mm256_sub_epi64(            //
 53                              _mm256_loadu_si256(aa),  //
 54                              _mm256_loadu_si256(bb)));
 55      ++rr;
 56      ++aa;
 57      ++bb;
 58    } while (rr < rrend);
 59  }
 60}
 61
 62EXPORT void znx_negate_i64_avx(uint64_t nn, int64_t* res, const int64_t* a) {
 63  if (nn <= 2) {
 64    if (nn == 1) {
 65      res[0] = -a[0];
 66    } else {
 67      _mm_storeu_si128((__m128i*)res,           //
 68                       _mm_sub_epi64(           //
 69                           _mm_set1_epi64x(0),  //
 70                           _mm_loadu_si128((__m128i*)a)));
 71    }
 72  } else {
 73    const __m256i* aa = (__m256i*)a;
 74    __m256i* rr = (__m256i*)res;
 75    __m256i* const rrend = (__m256i*)(res + nn);
 76    do {
 77      _mm256_storeu_si256(rr,                         //
 78                          _mm256_sub_epi64(           //
 79                              _mm256_set1_epi64x(0),  //
 80                              _mm256_loadu_si256(aa)));
 81      ++rr;
 82      ++aa;
 83    } while (rr < rrend);
 84  }
 85}
 86
 87EXPORT void rnx_divide_by_m_avx(uint64_t n, double m, double* res, const double* a) {
 88  // TODO: see if there is a faster way of dividing by a power of 2?
 89  const double invm = 1. / m;
 90  if (n < 8) {
 91    switch (n) {
 92      case 1:
 93        *res = *a * invm;
 94        break;
 95      case 2:
 96        _mm_storeu_pd(res,                         //
 97                      _mm_mul_pd(_mm_loadu_pd(a),  //
 98                                 _mm_set1_pd(invm)));
 99        break;
100      case 4:
101        _mm256_storeu_pd(res,                               //
102                         _mm256_mul_pd(_mm256_loadu_pd(a),  //
103                                       _mm256_set1_pd(invm)));
104        break;
105      default:
106        NOT_SUPPORTED();  // non-power of 2
107    }
108    return;
109  }
110  const __m256d invm256 = _mm256_set1_pd(invm);
111  double* rr = res;
112  const double* aa = a;
113  const double* const aaend = a + n;
114  do {
115    _mm256_storeu_pd(rr,                                 //
116                     _mm256_mul_pd(_mm256_loadu_pd(aa),  //
117                                   invm256));
118    _mm256_storeu_pd(rr + 4,                                 //
119                     _mm256_mul_pd(_mm256_loadu_pd(aa + 4),  //
120                                   invm256));
121    rr += 8;
122    aa += 8;
123  } while (aa < aaend);
124}