/*
* Copyright(c) 2021 Intel Corporation
*
* This source code is subject to the terms of the BSD 2 Clause License and
* the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
* was not distributed with this source code in the LICENSE file, you can
* obtain it at https://www.aomedia.org/license/software-license. If the Alliance for Open
* Media Patent License 1.0 was not distributed with this source code in the
* PATENTS file, you can obtain it at https://www.aomedia.org/license/patent-license.
*/

#include "definitions.h"

#if EN_AVX512_SUPPORT

#include <immintrin.h>

/* clang-format off */
DECLARE_ALIGNED(32, static const uint8_t, bilinear_filters_avx2[512]) = {
  16,  0, 16,  0, 16,  0, 16,  0, 16,  0, 16,  0, 16,  0, 16,  0,
  16,  0, 16,  0, 16,  0, 16,  0, 16,  0, 16,  0, 16,  0, 16,  0,
  14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2,
  14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2,
  12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4,
  12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4,
  10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6,
  10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6,
   8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,
   8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,  8,
   6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,
   6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,  6, 10,
   4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,
   4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,  4, 12,
   2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,
   2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,  2, 14,
};
/* clang-format on */

#define FILTER_SRC(filter)                                     \
    do {                                                       \
        /* filter the source */                                \
        exp_src_lo = _mm512_maddubs_epi16(exp_src_lo, filter); \
        exp_src_hi = _mm512_maddubs_epi16(exp_src_hi, filter); \
                                                               \
        /* add 8 to source */                                  \
        exp_src_lo = _mm512_add_epi16(exp_src_lo, pw8);        \
        exp_src_hi = _mm512_add_epi16(exp_src_hi, pw8);        \
                                                               \
        /* divide source by 16 */                              \
        exp_src_lo = _mm512_srai_epi16(exp_src_lo, 4);         \
        exp_src_hi = _mm512_srai_epi16(exp_src_hi, 4);         \
    } while (0)

#define MERGE_WITH_SRC(src_reg, reg)                     \
    do {                                                 \
        exp_src_lo = _mm512_unpacklo_epi8(src_reg, reg); \
        exp_src_hi = _mm512_unpackhi_epi8(src_reg, reg); \
    } while (0)

#define LOAD_SRC_DST                                         \
    do {                                                     \
        /* load source and destination */                    \
        src_reg = _mm512_loadu_si512((__m512i const*)(src)); \
        dst_reg = _mm512_loadu_si512((__m512i const*)(dst)); \
    } while (0)

#define AVG_NEXT_SRC(src_reg, size_stride)                                                           \
    do {                                                                                             \
        /* average between current and next stride source */                                         \
        src_reg = _mm512_avg_epu8(src_reg, _mm512_loadu_si512((__m512i const*)(src + size_stride))); \
    } while (0)

#define MERGE_NEXT_SRC(src_reg, size_stride)                                              \
    do {                                                                                  \
        MERGE_WITH_SRC(src_reg, _mm512_loadu_si512((__m512i const*)(src + size_stride))); \
    } while (0)

#define CALC_SUM_SSE_INSIDE_LOOP                                      \
    do {                                                              \
        /* expand each byte to 2 bytes */                             \
        __m512i exp_dst_lo = _mm512_unpacklo_epi8(dst_reg, zero_reg); \
        __m512i exp_dst_hi = _mm512_unpackhi_epi8(dst_reg, zero_reg); \
        /* source - dest */                                           \
        exp_src_lo = _mm512_sub_epi16(exp_src_lo, exp_dst_lo);        \
        exp_src_hi = _mm512_sub_epi16(exp_src_hi, exp_dst_hi);        \
        /* caculate sum */                                            \
        sum_reg    = _mm512_add_epi16(sum_reg, exp_src_lo);           \
        exp_src_lo = _mm512_madd_epi16(exp_src_lo, exp_src_lo);       \
        sum_reg    = _mm512_add_epi16(sum_reg, exp_src_hi);           \
        exp_src_hi = _mm512_madd_epi16(exp_src_hi, exp_src_hi);       \
        /* calculate sse */                                           \
        sse_reg = _mm512_add_epi32(sse_reg, exp_src_lo);              \
        sse_reg = _mm512_add_epi32(sse_reg, exp_src_hi);              \
    } while (0)

// final calculation to sse
#define CALC_SSE                                                                                                       \
    do {                                                                                                               \
        __m512i sse_reg_hi = _mm512_bsrli_epi128(sse_reg, 8);                                                          \
        sse_reg            = _mm512_add_epi32(sse_reg, sse_reg_hi);                                                    \
        sse_reg_hi         = _mm512_bsrli_epi128(sse_reg, 4);                                                          \
        sse_reg            = _mm512_add_epi32(sse_reg, sse_reg_hi);                                                    \
        __m256i sse_tmp256 = _mm256_add_epi32(_mm512_castsi512_si256(sse_reg), _mm512_extracti64x4_epi64(sse_reg, 1)); \
        *((int*)sse)       = _mm_cvtsi128_si32(_mm256_castsi256_si128(sse_tmp256)) +                                   \
            _mm_cvtsi128_si32(_mm256_extractf128_si256(sse_tmp256, 1));                                                \
    } while (0)

