/*
 * 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.
 */

#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/gpu/GpuIndex.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/impl/IndexUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/gpu/utils/StaticUtils.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/gpu/utils/CopyUtils.cuh>
#include <faiss/gpu/utils/Float16.cuh>

#include <algorithm>
#include <limits>
#include <memory>

namespace faiss {
namespace gpu {

/// Default CPU search size for which we use paged copies
constexpr idx_t kMinPageSize = (idx_t)256 * 1024 * 1024;

/// Size above which we page copies from the CPU to GPU (non-paged
/// memory usage)
constexpr idx_t kNonPinnedPageSize = (idx_t)256 * 1024 * 1024;

// Default size for which we page add or search
constexpr idx_t kAddPageSize = (idx_t)256 * 1024 * 1024;

// Or, maximum number of vectors to consider per page of add or search
constexpr idx_t kAddVecSize = (idx_t)512 * 1024;

// Use a smaller search size, as precomputed code usage on IVFPQ
// requires substantial amounts of memory
// FIXME: parameterize based on algorithm need
constexpr idx_t kSearchVecSize = (idx_t)32 * 1024;

bool should_use_cuvs(GpuIndexConfig config_) {
    auto prop = getDeviceProperties(config_.device);

    if (prop.major < 7)
        return false;

    return config_.use_cuvs;
}

GpuIndex::GpuIndex(
        std::shared_ptr<GpuResources> resources,
        int dims,
        faiss::MetricType metric,
        float metricArg,
        GpuIndexConfig config)
        : Index(dims, metric),
          resources_(resources),
          config_(config),
          minPagedSize_(kMinPageSize) {
    FAISS_THROW_IF_NOT_FMT(
            config_.device < getNumDevices(),
            "Invalid GPU device %d",
            config_.device);

    FAISS_THROW_IF_NOT_MSG(dims > 0, "Invalid number of dimensions");

    FAISS_THROW_IF_NOT_FMT(
            config_.memorySpace == MemorySpace::Device ||
                    (config_.memorySpace == MemorySpace::Unified &&
                     getFullUnifiedMemSupport(config_.device)),
            "Device %d does not support full CUDA 8 Unified Memory (CC 6.0+)",
            config_.device);

    metric_arg = metricArg;

    FAISS_ASSERT((bool)resources_);
    resources_->initializeForDevice(config_.device);
}

int GpuIndex::getDevice() const {
    return config_.device;
}

void GpuIndex::copyFrom(const faiss::Index* index) {
    d = index->d;
    metric_type = index->metric_type;
    metric_arg = index->metric_arg;
    ntotal = index->ntotal;
    is_trained = index->is_trained;
}

void GpuIndex::copyTo(faiss::Index* index) const {
    index->d = d;
    index->metric_type = metric_type;
    index->metric_arg = metric_arg;
    index->ntotal = ntotal;
    index->is_trained = is_trained;
}

void GpuIndex::setMinPagingSize(size_t size) {
    minPagedSize_ = size;
}

size_t GpuIndex::getMinPagingSize() const {
    return minPagedSize_;
}

void GpuIndex::add_ex(idx_t n, const void* x, NumericType numeric_type) {
    add_with_ids_ex(n, x, numeric_type, nullptr);
}

void GpuIndex::add(idx_t n, const float* x) {
    // Pass to add_with_ids
    add_ex(n, x, NumericType::Float32);
}

void GpuIndex::add_with_ids_ex(
        idx_t n,
        const void* x,
        NumericType numeric_type,
        const idx_t* ids) {
    DeviceScope scope(config_.device);
    FAISS_THROW_IF_NOT_MSG(this->is_trained, "Index not trained");

    if (n == 0) {
        // nothing to add
        return;
    }

    std::vector<idx_t> generatedIds;

    // Generate IDs if we need them
    if (!ids && addImplRequiresIDs_()) {
        generatedIds = std::vector<idx_t>(n);

        for (idx_t i = 0; i < n; ++i) {
            generatedIds[i] = this->ntotal + i;
        }
    }

    addPaged_ex_(n, x, numeric_type, ids ? ids : generatedIds.data());
}

void GpuIndex::add_with_ids(idx_t n, const float* x, const idx_t* ids) {
    add_with_ids_ex(n, static_cast<const void*>(x), NumericType::Float32, ids);
}

void GpuIndex::addPaged_ex_(
        idx_t n,
        const void* x,
        NumericType numeric_type,
        const idx_t* ids) {
    if (n <= 0) {
        return;
    }

    auto dispatch = [&](auto dummy_type) {
        using data_t = decltype(dummy_type);
        const data_t* typed_x = reinterpret_cast<const data_t*>(x);

        idx_t totalSize = n * this->d * sizeof(data_t);

        if (!should_use_cuvs(config_) &&
            (totalSize > kAddPageSize || n > kAddVecSize)) {
            // How many vectors fit into kAddPageSize?
            idx_t maxNumVecsForPageSize =
                    kAddPageSize / (this->d * sizeof(data_t));

            // Always add at least 1 vector, if we have huge vectors
            maxNumVecsForPageSize = std::max(maxNumVecsForPageSize, idx_t(1));

            auto tileSize = std::min(n, maxNumVecsForPageSize);
            tileSize = std::min(tileSize, kSearchVecSize);

            for (idx_t i = 0; i < n; i += tileSize) {
                auto curNum = std::min(tileSize, n - i);
                addPage_ex_(
                        curNum,
                        static_cast<const void*>(typed_x + i * this->d),
                        numeric_type,
                        ids ? ids + i : nullptr);
            }
        } else {
            addPage_ex_(
                    n, static_cast<const void*>(typed_x), numeric_type, ids);
        }
    };

    if (numeric_type == NumericType::Float32) {
        dispatch(float{});
    } else if (numeric_type == NumericType::Float16) {
        dispatch(half{});
    } else if (numeric_type == NumericType::Int8) {
        dispatch(int8_t{});
    } else {
        FAISS_THROW_MSG("GpuIndex::addPaged_: Unsupported numeric type");
    }
}

void GpuIndex::addPaged_(idx_t n, const float* x, const idx_t* ids) {
    addPaged_ex_(n, static_cast<const void*>(x), NumericType::Float32, ids);
}

void GpuIndex::addPage_ex_(
        idx_t n,
        const void* x,
        NumericType numeric_type,
        const idx_t* ids) {
    // At this point, `x` can be resident on CPU or GPU, and `ids` may be
    // resident on CPU, GPU or may be null.
    //
    // Before continuing, we guarantee that all data will be resident on the
    // GPU.
    auto stream = resources_->getDefaultStreamCurrentDevice();

    auto dispatch = [&](auto dummy_type) {
        using data_t = decltype(dummy_type);

        auto vecs = toDeviceTemporary<data_t, 2>(
                resources_.get(),
                config_.device,
                const_cast<data_t*>(reinterpret_cast<const data_t*>(x)),
                stream,
                {n, this->d});

        if (ids) {
            auto indices = toDeviceTemporary<idx_t, 1>(
                    resources_.get(),
                    config_.device,
                    const_cast<idx_t*>(ids),
                    stream,
                    {n});

            addImpl_ex_(
                    n,
                    static_cast<const void*>(vecs.data()),
                    numeric_type,
                    ids ? indices.data() : nullptr);
        } else {
            addImpl_ex_(
                    n,
                    static_cast<const void*>(vecs.data()),
                    numeric_type,
                    nullptr);
        }
    };

    if (numeric_type == NumericType::Float32) {
        dispatch(float{});
    } else if (numeric_type == NumericType::Float16) {
        dispatch(half{});
    } else if (numeric_type == NumericType::Int8) {
        dispatch(int8_t{});
    } else {
        FAISS_THROW_MSG("GpuIndex::addPage_: Unsupported numeric type");
    }
}

void GpuIndex::addPage_(idx_t n, const float* x, const idx_t* ids) {
    addPage_ex_(n, static_cast<const void*>(x), NumericType::Float32, ids);
}

void GpuIndex::assign(idx_t n, const float* x, idx_t* labels, idx_t k) const {
    DeviceScope scope(config_.device);
    FAISS_THROW_IF_NOT_MSG(this->is_trained, "Index not trained");

    validateKSelect(k, should_use_cuvs(config_));

    auto stream = resources_->getDefaultStream(config_.device);

    // We need to create a throw-away buffer for distances, which we don't use
    // but which we do need for the search call
    DeviceTensor<float, 2, true> distances(
            resources_.get(), makeTempAlloc(AllocType::Other, stream), {n, k});

    // Forward to search
    search(n, x, k, distances.data(), labels);
}

void GpuIndex::search_ex(
        idx_t n,
        const void* x,
        NumericType numeric_type,
        idx_t k,
        float* distances,
        idx_t* labels,
        const SearchParameters* params) const {
    DeviceScope scope(config_.device);
    FAISS_THROW_IF_NOT_MSG(this->is_trained, "Index not trained");

    validateKSelect(k, should_use_cuvs(config_));

    if (n == 0 || k == 0) {
        // nothing to search
        return;
    }

    auto stream = resources_->getDefaultStream(config_.device);

    // We guarantee that the searchImpl_ will be called with device-resident
    // pointers.

    // The input vectors may be too large for the GPU, but we still
    // assume that the output distances and labels are not.
    // Go ahead and make space for output distances and labels on the
    // GPU.
    // If we reach a point where all inputs are too big, we can add
    // another level of tiling.
    auto outDistances = toDeviceTemporary<float, 2>(
            resources_.get(), config_.device, distances, stream, {n, k});

    auto outLabels = toDeviceTemporary<idx_t, 2>(
            resources_.get(), config_.device, labels, stream, {n, k});

    bool usePaged = false;

    if (getDeviceForAddress(x) == -1) {
        // It is possible that the user is querying for a vector set size
        // `x` that won't fit on the GPU.
        // In this case, we will have to handle paging of the data from CPU
        // -> GPU.
        // Currently, we don't handle the case where the output data won't
        // fit on the GPU (e.g., n * k is too large for the GPU memory).
        size_t dataSize =
                (size_t)n * this->d * get_numeric_type_size(numeric_type);

        if (dataSize >= minPagedSize_) {
            searchFromCpuPaged_ex_(
                    n,
                    x,
                    numeric_type,
                    k,
                    outDistances.data(),
                    outLabels.data(),
                    params);
            usePaged = true;
        }
    }

    if (!usePaged) {
        searchNonPaged_ex_(
                n,
                x,
                numeric_type,
                k,
                outDistances.data(),
                outLabels.data(),
                params);
    }

    // Copy back if necessary
    fromDevice<float, 2>(outDistances, distances, stream);
    fromDevice<idx_t, 2>(outLabels, labels, stream);
}

void GpuIndex::search(
        idx_t n,
        const float* x,
        idx_t k,
        float* distances,
        idx_t* labels,
        const SearchParameters* params) const {
    search_ex(
            n,
            static_cast<const void*>(x),
            NumericType::Float32,
            k,
            distances,
            labels,
            params);
}

void GpuIndex::search_and_reconstruct(
        idx_t n,
        const float* x,
        idx_t k,
        float* distances,
        idx_t* labels,
        float* recons,
        const SearchParameters* params) const {
    search(n, x, k, distances, labels, params);
    reconstruct_batch(n * k, labels, recons);
}

void GpuIndex::searchNonPaged_ex_(
        idx_t n,
        const void* x,
        NumericType numeric_type,
        int k,
        float* outDistancesData,
        idx_t* outIndicesData,
        const SearchParameters* params) const {
    auto stream = resources_->getDefaultStream(config_.device);

    // Make sure arguments are on the device we desire; use temporary
    // memory allocations to move it if necessary
    if (numeric_type == NumericType::Float32) {
        auto vecs = toDeviceTemporary<float, 2>(
                resources_.get(),
                config_.device,
                const_cast<float*>(static_cast<const float*>(x)),
                stream,
                {n, this->d});

        searchImpl_(
                n, vecs.data(), k, outDistancesData, outIndicesData, params);
    } else if (numeric_type == NumericType::Float16) {
        auto vecs = toDeviceTemporary<half, 2>(
                resources_.get(),
                config_.device,
                const_cast<half*>(static_cast<const half*>(x)),
                stream,
                {n, this->d});

        searchImpl_ex_(
                n,
                static_cast<const void*>(vecs.data()),
                numeric_type,
                k,
                outDistancesData,
                outIndicesData,
                params);
    } else if (numeric_type == NumericType::Int8) {
        auto vecs = toDeviceTemporary<int8_t, 2>(
                resources_.get(),
                config_.device,
                const_cast<int8_t*>(static_cast<const int8_t*>(x)),
                stream,
                {n, this->d});

        searchImpl_ex_(
                n,
                static_cast<const void*>(vecs.data()),
                numeric_type,
                k,
                outDistancesData,
                outIndicesData,
                params);
    } else {
        FAISS_THROW_MSG("GpuIndex::search: Unsupported numeric type");
    }
}

void GpuIndex::searchNonPaged_(
        idx_t n,
        const float* x,
        int k,
        float* outDistancesData,
        idx_t* outIndicesData,
        const SearchParameters* params) const {
    searchNonPaged_ex_(
            n,
            static_cast<const void*>(x),
            NumericType::Float32,
            k,
            outDistancesData,
            outIndicesData,
            params);
}

void GpuIndex::searchFromCpuPaged_ex_(
        idx_t n,
        const void* x,
        NumericType numeric_type,
        int k,
        float* outDistancesData,
        idx_t* outIndicesData,
        const SearchParameters* params) const {
    Tensor<float, 2, true> outDistances(outDistancesData, {n, k});
    Tensor<idx_t, 2, true> outIndices(outIndicesData, {n, k});

    // Is pinned memory available?
    auto pinnedAlloc = resources_->getPinnedMemory();
    idx_t pageSizeInVecs =
            ((pinnedAlloc.second / 2) /
             (get_numeric_type_size(numeric_type) * this->d));

    if (!pinnedAlloc.first || pageSizeInVecs < 1) {
        // Just page without overlapping copy with compute
        idx_t batchSize = utils::nextHighestPowerOf2(
                (kNonPinnedPageSize /
                 (get_numeric_type_size(numeric_type) * this->d)));

        for (idx_t cur = 0; cur < n; cur += batchSize) {
            auto num = std::min(batchSize, n - cur);

            auto outDistancesSlice = outDistances.narrowOutermost(cur, num);
            auto outIndicesSlice = outIndices.narrowOutermost(cur, num);
            if (numeric_type == NumericType::Float32) {
                searchNonPaged_ex_(
                        num,
                        static_cast<const void*>(
                                static_cast<const float*>(x) + cur * this->d),
                        numeric_type,
                        k,
                        outDistancesSlice.data(),
                        outIndicesSlice.data(),
                        params);
            } else if (numeric_type == NumericType::Float16) {
                searchNonPaged_ex_(
                        num,
                        static_cast<const void*>(
                                static_cast<const half*>(x) + cur * this->d),
                        numeric_type,
                        k,
                        outDistancesSlice.data(),
                        outIndicesSlice.data(),
                        params);
            } else if (numeric_type == NumericType::Int8) {
                searchNonPaged_ex_(
                        num,
                        static_cast<const void*>(
                                static_cast<const int8_t*>(x) + cur * this->d),
                        numeric_type,
                        k,
                        outDistancesSlice.data(),
                        outIndicesSlice.data(),
                        params);
            }
        }

        return;
    }

    //
    // Pinned memory is available, so we can overlap copy with compute.
    // We use two pinned memory buffers, and triple-buffer the
    // procedure:
    //
    // 1 CPU copy -> pinned
    // 2 pinned copy -> GPU
    // 3 GPU compute
    //
    // 1 2 3 1 2 3 ...   (pinned buf A)
    //   1 2 3 1 2 ...   (pinned buf B)
    //     1 2 3 1 ...   (pinned buf A)
    // time ->
    //
    auto defaultStream = resources_->getDefaultStream(config_.device);
    auto copyStream = resources_->getAsyncCopyStream(config_.device);

    auto dispatch = [&](auto dummy_type) {
        using data_t = decltype(dummy_type);

        data_t* bufPinnedA = (data_t*)pinnedAlloc.first;
        data_t* bufPinnedB = bufPinnedA + (size_t)pageSizeInVecs * this->d;
        data_t* bufPinned[2] = {bufPinnedA, bufPinnedB};

        // Reserve space on the GPU for the destination of the pinned buffer
        // copy
        DeviceTensor<data_t, 2, true> bufGpuA(
                resources_.get(),
                makeTempAlloc(AllocType::Other, defaultStream),
                {pageSizeInVecs, this->d});
        DeviceTensor<data_t, 2, true> bufGpuB(
                resources_.get(),
                makeTempAlloc(AllocType::Other, defaultStream),
                {pageSizeInVecs, this->d});
        DeviceTensor<data_t, 2, true>* bufGpus[2] = {&bufGpuA, &bufGpuB};

        // Copy completion events for the pinned buffers
        std::unique_ptr<CudaEvent> eventPinnedCopyDone[2];

        // Execute completion events for the GPU buffers
        std::unique_ptr<CudaEvent> eventGpuExecuteDone[2];

        // All offsets are in terms of number of vectors

        // Current start offset for buffer 1
        idx_t cur1 = 0;
        idx_t cur1BufIndex = 0;

        // Current start offset for buffer 2
        idx_t cur2 = -1;
        idx_t cur2BufIndex = 0;

        // Current start offset for buffer 3
        idx_t cur3 = -1;
        idx_t cur3BufIndex = 0;

        while (cur3 < n) {
            // Start async pinned -> GPU copy first (buf 2)
            if (cur2 != -1 && cur2 < n) {
                // Copy pinned to GPU
                auto numToCopy = std::min(pageSizeInVecs, n - cur2);

                // Make sure any previous execution has completed before
                // continuing
                auto& eventPrev = eventGpuExecuteDone[cur2BufIndex];
                if (eventPrev.get()) {
                    eventPrev->streamWaitOnEvent(copyStream);
                }

                CUDA_VERIFY(cudaMemcpyAsync(
                        bufGpus[cur2BufIndex]->data(),
                        bufPinned[cur2BufIndex],
                        numToCopy * this->d * sizeof(data_t),
                        cudaMemcpyHostToDevice,
                        copyStream));

                // Mark a completion event in this stream
                eventPinnedCopyDone[cur2BufIndex].reset(
                        new CudaEvent(copyStream));

                // We pick up from here
                cur3 = cur2;
                cur2 += numToCopy;
                cur2BufIndex = (cur2BufIndex == 0) ? 1 : 0;
            }

            if (cur3 != idx_t(-1) && cur3 < n) {
                // Process on GPU
                auto numToProcess = std::min(pageSizeInVecs, n - cur3);

                // Make sure the previous copy has completed before continuing
                auto& eventPrev = eventPinnedCopyDone[cur3BufIndex];
                FAISS_ASSERT(eventPrev.get());

                eventPrev->streamWaitOnEvent(defaultStream);

                // Create tensor wrappers
                // DeviceTensor<float, 2, true>
                // input(bufGpus[cur3BufIndex]->data(),
                //                                    {numToProcess, this->d});
                auto outDistancesSlice =
                        outDistances.narrowOutermost(cur3, numToProcess);
                auto outIndicesSlice =
                        outIndices.narrowOutermost(cur3, numToProcess);

                searchImpl_ex_(
                        numToProcess,
                        static_cast<const void*>(bufGpus[cur3BufIndex]->data()),
                        numeric_type,
                        k,
                        outDistancesSlice.data(),
                        outIndicesSlice.data(),
                        params);

                // Create completion event
                eventGpuExecuteDone[cur3BufIndex].reset(
                        new CudaEvent(defaultStream));

                // We pick up from here
                cur3BufIndex = (cur3BufIndex == 0) ? 1 : 0;
                cur3 += numToProcess;
            }

            if (cur1 < n) {
                // Copy CPU mem to CPU pinned
                auto numToCopy = std::min(pageSizeInVecs, n - cur1);

                // Make sure any previous copy has completed before continuing
                auto& eventPrev = eventPinnedCopyDone[cur1BufIndex];
                if (eventPrev.get()) {
                    eventPrev->cpuWaitOnEvent();
                }

                memcpy(bufPinned[cur1BufIndex],
                       static_cast<const data_t*>(x) + cur1 * this->d,
                       numToCopy * this->d *
                               get_numeric_type_size(numeric_type));

                // We pick up from here
                cur2 = cur1;
                cur1 += numToCopy;
                cur1BufIndex = (cur1BufIndex == 0) ? 1 : 0;
            }
        }
    };

    if (numeric_type == NumericType::Float32) {
        dispatch(float{});
    } else if (numeric_type == NumericType::Float16) {
        dispatch(half{});
    } else if (numeric_type == NumericType::Int8) {
        dispatch(int8_t{});
    } else {
        FAISS_THROW_MSG(
                "GpuIndex::searchFromCpuPaged_: Unsupported numeric type");
    }
}

void GpuIndex::searchFromCpuPaged_(
        idx_t n,
        const float* x,
        int k,
        float* outDistancesData,
        idx_t* outIndicesData,
        const SearchParameters* params) const {
    searchFromCpuPaged_ex_(
            n,
            static_cast<const void*>(x),
            NumericType::Float32,
            k,
            outDistancesData,
            outIndicesData,
            params);
}

void GpuIndex::compute_residual(const float* x, float* residual, idx_t key)
        const {
    FAISS_THROW_MSG("compute_residual not implemented for this type of index");
}

void GpuIndex::compute_residual_n(
        idx_t n,
        const float* xs,
        float* residuals,
        const idx_t* keys) const {
    FAISS_THROW_MSG(
            "compute_residual_n not implemented for this type of index");
}

std::shared_ptr<GpuResources> GpuIndex::getResources() {
    return resources_;
}

GpuIndex* tryCastGpuIndex(faiss::Index* index) {
    return dynamic_cast<GpuIndex*>(index);
}

bool isGpuIndex(faiss::Index* index) {
    return tryCastGpuIndex(index) != nullptr;
}

bool isGpuIndexImplemented(faiss::Index* index) {
#define CHECK_INDEX(TYPE)                 \
    do {                                  \
        if (dynamic_cast<TYPE*>(index)) { \
            return true;                  \
        }                                 \
    } while (false)

    CHECK_INDEX(faiss::IndexFlat);
    // FIXME: do we want recursive checking of the IVF quantizer?
    CHECK_INDEX(faiss::IndexIVFFlat);
    CHECK_INDEX(faiss::IndexIVFPQ);
    CHECK_INDEX(faiss::IndexIVFScalarQuantizer);

    return false;
}

} // namespace gpu

// This is the one defined in utils.cpp
extern std::string& ref_gpu_compile_options();

struct InitGpuCompileOptions {
    InitGpuCompileOptions() {
        ref_gpu_compile_options() = std::string("GPU ") +
#ifdef USE_NVIDIA_CUVS
                "NVIDIA_CUVS " +
#endif

#ifdef USE_AMD_ROCM
                "AMD_ROCM " +
#endif
                "";
    }
};

InitGpuCompileOptions InitGpuCompileOptions_instance;

} // namespace faiss
