1 /** 2 * Copyright (c) Facebook, Inc. and its affiliates. 3 * 4 * This source code is licensed under the MIT license found in the 5 * LICENSE file in the root directory of this source tree. 6 */ 7 8 #pragma once 9 10 #include <faiss/Index.h> 11 #include <faiss/gpu/GpuIndicesOptions.h> 12 #include <thrust/device_vector.h> 13 #include <faiss/gpu/utils/Tensor.cuh> 14 15 // A collection of utility functions for IVFPQ and IVFFlat, for 16 // post-processing and k-selecting the results 17 namespace faiss { 18 namespace gpu { 19 20 class GpuResources; 21 22 /// Function for multi-pass scanning that collects the length of 23 /// intermediate results for all (query, probe) pair 24 void runCalcListOffsets( 25 GpuResources* res, 26 Tensor<int, 2, true>& topQueryToCentroid, 27 thrust::device_vector<int>& listLengths, 28 Tensor<int, 2, true>& prefixSumOffsets, 29 Tensor<char, 1, true>& thrustMem, 30 cudaStream_t stream); 31 32 /// Performs a first pass of k-selection on the results 33 void runPass1SelectLists( 34 Tensor<int, 2, true>& prefixSumOffsets, 35 Tensor<float, 1, true>& distance, 36 int nprobe, 37 int k, 38 bool chooseLargest, 39 Tensor<float, 3, true>& heapDistances, 40 Tensor<int, 3, true>& heapIndices, 41 cudaStream_t stream); 42 43 /// Performs a final pass of k-selection on the results, producing the 44 /// final indices 45 void runPass2SelectLists( 46 Tensor<float, 2, true>& heapDistances, 47 Tensor<int, 2, true>& heapIndices, 48 thrust::device_vector<void*>& listIndices, 49 IndicesOptions indicesOptions, 50 Tensor<int, 2, true>& prefixSumOffsets, 51 Tensor<int, 2, true>& topQueryToCentroid, 52 int k, 53 bool chooseLargest, 54 Tensor<float, 2, true>& outDistances, 55 Tensor<Index::idx_t, 2, true>& outIndices, 56 cudaStream_t stream); 57 58 } // namespace gpu 59 } // namespace faiss 60