1#include <memory.h>
 2
 3#include "immintrin.h"
 4#include "vec_rnx_arithmetic_private.h"
 5
 6/** @brief sets res = gadget_decompose(a) */
 7EXPORT void rnx_approxdecomp_from_tnxdbl_avx(         //
 8    const MOD_RNX* module,                            // N
 9    const TNXDBL_APPROXDECOMP_GADGET* gadget,         // output base 2^K
10    double* res, uint64_t res_size, uint64_t res_sl,  // res
11    const double* a                                   // a
12) {
13  const uint64_t nn = module->n;
14  if (nn < 4) return rnx_approxdecomp_from_tnxdbl_ref(module, gadget, res, res_size, res_sl, a);
15  const uint64_t ell = gadget->ell;
16  const __m256i k = _mm256_set1_epi64x(gadget->k);
17  const __m256d add_cst = _mm256_set1_pd(gadget->add_cst);
18  const __m256i and_mask = _mm256_set1_epi64x(gadget->and_mask);
19  const __m256i or_mask = _mm256_set1_epi64x(gadget->or_mask);
20  const __m256d sub_cst = _mm256_set1_pd(gadget->sub_cst);
21  const uint64_t msize = res_size <= ell ? res_size : ell;
22  // gadget decompose column by column
23  if (msize == ell) {
24    // this is the main scenario when msize == ell
25    double* const last_r = res + (msize - 1) * res_sl;
26    for (uint64_t j = 0; j < nn; j += 4) {
27      double* rr = last_r + j;
28      const double* aa = a + j;
29      __m256d t_dbl = _mm256_add_pd(_mm256_loadu_pd(aa), add_cst);
30      __m256i t_int = _mm256_castpd_si256(t_dbl);
31      do {
32        __m256i u_int = _mm256_or_si256(_mm256_and_si256(t_int, and_mask), or_mask);
33        _mm256_storeu_pd(rr, _mm256_sub_pd(_mm256_castsi256_pd(u_int), sub_cst));
34        t_int = _mm256_srlv_epi64(t_int, k);
35        rr -= res_sl;
36      } while (rr >= res);
37    }
38  } else if (msize > 0) {
39    // otherwise, if msize < ell: there is one additional rshift
40    const __m256i first_rsh = _mm256_set1_epi64x((ell - msize) * gadget->k);
41    double* const last_r = res + (msize - 1) * res_sl;
42    for (uint64_t j = 0; j < nn; j += 4) {
43      double* rr = last_r + j;
44      const double* aa = a + j;
45      __m256d t_dbl = _mm256_add_pd(_mm256_loadu_pd(aa), add_cst);
46      __m256i t_int = _mm256_srlv_epi64(_mm256_castpd_si256(t_dbl), first_rsh);
47      do {
48        __m256i u_int = _mm256_or_si256(_mm256_and_si256(t_int, and_mask), or_mask);
49        _mm256_storeu_pd(rr, _mm256_sub_pd(_mm256_castsi256_pd(u_int), sub_cst));
50        t_int = _mm256_srlv_epi64(t_int, k);
51        rr -= res_sl;
52      } while (rr >= res);
53    }
54  }
55  // zero-out the last slices (if any)
56  for (uint64_t i = msize; i < res_size; ++i) {
57    memset(res + i * res_sl, 0, nn * sizeof(double));
58  }
59}