faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp

@@ -16,6 +16,7 @@
 #include <faiss/gpu/GpuIndexIVFFlat.h>
 #include <faiss/gpu/GpuIndexIVFPQ.h>
 #include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
+#include <faiss/gpu/impl/IndexUtils.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
 #include <faiss/impl/FaissAssert.h>
 

data/vendor/faiss/faiss/gpu/GpuCloner.cpp

@@ -121,7 +121,6 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         GpuIndexFlatConfig config;
         config.device = device;
         config.useFloat16 = useFloat16;
-        config.storeTransposed = storeTransposed;
         return new GpuIndexFlat(provider, ifl, config);
     } else if (
             dynamic_cast<const IndexScalarQuantizer*>(index) &&
@@ -147,7 +146,6 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.device = device;
         config.indicesOptions = indicesOptions;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
-        config.flatConfig.storeTransposed = storeTransposed;
 
         GpuIndexIVFFlat* res = new GpuIndexIVFFlat(
                 provider, ifl->d, ifl->nlist, ifl->metric_type, config);
@@ -164,7 +162,6 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.device = device;
         config.indicesOptions = indicesOptions;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
-        config.flatConfig.storeTransposed = storeTransposed;
 
         GpuIndexIVFScalarQuantizer* res = new GpuIndexIVFScalarQuantizer(
                 provider,
@@ -195,7 +192,6 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.device = device;
         config.indicesOptions = indicesOptions;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
-        config.flatConfig.storeTransposed = storeTransposed;
         config.useFloat16LookupTables = useFloat16;
         config.usePrecomputedTables = usePrecomputed;
 

data/vendor/faiss/faiss/gpu/GpuIndex.h

@@ -74,7 +74,19 @@ class GpuIndex : public faiss::Index {
             const float* x,
             Index::idx_t k,
             float* distances,
-            Index::idx_t* labels) const override;
+            Index::idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    /// `x`, `distances` and `labels` and `recons` can be resident on the CPU or
+    /// any GPU; copies are performed as needed
+    void search_and_reconstruct(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            float* recons,
+            const SearchParameters* params = nullptr) const override;
 
     /// Overridden to force GPU indices to provide their own GPU-friendly
     /// implementation
@@ -111,7 +123,8 @@ class GpuIndex : public faiss::Index {
             const float* x,
             int k,
             float* distances,
-            Index::idx_t* labels) const = 0;
+            Index::idx_t* labels,
+            const SearchParameters* params) const = 0;
 
    private:
     /// Handles paged adds if the add set is too large, passes to
@@ -127,7 +140,8 @@ class GpuIndex : public faiss::Index {
             const float* x,
             int k,
             float* outDistancesData,
-            Index::idx_t* outIndicesData) const;
+            Index::idx_t* outIndicesData,
+            const SearchParameters* params) const;
 
     /// Calls searchImpl_ for a single page of GPU-resident data,
     /// handling paging of the data and copies from the CPU
@@ -136,7 +150,8 @@ class GpuIndex : public faiss::Index {
             const float* x,
             int k,
             float* outDistancesData,
-            Index::idx_t* outIndicesData) const;
+            Index::idx_t* outIndicesData,
+            const SearchParameters* params) const;
 
    protected:
     /// Manages streams, cuBLAS handles and scratch memory for devices
@@ -149,5 +164,14 @@ class GpuIndex : public faiss::Index {
     size_t minPagedSize_;
 };
 
+/// If the given index is a GPU index, this returns the index instance
+GpuIndex* tryCastGpuIndex(faiss::Index* index);
+
+/// Is the given index instance a GPU index?
+bool isGpuIndex(faiss::Index* index);
+
+/// Does the given CPU index instance have a corresponding GPU implementation?
+bool isGpuIndexImplemented(faiss::Index* index);
+
 } // namespace gpu
 } // namespace faiss

data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h

