/*

 *  Copyright (c) 2023 The WebM project authors. All Rights Reserved.

 *

 *  Use of this source code is governed by a BSD-style license

 *  that can be found in the LICENSE file in the root of the source

 *  tree. An additional intellectual property rights grant can be found

 *  in the file PATENTS.  All contributing project authors may

 *  be found in the AUTHORS file in the root of the source tree.

 */

#include <immintrin.h>  // AVX2

#include "./vpx_dsp_rtcd.h"

#include "vpx_dsp/txfm_common.h"

#define PAIR256_SET_EPI16(a, b)                                            \

  _mm256_set_epi16((int16_t)(b), (int16_t)(a), (int16_t)(b), (int16_t)(a), \

                   (int16_t)(b), (int16_t)(a), (int16_t)(b), (int16_t)(a), \

                   (int16_t)(b), (int16_t)(a), (int16_t)(b), (int16_t)(a), \

                   (int16_t)(b), (int16_t)(a), (int16_t)(b), (int16_t)(a))

static INLINE void idct_load16x16(const tran_low_t *input, __m256i *in,

                                  int stride) {

  int i;

  // Load 16x16 values

  for (i = 0; i < 16; i++) {

#if CONFIG_VP9_HIGHBITDEPTH

    const __m128i in0 = _mm_loadu_si128((const __m128i *)(input + i * stride));

    const __m128i in1 =

        _mm_loadu_si128((const __m128i *)((input + i * stride) + 4));

    const __m128i in2 =

        _mm_loadu_si128((const __m128i *)((input + i * stride) + 8));

    const __m128i in3 =

        _mm_loadu_si128((const __m128i *)((input + i * stride) + 12));

    const __m128i ls = _mm_packs_epi32(in0, in1);

    const __m128i rs = _mm_packs_epi32(in2, in3);

    in[i] = _mm256_inserti128_si256(_mm256_castsi128_si256(ls), rs, 1);

#else

    in[i] = _mm256_load_si256((const __m256i *)(input + i * stride));

#endif

  }

}

static INLINE __m256i dct_round_shift_avx2(__m256i in) {

  const __m256i t = _mm256_add_epi32(in, _mm256_set1_epi32(DCT_CONST_ROUNDING));

  return _mm256_srai_epi32(t, DCT_CONST_BITS);

}

static INLINE __m256i idct_madd_round_shift_avx2(__m256i *in, __m256i *cospi) {

  const __m256i t = _mm256_madd_epi16(*in, *cospi);

  return dct_round_shift_avx2(t);

}

// Calculate the dot product between in0/1 and x and wrap to short.

static INLINE __m256i idct_calc_wraplow_avx2(__m256i *in0, __m256i *in1,

                                             __m256i *x) {

  const __m256i t0 = idct_madd_round_shift_avx2(in0, x);

  const __m256i t1 = idct_madd_round_shift_avx2(in1, x);

  return _mm256_packs_epi32(t0, t1);

}

// Multiply elements by constants and add them together.

static INLINE void butterfly16(__m256i in0, __m256i in1, int c0, int c1,

                               __m256i *out0, __m256i *out1) {

  __m256i cst0 = PAIR256_SET_EPI16(c0, -c1);

  __m256i cst1 = PAIR256_SET_EPI16(c1, c0);

  __m256i lo = _mm256_unpacklo_epi16(in0, in1);

  __m256i hi = _mm256_unpackhi_epi16(in0, in1);

  *out0 = idct_calc_wraplow_avx2(&lo, &hi, &cst0);

  *out1 = idct_calc_wraplow_avx2(&lo, &hi, &cst1);

}

