faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h

@@ -76,6 +76,12 @@ bool getTensorCoreSupport(int device);
 /// Equivalent to getTensorCoreSupport(getCurrentDevice())
 bool getTensorCoreSupportCurrentDevice();
 
+/// Returns the warp size of the given GPU device
+int getWarpSize(int device);
+
+/// Equivalent to getWarpSize(getCurrentDevice())
+int getWarpSizeCurrentDevice();
+
 /// Returns the amount of currently available memory on the given device
 size_t getFreeMemory(int device);
 

data/vendor/faiss/faiss/gpu/utils/RaftUtils.h

@@ -0,0 +1,75 @@
+/**
+ * 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.
+ */
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <faiss/MetricType.h>
+#include <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/utils/Tensor.cuh>
+
+#include <raft/distance/distance_types.hpp>
+
+#pragma GCC visibility push(default)
+namespace faiss {
+namespace gpu {
+
+inline raft::distance::DistanceType metricFaissToRaft(
+        MetricType metric,
+        bool exactDistance) {
+    switch (metric) {
+        case MetricType::METRIC_INNER_PRODUCT:
+            return raft::distance::DistanceType::InnerProduct;
+        case MetricType::METRIC_L2:
+            return raft::distance::DistanceType::L2Expanded;
+        case MetricType::METRIC_L1:
+            return raft::distance::DistanceType::L1;
+        case MetricType::METRIC_Linf:
+            return raft::distance::DistanceType::Linf;
+        case MetricType::METRIC_Lp:
+            return raft::distance::DistanceType::LpUnexpanded;
+        case MetricType::METRIC_Canberra:
+            return raft::distance::DistanceType::Canberra;
+        case MetricType::METRIC_BrayCurtis:
+            return raft::distance::DistanceType::BrayCurtis;
+        case MetricType::METRIC_JensenShannon:
+            return raft::distance::DistanceType::JensenShannon;
+        default:
+            RAFT_FAIL("Distance type not supported");
+    }
+}
+
+/// Identify matrix rows containing non NaN values. validRows[i] is false if row
+/// i contains a NaN value and true otherwise.
+void validRowIndices(
+        GpuResources* res,
+        Tensor<float, 2, true>& vecs,
+        bool* validRows);
+
+/// Filter out matrix rows containing NaN values. The vectors and indices are
+/// updated in-place.
+idx_t inplaceGatherFilteredRows(
+        GpuResources* res,
+        Tensor<float, 2, true>& vecs,
+        Tensor<idx_t, 1, true>& indices);
+} // namespace gpu
+} // namespace faiss
+#pragma GCC visibility pop

data/vendor/faiss/faiss/gpu/utils/Timer.cpp