// final calculation to sum
#define CALC_SUM                                                                                                       \
    do {                                                                                                               \
        __m512i sum_reg_lo = _mm512_cvtepi16_epi32(_mm512_castsi512_si256(sum_reg));                                   \
        __m512i sum_reg_hi = _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(sum_reg, 1));                             \
        sum_reg            = _mm512_add_epi32(sum_reg_lo, sum_reg_hi);                                                 \
        sum_reg_hi         = _mm512_bsrli_epi128(sum_reg, 8);                                                          \
        sum_reg            = _mm512_add_epi32(sum_reg, sum_reg_hi);                                                    \
        sum_reg_hi         = _mm512_bsrli_epi128(sum_reg, 4);                                                          \
        sum_reg            = _mm512_add_epi32(sum_reg, sum_reg_hi);                                                    \
        __m256i sum_tmp256 = _mm256_add_epi32(_mm512_castsi512_si256(sum_reg), _mm512_extracti64x4_epi64(sum_reg, 1)); \
        sum                = _mm_cvtsi128_si32(_mm256_castsi256_si128(sum_tmp256)) +                                   \
            _mm_cvtsi128_si32(_mm256_extractf128_si256(sum_tmp256, 1));                                                \
    } while (0)

// Functions related to sub pixel variance width 16
#define LOAD_SRC_DST_INSERT(src_stride, dst_stride)                                               \
    do {                                                                                          \
        /* load source and destination of 2 rows and insert*/                                     \
        src_reg = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(src))), \
                                     _mm256_loadu_si256((__m256i*)(src + src_stride)),            \
                                     1);                                                          \
        dst_reg = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(dst))), \
                                     _mm256_loadu_si256((__m256i*)(dst + dst_stride)),            \
                                     1);                                                          \
    } while (0)

#define AVG_NEXT_SRC_INSERT(src_reg, size_stride)                                                         \
    do {                                                                                                  \
        /* average between current and next stride source */                                              \
        src_reg = _mm512_avg_epu8(                                                                        \
            src_reg,                                                                                      \
            _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(src + size_stride))), \
                               _mm256_loadu_si256((__m256i*)(src + (size_stride << 1))),                  \
                               1));                                                                       \
    } while (0)

#define MERGE_NEXT_SRC_INSERT(src_reg, size_stride)                                                                  \
    do {                                                                                                             \
        MERGE_WITH_SRC(src_reg,                                                                                      \
                       _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(src + size_stride))), \
                                          _mm256_loadu_si256((__m256i*)(src + (src_stride + size_stride))),          \
                                          1));                                                                       \
    } while (0)

#define LOAD_SRC_NEXT_BYTE_INSERT                                                                          \
    do {                                                                                                   \
        /* load source and another source from next row   */                                               \
        src_reg = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(src))),          \
                                     _mm256_loadu_si256((__m256i*)(src + src_stride)),                     \
                                     1);                                                                   \
        /* load source and next row source from 1 byte onwards   */                                        \
        src_next_reg = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(src + 1))), \
                                          _mm256_loadu_si256((__m256i*)(src + src_stride + 1)),            \
                                          1);                                                              \
    } while (0)

#define LOAD_DST_INSERT                                                                           \
    do {                                                                                          \
        dst_reg = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i*)(dst))), \
                                     _mm256_loadu_si256((__m256i*)(dst + dst_stride)),            \
                                     1);                                                          \
    } while (0)

#define LOAD_SRC_MERGE_FILTER_256BIT(filter)                                               \
    do {                                                                                   \
        __m256i src_reg_0     = _mm256_loadu_si256((__m256i*)(src));                       \
        __m256i src_reg_1     = _mm256_loadu_si256((__m256i*)(src + 1));                   \
        __m256i src_lo        = _mm256_unpacklo_epi8(src_reg_0, src_reg_1);                \
        __m256i src_hi        = _mm256_unpackhi_epi8(src_reg_0, src_reg_1);                \
        __m256i filter_256bit = _mm512_castsi512_si256(filter);                            \
        __m256i pw8_256bit    = _mm512_extracti64x4_epi64(pw8, 1);                         \
        LOAD_DST_INSERT;                                                                   \
        /* filter the source */                                                            \
        src_lo = _mm256_maddubs_epi16(src_lo, filter_256bit);                              \
        src_hi = _mm256_maddubs_epi16(src_hi, filter_256bit);                              \
                                                                                           \
        /* add 8 to source */                                                              \
        src_lo = _mm256_add_epi16(src_lo, pw8_256bit);                                     \
        src_hi = _mm256_add_epi16(src_hi, pw8_256bit);                                     \
                                                                                           \
        /* divide source by 16 */                                                          \
        src_lo       = _mm256_srai_epi16(src_lo, 4);                                       \
        src_hi       = _mm256_srai_epi16(src_hi, 4);                                       \
        src_reg_0    = _mm256_packus_epi16(src_lo, src_hi);                                \
        src_next_reg = _mm512_inserti64x4(                                                 \
            _mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_pack, 1)), src_reg_0, 1); \
    } while (0)