static INLINE void idct16_16col(__m256i *in, __m256i *out) {

  __m256i step1[16], step2[16];

  // stage 2

  butterfly16(in[1], in[15], cospi_30_64, cospi_2_64, &step2[8], &step2[15]);

  butterfly16(in[9], in[7], cospi_14_64, cospi_18_64, &step2[9], &step2[14]);

  butterfly16(in[5], in[11], cospi_22_64, cospi_10_64, &step2[10], &step2[13]);

  butterfly16(in[13], in[3], cospi_6_64, cospi_26_64, &step2[11], &step2[12]);

  // stage 3

  butterfly16(in[2], in[14], cospi_28_64, cospi_4_64, &step1[4], &step1[7]);

  butterfly16(in[10], in[6], cospi_12_64, cospi_20_64, &step1[5], &step1[6]);

  step1[8] = _mm256_add_epi16(step2[8], step2[9]);

  step1[9] = _mm256_sub_epi16(step2[8], step2[9]);

  step1[10] = _mm256_sub_epi16(step2[11], step2[10]);

  step1[11] = _mm256_add_epi16(step2[10], step2[11]);

  step1[12] = _mm256_add_epi16(step2[12], step2[13]);

  step1[13] = _mm256_sub_epi16(step2[12], step2[13]);

  step1[14] = _mm256_sub_epi16(step2[15], step2[14]);

  step1[15] = _mm256_add_epi16(step2[14], step2[15]);

  // stage 4

  butterfly16(in[0], in[8], cospi_16_64, cospi_16_64, &step2[1], &step2[0]);

  butterfly16(in[4], in[12], cospi_24_64, cospi_8_64, &step2[2], &step2[3]);

  butterfly16(step1[14], step1[9], cospi_24_64, cospi_8_64, &step2[9],

              &step2[14]);

  butterfly16(step1[10], step1[13], -cospi_8_64, -cospi_24_64, &step2[13],

              &step2[10]);

  step2[5] = _mm256_sub_epi16(step1[4], step1[5]);

  step1[4] = _mm256_add_epi16(step1[4], step1[5]);

  step2[6] = _mm256_sub_epi16(step1[7], step1[6]);

  step1[7] = _mm256_add_epi16(step1[6], step1[7]);

  step2[8] = step1[8];

  step2[11] = step1[11];

  step2[12] = step1[12];

  step2[15] = step1[15];

  // stage 5

  step1[0] = _mm256_add_epi16(step2[0], step2[3]);

  step1[1] = _mm256_add_epi16(step2[1], step2[2]);

  step1[2] = _mm256_sub_epi16(step2[1], step2[2]);

  step1[3] = _mm256_sub_epi16(step2[0], step2[3]);

  butterfly16(step2[6], step2[5], cospi_16_64, cospi_16_64, &step1[5],

              &step1[6]);

  step1[8] = _mm256_add_epi16(step2[8], step2[11]);

  step1[9] = _mm256_add_epi16(step2[9], step2[10]);

  step1[10] = _mm256_sub_epi16(step2[9], step2[10]);

  step1[11] = _mm256_sub_epi16(step2[8], step2[11]);

  step1[12] = _mm256_sub_epi16(step2[15], step2[12]);

  step1[13] = _mm256_sub_epi16(step2[14], step2[13]);

  step1[14] = _mm256_add_epi16(step2[14], step2[13]);

  step1[15] = _mm256_add_epi16(step2[15], step2[12]);

  // stage 6

  step2[0] = _mm256_add_epi16(step1[0], step1[7]);

  step2[1] = _mm256_add_epi16(step1[1], step1[6]);

  step2[2] = _mm256_add_epi16(step1[2], step1[5]);

  step2[3] = _mm256_add_epi16(step1[3], step1[4]);

  step2[4] = _mm256_sub_epi16(step1[3], step1[4]);

  step2[5] = _mm256_sub_epi16(step1[2], step1[5]);

  step2[6] = _mm256_sub_epi16(step1[1], step1[6]);

  step2[7] = _mm256_sub_epi16(step1[0], step1[7]);

  butterfly16(step1[13], step1[10], cospi_16_64, cospi_16_64, &step2[10],

              &step2[13]);

  butterfly16(step1[12], step1[11], cospi_16_64, cospi_16_64, &step2[11],

              &step2[12]);

  // stage 7

  out[0] = _mm256_add_epi16(step2[0], step1[15]);

  out[1] = _mm256_add_epi16(step2[1], step1[14]);

  out[2] = _mm256_add_epi16(step2[2], step2[13]);

  out[3] = _mm256_add_epi16(step2[3], step2[12]);

  out[4] = _mm256_add_epi16(step2[4], step2[11]);

  out[5] = _mm256_add_epi16(step2[5], step2[10]);

  out[6] = _mm256_add_epi16(step2[6], step1[9]);

  out[7] = _mm256_add_epi16(step2[7], step1[8]);

  out[8] = _mm256_sub_epi16(step2[7], step1[8]);

  out[9] = _mm256_sub_epi16(step2[6], step1[9]);

  out[10] = _mm256_sub_epi16(step2[5], step2[10]);

  out[11] = _mm256_sub_epi16(step2[4], step2[11]);

  out[12] = _mm256_sub_epi16(step2[3], step2[12]);

  out[13] = _mm256_sub_epi16(step2[2], step2[13]);

  out[14] = _mm256_sub_epi16(step2[1], step1[14]);

  out[15] = _mm256_sub_epi16(step2[0], step1[15]);

}

