/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <cstdint>
#include <string>

#include <immintrin.h>

#include <faiss/impl/platform_macros.h>
#include <faiss/impl/simdlib/simdlib.h>

#include <faiss/impl/simdlib/simdlib_avx2.h>

namespace faiss {

/** Simple wrapper around the AVX 512-bit registers
 *
 * The objective is to separate the different interpretations of the same
 * registers (as a vector of uint8, uint16 or uint32), to provide printing
 * functions, and to give more readable names to the AVX intrinsics. It does not
 * pretend to be exhaustive, functions are added as needed.
 */

/// 512-bit representation without interpretation as a vector
template <>
struct simd512bit_tpl<SIMDLevel::AVX512> {
    union {
        __m512i i;
        __m512 f;
    };

    simd512bit_tpl() {}

    explicit simd512bit_tpl(__m512i i) : i(i) {}

    explicit simd512bit_tpl(__m512 f) : f(f) {}

    explicit simd512bit_tpl(const void* x)
            : i(_mm512_loadu_si512((__m512i const*)x)) {}

    // sets up a lower half of the register while keeping upper one as zero
    explicit simd512bit_tpl(simd256bit_tpl<SIMDLevel::AVX2> lo)
            : simd512bit_tpl(_mm512_inserti32x8(
                      _mm512_castsi256_si512(lo.i),
                      _mm256_setzero_si256(),
                      1)) {}

    // constructs from lower and upper halves
    explicit simd512bit_tpl(
            simd256bit_tpl<SIMDLevel::AVX2> lo,
            simd256bit_tpl<SIMDLevel::AVX2> hi)
            : simd512bit_tpl(_mm512_inserti32x8(
                      _mm512_castsi256_si512(lo.i),
                      hi.i,
                      1)) {}

    void clear() {
        i = _mm512_setzero_si512();
    }

    void storeu(void* ptr) const {
        _mm512_storeu_si512((__m512i*)ptr, i);
    }

    void loadu(const void* ptr) {
        i = _mm512_loadu_si512((__m512i*)ptr);
    }

    void store(void* ptr) const {
        _mm512_storeu_si512((__m512i*)ptr, i);
    }

    void bin(char bits[513]) const {
        char bytes[64];
        storeu((void*)bytes);
        for (int i = 0; i < 512; i++) {
            bits[i] = '0' + ((bytes[i / 8] >> (i % 8)) & 1);
        }
        bits[512] = 0;
    }

    std::string bin() const {
        char bits[513];
        bin(bits);
        return std::string(bits);
    }
};

/// vector of 32 elements in uint16
template <>
struct simd32uint16_tpl<SIMDLevel::AVX512> : simd512bit_tpl<SIMDLevel::AVX512> {
    simd32uint16_tpl() {}

    explicit simd32uint16_tpl(__m512i i)
            : simd512bit_tpl<SIMDLevel::AVX512>(i) {}

    explicit simd32uint16_tpl(int x)
            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi16(x)) {}