static INLINE unsigned int svt_aom_sub_pixel_variance64xh_avx512(const uint8_t* src, int src_stride, int x_offset,
                                                                 int y_offset, const uint8_t* dst, int dst_stride,
                                                                 int height, unsigned int* sse) {
    __m512i dst_reg, exp_src_lo, exp_src_hi;
    int     sum;
    __m512i sum_reg  = _mm512_set1_epi16(0);
    __m512i sse_reg  = _mm512_set1_epi16(0);
    __m512i zero_reg = _mm512_set1_epi16(0);

    // x_offset = 0 and y_offset = 0
    if (x_offset == 0) {
        if (y_offset == 0) {
            for (int i = 0; i < height; i++) {
                __m512i src_reg;
                LOAD_SRC_DST;
                // expend each byte to 2 bytes
                MERGE_WITH_SRC(src_reg, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += src_stride;
                dst += dst_stride;
            }
            // x_offset = 0 and y_offset = 4
        } else if (y_offset == 4) {
            for (int i = 0; i < height; i++) {
                __m512i src_reg;
                LOAD_SRC_DST;
                AVG_NEXT_SRC(src_reg, src_stride);
                // expend each byte to 2 bytes
                MERGE_WITH_SRC(src_reg, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += src_stride;
                dst += dst_stride;
            }
            // x_offset = 0 and y_offset = bilin interpolation
        } else {
            y_offset <<= 5;
            __m256i filter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + y_offset));
            __m512i filter    = _mm512_castsi256_si512(filter256);
            filter            = _mm512_inserti64x4(filter, filter256, 1);
            __m512i pw8       = _mm512_set1_epi16(8);
            for (int i = 0; i < height; i++) {
                __m512i src_reg;
                LOAD_SRC_DST;
                MERGE_NEXT_SRC(src_reg, src_stride);
                FILTER_SRC(filter);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += src_stride;
                dst += dst_stride;
            }
        }
        // x_offset = 4  and y_offset = 0
    } else if (x_offset == 4) {
        if (y_offset == 0) {
            for (int i = 0; i < height; i++) {
                __m512i src_reg;
                LOAD_SRC_DST;
                AVG_NEXT_SRC(src_reg, 1);
                // expand each byte to 2 bytes
                MERGE_WITH_SRC(src_reg, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += src_stride;
                dst += dst_stride;
            }
            // x_offset = 4  and y_offset = 4
        } else if (y_offset == 4) {
            // load source and another source starting from the next
            // following byte
            __m512i src_reg = _mm512_loadu_si512((__m512i const*)(src));
            AVG_NEXT_SRC(src_reg, 1);
            for (int i = 0; i < height; i++) {
                __m512i src_avg = src_reg;
                src += src_stride;
                LOAD_SRC_DST;
                AVG_NEXT_SRC(src_reg, 1);
                // average between previous average to current average
                src_avg = _mm512_avg_epu8(src_avg, src_reg);
                // expand each byte to 2 bytes
                MERGE_WITH_SRC(src_avg, zero_reg);
                // save current source average
                CALC_SUM_SSE_INSIDE_LOOP;
                dst += dst_stride;
            }
            // x_offset = 4  and y_offset = bilin interpolation
        } else {
            y_offset <<= 5;
            __m256i filter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + y_offset));
            __m512i filter    = _mm512_castsi256_si512(filter256);
            filter            = _mm512_inserti64x4(filter, filter256, 1);
            __m512i pw8       = _mm512_set1_epi16(8);
            // load source and another source starting from the next
            // following byte
            __m512i src_reg = _mm512_loadu_si512((__m512i const*)(src));
            AVG_NEXT_SRC(src_reg, 1);
            for (int i = 0; i < height; i++) {
                // save current source average
                __m512i src_avg = src_reg;
                src += src_stride;
                LOAD_SRC_DST;
                AVG_NEXT_SRC(src_reg, 1);
                MERGE_WITH_SRC(src_avg, src_reg);
                FILTER_SRC(filter);
                CALC_SUM_SSE_INSIDE_LOOP;
                dst += dst_stride;
            }
        }
        // x_offset = bilin interpolation and y_offset = 0
    } else {
        if (y_offset == 0) {
            x_offset <<= 5;
            __m256i filter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + x_offset));
            __m512i filter    = _mm512_castsi256_si512(filter256);
            filter            = _mm512_inserti64x4(filter, filter256, 1);

            //Fiter fix
            __m512i pw8 = _mm512_set1_epi16(8);
            for (int i = 0; i < height; i++) {
                __m512i src_reg;
                LOAD_SRC_DST;
                MERGE_NEXT_SRC(src_reg, 1);
                FILTER_SRC(filter);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += src_stride;
                dst += dst_stride;
            }
            // x_offset = bilin interpolation and y_offset = 4
        } else if (y_offset == 4) {
            x_offset <<= 5;
            __m256i filter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + x_offset));
            __m512i filter    = _mm512_castsi256_si512(filter256);
            filter            = _mm512_inserti64x4(filter, filter256, 1);
            __m512i pw8       = _mm512_set1_epi16(8);
            __m512i src_reg   = _mm512_loadu_si512((__m512i const*)(src));
            MERGE_NEXT_SRC(src_reg, 1);
            FILTER_SRC(filter);
            // convert each 16 bit to 8 bit to each low and high lane source
            __m512i src_pack = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
            for (int i = 0; i < height; i++) {
                src += src_stride;
                LOAD_SRC_DST;
                MERGE_NEXT_SRC(src_reg, 1);
                FILTER_SRC(filter);
                src_reg = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
                // average between previous pack to the current
                src_pack = _mm512_avg_epu8(src_pack, src_reg);
                MERGE_WITH_SRC(src_pack, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src_pack = src_reg;
                dst += dst_stride;
            }
            // x_offset = bilin interpolation and y_offset = bilin interpolation
        } else {
            x_offset <<= 5;
            y_offset <<= 5;

            __m256i xfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + x_offset));
            __m512i xfilter    = _mm512_castsi256_si512(xfilter256);
            xfilter            = _mm512_inserti64x4(xfilter, xfilter256, 1);

            __m256i yfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + y_offset));
            __m512i yfilter    = _mm512_castsi256_si512(yfilter256);
            yfilter            = _mm512_inserti64x4(yfilter, yfilter256, 1);

            __m512i pw8 = _mm512_set1_epi16(8);
            // load source and another source starting from the next
            // following byte
            __m512i src_reg = _mm512_loadu_si512((__m256i const*)(src));
            MERGE_NEXT_SRC(src_reg, 1);

            FILTER_SRC(xfilter);
            // convert each 16 bit to 8 bit to each low and high lane source
            __m512i src_pack = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
            for (int i = 0; i < height; i++) {
                src += src_stride;
                LOAD_SRC_DST;
                MERGE_NEXT_SRC(src_reg, 1);
                FILTER_SRC(xfilter);
                src_reg = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
                // merge previous pack to current pack source
                MERGE_WITH_SRC(src_pack, src_reg);
                // filter the source
                FILTER_SRC(yfilter);
                src_pack = src_reg;
                CALC_SUM_SSE_INSIDE_LOOP;
                dst += dst_stride;
            }
        }
    }

    CALC_SSE;
    CALC_SUM;
    return sum;
}