static INLINE void recon_and_store16(uint8_t *dest, __m256i in_x) {

  const __m256i zero = _mm256_setzero_si256();

  __m256i d0 = _mm256_castsi128_si256(_mm_loadu_si128((__m128i *)(dest)));

  d0 = _mm256_permute4x64_epi64(d0, 0xd8);

  d0 = _mm256_unpacklo_epi8(d0, zero);

  d0 = _mm256_add_epi16(in_x, d0);

  d0 = _mm256_packus_epi16(

      d0, _mm256_castsi128_si256(_mm256_extractf128_si256(d0, 1)));

  _mm_storeu_si128((__m128i *)dest, _mm256_castsi256_si128(d0));

}

static INLINE void write_buffer_16x1(uint8_t *dest, __m256i in) {

  const __m256i final_rounding = _mm256_set1_epi16(1 << 5);

  __m256i out;

  out = _mm256_adds_epi16(in, final_rounding);

  out = _mm256_srai_epi16(out, 6);

  recon_and_store16(dest, out);

}

static INLINE void store_buffer_16x32(__m256i *in, uint8_t *dst, int stride) {

  const __m256i final_rounding = _mm256_set1_epi16(1 << 5);

  int j = 0;

  while (j < 32) {

    in[j] = _mm256_adds_epi16(in[j], final_rounding);

    in[j + 1] = _mm256_adds_epi16(in[j + 1], final_rounding);

    in[j] = _mm256_srai_epi16(in[j], 6);

    in[j + 1] = _mm256_srai_epi16(in[j + 1], 6);

    recon_and_store16(dst, in[j]);

    dst += stride;

    recon_and_store16(dst, in[j + 1]);

    dst += stride;

    j += 2;

  }

}

