All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFFlatScan.cu
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#include "IVFFlatScan.cuh"
#include "../GpuResources.h"
#include "IVFUtils.cuh"
#include "../utils/ConversionOperators.cuh"
#include "../utils/DeviceDefs.cuh"
#include "../utils/DeviceUtils.h"
#include "../utils/DeviceTensor.cuh"
#include "../utils/Float16.cuh"
#include "../utils/MathOperators.cuh"
#include "../utils/LoadStoreOperators.cuh"
#include "../utils/PtxUtils.cuh"
#include "../utils/Reductions.cuh"
#include "../utils/StaticUtils.h"
#include <thrust/host_vector.h>
namespace faiss { namespace gpu {
template <typename T>
inline __device__ typename Math<T>::ScalarType l2Distance(T a, T b) {
a = Math<T>::sub(a, b);
a = Math<T>::mul(a, a);
return Math<T>::reduceAdd(a);
}
template <typename T>
inline __device__ typename Math<T>::ScalarType ipDistance(T a, T b) {
return Math<T>::reduceAdd(Math<T>::mul(a, b));
}
// For list scanning, even if the input data is `half`, we perform all
// math in float32, because the code is memory b/w bound, and the
// added precision for accumulation is useful
/// The class that we use to provide scan specializations
template <int Dims, bool L2, typename T>
struct IVFFlatScan {
};
// Fallback implementation: works for any dimension size
template <bool L2, typename T>
struct IVFFlatScan<-1, L2, T> {
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
extern __shared__ float smem[];
T* vecs = (T*) vecData;
for (int vec = 0; vec < numVecs; ++vec) {
// Reduce in dist
float dist = 0.0f;
for (int d = threadIdx.x; d < dim; d += blockDim.x) {
float vecVal = ConvertTo<float>::to(vecs[vec * dim + d]);
float queryVal = query[d];
float curDist;
if (L2) {
curDist = l2Distance(queryVal, vecVal);
} else {
curDist = ipDistance(queryVal, vecVal);
}
dist += curDist;
}
// Reduce distance within block
dist = blockReduceAllSum<float, false, true>(dist, smem);
if (threadIdx.x == 0) {
distanceOut[vec] = dist;
}
}
}
};
// implementation: works for # dims == blockDim.x
template <bool L2, typename T>
struct IVFFlatScan<0, L2, T> {
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
extern __shared__ float smem[];
T* vecs = (T*) vecData;
float queryVal = query[threadIdx.x];
constexpr int kUnroll = 4;
int limit = utils::roundDown(numVecs, kUnroll);
for (int i = 0; i < limit; i += kUnroll) {
float vecVal[kUnroll];
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
vecVal[j] = ConvertTo<float>::to(vecs[(i + j) * dim + threadIdx.x]);
}
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
if (L2) {
vecVal[j] = l2Distance(queryVal, vecVal[j]);
} else {
vecVal[j] = ipDistance(queryVal, vecVal[j]);
}
}
blockReduceAllSum<kUnroll, float, false, true>(vecVal, smem);
if (threadIdx.x == 0) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
distanceOut[i + j] = vecVal[j];
}
}
}
// Handle remainder
for (int i = limit; i < numVecs; ++i) {
float vecVal = ConvertTo<float>::to(vecs[i * dim + threadIdx.x]);
if (L2) {
vecVal = l2Distance(queryVal, vecVal);
} else {
vecVal = ipDistance(queryVal, vecVal);
}
vecVal = blockReduceAllSum<float, false, true>(vecVal, smem);
if (threadIdx.x == 0) {
distanceOut[i] = vecVal;
}
}
}
};
// 64-d float32 implementation
template <bool L2>
struct IVFFlatScan<64, L2, float> {
static constexpr int kDims = 64;
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
// Each warp reduces a single 64-d vector; each lane loads a float2
float* vecs = (float*) vecData;
int laneId = getLaneId();
int warpId = threadIdx.x / kWarpSize;
int numWarps = blockDim.x / kWarpSize;
float2 queryVal = *(float2*) &query[laneId * 2];
constexpr int kUnroll = 4;
float2 vecVal[kUnroll];
int limit = utils::roundDown(numVecs, kUnroll * numWarps);
for (int i = warpId; i < limit; i += kUnroll * numWarps) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
// Vector we are loading from is i
// Dim we are loading from is laneId * 2
vecVal[j] = *(float2*) &vecs[(i + j * numWarps) * kDims + laneId * 2];
}
float dist[kUnroll];
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
if (L2) {
dist[j] = l2Distance(queryVal, vecVal[j]);
} else {
dist[j] = ipDistance(queryVal, vecVal[j]);
}
}
// Reduce within the warp
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
dist[j] = warpReduceAllSum(dist[j]);
}
if (laneId == 0) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
distanceOut[i + j * numWarps] = dist[j];
}
}
}
// Handle remainder
for (int i = limit + warpId; i < numVecs; i += numWarps) {
vecVal[0] = *(float2*) &vecs[i * kDims + laneId * 2];
float dist;
if (L2) {
dist = l2Distance(queryVal, vecVal[0]);
} else {
dist = ipDistance(queryVal, vecVal[0]);
}
dist = warpReduceAllSum(dist);
if (laneId == 0) {
distanceOut[i] = dist;
}
}
}
};
#ifdef FAISS_USE_FLOAT16
// float16 implementation
template <bool L2>
struct IVFFlatScan<64, L2, half> {
static constexpr int kDims = 64;
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
// Each warp reduces a single 64-d vector; each lane loads a half2
half* vecs = (half*) vecData;
int laneId = getLaneId();
int warpId = threadIdx.x / kWarpSize;
int numWarps = blockDim.x / kWarpSize;
float2 queryVal = *(float2*) &query[laneId * 2];
constexpr int kUnroll = 4;
half2 vecVal[kUnroll];
int limit = utils::roundDown(numVecs, kUnroll * numWarps);
for (int i = warpId; i < limit; i += kUnroll * numWarps) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
// Vector we are loading from is i
// Dim we are loading from is laneId * 2
vecVal[j] = *(half2*) &vecs[(i + j * numWarps) * kDims + laneId * 2];
}
float dist[kUnroll];
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
if (L2) {
dist[j] = l2Distance(queryVal, __half22float2(vecVal[j]));
} else {
dist[j] = ipDistance(queryVal, __half22float2(vecVal[j]));
}
}
// Reduce within the warp
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
dist[j] = warpReduceAllSum(dist[j]);
}
if (laneId == 0) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
distanceOut[i + j * numWarps] = dist[j];
}
}
}
// Handle remainder
for (int i = limit + warpId; i < numVecs; i += numWarps) {
vecVal[0] = *(half2*) &vecs[i * kDims + laneId * 2];
float dist;
if (L2) {
dist = l2Distance(queryVal, __half22float2(vecVal[0]));
} else {
dist = ipDistance(queryVal, __half22float2(vecVal[0]));
}
dist = warpReduceAllSum(dist);
if (laneId == 0) {
distanceOut[i] = dist;
}
}
}
};
#endif
// 128-d float32 implementation
template <bool L2>
struct IVFFlatScan<128, L2, float> {
static constexpr int kDims = 128;
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
// Each warp reduces a single 128-d vector; each lane loads a float4
float* vecs = (float*) vecData;
int laneId = getLaneId();
int warpId = threadIdx.x / kWarpSize;
int numWarps = blockDim.x / kWarpSize;
float4 queryVal = *(float4*) &query[laneId * 4];
constexpr int kUnroll = 4;
float4 vecVal[kUnroll];
int limit = utils::roundDown(numVecs, kUnroll * numWarps);
for (int i = warpId; i < limit; i += kUnroll * numWarps) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
// Vector we are loading from is i
// Dim we are loading from is laneId * 4
vecVal[j] = *(float4*) &vecs[(i + j * numWarps) * kDims + laneId * 4];
}
float dist[kUnroll];
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
if (L2) {
dist[j] = l2Distance(queryVal, vecVal[j]);
} else {
dist[j] = ipDistance(queryVal, vecVal[j]);
}
}
// Reduce within the warp
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
dist[j] = warpReduceAllSum(dist[j]);
}
if (laneId == 0) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
distanceOut[i + j * numWarps] = dist[j];
}
}
}
// Handle remainder
for (int i = limit + warpId; i < numVecs; i += numWarps) {
vecVal[0] = *(float4*) &vecs[i * kDims + laneId * 4];
float dist;
if (L2) {
dist = l2Distance(queryVal, vecVal[0]);
} else {
dist = ipDistance(queryVal, vecVal[0]);
}
dist = warpReduceAllSum(dist);
if (laneId == 0) {
distanceOut[i] = dist;
}
}
}
};
#ifdef FAISS_USE_FLOAT16
// float16 implementation
template <bool L2>
struct IVFFlatScan<128, L2, half> {
static constexpr int kDims = 128;
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
// Each warp reduces a single 128-d vector; each lane loads a Half4
half* vecs = (half*) vecData;
int laneId = getLaneId();
int warpId = threadIdx.x / kWarpSize;
int numWarps = blockDim.x / kWarpSize;
float4 queryVal = *(float4*) &query[laneId * 4];
constexpr int kUnroll = 4;
Half4 vecVal[kUnroll];
int limit = utils::roundDown(numVecs, kUnroll * numWarps);
for (int i = warpId; i < limit; i += kUnroll * numWarps) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
// Vector we are loading from is i
// Dim we are loading from is laneId * 4
vecVal[j] =
LoadStore<Half4>::load(
&vecs[(i + j * numWarps) * kDims + laneId * 4]);
}
float dist[kUnroll];
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
if (L2) {
dist[j] = l2Distance(queryVal, half4ToFloat4(vecVal[j]));
} else {
dist[j] = ipDistance(queryVal, half4ToFloat4(vecVal[j]));
}
}
// Reduce within the warp
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
dist[j] = warpReduceAllSum(dist[j]);
}
if (laneId == 0) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
distanceOut[i + j * numWarps] = dist[j];
}
}
}
// Handle remainder
for (int i = limit + warpId; i < numVecs; i += numWarps) {
vecVal[0] = LoadStore<Half4>::load(&vecs[i * kDims + laneId * 4]);
float dist;
if (L2) {
dist = l2Distance(queryVal, half4ToFloat4(vecVal[0]));
} else {
dist = ipDistance(queryVal, half4ToFloat4(vecVal[0]));
}
dist = warpReduceAllSum(dist);
if (laneId == 0) {
distanceOut[i] = dist;
}
}
}
};
#endif
// 256-d float32 implementation
template <bool L2>
struct IVFFlatScan<256, L2, float> {
static constexpr int kDims = 256;
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
// A specialization here to load per-warp seems to be worse, since
// we're already running at near memory b/w peak
IVFFlatScan<0, L2, float>::scan(query,
vecData,
numVecs,
dim,
distanceOut);
}
};
#ifdef FAISS_USE_FLOAT16
// float16 implementation
template <bool L2>
struct IVFFlatScan<256, L2, half> {
static constexpr int kDims = 256;
static __device__ void scan(float* query,
void* vecData,
int numVecs,
int dim,
float* distanceOut) {
// Each warp reduces a single 256-d vector; each lane loads a Half8
half* vecs = (half*) vecData;
int laneId = getLaneId();
int warpId = threadIdx.x / kWarpSize;
int numWarps = blockDim.x / kWarpSize;
// This is not a contiguous load, but we only have to load these two
// values, so that we can load by Half8 below
float4 queryValA = *(float4*) &query[laneId * 8];
float4 queryValB = *(float4*) &query[laneId * 8 + 4];
constexpr int kUnroll = 4;
Half8 vecVal[kUnroll];
int limit = utils::roundDown(numVecs, kUnroll * numWarps);
for (int i = warpId; i < limit; i += kUnroll * numWarps) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
// Vector we are loading from is i
// Dim we are loading from is laneId * 8
vecVal[j] =
LoadStore<Half8>::load(
&vecs[(i + j * numWarps) * kDims + laneId * 8]);
}
float dist[kUnroll];
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
if (L2) {
dist[j] = l2Distance(queryValA, half4ToFloat4(vecVal[j].a));
dist[j] += l2Distance(queryValB, half4ToFloat4(vecVal[j].b));
} else {
dist[j] = ipDistance(queryValA, half4ToFloat4(vecVal[j].a));
dist[j] += ipDistance(queryValB, half4ToFloat4(vecVal[j].b));
}
}
// Reduce within the warp
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
dist[j] = warpReduceAllSum(dist[j]);
}
if (laneId == 0) {
#pragma unroll
for (int j = 0; j < kUnroll; ++j) {
distanceOut[i + j * numWarps] = dist[j];
}
}
}
// Handle remainder
for (int i = limit + warpId; i < numVecs; i += numWarps) {
vecVal[0] = LoadStore<Half8>::load(&vecs[i * kDims + laneId * 8]);
float dist;
if (L2) {
dist = l2Distance(queryValA, half4ToFloat4(vecVal[0].a));
dist += l2Distance(queryValB, half4ToFloat4(vecVal[0].b));
} else {
dist = ipDistance(queryValA, half4ToFloat4(vecVal[0].a));
dist += ipDistance(queryValB, half4ToFloat4(vecVal[0].b));
}
dist = warpReduceAllSum(dist);
if (laneId == 0) {
distanceOut[i] = dist;
}
}
}
};
#endif
template <int Dims, bool L2, typename T>
__global__ void
ivfFlatScan(Tensor<float, 2, true> queries,
Tensor<int, 2, true> listIds,
void** allListData,
int* listLengths,
Tensor<int, 2, true> prefixSumOffsets,
Tensor<float, 1, true> distance) {
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);
auto listId = listIds[queryId][probeId];
// Safety guard in case NaNs in input cause no list ID to be generated
if (listId == -1) {
return;
}
auto query = queries[queryId].data();
auto vecs = allListData[listId];
auto numVecs = listLengths[listId];
auto dim = queries.getSize(1);
auto distanceOut = distance[outBase].data();
IVFFlatScan<Dims, L2, T>::scan(query, vecs, numVecs, dim, distanceOut);
}
void
runIVFFlatScanTile(Tensor<float, 2, true>& queries,
Tensor<int, 2, true>& listIds,
thrust::device_vector<void*>& listData,
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,
bool l2Distance,
bool useFloat16,
Tensor<float, 2, true>& outDistances,
Tensor<long, 2, true>& outIndices,
cudaStream_t stream) {
// Calculate offset lengths, so we know where to write out
// intermediate results
runCalcListOffsets(listIds, listLengths, prefixSumOffsets, thrustMem, stream);
// Calculate distances for vectors within our chunk of lists
constexpr int kMaxThreadsIVF = 512;
// FIXME: if `half` and # dims is multiple of 2, halve the
// threadblock size
int dim = queries.getSize(1);
int numThreads = std::min(dim, kMaxThreadsIVF);
auto grid = dim3(listIds.getSize(1),
listIds.getSize(0));
auto block = dim3(numThreads);
// All exact dim kernels are unrolled by 4, hence the `4`
auto smem = sizeof(float) * utils::divUp(numThreads, kWarpSize) * 4;
#define RUN_IVF_FLAT(DIMS, L2, T) \
do { \
ivfFlatScan<DIMS, L2, T> \
<<<grid, block, smem, stream>>>( \
queries, \
listIds, \
listData.data().get(), \
listLengths.data().get(), \
prefixSumOffsets, \
allDistances); \
} while (0)
#ifdef FAISS_USE_FLOAT16
#define HANDLE_DIM_CASE(DIMS) \
do { \
if (l2Distance) { \
if (useFloat16) { \
RUN_IVF_FLAT(DIMS, true, half); \
} else { \
RUN_IVF_FLAT(DIMS, true, float); \
} \
} else { \
if (useFloat16) { \
RUN_IVF_FLAT(DIMS, false, half); \
} else { \
RUN_IVF_FLAT(DIMS, false, float); \
} \
} \
} while (0)
#else
#define HANDLE_DIM_CASE(DIMS) \
do { \
if (l2Distance) { \
if (useFloat16) { \
FAISS_ASSERT(false); \
} else { \
RUN_IVF_FLAT(DIMS, true, float); \
} \
} else { \
if (useFloat16) { \
FAISS_ASSERT(false); \
} else { \
RUN_IVF_FLAT(DIMS, false, float); \
} \
} \
} while (0)
#endif // FAISS_USE_FLOAT16
if (dim == 64) {
HANDLE_DIM_CASE(64);
} else if (dim == 128) {
HANDLE_DIM_CASE(128);
} else if (dim == 256) {
HANDLE_DIM_CASE(256);
} else if (dim <= kMaxThreadsIVF) {
HANDLE_DIM_CASE(0);
} else {
HANDLE_DIM_CASE(-1);
}
#undef HANDLE_DIM_CASE
#undef RUN_IVF_FLAT
// k-select the output in chunks, to increase parallelism
runPass1SelectLists(prefixSumOffsets,
allDistances,
listIds.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,
listIds,
k,
!l2Distance, // L2 distance chooses smallest
outDistances,
outIndices,
stream);
CUDA_VERIFY(cudaGetLastError());
}
void
runIVFFlatScan(Tensor<float, 2, true>& queries,
Tensor<int, 2, true>& listIds,
thrust::device_vector<void*>& listData,
thrust::device_vector<void*>& listIndices,
IndicesOptions indicesOptions,
thrust::device_vector<int>& listLengths,
int maxListLength,
int k,
bool l2Distance,
bool useFloat16,
// output
Tensor<float, 2, true>& outDistances,
// output
Tensor<long, 2, true>& outIndices,
GpuResources* res) {
constexpr int kMinQueryTileSize = 8;
constexpr int kMaxQueryTileSize = 128;
constexpr int kThrustMemSize = 16384;
int nprobe = listIds.getSize(1);
auto& mem = res->getMemoryManagerCurrentDevice();
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(
mem, {kThrustMemSize}, stream);
DeviceTensor<char, 1, true> thrustMem2(
mem, {kThrustMemSize}, stream);
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 = mem.getSizeAvailable();
// 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
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(
mem, {queryTileSize * nprobe + 1}, stream);
DeviceTensor<int, 1, true> prefixSumOffsetSpace2(
mem, {queryTileSize * nprobe + 1}, stream);
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));
DeviceTensor<float, 1, true> allDistances1(
mem, {queryTileSize * nprobe * maxListLength}, stream);
DeviceTensor<float, 1, true> allDistances2(
mem, {queryTileSize * nprobe * maxListLength}, stream);
DeviceTensor<float, 1, true>* allDistances[2] =
{&allDistances1, &allDistances2};
DeviceTensor<float, 3, true> heapDistances1(
mem, {queryTileSize, pass2Chunks, k}, stream);
DeviceTensor<float, 3, true> heapDistances2(
mem, {queryTileSize, pass2Chunks, k}, stream);
DeviceTensor<float, 3, true>* heapDistances[2] =
{&heapDistances1, &heapDistances2};
DeviceTensor<int, 3, true> heapIndices1(
mem, {queryTileSize, pass2Chunks, k}, stream);
DeviceTensor<int, 3, true> heapIndices2(
mem, {queryTileSize, pass2Chunks, k}, stream);
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 listIdsView =
listIds.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);
runIVFFlatScanTile(queryView,
listIdsView,
listData,
listIndices,
indicesOptions,
listLengths,
*thrustMem[curStream],
prefixSumOffsetsView,
*allDistances[curStream],
heapDistancesView,
heapIndicesView,
k,
l2Distance,
useFloat16,
outDistanceView,
outIndicesView,
streams[curStream]);
curStream = (curStream + 1) % 2;
}
streamWait({stream}, streams);
}
} } // namespace
faiss::gpu::LoadStore
Definition: LoadStoreOperators.cuh:25
faiss::gpu::Math::reduceAdd
static __device__ T reduceAdd(T v)
For a vector type, this is a horizontal add, returning sum(v_i)
Definition: MathOperators.cuh:45
faiss::gpu::ConvertTo
Definition: ConversionOperators.cuh:24
faiss::gpu::IVFFlatScan
The class that we use to provide scan specializations.
Definition: IVFFlatScan.cu:47