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

#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/gpu/utils/StaticUtils.h>
#include <faiss/gpu/impl/IVFUtils.cuh>
#include <faiss/gpu/impl/PQCodeDistances.cuh>
#include <faiss/gpu/impl/PQCodeLoad.cuh>
#include <faiss/gpu/utils/ConversionOperators.cuh>
#include <faiss/gpu/utils/DeviceTensor.cuh>
#include <faiss/gpu/utils/Float16.cuh>
#include <faiss/gpu/utils/HostTensor.cuh>
#include <faiss/gpu/utils/LoadStoreOperators.cuh>
#include <faiss/gpu/utils/NoTypeTensor.cuh>
#include <faiss/gpu/utils/WarpPackedBits.cuh>

namespace faiss {
namespace gpu {

// A basic implementation that works for the interleaved by vector layout for
// any number of sub-quantizers
template <typename EncodeT, int EncodeBits, typename CodeDistanceT>
__global__ void pqScanInterleaved(
        Tensor<float, 2, true> queries,
        Tensor<float, 3, true> pqCentroids,
        Tensor<int, 2, true> topQueryToCentroid,
        Tensor<CodeDistanceT, 4, true> codeDistances,
        void** listCodes,
        int* listLengths,
        Tensor<int, 2, true> prefixSumOffsets,
        Tensor<float, 1, true> distance) {
    // Each block handles a single query
    auto queryId = blockIdx.y;
    auto probeId = blockIdx.x;

    auto listId = topQueryToCentroid[queryId][probeId];
    // Safety guard in case NaNs in input cause no list ID to be generated
    if (listId == -1) {
        return;
    }

    int numWarps = blockDim.x / kWarpSize;
    // FIXME: some issue with getLaneId() and CUDA 10.1 and P4 GPUs?
    int laneId = threadIdx.x % kWarpSize;
    int warpId = threadIdx.x / kWarpSize;

    int numSubQuantizers = codeDistances.getSize(2);

    // This is where we start writing out data
    // We ensure that before the array (at offset -1), there is a 0 value
    int outBase = *(prefixSumOffsets[queryId][probeId].data() - 1);
    float* distanceOut = distance[outBase].data();
    auto localCodeDistances = codeDistances[queryId][probeId];

    // This is where the codes for our list start
    auto vecsBase = (EncodeT*)listCodes[listId];
    int numVecs = listLengths[listId];

    // How many vector blocks of 32 are in this list?
    int numBlocks = utils::divUp(numVecs, 32);

    // Number of EncodeT words per each dimension of block of 32 vecs
    constexpr int bytesPerVectorBlockDim = EncodeBits * 32 / 8;
    constexpr int wordsPerVectorBlockDim =
            bytesPerVectorBlockDim / sizeof(EncodeT);
    int wordsPerVectorBlock = wordsPerVectorBlockDim * numSubQuantizers;

    for (int block = warpId; block < numBlocks; block += numWarps) {
        float dist = 0;

        // This is the vector a given lane/thread handles
        int vec = block * kWarpSize + laneId;
        bool valid = vec < numVecs;

        EncodeT* data = vecsBase + block * wordsPerVectorBlock;

        for (int sq = 0; sq < numSubQuantizers; ++sq) {
            EncodeT enc =
                    WarpPackedBits<EncodeT, EncodeBits>::read(laneId, data);
            EncodeT code =
                    WarpPackedBits<EncodeT, EncodeBits>::postRead(laneId, enc);

            dist += valid ? ConvertTo<float>::to(localCodeDistances[sq][code])
                          : 0;
            data += wordsPerVectorBlockDim;
        }

        if (valid) {
            distanceOut[vec] = dist;
        }
    }
}

template <typename LookupT, typename LookupVecT>
struct LoadCodeDistances {
    static inline __device__ void load(
            LookupT* smem,
            LookupT* codes,
            int numCodes) {
        constexpr int kWordSize = sizeof(LookupVecT) / sizeof(LookupT);

        // We can only use the vector type if the data is guaranteed to be
        // aligned. The codes are innermost, so if it is evenly divisible,
        // then any slice will be aligned.
        if (numCodes % kWordSize == 0) {
            // Load the data by float4 for efficiency, and then handle any
            // remainder limitVec is the number of whole vec words we can load,
            // in terms of whole blocks performing the load
            constexpr int kUnroll = 2;
            int limitVec = numCodes / (kUnroll * kWordSize * blockDim.x);
            limitVec *= kUnroll * blockDim.x;

            LookupVecT* smemV = (LookupVecT*)smem;
            LookupVecT* codesV = (LookupVecT*)codes;

            for (int i = threadIdx.x; i < limitVec; i += kUnroll * blockDim.x) {
                LookupVecT vals[kUnroll];

#pragma unroll
                for (int j = 0; j < kUnroll; ++j) {
                    vals[j] = LoadStore<LookupVecT>::load(
                            &codesV[i + j * blockDim.x]);
                }

#pragma unroll
                for (int j = 0; j < kUnroll; ++j) {
                    LoadStore<LookupVecT>::store(
                            &smemV[i + j * blockDim.x], vals[j]);
                }
            }

            // This is where we start loading the remainder that does not evenly
            // fit into kUnroll x blockDim.x
            int remainder = limitVec * kWordSize;

            for (int i = remainder + threadIdx.x; i < numCodes;
                 i += blockDim.x) {
                smem[i] = codes[i];
            }
        } else {
            // Potential unaligned load
            constexpr int kUnroll = 4;

            int limit = utils::roundDown(numCodes, kUnroll * blockDim.x);

            int i = threadIdx.x;
            for (; i < limit; i += kUnroll * blockDim.x) {
                LookupT vals[kUnroll];

#pragma unroll
                for (int j = 0; j < kUnroll; ++j) {
                    vals[j] = codes[i + j * blockDim.x];
                }

#pragma unroll
                for (int j = 0; j < kUnroll; ++j) {
                    smem[i + j * blockDim.x] = vals[j];
                }
            }

            for (; i < numCodes; i += blockDim.x) {
                smem[i] = codes[i];
            }
        }
    }
};

template <int NumSubQuantizers, typename LookupT, typename LookupVecT>
__global__ void pqScanNoPrecomputedMultiPass(
        Tensor<float, 2, true> queries,
        Tensor<float, 3, true> pqCentroids,
        Tensor<int, 2, true> topQueryToCentroid,
        Tensor<LookupT, 4, true> codeDistances,
        void** listCodes,
        int* listLengths,
        Tensor<int, 2, true> prefixSumOffsets,
        Tensor<float, 1, true> distance) {
    const auto codesPerSubQuantizer = pqCentroids.getSize(2);

    // Where the pq code -> residual distance is stored
    extern __shared__ char smemCodeDistances[];
    LookupT* codeDist = (LookupT*)smemCodeDistances;

    // Each block handles a single query
    auto queryId = blockIdx.y;
    auto probeId = blockIdx.x;

    // This is where we start writing out data
    // We ensure that before the array (at offset -1), there is a 0 value
    int outBase = *(prefixSumOffsets[queryId][probeId].data() - 1);
    float* distanceOut = distance[outBase].data();

    auto listId = topQueryToCentroid[queryId][probeId];
    // Safety guard in case NaNs in input cause no list ID to be generated
    if (listId == -1) {
        return;
    }

    uint8_t* codeList = (uint8_t*)listCodes[listId];
    int limit = listLengths[listId];

    constexpr int kNumCode32 =
            NumSubQuantizers <= 4 ? 1 : (NumSubQuantizers / 4);
    unsigned int code32[kNumCode32];
    unsigned int nextCode32[kNumCode32];

    // We double-buffer the code loading, which improves memory utilization
    if (threadIdx.x < limit) {
        LoadCode32<NumSubQuantizers>::load(code32, codeList, threadIdx.x);
    }

    LoadCodeDistances<LookupT, LookupVecT>::load(
            codeDist,
            codeDistances[queryId][probeId].data(),
            codeDistances.getSize(2) * codeDistances.getSize(3));

    // Prevent WAR dependencies
    __syncthreads();

    // Each thread handles one code element in the list, with a
    // block-wide stride
    for (int codeIndex = threadIdx.x; codeIndex < limit;
         codeIndex += blockDim.x) {
        // Prefetch next codes
        if (codeIndex + blockDim.x < limit) {
            LoadCode32<NumSubQuantizers>::load(
                    nextCode32, codeList, codeIndex + blockDim.x);
        }

        float dist = 0.0f;

#pragma unroll
        for (int word = 0; word < kNumCode32; ++word) {
            constexpr int kBytesPerCode32 =
                    NumSubQuantizers < 4 ? NumSubQuantizers : 4;

            if (kBytesPerCode32 == 1) {
                auto code = code32[0];
                dist = ConvertTo<float>::to(codeDist[code]);

            } else {
#pragma unroll
                for (int byte = 0; byte < kBytesPerCode32; ++byte) {
                    auto code = getByte(code32[word], byte * 8, 8);

                    auto offset = codesPerSubQuantizer *
                            (word * kBytesPerCode32 + byte);

                    dist += ConvertTo<float>::to(codeDist[offset + code]);
                }
            }
        }

        // Write out intermediate distance result
        // We do not maintain indices here, in order to reduce global
        // memory traffic. Those are recovered in the final selection step.
        distanceOut[codeIndex] = dist;

        // Rotate buffers
#pragma unroll
        for (int word = 0; word < kNumCode32; ++word) {
            code32[word] = nextCode32[word];
        }
    }
}

template <typename CentroidT>
void runMultiPassTile(
        GpuResources* res,
        Tensor<float, 2, true>& queries,
        Tensor<CentroidT, 2, true>& centroids,
        Tensor<float, 3, true>& pqCentroidsInnermostCode,
        NoTypeTensor<4, true>& codeDistances,
        Tensor<float, 2, true>& coarseDistances,
        Tensor<int, 2, true>& coarseIndices,
        Tensor<uint8_t, 1, true>& bitset,
        bool useFloat16Lookup,
        bool useMMCodeDistance,
        bool interleavedCodeLayout,
        int bitsPerSubQuantizer,
        int numSubQuantizers,
        int numSubQuantizerCodes,
        thrust::device_vector<void*>& listCodes,
        thrust::device_vector<void*>& listIndices,
        IndicesOptions indicesOptions,
        thrust::device_vector<int>& listLengths,
        Tensor<char, 1, true>& thrustMem,
        Tensor<int, 2, true>& prefixSumOffsets,
        Tensor<float, 1, true>& allDistances,
        Tensor<float, 3, true>& heapDistances,
        Tensor<int, 3, true>& heapIndices,
        int k,
        faiss::MetricType metric,
        Tensor<float, 2, true>& outDistances,
        Tensor<Index::idx_t, 2, true>& outIndices,
        cudaStream_t stream) {
    // We only support two metrics at the moment
    FAISS_ASSERT(
            metric == MetricType::METRIC_INNER_PRODUCT ||
            metric == MetricType::METRIC_L2);

    bool l2Distance = metric == MetricType::METRIC_L2;

    // Calculate offset lengths, so we know where to write out
    // intermediate results
    runCalcListOffsets(
            res,
            coarseIndices,
            listLengths,
            prefixSumOffsets,
            thrustMem,
            stream);

    // Calculate residual code distances, since this is without
    // precomputed codes
    runPQCodeDistances(
            res,
            pqCentroidsInnermostCode,
            queries,
            centroids,
            coarseDistances,
            coarseIndices,
            codeDistances,
            useMMCodeDistance,
            l2Distance,
            useFloat16Lookup,
            stream);

    if (interleavedCodeLayout) {
        // The vector interleaved layout implementation
        auto kThreadsPerBlock = 256;

        auto grid = dim3(coarseIndices.getSize(1), coarseIndices.getSize(0));
        auto block = dim3(kThreadsPerBlock);

#define RUN_INTERLEAVED(BITS_PER_CODE, CODE_DIST_T)            \
    do {                                                       \
        pqScanInterleaved<uint8_t, BITS_PER_CODE, CODE_DIST_T> \
                <<<grid, block, 0, stream>>>(                  \
                        queries,                               \
                        pqCentroidsInnermostCode,              \
                        coarseIndices,                         \
                        codeDistancesT,                        \
                        listCodes.data().get(),                \
                        listLengths.data().get(),              \
                        prefixSumOffsets,                      \
                        allDistances);                         \
    } while (0)

        if (useFloat16Lookup) {
            auto codeDistancesT = codeDistances.toTensor<half>();

            switch (bitsPerSubQuantizer) {
                case 4: {
                    RUN_INTERLEAVED(4, half);
                } break;
                case 5: {
                    RUN_INTERLEAVED(5, half);
                } break;
                case 6: {
                    RUN_INTERLEAVED(6, half);
                } break;
                case 8: {
                    RUN_INTERLEAVED(8, half);
                } break;
                default:
                    FAISS_ASSERT(false);
                    break;
            }
        } else {
            auto codeDistancesT = codeDistances.toTensor<float>();

            switch (bitsPerSubQuantizer) {
                case 4: {
                    RUN_INTERLEAVED(4, float);
                } break;
                case 5: {
                    RUN_INTERLEAVED(5, float);
                } break;
                case 6: {
                    RUN_INTERLEAVED(6, float);
                } break;
                case 8: {
                    RUN_INTERLEAVED(8, float);
                } break;
                default:
                    FAISS_ASSERT(false);
                    break;
            }
        }
    } else {
        // Convert all codes to a distance, and write out (distance,
        // index) values for all intermediate results
        auto kThreadsPerBlock = 256;

        auto grid = dim3(coarseIndices.getSize(1), coarseIndices.getSize(0));
        auto block = dim3(kThreadsPerBlock);

        // pq centroid distances
        auto smem = useFloat16Lookup ? sizeof(half) : sizeof(float);

        smem *= numSubQuantizers * numSubQuantizerCodes;
        FAISS_ASSERT(smem <= getMaxSharedMemPerBlockCurrentDevice());

#define RUN_PQ_OPT(NUM_SUB_Q, LOOKUP_T, LOOKUP_VEC_T)                   \
    do {                                                                \
        auto codeDistancesT = codeDistances.toTensor<LOOKUP_T>();       \
                                                                        \
        pqScanNoPrecomputedMultiPass<NUM_SUB_Q, LOOKUP_T, LOOKUP_VEC_T> \
                <<<grid, block, smem, stream>>>(                        \
                        queries,                                        \
                        pqCentroidsInnermostCode,                       \
                        coarseIndices,                                  \
                        codeDistancesT,                                 \
                        listCodes.data().get(),                         \
                        listLengths.data().get(),                       \
                        prefixSumOffsets,                               \
                        allDistances);                                  \
    } while (0)

#define RUN_PQ(NUM_SUB_Q)                         \
    do {                                          \
        if (useFloat16Lookup) {                   \
            RUN_PQ_OPT(NUM_SUB_Q, half, Half8);   \
        } else {                                  \
            RUN_PQ_OPT(NUM_SUB_Q, float, float4); \
        }                                         \
    } while (0)

        switch (numSubQuantizers) {
            case 1:
                RUN_PQ(1);
                break;
            case 2:
                RUN_PQ(2);
                break;
            case 3:
                RUN_PQ(3);
                break;
            case 4:
                RUN_PQ(4);
                break;
            case 8:
                RUN_PQ(8);
                break;
            case 12:
                RUN_PQ(12);
                break;
            case 16:
                RUN_PQ(16);
                break;
            case 20:
                RUN_PQ(20);
                break;
            case 24:
                RUN_PQ(24);
                break;
            case 28:
                RUN_PQ(28);
                break;
            case 32:
                RUN_PQ(32);
                break;
            case 40:
                RUN_PQ(40);
                break;
            case 48:
                RUN_PQ(48);
                break;
            case 56:
                RUN_PQ(56);
                break;
            case 64:
                RUN_PQ(64);
                break;
            case 96:
                RUN_PQ(96);
                break;
            default:
                FAISS_ASSERT(false);
                break;
        }

#undef RUN_PQ
#undef RUN_PQ_OPT
    }

    CUDA_TEST_ERROR();

    // k-select the output in chunks, to increase parallelism
    runPass1SelectLists(
            listIndices,
            indicesOptions,
            prefixSumOffsets,
            coarseIndices,
            bitset,
            allDistances,
            coarseIndices.getSize(1),
            k,
            !l2Distance, // L2 distance chooses smallest
            heapDistances,
            heapIndices,
            stream);

    // k-select final output
    auto flatHeapDistances = heapDistances.downcastInner<2>();
    auto flatHeapIndices = heapIndices.downcastInner<2>();

    runPass2SelectLists(
            flatHeapDistances,
            flatHeapIndices,
            listIndices,
            indicesOptions,
            prefixSumOffsets,
            coarseIndices,
            k,
            !l2Distance, // L2 distance chooses smallest
            outDistances,
            outIndices,
            stream);
}

template <typename CentroidT>
void runPQScanMultiPassNoPrecomputed(
        Tensor<float, 2, true>& queries,
        Tensor<CentroidT, 2, true>& centroids,
        Tensor<float, 3, true>& pqCentroidsInnermostCode,
        Tensor<float, 2, true>& coarseDistances,
        Tensor<int, 2, true>& coarseIndices,
        Tensor<uint8_t, 1, true>& bitset,
        bool useFloat16Lookup,
        bool useMMCodeDistance,
        bool interleavedCodeLayout,
        int bitsPerSubQuantizer,
        int numSubQuantizers,
        int numSubQuantizerCodes,
        thrust::device_vector<void*>& listCodes,
        thrust::device_vector<void*>& listIndices,
        IndicesOptions indicesOptions,
        thrust::device_vector<int>& listLengths,
        int maxListLength,
        int k,
        faiss::MetricType metric,
        // output
        Tensor<float, 2, true>& outDistances,
        // output
        Tensor<Index::idx_t, 2, true>& outIndices,
        GpuResources* res) {
    constexpr int kMinQueryTileSize = 8;
    constexpr int kMaxQueryTileSize = 128;
    constexpr int kThrustMemSize = 16384;

    int nprobe = coarseIndices.getSize(1);

    auto stream = res->getDefaultStreamCurrentDevice();

    // Make a reservation for Thrust to do its dirty work (global memory
    // cross-block reduction space); hopefully this is large enough.
    DeviceTensor<char, 1, true> thrustMem1(
            res, makeTempAlloc(AllocType::Other, stream), {kThrustMemSize});
    DeviceTensor<char, 1, true> thrustMem2(
            res, makeTempAlloc(AllocType::Other, stream), {kThrustMemSize});
    DeviceTensor<char, 1, true>* thrustMem[2] = {&thrustMem1, &thrustMem2};

    // How much temporary storage is available?
    // If possible, we'd like to fit within the space available.
    size_t sizeAvailable = res->getTempMemoryAvailableCurrentDevice();

    // We run two passes of heap selection
    // This is the size of the first-level heap passes
    constexpr int kNProbeSplit = 8;
    int pass2Chunks = std::min(nprobe, kNProbeSplit);

    size_t sizeForFirstSelectPass =
            pass2Chunks * k * (sizeof(float) + sizeof(int));

    // How much temporary storage we need per each query
    size_t sizePerQuery = 2 *                         // streams
            ((nprobe * sizeof(int) + sizeof(int)) +   // prefixSumOffsets
             nprobe * maxListLength * sizeof(float) + // allDistances
             // residual distances
             nprobe * numSubQuantizers * numSubQuantizerCodes * sizeof(float) +
             sizeForFirstSelectPass);

    int queryTileSize = (int)(sizeAvailable / sizePerQuery);

    if (queryTileSize < kMinQueryTileSize) {
        queryTileSize = kMinQueryTileSize;
    } else if (queryTileSize > kMaxQueryTileSize) {
        queryTileSize = kMaxQueryTileSize;
    }

    // FIXME: we should adjust queryTileSize to deal with this, since
    // indexing is in int32
    FAISS_ASSERT(
            queryTileSize * nprobe * maxListLength <
            std::numeric_limits<int>::max());

    // Temporary memory buffers
    // Make sure there is space prior to the start which will be 0, and
    // will handle the boundary condition without branches
    DeviceTensor<int, 1, true> prefixSumOffsetSpace1(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize * nprobe + 1});
    DeviceTensor<int, 1, true> prefixSumOffsetSpace2(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize * nprobe + 1});

