All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFPQ.cuh
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#pragma once
#include "IVFBase.cuh"
#include "../utils/Float16.cuh"
namespace faiss { namespace gpu {
/// Implementing class for IVFPQ on the GPU
class IVFPQ : public IVFBase {
public:
IVFPQ(GpuResources* resources,
/// We do not own this reference
FlatIndex* quantizer,
int numSubQuantizers,
int bitsPerSubQuantizer,
float* pqCentroidData,
IndicesOptions indicesOptions,
bool useFloat16LookupTables,
MemorySpace space);
/// Returns true if we support PQ in this size
static bool isSupportedPQCodeLength(int size);
/// For no precomputed codes, is this a supported sub-dimension
/// size?
/// FIXME: get MM implementation working again
static bool isSupportedNoPrecomputedSubDimSize(int dims);
~IVFPQ() override;
/// Enable or disable pre-computed codes
void setPrecomputedCodes(bool enable);
/// Adds a set of codes and indices to a list; the data can be
/// resident on either the host or the device
void addCodeVectorsFromCpu(int listId,
const void* codes,
const long* indices,
size_t numVecs);
/// Calcuates the residual and quantizes the vectors, adding them to
/// this index
/// The input data must be on our current device.
/// Returns the number of vectors successfully added. Vectors may
/// not be able to be added because they contain NaNs.
int classifyAndAddVectors(Tensor<float, 2, true>& vecs,
Tensor<long, 1, true>& indices);
/// Find the approximate k nearest neigbors for `queries` against
/// our database
void query(Tensor<float, 2, true>& queries,
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<long, 2, true>& outIndices);
/// Return the list codes of a particular list back to the CPU
std::vector<unsigned char> getListCodes(int listId) const;
/// Returns our set of sub-quantizers of the form
/// (sub q)(code id)(sub dim)
Tensor<float, 3, true> getPQCentroids();
private:
/// Sets the current product quantizer centroids; the data can be
/// resident on either the host or the device. It will be transposed
/// into our preferred data layout
/// Data must be a row-major, 3-d array of size
/// (numSubQuantizers, numSubQuantizerCodes, dim / numSubQuantizers)
void setPQCentroids_(float* data);
/// Calculate precomputed residual distance information
void precomputeCodes_();
/// Runs kernels for scanning inverted lists with precomputed codes
void runPQPrecomputedCodes_(Tensor<float, 2, true>& queries,
DeviceTensor<float, 2, true>& coarseDistances,
DeviceTensor<int, 2, true>& coarseIndices,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<long, 2, true>& outIndices);
/// Runs kernels for scanning inverted lists without precomputed codes
void runPQNoPrecomputedCodes_(Tensor<float, 2, true>& queries,
DeviceTensor<float, 2, true>& coarseDistances,
DeviceTensor<int, 2, true>& coarseIndices,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<long, 2, true>& outIndices);
private:
/// Number of sub-quantizers per vector
const int numSubQuantizers_;
/// Number of bits per sub-quantizer
const int bitsPerSubQuantizer_;
/// Number of per sub-quantizer codes (2^bits)
const int numSubQuantizerCodes_;
/// Number of dimensions per each sub-quantizer
const int dimPerSubQuantizer_;
/// Do we maintain precomputed terms and lookup tables in float16
/// form?
const bool useFloat16LookupTables_;
/// On the GPU, we prefer different PQ centroid data layouts for
/// different purposes.
///
/// (sub q)(sub dim)(code id)
DeviceTensor<float, 3, true> pqCentroidsInnermostCode_;
/// (sub q)(code id)(sub dim)
DeviceTensor<float, 3, true> pqCentroidsMiddleCode_;
/// Are precomputed codes enabled? (additional factoring and
/// precomputation of the residual distance, to reduce query-time work)
bool precomputedCodes_;
/// Precomputed term 2 in float form
/// (centroid id)(sub q)(code id)
DeviceTensor<float, 3, true> precomputedCode_;
#ifdef FAISS_USE_FLOAT16
/// Precomputed term 2 in half form
DeviceTensor<half, 3, true> precomputedCodeHalf_;
#endif
};
} } // namespace
faiss::gpu::IVFPQ::addCodeVectorsFromCpu
void addCodeVectorsFromCpu(int listId, const void *codes, const long *indices, size_t numVecs)
Definition: IVFPQ.cu:349
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:23
faiss::gpu::IVFBase
Base inverted list functionality for IVFFlat and IVFPQ.
Definition: IVFBase.cuh:27
faiss::gpu::GpuResources
Definition: GpuResources.h:23
faiss::gpu::IVFPQ::isSupportedPQCodeLength
static bool isSupportedPQCodeLength(int size)
Returns true if we support PQ in this size.
Definition: IVFPQ.cu:72
faiss::gpu::IVFPQ::classifyAndAddVectors
int classifyAndAddVectors(Tensor< float, 2, true > &vecs, Tensor< long, 1, true > &indices)
Definition: IVFPQ.cu:120
faiss::gpu::IVFPQ::query
void query(Tensor< float, 2, true > &queries, int nprobe, int k, Tensor< float, 2, true > &outDistances, Tensor< long, 2, true > &outIndices)
Definition: IVFPQ.cu:519
faiss::gpu::IVFPQ::getPQCentroids
Tensor< float, 3, true > getPQCentroids()
Definition: IVFPQ.cu:594
faiss::gpu::IVFPQ::setPrecomputedCodes
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
Definition: IVFPQ.cu:102
faiss::gpu::IVFPQ::getListCodes
std::vector< unsigned char > getListCodes(int listId) const
Return the list codes of a particular list back to the CPU.
Definition: IVFPQ.cu:586
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:30
faiss::gpu::DeviceTensor< float, 2, true >
faiss::gpu::IVFPQ::IVFPQ
IVFPQ(GpuResources *resources, FlatIndex *quantizer, int numSubQuantizers, int bitsPerSubQuantizer, float *pqCentroidData, IndicesOptions indicesOptions, bool useFloat16LookupTables, MemorySpace space)
Definition: IVFPQ.cu:35
faiss::gpu::IVFPQ
Implementing class for IVFPQ on the GPU.
Definition: IVFPQ.cuh:19
faiss::gpu::IVFPQ::isSupportedNoPrecomputedSubDimSize
static bool isSupportedNoPrecomputedSubDimSize(int dims)
Definition: IVFPQ.cu:97