static INLINE void transpose2_8x8_avx2(__m256i *in, __m256i *out) {

  int i;

  __m256i t[16], u[16];

  // (1st, 2nd) ==> (lo, hi)

  //   (0, 1)   ==>  (0, 1)

  //   (2, 3)   ==>  (2, 3)

  //   (4, 5)   ==>  (4, 5)

  //   (6, 7)   ==>  (6, 7)

  for (i = 0; i < 4; i++) {

    t[2 * i] = _mm256_unpacklo_epi16(in[2 * i], in[2 * i + 1]);

    t[2 * i + 1] = _mm256_unpackhi_epi16(in[2 * i], in[2 * i + 1]);

  }

  // (1st, 2nd) ==> (lo, hi)

  //   (0, 2)   ==>  (0, 2)

  //   (1, 3)   ==>  (1, 3)

  //   (4, 6)   ==>  (4, 6)

  //   (5, 7)   ==>  (5, 7)

  for (i = 0; i < 2; i++) {

    u[i] = _mm256_unpacklo_epi32(t[i], t[i + 2]);

    u[i + 2] = _mm256_unpackhi_epi32(t[i], t[i + 2]);

    u[i + 4] = _mm256_unpacklo_epi32(t[i + 4], t[i + 6]);

    u[i + 6] = _mm256_unpackhi_epi32(t[i + 4], t[i + 6]);

  }

  // (1st, 2nd) ==> (lo, hi)

  //   (0, 4)   ==>  (0, 1)

  //   (1, 5)   ==>  (4, 5)

  //   (2, 6)   ==>  (2, 3)

  //   (3, 7)   ==>  (6, 7)

  for (i = 0; i < 2; i++) {

    out[2 * i] = _mm256_unpacklo_epi64(u[2 * i], u[2 * i + 4]);

    out[2 * i + 1] = _mm256_unpackhi_epi64(u[2 * i], u[2 * i + 4]);

    out[2 * i + 4] = _mm256_unpacklo_epi64(u[2 * i + 1], u[2 * i + 5]);

    out[2 * i + 5] = _mm256_unpackhi_epi64(u[2 * i + 1], u[2 * i + 5]);

  }

}

static INLINE void transpose_16bit_16x16_avx2(__m256i *in, __m256i *out) {

  __m256i t[16];

#define LOADL(idx)                                                            \

  t[idx] = _mm256_castsi128_si256(_mm_load_si128((__m128i const *)&in[idx])); \

  t[idx] = _mm256_inserti128_si256(                                           \

      t[idx], _mm_load_si128((__m128i const *)&in[(idx) + 8]), 1);

#define LOADR(idx)                                                           \

  t[8 + (idx)] =                                                             \

      _mm256_castsi128_si256(_mm_load_si128((__m128i const *)&in[idx] + 1)); \

  t[8 + (idx)] = _mm256_inserti128_si256(                                    \

      t[8 + (idx)], _mm_load_si128((__m128i const *)&in[(idx) + 8] + 1), 1);

  // load left 8x16

  LOADL(0)

  LOADL(1)

  LOADL(2)

  LOADL(3)

  LOADL(4)

  LOADL(5)

  LOADL(6)

  LOADL(7)

  // load right 8x16

  LOADR(0)

  LOADR(1)

  LOADR(2)

  LOADR(3)

  LOADR(4)

  LOADR(5)

  LOADR(6)

  LOADR(7)

  // get the top 16x8 result

  transpose2_8x8_avx2(t, out);

  // get the bottom 16x8 result

  transpose2_8x8_avx2(&t[8], &out[8]);

}

void vpx_idct16x16_256_add_avx2(const tran_low_t *input, uint8_t *dest,

                                int stride) {

  int i;

  __m256i in[16];

  // Load 16x16 values

  idct_load16x16(input, in, 16);

  transpose_16bit_16x16_avx2(in, in);

  idct16_16col(in, in);

  transpose_16bit_16x16_avx2(in, in);

  idct16_16col(in, in);

  for (i = 0; i < 16; ++i) {

    write_buffer_16x1(dest + i * stride, in[i]);

  }

}

// Only do addition and subtraction butterfly, size = 16, 32

static INLINE void add_sub_butterfly_avx2(__m256i *in, __m256i *out, int size) {

  int i = 0;

  const int num = size >> 1;

  const int bound = size - 1;

  while (i < num) {

    out[i] = _mm256_add_epi16(in[i], in[bound - i]);

    out[bound - i] = _mm256_sub_epi16(in[i], in[bound - i]);

    i++;

  }

}

// For each 16x32 block __m256i in[32],

// Input with index, 0, 4, 8, 12, 16, 20, 24, 28

