/**
 * Copyright (c) Facebook, Inc. and its 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 <faiss/IndexScalarQuantizer.h>
#include <faiss/gpu/utils/ConversionOperators.cuh>
#include <faiss/gpu/utils/DeviceTensor.cuh>
#include <faiss/gpu/utils/HostTensor.cuh>
#include <faiss/gpu/utils/PtxUtils.cuh>
#include <faiss/gpu/utils/WarpShuffles.cuh>

namespace faiss {
namespace gpu {

inline bool isSQSupported(ScalarQuantizer::QuantizerType qtype) {
    switch (qtype) {
        case ScalarQuantizer::QuantizerType::QT_8bit:
        case ScalarQuantizer::QuantizerType::QT_8bit_uniform:
        case ScalarQuantizer::QuantizerType::QT_8bit_direct:
        case ScalarQuantizer::QuantizerType::QT_4bit:
        case ScalarQuantizer::QuantizerType::QT_4bit_uniform:
        case ScalarQuantizer::QuantizerType::QT_6bit:
        case ScalarQuantizer::QuantizerType::QT_fp16:
            return true;
        default:
            return false;
    }
}

// Wrapper around the CPU ScalarQuantizer that allows storage of parameters in
// GPU memory
struct GpuScalarQuantizer : public ScalarQuantizer {
    GpuScalarQuantizer(GpuResources* res, const ScalarQuantizer& sq)
            : ScalarQuantizer(sq),
              gpuTrained(DeviceTensor<float, 1, true>(
                      res,
                      makeDevAlloc(AllocType::Quantizer, 0),
                      {(int)sq.trained.size()})) {
        HostTensor<float, 1, true> cpuTrained(
                (float*)sq.trained.data(), {(int)sq.trained.size()});

        auto stream = res->getDefaultStreamCurrentDevice();
        gpuTrained.copyFrom(cpuTrained, stream);
    }

    // ScalarQuantizer::trained copied to GPU memory
    DeviceTensor<float, 1, true> gpuTrained;
};

//
// Quantizer codecs
//

// QT is the quantizer type implemented
// DimMultiple is the minimum guaranteed dimension multiple of the vectors
// encoded (used for ensuring alignment for memory load/stores)
template <int QT, int DimMultiple>
struct Codec {};

/////
//
// 32 bit encodings
// (does not use qtype)
//
/////

struct CodecFloat {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = 1;

    CodecFloat(int vecBytes) : bytesPerVec(vecBytes) {}

    size_t getSmemSize(int dim) {
        return 0;
    }
    inline __device__ void initKernel(float* smem, int dim) {}

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        float* p = (float*)&((uint8_t*)data)[vec * bytesPerVec];
        out[0] = p[d];
    }

    inline __device__ float decodePartial(void* data, int vec, int d, int subD)
            const {
        // doesn't need implementing (kDimPerIter == 1)
        return 0.0f;
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        float* p = (float*)&((uint8_t*)data)[vec * bytesPerVec];
        p[d] = v[0];
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining,
            float v[kDimPerIter]) const {
        // doesn't need implementing (kDimPerIter == 1)
    }

    //
    // new implementation
    //
    using EncodeT = float;
    static constexpr int kEncodeBits = 32;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return v;
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return v;
    }

    int bytesPerVec;
};

/////
//
// 16 bit encodings
//
/////

// Arbitrary dimension fp16
template <>
struct Codec<ScalarQuantizer::QuantizerType::QT_fp16, 1> {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = 1;

    Codec(int vecBytes) : bytesPerVec(vecBytes) {}

    size_t getSmemSize(int dim) {
        return 0;
    }
    inline __device__ void initKernel(float* smem, int dim) {}

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        half* p = (half*)&((uint8_t*)data)[vec * bytesPerVec];
        out[0] = Convert<half, float>()(p[d]);
    }

    inline __device__ float decodePartial(void* data, int vec, int d, int subD)
            const {
        // doesn't need implementing (kDimPerIter == 1)
        return 0.0f;
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        half* p = (half*)&((uint8_t*)data)[vec * bytesPerVec];
        p[d] = Convert<float, half>()(v[0]);
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining,
            float v[kDimPerIter]) const {
        // doesn't need implementing (kDimPerIter == 1)
    }

    //
    // new implementation
    //
    using EncodeT = half;
    static constexpr int kEncodeBits = 16;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return Convert<float, half>()(v);
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return Convert<half, float>()(v);
    }

    int bytesPerVec;
};

/////
//
// 8 bit encodings
//
/////

template <int DimPerIter>
struct Get8BitType {};

template <>
struct Get8BitType<1> {
    using T = uint8_t;
};

template <>
struct Get8BitType<2> {
    using T = uint16_t;
};

template <>
struct Get8BitType<4> {
    using T = uint32_t;
};

// Uniform quantization across all dimensions
template <int DimMultiple>
struct Codec<ScalarQuantizer::QuantizerType::QT_8bit_uniform, DimMultiple> {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = DimMultiple;
    using MemT = typename Get8BitType<DimMultiple>::T;

    Codec(int vecBytes, float min, float diff)
            : bytesPerVec(vecBytes), vmin(min), vdiff(diff) {}

    size_t getSmemSize(int dim) {
        return 0;
    }
    inline __device__ void initKernel(float* smem, int dim) {
        // We are performing vmin + vdiff * (v + 0.5) / (2^bits - 1)
        // This can be simplified to vmin' + vdiff' * v where:
        // vdiff' = vdiff / (2^bits - 1)
        // vmin' = vmin + 0.5 * vdiff'
        auto vd = vdiff * (1.0f / 255.0f);
        vmin = vmin + 0.5f * vd;
        vdiff = vd;
    }

    inline __device__ float decodeHelper(uint8_t v) const {
        return vmin + (float)v * vdiff;
    }

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        MemT* p = (MemT*)&((uint8_t*)data)[vec * bytesPerVec];
        MemT pv = p[d];

        uint8_t x[kDimPerIter];
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            x[i] = (uint8_t)((pv >> (i * 8)) & 0xffU);
        }

        float xDec[kDimPerIter];
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            xDec[i] = decodeHelper(x[i]);
        }

#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            out[i] = xDec[i];
        }
    }

    inline __device__ float decodePartial(void* data, int vec, int d, int subD)
            const {
        if (DimMultiple > 1) {
            // should not be called
            assert(false);
        }

        // otherwise does not need implementing
        return 0;
    }

    inline __device__ uint8_t encodeHelper(float v) const {
        float x = (v - vmin) / vdiff;
        x = fminf(1.0f, fmaxf(0.0f, x));
        return (uint8_t)(x * 255.0f);
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        MemT* p = (MemT*)&((uint8_t*)data)[vec * bytesPerVec];

        MemT x[kDimPerIter];
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            x[i] = encodeHelper(v[i]);
        }

        MemT out = 0;
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            out |= (x[i] << (i * 8));
        }

        p[d] = out;
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining,
            float v[kDimPerIter]) const {
        if (DimMultiple > 1) {
            // should not be called
            assert(false);
        }

        // otherwise does not need implementing
    }

    //
    // interleaved code implementation
    //
    using EncodeT = uint8_t;
    static constexpr int kEncodeBits = 8;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return encodeHelper(v);
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return decodeHelper(v);
    }

    int bytesPerVec;
    float vmin;
    float vdiff;
};

// Uniform quantization per each dimension
template <int DimMultiple>
struct Codec<ScalarQuantizer::QuantizerType::QT_8bit, DimMultiple> {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = DimMultiple;
    using MemT = typename Get8BitType<DimMultiple>::T;

    Codec(int vecBytes, float* min, float* diff)
            : bytesPerVec(vecBytes),
              vmin(min),
              vdiff(diff),
              smemVmin(nullptr),
              smemVdiff(nullptr) {}

    size_t getSmemSize(int dim) {
        return sizeof(float) * dim * 2;
    }

    // Initialize shared memory and local storage
    // It is up to the user to call a trailing syncthreads (after any other
    // initialization required has been done)
    inline __device__ void initKernel(float* smem, int dim) {
        smemVmin = smem;
        smemVdiff = smem + dim;

        for (int i = threadIdx.x; i < dim; i += blockDim.x) {
            // We are performing vmin + vdiff * (v + 0.5) / (2^bits - 1)
            // This can be simplified to vmin' + vdiff' * v where:
            // vdiff' = vdiff / (2^bits - 1)
            // vmin' = vmin + 0.5 * vdiff'
            auto vd = vdiff[i] * (1.0f / 255.0f);
            smemVmin[i] = vmin[i] + 0.5f * vd;
            smemVdiff[i] = vd;
        }
    }

    inline __device__ float decodeHelper(uint8_t v, int realDim) const {
        return smemVmin[realDim] + (float)v * smemVdiff[realDim];
    }

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        MemT* p = (MemT*)&((uint8_t*)data)[vec * bytesPerVec];
        MemT pv = p[d];
        int realDim = d * kDimPerIter;

        uint8_t x[kDimPerIter];
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            x[i] = (uint8_t)((pv >> (i * 8)) & 0xffU);
        }

        float xDec[kDimPerIter];
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            xDec[i] = decodeHelper(x[i], realDim + i);
        }

#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            out[i] = xDec[i];
        }
    }

    inline __device__ float decodePartial(void* data, int vec, int d, int subD)
            const {
        if (DimMultiple > 1) {
            // should not be called
            assert(false);
        }

        // otherwise does not need implementing
        return 0;
    }

    inline __device__ uint8_t encodeHelper(float v, int realDim) const {
        float x = (v - vmin[realDim]) / vdiff[realDim];
        x = fminf(1.0f, fmaxf(0.0f, x));
        return (uint8_t)(x * 255.0f);
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        MemT* p = (MemT*)&((uint8_t*)data)[vec * bytesPerVec];
        int realDim = d * kDimPerIter;

        MemT x[kDimPerIter];
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            x[i] = encodeHelper(v[i], realDim + i);
        }

        MemT out = 0;
#pragma unroll
        for (int i = 0; i < kDimPerIter; ++i) {
            out |= (x[i] << (i * 8));
        }

        p[d] = out;
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining,
            float v[kDimPerIter]) const {
        if (DimMultiple > 1) {
            // should not be called
            assert(false);
        }

        // otherwise does not need implementing
    }

    //
    // interleaved code implementation
    //
    using EncodeT = uint8_t;
    static constexpr int kEncodeBits = 8;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return encodeHelper(v, dim);
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return decodeHelper(v, dim);
    }

    int bytesPerVec;

    // gmem pointers
    const float* vmin;
    const float* vdiff;

    // smem pointers (configured in the kernel)
    float* smemVmin;
    float* smemVdiff;
};

template <>
struct Codec<ScalarQuantizer::QuantizerType::QT_8bit_direct, 1> {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = 1;

    Codec(int vecBytes) : bytesPerVec(vecBytes) {}

    size_t getSmemSize(int dim) {
        return 0;
    }
    inline __device__ void initKernel(float* smem, int dim) {}

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        out[0] = (float)p[d];
    }

    inline __device__ float decodePartial(void* data, int vec, int d, int subD)
            const {
        // doesn't need implementing (kDimPerIter == 1)
        return 0.0f;
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        p[d] = (uint8_t)v[0];
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining,
            float v[kDimPerIter]) const {
        // doesn't need implementing (kDimPerIter == 1)
    }

    //
    // interleaved code implementation
    //
    using EncodeT = uint8_t;
    static constexpr int kEncodeBits = 8;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return (uint8_t)v;
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return (float)v;
    }

    int bytesPerVec;
};

/////
//
// 6 bit encodings
//
/////

template <>
struct Codec<ScalarQuantizer::QuantizerType::QT_6bit, 1> {
    Codec(int vecBytes, float* min, float* diff)
            : bytesPerVec(vecBytes),
              vmin(min),
              vdiff(diff),
              smemVmin(nullptr),
              smemVdiff(nullptr) {}

    size_t getSmemSize(int dim) {
        return sizeof(float) * dim * 2;
    }

    // Initialize shared memory and local storage
    // It is up to the user to call a trailing syncthreads (after any other
    // initialization required has been done)
    inline __device__ void initKernel(float* smem, int dim) {
        smemVmin = smem;
        smemVdiff = smem + dim;

        for (int i = threadIdx.x; i < dim; i += blockDim.x) {
            // We are performing vmin + vdiff * (v + 0.5) / (2^bits - 1)
            // This can be simplified to vmin' + vdiff' * v where:
            // vdiff' = vdiff / (2^bits - 1)
            // vmin' = vmin + 0.5 * vdiff'
            auto vd = vdiff[i] * (1.0f / 63.0f);
            smemVmin[i] = vmin[i] + 0.5f * vd;
            smemVdiff[i] = vd;
        }
    }

    inline __device__ float decodeHelper(uint8_t v, int realDim) const {
        return smemVmin[realDim] + (float)v * smemVdiff[realDim];
    }

    inline __device__ uint8_t encodeHelper(float v, int realDim) const {
        float x = (v - vmin[realDim]) / vdiff[realDim];
        x = fminf(1.0f, fmaxf(0.0f, x));
        return (uint8_t)(x * 63.0f);
    }

    //
    // interleaved code implementation
    //
    using EncodeT = uint8_t;
    static constexpr int kEncodeBits = 6;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return encodeHelper(v, dim);
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return decodeHelper(v, dim);
    }

    int bytesPerVec;

    // gmem pointers
    const float* vmin;
    const float* vdiff;

    // smem pointers
    float* smemVmin;
    float* smemVdiff;
};

/////
//
// 4 bit encodings
//
/////

// Uniform quantization across all dimensions
template <>
struct Codec<ScalarQuantizer::QuantizerType::QT_4bit_uniform, 1> {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = 2;

    Codec(int vecBytes, float min, float diff)
            : bytesPerVec(vecBytes), vmin(min), vdiff(diff) {}

    size_t getSmemSize(int dim) {
        return 0;
    }
    inline __device__ void initKernel(float* smem, int dim) {
        // We are performing vmin + vdiff * (v + 0.5) / (2^bits - 1)
        // This can be simplified to vmin' + vdiff' * v where:
        // vdiff' = vdiff / (2^bits - 1)
        // vmin' = vmin + 0.5 * vdiff'
        auto vd = vdiff * (1.0f / 15.0f);
        vmin = vmin + 0.5f * vd;
        vdiff = vd;
    }

    inline __device__ float decodeHelper(uint8_t v) const {
        return vmin + (float)v * vdiff;
    }

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        uint8_t pv = p[d];

        out[0] = decodeHelper(pv & 0xf);
        out[1] = decodeHelper(pv >> 4);
    }

    inline __device__ float decodePartial(
            void* data,
            int vec,
            int d,
            int subD /* unused */) const {
        // We can only be called for a single input
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        uint8_t pv = p[d];

        return decodeHelper(pv & 0xf);
    }

    inline __device__ uint8_t encodeHelper(float v) const {
        float x = (v - vmin) / vdiff;
        x = fminf(1.0f, fmaxf(0.0f, x));
        return (uint8_t)(x * 15.0f);
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        p[d] = encodeHelper(v[0]) | (encodeHelper(v[1]) << 4);
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining, /* unused */
            float v[kDimPerIter]) const {
        // We can only be called for a single output
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        p[d] = encodeHelper(v[0]);
    }

    //
    // interleaved code implementation
    //
    using EncodeT = uint8_t;
    static constexpr int kEncodeBits = 4;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return encodeHelper(v);
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return decodeHelper(v);
    }

    int bytesPerVec;
    float vmin;
    float vdiff;
};

