Facebook sync 2020-07-31 (#1308)

AutoTune.h

@@ -28,7 +28,7 @@ struct AutoTuneCriterion {
     typedef Index::idx_t idx_t;
     idx_t nq;  ///< nb of queries this criterion is evaluated on
     idx_t nnn; ///< nb of NNs that the query should request
-    idx_t gt_nnn; ///< nb of GT NNs required to evaluate criterion
+    idx_t gt_nnn; ///< nb of GT NNs required to evaluate crterion
 
     std::vector<float> gt_D;  ///< Ground-truth distances (size nq * gt_nnn)
     std::vector<idx_t> gt_I;  ///< Ground-truth indexes (size nq * gt_nnn)

IndexBinaryHash.cpp

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Copyright 2004-present Facebook. All Rights Reserved
 // -*- c++ -*-
 
 #include <faiss/IndexBinaryHash.h>

IndexBinaryIVF.cpp

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Copyright 2004-present Facebook. All Rights Reserved
 // -*- c++ -*-
 
 #include <faiss/IndexBinaryIVF.h>

IndexHNSW.cpp

@@ -289,7 +289,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
 {
     FAISS_THROW_IF_NOT_MSG(storage,
        "Please use IndexHSNWFlat (or variants) instead of IndexHNSW directly");
-    size_t nreorder = 0;
+    HNSWStats stats;
 
     idx_t check_period = InterruptCallback::get_period_hint (
           hnsw.max_level * d * hnsw.efSearch);
@@ -297,21 +297,22 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
 
-#pragma omp parallel reduction(+ : nreorder)
+#pragma omp parallel
         {
             VisitedTable vt (ntotal);
 
             DistanceComputer *dis = storage_distance_computer(storage);
             ScopeDeleter1<DistanceComputer> del(dis);
 
-#pragma omp for
+#pragma omp declare reduction (stats_combine : HNSWStats : omp_out.combine(omp_in))
+#pragma omp for reduction (stats_combine : stats)
             for(idx_t i = i0; i < i1; i++) {
                 idx_t * idxi = labels + i * k;
                 float * simi = distances + i * k;
                 dis->set_query(x + i * d);
 
                 maxheap_heapify (k, simi, idxi);
-                hnsw.search(*dis, k, idxi, simi, vt);
+                stats.combine(hnsw.search(*dis, k, idxi, simi, vt));
 
                 maxheap_reorder (k, simi, idxi);
 
@@ -320,7 +321,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
                     int k_reorder = reconstruct_from_neighbors->k_reorder;
                     if (k_reorder == -1 || k_reorder > k) k_reorder = k;
 
-                    nreorder += reconstruct_from_neighbors->compute_distances(
+                    stats.nreorder += reconstruct_from_neighbors->compute_distances(
                              k_reorder, idxi, x + i * d, simi);
 
                     // sort top k_reorder
@@ -341,7 +342,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
         }
     }
 
-    hnsw_stats.nreorder += nreorder;
+    hnsw_stats.combine(stats);
 }
 
 
@@ -416,6 +417,8 @@ void IndexHNSW::search_level_0(
 {
 
     storage_idx_t ntotal = hnsw.levels.size();
+    HNSWStats stats;
+
 #pragma omp parallel
     {
         DistanceComputer *qdis = storage_distance_computer(storage);
@@ -423,7 +426,8 @@ void IndexHNSW::search_level_0(
 
         VisitedTable vt (ntotal);
 
-#pragma omp for
+#pragma omp declare reduction (stats_combine : HNSWStats : omp_out.combine(omp_in))
+#pragma omp for reduction (stats_combine : stats)
         for(idx_t i = 0; i < n; i++) {
             idx_t * idxi = labels + i * k;
             float * simi = distances + i * k;
@@ -449,7 +453,7 @@ void IndexHNSW::search_level_0(
 
                     nres = hnsw.search_from_candidates(
                       *qdis, k, idxi, simi,
-                      candidates, vt, 0, nres
+                      candidates, vt, stats, 0, nres
                     );
                 }
             } else if (search_type == 2) {
@@ -466,7 +470,7 @@ void IndexHNSW::search_level_0(
                 }
                 hnsw.search_from_candidates(
                   *qdis, k, idxi, simi,
-                  candidates, vt, 0
+                  candidates, vt, stats, 0
                 );
 
             }
@@ -477,7 +481,7 @@ void IndexHNSW::search_level_0(
         }
     }
 
-
+    hnsw_stats.combine(stats);
 }
 
 void IndexHNSW::init_level_0_from_knngraph(
@@ -948,6 +952,7 @@ int search_from_candidates_2(const HNSW & hnsw,
                              idx_t *I, float * D,
                              MinimaxHeap &candidates,
                              VisitedTable &vt,
+                             HNSWStats &stats,
                              int level, int nres_in = 0)
 {
     int nres = nres_in;
@@ -996,15 +1001,9 @@ int search_from_candidates_2(const HNSW & hnsw,
         }
     }
 
-    if (level == 0) {
-#pragma omp critical
-        {
-            hnsw_stats.n1 ++;
-            if (candidates.size() == 0)
-                hnsw_stats.n2 ++;
-        }
-    }
-
+    stats.n1 ++;
+    if (candidates.size() == 0)
+        stats.n2 ++;
 
     return nres;
 }
@@ -1019,6 +1018,7 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
         IndexHNSW::search (n, x, k, distances, labels);
 
     } else { // "mixed" search
+        HNSWStats stats;
 
         const IndexIVFPQ *index_ivfpq =
             dynamic_cast<const IndexIVFPQ*>(storage);
@@ -1044,7 +1044,8 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
             int candidates_size = hnsw.upper_beam;
             MinimaxHeap candidates(candidates_size);
 
-#pragma omp for
+#pragma omp declare reduction (stats_combine : HNSWStats : omp_out.combine(omp_in))
+#pragma omp for reduction (stats_combine : stats)
             for(idx_t i = 0; i < n; i++) {
                 idx_t * idxi = labels + i * k;
                 float * simi = distances + i * k;
@@ -1083,7 +1084,7 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
 
                     hnsw.search_from_candidates(
                       *dis, k, idxi, simi,
-                      candidates, vt, 0, k
+                      candidates, vt, stats, 0, k
                     );
 
                     vt.advance();
@@ -1100,7 +1101,7 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
 
                     search_from_candidates_2 (
                         hnsw, *dis, k, idxi, simi,
-                        candidates, vt, 0, k);
+                        candidates, vt, stats, 0, k);
                     vt.advance ();
                     vt.advance ();
 
@@ -1109,6 +1110,8 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
                 maxheap_reorder (k, simi, idxi);
             }
         }
+
+        hnsw_stats.combine(stats);
     }
 
 

IndexIVF.cpp

@@ -300,10 +300,10 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
     bool do_heap_init = !(this->parallel_mode & PARALLEL_MODE_NO_HEAP_INIT);
 
     // don't start parallel section if single query
-    bool do_parallel =
-        pmode == 0 ? n > 1 :
-        pmode == 1 ? nprobe > 1 :
-        nprobe * n > 1;
+    bool do_parallel = omp_get_max_threads() >= 2 && (
+            pmode == 0 ? n > 1 :
+            pmode == 1 ? nprobe > 1 :
+            nprobe * n > 1);
 
 #pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis, nheap)
     {
@@ -327,6 +327,19 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
             }
         };
 
+        auto add_local_results = [&](
+                const float * local_dis, const idx_t * local_idx,
+                float *simi, idx_t *idxi)
+        {
+            if (metric_type == METRIC_INNER_PRODUCT) {
+                heap_addn<HeapForIP>
+                    (k, simi, idxi, local_dis, local_idx, k);
+            } else {
+                heap_addn<HeapForL2>
+                    (k, simi, idxi, local_dis, local_idx, k);
+            }
+        };
+
         auto reorder_result = [&] (float *simi, idx_t *idxi) {
             if (!do_heap_init) return;
             if (metric_type == METRIC_INNER_PRODUCT) {
@@ -447,20 +460,42 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
 #pragma omp barrier
 #pragma omp critical
                 {
-                    if (metric_type == METRIC_INNER_PRODUCT) {
-                        heap_addn<HeapForIP>
-                            (k, simi, idxi,
-                             local_dis.data(), local_idx.data(), k);
-                    } else {
-                        heap_addn<HeapForL2>
-                            (k, simi, idxi,
-                             local_dis.data(), local_idx.data(), k);
-                    }
+                    add_local_results (local_dis.data(), local_idx.data(),
+                                       simi, idxi);
                 }
 #pragma omp barrier
 #pragma omp single
                 reorder_result (simi, idxi);
             }
+        } else if (pmode == 2) {
+            std::vector <idx_t> local_idx (k);
+            std::vector <float> local_dis (k);
+
+#pragma omp single
+            for (size_t i = 0; i < n; i++) {
+                init_result (distances + i * k, labels + i * k);
+            }
+
+#pragma omp for schedule(dynamic)
+            for (size_t ij = 0; ij < n * nprobe; ij++) {
+                size_t i = ij / nprobe;
+                size_t j = ij % nprobe;
+
+                scanner->set_query (x + i * d);
+                init_result (local_dis.data(), local_idx.data());
+                ndis += scan_one_list (
+                        keys [ij], coarse_dis[ij],
+                        local_dis.data(), local_idx.data());
+#pragma omp critical
+                {
+                    add_local_results (local_dis.data(), local_idx.data(),
+                                       distances + i * k, labels + i * k);
+                }
+            }
+#pragma omp single
+            for (size_t i = 0; i < n; i++) {
+                reorder_result (distances + i * k, labels + i * k);
+            }
         } else {
             FAISS_THROW_FMT ("parallel_mode %d not supported\n",
                              pmode);

IndexIVF.h

@@ -68,6 +68,7 @@ struct Level1Quantizer {
 struct IVFSearchParameters {
     size_t nprobe;            ///< number of probes at query time
     size_t max_codes;         ///< max nb of codes to visit to do a query
+    IVFSearchParameters(): nprobe(1), max_codes(0) {}
     virtual ~IVFSearchParameters () {}
 };
 

IndexIVFPQ.h

@@ -22,6 +22,7 @@ namespace faiss {
 struct IVFPQSearchParameters: IVFSearchParameters {
     size_t scan_table_threshold;   ///< use table computation or on-the-fly?
     int polysemous_ht;             ///< Hamming thresh for polysemous filtering
+    IVFPQSearchParameters (): scan_table_threshold(0), polysemous_ht(0) {}
     ~IVFPQSearchParameters () {}
 };
 

InvertedLists.h

@@ -208,6 +208,7 @@ struct ArrayInvertedLists: InvertedLists {
  * that can be stacked horizontally, vertically and sliced.
  *****************************************************************/
 
+/// invlists that fail for all write functions
 struct ReadOnlyInvertedLists: InvertedLists {
 
     ReadOnlyInvertedLists (size_t nlist, size_t code_size):

OnDiskInvertedLists.cpp

@@ -16,10 +16,14 @@
 #include <sys/mman.h>
 #include <unistd.h>
 #include <sys/types.h>
+#include <sys/stat.h>
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/utils.h>
 
+#include <faiss/impl/io.h>
+#include <faiss/impl/io_macros.h>
+
 
 namespace faiss {
 
@@ -341,7 +345,7 @@ void OnDiskInvertedLists::update_totsize (size_t new_size)
 
 #define INVALID_OFFSET (size_t)(-1)
 
-OnDiskInvertedLists::List::List ():
+OnDiskOneList::OnDiskOneList ():
     size (0), capacity (0), offset (INVALID_OFFSET)
 {}
 
@@ -656,6 +660,12 @@ size_t OnDiskInvertedLists::merge_from (const InvertedLists **ils, int n_il,
 }
 
 
+size_t OnDiskInvertedLists::merge_from_1 (const InvertedLists *ils, bool verbose)
+{
+    return merge_from (&ils, 1, verbose);
+}
+
+
 void OnDiskInvertedLists::crop_invlists(size_t l0, size_t l1)
 {
     FAISS_THROW_IF_NOT(0 <= l0 && l0 <= l1 && l1 <= nlist);
@@ -668,6 +678,134 @@ void OnDiskInvertedLists::crop_invlists(size_t l0, size_t l1)
     nlist = l1 - l0;
 }
 
+/*******************************************************
+ * I/O support via callbacks
+ *******************************************************/
+
+
+
+
+OnDiskInvertedListsIOHook::OnDiskInvertedListsIOHook():
+    InvertedListsIOHook("ilod", typeid(OnDiskInvertedLists).name())
+{}
+
+
+void OnDiskInvertedListsIOHook::write(const InvertedLists *ils, IOWriter *f) const
+{
+    uint32_t h = fourcc ("ilod");
+    WRITE1 (h);
+    WRITE1 (ils->nlist);
+    WRITE1 (ils->code_size);
+    const OnDiskInvertedLists *od = dynamic_cast<const OnDiskInvertedLists*> (ils);
+    // this is a POD object
+    WRITEVECTOR (od->lists);
+
+    {
+        std::vector<OnDiskInvertedLists::Slot> v(
+                od->slots.begin(), od->slots.end());
+        WRITEVECTOR(v);
+    }
+    {
+        std::vector<char> x(od->filename.begin(), od->filename.end());
+        WRITEVECTOR(x);
+    }
+    WRITE1(od->totsize);
+}
+
+InvertedLists * OnDiskInvertedListsIOHook::read(IOReader *f, int io_flags) const
+{
+    OnDiskInvertedLists *od = new OnDiskInvertedLists();
+    od->read_only = io_flags & IO_FLAG_READ_ONLY;
+    READ1 (od->nlist);
+    READ1 (od->code_size);
+    // this is a POD object
+    READVECTOR (od->lists);
+    {
+        std::vector<OnDiskInvertedLists::Slot> v;
+        READVECTOR(v);
+        od->slots.assign(v.begin(), v.end());
+    }
+    {
+        std::vector<char> x;
+        READVECTOR(x);
+        od->filename.assign(x.begin(), x.end());
+
+        if (io_flags & IO_FLAG_ONDISK_SAME_DIR) {
+            FileIOReader *reader = dynamic_cast<FileIOReader*>(f);
+            FAISS_THROW_IF_NOT_MSG (
+                    reader, "IO_FLAG_ONDISK_SAME_DIR only supported "
+                    "when reading from file");
+            std::string indexname = reader->name;
+            std::string dirname = "./";
+            size_t slash = indexname.find_last_of('/');
+            if (slash != std::string::npos) {
+                dirname = indexname.substr(0, slash + 1);
+            }
+            std::string filename = od->filename;
+            slash = filename.find_last_of('/');
+            if (slash != std::string::npos) {
+                filename = filename.substr(slash + 1);
+            }
+            filename = dirname + filename;
+            printf("IO_FLAG_ONDISK_SAME_DIR: "
+                   "updating ondisk filename from %s to %s\n",
+                   od->filename.c_str(), filename.c_str());
+            od->filename = filename;
+        }
+
+    }
+    READ1(od->totsize);
+    od->do_mmap();
+    return od;
+}
+
+/** read from a ArrayInvertedLists into this invertedlist type */
+InvertedLists * OnDiskInvertedListsIOHook::read_ArrayInvertedLists(
+        IOReader *f, int /* io_flags */,
+        size_t nlist, size_t code_size,
+        const std::vector<size_t> &sizes) const
+{
+    auto ails = new OnDiskInvertedLists ();
+    ails->nlist = nlist;
+    ails->code_size = code_size;
+    ails->read_only = true;
+    ails->lists.resize (nlist);
+
+    FileIOReader *reader = dynamic_cast<FileIOReader*>(f);
+    FAISS_THROW_IF_NOT_MSG(reader, "mmap only supported for File objects");
+    FILE *fdesc = reader->f;
+    size_t o0 = ftell(fdesc);
+    size_t o = o0;
+    { // do the mmap
+        struct stat buf;
+        int ret = fstat (fileno(fdesc), &buf);
+        FAISS_THROW_IF_NOT_FMT (ret == 0,
+                                "fstat failed: %s", strerror(errno));
+        ails->totsize = buf.st_size;
+        ails->ptr = (uint8_t*)mmap (nullptr, ails->totsize,
+                                    PROT_READ, MAP_SHARED,
+                                    fileno(fdesc), 0);
+        FAISS_THROW_IF_NOT_FMT (ails->ptr != MAP_FAILED,
+                                "could not mmap: %s",
+                                strerror(errno));
+    }
+
+    FAISS_THROW_IF_NOT(o <= ails->totsize);
+
+    for (size_t i = 0; i < ails->nlist; i++) {
+        OnDiskInvertedLists::List & l = ails->lists[i];
+        l.size = l.capacity = sizes[i];
+        l.offset = o;
+        o += l.size * (sizeof(OnDiskInvertedLists::idx_t) +
+                       ails->code_size);
+    }
+    // resume normal reading of file
+    fseek (fdesc, o, SEEK_SET);
+
+    return ails;
+}
+
+
 
 
 

OnDiskInvertedLists.h

@@ -12,14 +12,26 @@
 
 #include <vector>
 #include <list>
+#include <typeinfo>
 
 #include <faiss/IndexIVF.h>
 
+#include <faiss/index_io.h>
+
 namespace faiss {
 
 
 struct LockLevels;
 
+
+struct OnDiskOneList {
+    size_t size;     // size of inverted list (entries)
+    size_t capacity; // allocated size (entries)
+    size_t offset;   // offset in buffer (bytes)
+    OnDiskOneList ();
+};
+
+
 /** On-disk storage of inverted lists.
  *
  * The data is stored in a mmapped chunk of memory (base ptointer ptr,
@@ -49,13 +61,7 @@ struct LockLevels;
  * lists in parallel.
  */
 struct OnDiskInvertedLists: InvertedLists {
-
-    struct List {
-        size_t size;     // size of inverted list (entries)
-        size_t capacity; // allocated size (entries)
-        size_t offset;   // offset in buffer (bytes)
-        List ();
-    };
+    using List = OnDiskOneList;
 
     // size nlist
     std::vector<List> lists;
@@ -95,6 +101,9 @@ struct OnDiskInvertedLists: InvertedLists {
     // allocating slots)
     size_t merge_from (const InvertedLists **ils, int n_il, bool verbose=false);
 
+    /// same as merge_from for a single invlist
+    size_t merge_from_1 (const InvertedLists *il, bool verbose=false);
+
     /// restrict the inverted lists to l0:l1 without touching the mmapped region
     void crop_invlists(size_t l0, size_t l1);
 
@@ -121,6 +130,17 @@ struct OnDiskInvertedLists: InvertedLists {
     OnDiskInvertedLists ();
 };
 
+struct OnDiskInvertedListsIOHook: InvertedListsIOHook {
+    OnDiskInvertedListsIOHook();
+    void write(const InvertedLists *ils, IOWriter *f) const override;
+    InvertedLists * read(IOReader *f, int io_flags) const override;
+    InvertedLists * read_ArrayInvertedLists(
+            IOReader *f, int io_flags,
+            size_t nlist, size_t code_size,
+            const std::vector<size_t> &sizes) const override;
+};
+
+
 
 } // namespace faiss
 

benchs/bench_6bit_codec.cpp

@@ -0,0 +1,80 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+
+#include <cstdio>
+#include <omp.h>
+
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/utils/utils.h>
+#include <faiss/utils/random.h>
+#include <faiss/utils/distances.h>
+
+using namespace faiss;
+
+int main() {
+
+    int d = 128;
+    int n = 2000;
+
+    std::vector<float> x(d * n);
+
+    float_rand(x.data(), d * n, 12345);
+
+    // make sure it's idempotent
+    ScalarQuantizer sq(d, ScalarQuantizer::QT_6bit);
+
+    omp_set_num_threads(1);
+
+    sq.train(n, x.data());
+
+    size_t code_size = sq.code_size;
+    printf("code size: %ld\n", sq.code_size);
+
+    // encode
+    std::vector<uint8_t> codes(code_size * n);
+    sq.compute_codes(x.data(), codes.data(), n);
+
+    // decode
+    std::vector<float> x2(d * n);
+    sq.decode(codes.data(), x2.data(), n);
+
+    printf("sqL2 recons error: %g\n",
+           fvec_L2sqr(x.data(), x2.data(), n * d) / n);
+
+    // encode again
+    std::vector<uint8_t> codes2(code_size * n);
+    sq.compute_codes(x2.data(), codes2.data(), n);
+
+    size_t ndiff = 0;
+    for (size_t i = 0; i < codes.size(); i++) {
+        if (codes[i] != codes2[i]) ndiff++;
+    }
+
+    printf("ndiff for idempotence: %ld / %ld\n", ndiff, codes.size());
+
+    std::unique_ptr<ScalarQuantizer::SQDistanceComputer>
+        dc(sq.get_distance_computer());
+    dc->codes = codes.data();
+    dc->code_size = sq.code_size;
+    printf("code size: %ld\n", dc->code_size);
+
+    double sum_dis = 0;
+    double t0 = getmillisecs();
+    for (int i = 0; i < n; i++) {
+        dc->set_query(&x[i * d]);
+        for (int j = 0; j < n; j++) {
+            sum_dis += (*dc)(j);
+        }
+    }
+    printf("distances computed in %.3f ms, checksum=%g\n",
+           getmillisecs() - t0, sum_dis);
+
+
+
+    return 0;
+}

benchs/bench_all_ivf/bench_all_ivf.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 import os
 import sys
 import time

benchs/bench_all_ivf/bench_kmeans.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import os
 import numpy as np

benchs/bench_all_ivf/datasets.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 """
 Common functions to load datasets and compute their ground-truth
 """

benchs/bench_all_ivf/parse_bench_all_ivf.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 import os
 import numpy as np
 from matplotlib import pyplot

benchs/bench_for_interrupt.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python3
-
 from __future__ import print_function
 import numpy as np
 import faiss

benchs/bench_gpu_1bn.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 from __future__ import print_function
 import numpy as np
 import time

benchs/bench_gpu_sift1m.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import os
 import time

benchs/bench_hnsw.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import time
 import sys

benchs/bench_pairwise_distances.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python3
-
 """small test script to benchmark the SIMD implementation of the
 distance computations for the additional metrics. Call eg. with L1 to
 get L1 distance computations.

benchs/bench_polysemous_1bn.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import os
 import sys

benchs/bench_polysemous_sift1m.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import time
 import numpy as np

benchs/bench_scalar_quantizer.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import time
 import numpy as np

benchs/bench_vector_ops.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 from __future__ import print_function
 import numpy as np
 import faiss

benchs/distributed_ondisk/README.md

@@ -1,4 +1,3 @@
-
 # Distributed on-disk index for 1T-scale datasets 
 
 This is code corresponding to the description in [Indexing 1T vectors](https://github.com/facebookresearch/faiss/wiki/Indexing-1T-vectors). 

benchs/distributed_ondisk/combined_index.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python3
-
 import os
 import faiss
 import numpy as np

benchs/distributed_ondisk/distributed_kmeans.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python3
-
 """
 Simple distributed kmeans implementation Relies on an abstraction
 for the training matrix, that can be sharded over several machines.

benchs/distributed_ondisk/rpc.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python3
-
 """
 Simplistic RPC implementation.
 Exposes all functions of a Server object.

benchs/distributed_ondisk/run_on_cluster.bash

@@ -1,4 +1,5 @@
 #! /bin/bash
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the

benchs/kmeans_mnist.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 import numpy as np
 import time
 import faiss

benchs/link_and_code/README.md

@@ -1,6 +1,5 @@
 
 
-
 README for the link & code implementation
 =========================================
 

benchs/link_and_code/bench_link_and_code.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 import os
 import sys
 import time

benchs/link_and_code/datasets.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 """
 Common functions to load datasets and compute their ground-truth
 """

benchs/link_and_code/neighbor_codec.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python2
-
 """
 This is the training code for the link and code. Especially the
 neighbors_kmeans function implements the EM-algorithm to find the

contrib/client_server.py

@@ -0,0 +1,88 @@
+#! /usr/bin/env python3
+
+from multiprocessing.dummy import Pool as ThreadPool
+import faiss
+from typing import List, Tuple
+
+from . import rpc
+
+############################################################
+# Server implementation
+############################################################
+
+
+class SearchServer(rpc.Server):
+    """ Assign version that can be exposed via RPC """
+
+    def __init__(self, s: int, index: faiss.Index):
+        rpc.Server.__init__(self, s)
+        self.index = index
+        self.index_ivf = faiss.extract_index_ivf(index)
+
+    def set_nprobe(self, nprobe: int) -> int:
+        """ set nprobe field """
+        self.index_ivf.nprobe = nprobe
+
+    def get_ntotal(self) -> int:
+        return self.index.ntotal
+
+    def __getattr__(self, f):
+        # all other functions get forwarded to the index
+        return getattr(self.index, f)
+
+
+def run_index_server(index: faiss.Index, port: int, v6: bool = False):
+    """ serve requests for that index forerver """
+    rpc.run_server(
+        lambda s: SearchServer(s, index),
+        port, v6=v6)
+
+
+############################################################
+# Client implementation
+############################################################
+
+class ClientIndex:
+    """manages a set of distance sub-indexes. The sub_indexes search a
+    subset of the inverted lists. Searches are merged afterwards
+    """
+
+    def __init__(self, machine_ports: List[Tuple[str, int]], v6: bool = False):
+        """ connect to a series of (host, port) pairs """
+        self.sub_indexes = []
+        for machine, port in machine_ports:
+            self.sub_indexes.append(rpc.Client(machine, port, v6))
+
+        self.ni = len(self.sub_indexes)
+        # pool of threads. Each thread manages one sub-index.
+        self.pool = ThreadPool(self.ni)
+        # test connection...
+        self.ntotal = self.get_ntotal()
+        self.verbose = False
+
+    def set_nprobe(self, nprobe: int) -> None:
+        self.pool.map(
+            lambda idx: idx.set_nprobe(nprobe),
+            self.sub_indexes
+        )
+
+    def set_omp_num_threads(self, nt: int) -> None:
+        self.pool.map(
+            lambda idx: idx.set_omp_num_threads(nt),
+            self.sub_indexes
+        )
+
+    def get_ntotal(self) -> None:
+        return sum(self.pool.map(
+            lambda idx: idx.get_ntotal(),
+            self.sub_indexes
+        ))
+
+    def search(self, x, k: int):
+
+        rh = faiss.ResultHeap(x.shape[0], k)
+
+        for Di, Ii in self.pool.imap(lambda idx: idx.search(x, k), self.sub_indexes):
+            rh.add_result(Di, Ii)
+        rh.finalize()
+        return rh.D, rh.I

contrib/exhaustive_search.py

@@ -0,0 +1,39 @@
+#! /usr/bin/env python3
+
+import faiss
+import time
+
+import logging
+
+LOG = logging.getLogger(__name__)
+
+def knn_ground_truth(xq, db_iterator, k):
+    """Computes the exact KNN search results for a dataset that possibly
+    does not fit in RAM but for whihch we have an iterator that
+    returns it block by block.
+    """
+    t0 = time.time()
+    nq, d = xq.shape
+    rh = faiss.ResultHeap(nq, k)
+
+    index = faiss.IndexFlatL2(d)
+    if faiss.get_num_gpus():
+        LOG.info('running on %d GPUs' % faiss.get_num_gpus())
+        index = faiss.index_cpu_to_all_gpus(index)
+
+    # compute ground-truth by blocks of bs, and add to heaps
+    i0 = 0
+    for xbi in db_iterator:
+        ni = xbi.shape[0]
+        index.add(xbi)
+        D, I = index.search(xq, k)
+        I += i0
+        rh.add_result(D, I)
+        index.reset()
+        i0 += ni
+        LOG.info("%d db elements, %.3f s" % (i0, time.time() - t0))
+
+    rh.finalize()
+    LOG.info("GT time: %.3f s (%d vectors)" % (time.time() - t0, i0))
+
+    return rh.D, rh.I

contrib/ondisk.py

@@ -0,0 +1,52 @@
+#! /usr/bin/env python3
+
+from typing import List
+import faiss
+import logging
+
+LOG = logging.getLogger(__name__)
+
+def merge_ondisk(trained_index: faiss.Index,
+                 shard_fnames: List[str],
+                 ivfdata_fname: str) -> None:
+    """ Add the contents of the indexes stored in shard_fnames into the index
+    trained_index. The on-disk data is stored in ivfdata_fname """
+    # merge the images into an on-disk index
+    # first load the inverted lists
+    ivfs = []
+    for fname in shard_fnames:
+        # the IO_FLAG_MMAP is to avoid actually loading the data thus
+        # the total size of the inverted lists can exceed the
+        # available RAM
+        LOG.info("read " + fname)
+        index = faiss.read_index(fname, faiss.IO_FLAG_MMAP)
+        index_ivf = faiss.extract_index_ivf(index)
+        ivfs.append(index_ivf.invlists)
+
+        # avoid that the invlists get deallocated with the index
+        index_ivf.own_invlists = False
+
+    # construct the output index
+    index = trained_index
+    index_ivf = faiss.extract_index_ivf(index)
+
+    assert index.ntotal == 0, "works only on empty index"
+
+    # prepare the output inverted lists. They will be written
+    # to merged_index.ivfdata
+    invlists = faiss.OnDiskInvertedLists(
+        index_ivf.nlist, index_ivf.code_size,
+        ivfdata_fname)
+
+    # merge all the inverted lists
+    ivf_vector = faiss.InvertedListsPtrVector()
+    for ivf in ivfs:
+        ivf_vector.push_back(ivf)
+
+    LOG.info("merge %d inverted lists " % ivf_vector.size())
+    ntotal = invlists.merge_from(ivf_vector.data(), ivf_vector.size())
+
+    # now replace the inverted lists in the output index
+    index.ntotal = index_ivf.ntotal = ntotal
+    index_ivf.replace_invlists(invlists, True)
+    invlists.this.disown()

contrib/rpc.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+
+"""
+Simplistic RPC implementation.
+Exposes all functions of a Server object.
+
+Uses pickle for serialization and the socket interface.
+"""
+
+import os,pdb,pickle,time,errno,sys,_thread,traceback,socket,threading,gc
+
+import logging
+
+LOG = logging.getLogger(__name__)
+
+# default
+PORT=12032
+
+
+#########################################################################
+# simple I/O functions
+
+
+def inline_send_handle(f, conn):
+    st = os.fstat(f.fileno())
+    size = st.st_size
+    pickle.dump(size, conn)
+    conn.write(f.read(size))
+
+def inline_send_string(s, conn):
+    size = len(s)
+    pickle.dump(size, conn)
+    conn.write(s)
+
+
+class FileSock:
+    " wraps a socket so that it is usable by pickle/cPickle "
+
+    def __init__(self,sock):
+        self.sock = sock
+        self.nr=0
+
+    def write(self, buf):
+        # print("sending %d bytes"%len(buf))
+        #self.sock.sendall(buf)
+        # print("...done")
+        bs = 512 * 1024
+        ns = 0
+        while ns < len(buf):
+            sent = self.sock.send(buf[ns:ns + bs])
+            ns += sent
+
+    def read(self,bs=512*1024):
+        #if self.nr==10000: pdb.set_trace()
+        self.nr+=1
+        # print("read bs=%d"%bs)
+        b = []
+        nb = 0
+        while len(b)<bs:
+            # print('   loop')
+            rb = self.sock.recv(bs - nb)
+            if not rb: break
+            b.append(rb)
+            nb += len(rb)
+        return b''.join(b)
+
+    def readline(self):
+        # print("readline!")
+        """may be optimized..."""
+        s=bytes()
+        while True:
+            c=self.read(1)
+            s+=c
+        if len(c)==0 or chr(c[0])=='\n':
+            return s
+
+class ClientExit(Exception):
+    pass
+
+class ServerException(Exception):
+    pass
+
+
+class Server:
+    """
+    server protocol. Methods from classes that subclass Server can be called
+    transparently from a client
+    """
+
+    def __init__(self, s, logf=sys.stderr, log_prefix=''):
+        self.logf = logf
+        self.log_prefix = log_prefix
+
+        # connection
+
+        self.conn = s
+        self.fs = FileSock(s)
+
+
+    def log(self, s):
+        self.logf.write("Sever log %s: %s\n" % (self.log_prefix, s))
+
+    def one_function(self):
+        """
+        Executes a single function with associated I/O.
+        Protocol:
+        - the arguments and results are serialized with the pickle protocol
+        - client sends : (fname,args)
+            fname = method name to call
+            args = tuple of arguments
+        - server sends result: (rid,st,ret)
+            rid = request id
+            st = None, or exception if there was during execution
+            ret = return value or None if st!=None
+        """
+
+        try:
+            (fname,args)=pickle.load(self.fs)
+        except EOFError:
+            raise ClientExit("read args")
+        self.log("executing method %s"%(fname))
+        st = None
+        ret = None
+        try:
+            f=getattr(self,fname)
+        except AttributeError:
+            st = AttributeError("unknown method "+fname)
+            self.log("unknown method ")
+
+        try:
+            ret = f(*args)
+        except Exception as e:
+            # due to a bug (in mod_python?), ServerException cannot be
+            # unpickled, so send the string and make the exception on the client side
+
+            #st=ServerException(
+            #  "".join(traceback.format_tb(sys.exc_info()[2]))+
+            #  str(e))
+            st="".join(traceback.format_tb(sys.exc_info()[2]))+str(e)
+            self.log("exception in method")
+            traceback.print_exc(50,self.logf)
+            self.logf.flush()
+
+        LOG.info("return")
+        try:
+            pickle.dump((st ,ret), self.fs, protocol=4)
+        except EOFError:
+            raise ClientExit("function return")
+
+    def exec_loop(self):
+        """ main execution loop. Loops and handles exit states"""
+
+        self.log("in exec_loop")
+        try:
+            while True:
+                self.one_function()
+        except ClientExit as e:
+            self.log("ClientExit %s"%e)
+        except socket.error as e:
+            self.log("socket error %s"%e)
+            traceback.print_exc(50,self.logf)
+        except EOFError:
+            self.log("EOF during communication")
+            traceback.print_exc(50,self.logf)
+        except BaseException:
+            # unexpected
+            traceback.print_exc(50,sys.stderr)
+            sys.exit(1)
+
+        LOG.info("exit sever")
+
+    def exec_loop_cleanup(self):
+        pass
+
+    ###################################################################
+    # spying stuff
+
+    def get_ps_stats(self):
+        ret=''
+        f=os.popen("echo ============ `hostname` uptime:; uptime;"+
+                   "echo ============ self:; "+
+                   "ps -p %d -o pid,vsize,rss,%%cpu,nlwp,psr; "%os.getpid()+
+                   "echo ============ run queue:;"+
+                   "ps ar -o user,pid,%cpu,%mem,ni,nlwp,psr,vsz,rss,cputime,command")
+        for l in f:
+            ret+=l
+        return ret
+
+class Client:
+    """
+    Methods of the server object can be called transparently. Exceptions are
+    re-raised.
+    """
+    def __init__(self, HOST, port=PORT, v6=False):
+        socktype = socket.AF_INET6 if v6 else socket.AF_INET
+
+        sock = socket.socket(socktype, socket.SOCK_STREAM)
+        LOG.info("connecting", HOST, port, socktype)
+        sock.connect((HOST, port))
+        self.sock = sock
+        self.fs = FileSock(sock)
+
+    def generic_fun(self, fname, args):
+        # int "gen fun",fname
+        pickle.dump((fname, args), self.fs, protocol=4)
+        return self.get_result()
+
+    def get_result(self):
+        (st, ret) = pickle.load(self.fs)
+        if st!=None:
+            raise ServerException(st)
+        else:
+            return ret
+
+    def __getattr__(self,name):
+        return lambda *x: self.generic_fun(name,x)
+
+
+def run_server(new_handler, port=PORT, report_to_file=None, v6=False):
+
+    HOST = ''                 # Symbolic name meaning the local host
+    socktype = socket.AF_INET6 if v6 else socket.AF_INET
+    s = socket.socket(socktype, socket.SOCK_STREAM)
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+
+    LOG.info("bind %s:%d" % (HOST, port))
+    s.bind((HOST, port))
+    s.listen(5)
+
+    LOG.info("accepting connections")
+    if report_to_file is not None:
+        LOG.info('storing host+port in', report_to_file)
+        open(report_to_file, 'w').write('%s:%d ' % (socket.gethostname(), port))
+
+    while True:
+        try:
+            conn, addr = s.accept()
+        except socket.error as e:
+            if e[1]=='Interrupted system call': continue
+            raise
+
+        LOG.info('Connected by', addr, end=' ')
+
+        ibs = new_handler(conn)
+
+        tid = _thread.start_new_thread(ibs.exec_loop,())
+
+        LOG.info("tid",tid)

demos/demo_auto_tune.py

@@ -1,10 +1,10 @@
+#!/usr/bin/env python2
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 from __future__ import print_function
 import os
 import time

demos/demo_client_server_ivf.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+
+import sys
+import numpy as np
+import faiss
+
+from faiss_contrib.client_server import run_index_server, ClientIndex
+
+
+#################################################################
+# Small I/O functions
+#################################################################
+
+
+def ivecs_read(fname):
+    a = np.fromfile(fname, dtype='int32')
+    d = a[0]
+    return a.reshape(-1, d + 1)[:, 1:].copy()
+
+
+def fvecs_read(fname):
+    return ivecs_read(fname).view('float32')
+
+
+#################################################################
+#  Main program
+#################################################################
+
+stage = int(sys.argv[1])
+
+tmpdir = '/tmp/'
+
+if stage == 0:
+    # train the index
+    xt = fvecs_read("sift1M/sift_learn.fvecs")
+    index = faiss.index_factory(xt.shape[1], "IVF4096,Flat")
+    print("training index")
+    index.train(xt)
+    print("write " + tmpdir + "trained.index")
+    faiss.write_index(index, tmpdir + "trained.index")
+
+
+if 1 <= stage <= 4:
+    # add 1/4 of the database to 4 independent indexes
+    bno = stage - 1
+    xb = fvecs_read("sift1M/sift_base.fvecs")
+    i0, i1 = int(bno * xb.shape[0] / 4), int((bno + 1) * xb.shape[0] / 4)
+    index = faiss.read_index(tmpdir + "trained.index")
+    print("adding vectors %d:%d" % (i0, i1))
+    index.add_with_ids(xb[i0:i1], np.arange(i0, i1))
+    print("write " + tmpdir + "block_%d.index" % bno)
+    faiss.write_index(index, tmpdir + "block_%d.index" % bno)
+
+
+machine_ports = [
+    ('localhost', 12010),
+    ('localhost', 12011),
+    ('localhost', 12012),
+    ('localhost', 12013),
+]
+v6 = False
+
+if 5 <= stage <= 8:
+    # load an index slice and launch index
+    bno = stage - 5
+
+    fname = tmpdir + "block_%d.index" % bno
+    print("read " + fname)
+    index = faiss.read_index(fname)
+
+    port = machine_ports[bno][1]
+    run_index_server(index, port, v6=v6)
+
+
+if stage == 9:
+    client_index = ClientIndex(machine_ports)
+    print('index size:', client_index.ntotal)
+    client_index.set_nprobe(16)
+
+    # load query vectors and ground-truth
+    xq = fvecs_read("sift1M/sift_query.fvecs")
+    gt = ivecs_read("sift1M/sift_groundtruth.ivecs")
+
+    D, I = client_index.search(xq, 5)
+
+    recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / float(xq.shape[0])
+    print("recall@1: %.3f" % recall_at_1)

demos/demo_ondisk_ivf.py

@@ -1,14 +1,14 @@
+#!/usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#!/usr/bin/env python2
-
 import sys
 import numpy as np
 import faiss
-
+from faiss_contrib.ondisk import merge_ondisk
 
 #################################################################
 # Small I/O functions
@@ -26,7 +26,7 @@ def fvecs_read(fname):
 
 
 #################################################################
-#  Main program
+# Main program
 #################################################################
 
 stage = int(sys.argv[1])
@@ -54,43 +54,18 @@ if 1 <= stage <= 4:
     print("write " + tmpdir + "block_%d.index" % bno)
     faiss.write_index(index, tmpdir + "block_%d.index" % bno)
 
-
 if stage == 5:
-    # merge the images into an on-disk index
-    # first load the inverted lists
-    ivfs = []
-    for bno in range(4):
-        # the IO_FLAG_MMAP is to avoid actually loading the data thus
-        # the total size of the inverted lists can exceed the
-        # available RAM
-        print("read " + tmpdir + "block_%d.index" % bno)
-        index = faiss.read_index(tmpdir + "block_%d.index" % bno,
-                                 faiss.IO_FLAG_MMAP)
-        ivfs.append(index.invlists)
-
-        # avoid that the invlists get deallocated with the index
-        index.own_invlists = False
 
+    print('loading trained index')
     # construct the output index
     index = faiss.read_index(tmpdir + "trained.index")
 
-    # prepare the output inverted lists. They will be written
-    # to merged_index.ivfdata
-    invlists = faiss.OnDiskInvertedLists(
-        index.nlist, index.code_size,
-        tmpdir + "merged_index.ivfdata")
+    block_fnames = [
+        tmpdir + "block_%d.index" % bno
+        for bno in range(4)
+    ]
 
-    # merge all the inverted lists
-    ivf_vector = faiss.InvertedListsPtrVector()
-    for ivf in ivfs:
-        ivf_vector.push_back(ivf)
-
-    print("merge %d inverted lists " % ivf_vector.size())
-    ntotal = invlists.merge_from(ivf_vector.data(), ivf_vector.size())
-
-    # now replace the inverted lists in the output index
-    index.ntotal = ntotal
-    index.replace_invlists(invlists)
+    merge_ondisk(index, block_fnames, tmpdir + "merged_index.ivfdata")
 
     print("write " + tmpdir + "populated.index")
     faiss.write_index(index, tmpdir + "populated.index")

gpu/GpuCloner.cpp

@@ -114,9 +114,9 @@ faiss::Index * index_gpu_to_cpu(const faiss::Index *gpu_index)
  * Cloning to 1 GPU
  **********************************************************/
 
-ToGpuCloner::ToGpuCloner(GpuResources *resources, int device,
+ToGpuCloner::ToGpuCloner(GpuResourcesProvider *prov, int device,
                          const GpuClonerOptions &options):
-    GpuClonerOptions(options), resources(resources), device(device)
+    GpuClonerOptions(options), provider(prov), device(device)
 {}
 
 Index *ToGpuCloner::clone_Index(const Index *index)
@@ -127,7 +127,7 @@ Index *ToGpuCloner::clone_Index(const Index *index)
         config.useFloat16 = useFloat16;
         config.storeTransposed = storeTransposed;
 
-        return new GpuIndexFlat(resources, ifl, config);
+        return new GpuIndexFlat(provider, ifl, config);
     } else if(auto ifl = dynamic_cast<const faiss::IndexIVFFlat *>(index)) {
         GpuIndexIVFFlatConfig config;
         config.device = device;
@@ -136,7 +136,7 @@ Index *ToGpuCloner::clone_Index(const Index *index)
         config.flatConfig.storeTransposed = storeTransposed;
 
         GpuIndexIVFFlat *res =
-            new GpuIndexIVFFlat(resources,
+            new GpuIndexIVFFlat(provider,
                                 ifl->d,
                                 ifl->nlist,
                                 ifl->metric_type,
@@ -156,7 +156,7 @@ Index *ToGpuCloner::clone_Index(const Index *index)
         config.flatConfig.storeTransposed = storeTransposed;
 
         GpuIndexIVFScalarQuantizer *res =
-            new GpuIndexIVFScalarQuantizer(resources,
+            new GpuIndexIVFScalarQuantizer(provider,
                                            ifl->d,
                                            ifl->nlist,
                                            ifl->sq.qtype,
@@ -184,7 +184,7 @@ Index *ToGpuCloner::clone_Index(const Index *index)
         config.useFloat16LookupTables = useFloat16;
         config.usePrecomputedTables = usePrecomputed;
 
-        GpuIndexIVFPQ *res = new GpuIndexIVFPQ(resources, ipq, config);
+        GpuIndexIVFPQ *res = new GpuIndexIVFPQ(provider, ipq, config);
 
         if(reserveVecs > 0 && ipq->ntotal == 0) {
             res->reserveMemory(reserveVecs);
@@ -198,12 +198,12 @@ Index *ToGpuCloner::clone_Index(const Index *index)
 
 
 faiss::Index * index_cpu_to_gpu(
-       GpuResources* resources, int device,
+       GpuResourcesProvider* provider, int device,
        const faiss::Index *index,
        const GpuClonerOptions *options)
 {
     GpuClonerOptions defaults;
-    ToGpuCloner cl(resources, device, options ? *options : defaults);
+    ToGpuCloner cl(provider, device, options ? *options : defaults);
     return cl.clone_Index(index);
 }
 
@@ -213,14 +213,14 @@ faiss::Index * index_cpu_to_gpu(
  **********************************************************/
 
 ToGpuClonerMultiple::ToGpuClonerMultiple(
-                        std::vector<GpuResources *> & resources,
+                        std::vector<GpuResourcesProvider *> & provider,
                         std::vector<int>& devices,
                         const GpuMultipleClonerOptions &options):
     GpuMultipleClonerOptions(options)
 {
-    FAISS_ASSERT(resources.size() == devices.size());
-    for(int i = 0; i < resources.size(); i++) {
-        sub_cloners.push_back(ToGpuCloner(resources[i], devices[i], options));
+    FAISS_ASSERT(provider.size() == devices.size());
+    for(int i = 0; i < provider.size(); i++) {
+        sub_cloners.push_back(ToGpuCloner(provider[i], devices[i], options));
     }
 }
 
@@ -394,13 +394,13 @@ Index *ToGpuClonerMultiple::clone_Index(const Index *index)
 
 
 faiss::Index * index_cpu_to_gpu_multiple(
-       std::vector<GpuResources*> & resources,
+       std::vector<GpuResourcesProvider*> & provider,
        std::vector<int> &devices,
        const faiss::Index *index,
        const GpuMultipleClonerOptions *options)
 {
     GpuMultipleClonerOptions defaults;
-    ToGpuClonerMultiple cl(resources, devices, options ? *options : defaults);
+    ToGpuClonerMultiple cl(provider, devices, options ? *options : defaults);
     return cl.clone_Index(index);
 }
 

gpu/GpuCloner.h

@@ -17,8 +17,7 @@
 
 namespace faiss { namespace gpu {
 
-class GpuResources;
-
+class GpuResourcesProvider;
 
 /// Cloner specialized for GPU -> CPU
 struct ToCPUCloner: faiss::Cloner {
@@ -29,10 +28,10 @@ struct ToCPUCloner: faiss::Cloner {
 
 /// Cloner specialized for CPU -> 1 GPU
 struct ToGpuCloner: faiss::Cloner, GpuClonerOptions {
-    GpuResources *resources;
+    GpuResourcesProvider *provider;
     int device;
 
-    ToGpuCloner(GpuResources *resources, int device,
+    ToGpuCloner(GpuResourcesProvider *prov, int device,
                 const GpuClonerOptions &options);
 
     Index *clone_Index(const Index *index) override;
@@ -43,7 +42,7 @@ struct ToGpuCloner: faiss::Cloner, GpuClonerOptions {
 struct ToGpuClonerMultiple: faiss::Cloner, GpuMultipleClonerOptions {
     std::vector<ToGpuCloner> sub_cloners;
 
-    ToGpuClonerMultiple(std::vector<GpuResources *> & resources,
+    ToGpuClonerMultiple(std::vector<GpuResourcesProvider *> & provider,
                         std::vector<int>& devices,
                         const GpuMultipleClonerOptions &options);
 
@@ -67,12 +66,12 @@ faiss::Index * index_gpu_to_cpu(const faiss::Index *gpu_index);
 
 /// converts any CPU index that can be converted to GPU
 faiss::Index * index_cpu_to_gpu(
-       GpuResources* resources, int device,
+       GpuResourcesProvider* provider, int device,
        const faiss::Index *index,
        const GpuClonerOptions *options = nullptr);
 
 faiss::Index * index_cpu_to_gpu_multiple(
-       std::vector<GpuResources*> & resources,
+       std::vector<GpuResourcesProvider*> & provider,
        std::vector<int> &devices,
        const faiss::Index *index,
        const GpuMultipleClonerOptions *options = nullptr);

gpu/GpuDistance.cu

@@ -18,49 +18,56 @@
 namespace faiss { namespace gpu {
 
 template <typename T>
-void bfKnnConvert(GpuResources* resources, const GpuDistanceParams& args) {
+void bfKnnConvert(GpuResourcesProvider* prov, const GpuDistanceParams& args) {
+  // Don't let the resources go out of scope
+  auto resImpl = prov->getResources();
+  auto res = resImpl.get();
   auto device = getCurrentDevice();
-  auto stream = resources->getDefaultStreamCurrentDevice();
-  auto& mem = resources->getMemoryManagerCurrentDevice();
+  auto stream = res->getDefaultStreamCurrentDevice();
 
   auto tVectors =
-    toDevice<T, 2>(resources,
-                   device,
-                   const_cast<T*>(reinterpret_cast<const T*>(args.vectors)),
-                   stream,
-                   {args.vectorsRowMajor ? args.numVectors : args.dims,
-                    args.vectorsRowMajor ? args.dims : args.numVectors});
+    toDeviceTemporary<T, 2>(
+      res,
+      device,
+      const_cast<T*>(reinterpret_cast<const T*>(args.vectors)),
+      stream,
+      {args.vectorsRowMajor ? args.numVectors : args.dims,
+       args.vectorsRowMajor ? args.dims : args.numVectors});
   auto tQueries =
-    toDevice<T, 2>(resources,
-                   device,
-                   const_cast<T*>(reinterpret_cast<const T*>(args.queries)),
-                   stream,
-                   {args.queriesRowMajor ? args.numQueries : args.dims,
-                    args.queriesRowMajor ? args.dims : args.numQueries});
+    toDeviceTemporary<T, 2>(
+      res,
+      device,
+      const_cast<T*>(reinterpret_cast<const T*>(args.queries)),
+      stream,
+      {args.queriesRowMajor ? args.numQueries : args.dims,
+       args.queriesRowMajor ? args.dims : args.numQueries});
 
   DeviceTensor<float, 1, true> tVectorNorms;
   if (args.vectorNorms) {
-    tVectorNorms = toDevice<float, 1>(resources,
-                                      device,
-                                      const_cast<float*>(args.vectorNorms),
-                                      stream,
-                                      {args.numVectors});
+    tVectorNorms =
+      toDeviceTemporary<float, 1>(res,
+                                  device,
+                                  const_cast<float*>(args.vectorNorms),
+                                  stream,
+                                  {args.numVectors});
   }
 
   auto tOutDistances =
-    toDevice<float, 2>(resources,
-                       device,
-                       args.outDistances,
-                       stream,
-                       {args.numQueries, args.k});
+    toDeviceTemporary<float, 2>(res,
+                                device,
+                                args.outDistances,
+                                stream,
+                                {args.numQueries, args.k});
 
   // The brute-force API only supports an interface for integer indices
   DeviceTensor<int, 2, true>
-    tOutIntIndices(mem, {args.numQueries, args.k}, stream);
+    tOutIntIndices(res,
+                   makeTempAlloc(AllocType::Other, stream),
+                   {args.numQueries, args.k});
 
   // Since we've guaranteed that all arguments are on device, call the
   // implementation
-  bfKnnOnDevice<T>(resources,
+  bfKnnOnDevice<T>(res,
                    device,
                    stream,
                    tVectors,
@@ -77,11 +84,11 @@ void bfKnnConvert(GpuResources* resources, const GpuDistanceParams& args) {
 
   // Convert and copy int indices out
   auto tOutIndices =
-    toDevice<faiss::Index::idx_t, 2>(resources,
-                                     device,
-                                     args.outIndices,
-                                     stream,
-                                     {args.numQueries, args.k});
+    toDeviceTemporary<faiss::Index::idx_t, 2>(res,
+                                              device,
+                                              args.outIndices,
+                                              stream,
+                                              {args.numQueries, args.k});
 
   // Convert int to idx_t
   convertTensor<int, faiss::Index::idx_t, 2>(stream,
@@ -94,7 +101,7 @@ void bfKnnConvert(GpuResources* resources, const GpuDistanceParams& args) {
 }
 
 void
-bfKnn(GpuResources* resources, const GpuDistanceParams& args) {
+bfKnn(GpuResourcesProvider* res, const GpuDistanceParams& args) {
   // For now, both vectors and queries must be of the same data type
   FAISS_THROW_IF_NOT_MSG(
     args.vectorType == args.queryType,
@@ -102,9 +109,9 @@ bfKnn(GpuResources* resources, const GpuDistanceParams& args) {
     "be the same (F32 or F16");
 
   if (args.vectorType == DistanceDataType::F32) {
-    bfKnnConvert<float>(resources, args);
+    bfKnnConvert<float>(res, args);
   } else if (args.vectorType == DistanceDataType::F16) {
-    bfKnnConvert<half>(resources, args);
+    bfKnnConvert<half>(res, args);
   } else {
     FAISS_THROW_MSG("unknown vectorType");
   }
@@ -112,7 +119,7 @@ bfKnn(GpuResources* resources, const GpuDistanceParams& args) {
 
 // legacy version
 void
-bruteForceKnn(GpuResources* resources,
+bruteForceKnn(GpuResourcesProvider* res,
               faiss::MetricType metric,
               // A region of memory size numVectors x dims, with dims
               // innermost
@@ -147,7 +154,7 @@ bruteForceKnn(GpuResources* resources,
   args.outDistances = outDistances;
   args.outIndices = outIndices;
 
-  bfKnn(resources, args);
+  bfKnn(res, args);
 }
 
 } } // namespace

gpu/GpuDistance.h

@@ -12,7 +12,7 @@
 
 namespace faiss { namespace gpu {
 
-class GpuResources;
+class GpuResourcesProvider;
 
 // Scalar type of the vector data
 enum class DistanceDataType {
@@ -116,10 +116,10 @@ struct GpuDistanceParams {
 ///
 /// For each vector in `queries`, searches all of `vectors` to find its k
 /// nearest neighbors with respect to the given metric
-void bfKnn(GpuResources* resources, const GpuDistanceParams& args);
+void bfKnn(GpuResourcesProvider* resources, const GpuDistanceParams& args);
 
 /// Deprecated legacy implementation
-void bruteForceKnn(GpuResources* resources,
+void bruteForceKnn(GpuResourcesProvider* resources,
                    faiss::MetricType metric,
                    // If vectorsRowMajor is true, this is
                    // numVectors x dims, with dims innermost; otherwise,

gpu/GpuIndex.cu

@@ -35,7 +35,7 @@ constexpr size_t kAddVecSize = (size_t) 512 * 1024;
 // FIXME: parameterize based on algorithm need
 constexpr size_t kSearchVecSize = (size_t) 32 * 1024;
 
-GpuIndex::GpuIndex(GpuResources* resources,
+GpuIndex::GpuIndex(std::shared_ptr<GpuResources> resources,
                    int dims,
                    faiss::MetricType metric,
                    float metricArg,
@@ -50,21 +50,16 @@ GpuIndex::GpuIndex(GpuResources* resources,
 
   FAISS_THROW_IF_NOT_MSG(dims > 0, "Invalid number of dimensions");
 
-#ifdef FAISS_UNIFIED_MEM
   FAISS_THROW_IF_NOT_FMT(
     memorySpace_ == MemorySpace::Device ||
     (memorySpace_ == MemorySpace::Unified &&
      getFullUnifiedMemSupport(device_)),
     "Device %d does not support full CUDA 8 Unified Memory (CC 6.0+)",
     config.device);
-#else
-  FAISS_THROW_IF_NOT_MSG(memorySpace_ == MemorySpace::Device,
-                     "Must compile with CUDA 8+ for Unified Memory support");
-#endif
 
   metric_arg = metricArg;
 
-  FAISS_ASSERT(resources_);
+  FAISS_ASSERT((bool) resources_);
   resources_->initializeForDevice(device_);
 }
 
@@ -174,18 +169,20 @@ GpuIndex::addPage_(int n,
   // Before continuing, we guarantee that all data will be resident on the GPU.
   auto stream = resources_->getDefaultStreamCurrentDevice();
 
-  auto vecs = toDevice<float, 2>(resources_,
-                                 device_,
-                                 const_cast<float*>(x),
-                                 stream,
-                                 {n, this->d});
+  auto vecs =
+    toDeviceTemporary<float, 2>(resources_.get(),
+                                device_,
+                                const_cast<float*>(x),
+                                stream,
+                                {n, this->d});
 
   if (ids) {
-    auto indices = toDevice<Index::idx_t, 1>(resources_,
-                                             device_,
-                                             const_cast<Index::idx_t*>(ids),
-                                             stream,
-                                             {n});
+    auto indices =
+      toDeviceTemporary<Index::idx_t, 1>(resources_.get(),
+                                         device_,
+                                         const_cast<Index::idx_t*>(ids),
+                                         stream,
+                                         {n});
 
     addImpl_(n, vecs.data(), ids ? indices.data() : nullptr);
   } else {
@@ -230,12 +227,14 @@ GpuIndex::search(Index::idx_t n,
   // If we reach a point where all inputs are too big, we can add
   // another level of tiling.
   auto outDistances =
-    toDevice<float, 2>(resources_, device_, distances, stream,
-                       {(int) n, (int) k});
+    toDeviceTemporary<float, 2>(
+      resources_.get(), device_, distances, stream,
+      {(int) n, (int) k});
 
   auto outLabels =
-    toDevice<faiss::Index::idx_t, 2>(resources_, device_, labels, stream,
-                                     {(int) n, (int) k});
+    toDeviceTemporary<faiss::Index::idx_t, 2>(
+      resources_.get(), device_, labels, stream,
+      {(int) n, (int) k});
 
   bool usePaged = false;
 
@@ -277,11 +276,11 @@ GpuIndex::searchNonPaged_(int n,
 
   // Make sure arguments are on the device we desire; use temporary
   // memory allocations to move it if necessary
-  auto vecs = toDevice<float, 2>(resources_,
-                                 device_,
-                                 const_cast<float*>(x),
-                                 stream,
-                                 {n, (int) this->d});
+  auto vecs = toDeviceTemporary<float, 2>(resources_.get(),
+                                          device_,
+                                          const_cast<float*>(x),
+                                          stream,
+                                          {n, (int) this->d});
 
   searchImpl_(n, vecs.data(), k, outDistancesData, outIndicesData);
 }
@@ -349,13 +348,13 @@ GpuIndex::searchFromCpuPaged_(int n,
   // Reserve space on the GPU for the destination of the pinned buffer
   // copy
   DeviceTensor<float, 2, true> bufGpuA(
-    resources_->getMemoryManagerCurrentDevice(),
-    {(int) pageSizeInVecs, (int) this->d},
-    defaultStream);
+    resources_.get(),
+    makeTempAlloc(AllocType::Other, defaultStream),
+    {(int) pageSizeInVecs, (int) this->d});
   DeviceTensor<float, 2, true> bufGpuB(
-    resources_->getMemoryManagerCurrentDevice(),
-    {(int) pageSizeInVecs, (int) this->d},
-    defaultStream);
+    resources_.get(),
+    makeTempAlloc(AllocType::Other, defaultStream),
+    {(int) pageSizeInVecs, (int) this->d});
   DeviceTensor<float, 2, true>* bufGpus[2] = {&bufGpuA, &bufGpuB};
 
   // Copy completion events for the pinned buffers
@@ -398,8 +397,7 @@ GpuIndex::searchFromCpuPaged_(int n,
                                   copyStream));
 
       // Mark a completion event in this stream
-      eventPinnedCopyDone[cur2BufIndex] =
-        std::move(std::unique_ptr<CudaEvent>(new CudaEvent(copyStream)));
+      eventPinnedCopyDone[cur2BufIndex].reset(new CudaEvent(copyStream));
 
       // We pick up from here
       cur3 = cur2;
@@ -430,8 +428,7 @@ GpuIndex::searchFromCpuPaged_(int n,
                   outIndicesSlice.data());
 
       // Create completion event
-      eventGpuExecuteDone[cur3BufIndex] =
-        std::move(std::unique_ptr<CudaEvent>(new CudaEvent(defaultStream)));
+      eventGpuExecuteDone[cur3BufIndex].reset(new CudaEvent(defaultStream));
 
       // We pick up from here
       cur3BufIndex = (cur3BufIndex == 0) ? 1 : 0;

gpu/GpuIndex.h

@@ -9,12 +9,10 @@
 #pragma once
 
 #include <faiss/Index.h>
-#include <faiss/gpu/utils/MemorySpace.h>
+#include <faiss/gpu/GpuResources.h>
 
 namespace faiss { namespace gpu {
 
-class GpuResources;
-
 struct GpuIndexConfig {
   inline GpuIndexConfig()
       : device(0),
@@ -32,7 +30,7 @@ struct GpuIndexConfig {
 
 class GpuIndex : public faiss::Index {
  public:
-  GpuIndex(GpuResources* resources,
+  GpuIndex(std::shared_ptr<GpuResources> resources,
            int dims,
            faiss::MetricType metric,
            float metricArg,
@@ -42,10 +40,6 @@ class GpuIndex : public faiss::Index {
     return device_;
   }
 
-  inline GpuResources* getResources() {
-    return resources_;
-  }
-
   /// Set the minimum data size for searches (in MiB) for which we use
   /// CPU -> GPU paging
   void setMinPagingSize(size_t size);
@@ -140,7 +134,7 @@ private:
 
  protected:
   /// Manages streams, cuBLAS handles and scratch memory for devices
-  GpuResources* resources_;
+  std::shared_ptr<GpuResources> resources_;
 
   /// The GPU device we are resident on
   const int device_;

gpu/GpuIndexBinaryFlat.cu

@@ -18,13 +18,12 @@ namespace faiss { namespace gpu {
 /// Default CPU search size for which we use paged copies
 constexpr size_t kMinPageSize = (size_t) 256 * 1024 * 1024;
 
-GpuIndexBinaryFlat::GpuIndexBinaryFlat(GpuResources* resources,
+GpuIndexBinaryFlat::GpuIndexBinaryFlat(GpuResourcesProvider* provider,
                                        const faiss::IndexBinaryFlat* index,
                                        GpuIndexBinaryFlatConfig config)
     : IndexBinary(index->d),
-      resources_(resources),
-      config_(std::move(config)),
-      data_(nullptr) {
+      resources_(provider->getResources()),
+      config_(std::move(config)) {
   FAISS_THROW_IF_NOT_FMT(this->d % 8 == 0,
                          "vector dimension (number of bits) "
                          "must be divisible by 8 (passed %d)",
@@ -37,13 +36,12 @@ GpuIndexBinaryFlat::GpuIndexBinaryFlat(GpuResources* resources,
 }
 
 
-GpuIndexBinaryFlat::GpuIndexBinaryFlat(GpuResources* resources,
+GpuIndexBinaryFlat::GpuIndexBinaryFlat(GpuResourcesProvider* provider,
                                        int dims,
                                        GpuIndexBinaryFlatConfig config)
     : IndexBinary(dims),
-      resources_(resources),
-      config_(std::move(config)),
-      data_(nullptr) {
+      resources_(provider->getResources()),
+      config_(std::move(config)) {
   FAISS_THROW_IF_NOT_FMT(this->d % 8 == 0,
                          "vector dimension (number of bits) "
                          "must be divisible by 8 (passed %d)",
@@ -54,13 +52,11 @@ GpuIndexBinaryFlat::GpuIndexBinaryFlat(GpuResources* resources,
 
   // Construct index
   DeviceScope scope(config_.device);
-  data_ = new BinaryFlatIndex(resources,
-                              this->d,
-                              config_.memorySpace);
+  data_.reset(
+    new BinaryFlatIndex(resources_.get(), this->d, config_.memorySpace));
 }
 
 GpuIndexBinaryFlat::~GpuIndexBinaryFlat() {
-  delete data_;
 }
 
 void
@@ -78,10 +74,10 @@ GpuIndexBinaryFlat::copyFrom(const faiss::IndexBinaryFlat* index) {
                          (size_t) index->ntotal);
   this->ntotal = index->ntotal;
 
-  delete data_;
-  data_ = new BinaryFlatIndex(resources_,
-                              this->d,
-                              config_.memorySpace);
+  // destroy old first before allocating new
+  data_.reset();
+  data_.reset(
+    new BinaryFlatIndex(resources_.get(), this->d, config_.memorySpace));
 
   // The index could be empty
   if (index->ntotal > 0) {
@@ -168,16 +164,16 @@ GpuIndexBinaryFlat::search(faiss::IndexBinary::idx_t n,
   // GPU.
   // If we reach a point where all inputs are too big, we can add
   // another level of tiling.
-  auto outDistances = toDevice<int32_t, 2>(resources_,
-                                           config_.device,
-                                           distances,
-                                           stream,
-                                           {(int) n, (int) k});
+  auto outDistances = toDeviceTemporary<int32_t, 2>(resources_.get(),
+                                                    config_.device,
+                                                    distances,
+                                                    stream,
+                                                    {(int) n, (int) k});
 
   // FlatIndex only supports an interface returning int indices
   DeviceTensor<int, 2, true> outIntIndices(
-    resources_->getMemoryManagerCurrentDevice(),
-    {(int) n, (int) k}, stream);
+    resources_.get(), makeTempAlloc(AllocType::Other, stream),
+    {(int) n, (int) k});
 
   bool usePaged = false;
 
@@ -205,11 +201,12 @@ GpuIndexBinaryFlat::search(faiss::IndexBinary::idx_t n,
   }
 
   // Convert and copy int indices out
-  auto outIndices = toDevice<faiss::Index::idx_t, 2>(resources_,
-                                                     config_.device,
-                                                     labels,
-                                                     stream,
-                                                     {(int) n, (int) k});
+  auto outIndices =
+    toDeviceTemporary<faiss::Index::idx_t, 2>(resources_.get(),
+                                              config_.device,
+                                              labels,
+                                              stream,
+                                              {(int) n, (int) k});
 
   // Convert int to long
   convertTensor<int, faiss::Index::idx_t, 2>(stream,
@@ -234,11 +231,11 @@ GpuIndexBinaryFlat::searchNonPaged_(int n,
 
   // Make sure arguments are on the device we desire; use temporary
   // memory allocations to move it if necessary
-  auto vecs = toDevice<uint8_t, 2>(resources_,
-                                   config_.device,
-                                   const_cast<uint8_t*>(x),
-                                   stream,
-                                   {n, (int) (this->d / 8)});
+  auto vecs = toDeviceTemporary<uint8_t, 2>(resources_.get(),
+                                            config_.device,
+                                            const_cast<uint8_t*>(x),
+                                            stream,
+                                            {n, (int) (this->d / 8)});
 
   data_->query(vecs, k, outDistances, outIndices);
 }

gpu/GpuIndexBinaryFlat.h

@@ -9,11 +9,12 @@
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/gpu/GpuIndex.h>
+#include <faiss/gpu/GpuResources.h>
+#include <memory>
 
 namespace faiss { namespace gpu {
 
 class BinaryFlatIndex;
-class GpuResources;
 
 struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {
 };
@@ -24,13 +25,13 @@ class GpuIndexBinaryFlat : public IndexBinary {
  public:
   /// Construct from a pre-existing faiss::IndexBinaryFlat instance, copying
   /// data over to the given GPU
-  GpuIndexBinaryFlat(GpuResources* resources,
+  GpuIndexBinaryFlat(GpuResourcesProvider* resources,
                      const faiss::IndexBinaryFlat* index,
                      GpuIndexBinaryFlatConfig config =
                      GpuIndexBinaryFlatConfig());
 
   /// Construct an empty instance that can be added to
-  GpuIndexBinaryFlat(GpuResources* resources,
+  GpuIndexBinaryFlat(GpuResourcesProvider* resources,
                      int dims,
                      GpuIndexBinaryFlatConfig config =
                      GpuIndexBinaryFlatConfig());
@@ -76,14 +77,13 @@ class GpuIndexBinaryFlat : public IndexBinary {
 
  protected:
   /// Manages streans, cuBLAS handles and scratch memory for devices
-  GpuResources* resources_;
+  std::shared_ptr<GpuResources> resources_;
 
   /// Configuration options
   GpuIndexBinaryFlatConfig config_;
 
-  /// Holds our GPU data containing the list of vectors; is managed via raw
-  /// pointer so as to allow non-CUDA compilers to see this header
-  BinaryFlatIndex* data_;
+  /// Holds our GPU data containing the list of vectors
+  std::unique_ptr<BinaryFlatIndex> data_;
 };
 
 } } // namespace gpu

gpu/GpuIndexFlat.cu

@@ -19,7 +19,22 @@
 
 namespace faiss { namespace gpu {
 
-GpuIndexFlat::GpuIndexFlat(GpuResources* resources,
+GpuIndexFlat::GpuIndexFlat(GpuResourcesProvider* provider,
+                           const faiss::IndexFlat* index,
+                           GpuIndexFlatConfig config) :
+    GpuIndex(provider->getResources(),
+             index->d,
+             index->metric_type,
+             index->metric_arg,
+             config),
+    config_(std::move(config)) {
+  // Flat index doesn't need training
+  this->is_trained = true;
+
+  copyFrom(index);
+}
+
+GpuIndexFlat::GpuIndexFlat(std::shared_ptr<GpuResources> resources,
                            const faiss::IndexFlat* index,
                            GpuIndexFlatConfig config) :
     GpuIndex(resources,
@@ -27,35 +42,54 @@ GpuIndexFlat::GpuIndexFlat(GpuResources* resources,
              index->metric_type,
              index->metric_arg,
              config),
-    config_(std::move(config)),
-    data_(nullptr) {
+    config_(std::move(config)) {
   // Flat index doesn't need training
   this->is_trained = true;
 
   copyFrom(index);
 }
 
-GpuIndexFlat::GpuIndexFlat(GpuResources* resources,
+GpuIndexFlat::GpuIndexFlat(GpuResourcesProvider* provider,
                            int dims,
                            faiss::MetricType metric,
                            GpuIndexFlatConfig config) :
-    GpuIndex(resources, dims, metric, 0, config),
-    config_(std::move(config)),
-    data_(nullptr) {
+    GpuIndex(provider->getResources(),
+             dims,
+             metric,
+             0,
+             config),
+    config_(std::move(config)) {
   // Flat index doesn't need training
   this->is_trained = true;
 
   // Construct index
   DeviceScope scope(device_);
-  data_ = new FlatIndex(resources,
-                        dims,
-                        config_.useFloat16,
-                        config_.storeTransposed,
-                        memorySpace_);
+  data_.reset(new FlatIndex(resources_.get(),
+                            dims,
+                            config_.useFloat16,
+                            config_.storeTransposed,
+                            memorySpace_));
+}
+
+GpuIndexFlat::GpuIndexFlat(std::shared_ptr<GpuResources> resources,
+                           int dims,
+                           faiss::MetricType metric,
+                           GpuIndexFlatConfig config) :
+    GpuIndex(resources, dims, metric, 0, config),
+    config_(std::move(config)) {
+  // Flat index doesn't need training
+  this->is_trained = true;
+
+  // Construct index
+  DeviceScope scope(device_);
+  data_.reset(new FlatIndex(resources_.get(),
+                            dims,
+                            config_.useFloat16,
+                            config_.storeTransposed,
+                            memorySpace_));
 }
 
 GpuIndexFlat::~GpuIndexFlat() {
-  delete data_;
 }
 
 void
@@ -72,12 +106,12 @@ GpuIndexFlat::copyFrom(const faiss::IndexFlat* index) {
                          (size_t) std::numeric_limits<int>::max(),
                          (size_t) index->ntotal);
 
-  delete data_;
-  data_ = new FlatIndex(resources_,
-                        this->d,
-                        config_.useFloat16,
-                        config_.storeTransposed,
-                        memorySpace_);
+  data_.reset();
+  data_.reset(new FlatIndex(resources_.get(),
+                            this->d,
+                            config_.useFloat16,
+                            config_.storeTransposed,
+                            memorySpace_));
 
   // The index could be empty
   if (index->ntotal > 0) {
@@ -199,7 +233,9 @@ GpuIndexFlat::searchImpl_(int n,
 
   // FlatIndex only supports int indices
   DeviceTensor<int, 2, true> outIntLabels(
-    resources_->getMemoryManagerCurrentDevice(), {n, k}, stream);
+    resources_.get(),
+    makeTempAlloc(AllocType::Other, stream),
+    {n, k});
 
   data_->query(queries, k, metric_type, metric_arg,
                outDistances, outIntLabels, true);
@@ -270,21 +306,24 @@ GpuIndexFlat::compute_residual_n(faiss::Index::idx_t n,
   DeviceScope scope(device_);
 
   auto vecsDevice =
-    toDevice<float, 2>(resources_, device_,
-                       const_cast<float*>(xs), stream,
-                       {(int) n, (int) this->d});
+    toDeviceTemporary<float, 2>(resources_.get(), device_,
+                                const_cast<float*>(xs), stream,
+                                {(int) n, (int) this->d});
   auto idsDevice =
-    toDevice<faiss::Index::idx_t, 1>(resources_, device_,
-                                     const_cast<faiss::Index::idx_t*>(keys),
-                                     stream,
-                                     {(int) n});
+    toDeviceTemporary<faiss::Index::idx_t, 1>(
+      resources_.get(), device_,
+      const_cast<faiss::Index::idx_t*>(keys),
+      stream,
+      {(int) n});
+
   auto residualDevice =
-    toDevice<float, 2>(resources_, device_, residuals, stream,
-                       {(int) n, (int) this->d});
+    toDeviceTemporary<float, 2>(resources_.get(), device_, residuals, stream,
+                                {(int) n, (int) this->d});
 
   // Convert idx_t to int
   auto keysInt =
-    convertTensor<faiss::Index::idx_t, int, 1>(resources_, stream, idsDevice);
+    convertTensorTemporary<faiss::Index::idx_t, int, 1>(
+      resources_.get(), stream, idsDevice);
 
   FAISS_ASSERT(data_);
   data_->computeResidual(vecsDevice,
@@ -298,13 +337,25 @@ GpuIndexFlat::compute_residual_n(faiss::Index::idx_t n,
 // GpuIndexFlatL2
 //
 
-GpuIndexFlatL2::GpuIndexFlatL2(GpuResources* resources,
+GpuIndexFlatL2::GpuIndexFlatL2(GpuResourcesProvider* provider,
+                               faiss::IndexFlatL2* index,
+                               GpuIndexFlatConfig config) :
+    GpuIndexFlat(provider, index, config) {
+}
+
+GpuIndexFlatL2::GpuIndexFlatL2(std::shared_ptr<GpuResources> resources,
                                faiss::IndexFlatL2* index,
                                GpuIndexFlatConfig config) :
     GpuIndexFlat(resources, index, config) {
 }
 
-GpuIndexFlatL2::GpuIndexFlatL2(GpuResources* resources,
+GpuIndexFlatL2::GpuIndexFlatL2(GpuResourcesProvider* provider,
+                               int dims,
+                               GpuIndexFlatConfig config) :
+    GpuIndexFlat(provider, dims, faiss::METRIC_L2, config) {
+}
+
+GpuIndexFlatL2::GpuIndexFlatL2(std::shared_ptr<GpuResources> resources,
                                int dims,
                                GpuIndexFlatConfig config) :
     GpuIndexFlat(resources, dims, faiss::METRIC_L2, config) {
@@ -332,13 +383,25 @@ GpuIndexFlatL2::copyTo(faiss::IndexFlat* index) {
 // GpuIndexFlatIP
 //
 
-GpuIndexFlatIP::GpuIndexFlatIP(GpuResources* resources,
+GpuIndexFlatIP::GpuIndexFlatIP(GpuResourcesProvider* provider,
+                               faiss::IndexFlatIP* index,
+                               GpuIndexFlatConfig config) :
+    GpuIndexFlat(provider, index, config) {
+}
+
+GpuIndexFlatIP::GpuIndexFlatIP(std::shared_ptr<GpuResources> resources,
                                faiss::IndexFlatIP* index,
                                GpuIndexFlatConfig config) :
     GpuIndexFlat(resources, index, config) {
 }
 
-GpuIndexFlatIP::GpuIndexFlatIP(GpuResources* resources,
+GpuIndexFlatIP::GpuIndexFlatIP(GpuResourcesProvider* provider,
+                               int dims,
+                               GpuIndexFlatConfig config) :
+    GpuIndexFlat(provider, dims, faiss::METRIC_INNER_PRODUCT, config) {
+}
+
+GpuIndexFlatIP::GpuIndexFlatIP(std::shared_ptr<GpuResources> resources,
                                int dims,
                                GpuIndexFlatConfig config) :
     GpuIndexFlat(resources, dims, faiss::METRIC_INNER_PRODUCT, config) {

gpu/GpuIndexFlat.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <faiss/gpu/GpuIndex.h>
+#include <memory>
 
 namespace faiss {
 
@@ -47,12 +48,21 @@ class GpuIndexFlat : public GpuIndex {
  public:
   /// Construct from a pre-existing faiss::IndexFlat instance, copying
   /// data over to the given GPU
-  GpuIndexFlat(GpuResources* resources,
+  GpuIndexFlat(GpuResourcesProvider* provider,
+               const faiss::IndexFlat* index,
+               GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  GpuIndexFlat(std::shared_ptr<GpuResources> resources,
                const faiss::IndexFlat* index,
                GpuIndexFlatConfig config = GpuIndexFlatConfig());
 
   /// Construct an empty instance that can be added to
-  GpuIndexFlat(GpuResources* resources,
+  GpuIndexFlat(GpuResourcesProvider* provider,
+               int dims,
+               faiss::MetricType metric,
+               GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  GpuIndexFlat(std::shared_ptr<GpuResources> resources,
                int dims,
                faiss::MetricType metric,
                GpuIndexFlatConfig config = GpuIndexFlatConfig());
@@ -100,7 +110,7 @@ class GpuIndexFlat : public GpuIndex {
                           const faiss::Index::idx_t* keys) const override;
 
   /// For internal access
-  inline FlatIndex* getGpuData() { return data_; }
+  inline FlatIndex* getGpuData() { return data_.get(); }
 
  protected:
   /// Flat index does not require IDs as there is no storage available for them
@@ -122,9 +132,8 @@ class GpuIndexFlat : public GpuIndex {
   /// Our config object
   const GpuIndexFlatConfig config_;
 
-  /// Holds our GPU data containing the list of vectors; is managed via raw
-  /// pointer so as to allow non-CUDA compilers to see this header
-  FlatIndex* data_;
+  /// Holds our GPU data containing the list of vectors
+  std::unique_ptr<FlatIndex> data_;
 };
 
 /// Wrapper around the GPU implementation that looks like
@@ -134,12 +143,20 @@ class GpuIndexFlatL2 : public GpuIndexFlat {
  public:
   /// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
   /// data over to the given GPU
-  GpuIndexFlatL2(GpuResources* resources,
+  GpuIndexFlatL2(GpuResourcesProvider* provider,
+                 faiss::IndexFlatL2* index,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  GpuIndexFlatL2(std::shared_ptr<GpuResources> resources,
                  faiss::IndexFlatL2* index,
                  GpuIndexFlatConfig config = GpuIndexFlatConfig());
 
   /// Construct an empty instance that can be added to
-  GpuIndexFlatL2(GpuResources* resources,
+  GpuIndexFlatL2(GpuResourcesProvider* provider,
+                 int dims,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  GpuIndexFlatL2(std::shared_ptr<GpuResources> resources,
                  int dims,
                  GpuIndexFlatConfig config = GpuIndexFlatConfig());
 
@@ -159,12 +176,20 @@ class GpuIndexFlatIP : public GpuIndexFlat {
  public:
   /// Construct from a pre-existing faiss::IndexFlatIP instance, copying
   /// data over to the given GPU
-  GpuIndexFlatIP(GpuResources* resources,
+  GpuIndexFlatIP(GpuResourcesProvider* provider,
+                 faiss::IndexFlatIP* index,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  GpuIndexFlatIP(std::shared_ptr<GpuResources> resources,
                  faiss::IndexFlatIP* index,
                  GpuIndexFlatConfig config = GpuIndexFlatConfig());
 
   /// Construct an empty instance that can be added to
-  GpuIndexFlatIP(GpuResources* resources,
+  GpuIndexFlatIP(GpuResourcesProvider* provider,
+                 int dims,
+                 GpuIndexFlatConfig config = GpuIndexFlatConfig());
+
+  GpuIndexFlatIP(std::shared_ptr<GpuResources> resources,
                  int dims,
                  GpuIndexFlatConfig config = GpuIndexFlatConfig());
 

gpu/GpuIndexIVF.cu

@@ -16,17 +16,18 @@
 
 namespace faiss { namespace gpu {
 
-GpuIndexIVF::GpuIndexIVF(GpuResources* resources,
+GpuIndexIVF::GpuIndexIVF(GpuResourcesProvider* provider,
                          int dims,
                          faiss::MetricType metric,
                          float metricArg,
                          int nlistIn,
                          GpuIndexIVFConfig config) :
-    GpuIndex(resources, dims, metric, metricArg, config),
-    ivfConfig_(std::move(config)),
+    GpuIndex(provider->getResources(),
+             dims, metric, metricArg, config),
     nlist(nlistIn),
     nprobe(1),
-    quantizer(nullptr) {
+    quantizer(nullptr),
+    ivfConfig_(std::move(config)) {
   init_();
 
   // Only IP and L2 are supported for now

gpu/GpuIndexIVF.h

@@ -18,7 +18,6 @@ namespace faiss { struct IndexIVF; }
 namespace faiss { namespace gpu {
 
 class GpuIndexFlat;
-class GpuResources;
 
 struct GpuIndexIVFConfig : public GpuIndexConfig {
   inline GpuIndexIVFConfig()
@@ -34,7 +33,7 @@ struct GpuIndexIVFConfig : public GpuIndexConfig {
 
 class GpuIndexIVF : public GpuIndex {
  public:
-  GpuIndexIVF(GpuResources* resources,
+  GpuIndexIVF(GpuResourcesProvider* provider,
               int dims,
               faiss::MetricType metric,
               float metricArg,

gpu/GpuIndexIVFFlat.cu

@@ -20,30 +20,28 @@
 
 namespace faiss { namespace gpu {
 
-GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
+GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResourcesProvider* provider,
                                  const faiss::IndexIVFFlat* index,
                                  GpuIndexIVFFlatConfig config) :
-    GpuIndexIVF(resources,
+    GpuIndexIVF(provider,
                 index->d,
                 index->metric_type,
                 index->metric_arg,
                 index->nlist,
                 config),
     ivfFlatConfig_(config),
-    reserveMemoryVecs_(0),
-    index_(nullptr) {
+    reserveMemoryVecs_(0) {
   copyFrom(index);
 }
 
-GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
+GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResourcesProvider* provider,
                                  int dims,
                                  int nlist,
                                  faiss::MetricType metric,
                                  GpuIndexIVFFlatConfig config) :
-    GpuIndexIVF(resources, dims, metric, 0, nlist, config),
+    GpuIndexIVF(provider, dims, metric, 0, nlist, config),
     ivfFlatConfig_(config),
-    reserveMemoryVecs_(0),
-    index_(nullptr) {
+    reserveMemoryVecs_(0) {
 
   // faiss::Index params
   this->is_trained = false;
@@ -53,7 +51,6 @@ GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
 }
 
 GpuIndexIVFFlat::~GpuIndexIVFFlat() {
-  delete index_;
 }
 
 void
@@ -72,8 +69,7 @@ GpuIndexIVFFlat::copyFrom(const faiss::IndexIVFFlat* index) {
   GpuIndexIVF::copyFrom(index);
 
   // Clear out our old data
-  delete index_;
-  index_ = nullptr;
+  index_.reset();
 
   // The other index might not be trained
   if (!index->is_trained) {
@@ -85,14 +81,14 @@ GpuIndexIVFFlat::copyFrom(const faiss::IndexIVFFlat* index) {
   FAISS_ASSERT(is_trained);
 
   // Copy our lists as well
-  index_ = new IVFFlat(resources_,
-                       quantizer->getGpuData(),
-                       index->metric_type,
-                       index->metric_arg,
-                       false, // no residual
-                       nullptr, // no scalar quantizer
-                       ivfFlatConfig_.indicesOptions,
-                       memorySpace_);
+  index_.reset(new IVFFlat(resources_.get(),
+                           quantizer->getGpuData(),
+                           index->metric_type,
+                           index->metric_arg,
+                           false, // no residual
+                           nullptr, // no scalar quantizer
+                           ivfFlatConfig_.indicesOptions,
+                           memorySpace_));
   InvertedLists *ivf = index->invlists;
 
   for (size_t i = 0; i < ivf->nlist; ++i) {
@@ -181,14 +177,14 @@ GpuIndexIVFFlat::train(Index::idx_t n, const float* x) {
   trainQuantizer_(n, x);
 
   // The quantizer is now trained; construct the IVF index
-  index_ = new IVFFlat(resources_,
-                       quantizer->getGpuData(),
-                       this->metric_type,
-                       this->metric_arg,
-                       false, // no residual
-                       nullptr, // no scalar quantizer
-                       ivfFlatConfig_.indicesOptions,
-                       memorySpace_);
+  index_.reset(new IVFFlat(resources_.get(),
+                           quantizer->getGpuData(),
+                           this->metric_type,
+                           this->metric_arg,
+                           false, // no residual
+                           nullptr, // no scalar quantizer
+                           ivfFlatConfig_.indicesOptions,
+                           memorySpace_));
 
   if (reserveMemoryVecs_) {
     index_->reserveMemory(reserveMemoryVecs_);

gpu/GpuIndexIVFFlat.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <faiss/gpu/GpuIndexIVF.h>
+#include <memory>
 
 namespace faiss { struct IndexIVFFlat; }
 
@@ -26,13 +27,13 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
  public:
   /// Construct from a pre-existing faiss::IndexIVFFlat instance, copying
   /// data over to the given GPU, if the input index is trained.
-  GpuIndexIVFFlat(GpuResources* resources,
+  GpuIndexIVFFlat(GpuResourcesProvider* provider,
                   const faiss::IndexIVFFlat* index,
                   GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
 
   /// Constructs a new instance with an empty flat quantizer; the user
   /// provides the number of lists desired.
-  GpuIndexIVFFlat(GpuResources* resources,
+  GpuIndexIVFFlat(GpuResourcesProvider* provider,
                   int dims,
                   int nlist,
                   faiss::MetricType metric,
@@ -79,7 +80,7 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
   size_t reserveMemoryVecs_;
 
   /// Instance that we own; contains the inverted list
-  IVFFlat* index_;
+  std::unique_ptr<IVFFlat> index_;
 };
 
 } } // namespace

gpu/GpuIndexIVFPQ.cu

@@ -20,10 +20,10 @@
 
 namespace faiss { namespace gpu {
 
-GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResources* resources,
+GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResourcesProvider* provider,
                              const faiss::IndexIVFPQ* index,
                              GpuIndexIVFPQConfig config) :
-    GpuIndexIVF(resources,
+    GpuIndexIVF(provider,
                 index->d,
                 index->metric_type,
                 index->metric_arg,
@@ -32,19 +32,18 @@ GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResources* resources,
     ivfpqConfig_(config),
     subQuantizers_(0),
     bitsPerCode_(0),
-    reserveMemoryVecs_(0),
-    index_(nullptr) {
+    reserveMemoryVecs_(0) {
   copyFrom(index);
 }
 
-GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResources* resources,
+GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResourcesProvider* provider,
                              int dims,
                              int nlist,
                              int subQuantizers,
                              int bitsPerCode,
                              faiss::MetricType metric,
                              GpuIndexIVFPQConfig config) :
-    GpuIndexIVF(resources,
+    GpuIndexIVF(provider,
                 dims,
                 metric,
                 0,
@@ -53,8 +52,7 @@ GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResources* resources,
     ivfpqConfig_(config),
     subQuantizers_(subQuantizers),
     bitsPerCode_(bitsPerCode),
-    reserveMemoryVecs_(0),
-    index_(nullptr) {
+    reserveMemoryVecs_(0) {
   verifySettings_();
 
   // We haven't trained ourselves, so don't construct the PQ index yet
@@ -62,7 +60,6 @@ GpuIndexIVFPQ::GpuIndexIVFPQ(GpuResources* resources,
 }
 
 GpuIndexIVFPQ::~GpuIndexIVFPQ() {
-  delete index_;
 }
 
 void
@@ -72,8 +69,7 @@ GpuIndexIVFPQ::copyFrom(const faiss::IndexIVFPQ* index) {
   GpuIndexIVF::copyFrom(index);
 
   // Clear out our old data
-  delete index_;
-  index_ = nullptr;
+  index_.reset();
 
   subQuantizers_ = index->pq.M;
   bitsPerCode_ = index->pq.nbits;
@@ -98,16 +94,16 @@ GpuIndexIVFPQ::copyFrom(const faiss::IndexIVFPQ* index) {
   // Copy our lists as well
   // The product quantizer must have data in it
   FAISS_ASSERT(index->pq.centroids.size() > 0);
-  index_ = new IVFPQ(resources_,
-                     index->metric_type,
-                     index->metric_arg,
-                     quantizer->getGpuData(),
-                     subQuantizers_,
-                     bitsPerCode_,
-                     (float*) index->pq.centroids.data(),
-                     ivfpqConfig_.indicesOptions,
-                     ivfpqConfig_.useFloat16LookupTables,
-                     memorySpace_);
+  index_.reset(new IVFPQ(resources_.get(),
+                         index->metric_type,
+                         index->metric_arg,
+                         quantizer->getGpuData(),
+                         subQuantizers_,
+                         bitsPerCode_,
+                         (float*) index->pq.centroids.data(),
+                         ivfpqConfig_.indicesOptions,
+                         ivfpqConfig_.useFloat16LookupTables,
+                         memorySpace_));
   // Doesn't make sense to reserve memory here
   index_->setPrecomputedCodes(ivfpqConfig_.usePrecomputedTables);
 
@@ -126,7 +122,7 @@ GpuIndexIVFPQ::copyFrom(const faiss::IndexIVFPQ* index) {
                        list_size);
 
     index_->addCodeVectorsFromCpu(
-                       i, ivf->get_codes(i), ivf->get_ids(i), list_size);
+      i, ivf->get_codes(i), ivf->get_ids(i), list_size);
   }
 }
 
@@ -157,9 +153,7 @@ GpuIndexIVFPQ::copyTo(faiss::IndexIVFPQ* index) const {
   index->polysemous_ht = 0;
   index->precomputed_table.clear();
 
-  InvertedLists *ivf = new ArrayInvertedLists(
-      nlist, index->code_size);
-
+  InvertedLists *ivf = new ArrayInvertedLists(nlist, index->code_size);
   index->replace_invlists(ivf, true);
 
   if (index_) {
@@ -276,16 +270,16 @@ GpuIndexIVFPQ::trainResidualQuantizer_(Index::idx_t n, const float* x) {
   pq.verbose = this->verbose;
   pq.train(n, residuals.data());
 
-  index_ = new IVFPQ(resources_,
-                     metric_type,
-                     metric_arg,
-                     quantizer->getGpuData(),
-                     subQuantizers_,
-                     bitsPerCode_,
-                     pq.centroids.data(),
-                     ivfpqConfig_.indicesOptions,
-                     ivfpqConfig_.useFloat16LookupTables,
-                     memorySpace_);
+  index_.reset(new IVFPQ(resources_.get(),
+                         metric_type,
+                         metric_arg,
+                         quantizer->getGpuData(),
+                         subQuantizers_,
+                         bitsPerCode_,
+                         pq.centroids.data(),
+                         ivfpqConfig_.indicesOptions,
+                         ivfpqConfig_.useFloat16LookupTables,
+                         memorySpace_));
   if (reserveMemoryVecs_) {
     index_->reserveMemory(reserveMemoryVecs_);
   }

gpu/GpuIndexIVFPQ.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <faiss/gpu/GpuIndexIVF.h>
+#include <memory>
 #include <vector>
 
 namespace faiss { struct IndexIVFPQ; }
@@ -39,12 +40,12 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
  public:
   /// Construct from a pre-existing faiss::IndexIVFPQ instance, copying
   /// data over to the given GPU, if the input index is trained.
-  GpuIndexIVFPQ(GpuResources* resources,
+  GpuIndexIVFPQ(GpuResourcesProvider* provider,
                 const faiss::IndexIVFPQ* index,
                 GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
 
   /// Construct an empty index
-  GpuIndexIVFPQ(GpuResources* resources,
+  GpuIndexIVFPQ(GpuResourcesProvider* provider,
                 int dims,
                 int nlist,
                 int subQuantizers,
@@ -137,7 +138,7 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
 
   /// The product quantizer instance that we own; contains the
   /// inverted lists
-  IVFPQ* index_;
+  std::unique_ptr<IVFPQ> index_;
 };
 
 } } // namespace

gpu/GpuIndexIVFScalarQuantizer.cu

@@ -18,20 +18,19 @@
 namespace faiss { namespace gpu {
 
 GpuIndexIVFScalarQuantizer::GpuIndexIVFScalarQuantizer(
-  GpuResources* resources,
+  GpuResourcesProvider* provider,
   const faiss::IndexIVFScalarQuantizer* index,
   GpuIndexIVFScalarQuantizerConfig config) :
-    GpuIndexIVF(resources,
+    GpuIndexIVF(provider,
                 index->d,
                 index->metric_type,
                 index->metric_arg,
                 index->nlist,
                 config),
-    ivfSQConfig_(config),
     sq(index->sq),
     by_residual(index->by_residual),
-    reserveMemoryVecs_(0),
-    index_(nullptr) {
+    ivfSQConfig_(config),
+    reserveMemoryVecs_(0) {
   copyFrom(index);
 
   FAISS_THROW_IF_NOT_MSG(isSQSupported(sq.qtype),
@@ -39,20 +38,18 @@ GpuIndexIVFScalarQuantizer::GpuIndexIVFScalarQuantizer(
 }
 
 GpuIndexIVFScalarQuantizer::GpuIndexIVFScalarQuantizer(
-  GpuResources* resources,
+  GpuResourcesProvider* provider,
   int dims,
   int nlist,
   faiss::ScalarQuantizer::QuantizerType qtype,
   faiss::MetricType metric,
   bool encodeResidual,
   GpuIndexIVFScalarQuantizerConfig config) :
-    GpuIndexIVF(resources, dims, metric, 0, nlist, config),
-    ivfSQConfig_(config),
+    GpuIndexIVF(provider, dims, metric, 0, nlist, config),
     sq(dims, qtype),
     by_residual(encodeResidual),
-    reserveMemoryVecs_(0),
-    index_(nullptr) {
-
+    ivfSQConfig_(config),
+    reserveMemoryVecs_(0) {
   // faiss::Index params
   this->is_trained = false;
 
@@ -63,7 +60,6 @@ GpuIndexIVFScalarQuantizer::GpuIndexIVFScalarQuantizer(
 }
 
 GpuIndexIVFScalarQuantizer::~GpuIndexIVFScalarQuantizer() {
-  delete index_;
 }
 
 void
@@ -81,8 +77,7 @@ GpuIndexIVFScalarQuantizer::copyFrom(
   DeviceScope scope(device_);
 
   // Clear out our old data
-  delete index_;
-  index_ = nullptr;
+  index_.reset();
 
   // Copy what we need from the CPU index
   GpuIndexIVF::copyFrom(index);
@@ -99,14 +94,14 @@ GpuIndexIVFScalarQuantizer::copyFrom(
   this->is_trained = true;
 
   // Copy our lists as well
-  index_ = new IVFFlat(resources_,
-                       quantizer->getGpuData(),
-                       index->metric_type,
-                       index->metric_arg,
-                       by_residual,
-                       &sq,
-                       ivfSQConfig_.indicesOptions,
-                       memorySpace_);
+  index_.reset(new IVFFlat(resources_.get(),
+                           quantizer->getGpuData(),
+                           index->metric_type,
+                           index->metric_arg,
+                           by_residual,
+                           &sq,
+                           ivfSQConfig_.indicesOptions,
+                           memorySpace_));
 
   InvertedLists* ivf = index->invlists;
 
@@ -215,14 +210,14 @@ GpuIndexIVFScalarQuantizer::train(Index::idx_t n, const float* x) {
   trainResiduals_(n, hostData.data());
 
   // The quantizer is now trained; construct the IVF index
-  index_ = new IVFFlat(resources_,
-                       quantizer->getGpuData(),
-                       this->metric_type,
-                       this->metric_arg,
-                       by_residual,
-                       &sq,
-                       ivfSQConfig_.indicesOptions,
-                       memorySpace_);
+  index_.reset(new IVFFlat(resources_.get(),
+                           quantizer->getGpuData(),
+                           this->metric_type,
+                           this->metric_arg,
+                           by_residual,
+                           &sq,
+                           ivfSQConfig_.indicesOptions,
+                           memorySpace_));
 
   if (reserveMemoryVecs_) {
     index_->reserveMemory(reserveMemoryVecs_);

gpu/GpuIndexIVFScalarQuantizer.h

@@ -10,6 +10,7 @@
 
 #include <faiss/gpu/GpuIndexIVF.h>
 #include <faiss/IndexScalarQuantizer.h>
+#include <memory>
 
 namespace faiss { namespace gpu {
 
@@ -26,7 +27,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
   /// Construct from a pre-existing faiss::IndexIVFScalarQuantizer instance,
   /// copying data over to the given GPU, if the input index is trained.
   GpuIndexIVFScalarQuantizer(
-    GpuResources* resources,
+    GpuResourcesProvider* provider,
     const faiss::IndexIVFScalarQuantizer* index,
     GpuIndexIVFScalarQuantizerConfig config =
     GpuIndexIVFScalarQuantizerConfig());
@@ -34,7 +35,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
   /// Constructs a new instance with an empty flat quantizer; the user
   /// provides the number of lists desired.
   GpuIndexIVFScalarQuantizer(
-    GpuResources* resources,
+    GpuResourcesProvider* provider,
     int dims,
     int nlist,
     faiss::ScalarQuantizer::QuantizerType qtype,
@@ -94,7 +95,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
   size_t reserveMemoryVecs_;
 
   /// Instance that we own; contains the inverted list
-  IVFFlat* index_;
+  std::unique_ptr<IVFFlat> index_;
 };
 
 } } // namespace

gpu/GpuResources.cpp

@@ -11,6 +11,87 @@
 
 namespace faiss { namespace gpu {
 
+AllocInfo makeDevAlloc(AllocType at, cudaStream_t st) {
+  return AllocInfo(at, getCurrentDevice(), MemorySpace::Device, st);
+}
+
+AllocInfo makeTempAlloc(AllocType at, cudaStream_t st) {
+  return AllocInfo(at, getCurrentDevice(), MemorySpace::Temporary, st);
+}
+
+AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st) {
+  return AllocInfo(at, getCurrentDevice(), sp, st);
+}
+
+//
+// GpuMemoryReservation
+//
+
+GpuMemoryReservation::GpuMemoryReservation()
+    : res(nullptr),
+      device(0),
+      stream(nullptr),
+      data(nullptr),
+      size(0) {
+}
+
+GpuMemoryReservation::GpuMemoryReservation(GpuResources* r,
+                                           int dev,
+                                           cudaStream_t str,
+                                           void* p,
+                                           size_t sz)
+    : res(r),
+      device(dev),
+      stream(str),
+      data(p),
+      size(sz) {
+}
+
+GpuMemoryReservation::GpuMemoryReservation(GpuMemoryReservation&& m) noexcept {
+  res = m.res; m.res = nullptr;
+  device = m.device; m.device = 0;
+  stream = m.stream; m.stream = nullptr;
+  data = m.data; m.data = nullptr;
+  size = m.size; m.size = 0;
+}
+
+GpuMemoryReservation&
+GpuMemoryReservation::operator=(GpuMemoryReservation&& m) {
+  // Can't be both a valid allocation and the same allocation
+  FAISS_ASSERT(!(res && res == m.res && device == m.device && data == m.data));
+
+  release();
+  res = m.res; m.res = nullptr;
+  device = m.device; m.device = 0;
+  stream = m.stream; m.stream = nullptr;
+  data = m.data; m.data = nullptr;
+  size = m.size; m.size = 0;
+
+  return *this;
+}
+
+void
+GpuMemoryReservation::release() {
+  if (res) {
+    res->deallocMemory(device, data);
+    res = nullptr;
+    device = 0;
+    stream = nullptr;
+    data = nullptr;
+    size = 0;
+  }
+}
+
+GpuMemoryReservation::~GpuMemoryReservation() {
+  if (res) {
+    res->deallocMemory(device, data);
+  }
+}
+
+//
+// GpuResources
+//
+
 GpuResources::~GpuResources() {
 }
 
@@ -29,11 +110,6 @@ GpuResources::getAlternateStreamsCurrentDevice() {
   return getAlternateStreams(getCurrentDevice());
 }
 
-DeviceMemory&
-GpuResources::getMemoryManagerCurrentDevice() {
-  return getMemoryManager(getCurrentDevice());
-}
-
 cudaStream_t
 GpuResources::getAsyncCopyStreamCurrentDevice() {
   return getAsyncCopyStream(getCurrentDevice());
@@ -49,4 +125,22 @@ GpuResources::syncDefaultStreamCurrentDevice() {
   syncDefaultStream(getCurrentDevice());
 }
 
+GpuMemoryReservation
+GpuResources::allocMemoryHandle(const AllocRequest& req) {
+  return GpuMemoryReservation(
+    this, req.device, req.stream, allocMemory(req), req.size);
+}
+
+size_t
+GpuResources::getTempMemoryAvailableCurrentDevice() const {
+  return getTempMemoryAvailable(getCurrentDevice());
+}
+
+//
+// GpuResourcesProvider
+//
+
+GpuResourcesProvider::~GpuResourcesProvider() {
+}
+
 } } // namespace

gpu/GpuResources.h

@@ -8,16 +8,171 @@
 
 #pragma once
 
-#include <faiss/gpu/utils/DeviceMemory.h>
+#include <faiss/impl/FaissAssert.h>
 #include <cuda_runtime.h>
 #include <cublas_v2.h>
+#include <memory>
 #include <utility>
 #include <vector>
 
 namespace faiss { namespace gpu {
 
+class GpuResources;
+
+enum AllocType {
+  /// Unknown allocation type or miscellaneous (not currently categorized)
+  Other = 0,
+
+  /// Primary data storage for GpuIndexFlat (the raw matrix of vectors and
+  /// vector norms if needed)
+  FlatData = 1,
+
+  /// Primary data storage for GpuIndexIVF* (the storage for each individual IVF
+  /// list)
+  IVFLists = 2,
+
+  /// Quantizer (PQ, SQ) dictionary information
+  Quantizer = 3,
+
+  /// For GpuIndexIVFPQ, "precomputed codes" for more efficient PQ lookup
+  /// require the use of possibly large tables. These are marked separately from
+  /// Quantizer as these can frequently be 100s - 1000s of MiB in size
+  QuantizerPrecomputedCodes = 4,
+
+  ///
+  /// StandardGpuResources implementation specific types
+  ///
+
+  /// When using StandardGpuResources, temporary memory allocations
+  /// (MemorySpace::Temporary) come out of a stack region of memory that is
+  /// allocated up front for each gpu (e.g., 1.5 GiB upon initialization). This
+  /// allocation by StandardGpuResources is marked with this AllocType.
+  TemporaryMemoryBuffer = 10,
+
+  /// When using StandardGpuResources, any MemorySpace::Temporary allocations
+  /// that cannot be satisfied within the TemporaryMemoryBuffer region fall back
+  /// to calling cudaMalloc which are sized to just the request at hand. These
+  /// "overflow" temporary allocations are marked with this AllocType.
+  TemporaryMemoryOverflow = 11,
+};
+
+/// Convert an AllocType to string
+std::string allocTypeToString(AllocType t);
+
+/// Memory regions accessible to the GPU
+enum MemorySpace {
+  /// Temporary device memory (guaranteed to no longer be used upon exit of a
+  /// top-level index call, and where the streams using it have completed GPU
+  /// work). Typically backed by Device memory (cudaMalloc/cudaFree).
+  Temporary = 0,
+
+  /// Managed using cudaMalloc/cudaFree (typical GPU device memory)
+  Device = 1,
+
+  /// Managed using cudaMallocManaged/cudaFree (typical Unified CPU/GPU memory)
+  Unified = 2,
+};
+
+/// Information on what/where an allocation is
+struct AllocInfo {
+  inline AllocInfo()
+      : type(AllocType::Other),
+        device(0),
+        space(MemorySpace::Device),
+        stream(nullptr) {
+  }
+
+  inline AllocInfo(AllocType at,
+                   int dev,
+                   MemorySpace sp,
+                   cudaStream_t st)
+      : type(at),
+        device(dev),
+        space(sp),
+        stream(st) {
+  }
+
+  /// The internal category of the allocation
+  AllocType type;
+
+  /// The device on which the allocation is happening
+  int device;
+
+  /// The memory space of the allocation
+  MemorySpace space;
+
+  /// The stream on which new work on the memory will be ordered (e.g., if a
+  /// piece of memory cached and to be returned for this call was last used on
+  /// stream 3 and a new memory request is for stream 4, the memory manager will
+  /// synchronize stream 4 to wait for the completion of stream 3 via events or
+  /// other stream synchronization.
+  ///
+  /// The memory manager guarantees that the returned memory is free to use
+  /// without data races on this stream specified.
+  cudaStream_t stream;
+};
+
+/// Create an AllocInfo for the current device with MemorySpace::Device
+AllocInfo makeDevAlloc(AllocType at, cudaStream_t st);
+
+/// Create an AllocInfo for the current device with MemorySpace::Temporary
+AllocInfo makeTempAlloc(AllocType at, cudaStream_t st);
+
+/// Create an AllocInfo for the current device
+AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st);
+
+/// Information on what/where an allocation is, along with how big it should be
+struct AllocRequest : public AllocInfo {
+  inline AllocRequest()
+      : AllocInfo(),
+        size(0) {
+  }
+
+  inline AllocRequest(const AllocInfo& info,
+                      size_t sz)
+      : AllocInfo(info),
+        size(sz) {
+  }
+
+  inline AllocRequest(AllocType at,
+                      int dev,
+                      MemorySpace sp,
+                      cudaStream_t st,
+                      size_t sz)
+      : AllocInfo(at, dev, sp, st),
+        size(sz) {
+  }
+
+  /// The size in bytes of the allocation
+  size_t size;
+};
+
+/// A RAII object that manages a temporary memory request
+struct GpuMemoryReservation {
+  GpuMemoryReservation();
+  GpuMemoryReservation(GpuResources* r,
+                       int dev,
+                       cudaStream_t str,
+                       void* p,
+                       size_t sz);
+  GpuMemoryReservation(GpuMemoryReservation&& m) noexcept;
+  ~GpuMemoryReservation();
+
+  GpuMemoryReservation& operator=(GpuMemoryReservation&& m);
+
+  inline void* get() { return data; }
+
+  void release();
+
+  GpuResources* res;
+  int device;
+  cudaStream_t stream;
+  void* data;
+  size_t size;
+};
+
 /// Base class of GPU-side resource provider; hides provision of
-/// cuBLAS handles, CUDA streams and a temporary memory manager
+/// cuBLAS handles, CUDA streams and all device memory allocation performed
 class GpuResources {
  public:
   virtual ~GpuResources();
@@ -37,8 +192,19 @@ class GpuResources {
   /// Returns the set of alternative streams that we use for the given device
   virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0;
 
-  /// Returns the temporary memory manager for the given device
-  virtual DeviceMemory& getMemoryManager(int device) = 0;
+  /// Memory management
+  /// Returns an allocation from the given memory space, ordered with respect to
+  /// the given stream (i.e., the first user will be a kernel in this stream).
+  /// All allocations are sized internally to be the next highest multiple of 16
+  /// bytes, and all allocations returned are guaranteed to be 16 byte aligned.
+  virtual void* allocMemory(const AllocRequest& req) = 0;
+
+  /// Returns a previous allocation
+  virtual void deallocMemory(int device, void* in) = 0;
+
+  /// For MemorySpace::Temporary, how much space is immediately available
+  /// without cudaMalloc allocation?
+  virtual size_t getTempMemoryAvailable(int device) const = 0;
 
   /// Returns the available CPU pinned memory buffer
   virtual std::pair<void*, size_t> getPinnedMemory() = 0;
@@ -46,12 +212,22 @@ class GpuResources {
   /// Returns the stream on which we perform async CPU <-> GPU copies
   virtual cudaStream_t getAsyncCopyStream(int device) = 0;
 
+  ///
+  /// Functions provided by default
+  ///
+
   /// Calls getBlasHandle with the current device
   cublasHandle_t getBlasHandleCurrentDevice();
 
   /// Calls getDefaultStream with the current device
   cudaStream_t getDefaultStreamCurrentDevice();
 
+  /// Calls getTempMemoryAvailable with the current device
+  size_t getTempMemoryAvailableCurrentDevice() const;
+
+  /// Returns a temporary memory allocation via a RAII object
+  GpuMemoryReservation allocMemoryHandle(const AllocRequest& req);
+
   /// Synchronizes the CPU with respect to the default stream for the
   /// given device
   // equivalent to cudaDeviceSynchronize(getDefaultStream(device))
@@ -63,11 +239,17 @@ class GpuResources {
   /// Calls getAlternateStreams for the current device
   std::vector<cudaStream_t> getAlternateStreamsCurrentDevice();
 
-  /// Calls getMemoryManager for the current device
-  DeviceMemory& getMemoryManagerCurrentDevice();
-
   /// Calls getAsyncCopyStream for the current device
   cudaStream_t getAsyncCopyStreamCurrentDevice();
 };
 
+/// Interface for a provider of a shared resources object
+class GpuResourcesProvider {
+ public:
+  virtual ~GpuResourcesProvider();
+
+  /// Returns the shared resources object
+  virtual std::shared_ptr<GpuResources> getResources() = 0;
+};
+
 } } // namespace

gpu/StandardGpuResources.cpp

@@ -8,9 +8,11 @@
 
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
-#include <faiss/gpu/utils/MemorySpace.h>
+#include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/impl/FaissAssert.h>
 #include <limits>
+#include <iostream>
+#include <sstream>
 
 namespace faiss { namespace gpu {
 
@@ -31,9 +33,60 @@ constexpr size_t k8GiBTempMem = (size_t) 1024 * 1024 * 1024;
 // Maximum temporary memory allocation for all GPUs
 constexpr size_t kMaxTempMem = (size_t) 1536 * 1024 * 1024;
 
+std::string allocsToString(const std::unordered_map<void*, AllocRequest>& map) {
+  // Produce a sorted list of all outstanding allocations by type
+  std::unordered_map<AllocType, std::pair<int, size_t>> stats;
+
+  for (auto& entry : map) {
+    auto& a = entry.second;
+
+    auto it = stats.find(a.type);
+    if (it != stats.end()) {
+      stats[a.type].first++;
+      stats[a.type].second += a.size;
+    } else {
+      stats[a.type] = std::make_pair(1, a.size);
+    }
+  }
+
+  std::stringstream ss;
+  for (auto& entry : stats) {
+    ss << "Alloc type " << allocTypeToString(entry.first) << ": "
+       << entry.second.first << " allocations, "
+       << entry.second.second << " bytes\n";
+  }
+
+  return ss.str();
 }
 
-StandardGpuResources::StandardGpuResources() :
+}
+
+std::string allocTypeToString(AllocType t) {
+  switch (t) {
+    case AllocType::Other:
+      return "Other";
+    case AllocType::FlatData:
+      return "FlatData";
+    case AllocType::IVFLists:
+      return "IVFLists";
+    case AllocType::Quantizer:
+      return "Quantizer";
+    case AllocType::QuantizerPrecomputedCodes:
+      return "QuantizerPrecomputedCodes";
+    case AllocType::TemporaryMemoryBuffer:
+      return "TemporaryMemoryBuffer";
+    case AllocType::TemporaryMemoryOverflow:
+      return "TemporaryMemoryOverflow";
+    default:
+      return "Unknown";
+  }
+}
+
+//
+// StandardGpuResourcesImpl
+//
+
+StandardGpuResourcesImpl::StandardGpuResourcesImpl() :
     pinnedMemAlloc_(nullptr),
     pinnedMemAllocSize_(0),
     // let the adjustment function determine the memory size for us by passing
@@ -44,7 +97,28 @@ StandardGpuResources::StandardGpuResources() :
     cudaMallocWarning_(true) {
 }
 
-StandardGpuResources::~StandardGpuResources() {
+StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
+  // The temporary memory allocator has allocated memory through us, so clean
+  // that up before we finish fully de-initializing ourselves
+  tempMemory_.clear();
+
+  // Make sure all allocations have been freed
+  bool allocError = false;
+
+  for (auto& entry : allocs_) {
+    auto& map = entry.second;
+
+    if (!map.empty()) {
+      std::cerr
+        << "StandardGpuResources destroyed with allocations outstanding:\n"
+        << "Device " << entry.first << " outstanding allocations:\n";
+      std::cerr << allocsToString(map);
+      allocError = true;
+    }
+  }
+
+  FAISS_ASSERT_MSG(!allocError, "GPU memory allocations not properly cleaned up");
+
   for (auto& entry : defaultStreams_) {
     DeviceScope scope(entry.first);
 
@@ -78,13 +152,16 @@ StandardGpuResources::~StandardGpuResources() {
   }
 
   if (pinnedMemAlloc_) {
-    freeMemorySpace(MemorySpace::HostPinned, pinnedMemAlloc_);
+    auto err = cudaFreeHost(pinnedMemAlloc_);
+    FAISS_ASSERT_FMT(err == cudaSuccess,
+                     "Failed to cudaFreeHost pointer %p (error %d %s)",
+                     pinnedMemAlloc_, (int) err, cudaGetErrorString(err));
   }
 }
 
 size_t
-StandardGpuResources::getDefaultTempMemForGPU(int device,
-                                              size_t requested) {
+StandardGpuResourcesImpl::getDefaultTempMemForGPU(int device,
+                                                  size_t requested) {
   auto totalMem = device != -1 ?
     getDeviceProperties(device).totalGlobalMem :
     std::numeric_limits<size_t>::max();
@@ -113,13 +190,12 @@ StandardGpuResources::getDefaultTempMemForGPU(int device,
 }
 
 void
-StandardGpuResources::noTempMemory() {
+StandardGpuResourcesImpl::noTempMemory() {
   setTempMemory(0);
-  setCudaMallocWarning(false);
 }
 
 void
-StandardGpuResources::setTempMemory(size_t size) {
+StandardGpuResourcesImpl::setTempMemory(size_t size) {
   if (tempMemSize_ != size) {
     // adjust based on general limits
     tempMemSize_ = getDefaultTempMemForGPU(-1, size);
@@ -129,14 +205,15 @@ StandardGpuResources::setTempMemory(size_t size) {
     // This should be safe to do, even if we are currently running work, because
     // the cudaFree call that this implies will force-synchronize all GPUs with
     // the CPU
-    for (auto& p : memory_) {
+    for (auto& p : tempMemory_) {
       int device = p.first;
       // Free the existing memory first
       p.second.reset();
 
       // Allocate new
       p.second = std::unique_ptr<StackDeviceMemory>(
-        new StackDeviceMemory(p.first,
+        new StackDeviceMemory(this,
+                              p.first,
                               // adjust for this specific device
                               getDefaultTempMemForGPU(device, tempMemSize_)));
     }
@@ -144,7 +221,7 @@ StandardGpuResources::setTempMemory(size_t size) {
 }
 
 void
-StandardGpuResources::setPinnedMemory(size_t size) {
+StandardGpuResourcesImpl::setPinnedMemory(size_t size) {
   // Should not call this after devices have been initialized
   FAISS_ASSERT(defaultStreams_.size() == 0);
   FAISS_ASSERT(!pinnedMemAlloc_);
@@ -153,7 +230,7 @@ StandardGpuResources::setPinnedMemory(size_t size) {
 }
 
 void
-StandardGpuResources::setDefaultStream(int device, cudaStream_t stream) {
+StandardGpuResourcesImpl::setDefaultStream(int device, cudaStream_t stream) {
   auto it = defaultStreams_.find(device);
   if (it != defaultStreams_.end()) {
     // Replace this stream with the user stream
@@ -165,30 +242,21 @@ StandardGpuResources::setDefaultStream(int device, cudaStream_t stream) {
 }
 
 void
-StandardGpuResources::setDefaultNullStreamAllDevices() {
+StandardGpuResourcesImpl::setDefaultNullStreamAllDevices() {
   for (int dev = 0; dev < getNumDevices(); ++dev) {
     setDefaultStream(dev, nullptr);
   }
 }
 
-void
-StandardGpuResources::setCudaMallocWarning(bool b) {
-  cudaMallocWarning_ = b;
-
-  for (auto& v : memory_) {
-    v.second->setCudaMallocWarning(b);
-  }
-}
-
 bool
-StandardGpuResources::isInitialized(int device) const {
+StandardGpuResourcesImpl::isInitialized(int device) const {
   // Use default streams as a marker for whether or not a certain
   // device has been initialized
   return defaultStreams_.count(device) != 0;
 }
 
 void
-StandardGpuResources::initializeForDevice(int device) {
+StandardGpuResourcesImpl::initializeForDevice(int device) {
   if (isInitialized(device)) {
     return;
   }
@@ -196,7 +264,15 @@ StandardGpuResources::initializeForDevice(int device) {
   // If this is the first device that we're initializing, create our
   // pinned memory allocation
   if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
-    allocMemorySpace(MemorySpace::HostPinned, &pinnedMemAlloc_, pinnedMemSize_);
+    auto err =
+      cudaHostAlloc(&pinnedMemAlloc_, pinnedMemSize_, cudaHostAllocDefault);
+
+    FAISS_THROW_IF_NOT_FMT(
+      err == cudaSuccess,
+      "failed to cudaHostAlloc %zu bytes for CPU <-> GPU "
+      "async copy buffer (error %d %s)",
+      pinnedMemSize_, (int) err, cudaGetErrorString(err));
+
     pinnedMemAllocSize_ = pinnedMemSize_;
   }
 
@@ -255,49 +331,244 @@ StandardGpuResources::initializeForDevice(int device) {
   }
 #endif
 
-  FAISS_ASSERT(memory_.count(device) == 0);
+  FAISS_ASSERT(allocs_.count(device) == 0);
+  allocs_[device] = std::unordered_map<void*, AllocRequest>();
 
+  FAISS_ASSERT(tempMemory_.count(device) == 0);
   auto mem = std::unique_ptr<StackDeviceMemory>(
-    new StackDeviceMemory(device,
+    new StackDeviceMemory(this,
+                          device,
                           // adjust for this specific device
                           getDefaultTempMemForGPU(device, tempMemSize_)));
-  mem->setCudaMallocWarning(cudaMallocWarning_);
 
-  memory_.emplace(device, std::move(mem));
+  tempMemory_.emplace(device, std::move(mem));
 }
 
 cublasHandle_t
-StandardGpuResources::getBlasHandle(int device) {
+StandardGpuResourcesImpl::getBlasHandle(int device) {
   initializeForDevice(device);
   return blasHandles_[device];
 }
 
 cudaStream_t
-StandardGpuResources::getDefaultStream(int device) {
+StandardGpuResourcesImpl::getDefaultStream(int device) {
   initializeForDevice(device);
   return defaultStreams_[device];
 }
 
 std::vector<cudaStream_t>
-StandardGpuResources::getAlternateStreams(int device) {
+StandardGpuResourcesImpl::getAlternateStreams(int device) {
   initializeForDevice(device);
   return alternateStreams_[device];
 }
 
-DeviceMemory& StandardGpuResources::getMemoryManager(int device) {
-  initializeForDevice(device);
-  return *memory_[device];
-}
-
 std::pair<void*, size_t>
-StandardGpuResources::getPinnedMemory() {
+StandardGpuResourcesImpl::getPinnedMemory() {
   return std::make_pair(pinnedMemAlloc_, pinnedMemAllocSize_);
 }
 
 cudaStream_t
-StandardGpuResources::getAsyncCopyStream(int device) {
+StandardGpuResourcesImpl::getAsyncCopyStream(int device) {
   initializeForDevice(device);
   return asyncCopyStreams_[device];
 }
 
+void*
+StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
+  initializeForDevice(req.device);
+
+  // We don't allocate a placeholder for zero-sized allocations
+  if (req.size == 0) {
+    return nullptr;
+  }
+
+  // Make sure that the allocation is a multiple of 16 bytes for alignment
+  // purposes
+  auto adjReq = req;
+  adjReq.size = utils::roundUp(adjReq.size, (size_t) 16);
+
+  void* p = nullptr;
+
+  if (adjReq.space == MemorySpace::Temporary) {
+    // If we don't have enough space in our temporary memory manager, we need
+    // to allocate this request separately
+    auto& tempMem = tempMemory_[adjReq.device];
+
+    if (adjReq.size > tempMem->getSizeAvailable()) {
+      // We need to allocate this ourselves
+      AllocRequest newReq = adjReq;
+      newReq.space = MemorySpace::Device;
+      newReq.type = AllocType::TemporaryMemoryOverflow;
+
+      return allocMemory(newReq);
+    }
+
+    // Otherwise, we can handle this locally
+    p = tempMemory_[adjReq.device]->allocMemory(adjReq.stream, adjReq.size);
+
+  } else if (adjReq.space == MemorySpace::Device) {
+    auto err = cudaMalloc(&p, adjReq.size);
+
+    // Throw if we fail to allocate
+    if (err != cudaSuccess) {
+      auto& map = allocs_[req.device];
+
+      std::stringstream ss;
+      ss << "Failed to cudaMalloc " << adjReq.size << " bytes "
+         << "on device " << adjReq.device << " (error "
+         << (int) err << " " << cudaGetErrorString(err)
+         << "\nOutstanding allocations:\n" << allocsToString(map);
+      auto str = ss.str();
+
+      FAISS_THROW_IF_NOT_FMT(err == cudaSuccess, "%s", str.c_str());
+    }
+  } else if (adjReq.space == MemorySpace::Unified) {
+    auto err = cudaMallocManaged(&p, adjReq.size);
+
+    if (err != cudaSuccess) {
+      auto& map = allocs_[req.device];
+
+      std::stringstream ss;
+      ss << "Failed to cudaMallocManaged " << adjReq.size << " bytes "
+         << "(error " << (int) err << " " << cudaGetErrorString(err)
+         << "\nOutstanding allocations:\n" << allocsToString(map);
+      auto str = ss.str();
+
+      FAISS_THROW_IF_NOT_FMT(err == cudaSuccess, "%s", str.c_str());
+    }
+  } else {
+    FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) adjReq.space);
+  }
+
+  allocs_[adjReq.device][p] = adjReq;
+
+  return p;
+}
+
+void
+StandardGpuResourcesImpl::deallocMemory(int device, void* p) {
+  FAISS_ASSERT(isInitialized(device));
+
+  if (!p) {
+    return;
+  }
+
+  auto& a = allocs_[device];
+  auto it = a.find(p);
+  if (it == a.end()) {
+    FAISS_ASSERT(it != a.end());
+  }
+
+  auto& req = it->second;
+
+  if (req.space == MemorySpace::Temporary) {
+    tempMemory_[device]->deallocMemory(device, req.stream, req.size, p);
+
+  } else if (req.space == MemorySpace::Device ||
+             req.space == MemorySpace::Unified) {
+    auto err = cudaFree(p);
+    FAISS_ASSERT_FMT(err == cudaSuccess,
+                     "Failed to cudaFree pointer %p (error %d %s)",
+                     p, (int) err, cudaGetErrorString(err));
+
+  } else {
+    FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) req.space);
+  }
+
+  a.erase(it);
+}
+
+size_t
+StandardGpuResourcesImpl::getTempMemoryAvailable(int device) const {
+  FAISS_ASSERT(isInitialized(device));
+
+  auto it = tempMemory_.find(device);
+  FAISS_ASSERT(it != tempMemory_.end());
+
+  return it->second->getSizeAvailable();
+}
+
+std::map<int, std::map<std::string, std::pair<int, size_t>>>
+StandardGpuResourcesImpl::getMemoryInfo() const {
+  using AT = std::map<std::string, std::pair<int, size_t>>;
+
+  std::map<int, AT> out;
+
+  for (auto& entry : allocs_) {
+    AT outDevice;
+
+    for (auto& a : entry.second) {
+      auto& v = outDevice[allocTypeToString(a.second.type)];
+      v.first++;
+      v.second += a.second.size;
+    }
+
+    out[entry.first] = std::move(outDevice);
+  }
+
+  return out;
+}
+
+//
+// StandardGpuResources
+//
+
+StandardGpuResources::StandardGpuResources()
+    : res_(new StandardGpuResourcesImpl) {
+}
+
+StandardGpuResources::~StandardGpuResources() {
+}
+
+std::shared_ptr<GpuResources>
+StandardGpuResources::getResources() {
+  return res_;
+}
+
+void
+StandardGpuResources::noTempMemory() {
+  res_->noTempMemory();
+}
+
+void
+StandardGpuResources::setTempMemory(size_t size) {
+  res_->setTempMemory(size);
+}
+
+void
+StandardGpuResources::setPinnedMemory(size_t size) {
+  res_->setPinnedMemory(size);
+}
+
+void
+StandardGpuResources::setDefaultStream(int device, cudaStream_t stream) {
+  res_->setDefaultStream(device, stream);
+}
+
+void
+StandardGpuResources::setDefaultNullStreamAllDevices() {
+  res_->setDefaultNullStreamAllDevices();
+}
+
+std::map<int, std::map<std::string, std::pair<int, size_t>>>
+StandardGpuResources::getMemoryInfo() const {
+  return res_->getMemoryInfo();
+}
+
+cudaStream_t
+StandardGpuResources::getDefaultStream(int device) {
+  return res_->getDefaultStream(device);
+}
+
+size_t
+StandardGpuResources::getTempMemoryAvailable(int device) const {
+  return res_->getTempMemoryAvailable(device);
+}
+
+void
+StandardGpuResources::syncDefaultStreamCurrentDevice() {
+  res_->syncDefaultStreamCurrentDevice();
+}
+
+
 } } // namespace

gpu/StandardGpuResources.h

@@ -11,18 +11,20 @@
 #include <faiss/gpu/GpuResources.h>
 #include <faiss/gpu/utils/StackDeviceMemory.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
+#include <functional>
+#include <map>
 #include <unordered_map>
 #include <vector>
 
 namespace faiss { namespace gpu {
 
-/// Default implementation of GpuResources that allocates a cuBLAS
-/// stream and 2 streams for use, as well as temporary memory
-class StandardGpuResources : public GpuResources {
+/// Standard implementation of the GpuResources object that provides for a
+/// temporary memory manager
+class StandardGpuResourcesImpl : public GpuResources {
  public:
-  StandardGpuResources();
+  StandardGpuResourcesImpl();
 
-  ~StandardGpuResources() override;
+  ~StandardGpuResourcesImpl() override;
 
   /// Disable allocation of temporary memory; all temporary memory
   /// requests will call cudaMalloc / cudaFree at the point of use
@@ -46,10 +48,6 @@ class StandardGpuResources : public GpuResources {
   /// for all devices
   void setDefaultNullStreamAllDevices();
 
-  /// Enable or disable the warning about not having enough temporary memory
-  /// when cudaMalloc gets called
-  void setCudaMallocWarning(bool b);
-
  public:
   /// Internal system calls
 
@@ -62,7 +60,17 @@ class StandardGpuResources : public GpuResources {
 
   std::vector<cudaStream_t> getAlternateStreams(int device) override;
 
-  DeviceMemory& getMemoryManager(int device) override;
+  /// Allocate non-temporary GPU memory
+  void* allocMemory(const AllocRequest& req) override;
+
+  /// Returns a previous allocation
+  void deallocMemory(int device, void* in) override;
+
+  size_t getTempMemoryAvailable(int device) const override;
+
+  /// Export a description of memory used for Python
+  std::map<int, std::map<std::string, std::pair<int, size_t>>>
+  getMemoryInfo() const;
 
   std::pair<void*, size_t> getPinnedMemory() override;
 
@@ -77,6 +85,13 @@ class StandardGpuResources : public GpuResources {
   static size_t getDefaultTempMemForGPU(int device, size_t requested);
 
  private:
+  /// Set of currently outstanding memory allocations per device
+  /// device -> (alloc request, allocated ptr)
+  std::unordered_map<int, std::unordered_map<void*, AllocRequest>> allocs_;
+
+  /// Temporary memory provider, per each device
+  std::unordered_map<int, std::unique_ptr<StackDeviceMemory>> tempMemory_;
+
   /// Our default stream that work is ordered on, one per each device
   std::unordered_map<int, cudaStream_t> defaultStreams_;
 
@@ -85,7 +100,7 @@ class StandardGpuResources : public GpuResources {
   std::unordered_map<int, cudaStream_t> userDefaultStreams_;
 
   /// Other streams we can use, per each device
-  std::unordered_map<int, std::vector<cudaStream_t> > alternateStreams_;
+  std::unordered_map<int, std::vector<cudaStream_t>> alternateStreams_;
 
   /// Async copy stream to use for GPU <-> CPU pinned memory copies
   std::unordered_map<int, cudaStream_t> asyncCopyStreams_;
@@ -93,9 +108,6 @@ class StandardGpuResources : public GpuResources {
   /// cuBLAS handle for each device
   std::unordered_map<int, cublasHandle_t> blasHandles_;
 
-  /// Temporary memory provider, per each device
-  std::unordered_map<int, std::unique_ptr<StackDeviceMemory> > memory_;
-
   /// Pinned memory allocation for use with this GPU
   void* pinnedMemAlloc_;
   size_t pinnedMemAllocSize_;
@@ -111,4 +123,52 @@ class StandardGpuResources : public GpuResources {
   bool cudaMallocWarning_;
 };
 
+/// Default implementation of GpuResources that allocates a cuBLAS
+/// stream and 2 streams for use, as well as temporary memory
+class StandardGpuResources : public GpuResourcesProvider {
+ public:
+  StandardGpuResources();
+  ~StandardGpuResources() override;
+
+  std::shared_ptr<GpuResources> getResources() override;
+
+  /// Disable allocation of temporary memory; all temporary memory
+  /// requests will call cudaMalloc / cudaFree at the point of use
+  void noTempMemory();
+
+  /// Specify that we wish to use a certain fixed size of memory on
+  /// all devices as temporary memory. This is the upper bound for the GPU
+  /// memory that we will reserve. We will never go above 1.5 GiB on any GPU;
+  /// smaller GPUs (with <= 4 GiB or <= 8 GiB) will use less memory than that.
+  /// To avoid any temporary memory allocation, pass 0.
+  void setTempMemory(size_t size);
+
+  /// Set amount of pinned memory to allocate, for async GPU <-> CPU
+  /// transfers
+  void setPinnedMemory(size_t size);
+
+  /// Called to change the stream for work ordering
+  void setDefaultStream(int device, cudaStream_t stream);
+
+  /// Called to change the work ordering streams to the null stream
+  /// for all devices
+  void setDefaultNullStreamAllDevices();
+
+  /// Export a description of memory used for Python
+  std::map<int, std::map<std::string, std::pair<int, size_t>>>
+  getMemoryInfo() const;
+
+  /// Returns the current default stream
+  cudaStream_t getDefaultStream(int device);
+
+  /// Returns the current amount of temp memory available
+  size_t getTempMemoryAvailable(int device) const;
+
+  /// Synchronize our default stream with the CPU
+  void syncDefaultStreamCurrentDevice();
+
+ private:
+  std::shared_ptr<StandardGpuResourcesImpl> res_;
+};
+
 } } // namespace

gpu/impl/BinaryFlatIndex.cu

@@ -18,9 +18,12 @@ BinaryFlatIndex::BinaryFlatIndex(GpuResources* res,
                                  MemorySpace space) :
     resources_(res),
     dim_(dim),
-    space_(space),
     num_(0),
-    rawData_(space) {
+    rawData_(res,
+             makeSpaceAlloc(
+               AllocType::FlatData,
+               space,
+               res->getDefaultStreamCurrentDevice())) {
   FAISS_ASSERT(dim % 8 == 0);
 }
 
@@ -74,14 +77,14 @@ BinaryFlatIndex::add(const unsigned char* data,
   num_ += numVecs;
 
   DeviceTensor<unsigned char, 2, true> vectors(
-    (unsigned char*) rawData_.data(), {(int) num_, (dim_ / 8)}, space_);
+    (unsigned char*) rawData_.data(), {(int) num_, (dim_ / 8)});
   vectors_ = std::move(vectors);
 }
 
 void
 BinaryFlatIndex::reset() {
   rawData_.clear();
-  vectors_ = std::move(DeviceTensor<unsigned char, 2, true>());
+  vectors_ = DeviceTensor<unsigned char, 2, true>();
   num_ = 0;
 }
 

gpu/impl/BinaryFlatIndex.cuh

@@ -10,7 +10,6 @@
 
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 #include <faiss/gpu/utils/DeviceVector.cuh>
-#include <faiss/gpu/utils/MemorySpace.h>
 
 namespace faiss { namespace gpu {
 
@@ -53,9 +52,6 @@ class BinaryFlatIndex {
   /// Dimensionality of our vectors
   const int dim_;
 
-  /// Memory space for our allocations
-  MemorySpace space_;
-
   /// How many vectors we have
   int num_;
 

gpu/impl/Distance.cu

@@ -31,7 +31,7 @@ namespace faiss { namespace gpu {
 
 template <typename T>
 void runDistance(bool computeL2,
-                 GpuResources* resources,
+                 GpuResources* res,
                  Tensor<T, 2, true>& centroids,
                  bool centroidsRowMajor,
                  Tensor<float, 1, true>* centroidNorms,
@@ -57,8 +57,7 @@ void runDistance(bool computeL2,
   FAISS_ASSERT(outDistances.getSize(1) == k);
   FAISS_ASSERT(outIndices.getSize(1) == k);
 
-  auto& mem = resources->getMemoryManagerCurrentDevice();
-  auto defaultStream = resources->getDefaultStreamCurrentDevice();
+  auto defaultStream = res->getDefaultStreamCurrentDevice();
 
   // If we're quering against a 0 sized set, just return empty results
   if (centroids.numElements() == 0) {
@@ -76,9 +75,10 @@ void runDistance(bool computeL2,
   // L2: If ||c||^2 is not pre-computed, calculate it
   DeviceTensor<float, 1, true> cNorms;
   if (computeL2 && !centroidNorms) {
-    cNorms =
-      std::move(DeviceTensor<float, 1, true>(
-                  mem, {numCentroids}, defaultStream));
+    cNorms = DeviceTensor<float, 1, true>(
+      res,
+      makeTempAlloc(AllocType::Other, defaultStream),
+      {numCentroids});
     runL2Norm(centroids, centroidsRowMajor, cNorms, true, defaultStream);
     centroidNorms = &cNorms;
   }
@@ -86,8 +86,10 @@ void runDistance(bool computeL2,
   //
   // Prepare norm vector ||q||^2; ||c||^2 is already pre-computed
   //
-  int qNormSize[1] = {numQueries};
-  DeviceTensor<float, 1, true> queryNorms(mem, qNormSize, defaultStream);
+  DeviceTensor<float, 1, true> queryNorms(
+      res,
+      makeTempAlloc(AllocType::Other, defaultStream),
+      {(int) numQueries});
 
   // ||q||^2
   if (computeL2) {
@@ -102,7 +104,7 @@ void runDistance(bool computeL2,
                  numCentroids,
                  dim,
                  sizeof(T),
-                 mem.getSizeAvailable(),
+                 res->getTempMemoryAvailableCurrentDevice(),
                  tileRows,
                  tileCols);
 
@@ -114,27 +116,27 @@ void runDistance(bool computeL2,
 
   // Temporary output memory space we'll use
   DeviceTensor<float, 2, true> distanceBuf1(
-    mem, {tileRows, tileCols}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, tileCols});
   DeviceTensor<float, 2, true> distanceBuf2(
-    mem, {tileRows, tileCols}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, tileCols});
   DeviceTensor<float, 2, true>* distanceBufs[2] =
     {&distanceBuf1, &distanceBuf2};
 
   DeviceTensor<float, 2, true> outDistanceBuf1(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<float, 2, true> outDistanceBuf2(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<float, 2, true>* outDistanceBufs[2] =
     {&outDistanceBuf1, &outDistanceBuf2};
 
   DeviceTensor<int, 2, true> outIndexBuf1(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<int, 2, true> outIndexBuf2(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<int, 2, true>* outIndexBufs[2] =
     {&outIndexBuf1, &outIndexBuf2};
 
-  auto streams = resources->getAlternateStreamsCurrentDevice();
+  auto streams = res->getAlternateStreamsCurrentDevice();
   streamWait(streams, {defaultStream});
 
   int curStream = 0;
@@ -196,7 +198,7 @@ void runDistance(bool computeL2,
                     centroidsRowMajor, // transposed MM if row major
                     computeL2 ? -2.0f : 1.0f,
                     0.0f,
-                    resources->getBlasHandleCurrentDevice(),
+                    res->getBlasHandleCurrentDevice(),
                     streams[curStream]);
 
       if (computeL2) {
@@ -290,7 +292,7 @@ void runDistance(bool computeL2,
 }
 
 template <typename T>
-void runL2Distance(GpuResources* resources,
+void runL2Distance(GpuResources* res,
                    Tensor<T, 2, true>& centroids,
                    bool centroidsRowMajor,
                    Tensor<float, 1, true>* centroidNorms,
@@ -301,7 +303,7 @@ void runL2Distance(GpuResources* resources,
                    Tensor<int, 2, true>& outIndices,
                    bool ignoreOutDistances = false) {
   runDistance<T>(true, // L2
-                 resources,
+                 res,
                  centroids,
                  centroidsRowMajor,
                  centroidNorms,
@@ -314,7 +316,7 @@ void runL2Distance(GpuResources* resources,
 }
 
 template <typename T>
-void runIPDistance(GpuResources* resources,
+void runIPDistance(GpuResources* res,
                    Tensor<T, 2, true>& centroids,
                    bool centroidsRowMajor,
                    Tensor<T, 2, true>& queries,
@@ -323,7 +325,7 @@ void runIPDistance(GpuResources* resources,
                    Tensor<float, 2, true>& outDistances,
                    Tensor<int, 2, true>& outIndices) {
   runDistance<T>(false, // IP
-                 resources,
+                 res,
                  centroids,
                  centroidsRowMajor,
                  nullptr, // no centroid norms provided
@@ -340,7 +342,7 @@ void runIPDistance(GpuResources* resources,
 //
 
 void
-runIPDistance(GpuResources* resources,
+runIPDistance(GpuResources* res,
               Tensor<float, 2, true>& vectors,
               bool vectorsRowMajor,
               Tensor<float, 2, true>& queries,
@@ -348,7 +350,7 @@ runIPDistance(GpuResources* resources,
               int k,
               Tensor<float, 2, true>& outDistances,
               Tensor<int, 2, true>& outIndices) {
-  runIPDistance<float>(resources,
+  runIPDistance<float>(res,
                        vectors,
                        vectorsRowMajor,
                        queries,
@@ -359,7 +361,7 @@ runIPDistance(GpuResources* resources,
 }
 
 void
-runIPDistance(GpuResources* resources,
+runIPDistance(GpuResources* res,
               Tensor<half, 2, true>& vectors,
               bool vectorsRowMajor,
               Tensor<half, 2, true>& queries,
@@ -367,7 +369,7 @@ runIPDistance(GpuResources* resources,
               int k,
               Tensor<float, 2, true>& outDistances,
               Tensor<int, 2, true>& outIndices) {
-  runIPDistance<half>(resources,
+  runIPDistance<half>(res,
                       vectors,
                       vectorsRowMajor,
                       queries,
@@ -378,7 +380,7 @@ runIPDistance(GpuResources* resources,
 }
 
 void
-runL2Distance(GpuResources* resources,
+runL2Distance(GpuResources* res,
               Tensor<float, 2, true>& vectors,
               bool vectorsRowMajor,
               Tensor<float, 1, true>* vectorNorms,
@@ -388,7 +390,7 @@ runL2Distance(GpuResources* resources,
               Tensor<float, 2, true>& outDistances,
               Tensor<int, 2, true>& outIndices,
               bool ignoreOutDistances) {
-  runL2Distance<float>(resources,
+  runL2Distance<float>(res,
                        vectors,
                        vectorsRowMajor,
                        vectorNorms,
@@ -401,7 +403,7 @@ runL2Distance(GpuResources* resources,
 }
 
 void
-runL2Distance(GpuResources* resources,
+runL2Distance(GpuResources* res,
               Tensor<half, 2, true>& vectors,
               bool vectorsRowMajor,
               Tensor<float, 1, true>* vectorNorms,
@@ -411,7 +413,7 @@ runL2Distance(GpuResources* resources,
               Tensor<float, 2, true>& outDistances,
               Tensor<int, 2, true>& outIndices,
               bool ignoreOutDistances) {
-  runL2Distance<half>(resources,
+  runL2Distance<half>(res,
                       vectors,
                       vectorsRowMajor,
                       vectorNorms,

gpu/impl/FlatIndex.cu

@@ -28,7 +28,13 @@ FlatIndex::FlatIndex(GpuResources* res,
     storeTransposed_(storeTransposed),
     space_(space),
     num_(0),
-    rawData_(space) {
+    rawData_(
+      res,
+      AllocInfo(
+        AllocType::FlatData,
+        getCurrentDevice(),
+        space,
+        res->getDefaultStreamCurrentDevice())) {
 }
 
 bool
@@ -101,7 +107,10 @@ FlatIndex::getVectorsFloat32Copy(cudaStream_t stream) {
 
 DeviceTensor<float, 2, true>
 FlatIndex::getVectorsFloat32Copy(int from, int num, cudaStream_t stream) {
-  DeviceTensor<float, 2, true> vecFloat32({num, dim_}, space_);
+  DeviceTensor<float, 2, true> vecFloat32(
+    resources_,
+    makeDevAlloc(AllocType::Other, stream),
+    {num, dim_});
 
   if (useFloat16_) {
     auto halfNarrow = vectorsHalf_.narrowOutermost(from, num);
@@ -122,12 +131,11 @@ FlatIndex::query(Tensor<float, 2, true>& input,
                  Tensor<int, 2, true>& outIndices,
                  bool exactDistance) {
   auto stream = resources_->getDefaultStreamCurrentDevice();
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
 
   if (useFloat16_) {
     // We need to convert the input to float16 for comparison to ourselves
     auto inputHalf =
-      convertTensor<float, half, 2>(resources_, stream, input);
+      convertTensorTemporary<float, half, 2>(resources_, stream, input);
 
     query(inputHalf, k, metric, metricArg,
           outDistances, outIndices, exactDistance);
@@ -228,14 +236,14 @@ FlatIndex::add(const float* data, int numVecs, cudaStream_t stream) {
   if (useFloat16_) {
     // Make sure that `data` is on our device; we'll run the
     // conversion on our device
-    auto devData = toDevice<float, 2>(resources_,
-                                      getCurrentDevice(),
-                                      (float*) data,
-                                      stream,
-                                      {numVecs, dim_});
+    auto devData = toDeviceTemporary<float, 2>(resources_,
+                                               getCurrentDevice(),
+                                               (float*) data,
+                                               stream,
+                                               {numVecs, dim_});
 
     auto devDataHalf =
-      convertTensor<float, half, 2>(resources_, stream, devData);
+      convertTensorTemporary<float, half, 2>(resources_, stream, devData);
 
     rawData_.append((char*) devDataHalf.data(),
                     devDataHalf.getSizeInBytes(),
@@ -252,33 +260,45 @@ FlatIndex::add(const float* data, int numVecs, cudaStream_t stream) {
 
   if (useFloat16_) {
     DeviceTensor<half, 2, true> vectorsHalf(
-      (half*) rawData_.data(), {(int) num_, dim_}, space_);
+      (half*) rawData_.data(), {(int) num_, dim_});
     vectorsHalf_ = std::move(vectorsHalf);
   } else {
     DeviceTensor<float, 2, true> vectors(
-      (float*) rawData_.data(), {(int) num_, dim_}, space_);
+      (float*) rawData_.data(), {(int) num_, dim_});
     vectors_ = std::move(vectors);
   }
 
   if (storeTransposed_) {
     if (useFloat16_) {
       vectorsHalfTransposed_ =
-        std::move(DeviceTensor<half, 2, true>({dim_, (int) num_}, space_));
+        DeviceTensor<half, 2, true>(
+          resources_,
+          makeSpaceAlloc(AllocType::FlatData, space_, stream),
+          {dim_, (int) num_});
       runTransposeAny(vectorsHalf_, 0, 1, vectorsHalfTransposed_, stream);
     } else {
       vectorsTransposed_ =
-        std::move(DeviceTensor<float, 2, true>({dim_, (int) num_}, space_));
+        DeviceTensor<float, 2, true>(
+          resources_,
+          makeSpaceAlloc(AllocType::FlatData, space_, stream),
+          {dim_, (int)num_});
       runTransposeAny(vectors_, 0, 1, vectorsTransposed_, stream);
     }
   }
 
   // Precompute L2 norms of our database
   if (useFloat16_) {
-    DeviceTensor<float, 1, true> norms({(int) num_}, space_);
+    DeviceTensor<float, 1, true> norms(
+      resources_,
+      makeSpaceAlloc(AllocType::FlatData, space_, stream),
+      {(int) num_});
     runL2Norm(vectorsHalf_, true, norms, true, stream);
     norms_ = std::move(norms);
   } else {
-    DeviceTensor<float, 1, true> norms({(int) num_}, space_);
+    DeviceTensor<float, 1, true> norms(
+      resources_,
+      makeSpaceAlloc(AllocType::FlatData, space_, stream),
+      {(int) num_});
     runL2Norm(vectors_, true, norms, true, stream);
     norms_ = std::move(norms);
   }
@@ -287,11 +307,11 @@ FlatIndex::add(const float* data, int numVecs, cudaStream_t stream) {
 void
 FlatIndex::reset() {
   rawData_.clear();
-  vectors_ = std::move(DeviceTensor<float, 2, true>());
-  vectorsTransposed_ = std::move(DeviceTensor<float, 2, true>());
-  vectorsHalf_ = std::move(DeviceTensor<half, 2, true>());
-  vectorsHalfTransposed_ = std::move(DeviceTensor<half, 2, true>());
-  norms_ = std::move(DeviceTensor<float, 1, true>());
+  vectors_ = DeviceTensor<float, 2, true>();
+  vectorsTransposed_ = DeviceTensor<float, 2, true>();
+  vectorsHalf_ = DeviceTensor<half, 2, true>();
+  vectorsHalfTransposed_ = DeviceTensor<half, 2, true>();
+  norms_ = DeviceTensor<float, 1, true>();
   num_ = 0;
 }
 

gpu/impl/FlatIndex.cuh

@@ -11,7 +11,7 @@
 #include <faiss/MetricType.h>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 #include <faiss/gpu/utils/DeviceVector.cuh>
-#include <faiss/gpu/utils/MemorySpace.h>
+#include <faiss/gpu/GpuResources.h>
 
 namespace faiss { namespace gpu {
 

gpu/impl/GeneralDistance.cuh

@@ -249,7 +249,7 @@ void runGeneralDistanceKernel(Tensor<T, 2, InnerContig>& vecs,
 }
 
 template <typename T, typename DistanceOp, bool InnerContig>
-void runGeneralDistance(GpuResources* resources,
+void runGeneralDistance(GpuResources* res,
                         Tensor<T, 2, InnerContig>& centroids,
                         Tensor<T, 2, InnerContig>& queries,
                         int k,
@@ -272,8 +272,7 @@ void runGeneralDistance(GpuResources* resources,
   FAISS_ASSERT(outDistances.getSize(1) == k);
   FAISS_ASSERT(outIndices.getSize(1) == k);
 
-  auto& mem = resources->getMemoryManagerCurrentDevice();
-  auto defaultStream = resources->getDefaultStreamCurrentDevice();
+  auto defaultStream = res->getDefaultStreamCurrentDevice();
 
   // If we're quering against a 0 sized set, just return empty results
   if (centroids.numElements() == 0) {
@@ -296,7 +295,7 @@ void runGeneralDistance(GpuResources* resources,
                  numCentroids,
                  dim,
                  sizeof(T),
-                 mem.getSizeAvailable(),
+                 res->getTempMemoryAvailableCurrentDevice(),
                  tileRows,
                  tileCols);
 
@@ -308,27 +307,27 @@ void runGeneralDistance(GpuResources* resources,
 
   // Temporary output memory space we'll use
   DeviceTensor<float, 2, true> distanceBuf1(
-    mem, {tileRows, tileCols}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, tileCols});
   DeviceTensor<float, 2, true> distanceBuf2(
-    mem, {tileRows, tileCols}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, tileCols});
   DeviceTensor<float, 2, true>* distanceBufs[2] =
     {&distanceBuf1, &distanceBuf2};
 
   DeviceTensor<float, 2, true> outDistanceBuf1(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<float, 2, true> outDistanceBuf2(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<float, 2, true>* outDistanceBufs[2] =
     {&outDistanceBuf1, &outDistanceBuf2};
 
   DeviceTensor<int, 2, true> outIndexBuf1(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<int, 2, true> outIndexBuf2(
-    mem, {tileRows, numColTiles * k}, defaultStream);
+    res, makeTempAlloc(AllocType::Other, defaultStream), {tileRows, numColTiles * k});
   DeviceTensor<int, 2, true>* outIndexBufs[2] =
     {&outIndexBuf1, &outIndexBuf2};
 
-  auto streams = resources->getAlternateStreamsCurrentDevice();
+  auto streams = res->getAlternateStreamsCurrentDevice();
   streamWait(streams, {defaultStream});
 
   int curStream = 0;

gpu/impl/GpuScalarQuantizer.cuh

@@ -31,9 +31,13 @@ inline bool isSQSupported(ScalarQuantizer::QuantizerType qtype) {
 // Wrapper around the CPU ScalarQuantizer that allows storage of parameters in
 // GPU memory
 struct GpuScalarQuantizer : public ScalarQuantizer {
-  GpuScalarQuantizer(const ScalarQuantizer& sq)
+  GpuScalarQuantizer(GpuResources* res,
+                     const ScalarQuantizer& sq)
       : ScalarQuantizer(sq),
-        gpuTrained(DeviceTensor<float, 1, true>({(int) sq.trained.size()})) {
+        gpuTrained(DeviceTensor<float, 1, true>(
+                     res,
+                     makeDevAlloc(AllocType::Quantizer, 0),
+          {(int) sq.trained.size()})) {
     HostTensor<float, 1, true>
       cpuTrained((float*) sq.trained.data(), {(int) sq.trained.size()});
 

gpu/impl/IVFBase.cu

@@ -31,11 +31,11 @@ IVFBase::IVFBase(GpuResources* resources,
     metric_(metric),
     metricArg_(metricArg),
     quantizer_(quantizer),
+    dim_(quantizer->getDim()),
+    numLists_(quantizer->getSize()),
     bytesPerVector_(bytesPerVector),
     indicesOptions_(indicesOptions),
     space_(space),
-    dim_(quantizer->getDim()),
-    numLists_(quantizer->getSize()),
     maxListLength_(0) {
   reset();
 }
@@ -82,13 +82,18 @@ IVFBase::reset() {
   deviceListLengths_.clear();
   listOffsetToUserIndex_.clear();
 
+  auto info = AllocInfo(AllocType::IVFLists,
+                        getCurrentDevice(),
+                        space_,
+                        resources_->getDefaultStreamCurrentDevice());
+
   for (size_t i = 0; i < numLists_; ++i) {
     deviceListData_.emplace_back(
       std::unique_ptr<DeviceVector<unsigned char>>(
-        new DeviceVector<unsigned char>(space_)));
+        new DeviceVector<unsigned char>(resources_, info)));
     deviceListIndices_.emplace_back(
       std::unique_ptr<DeviceVector<unsigned char>>(
-        new DeviceVector<unsigned char>(space_)));
+        new DeviceVector<unsigned char>(resources_, info)));
     listOffsetToUserIndex_.emplace_back(std::vector<long>());
   }
 
@@ -149,8 +154,6 @@ IVFBase::updateDeviceListInfo_(cudaStream_t stream) {
 void
 IVFBase::updateDeviceListInfo_(const std::vector<int>& listIds,
                                cudaStream_t stream) {
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
-
   HostTensor<int, 1, true>
     hostListsToUpdate({(int) listIds.size()});
   HostTensor<int, 1, true>
@@ -173,13 +176,13 @@ IVFBase::updateDeviceListInfo_(const std::vector<int>& listIds,
 
   // Copy the above update sets to the GPU
   DeviceTensor<int, 1, true> listsToUpdate(
-    mem, hostListsToUpdate, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream), hostListsToUpdate);
   DeviceTensor<int, 1, true> newListLength(
-    mem,  hostNewListLength, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream), hostNewListLength);
   DeviceTensor<void*, 1, true> newDataPointers(
-    mem, hostNewDataPointers, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream), hostNewDataPointers);
   DeviceTensor<void*, 1, true> newIndexPointers(
-    mem, hostNewIndexPointers, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream), hostNewIndexPointers);
 
   // Update all pointers to the lists on the device that may have
   // changed

gpu/impl/IVFBase.cuh

@@ -12,7 +12,6 @@
 #include <faiss/gpu/GpuIndicesOptions.h>
 #include <faiss/gpu/utils/DeviceVector.cuh>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
-#include <faiss/gpu/utils/MemorySpace.h>
 #include <memory>
 #include <thrust/device_vector.h>
 #include <vector>

gpu/impl/IVFFlat.cu

@@ -25,7 +25,7 @@
 
 namespace faiss { namespace gpu {
 
-IVFFlat::IVFFlat(GpuResources* resources,
+IVFFlat::IVFFlat(GpuResources* res,
                  FlatIndex* quantizer,
                  faiss::MetricType metric,
                  float metricArg,
@@ -33,7 +33,7 @@ IVFFlat::IVFFlat(GpuResources* resources,
                  faiss::ScalarQuantizer* scalarQ,
                  IndicesOptions indicesOptions,
                  MemorySpace space) :
-    IVFBase(resources,
+    IVFBase(res,
             metric,
             metricArg,
             quantizer,
@@ -42,7 +42,7 @@ IVFFlat::IVFFlat(GpuResources* resources,
             indicesOptions,
             space),
     useResidual_(useResidual),
-    scalarQ_(scalarQ ? new GpuScalarQuantizer(*scalarQ) : nullptr) {
+    scalarQ_(scalarQ ? new GpuScalarQuantizer(res, *scalarQ) : nullptr) {
 }
 
 IVFFlat::~IVFFlat() {
@@ -106,25 +106,24 @@ IVFFlat::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
   FAISS_ASSERT(vecs.getSize(0) == indices.getSize(0));
   FAISS_ASSERT(vecs.getSize(1) == dim_);
 
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
   auto stream = resources_->getDefaultStreamCurrentDevice();
 
   // Number of valid vectors that we actually add; we return this
   int numAdded = 0;
 
-  DeviceTensor<float, 2, true>
-    listDistance2d(mem, {vecs.getSize(0), 1}, stream);
+  DeviceTensor<float, 2, true> listDistance2d(
+    resources_, makeTempAlloc(AllocType::Other, stream), {vecs.getSize(0), 1});
 
-  DeviceTensor<int, 2, true>
-    listIds2d(mem, {vecs.getSize(0), 1},  stream);
+  DeviceTensor<int, 2, true> listIds2d(
+    resources_, makeTempAlloc(AllocType::Other, stream), {vecs.getSize(0), 1});
   auto listIds = listIds2d.view<1>({vecs.getSize(0)});
 
   quantizer_->query(vecs, 1, metric_, metricArg_,
                     listDistance2d, listIds2d, false);
 
   // Calculate residuals for these vectors, if needed
-  DeviceTensor<float, 2, true>
-    residuals(mem, {vecs.getSize(0), dim_}, stream);
+  DeviceTensor<float, 2, true> residuals(
+    resources_, makeTempAlloc(AllocType::Other, stream), {vecs.getSize(0), dim_});
 
   if (useResidual_) {
     quantizer_->computeResidual(vecs, listIds, residuals);
@@ -251,7 +250,8 @@ IVFFlat::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
 
   // We similarly need to actually append the new vectors
   {
-    DeviceTensor<int, 1, true> listOffset(mem, listOffsetHost, stream);
+    DeviceTensor<int, 1, true> listOffset(
+      resources_, makeTempAlloc(AllocType::Other, stream), listOffsetHost);
 
     // Now, for each list to which a vector is being assigned, write it
     runIVFFlatInvertedListAppend(listIds,
@@ -276,7 +276,6 @@ IVFFlat::query(Tensor<float, 2, true>& queries,
                int k,
                Tensor<float, 2, true>& outDistances,
                Tensor<long, 2, true>& outIndices) {
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
   auto stream = resources_->getDefaultStreamCurrentDevice();
 
   // These are caught at a higher level
@@ -290,10 +289,10 @@ IVFFlat::query(Tensor<float, 2, true>& queries,
   FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
 
   // Reserve space for the quantized information
-  DeviceTensor<float, 2, true>
-    coarseDistances(mem, {queries.getSize(0), nprobe}, stream);
-  DeviceTensor<int, 2, true>
-    coarseIndices(mem, {queries.getSize(0), nprobe}, stream);
+  DeviceTensor<float, 2, true> coarseDistances(
+    resources_, makeTempAlloc(AllocType::Other, stream), {queries.getSize(0), nprobe});
+  DeviceTensor<int, 2, true> coarseIndices(
+    resources_, makeTempAlloc(AllocType::Other, stream), {queries.getSize(0), nprobe});
 
   // Find the `nprobe` closest lists; we can use int indices both
   // internally and externally
@@ -305,8 +304,9 @@ IVFFlat::query(Tensor<float, 2, true>& queries,
                     coarseIndices,
                     false);
 
-  DeviceTensor<float, 3, true>
-    residualBase(mem, {queries.getSize(0), nprobe, dim_}, stream);
+  DeviceTensor<float, 3, true> residualBase(
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {queries.getSize(0), nprobe, dim_});
 
   if (useResidual_) {
     // Reconstruct vectors from the quantizer

gpu/impl/IVFFlatScan.cu

@@ -176,7 +176,8 @@ ivfFlatScan(Tensor<float, 2, true> queries,
 }
 
 void
-runIVFFlatScanTile(Tensor<float, 2, true>& queries,
+runIVFFlatScanTile(GpuResources* res,
+                   Tensor<float, 2, true>& queries,
                    Tensor<int, 2, true>& listIds,
                    thrust::device_vector<void*>& listData,
                    thrust::device_vector<void*>& listIndices,
@@ -214,7 +215,8 @@ runIVFFlatScanTile(Tensor<float, 2, true>& queries,
 
   // Calculate offset lengths, so we know where to write out
   // intermediate results
-  runCalcListOffsets(listIds, listLengths, prefixSumOffsets, thrustMem, stream);
+  runCalcListOffsets(
+    res, listIds, listLengths, prefixSumOffsets, thrustMem, stream);
 
   auto grid = dim3(listIds.getSize(1), listIds.getSize(0));
   auto block = dim3(kWarpSize * kIVFFlatScanWarps);
@@ -383,21 +385,20 @@ runIVFFlatScan(Tensor<float, 2, true>& queries,
 
   int nprobe = listIds.getSize(1);
 
-  auto& mem = res->getMemoryManagerCurrentDevice();
   auto stream = res->getDefaultStreamCurrentDevice();
 
   // Make a reservation for Thrust to do its dirty work (global memory
   // cross-block reduction space); hopefully this is large enough.
   DeviceTensor<char, 1, true> thrustMem1(
-    mem, {kThrustMemSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream), {kThrustMemSize});
   DeviceTensor<char, 1, true> thrustMem2(
-    mem, {kThrustMemSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream), {kThrustMemSize});
   DeviceTensor<char, 1, true>* thrustMem[2] =
     {&thrustMem1, &thrustMem2};
 
   // How much temporary storage is available?
   // If possible, we'd like to fit within the space available.
-  size_t sizeAvailable = mem.getSizeAvailable();
+  size_t sizeAvailable = res->getTempMemoryAvailableCurrentDevice();
 
   // We run two passes of heap selection
   // This is the size of the first-level heap passes
@@ -431,9 +432,9 @@ runIVFFlatScan(Tensor<float, 2, true>& queries,
   // Make sure there is space prior to the start which will be 0, and
   // will handle the boundary condition without branches
   DeviceTensor<int, 1, true> prefixSumOffsetSpace1(
-    mem, {queryTileSize * nprobe + 1}, stream);
+    res, makeTempAlloc(AllocType::Other, stream), {queryTileSize * nprobe + 1});
   DeviceTensor<int, 1, true> prefixSumOffsetSpace2(
-    mem, {queryTileSize * nprobe + 1}, stream);
+    res, makeTempAlloc(AllocType::Other, stream), {queryTileSize * nprobe + 1});
 
   DeviceTensor<int, 2, true> prefixSumOffsets1(
     prefixSumOffsetSpace1[1].data(),
@@ -456,23 +457,29 @@ runIVFFlatScan(Tensor<float, 2, true>& queries,
                               stream));
 
   DeviceTensor<float, 1, true> allDistances1(
-    mem, {queryTileSize * nprobe * maxListLength}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe * maxListLength});
   DeviceTensor<float, 1, true> allDistances2(
-    mem, {queryTileSize * nprobe * maxListLength}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe * maxListLength});
   DeviceTensor<float, 1, true>* allDistances[2] =
     {&allDistances1, &allDistances2};
 
   DeviceTensor<float, 3, true> heapDistances1(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<float, 3, true> heapDistances2(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<float, 3, true>* heapDistances[2] =
     {&heapDistances1, &heapDistances2};
 
   DeviceTensor<int, 3, true> heapIndices1(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<int, 3, true> heapIndices2(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<int, 3, true>* heapIndices[2] =
     {&heapIndices1, &heapIndices2};
 
@@ -505,7 +512,8 @@ runIVFFlatScan(Tensor<float, 2, true>& queries,
     auto outIndicesView =
       outIndices.narrowOutermost(query, numQueriesInTile);
 
-    runIVFFlatScanTile(queryView,
+    runIVFFlatScanTile(res,
+                       queryView,
                        listIdsView,
                        listData,
                        listIndices,

gpu/impl/IVFPQ.cu

@@ -52,8 +52,8 @@ IVFPQ::IVFPQ(GpuResources* resources,
     bitsPerSubQuantizer_(bitsPerSubQuantizer),
     numSubQuantizerCodes_(utils::pow2(bitsPerSubQuantizer_)),
     dimPerSubQuantizer_(dim_ / numSubQuantizers),
-    precomputedCodes_(false),
-    useFloat16LookupTables_(useFloat16LookupTables) {
+    useFloat16LookupTables_(useFloat16LookupTables),
+    precomputedCodes_(false) {
   FAISS_ASSERT(pqCentroidData);
 
   FAISS_ASSERT(bitsPerSubQuantizer_ <= 8);
@@ -111,8 +111,8 @@ IVFPQ::setPrecomputedCodes(bool enable) {
       precomputeCodes_();
     } else {
       // Clear out old precomputed code data
-      precomputedCode_ = std::move(DeviceTensor<float, 3, true>());
-      precomputedCodeHalf_ = std::move(DeviceTensor<half, 3, true>());
+      precomputedCode_ = DeviceTensor<float, 3, true>();
+      precomputedCodeHalf_ = DeviceTensor<half, 3, true>();
     }
   }
 }
@@ -123,16 +123,19 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
   FAISS_ASSERT(vecs.getSize(0) == indices.getSize(0));
   FAISS_ASSERT(vecs.getSize(1) == dim_);
 
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
   auto stream = resources_->getDefaultStreamCurrentDevice();
 
   // Number of valid vectors that we actually add; we return this
   int numAdded = 0;
 
   // We don't actually need this
-  DeviceTensor<float, 2, true> listDistance(mem, {vecs.getSize(0), 1}, stream);
+  DeviceTensor<float, 2, true> listDistance(
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {vecs.getSize(0), 1});
   // We use this
-  DeviceTensor<int, 2, true> listIds2d(mem, {vecs.getSize(0), 1}, stream);
+  DeviceTensor<int, 2, true> listIds2d(
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {vecs.getSize(0), 1});
   auto listIds = listIds2d.view<1>({vecs.getSize(0)});
 
   quantizer_->query(vecs,
@@ -151,7 +154,8 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
 
   // Calculate the residual for each closest centroid
   DeviceTensor<float, 2, true> residuals(
-    mem, {vecs.getSize(0), vecs.getSize(1)}, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {vecs.getSize(0), vecs.getSize(1)});
 
   if (quantizer_->getUseFloat16()) {
     auto& coarseCentroids = quantizer_->getVectorsFloat16Ref();
@@ -169,9 +173,8 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
     {residuals.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
 
   DeviceTensor<float, 3, true> residualsTranspose(
-    mem,
-    {numSubQuantizers_, residuals.getSize(0), dimPerSubQuantizer_},
-    stream);
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {numSubQuantizers_, residuals.getSize(0), dimPerSubQuantizer_});
 
   runTransposeAny(residualsView, 0, 1, residualsTranspose, stream);
 
@@ -186,9 +189,11 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
   // => (numSubQuantizer x vec x 1)
 
   DeviceTensor<float, 3, true> closestSubQDistance(
-    mem, {numSubQuantizers_, residuals.getSize(0), 1}, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {numSubQuantizers_, residuals.getSize(0), 1});
   DeviceTensor<int, 3, true> closestSubQIndex(
-    mem, {numSubQuantizers_, residuals.getSize(0), 1}, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {numSubQuantizers_, residuals.getSize(0), 1});
 
   for (int subQ = 0; subQ < numSubQuantizers_; ++subQ) {
     auto closestSubQDistanceView = closestSubQDistance[subQ].view();
@@ -217,7 +222,8 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
 
   // Transpose this for easy use
   DeviceTensor<int, 2, true> encodings(
-    mem, {residuals.getSize(0), numSubQuantizers_}, stream);
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {residuals.getSize(0), numSubQuantizers_});
 
   runTransposeAny(closestSubQIndexView, 0, 1, encodings, stream);
 
@@ -337,7 +343,9 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
 
   // We similarly need to actually append the new encoded vectors
   {
-    DeviceTensor<int, 1, true> listOffset(mem, listOffsetHost, stream);
+    DeviceTensor<int, 1, true> listOffset(
+      resources_, makeTempAlloc(AllocType::Other, stream),
+      listOffsetHost);
 
     // This kernel will handle appending each encoded vector + index to
     // the appropriate list
@@ -409,6 +417,8 @@ IVFPQ::addCodeVectorsFromCpu(int listId,
 
 void
 IVFPQ::setPQCentroids_(float* data) {
+  auto stream = resources_->getDefaultStreamCurrentDevice();
+
   size_t pqSize =
     numSubQuantizers_ * numSubQuantizerCodes_ * dimPerSubQuantizer_;
 
@@ -420,23 +430,34 @@ IVFPQ::setPQCentroids_(float* data) {
   HostTensor<float, 3, true> pqHost(
     hostMemory.data(),
     {numSubQuantizers_, numSubQuantizerCodes_, dimPerSubQuantizer_});
-  DeviceTensor<float, 3, true> pqDevice(
-    pqHost,
-    resources_->getDefaultStreamCurrentDevice());
 
   DeviceTensor<float, 3, true> pqDeviceTranspose(
+    resources_,
+    makeDevAlloc(AllocType::Quantizer, stream),
     {numSubQuantizers_, dimPerSubQuantizer_, numSubQuantizerCodes_});
-  runTransposeAny(pqDevice, 1, 2, pqDeviceTranspose,
-                  resources_->getDefaultStreamCurrentDevice());
+
+  {
+    // Only needed for the duration of the transposition
+    DeviceTensor<float, 3, true> pqDevice(
+      resources_,
+      makeTempAlloc(AllocType::Quantizer, stream),
+      pqHost);
+
+    runTransposeAny(pqDevice, 1, 2, pqDeviceTranspose,
+                    stream);
+  }
 
   pqCentroidsInnermostCode_ = std::move(pqDeviceTranspose);
 
   // Also maintain the PQ centroids in the form
   // (sub q)(code id)(sub dim)
   DeviceTensor<float, 3, true> pqCentroidsMiddleCode(
+    resources_,
+    makeDevAlloc(AllocType::Quantizer, stream),
     {numSubQuantizers_, numSubQuantizerCodes_, dimPerSubQuantizer_});
+
   runTransposeAny(pqCentroidsInnermostCode_, 1, 2, pqCentroidsMiddleCode,
-                  resources_->getDefaultStreamCurrentDevice());
+                  stream);
 
   pqCentroidsMiddleCode_ = std::move(pqCentroidsMiddleCode);
 }
@@ -446,6 +467,8 @@ void
 IVFPQ::precomputeCodesT_() {
   FAISS_ASSERT(metric_ == MetricType::METRIC_L2);
 
+  auto stream = resources_->getDefaultStreamCurrentDevice();
+
   //
   //    d = || x - y_C ||^2 + || y_R ||^2 + 2 * (y_C|y_R) - 2 * (x|y_R)
   //        ---------------   ---------------------------       -------
@@ -455,18 +478,6 @@ IVFPQ::precomputeCodesT_() {
   // Terms 1 and 3 are available only at query time. We compute term 2
   // here.
 
-  // Compute ||y_R||^2 by treating
-  // (sub q)(code id)(sub dim) as (sub q * code id)(sub dim)
-  auto pqCentroidsMiddleCodeView =
-    pqCentroidsMiddleCode_.view<2>(
-      {numSubQuantizers_ * numSubQuantizerCodes_, dimPerSubQuantizer_});
-  DeviceTensor<float, 1, true> subQuantizerNorms(
-    {numSubQuantizers_ * numSubQuantizerCodes_});
-
-  runL2Norm(pqCentroidsMiddleCodeView, true,
-            subQuantizerNorms, true,
-            resources_->getDefaultStreamCurrentDevice());
-
   // Compute 2 * (y_C|y_R) via batch matrix multiplication
   // batch size (sub q) x {(centroid id)(sub dim) x (code id)(sub dim)'}
   //         => (sub q) x {(centroid id)(code id)}
@@ -477,30 +488,43 @@ IVFPQ::precomputeCodesT_() {
   // Transpose (centroid id)(sub q)(sub dim) to
   //           (sub q)(centroid id)(sub dim)
   auto& coarseCentroids = quantizer_->template getVectorsRef<CentroidT>();
-  auto centroidView = coarseCentroids.template view<3>(
-    {coarseCentroids.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
-  DeviceTensor<CentroidT, 3, true> centroidsTransposed(
-    {numSubQuantizers_, coarseCentroids.getSize(0), dimPerSubQuantizer_});
-
-  runTransposeAny(centroidView, 0, 1, centroidsTransposed,
-                  resources_->getDefaultStreamCurrentDevice());
 
+  // Create the coarse PQ product
   DeviceTensor<float, 3, true> coarsePQProduct(
+    resources_,
+    makeTempAlloc(AllocType::QuantizerPrecomputedCodes, stream),
     {numSubQuantizers_, coarseCentroids.getSize(0), numSubQuantizerCodes_});
 
-  runIteratedMatrixMult(coarsePQProduct, false,
-                        centroidsTransposed, false,
-                        pqCentroidsMiddleCode_, true,
-                        2.0f, 0.0f,
-                        resources_->getBlasHandleCurrentDevice(),
-                        resources_->getDefaultStreamCurrentDevice());
+  {
+    auto centroidView = coarseCentroids.template view<3>(
+      {coarseCentroids.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
+
+    // This is only needed temporarily
+    DeviceTensor<CentroidT, 3, true> centroidsTransposed(
+      resources_,
+      makeTempAlloc(AllocType::QuantizerPrecomputedCodes, stream),
+      {numSubQuantizers_, coarseCentroids.getSize(0), dimPerSubQuantizer_});
+
+    runTransposeAny(centroidView, 0, 1, centroidsTransposed,
+                    stream);
+
+    runIteratedMatrixMult(coarsePQProduct, false,
+                          centroidsTransposed, false,
+                          pqCentroidsMiddleCode_, true,
+                          2.0f, 0.0f,
+                          resources_->getBlasHandleCurrentDevice(),
+                          stream);
+  }
 
   // Transpose (sub q)(centroid id)(code id) to
   //           (centroid id)(sub q)(code id)
+  // This will become our precomputed code output
   DeviceTensor<float, 3, true> coarsePQProductTransposed(
+    resources_,
+    makeDevAlloc(AllocType::QuantizerPrecomputedCodes, stream),
     {coarseCentroids.getSize(0), numSubQuantizers_, numSubQuantizerCodes_});
   runTransposeAny(coarsePQProduct, 0, 1, coarsePQProductTransposed,
-                  resources_->getDefaultStreamCurrentDevice());
+                  stream);
 
   // View (centroid id)(sub q)(code id) as
   //      (centroid id)(sub q * code id)
@@ -510,16 +534,35 @@ IVFPQ::precomputeCodesT_() {
   // Sum || y_R ||^2 + 2 * (y_C|y_R)
   // i.e., add norms                              (sub q * code id)
   // along columns of inner product  (centroid id)(sub q * code id)
-  runSumAlongColumns(subQuantizerNorms, coarsePQProductTransposedView,
-                     resources_->getDefaultStreamCurrentDevice());
+  {
+    // Compute ||y_R||^2 by treating
+    // (sub q)(code id)(sub dim) as (sub q * code id)(sub dim)
+    auto pqCentroidsMiddleCodeView =
+      pqCentroidsMiddleCode_.view<2>(
+        {numSubQuantizers_ * numSubQuantizerCodes_, dimPerSubQuantizer_});
+    DeviceTensor<float, 1, true> subQuantizerNorms(
+      resources_,
+      makeTempAlloc(AllocType::QuantizerPrecomputedCodes, stream),
+      {numSubQuantizers_ * numSubQuantizerCodes_});
+
+    runL2Norm(pqCentroidsMiddleCodeView, true,
+              subQuantizerNorms, true,
+              stream);
+
+    runSumAlongColumns(subQuantizerNorms, coarsePQProductTransposedView,
+                       stream);
+  }
 
   // We added into the view, so `coarsePQProductTransposed` is now our
   // precomputed term 2.
   if (useFloat16LookupTables_) {
     precomputedCodeHalf_ =
-      convertTensor<float, half, 3>(resources_,
-                                    resources_->getDefaultStreamCurrentDevice(),
-                                    coarsePQProductTransposed);
+      DeviceTensor<half, 3, true>(
+        resources_,
+        makeDevAlloc(AllocType::QuantizerPrecomputedCodes, stream),
+        {coarseCentroids.getSize(0), numSubQuantizers_, numSubQuantizerCodes_});
+
+    convertTensor(stream, coarsePQProductTransposed, precomputedCodeHalf_);
   } else {
     precomputedCode_ = std::move(coarsePQProductTransposed);
   }
@@ -544,7 +587,6 @@ IVFPQ::query(Tensor<float, 2, true>& queries,
   FAISS_ASSERT(nprobe <= GPU_MAX_SELECTION_K);
   FAISS_ASSERT(k <= GPU_MAX_SELECTION_K);
 
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
   auto stream = resources_->getDefaultStreamCurrentDevice();
   nprobe = std::min(nprobe, quantizer_->getSize());
 
@@ -554,9 +596,13 @@ IVFPQ::query(Tensor<float, 2, true>& queries,
 
   // Reserve space for the closest coarse centroids
   DeviceTensor<float, 2, true>
-    coarseDistances(mem, {queries.getSize(0), nprobe}, stream);
+    coarseDistances(
+      resources_, makeTempAlloc(AllocType::Other, stream),
+      {queries.getSize(0), nprobe});
   DeviceTensor<int, 2, true>
-    coarseIndices(mem, {queries.getSize(0), nprobe}, stream);
+    coarseIndices(
+      resources_, makeTempAlloc(AllocType::Other, stream),
+      {queries.getSize(0), nprobe});
 
   // Find the `nprobe` closest coarse centroids; we can use int
   // indices both internally and externally
@@ -628,7 +674,6 @@ IVFPQ::runPQPrecomputedCodes_(
   Tensor<long, 2, true>& outIndices) {
   FAISS_ASSERT(metric_ == MetricType::METRIC_L2);
 
-  auto& mem = resources_->getMemoryManagerCurrentDevice();
   auto stream = resources_->getDefaultStreamCurrentDevice();
 
   // Compute precomputed code term 3, - 2 * (x|y_R)
@@ -636,9 +681,8 @@ IVFPQ::runPQPrecomputedCodes_(
   // {sub q} x {(query id)(sub dim) * (code id)(sub dim)'} =>
   // {sub q} x {(query id)(code id)}
   DeviceTensor<float, 3, true> term3Transposed(
-    mem,
-    {queries.getSize(0), numSubQuantizers_, numSubQuantizerCodes_},
-    stream);
+    resources_, makeTempAlloc(AllocType::Other, stream),
+    {queries.getSize(0), numSubQuantizers_, numSubQuantizerCodes_});
 
   // These allocations within are only temporary, so release them when
   // we're done to maximize free space
@@ -646,15 +690,13 @@ IVFPQ::runPQPrecomputedCodes_(
     auto querySubQuantizerView = queries.view<3>(
       {queries.getSize(0), numSubQuantizers_, dimPerSubQuantizer_});
     DeviceTensor<float, 3, true> queriesTransposed(
-      mem,
-      {numSubQuantizers_, queries.getSize(0), dimPerSubQuantizer_},
-      stream);
+      resources_, makeTempAlloc(AllocType::Other, stream),
+      {numSubQuantizers_, queries.getSize(0), dimPerSubQuantizer_});
     runTransposeAny(querySubQuantizerView, 0, 1, queriesTransposed, stream);
 
     DeviceTensor<float, 3, true> term3(
-      mem,
-      {numSubQuantizers_, queries.getSize(0), numSubQuantizerCodes_},
-      stream);
+      resources_, makeTempAlloc(AllocType::Other, stream),
+      {numSubQuantizers_, queries.getSize(0), numSubQuantizerCodes_});
 
     runIteratedMatrixMult(term3, false,
                           queriesTransposed, false,
@@ -672,7 +714,8 @@ IVFPQ::runPQPrecomputedCodes_(
 
   if (useFloat16LookupTables_) {
     term3Half =
-      convertTensor<float, half, 3>(resources_, stream, term3Transposed);
+      convertTensorTemporary<float, half, 3>(
+        resources_, stream, term3Transposed);
 
     term2 = NoTypeTensor<3, true>(precomputedCodeHalf_);
     term3 = NoTypeTensor<3, true>(term3Half);

gpu/impl/IVFUtils.cu

@@ -38,7 +38,8 @@ getResultLengths(Tensor<int, 2, true> topQueryToCentroid,
   length[queryId][listId] = (centroidId != -1) ? listLengths[centroidId] : 0;
 }
 
-void runCalcListOffsets(Tensor<int, 2, true>& topQueryToCentroid,
+void runCalcListOffsets(GpuResources* res,
+                        Tensor<int, 2, true>& topQueryToCentroid,
                         thrust::device_vector<int>& listLengths,
                         Tensor<int, 2, true>& prefixSumOffsets,
                         Tensor<char, 1, true>& thrustMem,
@@ -64,8 +65,10 @@ void runCalcListOffsets(Tensor<int, 2, true>& topQueryToCentroid,
   // Prefix sum of the indices, so we know where the intermediate
   // results should be maintained
   // Thrust wants a place for its temporary allocations, so provide
-  // one, so it won't call cudaMalloc/Free
-  GpuResourcesThrustAllocator alloc(thrustMem.data(),
+  // one, so it won't call cudaMalloc/Free if we size it sufficiently
+  GpuResourcesThrustAllocator alloc(res,
+                                    stream,
+                                    thrustMem.data(),
                                     thrustMem.getSizeInBytes());
 
   thrust::inclusive_scan(thrust::cuda::par(alloc).on(stream),

gpu/impl/IVFUtils.cuh

@@ -16,9 +16,12 @@
 // post-processing and k-selecting the results
 namespace faiss { namespace gpu {
 
+class GpuResources;
+
 /// Function for multi-pass scanning that collects the length of
 /// intermediate results for all (query, probe) pair
-void runCalcListOffsets(Tensor<int, 2, true>& topQueryToCentroid,
+void runCalcListOffsets(GpuResources* res,
+                        Tensor<int, 2, true>& topQueryToCentroid,
                         thrust::device_vector<int>& listLengths,
                         Tensor<int, 2, true>& prefixSumOffsets,
                         Tensor<char, 1, true>& thrustMem,

gpu/impl/L2Norm.cu

@@ -241,7 +241,7 @@ void runL2Norm(Tensor<T, 2, true, IndexType>& input,
 
       auto dim = inputV.getSize(1);
       bool normLoop = dim > maxThreads;
-      auto numThreads = min(dim, maxThreads);
+      auto numThreads = std::min(dim, maxThreads);
 
       auto grid = dim3(utils::divUp(inputV.getSize(0), rowTileSize));
       auto block = dim3(numThreads);
@@ -254,7 +254,7 @@ void runL2Norm(Tensor<T, 2, true, IndexType>& input,
 
       auto dim = input.getSize(1);
       bool normLoop = dim > maxThreads;
-      auto numThreads = min(dim, maxThreads);
+      auto numThreads = std::min(dim, maxThreads);
 
       auto grid = dim3(utils::divUp(input.getSize(0), rowTileSize));
       auto block = dim3(numThreads);

gpu/impl/PQCodeDistances-inl.cuh

@@ -320,18 +320,17 @@ runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
                      Tensor<int, 2, true>& topQueryToCentroid,
                      NoTypeTensor<4, true>& outCodeDistances,
                      bool useFloat16Lookup,
-                     DeviceMemory& mem,
+                     GpuResources* res,
                      cublasHandle_t handle,
                      cudaStream_t stream) {
   // Calculate (q - c) residual vector
   // (sub q)(query id)(centroid id)(sub dim)
   DeviceTensor<float, 4, true> residual(
-    mem,
+    res, makeTempAlloc(AllocType::Other, stream),
     {pqCentroids.getSize(0),
-        topQueryToCentroid.getSize(0),
-        topQueryToCentroid.getSize(1),
-        pqCentroids.getSize(1)},
-    stream);
+     topQueryToCentroid.getSize(0),
+     topQueryToCentroid.getSize(1),
+     pqCentroids.getSize(1)});
 
   runResidualVector(pqCentroids, queries,
                     coarseCentroids, topQueryToCentroid,
@@ -339,11 +338,10 @@ runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
 
   // Calculate ||q - c||^2
   DeviceTensor<float, 1, true> residualNorms(
-    mem,
+    res, makeTempAlloc(AllocType::Other, stream),
     {pqCentroids.getSize(0) *
-        topQueryToCentroid.getSize(0) *
-        topQueryToCentroid.getSize(1)},
-    stream);
+     topQueryToCentroid.getSize(0) *
+     topQueryToCentroid.getSize(1)});
 
   auto residualView2 = residual.view<2>(
     {pqCentroids.getSize(0) *
@@ -362,11 +360,10 @@ runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
         pqCentroids.getSize(1)});
 
   DeviceTensor<float, 3, true> residualDistance(
-    mem,
+    res, makeTempAlloc(AllocType::Other, stream),
     {pqCentroids.getSize(0),
-        topQueryToCentroid.getSize(0) * topQueryToCentroid.getSize(1),
-        pqCentroids.getSize(2)},
-    stream);
+     topQueryToCentroid.getSize(0) * topQueryToCentroid.getSize(1),
+     pqCentroids.getSize(2)});
 
   runIteratedMatrixMult(residualDistance, false,
                         residualView3, false,
@@ -389,11 +386,11 @@ runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
 
   if (useFloat16Lookup) {
     outCodeDistancesFloatMem = DeviceTensor<float, 4, true>(
-      mem, {outCodeDistances.getSize(0),
-          outCodeDistances.getSize(1),
-          outCodeDistances.getSize(2),
-          outCodeDistances.getSize(3)},
-      stream);
+      res, makeTempAlloc(AllocType::Other, stream),
+      {outCodeDistances.getSize(0),
+       outCodeDistances.getSize(1),
+       outCodeDistances.getSize(2),
+       outCodeDistances.getSize(3)});
 
     outCodeDistancesF = outCodeDistancesFloatMem;
   } else {
@@ -413,9 +410,8 @@ runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
   // (sub q)(sub dim)(code) is pqCentroids
   // transpose to (sub q)(code)(sub dim)
   DeviceTensor<float, 3, true> pqCentroidsTranspose(
-    mem,
-    {pqCentroids.getSize(0), pqCentroids.getSize(2), pqCentroids.getSize(1)},
-    stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {pqCentroids.getSize(0), pqCentroids.getSize(2), pqCentroids.getSize(1)});
 
   runTransposeAny(pqCentroids, 1, 2, pqCentroidsTranspose, stream);
 
@@ -424,9 +420,8 @@ runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
         pqCentroids.getSize(1)});
 
   DeviceTensor<float, 1, true> pqCentroidsNorm(
-    mem,
-    {pqCentroids.getSize(0) * pqCentroids.getSize(2)},
-    stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {pqCentroids.getSize(0) * pqCentroids.getSize(2)});
 
   runL2Norm(pqCentroidsTransposeView, true, pqCentroidsNorm, true, stream);
 

gpu/impl/PQCodeDistances.cuh

@@ -37,7 +37,7 @@ void runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
                           Tensor<int, 2, true>& topQueryToCentroid,
                           NoTypeTensor<4, true>& outCodeDistances,
                           bool useFloat16Lookup,
-                          DeviceMemory& mem,
+                          GpuResources* res,
                           cublasHandle_t handle,
                           cudaStream_t stream);
 

gpu/impl/PQScanMultiPassNoPrecomputed-inl.cuh

@@ -219,7 +219,8 @@ pqScanNoPrecomputedMultiPass(Tensor<float, 2, true> queries,
 
 template <typename CentroidT>
 void
-runMultiPassTile(Tensor<float, 2, true>& queries,
+runMultiPassTile(GpuResources* res,
+                 Tensor<float, 2, true>& queries,
                  Tensor<CentroidT, 2, true>& centroids,
                  Tensor<float, 3, true>& pqCentroidsInnermostCode,
                  NoTypeTensor<4, true>& codeDistances,
@@ -250,7 +251,7 @@ runMultiPassTile(Tensor<float, 2, true>& queries,
 
   // Calculate offset lengths, so we know where to write out
   // intermediate results
-  runCalcListOffsets(topQueryToCentroid, listLengths, prefixSumOffsets,
+  runCalcListOffsets(res, topQueryToCentroid, listLengths, prefixSumOffsets,
                      thrustMem, stream);
 
   // Calculate residual code distances, since this is without
@@ -419,21 +420,22 @@ runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,
 
   int nprobe = topQueryToCentroid.getSize(1);
 
-  auto& mem = res->getMemoryManagerCurrentDevice();
   auto stream = res->getDefaultStreamCurrentDevice();
 
   // Make a reservation for Thrust to do its dirty work (global memory
   // cross-block reduction space); hopefully this is large enough.
   DeviceTensor<char, 1, true> thrustMem1(
-    mem, {kThrustMemSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {kThrustMemSize});
   DeviceTensor<char, 1, true> thrustMem2(
-    mem, {kThrustMemSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {kThrustMemSize});
   DeviceTensor<char, 1, true>* thrustMem[2] =
     {&thrustMem1, &thrustMem2};
 
   // How much temporary storage is available?
   // If possible, we'd like to fit within the space available.
-  size_t sizeAvailable = mem.getSizeAvailable();
+  size_t sizeAvailable = res->getTempMemoryAvailableCurrentDevice();
 
   // We run two passes of heap selection
   // This is the size of the first-level heap passes
@@ -469,9 +471,11 @@ runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,
   // Make sure there is space prior to the start which will be 0, and
   // will handle the boundary condition without branches
   DeviceTensor<int, 1, true> prefixSumOffsetSpace1(
-    mem, {queryTileSize * nprobe + 1}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe + 1});
   DeviceTensor<int, 1, true> prefixSumOffsetSpace2(
-    mem, {queryTileSize * nprobe + 1}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe + 1});
 
   DeviceTensor<int, 2, true> prefixSumOffsets1(
     prefixSumOffsetSpace1[1].data(),
@@ -500,14 +504,16 @@ runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,
     codeDistanceTypeSize;
 
   DeviceTensor<char, 1, true> codeDistances1Mem(
-    mem, {totalCodeDistancesSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {totalCodeDistancesSize});
   NoTypeTensor<4, true> codeDistances1(
     codeDistances1Mem.data(),
     codeDistanceTypeSize,
     {queryTileSize, nprobe, numSubQuantizers, numSubQuantizerCodes});
 
   DeviceTensor<char, 1, true> codeDistances2Mem(
-    mem, {totalCodeDistancesSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {totalCodeDistancesSize});
   NoTypeTensor<4, true> codeDistances2(
     codeDistances2Mem.data(),
     codeDistanceTypeSize,
@@ -517,23 +523,29 @@ runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,
     {&codeDistances1, &codeDistances2};
 
   DeviceTensor<float, 1, true> allDistances1(
-    mem, {queryTileSize * nprobe * maxListLength}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe * maxListLength});
   DeviceTensor<float, 1, true> allDistances2(
-    mem, {queryTileSize * nprobe * maxListLength}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe * maxListLength});
   DeviceTensor<float, 1, true>* allDistances[2] =
     {&allDistances1, &allDistances2};
 
   DeviceTensor<float, 3, true> heapDistances1(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<float, 3, true> heapDistances2(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<float, 3, true>* heapDistances[2] =
     {&heapDistances1, &heapDistances2};
 
   DeviceTensor<int, 3, true> heapIndices1(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<int, 3, true> heapIndices2(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<int, 3, true>* heapIndices[2] =
     {&heapIndices1, &heapIndices2};
 
@@ -566,7 +578,8 @@ runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,
     auto outIndicesView =
       outIndices.narrowOutermost(query, numQueriesInTile);
 
-    runMultiPassTile(queryView,
+    runMultiPassTile(res,
+                     queryView,
                      centroids,
                      pqCentroidsInnermostCode,
                      codeDistancesView,

gpu/impl/PQScanMultiPassNoPrecomputed.cuh

@@ -19,7 +19,7 @@ class GpuResources;
 
 /// For no precomputed codes, is this a supported number of dimensions
 /// per subquantizer?
-bool isSupportedNoPrecomputedSubDimSize(int dims);
+inline bool isSupportedNoPrecomputedSubDimSize(int dims);
 
 template <typename CentroidT>
 void runPQScanMultiPassNoPrecomputed(Tensor<float, 2, true>& queries,

gpu/impl/PQScanMultiPassPrecomputed.cu

@@ -214,7 +214,8 @@ pqScanPrecomputedMultiPass(Tensor<float, 2, true> queries,
 }
 
 void
-runMultiPassTile(Tensor<float, 2, true>& queries,
+runMultiPassTile(GpuResources* res,
+                 Tensor<float, 2, true>& queries,
                  Tensor<float, 2, true>& precompTerm1,
                  NoTypeTensor<3, true>& precompTerm2,
                  NoTypeTensor<3, true>& precompTerm3,
@@ -238,7 +239,7 @@ runMultiPassTile(Tensor<float, 2, true>& queries,
                  cudaStream_t stream) {
   // Calculate offset lengths, so we know where to write out
   // intermediate results
-  runCalcListOffsets(topQueryToCentroid, listLengths, prefixSumOffsets,
+  runCalcListOffsets(res, topQueryToCentroid, listLengths, prefixSumOffsets,
                      thrustMem, stream);
 
   // Convert all codes to a distance, and write out (distance,
@@ -398,21 +399,22 @@ void runPQScanMultiPassPrecomputed(Tensor<float, 2, true>& queries,
 
   int nprobe = topQueryToCentroid.getSize(1);
 
-  auto& mem = res->getMemoryManagerCurrentDevice();
   auto stream = res->getDefaultStreamCurrentDevice();
 
   // Make a reservation for Thrust to do its dirty work (global memory
   // cross-block reduction space); hopefully this is large enough.
   DeviceTensor<char, 1, true> thrustMem1(
-    mem, {kThrustMemSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {kThrustMemSize});
   DeviceTensor<char, 1, true> thrustMem2(
-    mem, {kThrustMemSize}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {kThrustMemSize});
   DeviceTensor<char, 1, true>* thrustMem[2] =
     {&thrustMem1, &thrustMem2};
 
   // How much temporary storage is available?
   // If possible, we'd like to fit within the space available.
-  size_t sizeAvailable = mem.getSizeAvailable();
+  size_t sizeAvailable = res->getTempMemoryAvailableCurrentDevice();
 
   // We run two passes of heap selection
   // This is the size of the first-level heap passes
@@ -446,9 +448,11 @@ void runPQScanMultiPassPrecomputed(Tensor<float, 2, true>& queries,
  // Make sure there is space prior to the start which will be 0, and
   // will handle the boundary condition without branches
   DeviceTensor<int, 1, true> prefixSumOffsetSpace1(
-    mem, {queryTileSize * nprobe + 1}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe + 1});
   DeviceTensor<int, 1, true> prefixSumOffsetSpace2(
-    mem, {queryTileSize * nprobe + 1}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe + 1});
 
   DeviceTensor<int, 2, true> prefixSumOffsets1(
     prefixSumOffsetSpace1[1].data(),
@@ -471,23 +475,29 @@ void runPQScanMultiPassPrecomputed(Tensor<float, 2, true>& queries,
                               stream));
 
   DeviceTensor<float, 1, true> allDistances1(
-    mem, {queryTileSize * nprobe * maxListLength}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe * maxListLength});
   DeviceTensor<float, 1, true> allDistances2(
-    mem, {queryTileSize * nprobe * maxListLength}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize * nprobe * maxListLength});
   DeviceTensor<float, 1, true>* allDistances[2] =
     {&allDistances1, &allDistances2};
 
   DeviceTensor<float, 3, true> heapDistances1(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<float, 3, true> heapDistances2(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<float, 3, true>* heapDistances[2] =
     {&heapDistances1, &heapDistances2};
 
   DeviceTensor<int, 3, true> heapIndices1(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<int, 3, true> heapIndices2(
-    mem, {queryTileSize, pass2Chunks, k}, stream);
+    res, makeTempAlloc(AllocType::Other, stream),
+    {queryTileSize, pass2Chunks, k});
   DeviceTensor<int, 3, true>* heapIndices[2] =
     {&heapIndices1, &heapIndices2};
 
@@ -522,7 +532,8 @@ void runPQScanMultiPassPrecomputed(Tensor<float, 2, true>& queries,
     auto outIndicesView =
       outIndices.narrowOutermost(query, numQueriesInTile);
 
-    runMultiPassTile(queryView,
+    runMultiPassTile(res,
+                     queryView,
                      term1View,
                      precompTerm2,
                      term3View,

gpu/perf/IndexWrapper-inl.h

@@ -13,7 +13,7 @@ namespace faiss { namespace gpu {
 template <typename GpuIndex>
 IndexWrapper<GpuIndex>::IndexWrapper(
   int numGpus,
-  std::function<std::unique_ptr<GpuIndex>(GpuResources*, int)> init) {
+  std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init) {
   FAISS_ASSERT(numGpus <= faiss::gpu::getNumDevices());
   for (int i = 0; i < numGpus; ++i) {
     auto res = std::unique_ptr<faiss::gpu::StandardGpuResources>(

gpu/perf/IndexWrapper.h

@@ -27,7 +27,7 @@ struct IndexWrapper {
 
   IndexWrapper(
     int numGpus,
-    std::function<std::unique_ptr<GpuIndex>(GpuResources*, int)> init);
+    std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init);
   faiss::Index* getIndex();
 
   void runOnIndices(std::function<void(GpuIndex*)> f);

gpu/perf/PerfClustering.cpp

@@ -53,7 +53,7 @@ int main(int argc, char** argv) {
   printf("transposed storage %s\n", FLAGS_transposed ? "enabled" : "disabled");
   printf("verbose %s\n", FLAGS_verbose ? "enabled" : "disabled");
 
-  auto initFn = [](faiss::gpu::GpuResources* res, int dev) ->
+  auto initFn = [](faiss::gpu::GpuResourcesProvider* res, int dev) ->
     std::unique_ptr<faiss::gpu::GpuIndexFlat> {
     if (FLAGS_pinned_mem >= 0) {
       ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(

gpu/perf/PerfFlat.cu

@@ -68,7 +68,7 @@ int main(int argc, char** argv) {
   // Convert to GPU index
   printf("Copying index to %d GPU(s)...\n", FLAGS_num_gpus);
 
-  auto initFn = [&index](faiss::gpu::GpuResources* res, int dev) ->
+  auto initFn = [&index](faiss::gpu::GpuResourcesProvider* res, int dev) ->
     std::unique_ptr<faiss::gpu::GpuIndexFlat> {
     ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(
       FLAGS_pinned_mem);

gpu/perf/PerfIVFFlat.cu

@@ -63,7 +63,7 @@ int main(int argc, char** argv) {
   // Convert to GPU index
   printf("Copying index to %d GPU(s)...\n", FLAGS_num_gpus);
 
-  auto initFn = [&index](faiss::gpu::GpuResources* res, int dev) ->
+  auto initFn = [&index](faiss::gpu::GpuResourcesProvider* res, int dev) ->
     std::unique_ptr<faiss::gpu::GpuIndexIVFFlat> {
     GpuIndexIVFFlatConfig config;
     config.device = dev;
@@ -139,8 +139,6 @@ int main(int argc, char** argv) {
                "", true, FLAGS_diff, false);
 
   CUDA_VERIFY(cudaDeviceSynchronize());
-  // printf("\ncudaMalloc usage %zd\n",
-  //        resources.getMemoryManager().getHighWaterCudaMalloc());
 
   return 0;
 }

gpu/perf/PerfIVFPQ.cu

@@ -76,7 +76,7 @@ int main(int argc, char** argv) {
   auto useFloat16Lookup = FLAGS_float16_lookup;
 
   auto initFn = [precomp, indicesOpt, useFloat16Lookup, &index]
-    (faiss::gpu::GpuResources* res, int dev) ->
+    (faiss::gpu::GpuResourcesProvider* res, int dev) ->
     std::unique_ptr<faiss::gpu::GpuIndexIVFPQ> {
 
     faiss::gpu::GpuIndexIVFPQConfig config;
@@ -150,8 +150,6 @@ int main(int argc, char** argv) {
                "", true, FLAGS_diff, false);
 
   CUDA_VERIFY(cudaDeviceSynchronize());
-  // printf("\ncudaMalloc usage %zd\n",
-  //        resources.getMemoryManager().getHighWaterCudaMalloc());
 
   return 0;
 }

gpu/perf/PerfSelect.cu

@@ -6,6 +6,7 @@
  */
 
 
+#include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
 #include <faiss/gpu/utils/DeviceUtils.h>
 #include <faiss/gpu/utils/BlockSelectKernel.cuh>
@@ -29,10 +30,12 @@ DEFINE_int32(iter, 5, "iterations to run");
 DEFINE_bool(k_powers, false, "test k powers of 2 from 1 -> max k");
 
 int main(int argc, char** argv) {
+  using namespace faiss::gpu;
+
   gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-  std::vector<float> v = faiss::gpu::randVecs(FLAGS_rows, FLAGS_cols);
-  faiss::gpu::HostTensor<float, 2, true> hostVal({FLAGS_rows, FLAGS_cols});
+  std::vector<float> v = randVecs(FLAGS_rows, FLAGS_cols);
+  HostTensor<float, 2, true> hostVal({FLAGS_rows, FLAGS_cols});
 
   for (int r = 0; r < FLAGS_rows; ++r) {
     for (int c = 0; c < FLAGS_cols; ++c) {
@@ -40,8 +43,14 @@ int main(int argc, char** argv) {
     }
   }
 
+  StandardGpuResources res;
+  res.noTempMemory();
+
+  auto resUse = res.getResources();
+
   // Select top-k on GPU
-  faiss::gpu::DeviceTensor<float, 2, true> gpuVal(hostVal, 0);
+  DeviceTensor<float, 2, true>
+    gpuVal(resUse.get(), makeDevAlloc(AllocType::Other, 0), hostVal);
 
   int startK = FLAGS_k;
   int limitK = FLAGS_k;
@@ -52,16 +61,16 @@ int main(int argc, char** argv) {
   }
 
   for (int k = startK; k <= limitK; k *= 2) {
-    faiss::gpu::DeviceTensor<float, 2, true> gpuOutVal({FLAGS_rows, k});
-    faiss::gpu::DeviceTensor<int, 2, true> gpuOutInd({FLAGS_rows, k});
+    DeviceTensor<float, 2, true>
+      gpuOutVal(resUse.get(), makeDevAlloc(AllocType::Other, 0), {FLAGS_rows, k});
+    DeviceTensor<int, 2, true>
+      gpuOutInd(resUse.get(), makeDevAlloc(AllocType::Other, 0), {FLAGS_rows, k});
 
     for (int i = 0; i < FLAGS_iter; ++i) {
       if (FLAGS_warp) {
-        faiss::gpu::runWarpSelect(gpuVal, gpuOutVal, gpuOutInd,
-                                  FLAGS_dir, k, 0);
+        runWarpSelect(gpuVal, gpuOutVal, gpuOutInd, FLAGS_dir, k, 0);
       } else {
-        faiss::gpu::runBlockSelect(gpuVal, gpuOutVal, gpuOutInd,
-                                   FLAGS_dir, k, 0);
+        runBlockSelect(gpuVal, gpuOutVal, gpuOutInd, FLAGS_dir, k, 0);
       }
     }
   }

gpu/perf/slow.py

@@ -1,9 +1,10 @@
+#! /usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python3
 # this is a slow computation to test whether ctrl-C handling works
 import faiss
 import numpy as np

gpu/test/TestGpuDistance.cu

@@ -21,19 +21,21 @@ void testTransposition(bool colMajorVecs,
                        bool colMajorQueries,
                        faiss::MetricType metric,
                        float metricArg = 0) {
-  int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
+  using namespace faiss::gpu;
 
-  faiss::gpu::StandardGpuResources res;
+  int device = randVal(0, getNumDevices() - 1);
+
+  StandardGpuResources res;
   res.noTempMemory();
 
-  int dim = faiss::gpu::randVal(20, 150);
-  int numVecs = faiss::gpu::randVal(10, 30000);
-  int numQuery = faiss::gpu::randVal(1, 1024);
-  int k = std::min(numVecs, faiss::gpu::randVal(20, 70));
+  int dim = randVal(20, 150);
+  int numVecs = randVal(10, 30000);
+  int numQuery = randVal(1, 1024);
+  int k = std::min(numVecs, randVal(20, 70));
 
   // Input data for CPU
-  std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
-  std::vector<float> queries = faiss::gpu::randVecs(numQuery, dim);
+  std::vector<float> vecs = randVecs(numVecs, dim);
+  std::vector<float> queries = randVecs(numQuery, dim);
 
   if (metric == faiss::MetricType::METRIC_JensenShannon) {
     // make values positive
@@ -64,26 +66,30 @@ void testTransposition(bool colMajorVecs,
                     cpuDistance.data(), cpuIndices.data());
 
   // The transpose and distance code assumes the desired device is already set
-  faiss::gpu::DeviceScope scope(device);
+  DeviceScope scope(device);
   auto stream = res.getDefaultStream(device);
 
   // Copy input data to GPU, and pre-transpose both vectors and queries for
   // passing
-  auto gpuVecs = faiss::gpu::toDevice<float, 2>(
-    nullptr, device, vecs.data(), stream, {numVecs, dim});
-  auto gpuQueries = faiss::gpu::toDevice<float, 2>(
-    nullptr, device, queries.data(), stream, {numQuery, dim});
+  auto gpuVecs = toDeviceNonTemporary<float, 2>(
+    res.getResources().get(), device, vecs.data(), stream, {numVecs, dim});
+  auto gpuQueries = toDeviceNonTemporary<float, 2>(
+    res.getResources().get(), device, queries.data(), stream, {numQuery, dim});
 
-  faiss::gpu::DeviceTensor<float, 2, true> vecsT({dim, numVecs});
-  faiss::gpu::runTransposeAny(gpuVecs, 0, 1, vecsT, stream);
+  DeviceTensor<float, 2, true>
+    vecsT(res.getResources().get(),
+          makeDevAlloc(AllocType::Other, stream), {dim, numVecs});
+  runTransposeAny(gpuVecs, 0, 1, vecsT, stream);
 
-  faiss::gpu::DeviceTensor<float, 2, true> queriesT({dim, numQuery});
-  faiss::gpu::runTransposeAny(gpuQueries, 0, 1, queriesT, stream);
+  DeviceTensor<float, 2, true>
+    queriesT(res.getResources().get(),
+             makeDevAlloc(AllocType::Other, stream), {dim, numQuery});
+  runTransposeAny(gpuQueries, 0, 1, queriesT, stream);
 
   std::vector<float> gpuDistance(numQuery * k, 0);
   std::vector<faiss::Index::idx_t> gpuIndices(numQuery * k, -1);
 
-  faiss::gpu::GpuDistanceParams args;
+  GpuDistanceParams args;
   args.metric = metric;
   args.metricArg = metricArg;
   args.k = k;
@@ -97,21 +103,21 @@ void testTransposition(bool colMajorVecs,
   args.outDistances = gpuDistance.data();
   args.outIndices = gpuIndices.data();
 
-  faiss::gpu::bfKnn(&res, args);
+  bfKnn(&res, args);
 
   std::stringstream str;
   str << "metric " << metric
       << " colMajorVecs " << colMajorVecs
       << " colMajorQueries " << colMajorQueries;
 
-  faiss::gpu::compareLists(cpuDistance.data(),
-                           cpuIndices.data(),
-                           gpuDistance.data(),
-                           gpuIndices.data(),
-                           numQuery, k,
-                           str.str(),
-                           false, false, true,
-                           6e-3f, 0.1f, 0.015f);
+  compareLists(cpuDistance.data(),
+               cpuIndices.data(),
+               gpuDistance.data(),
+               gpuIndices.data(),
+               numQuery, k,
+               str.str(),
+               false, false, true,
+               6e-3f, 0.1f, 0.015f);
 }
 
 // Test different memory layouts for brute-force k-NN

gpu/test/TestGpuIndexIVFPQ.cpp

@@ -131,8 +131,8 @@ TEST(TestGpuIndexIVFPQ, Query_L2) {
     cpuIndex.train(opt.numTrain, trainVecs.data());
     cpuIndex.add(opt.numAdd, addVecs.data());
 
+    // Use the default temporary memory management to test the memory manager
     faiss::gpu::StandardGpuResources res;
-    res.noTempMemory();
 
     faiss::gpu::GpuIndexIVFPQConfig config;
     config.device = opt.device;
@@ -167,8 +167,8 @@ TEST(TestGpuIndexIVFPQ, Query_IP) {
     cpuIndex.train(opt.numTrain, trainVecs.data());
     cpuIndex.add(opt.numAdd, addVecs.data());
 
+    // Use the default temporary memory management to test the memory manager
     faiss::gpu::StandardGpuResources res;
-    res.noTempMemory();
 
     faiss::gpu::GpuIndexIVFPQConfig config;
     config.device = opt.device;
@@ -199,8 +199,8 @@ TEST(TestGpuIndexIVFPQ, Float16Coarse) {
   cpuIndex.nprobe = opt.nprobe;
   cpuIndex.train(opt.numTrain, trainVecs.data());
 
+  // Use the default temporary memory management to test the memory manager
   faiss::gpu::StandardGpuResources res;
-  res.noTempMemory();
 
   faiss::gpu::GpuIndexIVFPQConfig config;
   config.device = opt.device;
@@ -235,8 +235,8 @@ TEST(TestGpuIndexIVFPQ, Add_L2) {
     cpuIndex.nprobe = opt.nprobe;
     cpuIndex.train(opt.numTrain, trainVecs.data());
 
+    // Use the default temporary memory management to test the memory manager
     faiss::gpu::StandardGpuResources res;
-    res.noTempMemory();
 
     faiss::gpu::GpuIndexIVFPQConfig config;
     config.device = opt.device;
@@ -272,8 +272,8 @@ TEST(TestGpuIndexIVFPQ, Add_IP) {
     cpuIndex.nprobe = opt.nprobe;
     cpuIndex.train(opt.numTrain, trainVecs.data());
 
+    // Use the default temporary memory management to test the memory manager
     faiss::gpu::StandardGpuResources res;
-    res.noTempMemory();
 
     faiss::gpu::GpuIndexIVFPQConfig config;
     config.device = opt.device;
@@ -300,8 +300,8 @@ TEST(TestGpuIndexIVFPQ, CopyTo) {
   std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
   std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
 
+  // Use the default temporary memory management to test the memory manager
   faiss::gpu::StandardGpuResources res;
-  res.noTempMemory();
 
   faiss::gpu::GpuIndexIVFPQConfig config;
   config.device = opt.device;
@@ -358,9 +358,8 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
   cpuIndex.train(opt.numTrain, trainVecs.data());
   cpuIndex.add(opt.numAdd, addVecs.data());
 
-  // Use garbage values to see if we overwrite them
+  // Use the default temporary memory management to test the memory manager
   faiss::gpu::StandardGpuResources res;
-  res.noTempMemory();
 
   faiss::gpu::GpuIndexIVFPQConfig config;
   config.device = opt.device;
@@ -368,6 +367,7 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
   config.indicesOptions = opt.indicesOpt;
   config.useFloat16LookupTables = opt.useFloat16;
 
+  // Use garbage values to see if we overwrite them
   faiss::gpu::GpuIndexIVFPQ
     gpuIndex(&res, 1, 1, 1, 1, faiss::METRIC_L2, config);
   gpuIndex.setNumProbes(1);
@@ -401,8 +401,8 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
   std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
   std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
 
+  // Use the default temporary memory management to test the memory manager
   faiss::gpu::StandardGpuResources res;
-  res.noTempMemory();
 
   faiss::gpu::GpuIndexIVFPQConfig config;
   config.device = opt.device;
@@ -447,8 +447,8 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
 TEST(TestGpuIndexIVFPQ, AddNaN) {
   Options opt;
 
+  // Use the default temporary memory management to test the memory manager
   faiss::gpu::StandardGpuResources res;
-  res.noTempMemory();
 
   faiss::gpu::GpuIndexIVFPQConfig config;
   config.device = opt.device;

gpu/test/TestGpuSelect.cu

@@ -6,6 +6,7 @@
  */
 
 
+#include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/test/TestUtils.h>
 #include <faiss/gpu/utils/BlockSelectKernel.cuh>
 #include <faiss/gpu/utils/DeviceDefs.cuh>
@@ -20,8 +21,12 @@
 #include <vector>
 
 void testForSize(int rows, int cols, int k, bool dir, bool warp) {
-  std::vector<float> v = faiss::gpu::randVecs(rows, cols);
-  faiss::gpu::HostTensor<float, 2, true> hostVal({rows, cols});
+  using namespace faiss::gpu;
+
+  StandardGpuResources res;
+
+  std::vector<float> v = randVecs(rows, cols);
+  HostTensor<float, 2, true> hostVal({rows, cols});
 
   for (int r = 0; r < rows; ++r) {
     for (int c = 0; c < cols; ++c) {
@@ -52,19 +57,25 @@ void testForSize(int rows, int cols, int k, bool dir, bool warp) {
   }
 
   // Select top-k on GPU
-  faiss::gpu::DeviceTensor<float, 2, true> gpuVal(hostVal, 0);
-  faiss::gpu::DeviceTensor<float, 2, true> gpuOutVal({rows, k});
-  faiss::gpu::DeviceTensor<int, 2, true> gpuOutInd({rows, k});
+  DeviceTensor<float, 2, true>
+    gpuVal(res.getResources().get(),
+           makeDevAlloc(AllocType::Other, 0), hostVal);
+  DeviceTensor<float, 2, true>
+    gpuOutVal(res.getResources().get(),
+              makeDevAlloc(AllocType::Other, 0), {rows, k});
+  DeviceTensor<int, 2, true>
+    gpuOutInd(res.getResources().get(),
+              makeDevAlloc(AllocType::Other, 0), {rows, k});
 
   if (warp) {
-    faiss::gpu::runWarpSelect(gpuVal, gpuOutVal, gpuOutInd, dir, k, 0);
+    runWarpSelect(gpuVal, gpuOutVal, gpuOutInd, dir, k, 0);
   } else {
-    faiss::gpu::runBlockSelect(gpuVal, gpuOutVal, gpuOutInd, dir, k, 0);
+    runBlockSelect(gpuVal, gpuOutVal, gpuOutInd, dir, k, 0);
   }
 
   // Copy back to CPU
-  faiss::gpu::HostTensor<float, 2, true> outVal(gpuOutVal, 0);
-  faiss::gpu::HostTensor<int, 2, true> outInd(gpuOutInd, 0);
+  HostTensor<float, 2, true> outVal(gpuOutVal, 0);
+  HostTensor<int, 2, true> outInd(gpuOutInd, 0);
 
   for (int r = 0; r < rows; ++r) {
     std::unordered_map<int, int> seenIndices;

gpu/test/demo_ivfpq_indexing_gpu.cpp

@@ -5,9 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Copyright 2004-present Facebook. All Rights Reserved
-
-
 #include <cmath>
 #include <cstdio>
 #include <cstdlib>
@@ -17,7 +14,6 @@
 
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/GpuIndexIVFPQ.h>
-#include <faiss/gpu/GpuCloner.h>
 
 #include <faiss/gpu/GpuAutoTune.h>
 #include <faiss/index_io.h>

gpu/test/test_gpu_index.py

@@ -73,13 +73,22 @@ class EvalIVFPQAccuracy(unittest.TestCase):
         gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
         ts.append(time.time())
 
+        # Validate the layout of the memory info
+        mem_info = res.getMemoryInfo()
+
+        assert type(mem_info) == dict
+        assert type(mem_info[0]['FlatData']) == tuple
+        assert type(mem_info[0]['FlatData'][0]) == int
+        assert type(mem_info[0]['FlatData'][1]) == int
+
         gpu_index.setNumProbes(4)
 
         D, Inew = gpu_index.search(xq, 10)
         ts.append(time.time())
         print('times:', [t - ts[0] for t in ts])
 
-        self.assertGreaterEqual((Iref == Inew).sum(), Iref.size)
+        # Give us some margin of error
+        self.assertGreaterEqual((Iref == Inew).sum(), Iref.size - 50)
 
         if faiss.get_num_gpus() == 1:
             return
@@ -310,5 +319,32 @@ class TestAlternativeDistances(unittest.TestCase):
         self.do_test(faiss.METRIC_Lp, 0.7)
 
 
+class TestGpuRef(unittest.TestCase):
+
+    def test_gpu_ref(self):
+        # this crashes
+        dim = 256
+        training_data = np.random.randint(256, size=(10000, dim // 8)).astype('uint8')
+        centroids = 330
+
+        def create_cpu(dim):
+            quantizer = faiss.IndexBinaryFlat(dim)
+            return faiss.IndexBinaryIVF(quantizer, dim, centroids)
+
+        def create_gpu(dim):
+            gpu_quantizer = faiss.index_cpu_to_all_gpus(faiss.IndexFlatL2(dim))
+
+            index = create_cpu(dim)
+            index.clustering_index = gpu_quantizer
+            index.dont_dealloc_me = gpu_quantizer
+            return index
+
+        index = create_gpu(dim)
+
+        index.verbose = True
+        index.cp.verbose = True
+
+        index.train(training_data)
+
 if __name__ == '__main__':
     unittest.main()

gpu/test/test_gpu_index_ivfsq.py

@@ -1,10 +1,10 @@
+#! /usr/bin/env python3
+
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-#! /usr/bin/env python3
-
 from __future__ import print_function
 import unittest
 import numpy as np

gpu/test/test_pytorch_faiss.py

@@ -91,10 +91,19 @@ def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None,
     D_ptr = swig_ptr_from_FloatTensor(D)
     I_ptr = swig_ptr_from_LongTensor(I)
 
-    faiss.bruteForceKnn(res, metric,
-                        xb_ptr, xb_row_major, nb,
-                        xq_ptr, xq_row_major, nq,
-                        d, k, D_ptr, I_ptr)
+    args = faiss.GpuDistanceParams()
+    args.metric = metric
+    args.k = k
+    args.dims = d
+    args.vectors = xb_ptr
+    args.vectorsRowMajor = xb_row_major
+    args.numVectors = nb
+    args.queries = xq_ptr
+    args.queriesRowMajor = xq_row_major
+    args.numQueries = nq
+    args.outDistances = D_ptr
+    args.outIndices = I_ptr
+    faiss.bfKnn(res, args)
 
     return D, I
 

gpu/utils/ConversionOperators.cuh

@@ -103,22 +103,32 @@ void convertTensor(cudaStream_t stream,
 }
 
 template <typename From, typename To, int Dim>
-DeviceTensor<To, Dim, true> convertTensor(GpuResources* res,
-                                          cudaStream_t stream,
-                                          Tensor<From, Dim, true>& in) {
-  DeviceTensor<To, Dim, true> out;
-
-  if (res) {
-    out = std::move(DeviceTensor<To, Dim, true>(
-                      res->getMemoryManagerCurrentDevice(),
-                      in.sizes(),
-                      stream));
-  } else {
-    out = std::move(DeviceTensor<To, Dim, true>(in.sizes()));
-  }
+DeviceTensor<To, Dim, true>
+convertTensorTemporary(GpuResources* res,
+                       cudaStream_t stream,
+                       Tensor<From, Dim, true>& in) {
+  FAISS_ASSERT(res);
+  DeviceTensor<To, Dim, true> out(
+    res, makeTempAlloc(AllocType::Other, stream),
+    in.sizes());
 
   convertTensor(stream, in, out);
   return out;
 }
 
+template <typename From, typename To, int Dim>
+DeviceTensor<To, Dim, true>
+convertTensorNonTemporary(GpuResources* res,
+                          cudaStream_t stream,
+                          Tensor<From, Dim, true>& in) {
+  FAISS_ASSERT(res);
+  DeviceTensor<To, Dim, true> out(
+    res, makeDevAlloc(AllocType::Other, stream),
+    in.sizes());
+
+  convertTensor(stream, in, out);
+  return out;
+}
+
+
 } } // namespace