// output pixels: 0-7 in __m256i out[32]

static INLINE void idct32_1024_16x32_quarter_1(__m256i *in, __m256i *out) {

  __m256i step1[8], step2[8];

  // stage 3

  butterfly16(in[4], in[28], cospi_28_64, cospi_4_64, &step1[4], &step1[7]);

  butterfly16(in[20], in[12], cospi_12_64, cospi_20_64, &step1[5], &step1[6]);

  // stage 4

  butterfly16(in[0], in[16], cospi_16_64, cospi_16_64, &step2[1], &step2[0]);

  butterfly16(in[8], in[24], cospi_24_64, cospi_8_64, &step2[2], &step2[3]);

  step2[4] = _mm256_add_epi16(step1[4], step1[5]);

  step2[5] = _mm256_sub_epi16(step1[4], step1[5]);

  step2[6] = _mm256_sub_epi16(step1[7], step1[6]);

  step2[7] = _mm256_add_epi16(step1[7], step1[6]);

  // stage 5

  step1[0] = _mm256_add_epi16(step2[0], step2[3]);

  step1[1] = _mm256_add_epi16(step2[1], step2[2]);

  step1[2] = _mm256_sub_epi16(step2[1], step2[2]);

  step1[3] = _mm256_sub_epi16(step2[0], step2[3]);

  step1[4] = step2[4];

  butterfly16(step2[6], step2[5], cospi_16_64, cospi_16_64, &step1[5],

              &step1[6]);

  step1[7] = step2[7];

  // stage 6

  out[0] = _mm256_add_epi16(step1[0], step1[7]);

  out[1] = _mm256_add_epi16(step1[1], step1[6]);

  out[2] = _mm256_add_epi16(step1[2], step1[5]);

  out[3] = _mm256_add_epi16(step1[3], step1[4]);

  out[4] = _mm256_sub_epi16(step1[3], step1[4]);

  out[5] = _mm256_sub_epi16(step1[2], step1[5]);

  out[6] = _mm256_sub_epi16(step1[1], step1[6]);

  out[7] = _mm256_sub_epi16(step1[0], step1[7]);

}

static INLINE void idct32_16x32_quarter_2_stage_4_to_6(__m256i *step1,

                                                       __m256i *out) {

  __m256i step2[32];

  // stage 4

  step2[8] = step1[8];

  step2[15] = step1[15];

  butterfly16(step1[14], step1[9], cospi_24_64, cospi_8_64, &step2[9],

              &step2[14]);

  butterfly16(step1[13], step1[10], -cospi_8_64, cospi_24_64, &step2[10],

              &step2[13]);

  step2[11] = step1[11];

  step2[12] = step1[12];

  // stage 5

  step1[8] = _mm256_add_epi16(step2[8], step2[11]);

  step1[9] = _mm256_add_epi16(step2[9], step2[10]);

  step1[10] = _mm256_sub_epi16(step2[9], step2[10]);

  step1[11] = _mm256_sub_epi16(step2[8], step2[11]);

  step1[12] = _mm256_sub_epi16(step2[15], step2[12]);

  step1[13] = _mm256_sub_epi16(step2[14], step2[13]);

  step1[14] = _mm256_add_epi16(step2[14], step2[13]);

  step1[15] = _mm256_add_epi16(step2[15], step2[12]);

  // stage 6

  out[8] = step1[8];

  out[9] = step1[9];

  butterfly16(step1[13], step1[10], cospi_16_64, cospi_16_64, &out[10],

              &out[13]);

  butterfly16(step1[12], step1[11], cospi_16_64, cospi_16_64, &out[11],

              &out[12]);

  out[14] = step1[14];

  out[15] = step1[15];

}

// For each 16x32 block __m256i in[32],

// Input with index, 2, 6, 10, 14, 18, 22, 26, 30

// output pixels: 8-15 in __m256i out[32]

