Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
BinaryFlatIndex.cu
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 #include "BinaryFlatIndex.cuh"
10 #include "BinaryDistance.cuh"
11 #include "../utils/DeviceUtils.h"
12 #include "../GpuResources.h"
13
14 namespace faiss { namespace gpu {
15
16 BinaryFlatIndex::BinaryFlatIndex(GpuResources* res,
17 int dim,
18 MemorySpace space) :
19 resources_(res),
20 dim_(dim),
21 space_(space),
22 num_(0),
23 rawData_(space) {
24 FAISS_ASSERT(dim % 8 == 0);
25 }
26
27 /// Returns the number of vectors we contain
28 int BinaryFlatIndex::getSize() const {
29 return vectors_.getSize(0);
30 }
31
32 int BinaryFlatIndex::getDim() const {
33 return vectors_.getSize(1) * 8;
34 }
35
36 void
37 BinaryFlatIndex::reserve(size_t numVecs, cudaStream_t stream) {
38 rawData_.reserve(numVecs * (dim_ / 8) * sizeof(unsigned int), stream);
39 }
40
41 Tensor<unsigned char, 2, true>&
42 BinaryFlatIndex::getVectorsRef() {
43 return vectors_;
44 }
45
46 void
47 BinaryFlatIndex::query(Tensor<unsigned char, 2, true>& input,
48 int k,
49 Tensor<int, 2, true>& outDistances,
50 Tensor<int, 2, true>& outIndices) {
51 auto stream = resources_->getDefaultStreamCurrentDevice();
52
53 runBinaryDistance(vectors_,
54 input,
55 outDistances,
56 outIndices,
57 k,
58 stream);
59 }
60
61 void
62 BinaryFlatIndex::add(const unsigned char* data,
63 int numVecs,
64 cudaStream_t stream) {
65 if (numVecs == 0) {
66 return;
67 }
68
69 rawData_.append((char*) data,
70 (size_t) (dim_ / 8) * numVecs * sizeof(unsigned char),
71 stream,
72 true /* reserve exactly */);
73
74 num_ += numVecs;
75
76 DeviceTensor<unsigned char, 2, true> vectors(
77 (unsigned char*) rawData_.data(), {(int) num_, (dim_ / 8)}, space_);
78 vectors_ = std::move(vectors);
79 }
80
81 void
82 BinaryFlatIndex::reset() {
83 rawData_.clear();
84 vectors_ = std::move(DeviceTensor<unsigned char, 2, true>());
85 num_ = 0;
86 }
87
88 } }
faiss::gpu::BinaryFlatIndex::reset
void reset()
Free all storage.
Definition: BinaryFlatIndex.cu:82
faiss::gpu::GpuResources::getDefaultStreamCurrentDevice
cudaStream_t getDefaultStreamCurrentDevice()
Calls getDefaultStream with the current device.
Definition: GpuResources.cpp:23
faiss::gpu::Tensor::getSize
__host__ __device__ IndexT getSize(int i) const
Definition: Tensor.cuh:222
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