@@ -14,7 +14,10 @@ namespace faiss {
 namespace gpu {
 
 KernelTimer::KernelTimer(cudaStream_t stream)
-        : startEvent_(0), stopEvent_(0), stream_(stream), valid_(true) {
+        : startEvent_(nullptr),
+          stopEvent_(nullptr),
+          stream_(stream),
+          valid_(true) {
     CUDA_VERIFY(cudaEventCreate(&startEvent_));
     CUDA_VERIFY(cudaEventCreate(&stopEvent_));
 

data/vendor/faiss/faiss/gpu/utils/Timer.h

@@ -18,7 +18,7 @@ class KernelTimer {
    public:
     /// Constructor starts the timer and adds an event into the current
     /// device stream
-    KernelTimer(cudaStream_t stream = 0);
+    KernelTimer(cudaStream_t stream = nullptr);
 
     /// Destructor releases event resources
     ~KernelTimer();

data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp

@@ -261,7 +261,7 @@ void AdditiveQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
             is_trained, "The additive quantizer is not trained yet.");
 
     // standard additive quantizer decoding
-#pragma omp parallel for if (n > 1000)
+#pragma omp parallel for if (n > 100)
     for (int64_t i = 0; i < n; i++) {
         BitstringReader bsr(code + i * code_size, code_size);
         float* xi = x + i * d;
@@ -370,6 +370,8 @@ void AdditiveQuantizer::compute_LUT(
 
 namespace {
 
+/* compute inner products of one query with all centroids, given a look-up
+ * table of all inner producst with codebook entries */
 void compute_inner_prod_with_LUT(
         const AdditiveQuantizer& aq,
         const float* LUT,

data/vendor/faiss/faiss/impl/AdditiveQuantizer.h

@@ -49,11 +49,11 @@ struct AdditiveQuantizer : Quantizer {
     /// encode a norm into norm_bits bits
     uint64_t encode_norm(float norm) const;
 
-    uint32_t encode_qcint(
-            float x) const; ///< encode norm by non-uniform scalar quantization
+    /// encode norm by non-uniform scalar quantization
+    uint32_t encode_qcint(float x) const;
 
-    float decode_qcint(uint32_t c)
-            const; ///< decode norm by non-uniform scalar quantization
+    /// decode norm by non-uniform scalar quantization
+    float decode_qcint(uint32_t c) const;
 
     /// Encodes how search is performed and how vectors are encoded
     enum Search_type_t {
@@ -203,4 +203,4 @@ struct AdditiveQuantizer : Quantizer {
     virtual ~AdditiveQuantizer();
 };
 
-}; // namespace faiss
+} // namespace faiss

data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp

@@ -230,10 +230,35 @@ bool InterruptCallback::is_interrupted() {
 
 size_t InterruptCallback::get_period_hint(size_t flops) {
     if (!instance.get()) {
-        return 1L << 30; // never check
+        return (size_t)1 << 30; // never check
     }
     // for 10M flops, it is reasonable to check once every 10 iterations
     return std::max((size_t)10 * 10 * 1000 * 1000 / (flops + 1), (size_t)1);
 }
 
+void TimeoutCallback::set_timeout(double timeout_in_seconds) {
+    timeout = timeout_in_seconds;
+    start = std::chrono::steady_clock::now();
+}
+
+bool TimeoutCallback::want_interrupt() {
+    if (timeout == 0) {
+        return false;
+    }
+    auto end = std::chrono::steady_clock::now();
+    std::chrono::duration<float, std::milli> duration = end - start;
+    float elapsed_in_seconds = duration.count() / 1000.0;
+    if (elapsed_in_seconds > timeout) {
+        timeout = 0;
+        return true;
+    }
+    return false;
+}
+
+void TimeoutCallback::reset(double timeout_in_seconds) {
+    auto tc(new faiss::TimeoutCallback());
+    faiss::InterruptCallback::instance.reset(tc);
+    tc->set_timeout(timeout_in_seconds);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/AuxIndexStructures.h

@@ -41,7 +41,6 @@ struct RangeSearchResult {
 
     /// called when lims contains the nb of elements result entries
     /// for each query
-
     virtual void do_allocation();
 
     virtual ~RangeSearchResult();
@@ -123,7 +122,7 @@ struct RangeSearchPartialResult : BufferList {
     void copy_result(bool incremental = false);
 
     /// merge a set of PartialResult's into one RangeSearchResult
-    /// on ouptut the partialresults are empty!
+    /// on output the partialresults are empty!
     static void merge(
             std::vector<RangeSearchPartialResult*>& partial_results,
             bool do_delete = true);
@@ -162,10 +161,18 @@ struct FAISS_API InterruptCallback {
     static size_t get_period_hint(size_t flops);
 };
 
+struct TimeoutCallback : InterruptCallback {
+    std::chrono::time_point<std::chrono::steady_clock> start;
+    double timeout;
+    bool want_interrupt() override;
+    void set_timeout(double timeout_in_seconds);
+    static void reset(double timeout_in_seconds);
+};
+
 /// set implementation optimized for fast access.
 struct VisitedTable {
     std::vector<uint8_t> visited;
-    int visno;
+    uint8_t visno;
 
     explicit VisitedTable(int size) : visited(size), visno(1) {}
 

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -30,12 +30,81 @@ struct DistanceComputer {
     /// compute distance of vector i to current query
     virtual float operator()(idx_t i) = 0;
 
+    /// compute distances of current query to 4 stored vectors.
+    /// certain DistanceComputer implementations may benefit
+    /// heavily from this.
+    virtual void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) {
+        // compute first, assign next
+        const float d0 = this->operator()(idx0);
+        const float d1 = this->operator()(idx1);
+        const float d2 = this->operator()(idx2);
+        const float d3 = this->operator()(idx3);
+        dis0 = d0;
+        dis1 = d1;
+        dis2 = d2;
+        dis3 = d3;
+    }
+
     /// compute distance between two stored vectors
     virtual float symmetric_dis(idx_t i, idx_t j) = 0;
 
     virtual ~DistanceComputer() {}
 };
 
+/* Wrap the distance computer into one that negates the
+   distances. This makes supporting INNER_PRODUCE search easier */
+
+struct NegativeDistanceComputer : DistanceComputer {
+    /// owned by this
+    DistanceComputer* basedis;
+
+    explicit NegativeDistanceComputer(DistanceComputer* basedis)
+            : basedis(basedis) {}
+
+    void set_query(const float* x) override {
+        basedis->set_query(x);
+    }
+
+    /// compute distance of vector i to current query
+    float operator()(idx_t i) override {
+        return -(*basedis)(i);
+    }
+
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        basedis->distances_batch_4(
+                idx0, idx1, idx2, idx3, dis0, dis1, dis2, dis3);
+        dis0 = -dis0;
+        dis1 = -dis1;
+        dis2 = -dis2;
+        dis3 = -dis3;
+    }
+
+    /// compute distance between two stored vectors
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return -basedis->symmetric_dis(i, j);
+    }
+
+    virtual ~NegativeDistanceComputer() {
+        delete basedis;
+    }
+};
+
 /*************************************************************
  * Specialized version of the DistanceComputer when we know that codes are
  * laid out in a flat index.
@@ -49,7 +118,7 @@ struct FlatCodesDistanceComputer : DistanceComputer {
 
     FlatCodesDistanceComputer() : codes(nullptr), code_size(0) {}
 
-    float operator()(idx_t i) final {
+    float operator()(idx_t i) override {
         return distance_to_code(codes + i * code_size);
     }
 

data/vendor/faiss/faiss/impl/FaissAssert.h

@@ -94,13 +94,15 @@
         }                                              \
     } while (false)
 
-#define FAISS_THROW_IF_NOT_MSG(X, MSG)                       \
+#define FAISS_THROW_IF_MSG(X, MSG)                           \
     do {                                                     \
-        if (!(X)) {                                          \
+        if (X) {                                             \
             FAISS_THROW_FMT("Error: '%s' failed: " MSG, #X); \
         }                                                    \
     } while (false)
 
+#define FAISS_THROW_IF_NOT_MSG(X, MSG) FAISS_THROW_IF_MSG(!(X), MSG)
+
 #define FAISS_THROW_IF_NOT_FMT(X, FMT, ...)                               \
     do {                                                                  \
         if (!(X)) {                                                       \

data/vendor/faiss/faiss/impl/FaissException.h

@@ -1,3 +1,4 @@
+
 /**
  * Copyright (c) Facebook, Inc. and its affiliates.
  *
@@ -40,42 +41,20 @@ class FaissException : public std::exception {
 void handleExceptions(
         std::vector<std::pair<int, std::exception_ptr>>& exceptions);
 
-/** bare-bones unique_ptr
- * this one deletes with delete [] */
-template <class T>
-struct ScopeDeleter {
-    const T* ptr;
-    explicit ScopeDeleter(const T* ptr = nullptr) : ptr(ptr) {}
-    void release() {
-        ptr = nullptr;
-    }
-    void set(const T* ptr_in) {
-        ptr = ptr_in;
-    }
-    void swap(ScopeDeleter<T>& other) {
-        std::swap(ptr, other.ptr);
-    }
-    ~ScopeDeleter() {
-        delete[] ptr;
+/** RAII object for a set of possibly transformed vectors (deallocated only if
+ * they are indeed transformed)
+ */
+struct TransformedVectors {
+    const float* x;
+    bool own_x;
+    TransformedVectors(const float* x_orig, const float* x) : x(x) {
+        own_x = x_orig != x;
     }
-};
 
-/** same but deletes with the simple delete (least common case) */
-template <class T>
-struct ScopeDeleter1 {
-    const T* ptr;
-    explicit ScopeDeleter1(const T* ptr = nullptr) : ptr(ptr) {}
-    void release() {
-        ptr = nullptr;
-    }
-    void set(const T* ptr_in) {
-        ptr = ptr_in;
-    }
-    void swap(ScopeDeleter1<T>& other) {
-        std::swap(ptr, other.ptr);
-    }
-    ~ScopeDeleter1() {
-        delete ptr;
+    ~TransformedVectors() {
+        if (own_x) {
+            delete[] x;
+        }
     }
 };