#include "transpose.h"

#include <immintrin.h>

namespace fbgemm {

void transpose_4rows(int N, const std::uint8_t* src, std::uint8_t* dst) {
  constexpr int M = 4;
  // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
  int j;
  // vectorized loop
  for (j = 0; j < N / 32 * 32; j += 32) {
    // a : a0 a1 ... a31
    // b : b0 b1 ... b31
    // c : c0 c1 ... c31
    // d : d0 d1 ... d31
    __m256i a = _mm256_lddqu_si256((const __m256i*)(src + j + 0 * N));
    __m256i b = _mm256_lddqu_si256((const __m256i*)(src + j + 1 * N));
    __m256i c = _mm256_lddqu_si256((const __m256i*)(src + j + 2 * N));
    __m256i d = _mm256_lddqu_si256((const __m256i*)(src + j + 3 * N));

    // even-odd interleaving
    // ab_lo : a0 b0 a1 b1 ...  a7  b7 | a16 b16 ... a23 b23
    // ab_hi : a8 b8 a9 b9 ... a15 b15 | a24 b24 ... a31 b31
    // cd_lo : c0 d0 c1 d1 ...  c7  d7 | c16 d16 ... c23 d23
    // cd_hi : c8 d8 c9 d9 ... c15 d15 | c24 d24 ... c31 d31
    __m256i ab_lo = _mm256_unpacklo_epi8(a, b);
    __m256i ab_hi = _mm256_unpackhi_epi8(a, b);
    __m256i cd_lo = _mm256_unpacklo_epi8(c, d);
    __m256i cd_hi = _mm256_unpackhi_epi8(c, d);

    // 4-row interleaving but permuted at 128-bit granularity
    // y0 :  a0  b0  c0  d0 ...  a-d3 | a-d16 ... a-d19
    // y1 :  a4  b4  c4  d4 ...  a-d7 | a-d20 ... a-d23
    // y2 :  a8  b8  c8  d8 ... a-d11 | a-d24 ... a-d27
    // y3 : a12 b12 c12 d12 ... a-d15 | a-d28 ... a-d31
    __m256i y0 = _mm256_unpacklo_epi16(ab_lo, cd_lo);
    __m256i y1 = _mm256_unpackhi_epi16(ab_lo, cd_lo);
    __m256i y2 = _mm256_unpacklo_epi16(ab_hi, cd_hi);
    __m256i y3 = _mm256_unpackhi_epi16(ab_hi, cd_hi);

    // Storing with 128-bit lanes are permuted so that everything is in order
    _mm256_storeu_si256(
        (__m256i*)(dst + j * M + 0 * 32),
        _mm256_permute2f128_si256(y0, y1, 0x20));
    _mm256_storeu_si256(
        (__m256i*)(dst + j * M + 1 * 32),
        _mm256_permute2f128_si256(y2, y3, 0x20));
    _mm256_storeu_si256(
        (__m256i*)(dst + j * M + 2 * 32),
        _mm256_permute2f128_si256(y0, y1, 0x31));
    _mm256_storeu_si256(
        (__m256i*)(dst + j * M + 3 * 32),
        _mm256_permute2f128_si256(y2, y3, 0x31));
  }
  // scalar loop for remainder
  for (; j < N; ++j) {
    for (int i = 0; i < M; ++i) {
      dst[j * M + i] = src[j + i * N];
    }
  }
}

} // namespace fbgemm