Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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