faiss 0.1.0 → 0.1.1

data/vendor/faiss/gpu/StandardGpuResources.h

@@ -0,0 +1,114 @@
+/**
+ * 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 <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/utils/StackDeviceMemory.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#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 {
+ public:
+  StandardGpuResources();
+
+  ~StandardGpuResources() 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();
+
+  /// Enable or disable the warning about not having enough temporary memory
+  /// when cudaMalloc gets called
+  void setCudaMallocWarning(bool b);
+
+ public:
+  /// Internal system calls
+
+  /// Initialize resources for this device
+  void initializeForDevice(int device) override;
+
+  cublasHandle_t getBlasHandle(int device) override;
+
+  cudaStream_t getDefaultStream(int device) override;
+
+  std::vector<cudaStream_t> getAlternateStreams(int device) override;
+
+  DeviceMemory& getMemoryManager(int device) override;
+
+  std::pair<void*, size_t> getPinnedMemory() override;
+
+  cudaStream_t getAsyncCopyStream(int device) override;
+
+ private:
+  /// Have GPU resources been initialized for this device yet?
+  bool isInitialized(int device) const;
+
+  /// Adjust the default temporary memory allocation based on the total GPU
+  /// memory size
+  static size_t getDefaultTempMemForGPU(int device, size_t requested);
+
+ private:
+  /// Our default stream that work is ordered on, one per each device
+  std::unordered_map<int, cudaStream_t> defaultStreams_;
+
+  /// This contains particular streams as set by the user for
+  /// ordering, if any
+  std::unordered_map<int, cudaStream_t> userDefaultStreams_;
+
+  /// Other streams we can use, per each device
+  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_;
+
+  /// 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_;
+
+  /// Another option is to use a specified amount of memory on all
+  /// devices
+  size_t tempMemSize_;
+
+  /// Amount of pinned memory we should allocate
+  size_t pinnedMemSize_;
+
+  /// Whether or not a warning upon cudaMalloc is generated
+  bool cudaMallocWarning_;
+};
+
+} } // namespace

data/vendor/faiss/gpu/impl/RemapIndices.cpp

@@ -0,0 +1,43 @@
+/**
+ * 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/impl/RemapIndices.h>
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss { namespace gpu {
+
+// Utility function to translate (list id, offset) to a user index on
+// the CPU. In a cpp in order to use OpenMP
+void ivfOffsetToUserIndex(
+  long* indices,
+  int numLists,
+  int queries,
+  int k,
+  const std::vector<std::vector<long>>& listOffsetToUserIndex) {
+  FAISS_ASSERT(numLists == listOffsetToUserIndex.size());
+
+#pragma omp parallel for
+  for (int q = 0; q < queries; ++q) {
+    for (int r = 0; r < k; ++r) {
+      long offsetIndex = indices[q * k + r];
+
+      if (offsetIndex < 0) continue;
+
+      int listId = (int) (offsetIndex >> 32);
+      int listOffset = (int) (offsetIndex & 0xffffffff);
+
+      FAISS_ASSERT(listId < numLists);
+      auto& listIndices = listOffsetToUserIndex[listId];
+
+      FAISS_ASSERT(listOffset < listIndices.size());
+      indices[q * k + r] = listIndices[listOffset];
+    }
+  }
+}
+
+} } // namespace

data/vendor/faiss/gpu/impl/RemapIndices.h

@@ -0,0 +1,24 @@
+/**
+ * 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 <vector>
+
+namespace faiss { namespace gpu {
+
+/// Utility function to translate (list id, offset) to a user index on
+/// the CPU. In a cpp in order to use OpenMP.
+void ivfOffsetToUserIndex(
+  long* indices,
+  int numLists,
+  int queries,
+  int k,
+  const std::vector<std::vector<long>>& listOffsetToUserIndex);
+
+} } // namespace

data/vendor/faiss/gpu/perf/IndexWrapper-inl.h

@@ -0,0 +1,71 @@
+/**
+ * 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/impl/FaissAssert.h>
+
+namespace faiss { namespace gpu {
+
+template <typename GpuIndex>
+IndexWrapper<GpuIndex>::IndexWrapper(
+  int numGpus,
+  std::function<std::unique_ptr<GpuIndex>(GpuResources*, int)> init) {
+  FAISS_ASSERT(numGpus <= faiss::gpu::getNumDevices());
+  for (int i = 0; i < numGpus; ++i) {
+    auto res = std::unique_ptr<faiss::gpu::StandardGpuResources>(
+      new StandardGpuResources);
+
+    subIndex.emplace_back(init(res.get(), i));
+    resources.emplace_back(std::move(res));
+  }
+
+  if (numGpus > 1) {
+    // create proxy
+    replicaIndex =
+      std::unique_ptr<faiss::IndexReplicas>(new faiss::IndexReplicas);
+
+    for (auto& index : subIndex) {
+      replicaIndex->addIndex(index.get());
+    }
+  }
+}
+
+template <typename GpuIndex>
+faiss::Index*
+IndexWrapper<GpuIndex>::getIndex() {
+  if ((bool) replicaIndex) {
+    return replicaIndex.get();
+  } else {
+    FAISS_ASSERT(!subIndex.empty());
+    return subIndex.front().get();
+  }
+}
+
+template <typename GpuIndex>
+void
+IndexWrapper<GpuIndex>::runOnIndices(std::function<void(GpuIndex*)> f) {
+
+  if ((bool) replicaIndex) {
+    replicaIndex->runOnIndex(
+      [f](int, faiss::Index* index) {
+        f(dynamic_cast<GpuIndex*>(index));
+      });
+  } else {
+    FAISS_ASSERT(!subIndex.empty());
+    f(subIndex.front().get());
+  }
+}
+
+template <typename GpuIndex>
+void
+IndexWrapper<GpuIndex>::setNumProbes(int nprobe) {
+  runOnIndices([nprobe](GpuIndex* index) {
+      index->setNumProbes(nprobe);
+    });
+}
+
+} }

data/vendor/faiss/gpu/perf/IndexWrapper.h

@@ -0,0 +1,39 @@
+/**
+ * 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 <faiss/IndexReplicas.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <functional>
+#include <memory>
+#include <vector>
+
+namespace faiss { namespace gpu {
+
+// If we want to run multi-GPU, create a proxy to wrap the indices.
+// If we don't want multi-GPU, don't involve the proxy, so it doesn't
+// affect the timings.
+template <typename GpuIndex>
+struct IndexWrapper {
+  std::vector<std::unique_ptr<faiss::gpu::StandardGpuResources>> resources;
+  std::vector<std::unique_ptr<GpuIndex>> subIndex;
+  std::unique_ptr<faiss::IndexReplicas> replicaIndex;
+
+  IndexWrapper(
+    int numGpus,
+    std::function<std::unique_ptr<GpuIndex>(GpuResources*, int)> init);
+  faiss::Index* getIndex();
+
+  void runOnIndices(std::function<void(GpuIndex*)> f);
+  void setNumProbes(int nprobe);
+};
+
+} }
+
+#include <faiss/gpu/perf/IndexWrapper-inl.h>

data/vendor/faiss/gpu/perf/PerfClustering.cpp

@@ -0,0 +1,115 @@
+/**
+ * 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/utils/random.h>
+#include <faiss/Clustering.h>
+#include <faiss/gpu/GpuIndexFlat.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/perf/IndexWrapper.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/gpu/utils/Timer.h>
+#include <gflags/gflags.h>
+#include <memory>
+#include <vector>
+
+#include <cuda_profiler_api.h>
+
+DEFINE_int32(num, 10000, "# of vecs");
+DEFINE_int32(k, 100, "# of clusters");
+DEFINE_int32(dim, 128, "# of dimensions");
+DEFINE_int32(niter, 10, "# of iterations");
+DEFINE_bool(L2_metric, true, "If true, use L2 metric. If false, use IP metric");
+DEFINE_bool(use_float16, false, "use float16 vectors and math");
+DEFINE_bool(transposed, false, "transposed vector storage");
+DEFINE_bool(verbose, false, "turn on clustering logging");
+DEFINE_int64(seed, -1, "specify random seed");
+DEFINE_int32(num_gpus, 1, "number of gpus to use");
+DEFINE_int64(min_paging_size, -1, "minimum size to use CPU -> GPU paged copies");
+DEFINE_int64(pinned_mem, -1, "pinned memory allocation to use");
+DEFINE_int32(max_points, -1, "max points per centroid");
+
+using namespace faiss::gpu;
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+
+  cudaProfilerStop();
+
+  auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
+  printf("using seed %ld\n", seed);
+
+  std::vector<float> vecs((size_t) FLAGS_num * FLAGS_dim);
+  faiss::float_rand(vecs.data(), vecs.size(), seed);
+
+  printf("K-means metric %s dim %d centroids %d num train %d niter %d\n",
+         FLAGS_L2_metric ? "L2" : "IP",
+         FLAGS_dim, FLAGS_k, FLAGS_num, FLAGS_niter);
+  printf("float16 math %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
+  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) ->
+    std::unique_ptr<faiss::gpu::GpuIndexFlat> {
+    if (FLAGS_pinned_mem >= 0) {
+      ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(
+        FLAGS_pinned_mem);
+    }
+
+    GpuIndexFlatConfig config;
+    config.device = dev;
+    config.useFloat16 = FLAGS_use_float16;
+    config.storeTransposed = FLAGS_transposed;
+
+    auto p = std::unique_ptr<faiss::gpu::GpuIndexFlat>(
+      FLAGS_L2_metric ?
+      (faiss::gpu::GpuIndexFlat*)
+      new faiss::gpu::GpuIndexFlatL2(res, FLAGS_dim, config) :
+      (faiss::gpu::GpuIndexFlat*)
+      new faiss::gpu::GpuIndexFlatIP(res, FLAGS_dim, config));
+
+    if (FLAGS_min_paging_size >= 0) {
+      p->setMinPagingSize(FLAGS_min_paging_size);
+    }
+    return p;
+  };
+
+  IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
+
+  CUDA_VERIFY(cudaProfilerStart());
+  faiss::gpu::synchronizeAllDevices();
+
+  float gpuTime = 0.0f;
+
+  faiss::ClusteringParameters cp;
+  cp.niter = FLAGS_niter;
+  cp.verbose = FLAGS_verbose;
+
+  if (FLAGS_max_points > 0) {
+    cp.max_points_per_centroid = FLAGS_max_points;
+  }
+
+  faiss::Clustering kmeans(FLAGS_dim, FLAGS_k, cp);
+
+  // Time k-means
+  {
+    CpuTimer timer;
+
+    kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));
+
+    // There is a device -> host copy above, so no need to time
+    // additional synchronization with the GPU
+    gpuTime = timer.elapsedMilliseconds();
+  }
+
+  CUDA_VERIFY(cudaProfilerStop());
+  printf("k-means time %.3f ms\n", gpuTime);
+
+  CUDA_VERIFY(cudaDeviceSynchronize());
+
+  return 0;
+}

data/vendor/faiss/gpu/perf/PerfIVFPQAdd.cpp

@@ -0,0 +1,139 @@
+/**
+ * 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 <cuda_profiler_api.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/IndexIVFPQ.h>
+#include <faiss/gpu/GpuIndexIVFPQ.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/test/TestUtils.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/gpu/utils/Timer.h>
+#include <gflags/gflags.h>
+#include <map>
+#include <vector>
+
+DEFINE_int32(batches, 10, "number of batches of vectors to add");
+DEFINE_int32(batch_size, 10000, "number of vectors in each batch");
+DEFINE_int32(dim, 256, "dimension of vectors");
+DEFINE_int32(centroids, 4096, "num coarse centroids to use");
+DEFINE_int32(bytes_per_vec, 32, "bytes per encoded vector");
+DEFINE_int32(bits_per_code, 8, "bits per PQ code");
+DEFINE_int32(index, 2, "0 = no indices on GPU; 1 = 32 bit, 2 = 64 bit on GPU");
+DEFINE_bool(time_gpu, true, "time add to GPU");
+DEFINE_bool(time_cpu, false, "time add to CPU");
+DEFINE_bool(per_batch_time, false, "print per-batch times");
+DEFINE_bool(reserve_memory, false, "whether or not to pre-reserve memory");
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+
+  cudaProfilerStop();
+
+  int dim = FLAGS_dim;
+  int numCentroids = FLAGS_centroids;
+  int bytesPerVec = FLAGS_bytes_per_vec;
+  int bitsPerCode = FLAGS_bits_per_code;
+
+  faiss::gpu::StandardGpuResources res;
+
+  // IndexIVFPQ will complain, but just give us enough to get through this
+  int numTrain = 4 * numCentroids;
+  std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
+
+  faiss::IndexFlatL2 coarseQuantizer(dim);
+  faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, dim, numCentroids,
+                             bytesPerVec, bitsPerCode);
+  if (FLAGS_time_cpu) {
+    cpuIndex.train(numTrain, trainVecs.data());
+  }
+
+  faiss::gpu::GpuIndexIVFPQConfig config;
+  config.device = 0;
+  config.indicesOptions = (faiss::gpu::IndicesOptions) FLAGS_index;
+
+  faiss::gpu::GpuIndexIVFPQ gpuIndex(
+    &res, dim, numCentroids, bytesPerVec, bitsPerCode,
+    faiss::METRIC_L2, config);
+
+  if (FLAGS_time_gpu) {
+    gpuIndex.train(numTrain, trainVecs.data());
+    if (FLAGS_reserve_memory) {
+      size_t numVecs = (size_t) FLAGS_batches * (size_t) FLAGS_batch_size;
+      gpuIndex.reserveMemory(numVecs);
+    }
+  }
+
+  cudaDeviceSynchronize();
+  CUDA_VERIFY(cudaProfilerStart());
+
+  float totalGpuTime = 0.0f;
+  float totalCpuTime = 0.0f;
+
+  for (int i = 0; i < FLAGS_batches; ++i) {
+    if (!FLAGS_per_batch_time) {
+      if (i % 10 == 0) {
+        printf("Adding batch %d\n", i + 1);
+      }
+    }
+
+    auto addVecs = faiss::gpu::randVecs(FLAGS_batch_size, dim);
+
+    if (FLAGS_time_gpu) {
+      faiss::gpu::CpuTimer timer;
+      gpuIndex.add(FLAGS_batch_size, addVecs.data());
+      CUDA_VERIFY(cudaDeviceSynchronize());
+      auto time = timer.elapsedMilliseconds();
+
+      totalGpuTime += time;
+
+      if (FLAGS_per_batch_time) {
+      printf("Batch %d | GPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
+             i + 1, FLAGS_batch_size, time, time / (float) FLAGS_batch_size);
+      }
+    }
+
+    if (FLAGS_time_cpu) {
+      faiss::gpu::CpuTimer timer;
+      cpuIndex.add(FLAGS_batch_size, addVecs.data());
+      auto time = timer.elapsedMilliseconds();
+
+      totalCpuTime += time;
+
+      if (FLAGS_per_batch_time) {
+        printf("Batch %d | CPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
+               i + 1, FLAGS_batch_size, time, time / (float) FLAGS_batch_size);
+      }
+    }
+  }
+
+  CUDA_VERIFY(cudaProfilerStop());
+
+  int total = FLAGS_batch_size * FLAGS_batches;
+
+  if (FLAGS_time_gpu) {
+    printf("%d dim, %d centroids, %d x %d encoding\n"
+           "GPU time to add %d vectors (%d batches, %d per batch): "
+           "%.3f ms (%.3f us per)\n",
+           dim, numCentroids, bytesPerVec, bitsPerCode,
+           total, FLAGS_batches, FLAGS_batch_size,
+           totalGpuTime, totalGpuTime * 1000.0f / (float) total);
+  }
+
+  if (FLAGS_time_cpu) {
+    printf("%d dim, %d centroids, %d x %d encoding\n"
+           "CPU time to add %d vectors (%d batches, %d per batch): "
+           "%.3f ms (%.3f us per)\n",
+           dim, numCentroids, bytesPerVec, bitsPerCode,
+           total, FLAGS_batches, FLAGS_batch_size,
+           totalCpuTime, totalCpuTime * 1000.0f / (float) total);
+  }
+
+  return 0;
+}