faiss 0.1.0 → 0.1.1

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

@@ -0,0 +1,60 @@
+/**
+ * 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 <faiss/gpu/utils/Timer.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss { namespace gpu {
+
+KernelTimer::KernelTimer(cudaStream_t stream)
+    : startEvent_(0),
+      stopEvent_(0),
+      stream_(stream),
+      valid_(true) {
+  CUDA_VERIFY(cudaEventCreate(&startEvent_));
+  CUDA_VERIFY(cudaEventCreate(&stopEvent_));
+
+  CUDA_VERIFY(cudaEventRecord(startEvent_, stream_));
+}
+
+KernelTimer::~KernelTimer() {
+  CUDA_VERIFY(cudaEventDestroy(startEvent_));
+  CUDA_VERIFY(cudaEventDestroy(stopEvent_));
+}
+
+float
+KernelTimer::elapsedMilliseconds() {
+  FAISS_ASSERT(valid_);
+
+  CUDA_VERIFY(cudaEventRecord(stopEvent_, stream_));
+  CUDA_VERIFY(cudaEventSynchronize(stopEvent_));
+
+  auto time = 0.0f;
+  CUDA_VERIFY(cudaEventElapsedTime(&time, startEvent_, stopEvent_));
+  valid_ = false;
+
+  return time;
+}
+
+CpuTimer::CpuTimer() {
+  clock_gettime(CLOCK_REALTIME, &start_);
+}
+
+float
+CpuTimer::elapsedMilliseconds() {
+  struct timespec end;
+  clock_gettime(CLOCK_REALTIME, &end);
+
+  auto diffS = end.tv_sec - start_.tv_sec;
+  auto diffNs = end.tv_nsec - start_.tv_nsec;
+
+  return 1000.0f * (float) diffS + ((float) diffNs) / 1000000.0f;
+}
+
+} } // namespace

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

@@ -0,0 +1,52 @@
+/**
+ * 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.
+ */
+
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include <time.h>
+
+namespace faiss { namespace gpu {
+
+/// Utility class for timing execution of a kernel
+class KernelTimer {
+ public:
+  /// Constructor starts the timer and adds an event into the current
+  /// device stream
+  KernelTimer(cudaStream_t stream = 0);
+
+  /// Destructor releases event resources
+  ~KernelTimer();
+
+  /// Adds a stop event then synchronizes on the stop event to get the
+  /// actual GPU-side kernel timings for any kernels launched in the
+  /// current stream. Returns the number of milliseconds elapsed.
+  /// Can only be called once.
+  float elapsedMilliseconds();
+
+ private:
+  cudaEvent_t startEvent_;
+  cudaEvent_t stopEvent_;
+  cudaStream_t stream_;
+  bool valid_;
+};
+
+/// CPU wallclock elapsed timer
+class CpuTimer {
+ public:
+  /// Creates and starts a new timer
+  CpuTimer();
+
+  /// Returns elapsed time in milliseconds
+  float elapsedMilliseconds();
+
+ private:
+  struct timespec start_;
+};
+
+} } // namespace

data/vendor/faiss/impl/AuxIndexStructures.cpp