static INLINE void idct32_1024_16x32_quarter_2(__m256i *in, __m256i *out) {

  __m256i step1[16], step2[16];

  // stage 2

  butterfly16(in[2], in[30], cospi_30_64, cospi_2_64, &step2[8], &step2[15]);

  butterfly16(in[18], in[14], cospi_14_64, cospi_18_64, &step2[9], &step2[14]);

  butterfly16(in[10], in[22], cospi_22_64, cospi_10_64, &step2[10], &step2[13]);

  butterfly16(in[26], in[6], cospi_6_64, cospi_26_64, &step2[11], &step2[12]);

  // stage 3

  step1[8] = _mm256_add_epi16(step2[8], step2[9]);

  step1[9] = _mm256_sub_epi16(step2[8], step2[9]);

  step1[10] = _mm256_sub_epi16(step2[11], step2[10]);

  step1[11] = _mm256_add_epi16(step2[11], step2[10]);

  step1[12] = _mm256_add_epi16(step2[12], step2[13]);

  step1[13] = _mm256_sub_epi16(step2[12], step2[13]);

  step1[14] = _mm256_sub_epi16(step2[15], step2[14]);

  step1[15] = _mm256_add_epi16(step2[15], step2[14]);

  idct32_16x32_quarter_2_stage_4_to_6(step1, out);

}

static INLINE void idct32_16x32_quarter_3_4_stage_4_to_7(__m256i *step1,

                                                         __m256i *out) {

  __m256i step2[32];

  // stage 4

  step2[16] = _mm256_add_epi16(step1[16], step1[19]);

  step2[17] = _mm256_add_epi16(step1[17], step1[18]);

  step2[18] = _mm256_sub_epi16(step1[17], step1[18]);

  step2[19] = _mm256_sub_epi16(step1[16], step1[19]);

  step2[20] = _mm256_sub_epi16(step1[23], step1[20]);

  step2[21] = _mm256_sub_epi16(step1[22], step1[21]);

  step2[22] = _mm256_add_epi16(step1[22], step1[21]);

  step2[23] = _mm256_add_epi16(step1[23], step1[20]);

  step2[24] = _mm256_add_epi16(step1[24], step1[27]);

  step2[25] = _mm256_add_epi16(step1[25], step1[26]);

  step2[26] = _mm256_sub_epi16(step1[25], step1[26]);

  step2[27] = _mm256_sub_epi16(step1[24], step1[27]);

  step2[28] = _mm256_sub_epi16(step1[31], step1[28]);

  step2[29] = _mm256_sub_epi16(step1[30], step1[29]);

  step2[30] = _mm256_add_epi16(step1[29], step1[30]);

  step2[31] = _mm256_add_epi16(step1[28], step1[31]);

  // stage 5

  step1[16] = step2[16];

  step1[17] = step2[17];

  butterfly16(step2[29], step2[18], cospi_24_64, cospi_8_64, &step1[18],

              &step1[29]);

  butterfly16(step2[28], step2[19], cospi_24_64, cospi_8_64, &step1[19],

              &step1[28]);

  butterfly16(step2[27], step2[20], -cospi_8_64, cospi_24_64, &step1[20],

              &step1[27]);

  butterfly16(step2[26], step2[21], -cospi_8_64, cospi_24_64, &step1[21],

              &step1[26]);

  step1[22] = step2[22];

  step1[23] = step2[23];

  step1[24] = step2[24];

  step1[25] = step2[25];

  step1[30] = step2[30];

  step1[31] = step2[31];

  // stage 6

  out[16] = _mm256_add_epi16(step1[16], step1[23]);

  out[17] = _mm256_add_epi16(step1[17], step1[22]);

  out[18] = _mm256_add_epi16(step1[18], step1[21]);

  out[19] = _mm256_add_epi16(step1[19], step1[20]);

  step2[20] = _mm256_sub_epi16(step1[19], step1[20]);

  step2[21] = _mm256_sub_epi16(step1[18], step1[21]);

  step2[22] = _mm256_sub_epi16(step1[17], step1[22]);

  step2[23] = _mm256_sub_epi16(step1[16], step1[23]);

  step2[24] = _mm256_sub_epi16(step1[31], step1[24]);

  step2[25] = _mm256_sub_epi16(step1[30], step1[25]);

  step2[26] = _mm256_sub_epi16(step1[29], step1[26]);

  step2[27] = _mm256_sub_epi16(step1[28], step1[27]);

  out[28] = _mm256_add_epi16(step1[27], step1[28]);

  out[29] = _mm256_add_epi16(step1[26], step1[29]);

  out[30] = _mm256_add_epi16(step1[25], step1[30]);

  out[31] = _mm256_add_epi16(step1[24], step1[31]);

  // stage 7

  butterfly16(step2[27], step2[20], cospi_16_64, cospi_16_64, &out[20],

              &out[27]);

  butterfly16(step2[26], step2[21], cospi_16_64, cospi_16_64, &out[21],

              &out[26]);

  butterfly16(step2[25], step2[22], cospi_16_64, cospi_16_64, &out[22],

              &out[25]);

  butterfly16(step2[24], step2[23], cospi_16_64, cospi_16_64, &out[23],

              &out[24]);

}

