Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
BinaryFlatIndex.cuh
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
9 #pragma once
10
11 #include "../utils/DeviceTensor.cuh"
12 #include "../utils/DeviceVector.cuh"
13 #include "../utils/MemorySpace.h"
14
15 namespace faiss { namespace gpu {
16
17 class GpuResources;
18
19 /// Holder of GPU resources for a particular flat index
20 class BinaryFlatIndex {
21 public:
22 BinaryFlatIndex(GpuResources* res,
23 int dim,
24 MemorySpace space);
25
26 /// Returns the number of vectors we contain
27 int getSize() const;
28
29 int getDim() const;
30
31 /// Reserve storage that can contain at least this many vectors
32 void reserve(size_t numVecs, cudaStream_t stream);
33
34 /// Returns a reference to our vectors currently in use
35 Tensor<unsigned char, 2, true>& getVectorsRef();
36
37 void query(Tensor<unsigned char, 2, true>& vecs,
38 int k,
39 Tensor<int, 2, true>& outDistances,
40 Tensor<int, 2, true>& outIndices);
41
42 /// Add vectors to ourselves; the pointer passed can be on the host
43 /// or the device
44 void add(const unsigned char* data, int numVecs, cudaStream_t stream);
45
46 /// Free all storage
47 void reset();
48
49 private:
50 /// Collection of GPU resources that we use
51 GpuResources* resources_;
52
53 /// Dimensionality of our vectors
54 const int dim_;
55
56 /// Memory space for our allocations
57 MemorySpace space_;
58
59 /// How many vectors we have
60 int num_;
61
62 /// The underlying expandable storage
63 DeviceVector<char> rawData_;
64
65 /// Vectors currently in rawData_
66 DeviceTensor<unsigned char, 2, true> vectors_;
67 };
68
69 } } // namespace
faiss::gpu::BinaryFlatIndex::reset
void reset()
Free all storage.
Definition: BinaryFlatIndex.cu:82
faiss::gpu::BinaryFlatIndex
Holder of GPU resources for a particular flat index.
Definition: BinaryFlatIndex.cuh:20
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::DeviceVector< char >
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::BinaryFlatIndex::getSize
int getSize() const
Returns the number of vectors we contain.
Definition: BinaryFlatIndex.cu:28
faiss::gpu::BinaryFlatIndex::reserve
void reserve(size_t numVecs, cudaStream_t stream)
Reserve storage that can contain at least this many vectors.
Definition: BinaryFlatIndex.cu:37
faiss::gpu::DeviceTensor< unsigned char, 2, true >
faiss::gpu::BinaryFlatIndex::getVectorsRef
Tensor< unsigned char, 2, true > & getVectorsRef()
Returns a reference to our vectors currently in use.
Definition: BinaryFlatIndex.cu:42
faiss::gpu::BinaryFlatIndex::add
void add(const unsigned char *data, int numVecs, cudaStream_t stream)
Definition: BinaryFlatIndex.cu:62