static INLINE unsigned int svt_aom_sub_pixel_variance32xh_avx512(const uint8_t* src, int src_stride, int x_offset,
                                                                 int y_offset, const uint8_t* dst, int dst_stride,
                                                                 int height, unsigned int* sse) {
    __m512i src_reg, dst_reg, exp_src_lo, exp_src_hi;
    int     sum;
    __m512i sum_reg  = _mm512_set1_epi16(0);
    __m512i sse_reg  = _mm512_set1_epi16(0);
    __m512i zero_reg = _mm512_set1_epi16(0);

    // x_offset = 0 and y_offset = 0
    if (x_offset == 0) {
        if (y_offset == 0) {
            for (int i = 0; i < height; i += 2) {
                LOAD_SRC_DST_INSERT(src_stride, dst_stride);
                // expend each byte to 2 bytes
                MERGE_WITH_SRC(src_reg, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += (src_stride << 1);
                dst += (dst_stride << 1);
            }
            // x_offset = 0 and y_offset = 4
        } else if (y_offset == 4) {
            for (int i = 0; i < height; i += 2) {
                LOAD_SRC_DST_INSERT(src_stride, dst_stride);
                AVG_NEXT_SRC_INSERT(src_reg, src_stride);
                // expend each byte to 2 bytes
                MERGE_WITH_SRC(src_reg, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += (src_stride << 1);
                dst += (dst_stride << 1);
            }
            // x_offset = 0 and y_offset = bilin interpolation
        } else {
            y_offset <<= 5;
            __m256i yfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + y_offset));
            __m512i filter     = _mm512_castsi256_si512(yfilter256);
            filter             = _mm512_inserti64x4(filter, yfilter256, 1);
            __m512i pw8        = _mm512_set1_epi16(8);
            for (int i = 0; i < height; i += 2) {
                LOAD_SRC_DST_INSERT(src_stride, dst_stride);
                MERGE_NEXT_SRC_INSERT(src_reg, src_stride);
                FILTER_SRC(filter);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += (src_stride << 1);
                dst += (dst_stride << 1);
            }
        }
        // x_offset = 4  and y_offset = 0
    } else if (x_offset == 4) {
        if (y_offset == 0) {
            for (int i = 0; i < height; i += 2) {
                __m512i src_next_reg;
                LOAD_SRC_NEXT_BYTE_INSERT;
                LOAD_DST_INSERT;
                /* average between current and next stride source */
                src_reg = _mm512_avg_epu8(src_reg, src_next_reg);
                // expand each byte to 2 bytes
                MERGE_WITH_SRC(src_reg, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += (src_stride << 1);
                dst += (dst_stride << 1);
            }
            // x_offset = 4  and y_offset = 4
        } else if (y_offset == 4) {
            __m512i src_next_reg;
            // load and insert source and next row source
            LOAD_SRC_NEXT_BYTE_INSERT;
            __m512i src_avg = _mm512_avg_epu8(src_reg, src_next_reg);
            src += src_stride << 1;
            for (int i = 0; i < height - 2; i += 2) {
                LOAD_SRC_NEXT_BYTE_INSERT;
                src_next_reg     = _mm512_avg_epu8(src_reg, src_next_reg);
                __m512i src_temp = _mm512_inserti64x4(_mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_avg, 1)),
                                                      _mm512_castsi512_si256(src_next_reg),
                                                      1);
                src_temp         = _mm512_avg_epu8(src_avg, src_temp);
                LOAD_DST_INSERT;
                // expand each byte to 2 bytes
                MERGE_WITH_SRC(src_temp, zero_reg);
                // save current source average
                src_avg = src_next_reg;
                CALC_SUM_SSE_INSIDE_LOOP;
                dst += dst_stride << 1;
                src += src_stride << 1;
            }
            // last 2 rows processing happens here
            __m256i src_reg_0 = _mm256_avg_epu8(_mm256_loadu_si256((__m256i*)(src)),
                                                _mm256_loadu_si256((__m256i*)(src + 1)));
            src_next_reg      = _mm512_inserti64x4(
                _mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_avg, 1)), src_reg_0, 1);
            LOAD_DST_INSERT;
            src_avg = _mm512_avg_epu8(src_avg, src_next_reg);
            MERGE_WITH_SRC(src_avg, zero_reg);
            CALC_SUM_SSE_INSIDE_LOOP;
        } else {
            // x_offset = 4  and y_offset = bilin interpolation
            y_offset <<= 5;
            __m256i yfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + y_offset));
            __m512i filter     = _mm512_castsi256_si512(yfilter256);
            filter             = _mm512_inserti64x4(filter, yfilter256, 1);
            __m512i pw8        = _mm512_set1_epi16(8);
            // load and insert source and next row source
            __m512i src_next_reg;
            LOAD_SRC_NEXT_BYTE_INSERT;
            __m512i src_avg = _mm512_avg_epu8(src_reg, src_next_reg);
            src += src_stride << 1;
            for (int i = 0; i < height - 2; i += 2) {
                LOAD_SRC_NEXT_BYTE_INSERT;
                src_next_reg     = _mm512_avg_epu8(src_reg, src_next_reg);
                __m512i src_temp = _mm512_inserti64x4(_mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_avg, 1)),
                                                      _mm512_castsi512_si256(src_next_reg),
                                                      1);
                LOAD_DST_INSERT;
                MERGE_WITH_SRC(src_avg, src_temp);
                // save current source average
                src_avg = src_next_reg;
                FILTER_SRC(filter);
                CALC_SUM_SSE_INSIDE_LOOP;
                dst += dst_stride << 1;
                src += src_stride << 1;
            }
            // last 2 rows processing happens here
            __m256i src_reg_0 = _mm256_loadu_si256((__m256i*)(src));
            __m256i src_reg_1 = _mm256_loadu_si256((__m256i*)(src + 1));
            src_reg_0         = _mm256_avg_epu8(src_reg_0, src_reg_1);
            src_next_reg      = _mm512_inserti64x4(
                _mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_avg, 1)), src_reg_0, 1);
            LOAD_DST_INSERT;
            MERGE_WITH_SRC(src_avg, src_next_reg);
            FILTER_SRC(filter);
            CALC_SUM_SSE_INSIDE_LOOP;
        }
        //  // x_offset = bilin interpolation and y_offset = 0
    } else {
        if (y_offset == 0) {
            __m512i filter, pw8;
            __m256i xfilter256;
            x_offset <<= 5;
            xfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + x_offset));
            filter     = _mm512_castsi256_si512(xfilter256);
            filter     = _mm512_inserti64x4(filter, xfilter256, 1);
            pw8        = _mm512_set1_epi16(8);
            for (int i = 0; i < height; i += 2) {
                LOAD_SRC_DST_INSERT(src_stride, dst_stride);
                MERGE_NEXT_SRC_INSERT(src_reg, 1);
                FILTER_SRC(filter);
                CALC_SUM_SSE_INSIDE_LOOP;
                src += (src_stride << 1);
                dst += (dst_stride << 1);
            }
            // x_offset = bilin interpolation and y_offset = 4
        } else if (y_offset == 4) {
            __m512i filter, pw8, src_next_reg, src_pack;
            __m256i xfilter256;
            x_offset <<= 5;
            xfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + x_offset));
            filter     = _mm512_castsi256_si512(xfilter256);
            filter     = _mm512_inserti64x4(filter, xfilter256, 1);
            pw8        = _mm512_set1_epi16(8);
            // load and insert source and next row source
            LOAD_SRC_NEXT_BYTE_INSERT;
            MERGE_WITH_SRC(src_reg, src_next_reg);
            FILTER_SRC(filter);
            // convert each 16 bit to 8 bit to each low and high lane source
            src_pack = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
            src += src_stride << 1;
            for (int i = 0; i < height - 2; i += 2) {
                LOAD_SRC_NEXT_BYTE_INSERT;
                LOAD_DST_INSERT;
                MERGE_WITH_SRC(src_reg, src_next_reg);
                FILTER_SRC(filter);
                src_reg      = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
                src_next_reg = _mm512_inserti64x4(
                    _mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_pack, 1)), _mm512_castsi512_si256(src_reg), 1);

                // average between previous pack to the current
                src_pack = _mm512_avg_epu8(src_pack, src_next_reg);
                MERGE_WITH_SRC(src_pack, zero_reg);
                CALC_SUM_SSE_INSIDE_LOOP;
                src_pack = src_reg;
                src += src_stride << 1;
                dst += dst_stride << 1;
            }
            // last 2 rows processing happens here
            LOAD_SRC_MERGE_FILTER_256BIT(filter);
            // average between previous pack to the current
            src_pack = _mm512_avg_epu8(src_pack, src_next_reg);
            MERGE_WITH_SRC(src_pack, zero_reg);
            CALC_SUM_SSE_INSIDE_LOOP;
        } else {
            // x_offset = bilin interpolation and y_offset = bilin interpolation
            __m512i xfilter, yfilter, pw8, src_next_reg, src_pack;
            __m256i xfilter256, yfilter256;
            x_offset <<= 5;
            y_offset <<= 5;

            xfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + x_offset));
            xfilter    = _mm512_castsi256_si512(xfilter256);
            xfilter    = _mm512_inserti64x4(xfilter, xfilter256, 1);

            yfilter256 = _mm256_load_si256((__m256i const*)(bilinear_filters_avx2 + y_offset));
            yfilter    = _mm512_castsi256_si512(yfilter256);
            yfilter    = _mm512_inserti64x4(yfilter, yfilter256, 1);
            pw8        = _mm512_set1_epi16(8);

            // load and insert source and next row source
            LOAD_SRC_NEXT_BYTE_INSERT;
            MERGE_WITH_SRC(src_reg, src_next_reg);
            FILTER_SRC(xfilter);
            // convert each 16 bit to 8 bit to each low and high lane source
            src_pack = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
            src += src_stride << 1;
            for (int i = 0; i < height - 2; i += 2) {
                LOAD_SRC_NEXT_BYTE_INSERT;
                LOAD_DST_INSERT;
                MERGE_WITH_SRC(src_reg, src_next_reg);
                FILTER_SRC(xfilter);
                src_reg      = _mm512_packus_epi16(exp_src_lo, exp_src_hi);
                src_next_reg = _mm512_inserti64x4(
                    _mm512_castsi256_si512(_mm512_extracti64x4_epi64(src_pack, 1)), _mm512_castsi512_si256(src_reg), 1);

                // average between previous pack to the current
                MERGE_WITH_SRC(src_pack, src_next_reg);
                // filter the source
                FILTER_SRC(yfilter);
                src_pack = src_reg;
                CALC_SUM_SSE_INSIDE_LOOP;
                src += src_stride << 1;
                dst += dst_stride << 1;
            }
            // last 2 rows processing happens here
            LOAD_SRC_MERGE_FILTER_256BIT(xfilter);

            MERGE_WITH_SRC(src_pack, src_next_reg);
            FILTER_SRC(yfilter);
            CALC_SUM_SSE_INSIDE_LOOP;
        }
    }

    CALC_SSE;
    CALC_SUM;
    return sum;
}