static INLINE void idct32_1024_16x32_quarter_1_2(__m256i *in, __m256i *out) {

  __m256i temp[16];

  // For each 16x32 block __m256i in[32],

  // Input with index, 0, 4, 8, 12, 16, 20, 24, 28

  // output pixels: 0-7 in __m256i out[32]

  idct32_1024_16x32_quarter_1(in, temp);

  // Input with index, 2, 6, 10, 14, 18, 22, 26, 30

  // output pixels: 8-15 in __m256i out[32]

  idct32_1024_16x32_quarter_2(in, temp);

  // stage 7

  add_sub_butterfly_avx2(temp, out, 16);

}

// For each 16x32 block __m256i in[32],

// Input with odd index,

// 1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31

// output pixels: 16-23, 24-31 in __m256i out[32]

static INLINE void idct32_1024_16x32_quarter_3_4(__m256i *in, __m256i *out) {

  __m256i step1[32], step2[32];

  // stage 1

  butterfly16(in[1], in[31], cospi_31_64, cospi_1_64, &step1[16], &step1[31]);

  butterfly16(in[17], in[15], cospi_15_64, cospi_17_64, &step1[17], &step1[30]);

  butterfly16(in[9], in[23], cospi_23_64, cospi_9_64, &step1[18], &step1[29]);

  butterfly16(in[25], in[7], cospi_7_64, cospi_25_64, &step1[19], &step1[28]);

  butterfly16(in[5], in[27], cospi_27_64, cospi_5_64, &step1[20], &step1[27]);

  butterfly16(in[21], in[11], cospi_11_64, cospi_21_64, &step1[21], &step1[26]);

  butterfly16(in[13], in[19], cospi_19_64, cospi_13_64, &step1[22], &step1[25]);

  butterfly16(in[29], in[3], cospi_3_64, cospi_29_64, &step1[23], &step1[24]);

  // stage 2

  step2[16] = _mm256_add_epi16(step1[16], step1[17]);

  step2[17] = _mm256_sub_epi16(step1[16], step1[17]);

  step2[18] = _mm256_sub_epi16(step1[19], step1[18]);

  step2[19] = _mm256_add_epi16(step1[19], step1[18]);

  step2[20] = _mm256_add_epi16(step1[20], step1[21]);

  step2[21] = _mm256_sub_epi16(step1[20], step1[21]);

  step2[22] = _mm256_sub_epi16(step1[23], step1[22]);

  step2[23] = _mm256_add_epi16(step1[23], step1[22]);

  step2[24] = _mm256_add_epi16(step1[24], step1[25]);

  step2[25] = _mm256_sub_epi16(step1[24], step1[25]);

  step2[26] = _mm256_sub_epi16(step1[27], step1[26]);

  step2[27] = _mm256_add_epi16(step1[27], step1[26]);

  step2[28] = _mm256_add_epi16(step1[28], step1[29]);

  step2[29] = _mm256_sub_epi16(step1[28], step1[29]);

  step2[30] = _mm256_sub_epi16(step1[31], step1[30]);

  step2[31] = _mm256_add_epi16(step1[31], step1[30]);

  // stage 3

  step1[16] = step2[16];

  step1[31] = step2[31];

  butterfly16(step2[30], step2[17], cospi_28_64, cospi_4_64, &step1[17],

              &step1[30]);

  butterfly16(step2[29], step2[18], -cospi_4_64, cospi_28_64, &step1[18],

              &step1[29]);

  step1[19] = step2[19];

  step1[20] = step2[20];

  butterfly16(step2[26], step2[21], cospi_12_64, cospi_20_64, &step1[21],

              &step1[26]);

  butterfly16(step2[25], step2[22], -cospi_20_64, cospi_12_64, &step1[22],

              &step1[25]);

  step1[23] = step2[23];

  step1[24] = step2[24];

  step1[27] = step2[27];

  step1[28] = step2[28];

  idct32_16x32_quarter_3_4_stage_4_to_7(step1, out);

}

