Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFPQ.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 "IVFPQ.cuh"
10 #include "../GpuResources.h"
11 #include "BroadcastSum.cuh"
12 #include "Distance.cuh"
13 #include "FlatIndex.cuh"
14 #include "InvertedListAppend.cuh"
15 #include "L2Norm.cuh"
16 #include "PQCodeDistances.cuh"
17 #include "PQScanMultiPassNoPrecomputed.cuh"
18 #include "PQScanMultiPassPrecomputed.cuh"
19 #include "RemapIndices.h"
20 #include "VectorResidual.cuh"
21 #include "../utils/DeviceDefs.cuh"
22 #include "../utils/DeviceUtils.h"
23 #include "../utils/HostTensor.cuh"
24 #include "../utils/MatrixMult.cuh"
25 #include "../utils/NoTypeTensor.cuh"
26 #include "../utils/Transpose.cuh"
27 #include <limits>
28 #include <thrust/host_vector.h>
29 #include <unordered_map>
30
31 namespace faiss { namespace gpu {
32
33 IVFPQ::IVFPQ(GpuResources* resources,
34 FlatIndex* quantizer,
35 int numSubQuantizers,
36 int bitsPerSubQuantizer,
37 float* pqCentroidData,
38 IndicesOptions indicesOptions,
39 bool useFloat16LookupTables,
40 MemorySpace space) :
41 IVFBase(resources,
42 quantizer,
43 numSubQuantizers,
44 indicesOptions,
45 space),
46 numSubQuantizers_(numSubQuantizers),
47 bitsPerSubQuantizer_(bitsPerSubQuantizer),
48 numSubQuantizerCodes_(utils::pow2(bitsPerSubQuantizer_)),
49 dimPerSubQuantizer_(dim_ / numSubQuantizers),
50 precomputedCodes_(false),
51 useFloat16LookupTables_(useFloat16LookupTables) {
52 FAISS_ASSERT(pqCentroidData);
53
54 FAISS_ASSERT(bitsPerSubQuantizer_ <= 8);
55 FAISS_ASSERT(dim_ % numSubQuantizers_ == 0);
56 FAISS_ASSERT(isSupportedPQCodeLength(bytesPerVector_));
57
58 #ifndef FAISS_USE_FLOAT16
59 FAISS_ASSERT(!useFloat16LookupTables_);
60 #endif
61
62 setPQCentroids_(pqCentroidData);
63 }
64
65 IVFPQ::~IVFPQ() {
66 }
67
68
69 bool
70 IVFPQ::isSupportedPQCodeLength(int size) {
71 switch (size) {
72 case 1:
73 case 2:
74 case 3:
75 case 4:
76 case 8:
77 case 12:
78 case 16:
79 case 20:
80 case 24:
81 case 28:
82 case 32:
83 case 40:
84 case 48:
85 case 56: // only supported with float16
86 case 64: // only supported with float16
87 case 96: // only supported with float16
88 return true;
89 default:
90 return false;
91 }
92 }
93
94 bool
95 IVFPQ::isSupportedNoPrecomputedSubDimSize(int dims) {
96 return faiss::gpu::isSupportedNoPrecomputedSubDimSize(dims);
97 }
98
99 void
100 IVFPQ::setPrecomputedCodes(bool enable) {
101 if (precomputedCodes_ != enable) {
102 precomputedCodes_ = enable;
103
104 if (precomputedCodes_) {
105 precomputeCodes_();
106 } else {
107 // Clear out old precomputed code data
108 precomputedCode_ = std::move(DeviceTensor<float, 3, true>());
109
110 #ifdef FAISS_USE_FLOAT16
111 precomputedCodeHalf_ = std::move(DeviceTensor<half, 3, true>());
112 #endif
113 }
114 }
115 }
116
117 int
118 IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
119 Tensor<long, 1, true>& indices) {
120 FAISS_ASSERT(vecs.getSize(0) == indices.getSize(0));
121 FAISS_ASSERT(vecs.getSize(1) == dim_);
122
123 FAISS_ASSERT(!quantizer_->getUseFloat16());
124 auto& coarseCentroids = quantizer_->getVectorsFloat32Ref();
125 auto& mem = resources_->getMemoryManagerCurrentDevice();
126 auto stream = resources_->getDefaultStreamCurrentDevice();
127
128 // Number of valid vectors that we actually add; we return this
129 int numAdded = 0;
130
131 // We don't actually need this
132 DeviceTensor<float, 2, true> listDistance(mem, {vecs.getSize(0), 1}, stream);
133 // We use this
134 DeviceTensor<int, 2, true> listIds2d(mem, {vecs.getSize(0), 1}, stream);
135 auto listIds = listIds2d.view<1>({vecs.getSize(0)});
136
137 quantizer_->query(vecs, 1, listDistance, listIds2d, false);
138
139 // Copy the lists that we wish to append to back to the CPU
140 // FIXME: really this can be into pinned memory and a true async
141 // copy on a different stream; we can start the copy early, but it's
142 // tiny
143 HostTensor<int, 1, true> listIdsHost(listIds, stream);
144
145 // Calculate the residual for each closest centroid
146 DeviceTensor<float, 2, true> residuals(
147 mem, {vecs.getSize(0), vecs.getSize(1)}, stream);
148
149 runCalcResidual(vecs, coarseCentroids, listIds, residuals, stream);
150
151 // Residuals are in the form
152 // (vec x numSubQuantizer x dimPerSubQuantizer)
153 // transpose to
154 // (numSubQuantizer x vec x dimPerSubQuantizer)
155 auto residualsView = residuals.view<3>(
156 {residuals.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
157
158 DeviceTensor<float, 3, true> residualsTranspose(
159 mem,
160 {numSubQuantizers_, residuals.getSize(0), dimPerSubQuantizer_},
161 stream);
162
163 runTransposeAny(residualsView, 0, 1, residualsTranspose, stream);
164
165 // Get the product quantizer centroids in the form
166 // (numSubQuantizer x numSubQuantizerCodes x dimPerSubQuantizer)
167 // which is pqCentroidsMiddleCode_
168
169 // We now have a batch operation to find the top-1 distances:
170 // batch size: numSubQuantizer
171 // centroids: (numSubQuantizerCodes x dimPerSubQuantizer)
172 // residuals: (vec x dimPerSubQuantizer)
173 // => (numSubQuantizer x vec x 1)
174
175 DeviceTensor<float, 3, true> closestSubQDistance(
176 mem, {numSubQuantizers_, residuals.getSize(0), 1}, stream);
177 DeviceTensor<int, 3, true> closestSubQIndex(
178 mem, {numSubQuantizers_, residuals.getSize(0), 1}, stream);
179
180 for (int subQ = 0; subQ < numSubQuantizers_; ++subQ) {
181 auto closestSubQDistanceView = closestSubQDistance[subQ].view();
182 auto closestSubQIndexView = closestSubQIndex[subQ].view();
183
184 auto pqCentroidsMiddleCodeView = pqCentroidsMiddleCode_[subQ].view();
185 auto residualsTransposeView = residualsTranspose[subQ].view();
186
187 runL2Distance(resources_,
188 pqCentroidsMiddleCodeView,
189 true, // pqCentroidsMiddleCodeView is row major
190 nullptr, // no precomputed norms
191 residualsTransposeView,
192 true, // residualsTransposeView is row major
193 1,
194 closestSubQDistanceView,
195 closestSubQIndexView,
196 // We don't care about distances
197 true);
198 }
199
200 // Now, we have the nearest sub-q centroid for each slice of the
201 // residual vector.
202 auto closestSubQIndexView = closestSubQIndex.view<2>(
203 {numSubQuantizers_, residuals.getSize(0)});
204
205 // Transpose this for easy use
206 DeviceTensor<int, 2, true> encodings(
207 mem, {residuals.getSize(0), numSubQuantizers_}, stream);
208
209 runTransposeAny(closestSubQIndexView, 0, 1, encodings, stream);
210
211 // Now we add the encoded vectors to the individual lists
212 // First, make sure that there is space available for adding the new
213 // encoded vectors and indices
214
215 // list id -> # being added
216 std::unordered_map<int, int> assignCounts;
217
218 // vector id -> offset in list
219 // (we already have vector id -> list id in listIds)
220 HostTensor<int, 1, true> listOffsetHost({listIdsHost.getSize(0)});
221
222 for (int i = 0; i < listIdsHost.getSize(0); ++i) {
223 int listId = listIdsHost[i];
224
225 // Add vector could be invalid (contains NaNs etc)
226 if (listId < 0) {
227 listOffsetHost[i] = -1;
228 continue;
229 }
230
231 FAISS_ASSERT(listId < numLists_);
232 ++numAdded;
233
234 int offset = deviceListData_[listId]->size() / bytesPerVector_;
235
236 auto it = assignCounts.find(listId);
237 if (it != assignCounts.end()) {
238 offset += it->second;
239 it->second++;
240 } else {
241 assignCounts[listId] = 1;
242 }
243
244 listOffsetHost[i] = offset;
245 }
246
247 // If we didn't add anything (all invalid vectors), no need to
248 // continue
249 if (numAdded == 0) {
250 return 0;
251 }
252
253 // We need to resize the data structures for the inverted lists on
254 // the GPUs, which means that they might need reallocation, which
255 // means that their base address may change. Figure out the new base
256 // addresses, and update those in a batch on the device
257 {
258 // Resize all of the lists that we are appending to
259 for (auto& counts : assignCounts) {
260 auto& codes = deviceListData_[counts.first];
261 codes->resize(codes->size() + counts.second * bytesPerVector_,
262 stream);
263 int newNumVecs = (int) (codes->size() / bytesPerVector_);
264
265 auto& indices = deviceListIndices_[counts.first];
266 if ((indicesOptions_ == INDICES_32_BIT) ||
267 (indicesOptions_ == INDICES_64_BIT)) {
268 size_t indexSize =
269 (indicesOptions_ == INDICES_32_BIT) ? sizeof(int) : sizeof(long);
270
271 indices->resize(indices->size() + counts.second * indexSize, stream);
272 } else if (indicesOptions_ == INDICES_CPU) {
273 // indices are stored on the CPU side
274 FAISS_ASSERT(counts.first < listOffsetToUserIndex_.size());
275
276 auto& userIndices = listOffsetToUserIndex_[counts.first];
277 userIndices.resize(newNumVecs);
278 } else {
279 // indices are not stored on the GPU or CPU side
280 FAISS_ASSERT(indicesOptions_ == INDICES_IVF);
281 }
282
283 // This is used by the multi-pass query to decide how much scratch
284 // space to allocate for intermediate results
285 maxListLength_ = std::max(maxListLength_, newNumVecs);
286 }
287
288 // Update all pointers and sizes on the device for lists that we
289 // appended to
290 {
291 std::vector<int> listIds(assignCounts.size());
292 int i = 0;
293 for (auto& counts : assignCounts) {
294 listIds[i++] = counts.first;
295 }
296
297 updateDeviceListInfo_(listIds, stream);
298 }
299 }
300
301 // If we're maintaining the indices on the CPU side, update our
302 // map. We already resized our map above.
303 if (indicesOptions_ == INDICES_CPU) {
304 // We need to maintain the indices on the CPU side
305 HostTensor<long, 1, true> hostIndices(indices, stream);
306
307 for (int i = 0; i < hostIndices.getSize(0); ++i) {
308 int listId = listIdsHost[i];
309
310 // Add vector could be invalid (contains NaNs etc)
311 if (listId < 0) {
312 continue;
313 }
314
315 int offset = listOffsetHost[i];
316
317 FAISS_ASSERT(listId < listOffsetToUserIndex_.size());
318 auto& userIndices = listOffsetToUserIndex_[listId];
319
320 FAISS_ASSERT(offset < userIndices.size());
321 userIndices[offset] = hostIndices[i];
322 }
323 }
324
325 // We similarly need to actually append the new encoded vectors
326 {
327 DeviceTensor<int, 1, true> listOffset(mem, listOffsetHost, stream);
328
329 // This kernel will handle appending each encoded vector + index to
330 // the appropriate list
331 runIVFPQInvertedListAppend(listIds,
332 listOffset,
333 encodings,
334 indices,
335 deviceListDataPointers_,
336 deviceListIndexPointers_,
337 indicesOptions_,
338 stream);
339 }
340
341 return numAdded;
342 }
343
344 void
345 IVFPQ::addCodeVectorsFromCpu(int listId,
346 const void* codes,
347 const long* indices,
348 size_t numVecs) {
349 // This list must already exist
350 FAISS_ASSERT(listId < deviceListData_.size());
351 auto stream = resources_->getDefaultStreamCurrentDevice();
352
353 // If there's nothing to add, then there's nothing we have to do
354 if (numVecs == 0) {
355 return;
356 }
357
358 size_t lengthInBytes = numVecs * bytesPerVector_;
359
360 auto& listCodes = deviceListData_[listId];
361 auto prevCodeData = listCodes->data();
362
363 // We only have int32 length representations on the GPU per each
364 // list; the length is in sizeof(char)
365 FAISS_ASSERT(listCodes->size() % bytesPerVector_ == 0);
366 FAISS_ASSERT(listCodes->size() + lengthInBytes <=
367 (size_t) std::numeric_limits<int>::max());
368
369 listCodes->append((unsigned char*) codes,
370 lengthInBytes,
371 stream,
372 true /* exact reserved size */);
373
374 // Handle the indices as well
375 addIndicesFromCpu_(listId, indices, numVecs);
376
377 // This list address may have changed due to vector resizing, but
378 // only bother updating it on the device if it has changed
379 if (prevCodeData != listCodes->data()) {
380 deviceListDataPointers_[listId] = listCodes->data();
381 }
382
383 // And our size has changed too
384 int listLength = listCodes->size() / bytesPerVector_;
385 deviceListLengths_[listId] = listLength;
386
387 // We update this as well, since the multi-pass algorithm uses it
388 maxListLength_ = std::max(maxListLength_, listLength);
389
390 // device_vector add is potentially happening on a different stream
391 // than our default stream
392 if (resources_->getDefaultStreamCurrentDevice() != 0) {
393 streamWait({stream}, {0});
394 }
395 }
396
397 void
398 IVFPQ::setPQCentroids_(float* data) {
399 size_t pqSize =
400 numSubQuantizers_ * numSubQuantizerCodes_ * dimPerSubQuantizer_;
401
402 // Make sure the data is on the host
403 // FIXME: why are we doing this?
404 thrust::host_vector<float> hostMemory;
405 hostMemory.insert(hostMemory.end(), data, data + pqSize);
406
407 HostTensor<float, 3, true> pqHost(
408 hostMemory.data(),
409 {numSubQuantizers_, numSubQuantizerCodes_, dimPerSubQuantizer_});
410 DeviceTensor<float, 3, true> pqDevice(
411 pqHost,
412 resources_->getDefaultStreamCurrentDevice());
413
414 DeviceTensor<float, 3, true> pqDeviceTranspose(
415 {numSubQuantizers_, dimPerSubQuantizer_, numSubQuantizerCodes_});
416 runTransposeAny(pqDevice, 1, 2, pqDeviceTranspose,
417 resources_->getDefaultStreamCurrentDevice());
418
419 pqCentroidsInnermostCode_ = std::move(pqDeviceTranspose);
420
421 // Also maintain the PQ centroids in the form
422 // (sub q)(code id)(sub dim)
423 DeviceTensor<float, 3, true> pqCentroidsMiddleCode(
424 {numSubQuantizers_, numSubQuantizerCodes_, dimPerSubQuantizer_});
425 runTransposeAny(pqCentroidsInnermostCode_, 1, 2, pqCentroidsMiddleCode,
426 resources_->getDefaultStreamCurrentDevice());
427
428 pqCentroidsMiddleCode_ = std::move(pqCentroidsMiddleCode);
429 }
430
431 void
432 IVFPQ::precomputeCodes_() {
433 //
434 // d = || x - y_C ||^2 + || y_R ||^2 + 2 * (y_C|y_R) - 2 * (x|y_R)
435 // --------------- --------------------------- -------
436 // term 1 term 2 term 3
437 //
438
439 // Terms 1 and 3 are available only at query time. We compute term 2
440 // here.
441 FAISS_ASSERT(!quantizer_->getUseFloat16());
442 auto& coarseCentroids = quantizer_->getVectorsFloat32Ref();
443
444 // Compute ||y_R||^2 by treating
445 // (sub q)(code id)(sub dim) as (sub q * code id)(sub dim)
446 auto pqCentroidsMiddleCodeView =
447 pqCentroidsMiddleCode_.view<2>(
448 {numSubQuantizers_ * numSubQuantizerCodes_, dimPerSubQuantizer_});
449 DeviceTensor<float, 1, true> subQuantizerNorms(
450 {numSubQuantizers_ * numSubQuantizerCodes_});
451
452 runL2Norm(pqCentroidsMiddleCodeView, true,
453 subQuantizerNorms, true,
454 resources_->getDefaultStreamCurrentDevice());
455
456 // Compute 2 * (y_C|y_R) via batch matrix multiplication
457 // batch size (sub q) x {(centroid id)(sub dim) x (code id)(sub dim)'}
458 // => (sub q) x {(centroid id)(code id)}
459 // => (sub q)(centroid id)(code id)
460
461 // View (centroid id)(dim) as
462 // (centroid id)(sub q)(dim)
463 // Transpose (centroid id)(sub q)(sub dim) to
464 // (sub q)(centroid id)(sub dim)
465 auto centroidView = coarseCentroids.view<3>(
466 {coarseCentroids.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
467 DeviceTensor<float, 3, true> centroidsTransposed(
468 {numSubQuantizers_, coarseCentroids.getSize(0), dimPerSubQuantizer_});
469
470 runTransposeAny(centroidView, 0, 1, centroidsTransposed,
471 resources_->getDefaultStreamCurrentDevice());
472
473 DeviceTensor<float, 3, true> coarsePQProduct(
474 {numSubQuantizers_, coarseCentroids.getSize(0), numSubQuantizerCodes_});
475
476 runIteratedMatrixMult(coarsePQProduct, false,
477 centroidsTransposed, false,
478 pqCentroidsMiddleCode_, true,
479 2.0f, 0.0f,
480 resources_->getBlasHandleCurrentDevice(),
481 resources_->getDefaultStreamCurrentDevice());
482
483 // Transpose (sub q)(centroid id)(code id) to
484 // (centroid id)(sub q)(code id)
485 DeviceTensor<float, 3, true> coarsePQProductTransposed(
486 {coarseCentroids.getSize(0), numSubQuantizers_, numSubQuantizerCodes_});
487 runTransposeAny(coarsePQProduct, 0, 1, coarsePQProductTransposed,
488 resources_->getDefaultStreamCurrentDevice());
489
490 // View (centroid id)(sub q)(code id) as
491 // (centroid id)(sub q * code id)
492 auto coarsePQProductTransposedView = coarsePQProductTransposed.view<2>(
493 {coarseCentroids.getSize(0), numSubQuantizers_ * numSubQuantizerCodes_});
494
495 // Sum || y_R ||^2 + 2 * (y_C|y_R)
496 // i.e., add norms (sub q * code id)
497 // along columns of inner product (centroid id)(sub q * code id)
498 runSumAlongColumns(subQuantizerNorms, coarsePQProductTransposedView,
499 resources_->getDefaultStreamCurrentDevice());
500
501 #ifdef FAISS_USE_FLOAT16
502 if (useFloat16LookupTables_) {
503 precomputedCodeHalf_ = toHalf(resources_,
504 resources_->getDefaultStreamCurrentDevice(),
505 coarsePQProductTransposed);
506 return;
507 }
508 #endif
509
510 // We added into the view, so `coarsePQProductTransposed` is now our
511 // precomputed term 2.
512 precomputedCode_ = std::move(coarsePQProductTransposed);
513 }
514
515 void
516 IVFPQ::query(Tensor<float, 2, true>& queries,
517 int nprobe,
518 int k,
519 Tensor<float, 2, true>& outDistances,
520 Tensor<long, 2, true>& outIndices) {
521 // These are caught at a higher level
522 FAISS_ASSERT(nprobe <= GPU_MAX_SELECTION_K);
523 FAISS_ASSERT(k <= GPU_MAX_SELECTION_K);
524
525 auto& mem = resources_->getMemoryManagerCurrentDevice();
526 auto stream = resources_->getDefaultStreamCurrentDevice();
527 nprobe = std::min(nprobe, quantizer_->getSize());
528
529 FAISS_ASSERT(queries.getSize(1) == dim_);
530 FAISS_ASSERT(outDistances.getSize(0) == queries.getSize(0));
531 FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
532
533 // Reserve space for the closest coarse centroids
534 DeviceTensor<float, 2, true>
535 coarseDistances(mem, {queries.getSize(0), nprobe}, stream);
536 DeviceTensor<int, 2, true>
537 coarseIndices(mem, {queries.getSize(0), nprobe}, stream);
538
539 // Find the `nprobe` closest coarse centroids; we can use int
540 // indices both internally and externally
541 quantizer_->query(queries,
542 nprobe,
543 coarseDistances,
544 coarseIndices,
545 true);
546
547 if (precomputedCodes_) {
548 runPQPrecomputedCodes_(queries,
549 coarseDistances,
550 coarseIndices,
551 k,
552 outDistances,
553 outIndices);
554 } else {
555 runPQNoPrecomputedCodes_(queries,
556 coarseDistances,
557 coarseIndices,
558 k,
559 outDistances,
560 outIndices);
561 }
562
563 // If the GPU isn't storing indices (they are on the CPU side), we
564 // need to perform the re-mapping here
565 // FIXME: we might ultimately be calling this function with inputs
566 // from the CPU, these are unnecessary copies
567 if (indicesOptions_ == INDICES_CPU) {
568 HostTensor<long, 2, true> hostOutIndices(outIndices, stream);
569
570 ivfOffsetToUserIndex(hostOutIndices.data(),
571 numLists_,
572 hostOutIndices.getSize(0),
573 hostOutIndices.getSize(1),
574 listOffsetToUserIndex_);
575
576 // Copy back to GPU, since the input to this function is on the
577 // GPU
578 outIndices.copyFrom(hostOutIndices, stream);
579 }
580 }
581
582 std::vector<unsigned char>
583 IVFPQ::getListCodes(int listId) const {
584 FAISS_ASSERT(listId < deviceListData_.size());
585
586 return deviceListData_[listId]->copyToHost<unsigned char>(
587 resources_->getDefaultStreamCurrentDevice());
588 }
589
590 Tensor<float, 3, true>
591 IVFPQ::getPQCentroids() {
592 return pqCentroidsMiddleCode_;
593 }
594
595 void
596 IVFPQ::runPQPrecomputedCodes_(
597 Tensor<float, 2, true>& queries,
598 DeviceTensor<float, 2, true>& coarseDistances,
599 DeviceTensor<int, 2, true>& coarseIndices,
600 int k,
601 Tensor<float, 2, true>& outDistances,
602 Tensor<long, 2, true>& outIndices) {
603 auto& mem = resources_->getMemoryManagerCurrentDevice();
604 auto stream = resources_->getDefaultStreamCurrentDevice();
605
606 // Compute precomputed code term 3, - 2 * (x|y_R)
607 // This is done via batch MM
608 // {sub q} x {(query id)(sub dim) * (code id)(sub dim)'} =>
609 // {sub q} x {(query id)(code id)}
610 DeviceTensor<float, 3, true> term3Transposed(
611 mem,
612 {queries.getSize(0), numSubQuantizers_, numSubQuantizerCodes_},
613 stream);
614
615 // These allocations within are only temporary, so release them when
616 // we're done to maximize free space
617 {
618 auto querySubQuantizerView = queries.view<3>(
619 {queries.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
620 DeviceTensor<float, 3, true> queriesTransposed(
621 mem,
622 {numSubQuantizers_, queries.getSize(0), dimPerSubQuantizer_},
623 stream);
624 runTransposeAny(querySubQuantizerView, 0, 1, queriesTransposed, stream);
625
626 DeviceTensor<float, 3, true> term3(
627 mem,
628 {numSubQuantizers_, queries.getSize(0), numSubQuantizerCodes_},
629 stream);
630
631 runIteratedMatrixMult(term3, false,
632 queriesTransposed, false,
633 pqCentroidsMiddleCode_, true,
634 -2.0f, 0.0f,
635 resources_->getBlasHandleCurrentDevice(),
636 stream);
637
638 runTransposeAny(term3, 0, 1, term3Transposed, stream);
639 }
640
641 NoTypeTensor<3, true> term2;
642 NoTypeTensor<3, true> term3;
643 #ifdef FAISS_USE_FLOAT16
644 DeviceTensor<half, 3, true> term3Half;
645
646 if (useFloat16LookupTables_) {
647 term3Half = toHalf(resources_, stream, term3Transposed);
648 term2 = NoTypeTensor<3, true>(precomputedCodeHalf_);
649 term3 = NoTypeTensor<3, true>(term3Half);
650 }
651 #endif
652
653 if (!useFloat16LookupTables_) {
654 term2 = NoTypeTensor<3, true>(precomputedCode_);
655 term3 = NoTypeTensor<3, true>(term3Transposed);
656 }
657
658 runPQScanMultiPassPrecomputed(queries,
659 coarseDistances, // term 1
660 term2, // term 2
661 term3, // term 3
662 coarseIndices,
663 useFloat16LookupTables_,
664 bytesPerVector_,
665 numSubQuantizers_,
666 numSubQuantizerCodes_,
667 deviceListDataPointers_,
668 deviceListIndexPointers_,
669 indicesOptions_,
670 deviceListLengths_,
671 maxListLength_,
672 k,
673 outDistances,
674 outIndices,
675 resources_);
676 }
677
678 void
679 IVFPQ::runPQNoPrecomputedCodes_(
680 Tensor<float, 2, true>& queries,
681 DeviceTensor<float, 2, true>& coarseDistances,
682 DeviceTensor<int, 2, true>& coarseIndices,
683 int k,
684 Tensor<float, 2, true>& outDistances,
685 Tensor<long, 2, true>& outIndices) {
686 FAISS_ASSERT(!quantizer_->getUseFloat16());
687 auto& coarseCentroids = quantizer_->getVectorsFloat32Ref();
688
689 runPQScanMultiPassNoPrecomputed(queries,
690 coarseCentroids,
691 pqCentroidsInnermostCode_,
692 coarseIndices,
693 useFloat16LookupTables_,
694 bytesPerVector_,
695 numSubQuantizers_,
696 numSubQuantizerCodes_,
697 deviceListDataPointers_,
698 deviceListIndexPointers_,
699 indicesOptions_,
700 deviceListLengths_,
701 maxListLength_,
702 k,
703 outDistances,
704 outIndices,
705 resources_);
706 }
707
708 } } // namespace
faiss::gpu::IVFBase::numLists_
const int numLists_
Number of inverted lists we maintain.
Definition: IVFBase.cuh:89
faiss::gpu::IVFBase::maxListLength_
int maxListLength_
Maximum list length seen.
Definition: IVFBase.cuh:113
faiss::gpu::GpuResources::getDefaultStreamCurrentDevice
cudaStream_t getDefaultStreamCurrentDevice()
Calls getDefaultStream with the current device.
Definition: GpuResources.cpp:23
faiss::gpu::IVFPQ::addCodeVectorsFromCpu
void addCodeVectorsFromCpu(int listId, const void *codes, const long *indices, size_t numVecs)
Definition: IVFPQ.cu:345
faiss::gpu::FlatIndex::getSize
int getSize() const
Returns the number of vectors we contain.
Definition: FlatIndex.cu:45
faiss::gpu::IVFBase::listOffsetToUserIndex_
std::vector< std::vector< long > > listOffsetToUserIndex_
Definition: IVFBase.cuh:125
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:21
faiss::gpu::Tensor::view
__host__ __device__ Tensor< T, SubDim, InnerContig, IndexT, PtrTraits > view(DataPtrType at)
Definition: Tensor-inl.cuh:632
faiss::gpu::IVFBase
Base inverted list functionality for IVFFlat and IVFPQ.
Definition: IVFBase.cuh:25
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::IVFBase::deviceListLengths_
thrust::device_vector< int > deviceListLengths_
Definition: IVFBase.cuh:110
faiss::gpu::IVFPQ::isSupportedPQCodeLength
static bool isSupportedPQCodeLength(int size)
Returns true if we support PQ in this size.
Definition: IVFPQ.cu:70
faiss::gpu::IVFBase::deviceListIndexPointers_
thrust::device_vector< void * > deviceListIndexPointers_
Definition: IVFBase.cuh:106
faiss::gpu::GpuResources::getBlasHandleCurrentDevice
cublasHandle_t getBlasHandleCurrentDevice()
Calls getBlasHandle with the current device.
Definition: GpuResources.cpp:18
faiss::gpu::GpuResources::getMemoryManagerCurrentDevice
DeviceMemory & getMemoryManagerCurrentDevice()
Calls getMemoryManager for the current device.
Definition: GpuResources.cpp:33
faiss::gpu::IVFPQ::classifyAndAddVectors
int classifyAndAddVectors(Tensor< float, 2, true > &vecs, Tensor< long, 1, true > &indices)
Definition: IVFPQ.cu:118
faiss::gpu::Tensor::copyFrom
__host__ void copyFrom(Tensor< T, Dim, InnerContig, IndexT, PtrTraits > &t, cudaStream_t stream)
Copies a tensor into ourselves; sizes must match.
Definition: Tensor-inl.cuh:130
faiss::gpu::IVFPQ::query
void query(Tensor< float, 2, true > &queries, int nprobe, int k, Tensor< float, 2, true > &outDistances, Tensor< long, 2, true > &outIndices)
Definition: IVFPQ.cu:516
faiss::gpu::IVFPQ::getPQCentroids
Tensor< float, 3, true > getPQCentroids()
Definition: IVFPQ.cu:591
faiss::gpu::IVFBase::quantizer_
FlatIndex * quantizer_
Quantizer object.
Definition: IVFBase.cuh:83
faiss::gpu::IVFPQ::setPrecomputedCodes
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
Definition: IVFPQ.cu:100
faiss::gpu::IVFPQ::getListCodes
std::vector< unsigned char > getListCodes(int listId) const
Return the list codes of a particular list back to the CPU.
Definition: IVFPQ.cu:583
faiss::gpu::Tensor::getSize
__host__ __device__ IndexT getSize(int i) const
Definition: Tensor.cuh:222
faiss::gpu::IVFBase::deviceListDataPointers_
thrust::device_vector< void * > deviceListDataPointers_
Definition: IVFBase.cuh:102
faiss::gpu::Tensor::data
__host__ __device__ DataPtrType data()
Returns a raw pointer to the start of our data.
Definition: Tensor.cuh:174
faiss::gpu::IVFBase::resources_
GpuResources * resources_
Collection of GPU resources that we use.
Definition: IVFBase.cuh:80
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::IVFBase::bytesPerVector_
const int bytesPerVector_
Number of bytes per vector in the list.
Definition: IVFBase.cuh:92
faiss::gpu::IVFBase::updateDeviceListInfo_
void updateDeviceListInfo_(cudaStream_t stream)
Update all device-side list pointer and size information.
Definition: IVFBase.cu:136
faiss::gpu::FlatIndex::getVectorsFloat32Ref
Tensor< float, 2, true > & getVectorsFloat32Ref()
Returns a reference to our vectors currently in use.
Definition: FlatIndex.cu:77
faiss::gpu::IVFBase::indicesOptions_
const IndicesOptions indicesOptions_
How are user indices stored on the GPU?
Definition: IVFBase.cuh:95
faiss::gpu::HostTensor
Definition: HostTensor.cuh:20
faiss::gpu::IVFBase::deviceListData_
std::vector< std::unique_ptr< DeviceVector< unsigned char > > > deviceListData_
Definition: IVFBase.cuh:119
faiss::gpu::DeviceTensor< float, 3, true >
faiss::gpu::IVFPQ::IVFPQ
IVFPQ(GpuResources *resources, FlatIndex *quantizer, int numSubQuantizers, int bitsPerSubQuantizer, float *pqCentroidData, IndicesOptions indicesOptions, bool useFloat16LookupTables, MemorySpace space)
Definition: IVFPQ.cu:33
faiss::gpu::IVFBase::dim_
const int dim_
Expected dimensionality of the vectors.
Definition: IVFBase.cuh:86
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:243
faiss::gpu::IVFPQ::isSupportedNoPrecomputedSubDimSize
static bool isSupportedNoPrecomputedSubDimSize(int dims)
Definition: IVFPQ.cu:95