template <>
struct Codec<ScalarQuantizer::QuantizerType::QT_4bit, 1> {
    /// How many dimensions per iteration we are handling for encoding or
    /// decoding
    static constexpr int kDimPerIter = 2;

    Codec(int vecBytes, float* min, float* diff)
            : bytesPerVec(vecBytes),
              vmin(min),
              vdiff(diff),
              smemVmin(nullptr),
              smemVdiff(nullptr) {}

    size_t getSmemSize(int dim) {
        return sizeof(float) * dim * 2;
    }

    inline __device__ void initKernel(float* smem, int dim) {
        smemVmin = smem;
        smemVdiff = smem + dim;

        for (int i = threadIdx.x; i < dim; i += blockDim.x) {
            // We are performing vmin + vdiff * (v + 0.5) / (2^bits - 1)
            // This can be simplified to vmin' + vdiff' * v where:
            // vdiff' = vdiff / (2^bits - 1)
            // vmin' = vmin + 0.5 * vdiff'
            auto vd = vdiff[i] / 15.0f;
            smemVmin[i] = vmin[i] + 0.5f * vd;
            smemVdiff[i] = vd;
        }

        __syncthreads();
    }

    inline __device__ float decodeHelper(uint8_t v, int realDim) const {
        return smemVmin[realDim] + (float)v * smemVdiff[realDim];
    }