@@ -62,7 +62,8 @@ class GpuIndexBinaryFlat : public IndexBinary {
             const uint8_t* x,
             faiss::IndexBinary::idx_t k,
             int32_t* distances,
-            faiss::IndexBinary::idx_t* labels) const override;
+            faiss::IndexBinary::idx_t* labels,
+            const faiss::SearchParameters* params = nullptr) const override;
 
     void reconstruct(faiss::IndexBinary::idx_t key, uint8_t* recons)
             const override;

data/vendor/faiss/faiss/gpu/GpuIndexFlat.h

@@ -24,17 +24,14 @@ namespace gpu {
 class FlatIndex;
 
 struct GpuIndexFlatConfig : public GpuIndexConfig {
-    inline GpuIndexFlatConfig() : useFloat16(false), storeTransposed(false) {}
+    inline GpuIndexFlatConfig() : useFloat16(false) {}
 
     /// Whether or not data is stored as float16
     bool useFloat16;
 
-    /// Whether or not data is stored (transparently) in a transposed
-    /// layout, enabling use of the NN GEMM call, which is ~10% faster.
-    /// This will improve the speed of the flat index, but will
-    /// substantially slow down any add() calls made, as all data must
-    /// be transposed, and will increase storage requirements (we store
-    /// data in both transposed and non-transposed layouts).
+    /// Deprecated: no longer used
+    /// Previously used to indicate whether internal storage of vectors is
+    /// transposed
     bool storeTransposed;
 };
 
@@ -98,6 +95,10 @@ class GpuIndexFlat : public GpuIndex {
     void reconstruct_n(Index::idx_t i0, Index::idx_t num, float* out)
             const override;
 
+    /// Batch reconstruction method
+    void reconstruct_batch(Index::idx_t n, const Index::idx_t* keys, float* out)
+            const override;
+
     /// Compute residual
     void compute_residual(const float* x, float* residual, Index::idx_t key)
             const override;
@@ -128,7 +129,8 @@ class GpuIndexFlat : public GpuIndex {
             const float* x,
             int k,
             float* distances,
-            Index::idx_t* labels) const override;
+            Index::idx_t* labels,
+            const SearchParameters* params) const override;
 
    protected:
     /// Our configuration options

data/vendor/faiss/faiss/gpu/GpuIndexIVF.h

@@ -8,18 +8,17 @@
 #pragma once
 
 #include <faiss/Clustering.h>
+#include <faiss/IndexIVF.h> // for SearchParametersIVF
 #include <faiss/gpu/GpuIndex.h>
 #include <faiss/gpu/GpuIndexFlat.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
-
-namespace faiss {
-struct IndexIVF;
-}
+#include <memory>
 
 namespace faiss {
 namespace gpu {
 
 class GpuIndexFlat;
+class IVFBase;
 
 struct GpuIndexIVFConfig : public GpuIndexConfig {
     inline GpuIndexIVFConfig() : indicesOptions(INDICES_64_BIT) {}
@@ -31,8 +30,13 @@ struct GpuIndexIVFConfig : public GpuIndexConfig {
     GpuIndexFlatConfig flatConfig;
 };
 
+/// Base class of all GPU IVF index types. This (for now) deliberately does not
+/// inherit from IndexIVF, as many of the public data members and functionality
+/// in IndexIVF is not supported in the same manner on the GPU.
 class GpuIndexIVF : public GpuIndex {
    public:
+    /// Version that auto-constructs a flat coarse quantizer based on the
+    /// desired metric
     GpuIndexIVF(
             GpuResourcesProvider* provider,
             int dims,
@@ -41,6 +45,17 @@ class GpuIndexIVF : public GpuIndex {
             int nlist,
             GpuIndexIVFConfig config = GpuIndexIVFConfig());
 
+    /// Version that takes a coarse quantizer instance. The GpuIndexIVF does not
+    /// own the coarseQuantizer instance by default (functions like IndexIVF).
+    GpuIndexIVF(
+            GpuResourcesProvider* provider,
+            Index* coarseQuantizer,
+            int dims,
+            faiss::MetricType metric,
+            float metricArg,
+            int nlist,
+            GpuIndexIVFConfig config = GpuIndexIVFConfig());
+
     ~GpuIndexIVF() override;
 
    private:
@@ -54,11 +69,16 @@ class GpuIndexIVF : public GpuIndex {
     /// Copy what we have to the CPU equivalent
     void copyTo(faiss::IndexIVF* index) const;
 
+    /// Should be called if the user ever changes the state of the IVF coarse
+    /// quantizer manually (e.g., substitutes a new instance or changes vectors
+    /// in the coarse quantizer outside the scope of training)
+    virtual void updateQuantizer() = 0;
+
     /// Returns the number of inverted lists we're managing
     int getNumLists() const;
 
     /// Returns the number of vectors present in a particular inverted list
-    virtual int getListLength(int listId) const = 0;
+    int getListLength(int listId) const;
 
     /// Return the encoded vector data contained in a particular inverted list,
     /// for debugging purposes.
@@ -66,16 +86,12 @@ class GpuIndexIVF : public GpuIndex {
     /// GPU-side representation.
     /// Otherwise, it is converted to the CPU format.
     /// compliant format, while the native GPU format may differ.
-    virtual std::vector<uint8_t> getListVectorData(
-            int listId,
-            bool gpuFormat = false) const = 0;
+    std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false)
+            const;
 
     /// Return the vector indices contained in a particular inverted list, for
     /// debugging purposes.
-    virtual std::vector<Index::idx_t> getListIndices(int listId) const = 0;
-
-    /// Return the quantizer we're using
-    GpuIndexFlat* getQuantizer();
+    std::vector<Index::idx_t> getListIndices(int listId) const;
 
     /// Sets the number of list probes per query
     void setNumProbes(int nprobe);
@@ -83,10 +99,49 @@ class GpuIndexIVF : public GpuIndex {
     /// Returns our current number of list probes per query
     int getNumProbes() const;
 
+    /// Same interface as faiss::IndexIVF, in order to search a set of vectors
+    /// pre-quantized by the IVF quantizer. Does not include IndexIVFStats as
+    /// that can only be obtained on the host via a GPU d2h copy.
+    /// @param n      nb of vectors to query
+    /// @param x      query vectors, size nx * d
+    /// @param assign coarse quantization indices, size nx * nprobe
+    /// @param centroid_dis
+    ///             distances to coarse centroids, size nx * nprobe
+    /// @param distance
+    ///             output distances, size n * k
+    /// @param labels output labels, size n * k
+    /// @param store_pairs store inv list index + inv list offset
+    ///                   instead in upper/lower 32 bit of result,
+    ///                   instead of ids (used for reranking).
+    /// @param params used to override the object's search parameters
+    void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const SearchParametersIVF* params = nullptr) const;
+
    protected:
+    void verifyIVFSettings_() const;
     bool addImplRequiresIDs_() const override;
     void trainQuantizer_(Index::idx_t n, const float* x);
 
+    /// Called from GpuIndex for add/add_with_ids
+    void addImpl_(int n, const float* x, const Index::idx_t* ids) override;
+
+    /// Called from GpuIndex for search
+    void searchImpl_(
+            int n,
+            const float* x,
+            int k,
+            float* distances,
+            Index::idx_t* labels,
+            const SearchParameters* params) const override;
+
    public:
     /// Exposing this like the CPU version for manipulation
     ClusteringParameters cp;
@@ -97,12 +152,18 @@ class GpuIndexIVF : public GpuIndex {
     /// Exposing this like the CPU version for manipulation
     int nprobe;
 
-    /// Exposeing this like the CPU version for query
-    GpuIndexFlat* quantizer;
+    /// A user-pluggable coarse quantizer
+    Index* quantizer;
+
+    /// Whether or not we own the coarse quantizer
+    bool own_fields;
 
    protected:
     /// Our configuration options
     const GpuIndexIVFConfig ivfConfig_;
+
+    /// For a trained/initialized index, this is a reference to the base class
+    std::shared_ptr<IVFBase> baseIndex_;
 };
 
 } // namespace gpu

data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h

@@ -40,12 +40,22 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
             GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
 
     /// Constructs a new instance with an empty flat quantizer; the user
-    /// provides the number of lists desired.
+    /// provides the number of IVF lists desired.
     GpuIndexIVFFlat(
             GpuResourcesProvider* provider,
             int dims,
             int nlist,
-            faiss::MetricType metric,
+            faiss::MetricType metric = faiss::METRIC_L2,
+            GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
+
+    /// Constructs a new instance with a provided CPU or GPU coarse quantizer;
+    /// the user provides the number of IVF lists desired.
+    GpuIndexIVFFlat(
+            GpuResourcesProvider* provider,
+            Index* coarseQuantizer,
+            int dims,
+            int nlist,
+            faiss::MetricType metric = faiss::METRIC_L2,
             GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
 
     ~GpuIndexIVFFlat() override;
@@ -69,37 +79,14 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
     /// information
     void reset() override;
 
+    /// Should be called if the user ever changes the state of the IVF coarse
+    /// quantizer manually (e.g., substitutes a new instance or changes vectors
+    /// in the coarse quantizer outside the scope of training)
+    void updateQuantizer() override;
+
     /// Trains the coarse quantizer based on the given vector data
     void train(Index::idx_t n, const float* x) override;
 
-    /// Returns the number of vectors present in a particular inverted list
-    int getListLength(int listId) const override;
-
-    /// Return the encoded vector data contained in a particular inverted list,
-    /// for debugging purposes.
-    /// If gpuFormat is true, the data is returned as it is encoded in the
-    /// GPU-side representation.
-    /// Otherwise, it is converted to the CPU format.
-    /// compliant format, while the native GPU format may differ.
-    std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false)
-            const override;
-
-    /// Return the vector indices contained in a particular inverted list, for
-    /// debugging purposes.
-    std::vector<Index::idx_t> getListIndices(int listId) const override;
-
-   protected:
-    /// Called from GpuIndex for add/add_with_ids
-    void addImpl_(int n, const float* x, const Index::idx_t* ids) override;
-
-    /// Called from GpuIndex for search
-    void searchImpl_(
-            int n,
-            const float* x,
-            int k,
-            float* distances,
-            Index::idx_t* labels) const override;
-
    protected:
     /// Our configuration options
     const GpuIndexIVFFlatConfig ivfFlatConfig_;
@@ -107,8 +94,8 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
     /// Desired inverted list memory reservation
     size_t reserveMemoryVecs_;
 
-    /// Instance that we own; contains the inverted list
-    std::unique_ptr<IVFFlat> index_;
+    /// Instance that we own; contains the inverted lists
+    std::shared_ptr<IVFFlat> index_;
 };
 
 } // namespace gpu

data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h

@@ -63,14 +63,27 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
             const faiss::IndexIVFPQ* index,
             GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
 
-    /// Construct an empty index
+    /// Constructs a new instance with an empty flat quantizer; the user
+    /// provides the number of IVF lists desired.
     GpuIndexIVFPQ(
             GpuResourcesProvider* provider,
             int dims,
             int nlist,
             int subQuantizers,
             int bitsPerCode,
-            faiss::MetricType metric,
+            faiss::MetricType metric = faiss::METRIC_L2,
+            GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
+
+    /// Constructs a new instance with a provided CPU or GPU coarse quantizer;
+    /// the user provides the number of IVF lists desired.
+    GpuIndexIVFPQ(
+            GpuResourcesProvider* provider,
+            Index* coarseQuantizer,
+            int dims,
+            int nlist,
+            int subQuantizers,
+            int bitsPerCode,
+            faiss::MetricType metric = faiss::METRIC_L2,
             GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
 
     ~GpuIndexIVFPQ() override;
@@ -112,44 +125,22 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
     /// product centroid information
     void reset() override;
 
+    /// Should be called if the user ever changes the state of the IVF coarse
+    /// quantizer manually (e.g., substitutes a new instance or changes vectors
+    /// in the coarse quantizer outside the scope of training)
+    void updateQuantizer() override;
+
     /// Trains the coarse and product quantizer based on the given vector data
     void train(Index::idx_t n, const float* x) override;
 
-    /// Returns the number of vectors present in a particular inverted list
-    int getListLength(int listId) const override;
-
-    /// Return the encoded vector data contained in a particular inverted list,
-    /// for debugging purposes.
-    /// If gpuFormat is true, the data is returned as it is encoded in the
-    /// GPU-side representation.
-    /// Otherwise, it is converted to the CPU format.
-    /// compliant format, while the native GPU format may differ.
-    std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false)
-            const override;
-
-    /// Return the vector indices contained in a particular inverted list, for
-    /// debugging purposes.
-    std::vector<Index::idx_t> getListIndices(int listId) const override;
-
    public:
     /// Like the CPU version, we expose a publically-visible ProductQuantizer
     /// for manipulation
     ProductQuantizer pq;
 
    protected:
-    /// Called from GpuIndex for add/add_with_ids
-    void addImpl_(int n, const float* x, const Index::idx_t* ids) override;
-
-    /// Called from GpuIndex for search
-    void searchImpl_(
-            int n,
-            const float* x,
-            int k,
-            float* distances,
-            Index::idx_t* labels) const override;
-
     /// Throws errors if configuration settings are improper
-    void verifySettings_() const;
+    void verifyPQSettings_() const;
 
     /// Trains the PQ quantizer based on the given vector data
     void trainResidualQuantizer_(Index::idx_t n, const float* x);
@@ -172,7 +163,7 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
 
     /// The product quantizer instance that we own; contains the
     /// inverted lists
-    std::unique_ptr<IVFPQ> index_;
+    std::shared_ptr<IVFPQ> index_;
 };
 
 } // namespace gpu

data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h

@@ -38,7 +38,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
                     GpuIndexIVFScalarQuantizerConfig());
 
     /// Constructs a new instance with an empty flat quantizer; the user
-    /// provides the number of lists desired.
+    /// provides the number of IVF lists desired.
     GpuIndexIVFScalarQuantizer(
             GpuResourcesProvider* provider,
             int dims,
@@ -49,6 +49,19 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
             GpuIndexIVFScalarQuantizerConfig config =
                     GpuIndexIVFScalarQuantizerConfig());
 
+    /// Constructs a new instance with a provided CPU or GPU coarse quantizer;
+    /// the user provides the number of IVF lists desired.
+    GpuIndexIVFScalarQuantizer(
+            GpuResourcesProvider* provider,
+            Index* coarseQuantizer,
+            int dims,
+            int nlist,
+            faiss::ScalarQuantizer::QuantizerType qtype,
+            faiss::MetricType metric = MetricType::METRIC_L2,
+            bool encodeResidual = true,
+            GpuIndexIVFScalarQuantizerConfig config =
+                    GpuIndexIVFScalarQuantizerConfig());
+
     ~GpuIndexIVFScalarQuantizer() override;
 
     /// Reserve GPU memory in our inverted lists for this number of vectors
@@ -70,36 +83,17 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
     /// quantizer information
     void reset() override;
 
+    /// Should be called if the user ever changes the state of the IVF coarse
+    /// quantizer manually (e.g., substitutes a new instance or changes vectors
+    /// in the coarse quantizer outside the scope of training)
+    void updateQuantizer() override;
+
     /// Trains the coarse and scalar quantizer based on the given vector data
     void train(Index::idx_t n, const float* x) override;
 
-    /// Returns the number of vectors present in a particular inverted list
-    int getListLength(int listId) const override;
-
-    /// Return the encoded vector data contained in a particular inverted list,
-    /// for debugging purposes.
-    /// If gpuFormat is true, the data is returned as it is encoded in the
-    /// GPU-side representation.
-    /// Otherwise, it is converted to the CPU format.
-    /// compliant format, while the native GPU format may differ.
-    std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false)
-            const override;
-
-    /// Return the vector indices contained in a particular inverted list, for
-    /// debugging purposes.
-    std::vector<Index::idx_t> getListIndices(int listId) const override;
-
    protected:
-    /// Called from GpuIndex for add/add_with_ids
-    void addImpl_(int n, const float* x, const Index::idx_t* ids) override;
-
-    /// Called from GpuIndex for search
-    void searchImpl_(
-            int n,
-            const float* x,
-            int k,
-            float* distances,
-            Index::idx_t* labels) const override;
+    /// Validates index SQ parameters
+    void verifySQSettings_() const;
 
     /// Called from train to handle SQ residual training
     void trainResiduals_(Index::idx_t n, const float* x);
@@ -119,7 +113,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
     size_t reserveMemoryVecs_;
 
     /// Instance that we own; contains the inverted list
-    std::unique_ptr<IVFFlat> index_;
+    std::shared_ptr<IVFFlat> index_;
 };
 
 } // namespace gpu