#define AOM_SUB_PIXEL_VAR_AVX512(w, h, wf, wlog2, hlog2)                                         \
    unsigned int svt_aom_sub_pixel_variance##w##x##h##_avx512(const uint8_t* src,                \
                                                              int            src_stride,         \
                                                              int            x_offset,           \
                                                              int            y_offset,           \
                                                              const uint8_t* dst,                \
                                                              int            dst_stride,         \
                                                              unsigned int*  sse_ptr) {           \
        /*Avoid overflow in helper by capping height.*/                                          \
        const int    hf  = AOMMIN(h, 64);                                                        \
        const int    wf2 = AOMMIN(wf, 128);                                                      \
        unsigned int sse = 0;                                                                    \
        int          se  = 0;                                                                    \
        for (int i = 0; i < (w / wf2); ++i) {                                                    \
            const uint8_t* src_ptr = src;                                                        \
            const uint8_t* dst_ptr = dst;                                                        \
            for (int j = 0; j < (h / hf); ++j) {                                                 \
                unsigned int sse2;                                                               \
                const int    se2 = svt_aom_sub_pixel_variance##wf##xh_avx512(                    \
                    src_ptr, src_stride, x_offset, y_offset, dst_ptr, dst_stride, hf, &sse2); \
                dst_ptr += hf * dst_stride;                                                      \
                src_ptr += hf * src_stride;                                                      \
                se += se2;                                                                       \
                sse += sse2;                                                                     \
            }                                                                                    \
            src += wf;                                                                           \
            dst += wf;                                                                           \
        }                                                                                        \
        *sse_ptr = sse;                                                                          \
        return sse - (unsigned int)(((int64_t)se * se) >> (wlog2 + hlog2));                      \
    }