    inline __device__ void decode(void* data, int vec, int d, float* out)
            const {
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        uint8_t pv = p[d];
        int realDim = d * kDimPerIter;

        out[0] = decodeHelper(pv & 0xf, realDim);
        out[1] = decodeHelper(pv >> 4, realDim + 1);
    }

    inline __device__ float decodePartial(
            void* data,
            int vec,
            int d,
            int subD /* unused */) const {
        // We can only be called for a single input
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        uint8_t pv = p[d];
        int realDim = d * kDimPerIter;

        return decodeHelper(pv & 0xf, realDim);
    }

    inline __device__ uint8_t encodeHelper(float v, int realDim) const {
        float x = (v - vmin[realDim]) / vdiff[realDim];
        x = fminf(1.0f, fmaxf(0.0f, x));
        return (uint8_t)(x * 15.0f);
    }

    inline __device__ void encode(
            void* data,
            int vec,
            int d,
            float v[kDimPerIter]) const {
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        int realDim = d * kDimPerIter;
        p[d] = encodeHelper(v[0], realDim) |
                (encodeHelper(v[1], realDim + 1) << 4);
    }

    inline __device__ void encodePartial(
            void* data,
            int vec,
            int d,
            int remaining, /* unused */
            float v[kDimPerIter]) const {
        // We can only be called for a single output
        uint8_t* p = &((uint8_t*)data)[vec * bytesPerVec];
        int realDim = d * kDimPerIter;

        p[d] = encodeHelper(v[0], realDim);
    }

    //
    // interleaved code implementation
    //
    using EncodeT = uint8_t;
    static constexpr int kEncodeBits = 4;

    inline __device__ EncodeT encodeNew(int dim, float v) const {
        return encodeHelper(v, dim);
    }

    inline __device__ float decodeNew(int dim, EncodeT v) const {
        return decodeHelper(v, dim);
    }

    int bytesPerVec;

    // gmem pointers
    const float* vmin;
    const float* vdiff;

    // smem pointers
    float* smemVmin;
    float* smemVdiff;
};

} // namespace gpu
} // namespace faiss