    DeviceTensor<int, 2, true> prefixSumOffsets1(
            prefixSumOffsetSpace1[1].data(), {queryTileSize, nprobe});
    DeviceTensor<int, 2, true> prefixSumOffsets2(
            prefixSumOffsetSpace2[1].data(), {queryTileSize, nprobe});
    DeviceTensor<int, 2, true>* prefixSumOffsets[2] = {
            &prefixSumOffsets1, &prefixSumOffsets2};

    // Make sure the element before prefixSumOffsets is 0, since we
    // depend upon simple, boundary-less indexing to get proper results
    CUDA_VERIFY(cudaMemsetAsync(
            prefixSumOffsetSpace1.data(), 0, sizeof(int), stream));
    CUDA_VERIFY(cudaMemsetAsync(
            prefixSumOffsetSpace2.data(), 0, sizeof(int), stream));

    int codeDistanceTypeSize = useFloat16Lookup ? sizeof(half) : sizeof(float);

    int totalCodeDistancesSize = queryTileSize * nprobe * numSubQuantizers *
            numSubQuantizerCodes * codeDistanceTypeSize;

    DeviceTensor<char, 1, true> codeDistances1Mem(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {totalCodeDistancesSize});
    NoTypeTensor<4, true> codeDistances1(
            codeDistances1Mem.data(),
            codeDistanceTypeSize,
            {queryTileSize, nprobe, numSubQuantizers, numSubQuantizerCodes});

