All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFBase.cuh
/**
* 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.
#pragma once
#include "../GpuIndicesOptions.h"
#include "../utils/DeviceVector.cuh"
#include "../utils/DeviceTensor.cuh"
#include <memory>
#include <thrust/device_vector.h>
#include <vector>
namespace faiss { namespace gpu {
class GpuResources;
struct FlatIndex;
/// Base inverted list functionality for IVFFlat and IVFPQ
class IVFBase {
public:
IVFBase(GpuResources* resources,
/// We do not own this reference
FlatIndex* quantizer,
int bytesPerVector,
IndicesOptions indicesOptions);
virtual ~IVFBase();
/// Reserve GPU memory in our inverted lists for this number of vectors
void reserveMemory(size_t numVecs);
/// Clear out all inverted lists, but retain the coarse quantizer
/// and the product quantizer info
void reset();
/// Return the number of dimensions we are indexing
int getDim() const;
/// After adding vectors, one can call this to reclaim device memory
/// to exactly the amount needed. Returns space reclaimed in bytes
size_t reclaimMemory();
/// Returns the number of inverted lists
size_t getNumLists() const;
/// For debugging purposes, return the list length of a particular
/// list
int getListLength(int listId) const;
/// Return the list indices of a particular list back to the CPU
std::vector<long> getListIndices(int listId) const;
protected:
/// Reclaim memory consumed on the device for our inverted lists
/// `exact` means we trim exactly to the memory needed
size_t reclaimMemory_(bool exact);
/// Update all device-side list pointer and size information
void updateDeviceListInfo_(cudaStream_t stream);
/// For a set of list IDs, update device-side list pointer and size
/// information
void updateDeviceListInfo_(const std::vector<int>& listIds,
cudaStream_t stream);
/// Shared function to copy indices from CPU to GPU
void addIndicesFromCpu_(int listId,
const long* indices,
size_t numVecs);
protected:
/// Collection of GPU resources that we use
GpuResources* resources_;
/// Quantizer object
FlatIndex* quantizer_;
/// Expected dimensionality of the vectors
const int dim_;
/// Number of inverted lists we maintain
const int numLists_;
/// Number of bytes per vector in the list
const int bytesPerVector_;
/// How are user indices stored on the GPU?
const IndicesOptions indicesOptions_;
/// Device representation of all inverted list data
/// id -> data
thrust::device_vector<void*> deviceListDataPointers_;
/// Device representation of all inverted list index pointers
/// id -> data
thrust::device_vector<void*> deviceListIndexPointers_;
/// Device representation of all inverted list lengths
/// id -> length
thrust::device_vector<int> deviceListLengths_;
/// Maximum list length seen
int maxListLength_;
/// Device memory for each separate list, as managed by the host.
/// Device memory as stored in DeviceVector is stored as unique_ptr
/// since deviceListSummary_ pointers must remain valid despite
/// resizing of deviceLists_
std::vector<std::unique_ptr<DeviceVector<unsigned char>>> deviceListData_;
std::vector<std::unique_ptr<DeviceVector<unsigned char>>> deviceListIndices_;
/// If we are storing indices on the CPU (indicesOptions_ is
/// INDICES_CPU), then this maintains a CPU-side map of what
/// (inverted list id, offset) maps to which user index
std::vector<std::vector<long>> listOffsetToUserIndex_;
};
} } // namespace
faiss::gpu::IVFBase::numLists_
const int numLists_
Number of inverted lists we maintain.
Definition: IVFBase.cuh:90
faiss::gpu::IVFBase::maxListLength_
int maxListLength_
Maximum list length seen.
Definition: IVFBase.cuh:111
faiss::gpu::IVFBase::listOffsetToUserIndex_
std::vector< std::vector< long > > listOffsetToUserIndex_
Definition: IVFBase.cuh:123
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:23
faiss::gpu::IVFBase::getDim
int getDim() const
Return the number of dimensions we are indexing.
Definition: IVFBase.cu:99
faiss::gpu::IVFBase::getListLength
int getListLength(int listId) const
Definition: IVFBase.cu:199
faiss::gpu::IVFBase::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: IVFBase.cu:44
faiss::gpu::IVFBase::reclaimMemory_
size_t reclaimMemory_(bool exact)
Definition: IVFBase.cu:110
faiss::gpu::IVFBase
Base inverted list functionality for IVFFlat and IVFPQ.
Definition: IVFBase.cuh:27
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::IVFBase::deviceListLengths_
thrust::device_vector< int > deviceListLengths_
Definition: IVFBase.cuh:108
faiss::gpu::IVFBase::deviceListIndexPointers_
thrust::device_vector< void * > deviceListIndexPointers_
Definition: IVFBase.cuh:104
faiss::gpu::IVFBase::IVFBase
IVFBase(GpuResources *resources, FlatIndex *quantizer, int bytesPerVector, IndicesOptions indicesOptions)
Definition: IVFBase.cu:26
faiss::gpu::IVFBase::quantizer_
FlatIndex * quantizer_
Quantizer object.
Definition: IVFBase.cuh:84
faiss::gpu::IVFBase::deviceListDataPointers_
thrust::device_vector< void * > deviceListDataPointers_
Definition: IVFBase.cuh:100
faiss::gpu::IVFBase::reset
void reset()
Definition: IVFBase.cu:74
faiss::gpu::IVFBase::resources_
GpuResources * resources_
Collection of GPU resources that we use.
Definition: IVFBase.cuh:81
faiss::gpu::IVFBase::bytesPerVector_
const int bytesPerVector_
Number of bytes per vector in the list.
Definition: IVFBase.cuh:93
faiss::gpu::IVFBase::updateDeviceListInfo_
void updateDeviceListInfo_(cudaStream_t stream)
Update all device-side list pointer and size information.
Definition: IVFBase.cu:137
faiss::gpu::IVFBase::getListIndices
std::vector< long > getListIndices(int listId) const
Return the list indices of a particular list back to the CPU.
Definition: IVFBase.cu:206
faiss::gpu::IVFBase::indicesOptions_
const IndicesOptions indicesOptions_
How are user indices stored on the GPU?
Definition: IVFBase.cuh:96
faiss::gpu::IVFBase::deviceListData_
std::vector< std::unique_ptr< DeviceVector< unsigned char > > > deviceListData_
Definition: IVFBase.cuh:117
faiss::gpu::IVFBase::dim_
const int dim_
Expected dimensionality of the vectors.
Definition: IVFBase.cuh:87
faiss::gpu::IVFBase::addIndicesFromCpu_
void addIndicesFromCpu_(int listId, const long *indices, size_t numVecs)
Shared function to copy indices from CPU to GPU.
Definition: IVFBase.cu:244
faiss::gpu::IVFBase::reclaimMemory
size_t reclaimMemory()
Definition: IVFBase.cu:104
faiss::gpu::IVFBase::getNumLists
size_t getNumLists() const
Returns the number of inverted lists.
Definition: IVFBase.cu:194