Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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