@@ -0,0 +1,305 @@
+/**
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#include <cstring>
+
+#include <faiss/impl/AuxIndexStructures.h>
+
+#include <faiss/impl/FaissAssert.h>
+
+
+namespace faiss {
+
+
+/***********************************************************************
+ * RangeSearchResult
+ ***********************************************************************/
+
+RangeSearchResult::RangeSearchResult (idx_t nq, bool alloc_lims): nq (nq) {
+    if (alloc_lims) {
+        lims = new size_t [nq + 1];
+        memset (lims, 0, sizeof(*lims) * (nq + 1));
+    } else {
+        lims = nullptr;
+    }
+    labels = nullptr;
+    distances = nullptr;
+    buffer_size = 1024 * 256;
+}
+
+/// called when lims contains the nb of elements result entries
+/// for each query
+void RangeSearchResult::do_allocation () {
+    size_t ofs = 0;
+    for (int i = 0; i < nq; i++) {
+        size_t n = lims[i];
+        lims [i] = ofs;
+        ofs += n;
+    }
+    lims [nq] = ofs;
+    labels = new idx_t [ofs];
+    distances = new float [ofs];
+}
+
+RangeSearchResult::~RangeSearchResult () {
+    delete [] labels;
+    delete [] distances;
+    delete [] lims;
+}
+
+
+
+
+
+/***********************************************************************
+ * BufferList
+ ***********************************************************************/
+
+
+BufferList::BufferList (size_t buffer_size):
+    buffer_size (buffer_size)
+{
+    wp = buffer_size;
+}
+
+BufferList::~BufferList ()
+{
+    for (int i = 0; i < buffers.size(); i++) {
+        delete [] buffers[i].ids;
+        delete [] buffers[i].dis;
+    }
+}
+
+void BufferList::add (idx_t id, float dis) {
+    if (wp == buffer_size) { // need new buffer
+        append_buffer();
+    }
+    Buffer & buf = buffers.back();
+    buf.ids [wp] = id;
+    buf.dis [wp] = dis;
+    wp++;
+}
+
+
+void BufferList::append_buffer ()
+{
+    Buffer buf = {new idx_t [buffer_size], new float [buffer_size]};
+    buffers.push_back (buf);
+    wp = 0;
+}
+
+/// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
+/// tables dest_ids, dest_dis
+void BufferList::copy_range (size_t ofs, size_t n,
+                             idx_t * dest_ids, float *dest_dis)
+{
+    size_t bno = ofs / buffer_size;
+    ofs -= bno * buffer_size;
+    while (n > 0) {
+        size_t ncopy = ofs + n < buffer_size ? n : buffer_size - ofs;
+        Buffer buf = buffers [bno];
+        memcpy (dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
+        memcpy (dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
+        dest_ids += ncopy;
+        dest_dis += ncopy;
+        ofs = 0;
+        bno ++;
+        n -= ncopy;
+    }
+}
+
+
+/***********************************************************************
+ * RangeSearchPartialResult
+ ***********************************************************************/
+
+void RangeQueryResult::add (float dis, idx_t id) {
+    nres++;
+    pres->add (id, dis);
+}
+
+
+
+RangeSearchPartialResult::RangeSearchPartialResult (RangeSearchResult * res_in):
+    BufferList(res_in->buffer_size),
+    res(res_in)
+{}
+
+
+/// begin a new result
+RangeQueryResult &
+    RangeSearchPartialResult::new_result (idx_t qno)
+{
+    RangeQueryResult qres = {qno, 0, this};
+    queries.push_back (qres);
+    return queries.back();
+}
+
+
+void RangeSearchPartialResult::finalize ()
+{
+    set_lims ();
+#pragma omp barrier
+
+#pragma omp single
+    res->do_allocation ();
+
+#pragma omp barrier
+    copy_result ();
+}
+
+
+/// called by range_search before do_allocation
+void RangeSearchPartialResult::set_lims ()
+{
+    for (int i = 0; i < queries.size(); i++) {
+        RangeQueryResult & qres = queries[i];
+        res->lims[qres.qno] = qres.nres;
+    }
+}
+
+/// called by range_search after do_allocation
+void RangeSearchPartialResult::copy_result (bool incremental)
+{
+    size_t ofs = 0;
+    for (int i = 0; i < queries.size(); i++) {
+        RangeQueryResult & qres = queries[i];
+
+        copy_range (ofs, qres.nres,
+                    res->labels + res->lims[qres.qno],
+                    res->distances + res->lims[qres.qno]);
+        if (incremental) {
+            res->lims[qres.qno] += qres.nres;
+        }
+        ofs += qres.nres;
+    }
+}
+
+void RangeSearchPartialResult::merge (std::vector <RangeSearchPartialResult *> &
+                                      partial_results, bool do_delete)
+{
+
+    int npres = partial_results.size();
+    if (npres == 0) return;
+    RangeSearchResult *result = partial_results[0]->res;
+    size_t nx = result->nq;
+
+    // count
+    for (const RangeSearchPartialResult * pres : partial_results) {
+        if (!pres) continue;
+        for (const RangeQueryResult &qres : pres->queries) {
+            result->lims[qres.qno] += qres.nres;
+        }
+    }
+    result->do_allocation ();
+    for (int j = 0; j < npres; j++) {
+        if (!partial_results[j]) continue;
+        partial_results[j]->copy_result (true);
+        if (do_delete) {
+            delete partial_results[j];
+            partial_results[j] = nullptr;
+        }
+    }
+
+    // reset the limits
+    for (size_t i = nx; i > 0; i--) {
+        result->lims [i] = result->lims [i - 1];
+    }
+    result->lims [0] = 0;
+}
+
+/***********************************************************************
+ * IDSelectorRange
+ ***********************************************************************/
+
+IDSelectorRange::IDSelectorRange (idx_t imin, idx_t imax):
+    imin (imin), imax (imax)
+{
+}
+
+bool IDSelectorRange::is_member (idx_t id) const
+{
+    return id >= imin && id < imax;
+}
+
+
+/***********************************************************************
+ * IDSelectorBatch
+ ***********************************************************************/
+
+IDSelectorBatch::IDSelectorBatch (size_t n, const idx_t *indices)
+{
+    nbits = 0;
+    while (n > (1L << nbits)) nbits++;
+    nbits += 5;
+    // for n = 1M, nbits = 25 is optimal, see P56659518
+
+    mask = (1L << nbits) - 1;
+    bloom.resize (1UL << (nbits - 3), 0);
+    for (long i = 0; i < n; i++) {
+        Index::idx_t id = indices[i];
+        set.insert(id);
+        id &= mask;
+        bloom[id >> 3] |= 1 << (id & 7);
+    }
+}
+
+bool IDSelectorBatch::is_member (idx_t i) const
+{
+    long im = i & mask;
+    if(!(bloom[im>>3] & (1 << (im & 7)))) {
+        return 0;
+    }
+    return set.count(i);
+}
+
+
+/***********************************************************
+ * Interrupt callback
+ ***********************************************************/
+
+
+std::unique_ptr<InterruptCallback> InterruptCallback::instance;
+
+std::mutex InterruptCallback::lock;
+
+void InterruptCallback::clear_instance () {
+    delete instance.release ();
+}
+
+void InterruptCallback::check () {
+    if (!instance.get()) {
+        return;
+    }
+    if (instance->want_interrupt ()) {
+        FAISS_THROW_MSG ("computation interrupted");
+    }
+}
+
+bool InterruptCallback::is_interrupted () {
+    if (!instance.get()) {
+        return false;
+    }
+    std::lock_guard<std::mutex> guard(lock);
+    return instance->want_interrupt();
+}
+
+
+size_t InterruptCallback::get_period_hint (size_t flops) {
+    if (!instance.get()) {
+        return 1L << 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);
+}
+
+
+
+
+} // namespace faiss

data/vendor/faiss/impl/AuxIndexStructures.h

@@ -0,0 +1,246 @@
+/**
+ * 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.
+ */
+
+// -*- c++ -*-
+
+// Auxiliary index structures, that are used in indexes but that can
+// be forward-declared
+
+#ifndef FAISS_AUX_INDEX_STRUCTURES_H
+#define FAISS_AUX_INDEX_STRUCTURES_H
+
+#include <stdint.h>
+
+#include <vector>
+#include <unordered_set>
+#include <memory>
+#include <mutex>
+
+#include <faiss/Index.h>
+
+namespace faiss {
+
+/** The objective is to have a simple result structure while
+ *  minimizing the number of mem copies in the result. The method
+ *  do_allocation can be overloaded to allocate the result tables in
+ *  the matrix type of a scripting language like Lua or Python. */
+struct RangeSearchResult {
+    size_t nq;      ///< nb of queries
+    size_t *lims;   ///< size (nq + 1)
+
+    typedef Index::idx_t idx_t;
+
+    idx_t *labels;     ///< result for query i is labels[lims[i]:lims[i+1]]
+    float *distances;  ///< corresponding distances (not sorted)
+
+    size_t buffer_size; ///< size of the result buffers used
+
+    /// lims must be allocated on input to range_search.
+    explicit RangeSearchResult (idx_t nq, bool alloc_lims=true);
+
+    /// called when lims contains the nb of elements result entries
+    /// for each query
+
+    virtual void do_allocation ();
+
+    virtual ~RangeSearchResult ();
+};
+
+
+/**
+
+ Encapsulates a set of ids to remove. */
+struct IDSelector {
+    typedef Index::idx_t idx_t;
+    virtual bool is_member (idx_t id) const = 0;
+    virtual ~IDSelector() {}
+};
+
+
+
+/** remove ids between [imni, imax) */
+struct IDSelectorRange: IDSelector {
+    idx_t imin, imax;
+
+    IDSelectorRange (idx_t imin, idx_t imax);
+    bool is_member(idx_t id) const override;
+    ~IDSelectorRange() override {}
+};
+
+
+/** Remove ids from a set. Repetitions of ids in the indices set
+ * passed to the constructor does not hurt performance. The hash
+ * function used for the bloom filter and GCC's implementation of
+ * unordered_set are just the least significant bits of the id. This
+ * works fine for random ids or ids in sequences but will produce many
+ * hash collisions if lsb's are always the same */
+struct IDSelectorBatch: IDSelector {
+
+    std::unordered_set<idx_t> set;
+
+    typedef unsigned char uint8_t;
+    std::vector<uint8_t> bloom; // assumes low bits of id are a good hash value
+    int nbits;
+    idx_t mask;
+
+    IDSelectorBatch (size_t n, const idx_t *indices);
+    bool is_member(idx_t id) const override;
+    ~IDSelectorBatch() override {}
+};
+
+/****************************************************************
+ * Result structures for range search.
+ *
+ * The main constraint here is that we want to support parallel
+ * queries from different threads in various ways: 1 thread per query,
+ * several threads per query. We store the actual results in blocks of
+ * fixed size rather than exponentially increasing memory. At the end,
+ * we copy the block content to a linear result array.
+ *****************************************************************/
+
+/** List of temporary buffers used to store results before they are
+ *  copied to the RangeSearchResult object. */
+struct BufferList {
+    typedef Index::idx_t idx_t;
+
+    // buffer sizes in # entries
+    size_t buffer_size;
+
+    struct Buffer {
+        idx_t *ids;
+        float *dis;
+    };
+
+    std::vector<Buffer> buffers;
+    size_t wp; ///< write pointer in the last buffer.
+
+    explicit BufferList (size_t buffer_size);
+
+    ~BufferList ();
+
+    /// create a new buffer
+    void append_buffer ();
+
+    /// add one result, possibly appending a new buffer if needed
+    void add (idx_t id, float dis);
+
+    /// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
+    /// tables dest_ids, dest_dis
+    void copy_range (size_t ofs, size_t n,
+                     idx_t * dest_ids, float *dest_dis);
+
+};
+
+struct RangeSearchPartialResult;
+
+/// result structure for a single query
+struct RangeQueryResult {
+    using idx_t = Index::idx_t;
+    idx_t qno;    //< id of the query
+    size_t nres;  //< nb of results for this query
+    RangeSearchPartialResult * pres;
+
+    /// called by search function to report a new result
+    void add (float dis, idx_t id);
+};
+
+/// the entries in the buffers are split per query
+struct RangeSearchPartialResult: BufferList {
+    RangeSearchResult * res;
+
+    /// eventually the result will be stored in res_in
+    explicit RangeSearchPartialResult (RangeSearchResult * res_in);
+
+    /// query ids + nb of results per query.
+    std::vector<RangeQueryResult> queries;
+
+    /// begin a new result
+    RangeQueryResult & new_result (idx_t qno);
+
+    /*****************************************
+     * functions used at the end of the search to merge the result
+     * lists */
+    void finalize ();
+
+    /// called by range_search before do_allocation
+    void set_lims ();
+
+    /// called by range_search after do_allocation
+    void copy_result (bool incremental = false);
+
+    /// merge a set of PartialResult's into one RangeSearchResult
+    /// on ouptut the partialresults are empty!
+    static void merge (std::vector <RangeSearchPartialResult *> &
+                       partial_results, bool do_delete=true);
+
+};
+
+
+/***********************************************************
+ * The distance computer maintains a current query and computes
+ * distances to elements in an index that supports random access.
+ *
+ * The DistanceComputer is not intended to be thread-safe (eg. because
+ * it maintains counters) so the distance functions are not const,
+ * instanciate one from each thread if needed.
+ ***********************************************************/
+struct DistanceComputer {
+     using idx_t = Index::idx_t;
+
+     /// called before computing distances
+     virtual void set_query(const float *x) = 0;
+
+     /// compute distance of vector i to current query
+     virtual float operator () (idx_t i) = 0;
+
+     /// compute distance between two stored vectors
+     virtual float symmetric_dis (idx_t i, idx_t j) = 0;
+
+     virtual ~DistanceComputer() {}
+};
+
+/***********************************************************
+ * Interrupt callback
+ ***********************************************************/
+
+struct InterruptCallback {
+    virtual bool want_interrupt () = 0;
+    virtual ~InterruptCallback() {}
+
+    // lock that protects concurrent calls to is_interrupted
+    static std::mutex lock;
+
+    static std::unique_ptr<InterruptCallback> instance;
+
+    static void clear_instance ();
+
+    /** check if:
+     * - an interrupt callback is set
+     * - the callback retuns true
+     * if this is the case, then throw an exception. Should not be called
+     * from multiple threds.
+     */
+    static void check ();
+
+    /// same as check() but return true if is interrupted instead of
+    /// throwing. Can be called from multiple threads.
+    static bool is_interrupted ();
+
+    /** assuming each iteration takes a certain number of flops, what
+     * is a reasonable interval to check for interrupts?
+     */
+    static size_t get_period_hint (size_t flops);
+
+};
+
+
+
+}; // namespace faiss
+
+
+
+#endif