static INLINE void idct32_1024_16x32(__m256i *in, __m256i *out) {

  __m256i temp[32];

  // For each 16x32 block __m256i in[32],

  // Input with index, 0, 4, 8, 12, 16, 20, 24, 28

  // output pixels: 0-7 in __m256i out[32]

  // AND

  // Input with index, 2, 6, 10, 14, 18, 22, 26, 30

  // output pixels: 8-15 in __m256i out[32]

  idct32_1024_16x32_quarter_1_2(in, temp);

  // For each 16x32 block __m256i in[32],

  // Input with odd index,

  // 1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31

  // output pixels: 16-23, 24-31 in __m256i out[32]

  idct32_1024_16x32_quarter_3_4(in, temp);

  // final stage

  add_sub_butterfly_avx2(temp, out, 32);

}

void vpx_idct32x32_1024_add_avx2(const tran_low_t *input, uint8_t *dest,

                                 int stride) {

  __m256i l[32], r[32], out[32], *in;

  int i;

  in = l;

  for (i = 0; i < 2; i++) {

    idct_load16x16(input, in, 32);

    transpose_16bit_16x16_avx2(in, in);

    idct_load16x16(input + 16, in + 16, 32);

    transpose_16bit_16x16_avx2(in + 16, in + 16);

    idct32_1024_16x32(in, in);

    in = r;

    input += 32 << 4;

  }

  for (i = 0; i < 32; i += 16) {

    transpose_16bit_16x16_avx2(l + i, out);

    transpose_16bit_16x16_avx2(r + i, out + 16);

    idct32_1024_16x32(out, out);

    store_buffer_16x32(out, dest, stride);

    dest += 16;

  }

}

// Case when only upper-left 16x16 has non-zero coeff

void vpx_idct32x32_135_add_avx2(const tran_low_t *input, uint8_t *dest,

                                int stride) {

  __m256i in[32], io[32], out[32];

  int i;

  for (i = 16; i < 32; i++) {

    in[i] = _mm256_setzero_si256();

  }

  // rows

  idct_load16x16(input, in, 32);

  transpose_16bit_16x16_avx2(in, in);

  idct32_1024_16x32(in, io);

  // columns

  for (i = 0; i < 32; i += 16) {

    transpose_16bit_16x16_avx2(io + i, in);

    idct32_1024_16x32(in, out);

    store_buffer_16x32(out, dest, stride);

    dest += 16;

  }

}