    explicit simd32uint16_tpl(uint16_t x)
            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi16(x)) {}

    explicit simd32uint16_tpl(simd512bit_tpl<SIMDLevel::AVX512> x)
            : simd512bit_tpl<SIMDLevel::AVX512>(x) {}

    explicit simd32uint16_tpl(const uint16_t* x)
            : simd512bit_tpl<SIMDLevel::AVX512>((const void*)x) {}

    // sets up a lower half of the register
    explicit simd32uint16_tpl(simd256bit_tpl<SIMDLevel::AVX2> lo)
            : simd512bit_tpl<SIMDLevel::AVX512>(lo) {}

    // constructs from lower and upper halves
    explicit simd32uint16_tpl(
            simd256bit_tpl<SIMDLevel::AVX2> lo,
            simd256bit_tpl<SIMDLevel::AVX2> hi)
            : simd512bit_tpl<SIMDLevel::AVX512>(lo, hi) {}

    std::string elements_to_string(const char* fmt) const {
        uint16_t bytes[32];
        storeu((void*)bytes);
        char res[2000];
        char* ptr = res;
#if defined(__GNUC__) || defined(__clang__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wformat-nonliteral"
#endif
        for (int i = 0; i < 32; i++) {
            ptr += snprintf(
                    ptr, (size_t)(res + sizeof(res) - ptr), fmt, bytes[i]);
        }
#if defined(__GNUC__) || defined(__clang__)
#pragma GCC diagnostic pop
#endif
        // strip last ,
        ptr[-1] = 0;
        return std::string(res);
    }

    std::string hex() const {
        return elements_to_string("%02x,");
    }

    std::string dec() const {
        return elements_to_string("%3d,");
    }

    void set1(uint16_t x) {
        i = _mm512_set1_epi16((short)x);
    }

    simd32uint16_tpl operator*(const simd32uint16_tpl& other) const {
        return simd32uint16_tpl(_mm512_mullo_epi16(i, other.i));
    }

    // shift must be known at compile time
    simd32uint16_tpl operator>>(const int shift) const {
        return simd32uint16_tpl(_mm512_srli_epi16(i, shift));
    }

    // shift must be known at compile time
    simd32uint16_tpl operator<<(const int shift) const {
        return simd32uint16_tpl(_mm512_slli_epi16(i, shift));
    }

    simd32uint16_tpl& operator+=(simd32uint16_tpl other) {
        i = _mm512_add_epi16(i, other.i);
        return *this;
    }

    simd32uint16_tpl& operator-=(simd32uint16_tpl other) {
        i = _mm512_sub_epi16(i, other.i);
        return *this;
    }

    simd32uint16_tpl operator+(simd32uint16_tpl other) const {
        return simd32uint16_tpl(_mm512_add_epi16(i, other.i));
    }

    simd32uint16_tpl operator-(simd32uint16_tpl other) const {
        return simd32uint16_tpl(_mm512_sub_epi16(i, other.i));
    }

    simd32uint16_tpl operator&(simd512bit_tpl<SIMDLevel::AVX512> other) const {
        return simd32uint16_tpl(_mm512_and_si512(i, other.i));
    }

    simd32uint16_tpl operator|(simd512bit_tpl<SIMDLevel::AVX512> other) const {
        return simd32uint16_tpl(_mm512_or_si512(i, other.i));
    }

    simd32uint16_tpl operator^(simd512bit_tpl<SIMDLevel::AVX512> other) const {
        return simd32uint16_tpl(_mm512_xor_si512(i, other.i));
    }

    simd32uint16_tpl operator~() const {
        return simd32uint16_tpl(_mm512_xor_si512(i, _mm512_set1_epi32(-1)));
    }

    simd16uint16_tpl<SIMDLevel::AVX2> low() const {
        return simd16uint16_tpl<SIMDLevel::AVX2>(_mm512_castsi512_si256(i));
    }

    simd16uint16_tpl<SIMDLevel::AVX2> high() const {
        return simd16uint16_tpl<SIMDLevel::AVX2>(
                _mm512_extracti32x8_epi32(i, 1));
    }

    // for debugging only
    uint16_t operator[](int i) const {
        ALIGNED(64) uint16_t tab[32];
        store(tab);
        return tab[i];
    }

    void accu_min(simd32uint16_tpl incoming) {
        i = _mm512_min_epu16(i, incoming.i);
    }

    void accu_max(simd32uint16_tpl incoming) {
        i = _mm512_max_epu16(i, incoming.i);
    }
};

// decompose in 128-lanes: a = (a0, a1, a2, a3), b = (b0, b1, b2, b3)
// return (a0 + a1 + a2 + a3, b0 + b1 + b2 + b3)
inline simd16uint16_tpl<SIMDLevel::AVX2> combine4x2(
        simd32uint16_tpl<SIMDLevel::AVX512> a,
        simd32uint16_tpl<SIMDLevel::AVX512> b) {
    return combine2x2(a.low(), b.low()) + combine2x2(a.high(), b.high());
}

// vector of 32 unsigned 8-bit integers
template <>
struct simd64uint8_tpl<SIMDLevel::AVX512> : simd512bit_tpl<SIMDLevel::AVX512> {
    simd64uint8_tpl() {}

    explicit simd64uint8_tpl(__m512i i)
            : simd512bit_tpl<SIMDLevel::AVX512>(i) {}