AOM_SUB_PIXEL_VAR_AVX512(128, 128, 64, 7, 7);
AOM_SUB_PIXEL_VAR_AVX512(128, 64, 64, 7, 6);
AOM_SUB_PIXEL_VAR_AVX512(64, 128, 64, 6, 7);
AOM_SUB_PIXEL_VAR_AVX512(64, 64, 64, 6, 6);
AOM_SUB_PIXEL_VAR_AVX512(64, 32, 64, 6, 5);

AOM_SUB_PIXEL_VAR_AVX512(32, 64, 32, 5, 6);
AOM_SUB_PIXEL_VAR_AVX512(32, 32, 32, 5, 5);
AOM_SUB_PIXEL_VAR_AVX512(32, 16, 32, 5, 4);

static INLINE __m128i mm512_add_hi_lo_epi16(const __m512i val) {
    const __m256i val_lo = _mm512_castsi512_si256(val);
    const __m256i val_hi = _mm512_extracti64x4_epi64(val, 1);
    const __m256i val2   = _mm256_add_epi16(val_lo, val_hi);

    return _mm_add_epi16(_mm256_castsi256_si128(val2), _mm256_extractf128_si256(val2, 1));
}

static INLINE __m128i mm512_sum_to_128_epi32(const __m512i val) {
    const __m256i val_lo = _mm512_castsi512_si256(val);
    const __m256i val_hi = _mm512_extracti64x4_epi64(val, 1);
    const __m256i val2   = _mm256_add_epi32(val_lo, val_hi);

    return _mm_add_epi32(_mm256_castsi256_si128(val2), _mm256_extractf128_si256(val2, 1));
}