data/vendor/faiss/faiss/gpu/GpuResources.cpp

@@ -184,5 +184,19 @@ size_t GpuResources::getTempMemoryAvailableCurrentDevice() const {
 
 GpuResourcesProvider::~GpuResourcesProvider() {}
 
+//
+// GpuResourcesProviderFromResourceInstance
+//
+
+GpuResourcesProviderFromInstance::GpuResourcesProviderFromInstance(
+        std::shared_ptr<GpuResources> p)
+        : res_(p) {}
+
+GpuResourcesProviderFromInstance::~GpuResourcesProviderFromInstance() {}
+
+std::shared_ptr<GpuResources> GpuResourcesProviderFromInstance::getResources() {
+    return res_;
+}
+
 } // namespace gpu
 } // namespace faiss

data/vendor/faiss/faiss/gpu/GpuResources.h

@@ -28,8 +28,7 @@ enum AllocType {
     FlatData = 1,
 
     /// Primary data storage for GpuIndexIVF* (the storage for each individual
-    /// IVF
-    /// list)
+    /// IVF list)
     IVFLists = 2,
 
     /// Quantizer (PQ, SQ) dictionary information
@@ -251,7 +250,8 @@ class GpuResources {
     cudaStream_t getAsyncCopyStreamCurrentDevice();
 };
 
-/// Interface for a provider of a shared resources object
+/// Interface for a provider of a shared resources object. This is to avoid
+/// interfacing std::shared_ptr to Python
 class GpuResourcesProvider {
    public:
     virtual ~GpuResourcesProvider();
@@ -260,5 +260,18 @@ class GpuResourcesProvider {
     virtual std::shared_ptr<GpuResources> getResources() = 0;
 };
 
+/// A simple wrapper for a GpuResources object to make a GpuResourcesProvider
+/// out of it again
+class GpuResourcesProviderFromInstance : public GpuResourcesProvider {
+   public:
+    explicit GpuResourcesProviderFromInstance(std::shared_ptr<GpuResources> p);
+    ~GpuResourcesProviderFromInstance() override;
+
+    std::shared_ptr<GpuResources> getResources() override;
+
+   private:
+    std::shared_ptr<GpuResources> res_;
+};
+
 } // namespace gpu
 } // namespace faiss

data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp

@@ -268,6 +268,9 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
         return;
     }
 
+    FAISS_ASSERT(device < getNumDevices());
+    DeviceScope scope(device);
+
     // If this is the first device that we're initializing, create our
     // pinned memory allocation
     if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
@@ -285,9 +288,6 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
         pinnedMemAllocSize_ = pinnedMemSize_;
     }
 
-    FAISS_ASSERT(device < getNumDevices());
-    DeviceScope scope(device);
-
     // Make sure that device properties for all devices are cached
     auto& prop = getDeviceProperties(device);