/* Copyright 2020 Google LLC. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <cstdint>
#include <cstring>

#include "ruy/check_macros.h"
#include "ruy/opt_set.h"
#include "ruy/pack_x86.h"
#include "ruy/path.h"
#include "ruy/platform.h"
#include "ruy/profiler/instrumentation.h"

#if RUY_PLATFORM_AVX && RUY_OPT(INTRINSICS)
#include <immintrin.h>  // IWYU pragma: keep
#endif

namespace ruy {

#if !(RUY_PLATFORM_AVX && RUY_OPT(ASM))

void Pack8bitColMajorForAvx(const std::int8_t*, std::int8_t, const std::int8_t*,
                            int, int, int, std::int8_t*, std::int32_t*) {
  // CPU-ID-based checks should disable the path that would reach this point.
  RUY_DCHECK(false);
}

void PackFloatColMajorForAvx(const float*, const float*, int, int, int,
                             float*) {
  // CPU-ID-based checks should disable the path that would reach this point.
  RUY_DCHECK(false);
}

void Pack8bitRowMajorForAvx(const std::uint8_t*, int, int, std::int8_t*, int,
                            int, int, int, int, int, int, std::int32_t*) {
  RUY_DCHECK(false);
}

#else  // RUY_PLATFORM_AVX && RUY_OPT(ASM)

// The first int8_t template parameter is arbitrary: this routine is common to
// all 8-bit source matrix types.
using PackImpl8bitAvx =
    PackImpl<Path::kAvx, FixedKernelLayout<Order::kColMajor, 4, 8>, std::int8_t,
             std::int8_t, std::int32_t, Order::kColMajor>;

using PackImplFloatAvx =
    PackImpl<Path::kAvx, FixedKernelLayout<Order::kRowMajor, 1, 8>, float,
             float, float, Order::kColMajor>;

namespace {

// Perform the equivalent of mm256_permutevar8x32 with
// a second argument of {7, 5, 3, 1, 6, 4, 2, 0}
inline __m256i PermuteEpi32EvenOdds(const __m256i& a) {
  // a_lo = 3 2 1 0
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  // a_hi = 7 6 5 4
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  // shuffle a_lo to get 3 1 2 0
  __m128i tmp_lo = _mm_shuffle_epi32(a_lo, 0xd8);
  // shuffle a_hi to get 7 5 6 4
  __m128i tmp_hi = _mm_shuffle_epi32(a_hi, 0xd8);
  // unpack lo 64 of res_lo and res hi to get 6 4 2 0
  __m128i res_lo = _mm_unpacklo_epi64(tmp_lo, tmp_hi);
  // unpack hi 64 of res_lo and res hi to get 7 5 3 1
  __m128i res_hi = _mm_unpackhi_epi64(tmp_lo, tmp_hi);
  return _mm256_set_m128i(res_hi, res_lo);
}

inline __m128i mm256_extracti128_si256(const __m256i& a, const int imm) {
  switch (imm) {
    case 0:
      return _mm256_extractf128_si256(a, 0);
    case 1:
      return _mm256_extractf128_si256(a, 1);
    default:
      RUY_DCHECK_LT(imm, 2);
      return _mm_setzero_si128();
  }
}

inline __m256i mm256_cvtepi8_epi16(const __m128i& a) {
  // Take the upper 64 bits of a and put in the first 64 bits of 'hi'
  __m128i hi = _mm_unpackhi_epi64(a, _mm_setzero_si128());
  return _mm256_set_m128i(_mm_cvtepi8_epi16(hi), _mm_cvtepi8_epi16(a));
}

inline __m256i mm256_cvtepi16_epi32(const __m128i& a) {
  // Take the upper 64 bits of a and put in the first 64 bits of 'hi'
  __m128i hi = _mm_unpackhi_epi64(a, _mm_setzero_si128());
  return _mm256_set_m128i(_mm_cvtepi16_epi32(hi), _mm_cvtepi16_epi32(a));
}

inline __m256i mm256_xor_si256(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_xor_si128(a_lo, b_lo);
  __m128i hi = _mm_xor_si128(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_unpacklo_epi32(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_unpacklo_epi32(a_lo, b_lo);
  __m128i hi = _mm_unpacklo_epi32(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_unpacklo_epi64(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_unpacklo_epi64(a_lo, b_lo);
  __m128i hi = _mm_unpacklo_epi64(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_unpackhi_epi32(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_unpackhi_epi32(a_lo, b_lo);
  __m128i hi = _mm_unpackhi_epi32(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_unpackhi_epi64(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_unpackhi_epi64(a_lo, b_lo);
  __m128i hi = _mm_unpackhi_epi64(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_add_epi32(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_add_epi32(a_lo, b_lo);
  __m128i hi = _mm_add_epi32(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_add_epi16(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_add_epi16(a_lo, b_lo);
  __m128i hi = _mm_add_epi16(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m256i mm256_madd_epi16(const __m256i& a, const __m256i& b) {
  __m128i a_lo = _mm256_extractf128_si256(a, 0);
  __m128i a_hi = _mm256_extractf128_si256(a, 1);
  __m128i b_lo = _mm256_extractf128_si256(b, 0);
  __m128i b_hi = _mm256_extractf128_si256(b, 1);
  __m128i lo = _mm_madd_epi16(a_lo, b_lo);
  __m128i hi = _mm_madd_epi16(a_hi, b_hi);
  return _mm256_set_m128i(hi, lo);
}

inline __m128i mm_permute_helper(const __m256i& a, const __m256i& b,
                                 const int imm) {
  __m128i tmp = _mm_setzero_si128();
  if (!(imm & 8)) {
    switch (imm & 3) {
      case 0:
        return _mm256_extractf128_si256(a, 0);
      case 1:
        return _mm256_extractf128_si256(a, 1);
      case 2:
        return _mm256_extractf128_si256(b, 0);
      case 3:
        return _mm256_extractf128_si256(b, 1);
    }
  }
  return tmp;
}

inline __m256i mm256_permute2x128_si256(const __m256i& a, const __m256i& b,
                                        const int imm) {
  const int lo_imm = imm & 15;
  __m128i lo = mm_permute_helper(a, b, lo_imm);
  const int hi_imm = (imm >> 4) & 15;
  __m128i hi = mm_permute_helper(a, b, hi_imm);
  return _mm256_set_m128i(hi, lo);
}

inline void Pack8bitColMajorForAvxPacker(const std::int8_t* src_ptr,
                                         std::int8_t input_xor,
                                         const std::int8_t* zerobuf,
                                         int src_stride, int remaining_src_cols,
                                         int src_rows, std::int8_t* packed_ptr,
                                         std::int32_t* sums_ptr,
                                         std::int8_t* trailing_buf) {
  using Layout = PackImpl8bitAvx::Layout;
  RUY_DCHECK_EQ(Layout::kCols, 8);
  RUY_DCHECK_EQ(Layout::kRows, 4);
  // Each Layout::Rows is 4 contiguous input, contiguous packed elements.
  // We process 8 of these chunks at a time, padding short input chunks.
  constexpr int kNumRowChunks = 8;
  constexpr int kNumChunkedSrcRows = kNumRowChunks * Layout::kRows;

  const std::int8_t* src_ptr0 = src_ptr;
  const std::int8_t* src_ptr1 = src_ptr0 + src_stride;
  const std::int8_t* src_ptr2 = src_ptr1 + src_stride;
  const std::int8_t* src_ptr3 = src_ptr2 + src_stride;
  const std::int8_t* src_ptr4 = src_ptr3 + src_stride;
  const std::int8_t* src_ptr5 = src_ptr4 + src_stride;
  const std::int8_t* src_ptr6 = src_ptr5 + src_stride;
  const std::int8_t* src_ptr7 = src_ptr6 + src_stride;
  std::int64_t src_inc0 = kNumChunkedSrcRows;
  std::int64_t src_inc1 = kNumChunkedSrcRows;
  std::int64_t src_inc2 = kNumChunkedSrcRows;
  std::int64_t src_inc3 = kNumChunkedSrcRows;
  std::int64_t src_inc4 = kNumChunkedSrcRows;
  std::int64_t src_inc5 = kNumChunkedSrcRows;
  std::int64_t src_inc6 = kNumChunkedSrcRows;
  std::int64_t src_inc7 = kNumChunkedSrcRows;
  // Handle cases where source does not have Layout::kCols (8) columns.
  if (remaining_src_cols < 8) {
    if (remaining_src_cols <= 0) {
      src_ptr0 = zerobuf;
      src_inc0 = 0;
    }
    if (remaining_src_cols <= 1) {
      src_ptr1 = zerobuf;
      src_inc1 = 0;
    }
    if (remaining_src_cols <= 2) {
      src_ptr2 = zerobuf;
      src_inc2 = 0;
    }
    if (remaining_src_cols <= 3) {
      src_ptr3 = zerobuf;
      src_inc3 = 0;
    }
    if (remaining_src_cols <= 4) {
      src_ptr4 = zerobuf;
      src_inc4 = 0;
    }
    if (remaining_src_cols <= 5) {
      src_ptr5 = zerobuf;
      src_inc5 = 0;
    }
    if (remaining_src_cols <= 6) {
      src_ptr6 = zerobuf;
      src_inc6 = 0;
    }
    src_ptr7 = zerobuf;
    src_inc7 = 0;
  }

  const std::int8_t zero_point = zerobuf[0];

  if (sums_ptr) {
    // i: Layout::kCols.
    for (int i = 0; i < 8; ++i) {
      sums_ptr[i] = 0;
    }
  }
  std::int32_t sums_adjustment = 0;
  const __m256i ones_16bit = _mm256_set1_epi16(1);
  __m256i sums_4x2_32bit_lo = _mm256_set1_epi32(0);
  __m256i sums_4x2_32bit_hi = _mm256_set1_epi32(0);

  // The overall packing effectively pads the source rows to
  // (src_rows + 63) & ~63. The iteration over k may skip when m=1, and then we
  // only pack for (src_rows + 31) & ~31. When there is an incomplete
  // destination block, this is stored into trailing_buf instead of packed_ptr.
  for (int k = 0; k < src_rows; k += kNumChunkedSrcRows) {
    // Available source rows.
    // If this is less than 0 (for m=1), we skip, having filled trailing
    // buffer for m=0. Also, if source rows is zero on m=1, then we filled
    // exactly to the end of the column in the packed buffer.
    const int available_src_rows = src_rows - k;
    // Effectively,
    // available rows = std::max(0, std::min(8, src_rows - k));
    // treat each case separately.
    if (available_src_rows >= kNumChunkedSrcRows) {
      if (sums_ptr) {
        __m256i t0, t1, t2, t3, t4, t5, t6, t7;
        __m256i r0, r1, r2, r3, r4, r5, r6, r7;
        const __m256i input_xor_v = _mm256_set1_epi8(input_xor);

        t0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr0));
        t4 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr4));
        t1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr1));
        t5 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr5));
        t2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr2));
        t6 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr6));
        t3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr3));
        t7 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr7));

        r0 = mm256_unpacklo_epi32(t0, t1);
        r4 = mm256_unpacklo_epi32(t4, t5);
        r2 = mm256_unpackhi_epi32(t0, t1);
        r6 = mm256_unpackhi_epi32(t4, t5);
        r1 = mm256_unpacklo_epi32(t2, t3);
        r5 = mm256_unpacklo_epi32(t6, t7);
        r3 = mm256_unpackhi_epi32(t2, t3);
        r7 = mm256_unpackhi_epi32(t6, t7);

        t0 = mm256_unpacklo_epi64(r0, r1);
        t4 = mm256_unpacklo_epi64(r4, r5);
        t2 = mm256_unpackhi_epi64(r0, r1);
        t6 = mm256_unpackhi_epi64(r4, r5);
        t1 = mm256_unpacklo_epi64(r2, r3);
        t5 = mm256_unpacklo_epi64(r6, r7);
        t3 = mm256_unpackhi_epi64(r2, r3);
        t7 = mm256_unpackhi_epi64(r6, r7);

        // The preceding sets of rearrangement operations interleaved by 4 bytes
        // and then by 8 bytes *within* lanes. The following set interleave by
        // 16 bytes (128-bit), operating *between* AVX lanes. For instance (t0,
        // t4) are interleaved to create (r0, r1). This complexity follows from
        // the way that AVX is centered around MM 128-bit lanes.
        r0 = mm256_permute2x128_si256(t0, t4, 0x20);
        r4 = mm256_permute2x128_si256(t1, t5, 0x20);
        r1 = mm256_permute2x128_si256(t0, t4, 0x31);
        r5 = mm256_permute2x128_si256(t1, t5, 0x31);
        r2 = mm256_permute2x128_si256(t2, t6, 0x20);
        r6 = mm256_permute2x128_si256(t3, t7, 0x20);
        r3 = mm256_permute2x128_si256(t2, t6, 0x31);
        r7 = mm256_permute2x128_si256(t3, t7, 0x31);

        r0 = mm256_xor_si256(r0, input_xor_v);
        r1 = mm256_xor_si256(r1, input_xor_v);
        r2 = mm256_xor_si256(r2, input_xor_v);
        r3 = mm256_xor_si256(r3, input_xor_v);
        r4 = mm256_xor_si256(r4, input_xor_v);
        r5 = mm256_xor_si256(r5, input_xor_v);
        r6 = mm256_xor_si256(r6, input_xor_v);
        r7 = mm256_xor_si256(r7, input_xor_v);

        __m256i sums_4x4_16bit_lo;
        sums_4x4_16bit_lo = mm256_cvtepi8_epi16(_mm256_castsi256_si128(r0));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r1)));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r2)));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r3)));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r4)));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r5)));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r6)));
        sums_4x4_16bit_lo = mm256_add_epi16(
            sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r7)));

        // The sums have been performed across columns, and now we have 4x16-bit
        // sums packed together. We use madd for pairwise 32-bit sums.
        const __m256i sums_4x2_32bit_lo_new =
            mm256_madd_epi16(sums_4x4_16bit_lo, ones_16bit);
        sums_4x2_32bit_lo =
            mm256_add_epi32(sums_4x2_32bit_lo, sums_4x2_32bit_lo_new);

        __m256i sums_4x4_16bit_hi;
        sums_4x4_16bit_hi = mm256_cvtepi8_epi16(mm256_extracti128_si256(r0, 1));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r1, 1)));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r2, 1)));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r3, 1)));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r4, 1)));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r5, 1)));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r6, 1)));
        sums_4x4_16bit_hi = mm256_add_epi16(
            sums_4x4_16bit_hi,
            mm256_cvtepi8_epi16(mm256_extracti128_si256(r7, 1)));

        const __m256i sums_4x2_32bit_hi_new =
            mm256_madd_epi16(sums_4x4_16bit_hi, ones_16bit);
        sums_4x2_32bit_hi =
            mm256_add_epi32(sums_4x2_32bit_hi, sums_4x2_32bit_hi_new);

        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 0 * 8 * 4),
                            r0);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 2 * 8 * 4),
                            r4);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 4 * 8 * 4),
                            r1);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 6 * 8 * 4),
                            r5);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 1 * 8 * 4),
                            r2);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 3 * 8 * 4),
                            r6);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 5 * 8 * 4),
                            r3);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 7 * 8 * 4),
                            r7);
      } else {
        __m256i t0, t1, t2, t3, t4, t5, t6, t7;
        __m256i r0, r1, r2, r3, r4, r5, r6, r7;
        const __m256i input_xor_v = _mm256_set1_epi8(input_xor);

        t0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr0));
        t4 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr4));
        t1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr1));
        t5 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr5));
        t2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr2));
        t6 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr6));
        t3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr3));
        t7 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src_ptr7));

        r0 = mm256_unpacklo_epi32(t0, t1);
        r4 = mm256_unpacklo_epi32(t4, t5);
        r2 = mm256_unpackhi_epi32(t0, t1);
        r6 = mm256_unpackhi_epi32(t4, t5);
        r1 = mm256_unpacklo_epi32(t2, t3);
        r5 = mm256_unpacklo_epi32(t6, t7);
        r3 = mm256_unpackhi_epi32(t2, t3);
        r7 = mm256_unpackhi_epi32(t6, t7);

        t0 = mm256_unpacklo_epi64(r0, r1);
        t4 = mm256_unpacklo_epi64(r4, r5);
        t2 = mm256_unpackhi_epi64(r0, r1);
        t6 = mm256_unpackhi_epi64(r4, r5);
        t1 = mm256_unpacklo_epi64(r2, r3);
        t5 = mm256_unpacklo_epi64(r6, r7);
        t3 = mm256_unpackhi_epi64(r2, r3);
        t7 = mm256_unpackhi_epi64(r6, r7);

        // The preceding sets of rearrangement operations interleaved by 4 bytes
        // and then by 8 bytes *within* lanes. The following set interleave by
        // 16 bytes (128-bit), operating *between* AVX lanes. For instance (t0,
        // t4) are interleaved to create (r0, r1). This complexity follows from
        // the way that AVX is centered around MM 128-bit lanes.
        r0 = mm256_permute2x128_si256(t0, t4, 0x20);
        r4 = mm256_permute2x128_si256(t1, t5, 0x20);
        r1 = mm256_permute2x128_si256(t0, t4, 0x31);
        r5 = mm256_permute2x128_si256(t1, t5, 0x31);
        r2 = mm256_permute2x128_si256(t2, t6, 0x20);
        r6 = mm256_permute2x128_si256(t3, t7, 0x20);
        r3 = mm256_permute2x128_si256(t2, t6, 0x31);
        r7 = mm256_permute2x128_si256(t3, t7, 0x31);

        r0 = mm256_xor_si256(r0, input_xor_v);
        r1 = mm256_xor_si256(r1, input_xor_v);
        r2 = mm256_xor_si256(r2, input_xor_v);
        r3 = mm256_xor_si256(r3, input_xor_v);
        r4 = mm256_xor_si256(r4, input_xor_v);
        r5 = mm256_xor_si256(r5, input_xor_v);
        r6 = mm256_xor_si256(r6, input_xor_v);
        r7 = mm256_xor_si256(r7, input_xor_v);

        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 0 * 8 * 4),
                            r0);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 2 * 8 * 4),
                            r4);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 4 * 8 * 4),
                            r1);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 6 * 8 * 4),
                            r5);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 1 * 8 * 4),
                            r2);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 3 * 8 * 4),
                            r6);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 5 * 8 * 4),
                            r3);
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(packed_ptr + 7 * 8 * 4),
                            r7);
      }
    } else if (available_src_rows > 0) {
      RUY_DCHECK_LT(available_src_rows, kNumChunkedSrcRows);
      // We do not care what goes into the trailing buffer, but we want
      // in_data[...] ^ input_xor == 0 for irrelevant values in the summation.
      //
      // We compensate for padding-with-zero_point by initializing the
      // summations with the compensating offset, effectively
      // ((input_xor ^ input_xor) - (zero_point ^ input_xor)) *
      //                         4 * (8 - ((available_src_rows + 3) >> 2)).
      //
      // Note that (zero_point ^ input_xor) is performed in 8-bits and then
      // cast.
      sums_adjustment +=
          -(zero_point ^ input_xor) * 4 * (8 - ((available_src_rows + 3) >> 2));

      __m256i t0, t1, t2, t3, t4, t5, t6, t7;
      __m256i r0, r1, r2, r3, r4, r5, r6, r7;
      const __m256i input_xor_v = _mm256_set1_epi8(input_xor);

      t0 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr0);
      t4 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr4);
      t1 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr1);
      t5 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr5);
      t2 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr2);
      t6 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr6);
      t3 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr3);
      t7 = MaskLoadu<Path::kAvx>(available_src_rows, zero_point, src_ptr7);

      r0 = mm256_unpacklo_epi32(t0, t1);
      r4 = mm256_unpacklo_epi32(t4, t5);
      r2 = mm256_unpackhi_epi32(t0, t1);
      r6 = mm256_unpackhi_epi32(t4, t5);
      r1 = mm256_unpacklo_epi32(t2, t3);
      r5 = mm256_unpacklo_epi32(t6, t7);
      r3 = mm256_unpackhi_epi32(t2, t3);
      r7 = mm256_unpackhi_epi32(t6, t7);

      t0 = mm256_unpacklo_epi64(r0, r1);
      t4 = mm256_unpacklo_epi64(r4, r5);
      t2 = mm256_unpackhi_epi64(r0, r1);
      t6 = mm256_unpackhi_epi64(r4, r5);
      t1 = mm256_unpacklo_epi64(r2, r3);
      t5 = mm256_unpacklo_epi64(r6, r7);
      t3 = mm256_unpackhi_epi64(r2, r3);
      t7 = mm256_unpackhi_epi64(r6, r7);

      // The preceding sets of rearrangement operations interleaved by 4 bytes
      // and then by 8 bytes *within* lanes. The following set interleave by
      // 16 bytes (128-bit), operating *between* AVX lanes. For instance (t0,
      // t4) are interleaved to create (r0, r1). This complexity follows from
      // the way that AVX is centered around MM 128-bit lanes.
      r0 = mm256_permute2x128_si256(t0, t4, 0x20);
      r4 = mm256_permute2x128_si256(t1, t5, 0x20);
      r1 = mm256_permute2x128_si256(t0, t4, 0x31);
      r5 = mm256_permute2x128_si256(t1, t5, 0x31);
      r2 = mm256_permute2x128_si256(t2, t6, 0x20);
      r6 = mm256_permute2x128_si256(t3, t7, 0x20);
      r3 = mm256_permute2x128_si256(t2, t6, 0x31);
      r7 = mm256_permute2x128_si256(t3, t7, 0x31);

      r0 = mm256_xor_si256(r0, input_xor_v);
      r1 = mm256_xor_si256(r1, input_xor_v);
      r2 = mm256_xor_si256(r2, input_xor_v);
      r3 = mm256_xor_si256(r3, input_xor_v);
      r4 = mm256_xor_si256(r4, input_xor_v);
      r5 = mm256_xor_si256(r5, input_xor_v);
      r6 = mm256_xor_si256(r6, input_xor_v);
      r7 = mm256_xor_si256(r7, input_xor_v);

      __m256i sums_4x4_16bit_lo;
      sums_4x4_16bit_lo = mm256_cvtepi8_epi16(_mm256_castsi256_si128(r0));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r1)));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r2)));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r3)));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r4)));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r5)));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r6)));
      sums_4x4_16bit_lo = mm256_add_epi16(
          sums_4x4_16bit_lo, mm256_cvtepi8_epi16(_mm256_castsi256_si128(r7)));

      // The sums have been performed across columns, and now we have 4x16-bit
      // sums packed together. We use madd for pairwise 32-bit sums.
      const __m256i sums_4x2_32bit_lo_new =
          mm256_madd_epi16(sums_4x4_16bit_lo, ones_16bit);
      sums_4x2_32bit_lo =
          mm256_add_epi32(sums_4x2_32bit_lo, sums_4x2_32bit_lo_new);

      __m256i sums_4x4_16bit_hi;
      sums_4x4_16bit_hi = mm256_cvtepi8_epi16(mm256_extracti128_si256(r0, 1));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r1, 1)));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r2, 1)));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r3, 1)));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r4, 1)));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r5, 1)));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r6, 1)));
      sums_4x4_16bit_hi =
          mm256_add_epi16(sums_4x4_16bit_hi,
                          mm256_cvtepi8_epi16(mm256_extracti128_si256(r7, 1)));

      const __m256i sums_4x2_32bit_hi_new =
          mm256_madd_epi16(sums_4x4_16bit_hi, ones_16bit);
      sums_4x2_32bit_hi =
          mm256_add_epi32(sums_4x2_32bit_hi, sums_4x2_32bit_hi_new);

      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 0 * 8 * 4),
                          r0);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 2 * 8 * 4),
                          r4);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 4 * 8 * 4),
                          r1);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 6 * 8 * 4),
                          r5);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 1 * 8 * 4),
                          r2);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 3 * 8 * 4),
                          r6);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 5 * 8 * 4),
                          r3);
      _mm256_storeu_si256(reinterpret_cast<__m256i*>(trailing_buf + 7 * 8 * 4),
                          r7);
    }

    packed_ptr += 8 * kNumChunkedSrcRows;
    src_ptr0 += src_inc0;
    src_ptr1 += src_inc1;
    src_ptr2 += src_inc2;
    src_ptr3 += src_inc3;
    src_ptr4 += src_inc4;
    src_ptr5 += src_inc5;
    src_ptr6 += src_inc6;
    src_ptr7 += src_inc7;
  }

  if (sums_ptr) {
    const __m256i sums_adjustment_v = _mm256_set1_epi32(sums_adjustment);

    __m256i sums =
        _mm256_loadu_si256(reinterpret_cast<const __m256i*>(sums_ptr));

    // We earlier used madd for pairwise 32-bit sums, and now we deinterlace the
    // neighbours, finshing up by adding them to the stored accumulated sums.
    const __m256i sums_2x4_32bit_lo = PermuteEpi32EvenOdds(sums_4x2_32bit_lo);
    const __m256i sums_2x4_32bit_hi = PermuteEpi32EvenOdds(sums_4x2_32bit_hi);
    const __m256i sums_2x4_32bit_a =
        mm256_permute2x128_si256(sums_2x4_32bit_lo, sums_2x4_32bit_hi, 0x20);
    const __m256i sums_2x4_32bit_b =
        mm256_permute2x128_si256(sums_2x4_32bit_lo, sums_2x4_32bit_hi, 0x31);
    sums = mm256_add_epi32(sums, sums_adjustment_v);
    sums = mm256_add_epi32(sums, sums_2x4_32bit_a);
    sums = mm256_add_epi32(sums, sums_2x4_32bit_b);

    _mm256_storeu_si256(reinterpret_cast<__m256i*>(sums_ptr), sums);
  }
}

// Use a generic AVX intrinsic for greater-than comparison.
template <>
inline __m256i CompareGreaterThan<Path::kAvx>(const __m256i& a,
                                              const __m256i& b) {
  constexpr int kGreaterThanSignalling = 14;
  const __m256 v = _mm256_cmp_ps(_mm256_cvtepi32_ps(a), _mm256_cvtepi32_ps(b),
                                 kGreaterThanSignalling);
  return _mm256_cvtps_epi32(v);
}

}  // namespace.

void Pack8bitColMajorForAvx(const std::int8_t* src_ptr, std::int8_t input_xor,
                            const std::int8_t* zerobuf, int src_stride,
                            int remaining_src_cols, int src_rows,
                            std::int8_t* packed_ptr, std::int32_t* sums_ptr) {
  profiler::ScopeLabel label("Pack kAvx 8bit");

  using Layout = PackImpl8bitAvx::Layout;
  RUY_DCHECK_EQ(Layout::kCols, 8);
  RUY_DCHECK_EQ(Layout::kRows, 4);

  // Each Layout::Rows is 4 contiguous input, contiguous packed elements.
  // We process 8 of these chunks at a time, padding short input chunks.
  static constexpr int kNumRowChunks = 8;  // Short input is padded.

  // Each packed block is 4*8, and there are normally 8. The trailing block is
  // only slightly shorter.
  constexpr int kTrailingBufSize =
      kNumRowChunks * Layout::kCols * Layout::kRows;
  std::int8_t trailing_buf[kTrailingBufSize];
  memset(trailing_buf, 0, kTrailingBufSize * sizeof(std::int8_t));

  Pack8bitColMajorForAvxPacker(src_ptr, input_xor, zerobuf, src_stride,
                               remaining_src_cols, src_rows, packed_ptr,
                               sums_ptr, trailing_buf);

  constexpr int kChunkedRowMask = kNumRowChunks * Layout::kRows - 1;
  const bool trailing_data = (src_rows & kChunkedRowMask) > 0;
  // If the number of source rows is not a multiple of kChunkedRowMask, there
  // will be data in the trailing buffer,
  if (trailing_data) {
    const int non_trailing_rows = src_rows & ~kChunkedRowMask;
    // Destination "rows" are padded to next highest multiple of Layout::kRows.
    const int dst_rows = (src_rows + 3) & ~3;
    const int trailing_rows = dst_rows - non_trailing_rows;
    memcpy(packed_ptr + Layout::kCols * non_trailing_rows, trailing_buf,
           Layout::kCols * trailing_rows * sizeof(std::int8_t));
  }
}

void PackFloatColMajorForAvx(const float* src_ptr, const float* zerobuf,
                             int src_stride, int remaining_src_cols,
                             int src_rows, float* packed_ptr) {
  profiler::ScopeLabel label("Pack kAvx float");
  static constexpr int kPackCols = 8;  // Source cols packed together.
  static constexpr int kPackRows = 8;  // Short input is padded.
  float trailing_buf[(kPackRows - 1) * kPackCols];
  if (remaining_src_cols < 8) {
    memset(trailing_buf, 0, sizeof(trailing_buf));
  }
  PackFloatColMajorForAvxCommonPacker<PackImplFloatAvx, Path::kAvx>(
      src_ptr, zerobuf, src_stride, remaining_src_cols, src_rows, packed_ptr,
      trailing_buf);

  const int trailing_rows = src_rows & (kPackRows - 1);
  if (trailing_rows > 0) {
    const int non_trailing_rows = src_rows & ~(kPackRows - 1);
    memcpy(packed_ptr + kPackCols * non_trailing_rows, trailing_buf,
           kPackCols * trailing_rows * sizeof(float));
  }
}

void Pack8bitRowMajorForAvx(const std::uint8_t* src_ptr, int src_stride,
                            int src_zero_point, std::int8_t* packed_ptr,
                            int packed_stride, int start_col, int end_col,
                            int src_cols, int block_row, int src_rows,
                            int input_xor, std::int32_t* sums) {
  int col = start_col;
  int src_end_col = std::min(end_col, src_cols);

  for (; col <= src_end_col - 8; col += 8) {
    std::int8_t* dst_ptr = packed_ptr;
    __m128i val0, val1, val2, val3;
    __m128i input_xor_dup = _mm_set1_epi8(input_xor);
    // Load a 4x8 block.
    if (block_row + 4 <= src_rows) {
      val0 = _mm_loadu_si64(src_ptr + 0 * src_stride);
      val1 = _mm_loadu_si64(src_ptr + 1 * src_stride);
      val2 = _mm_loadu_si64(src_ptr + 2 * src_stride);
      val3 = _mm_loadu_si64(src_ptr + 3 * src_stride);
    } else {
      val0 = _mm_set1_epi8(src_zero_point);
      val1 = val0;
      val2 = val0;
      val3 = val0;
      if (block_row + 0 < src_rows)
        val0 = _mm_loadu_si64(src_ptr + 0 * src_stride);
      if (block_row + 1 < src_rows)
        val1 = _mm_loadu_si64(src_ptr + 1 * src_stride);
      if (block_row + 2 < src_rows)
        val2 = _mm_loadu_si64(src_ptr + 2 * src_stride);
      if (block_row + 3 < src_rows)
        val3 = _mm_loadu_si64(src_ptr + 3 * src_stride);
    }
    // Maybe xor the sign bit to convert from uint8 to int8.
    val0 = _mm_xor_si128(val0, input_xor_dup);
    val1 = _mm_xor_si128(val1, input_xor_dup);
    val2 = _mm_xor_si128(val2, input_xor_dup);
    val3 = _mm_xor_si128(val3, input_xor_dup);
    // Update the sums.
    __m128i val16_0 = _mm_cvtepi8_epi16(val0);
    __m128i val16_1 = _mm_cvtepi8_epi16(val1);
    __m128i val16_2 = _mm_cvtepi8_epi16(val2);
    __m128i val16_3 = _mm_cvtepi8_epi16(val3);
    __m128i new_sum16 = _mm_add_epi16(_mm_add_epi16(val16_0, val16_1),
                                      _mm_add_epi16(val16_2, val16_3));
    __m256i sum =
        _mm256_loadu_si256(reinterpret_cast<const __m256i*>(sums + col));
    sum = mm256_add_epi32(sum, mm256_cvtepi16_epi32(new_sum16));
    _mm256_storeu_si256(reinterpret_cast<__m256i*>(sums + col), sum);
    // Perform the transposition of 4x4 blocks
    __m128i t2_val0 = _mm_unpacklo_epi8(val0, val1);
    __m128i t2_val1 = _mm_unpacklo_epi8(val2, val3);
    __m128i t4_val0 = _mm_unpacklo_epi16(t2_val0, t2_val1);
    __m128i t4_val1 = _mm_unpackhi_epi16(t2_val0, t2_val1);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst_ptr), t4_val0);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst_ptr + 16), t4_val1);
    src_ptr += 8;
    packed_ptr += packed_stride * 8;
  }
  for (; col < src_end_col; col++) {
    std::int32_t accum = 0;
    for (int r = 0; r < 4; r++) {
      std::int8_t packed_val;
      if (block_row + r < src_rows) {
        packed_val = input_xor ^ src_ptr[r * src_stride];
      } else {
        packed_val = input_xor ^ src_zero_point;
      }
      accum += packed_val;
      *packed_ptr++ = packed_val;
    }
    if (sums) {
      sums[col] += accum;
    }
    src_ptr++;
  }
  for (; col < end_col; col++) {
    std::memset(packed_ptr, 0, 4);
    packed_ptr += 4;
  }
}

#endif  // RUY_PLATFORM_AVX && RUY_OPT(INTRINSICS)

}  // namespace ruy