    DeviceTensor<char, 1, true> codeDistances2Mem(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {totalCodeDistancesSize});
    NoTypeTensor<4, true> codeDistances2(
            codeDistances2Mem.data(),
            codeDistanceTypeSize,
            {queryTileSize, nprobe, numSubQuantizers, numSubQuantizerCodes});

    NoTypeTensor<4, true>* codeDistances[2] = {
            &codeDistances1, &codeDistances2};

    DeviceTensor<float, 1, true> allDistances1(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize * nprobe * maxListLength});
    DeviceTensor<float, 1, true> allDistances2(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize * nprobe * maxListLength});
    DeviceTensor<float, 1, true>* allDistances[2] = {
            &allDistances1, &allDistances2};

    DeviceTensor<float, 3, true> heapDistances1(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize, pass2Chunks, k});
    DeviceTensor<float, 3, true> heapDistances2(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize, pass2Chunks, k});
    DeviceTensor<float, 3, true>* heapDistances[2] = {
            &heapDistances1, &heapDistances2};

    DeviceTensor<int, 3, true> heapIndices1(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize, pass2Chunks, k});
    DeviceTensor<int, 3, true> heapIndices2(
            res,
            makeTempAlloc(AllocType::Other, stream),
            {queryTileSize, pass2Chunks, k});
    DeviceTensor<int, 3, true>* heapIndices[2] = {&heapIndices1, &heapIndices2};

    auto streams = res->getAlternateStreamsCurrentDevice();
    streamWait(streams, {stream});

    int curStream = 0;

    for (int query = 0; query < queries.getSize(0); query += queryTileSize) {
        int numQueriesInTile =
                std::min(queryTileSize, queries.getSize(0) - query);

        auto prefixSumOffsetsView =
                prefixSumOffsets[curStream]->narrowOutermost(
                        0, numQueriesInTile);

        auto codeDistancesView =
                codeDistances[curStream]->narrowOutermost(0, numQueriesInTile);
        auto coarseDistancesView =
                coarseDistances.narrowOutermost(query, numQueriesInTile);
        auto coarseIndicesView =
                coarseIndices.narrowOutermost(query, numQueriesInTile);
        auto queryView = queries.narrowOutermost(query, numQueriesInTile);

        auto heapDistancesView =
                heapDistances[curStream]->narrowOutermost(0, numQueriesInTile);
        auto heapIndicesView =
                heapIndices[curStream]->narrowOutermost(0, numQueriesInTile);

        auto outDistanceView =
                outDistances.narrowOutermost(query, numQueriesInTile);
        auto outIndicesView =
                outIndices.narrowOutermost(query, numQueriesInTile);

        runMultiPassTile(
                res,
                queryView,
                centroids,
                pqCentroidsInnermostCode,
                codeDistancesView,
                coarseDistancesView,
                coarseIndicesView,
                bitset,
                useFloat16Lookup,
                useMMCodeDistance,
                interleavedCodeLayout,
                bitsPerSubQuantizer,
                numSubQuantizers,
                numSubQuantizerCodes,
                listCodes,
                listIndices,
                indicesOptions,
                listLengths,
                *thrustMem[curStream],
                prefixSumOffsetsView,
                *allDistances[curStream],
                heapDistancesView,
                heapIndicesView,
                k,
                metric,
                outDistanceView,
                outIndicesView,
                streams[curStream]);

        curStream = (curStream + 1) % 2;
    }

    streamWait({stream}, streams);
}

} // namespace gpu
} // namespace faiss