static INLINE int variance_final_from_32bit_sum_avx512(__m512i vsse, __m128i vsum, unsigned int* const sse) {
    // extract the low lane and add it to the high lane
    const __m128i sse_reg_128 = mm512_sum_to_128_epi32(vsse);

    // unpack sse and sum registers and add
    const __m128i sse_sum_lo = _mm_unpacklo_epi32(sse_reg_128, vsum);
    const __m128i sse_sum_hi = _mm_unpackhi_epi32(sse_reg_128, vsum);
    const __m128i sse_sum    = _mm_add_epi32(sse_sum_lo, sse_sum_hi);

    // perform the final summation and extract the results
    const __m128i res = _mm_add_epi32(sse_sum, _mm_srli_si128(sse_sum, 8));
    *((int*)sse)      = _mm_cvtsi128_si32(res);
    return _mm_extract_epi32(res, 1);
}

static INLINE __m512i sum_to_32bit_avx512(const __m512i sum) {
    const __m512i sum_lo = _mm512_cvtepi16_epi32(_mm512_castsi512_si256(sum));
    const __m512i sum_hi = _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(sum, 1));
    return _mm512_add_epi32(sum_lo, sum_hi);
}

static INLINE void variance_kernel_avx512(const __m512i src, const __m512i ref, __m512i* const sse,
                                          __m512i* const sum) {
    const __m512i adj_sub = _mm512_set1_epi16(0xff01); // (1,-1)

    // unpack into pairs of source and reference values
    const __m512i src_ref0 = _mm512_unpacklo_epi8(src, ref);
    const __m512i src_ref1 = _mm512_unpackhi_epi8(src, ref);

    // subtract adjacent elements using src*1 + ref*-1
    const __m512i diff0 = _mm512_maddubs_epi16(src_ref0, adj_sub);
    const __m512i diff1 = _mm512_maddubs_epi16(src_ref1, adj_sub);
    const __m512i madd0 = _mm512_madd_epi16(diff0, diff0);
    const __m512i madd1 = _mm512_madd_epi16(diff1, diff1);

    // add to the running totals
    *sum = _mm512_add_epi16(*sum, _mm512_add_epi16(diff0, diff1));
    *sse = _mm512_add_epi32(*sse, _mm512_add_epi32(madd0, madd1));
}

static INLINE int variance_final_512_avx512(__m512i vsse, __m512i vsum, unsigned int* const sse) {
    // extract the low lane and add it to the high lane
    const __m128i vsum_128  = mm512_add_hi_lo_epi16(vsum);
    const __m128i vsum_64   = _mm_add_epi16(vsum_128, _mm_srli_si128(vsum_128, 8));
    const __m128i sum_int32 = _mm_cvtepi16_epi32(vsum_64);
    return variance_final_from_32bit_sum_avx512(vsse, sum_int32, sse);
}

static INLINE int variance_final_1024_avx512(__m512i vsse, __m512i vsum, unsigned int* const sse) {
    // extract the low lane and add it to the high lane
    const __m128i vsum_128 = mm512_add_hi_lo_epi16(vsum);
    const __m128i vsum_64  = _mm_add_epi32(_mm_cvtepi16_epi32(vsum_128),
                                          _mm_cvtepi16_epi32(_mm_srli_si128(vsum_128, 8)));
    return variance_final_from_32bit_sum_avx512(vsse, vsum_64, sse);
}

static INLINE int variance_final_2048_avx512(__m512i vsse, __m512i vsum, unsigned int* const sse) {
    vsum                   = sum_to_32bit_avx512(vsum);
    const __m128i vsum_128 = mm512_sum_to_128_epi32(vsum);
    return variance_final_from_32bit_sum_avx512(vsse, vsum_128, sse);
}

