1#include <assert.h>
  2#include <immintrin.h>
  3
  4#include "reim4_arithmetic.h"
  5
  6void reim4_extract_1blk_from_reim_avx(uint64_t m, uint64_t blk,
  7                                      double* const dst,       // nrows * 8 doubles
  8                                      const double* const src  // a contiguous array of nrows reim vectors
  9) {
 10  assert(blk < (m >> 2));
 11  const double* src_ptr = src + (blk << 2);
 12  double* dst_ptr = dst;
 13  _mm256_storeu_pd(dst_ptr, _mm256_loadu_pd(src_ptr));
 14  _mm256_storeu_pd(dst_ptr + 4, _mm256_loadu_pd(src_ptr + m));
 15}
 16
 17void reim4_extract_1blk_from_contiguous_reim_avx(uint64_t m, uint64_t nrows, uint64_t blk, double* const dst,
 18                                                 const double* const src) {
 19  assert(blk < (m >> 2));
 20  const double* src_ptr = src + (blk << 2);
 21  double* dst_ptr = dst;
 22  for (uint64_t i = 0; i < nrows * 2; ++i) {
 23    _mm256_storeu_pd(dst_ptr, _mm256_loadu_pd(src_ptr));
 24    dst_ptr += 4;
 25    src_ptr += m;
 26  }
 27}
 28
 29EXPORT void reim4_extract_1blk_from_contiguous_reim_sl_avx(uint64_t m, uint64_t sl, uint64_t nrows, uint64_t blk,
 30                                                           double* const dst, const double* const src) {
 31  assert(blk < (m >> 2));
 32  const double* src_ptr = src + (blk << 2);
 33  double* dst_ptr = dst;
 34  for (uint64_t i = 0; i < nrows; ++i) {
 35    _mm256_storeu_pd(dst_ptr, _mm256_loadu_pd(src_ptr));
 36    _mm256_storeu_pd(dst_ptr + 4, _mm256_loadu_pd(src_ptr + m));
 37    dst_ptr += 8;
 38    src_ptr += sl;
 39  }
 40}
 41
 42void reim4_save_1blk_to_reim_avx(uint64_t m, uint64_t blk,
 43                                 double* dst,       // 1 reim vector of length m
 44                                 const double* src  // 8 doubles
 45) {
 46  assert(blk < (m >> 2));
 47  const double* src_ptr = src;
 48  double* dst_ptr = dst + (blk << 2);
 49  _mm256_storeu_pd(dst_ptr, _mm256_loadu_pd(src_ptr));
 50  _mm256_storeu_pd(dst_ptr + m, _mm256_loadu_pd(src_ptr + 4));
 51}
 52
 53void reim4_add_1blk_to_reim_avx(uint64_t m, uint64_t blk,
 54                                double* dst,       // 1 reim vector of length m
 55                                const double* src  // 8 doubles
 56) {
 57  assert(blk < (m >> 2));
 58  const double* src_ptr = src;
 59  double* dst_ptr = dst + (blk << 2);
 60
 61  __m256d s0 = _mm256_loadu_pd(src_ptr);
 62  __m256d d0 = _mm256_loadu_pd(dst_ptr);
 63  _mm256_storeu_pd(dst_ptr, _mm256_add_pd(s0, d0));
 64
 65  __m256d s1 = _mm256_loadu_pd(src_ptr + 4);
 66  __m256d d1 = _mm256_loadu_pd(dst_ptr + m);
 67  _mm256_storeu_pd(dst_ptr + m, _mm256_add_pd(s1, d1));
 68}
 69
 70__always_inline void cplx_prod(__m256d* re1, __m256d* re2, __m256d* im, const double* const u_ptr,
 71                               const double* const v_ptr) {
 72  const __m256d a = _mm256_loadu_pd(u_ptr);
 73  const __m256d c = _mm256_loadu_pd(v_ptr);
 74  *re1 = _mm256_fmadd_pd(a, c, *re1);
 75
 76  const __m256d b = _mm256_loadu_pd(u_ptr + 4);
 77  *im = _mm256_fmadd_pd(b, c, *im);
 78
 79  const __m256d d = _mm256_loadu_pd(v_ptr + 4);
 80  *re2 = _mm256_fmadd_pd(b, d, *re2);
 81  *im = _mm256_fmadd_pd(a, d, *im);
 82}
 83
 84void reim4_vec_mat1col_product_avx2(const uint64_t nrows,
 85                                    double* const dst,      // 8 doubles
 86                                    const double* const u,  // nrows * 8 doubles
 87                                    const double* const v   // nrows * 8 doubles
 88) {
 89  __m256d re1 = _mm256_setzero_pd();
 90  __m256d re2 = _mm256_setzero_pd();
 91  __m256d im1 = _mm256_setzero_pd();
 92  __m256d im2 = _mm256_setzero_pd();
 93
 94  const double* u_ptr = u;
 95  const double* v_ptr = v;
 96  for (uint64_t i = 0; i < nrows; ++i) {
 97    const __m256d a = _mm256_loadu_pd(u_ptr);
 98    const __m256d c = _mm256_loadu_pd(v_ptr);
 99    re1 = _mm256_fmadd_pd(a, c, re1);
100
101    const __m256d b = _mm256_loadu_pd(u_ptr + 4);
102    im2 = _mm256_fmadd_pd(b, c, im2);
103
104    const __m256d d = _mm256_loadu_pd(v_ptr + 4);
105    re2 = _mm256_fmadd_pd(b, d, re2);
106    im1 = _mm256_fmadd_pd(a, d, im1);
107
108    u_ptr += 8;
109    v_ptr += 8;
110  }
111
112  _mm256_storeu_pd(dst, _mm256_sub_pd(re1, re2));
113  _mm256_storeu_pd(dst + 4, _mm256_add_pd(im1, im2));
114}
115
116EXPORT void reim4_vec_mat2cols_product_avx2(const uint64_t nrows,
117                                            double* const dst,      // 16 doubles
118                                            const double* const u,  // nrows * 16 doubles
119                                            const double* const v   // nrows * 16 doubles
120) {
121  __m256d re1 = _mm256_setzero_pd();
122  __m256d im1 = _mm256_setzero_pd();
123  __m256d re2 = _mm256_setzero_pd();
124  __m256d im2 = _mm256_setzero_pd();
125
126  __m256d ur, ui, ar, ai, br, bi;
127  for (uint64_t i = 0; i < nrows; ++i) {
128    ur = _mm256_loadu_pd(u + 8 * i);
129    ui = _mm256_loadu_pd(u + 8 * i + 4);
130    ar = _mm256_loadu_pd(v + 16 * i);
131    ai = _mm256_loadu_pd(v + 16 * i + 4);
132    br = _mm256_loadu_pd(v + 16 * i + 8);
133    bi = _mm256_loadu_pd(v + 16 * i + 12);
134    re1 = _mm256_fmsub_pd(ui, ai, re1);
135    re2 = _mm256_fmsub_pd(ui, bi, re2);
136    im1 = _mm256_fmadd_pd(ur, ai, im1);
137    im2 = _mm256_fmadd_pd(ur, bi, im2);
138    re1 = _mm256_fmsub_pd(ur, ar, re1);
139    re2 = _mm256_fmsub_pd(ur, br, re2);
140    im1 = _mm256_fmadd_pd(ui, ar, im1);
141    im2 = _mm256_fmadd_pd(ui, br, im2);
142  }
143  _mm256_storeu_pd(dst, re1);
144  _mm256_storeu_pd(dst + 4, im1);
145  _mm256_storeu_pd(dst + 8, re2);
146  _mm256_storeu_pd(dst + 12, im2);
147}