    explicit simd64uint8_tpl(int x)
            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi8(x)) {}

    explicit simd64uint8_tpl(uint8_t x)
            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi8(x)) {}

    // sets up a lower half of the register
    explicit simd64uint8_tpl(simd256bit_tpl<SIMDLevel::AVX2> lo)
            : simd512bit_tpl<SIMDLevel::AVX512>(lo) {}

    // constructs from lower and upper halves
    explicit simd64uint8_tpl(
            simd256bit_tpl<SIMDLevel::AVX2> lo,
            simd256bit_tpl<SIMDLevel::AVX2> hi)
            : simd512bit_tpl<SIMDLevel::AVX512>(lo, hi) {}

    explicit simd64uint8_tpl(simd512bit_tpl<SIMDLevel::AVX512> x)
            : simd512bit_tpl<SIMDLevel::AVX512>(x) {}

    explicit simd64uint8_tpl(const uint8_t* x)
            : simd512bit_tpl<SIMDLevel::AVX512>((const void*)x) {}

    std::string elements_to_string(const char* fmt) const {
        uint8_t bytes[64];
        storeu((void*)bytes);
        char res[2000];
        char* ptr = res;
#if defined(__GNUC__) || defined(__clang__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wformat-nonliteral"
#endif
        for (int i = 0; i < 64; i++) {
            ptr += snprintf(
                    ptr, (size_t)(res + sizeof(res) - ptr), fmt, bytes[i]);
        }
#if defined(__GNUC__) || defined(__clang__)
#pragma GCC diagnostic pop
#endif
        // strip last ,
        ptr[-1] = 0;
        return std::string(res);
    }

    std::string hex() const {
        return elements_to_string("%02x,");
    }

    std::string dec() const {
        return elements_to_string("%3d,");
    }

    void set1(uint8_t x) {
        i = _mm512_set1_epi8((char)x);
    }

    simd64uint8_tpl operator&(simd512bit_tpl<SIMDLevel::AVX512> other) const {
        return simd64uint8_tpl(_mm512_and_si512(i, other.i));
    }

    simd64uint8_tpl operator+(simd64uint8_tpl other) const {
        return simd64uint8_tpl(_mm512_add_epi8(i, other.i));
    }

    simd64uint8_tpl lookup_4_lanes(simd64uint8_tpl idx) const {
        return simd64uint8_tpl(_mm512_shuffle_epi8(i, idx.i));
    }

    // extract + 0-extend lane
    // this operation is slow (3 cycles)
    simd32uint16_tpl<SIMDLevel::AVX512> lane0_as_uint16() const {
        __m256i x = _mm512_extracti32x8_epi32(i, 0);
        return simd32uint16_tpl<SIMDLevel::AVX512>(_mm512_cvtepu8_epi16(x));
    }

    simd32uint16_tpl<SIMDLevel::AVX512> lane1_as_uint16() const {
        __m256i x = _mm512_extracti32x8_epi32(i, 1);
        return simd32uint16_tpl<SIMDLevel::AVX512>(_mm512_cvtepu8_epi16(x));
    }

    simd64uint8_tpl& operator+=(simd64uint8_tpl other) {
        i = _mm512_add_epi8(i, other.i);
        return *this;
    }

    // for debugging only
    uint8_t operator[](int i) const {
        ALIGNED(64) uint8_t tab[64];
        store(tab);
        return tab[i];
    }
};

/// vector of 16 32-bit floats
template <>
struct simd16float32_tpl<SIMDLevel::AVX512>
        : simd512bit_tpl<SIMDLevel::AVX512> {
    simd16float32_tpl() {}

    explicit simd16float32_tpl(__m512 f)
            : simd512bit_tpl<SIMDLevel::AVX512>(f) {}

    explicit simd16float32_tpl(float x)
            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_ps(x)) {}

    explicit simd16float32_tpl(const float* x)
            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_loadu_ps(x)) {}

    void clear() {
        f = _mm512_setzero_ps();
    }

    void storeu(float* ptr) const {
        _mm512_storeu_ps(ptr, f);
    }

    void loadu(const float* ptr) {
        f = _mm512_loadu_ps(ptr);
    }

    void store(float* ptr) const {
        _mm512_storeu_ps(ptr, f);
    }

    simd16float32_tpl operator*(const simd16float32_tpl& other) const {
        return simd16float32_tpl(_mm512_mul_ps(f, other.f));
    }

    simd16float32_tpl operator+(const simd16float32_tpl& other) const {
        return simd16float32_tpl(_mm512_add_ps(f, other.f));
    }

    simd16float32_tpl operator-(const simd16float32_tpl& other) const {
        return simd16float32_tpl(_mm512_sub_ps(f, other.f));
    }

    simd16float32_tpl& operator+=(const simd16float32_tpl& other) {
        f = _mm512_add_ps(f, other.f);
        return *this;
    }

    std::string tostring() const {
        float tab[16];
        storeu(tab);
        char res[1000];
        char* ptr = res;
        for (int i = 0; i < 16; i++) {
            ptr += snprintf(
                    ptr, (size_t)(res + sizeof(res) - ptr), "%g,", tab[i]);
        }
        ptr[-1] = 0;
        return std::string(res);
    }
};

// compute a * b + c
inline simd16float32_tpl<SIMDLevel::AVX512> fmadd(
        const simd16float32_tpl<SIMDLevel::AVX512>& a,
        const simd16float32_tpl<SIMDLevel::AVX512>& b,
        const simd16float32_tpl<SIMDLevel::AVX512>& c) {
    return simd16float32_tpl<SIMDLevel::AVX512>(_mm512_fmadd_ps(a.f, b.f, c.f));
}

// horizontal add: sum all 16 floats in the register
inline float horizontal_add(const simd16float32_tpl<SIMDLevel::AVX512>& a) {
    return _mm512_reduce_add_ps(a.f);
}

} // namespace faiss