static INLINE void variance32_avx512(const uint8_t* src, const int src_stride, const uint8_t* ref, const int ref_stride,
                                     const int h, __m512i* const vsse, __m512i* const vsum) {
    *vsum = _mm512_setzero_si512();

    for (int i = 0; i < (h >> 1); i++) {
        const __m512i s = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i const*)src)),
                                             _mm256_loadu_si256((__m256i const*)(src + src_stride)),
                                             1);
        const __m512i r = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((__m256i const*)ref)),
                                             _mm256_loadu_si256((__m256i const*)(ref + ref_stride)),
                                             1);

        variance_kernel_avx512(s, r, vsse, vsum);

        src += (src_stride << 1);
        ref += (ref_stride << 1);
    }
}

static INLINE void variance64_avx512(const uint8_t* src, const int src_stride, const uint8_t* ref, const int ref_stride,
                                     const int h, __m512i* const vsse, __m512i* const vsum) {
    *vsum = _mm512_setzero_si512();

    for (int i = 0; i < h; i++) {
        const __m512i s = _mm512_loadu_si512((__m512i const*)(src));
        const __m512i r = _mm512_loadu_si512((__m512i const*)(ref));
        variance_kernel_avx512(s, r, vsse, vsum);

        src += src_stride;
        ref += ref_stride;
    }
}

static INLINE void variance128_avx512(const uint8_t* src, const int src_stride, const uint8_t* ref,
                                      const int ref_stride, const int h, __m512i* const vsse, __m512i* const vsum) {
    *vsum = _mm512_setzero_si512();

    for (int i = 0; i < h; i++) {
        __m512i s = _mm512_loadu_si512((__m512i const*)(src));
        __m512i r = _mm512_loadu_si512((__m512i const*)(ref));
        variance_kernel_avx512(s, r, vsse, vsum);

        s = _mm512_loadu_si512((__m512i const*)(src + 64));
        r = _mm512_loadu_si512((__m512i const*)(ref + 64));
        variance_kernel_avx512(s, r, vsse, vsum);

        src += src_stride;
        ref += ref_stride;
    }
}

#define AOM_VAR_LOOP_AVX512(bw, bh, bits, uh)                                                        \
    unsigned int svt_aom_variance##bw##x##bh##_avx512(                                               \
        const uint8_t* src, int src_stride, const uint8_t* ref, int ref_stride, unsigned int* sse) { \
        __m512i vsse = _mm512_setzero_si512();                                                       \
        __m512i vsum = _mm512_setzero_si512();                                                       \
        for (int i = 0; i < (bh / uh); i++) {                                                        \
            __m512i vsum16;                                                                          \
            variance##bw##_avx512(src, src_stride, ref, ref_stride, uh, &vsse, &vsum16);             \
            vsum = _mm512_add_epi32(vsum, sum_to_32bit_avx512(vsum16));                              \
            src += uh * src_stride;                                                                  \
            ref += uh * ref_stride;                                                                  \
        }                                                                                            \
        const __m128i vsum_128 = mm512_sum_to_128_epi32(vsum);                                       \
        const int     sum      = variance_final_from_32bit_sum_avx512(vsse, vsum_128, sse);          \
        return *sse - (unsigned int)(((int64_t)sum * sum) >> bits);                                  \
    }

AOM_VAR_LOOP_AVX512(64, 64, 12, 32); // 64x32 * ( 64/32)
AOM_VAR_LOOP_AVX512(64, 128, 13, 32); // 64x32 * (128/32)
AOM_VAR_LOOP_AVX512(128, 64, 13, 16); // 128x16 * ( 64/16)
AOM_VAR_LOOP_AVX512(128, 128, 14, 16); // 128x16 * (128/16)

#define AOM_VAR_NO_LOOP_AVX512(bw, bh, bits, max_pixel)                                              \
    unsigned int svt_aom_variance##bw##x##bh##_avx512(                                               \
        const uint8_t* src, int src_stride, const uint8_t* ref, int ref_stride, unsigned int* sse) { \
        __m512i vsse = _mm512_setzero_si512();                                                       \
        __m512i vsum = _mm512_setzero_si512();                                                       \
        variance##bw##_avx512(src, src_stride, ref, ref_stride, bh, &vsse, &vsum);                   \
        const int sum = variance_final_##max_pixel##_avx512(vsse, vsum, sse);                        \
        return *sse - (uint32_t)(((int64_t)sum * sum) >> bits);                                      \
    }

AOM_VAR_NO_LOOP_AVX512(32, 8, 8, 512);
AOM_VAR_NO_LOOP_AVX512(32, 16, 9, 512);
AOM_VAR_NO_LOOP_AVX512(32, 32, 10, 1024);
AOM_VAR_NO_LOOP_AVX512(32, 64, 11, 2048);

AOM_VAR_NO_LOOP_AVX512(64, 16, 10, 1024);
AOM_VAR_NO_LOOP_AVX512(64, 32, 11, 2048);

#endif // !EN_AVX512_SUPPORT
