faiss 0.1.3 → 0.1.4

data/vendor/faiss/faiss/gpu/StandardGpuResources.h

@@ -41,8 +41,22 @@ class StandardGpuResourcesImpl : public GpuResources {
   /// 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 stream for work ordering. We do not own `stream`;
+  /// i.e., it will not be destroyed when the GpuResources object gets cleaned
+  /// up.
+  /// We are guaranteed that all Faiss GPU work is ordered with respect to
+  /// this stream upon exit from an index or other Faiss GPU call.
+  void setDefaultStream(int device, cudaStream_t stream) override;
+
+  /// Revert the default stream to the original stream managed by this resources
+  /// object, in case someone called `setDefaultStream`.
+  void revertDefaultStream(int device);
+
+  /// Returns the stream for the given device on which all Faiss GPU work is
+  /// ordered.
+  /// We are guaranteed that all Faiss GPU work is ordered with respect to
+  /// this stream upon exit from an index or other Faiss GPU call.
+  cudaStream_t getDefaultStream(int device) override;
 
   /// Called to change the work ordering streams to the null stream
   /// for all devices
@@ -60,8 +74,6 @@ class StandardGpuResourcesImpl : public GpuResources {
 
   cublasHandle_t getBlasHandle(int device) override;
 
-  cudaStream_t getDefaultStream(int device) override;
-
   std::vector<cudaStream_t> getAlternateStreams(int device) override;
 
   /// Allocate non-temporary GPU memory
@@ -128,7 +140,9 @@ class StandardGpuResourcesImpl : public GpuResources {
 };
 
 /// Default implementation of GpuResources that allocates a cuBLAS
-/// stream and 2 streams for use, as well as temporary memory
+/// stream and 2 streams for use, as well as temporary memory.
+/// Internally, the Faiss GPU code uses the instance managed by getResources,
+/// but this is the user-facing object that is internally reference counted.
 class StandardGpuResources : public GpuResourcesProvider {
  public:
   StandardGpuResources();
@@ -151,9 +165,17 @@ class StandardGpuResources : public GpuResourcesProvider {
   /// transfers
   void setPinnedMemory(size_t size);
 
-  /// Called to change the stream for work ordering
+  /// Called to change the stream for work ordering. We do not own `stream`;
+  /// i.e., it will not be destroyed when the GpuResources object gets cleaned
+  /// up.
+  /// We are guaranteed that all Faiss GPU work is ordered with respect to
+  /// this stream upon exit from an index or other Faiss GPU call.
   void setDefaultStream(int device, cudaStream_t stream);
 
+  /// Revert the default stream to the original stream managed by this resources
+  /// object, in case someone called `setDefaultStream`.
+  void revertDefaultStream(int device);
+
   /// Called to change the work ordering streams to the null stream
   /// for all devices
   void setDefaultNullStreamAllDevices();

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

@@ -0,0 +1,547 @@
+/**
+ * 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/InterleavedCodes.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/gpu/utils/StaticUtils.h>
+
+namespace faiss { namespace gpu {
+
+inline uint8_t unpack5(int i, uint8_t vLower, uint8_t vUpper) {
+  uint8_t v = 0;
+
+  // lsb     ...    msb
+  // 0: 0 0 0 0 0 1 1 1
+  // 1: 1 1 2 2 2 2 2 3
+  // 2: 3 3 3 3 4 4 4 4
+  // 3: 4 5 5 5 5 5 6 6
+  // 4: 6 6 6 7 7 7 7 7
+  switch (i % 8) {
+    case 0:
+      // 5 lsbs of lower
+      v = vLower & 0x1f;
+      break;
+    case 1:
+      // 3 msbs of lower as v lsbs
+      // 2 msbs of upper as v msbs
+      v = (vLower >> 5) | ((vUpper & 0x3) << 3);
+      break;
+    case 2:
+      // 5 of lower
+      v = (vLower >> 2) & 0x1f;
+      break;
+    case 3:
+      // 1 msbs of lower as v lsbs
+      // 4 lsbs of upper as v msbs
+      v = (vLower >> 7) | ((vUpper & 0xf) << 1);
+      break;
+    case 4:
+      // 4 msbs of lower as v lsbs
+      // 1 lsbs of upper as v msbs
+      v = (vLower >> 4) | ((vUpper & 0x1) << 4);
+      break;
+    case 5:
+      // 5 of lower
+      v = (vLower >> 1) & 0x1f;
+      break;
+    case 6:
+      // 2 msbs of lower as v lsbs
+      // 3 lsbs of upper as v msbs
+      v = (vLower >> 6) | ((vUpper & 0x7) << 2);
+      break;
+    case 7:
+      // 5 of lower
+      v = (vLower >> 3);
+      break;
+  }
+
+  return v;
+}
+
+inline uint8_t unpack6(int i, uint8_t vLower, uint8_t vUpper) {
+  uint8_t v = 0;
+
+  switch (i % 4) {
+    case 0:
+      // 6 lsbs of lower
+      v = vLower & 0x3f;
+      break;
+    case 1:
+      // 2 msbs of lower as v lsbs
+      // 4 lsbs of upper as v msbs
+      v = (vLower >> 6) | ((vUpper & 0xf) << 2);
+      break;
+    case 2:
+      // 4 msbs of lower as v lsbs
+      // 2 lsbs of upper as v msbs
+      v = (vLower >> 4) | ((vUpper & 0x3) << 4);
+      break;
+    case 3:
+      // 6 msbs of lower
+      v = (vLower >> 2);
+      break;
+  }
+
+  return v;
+}
+
+
+std::vector<uint8_t>
+unpackNonInterleaved(std::vector<uint8_t> data,
+                     int numVecs,
+                     int dims,
+                     int bitsPerCode) {
+  int srcVecSize = utils::divUp(dims * bitsPerCode, 8);
+  FAISS_ASSERT(data.size() == numVecs * srcVecSize);
+
+  if (bitsPerCode == 8 ||
+      bitsPerCode == 16 ||
+      bitsPerCode == 32) {
+    // nothing to do
+    return data;
+  }
+
+  // bit codes padded to whole bytes
+  std::vector<uint8_t> out(numVecs * dims * utils::divUp(bitsPerCode, 8));
+
+  if (bitsPerCode == 4) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      for (int j = 0; j < dims; ++j) {
+        int srcIdx = i * srcVecSize + (j / 2);
+        FAISS_ASSERT(srcIdx < data.size());
+
+        uint8_t v = data[srcIdx];
+        v = (j % 2 == 0) ? v & 0xf : v >> 4;
+
+        out[i * dims + j] = v;
+      }
+    }
+  } else if (bitsPerCode == 5) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      for (int j = 0; j < dims; ++j) {
+        int lo = i * srcVecSize + (j * 5) / 8;
+        int hi = lo + 1;
+
+        FAISS_ASSERT(lo < data.size());
+        FAISS_ASSERT(hi <= data.size());
+
+        auto vLower = data[lo];
+        auto vUpper = hi < data.size() ? data[hi] : 0;
+
+        out[i * dims + j] = unpack5(j, vLower, vUpper);
+      }
+    }
+  } else if (bitsPerCode == 6) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      for (int j = 0; j < dims; ++j) {
+        int lo = i * srcVecSize + (j * 6) / 8;
+        int hi = lo + 1;
+
+        FAISS_ASSERT(lo < data.size());
+        FAISS_ASSERT(hi <= data.size());
+
+        auto vLower = data[lo];
+        auto vUpper = hi < data.size() ? data[hi] : 0;
+
+        out[i * dims + j] = unpack6(j, vLower, vUpper);
+      }
+    }
+  } else {
+    // unhandled
+    FAISS_ASSERT(false);
+  }
+
+  return out;
+}
+
+template <typename T>
+void
+unpackInterleavedWord(const T* in,
+                      T* out,
+                      int numVecs,
+                      int dims,
+                      int bitsPerCode) {
+  int wordsPerDimBlock = 32 * bitsPerCode / (8 * sizeof(T));
+  int wordsPerBlock = wordsPerDimBlock * dims;
+  int numBlocks = utils::divUp(numVecs, 32);
+
+#pragma omp parallel for
+  for (int i = 0; i < numVecs; ++i) {
+    int block = i / 32;
+    FAISS_ASSERT(block < numBlocks);
+    int lane = i % 32;
+
+    for (int j = 0; j < dims; ++j) {
+      int srcOffset = block * wordsPerBlock + j * wordsPerDimBlock + lane;
+      out[i * dims + j] = in[srcOffset];
+    }
+  }
+}
+
+std::vector<uint8_t>
+unpackInterleaved(std::vector<uint8_t> data,
+                  int numVecs,
+                  int dims,
+                  int bitsPerCode) {
+  int bytesPerDimBlock = 32 * bitsPerCode / 8;
+  int bytesPerBlock = bytesPerDimBlock * dims;
+  int numBlocks = utils::divUp(numVecs, 32);
+  size_t totalSize = (size_t) bytesPerBlock * numBlocks;
+  FAISS_ASSERT(data.size() == totalSize);
+
+  // bit codes padded to whole bytes
+  std::vector<uint8_t> out(numVecs * dims * utils::divUp(bitsPerCode, 8));
+
+  if (bitsPerCode == 8) {
+    unpackInterleavedWord<uint8_t>(data.data(), out.data(),
+                                   numVecs, dims, bitsPerCode);
+  } else if (bitsPerCode == 16) {
+    unpackInterleavedWord<uint16_t>((uint16_t*) data.data(),
+                                    (uint16_t*) out.data(),
+                                    numVecs, dims, bitsPerCode);
+  } else if (bitsPerCode == 32) {
+    unpackInterleavedWord<uint32_t>((uint32_t*) data.data(),
+                                    (uint32_t*) out.data(),
+                                    numVecs, dims, bitsPerCode);
+  } else if (bitsPerCode == 4) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      int block = i / 32;
+      int lane = i % 32;
+
+      int word = lane / 2;
+      int subWord = lane % 2;
+
+      for (int j = 0; j < dims; ++j) {
+        auto v =
+          data[block * bytesPerBlock + j * bytesPerDimBlock + word];
+
+        v = (subWord == 0) ? v & 0xf : v >> 4;
+        out[i * dims + j] = v;
+      }
+    }
+  } else if (bitsPerCode == 5) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      int block = i / 32;
+      int blockVector = i % 32;
+
+      for (int j = 0; j < dims; ++j) {
+        uint8_t* dimBlock =
+          &data[block * bytesPerBlock + j * bytesPerDimBlock];
+
+        int lo = (blockVector * 5) / 8;
+        int hi = lo + 1;
+
+        FAISS_ASSERT(lo < bytesPerDimBlock);
+        FAISS_ASSERT(hi <= bytesPerDimBlock);
+
+        auto vLower = dimBlock[lo];
+        auto vUpper = hi < bytesPerDimBlock ? dimBlock[hi] : 0;
+
+        out[i * dims + j] = unpack5(blockVector, vLower, vUpper);
+      }
+    }
+  } else if (bitsPerCode == 6) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      int block = i / 32;
+      int blockVector = i % 32;
+
+      for (int j = 0; j < dims; ++j) {
+        uint8_t* dimBlock =
+          &data[block * bytesPerBlock + j * bytesPerDimBlock];
+
+        int lo = (blockVector * 6) / 8;
+        int hi = lo + 1;
+
+        FAISS_ASSERT(lo < bytesPerDimBlock);
+        FAISS_ASSERT(hi <= bytesPerDimBlock);
+
+        auto vLower = dimBlock[lo];
+        auto vUpper = hi < bytesPerDimBlock ? dimBlock[hi] : 0;
+
+        out[i * dims + j] = unpack6(blockVector, vLower, vUpper);
+      }
+    }
+  } else {
+    // unimplemented
+    FAISS_ASSERT(false);
+  }
+
+  return out;
+}
+
+inline uint8_t pack5(int i, uint8_t lo, uint8_t hi, uint8_t hi2) {
+  FAISS_ASSERT((lo & 0x1f) == lo);
+  FAISS_ASSERT((hi & 0x1f) == hi);
+  FAISS_ASSERT((hi2 & 0x1f) == hi2);
+
+  uint8_t v = 0;
+
+  // lsb     ...    msb
+  // 0: 0 0 0 0 0 1 1 1
+  // 1: 1 1 2 2 2 2 2 3
+  // 2: 3 3 3 3 4 4 4 4
+  // 3: 4 5 5 5 5 5 6 6
+  // 4: 6 6 6 7 7 7 7 7
+  switch (i % 5) {
+    case 0:
+      // 5 msbs of lower as vOut lsbs
+      // 3 lsbs of upper as vOut msbs
+      v = (lo & 0x1f) | (hi << 5);
+      break;
+    case 1:
+      // 2 msbs of lower as vOut lsbs
+      // 5 lsbs of upper as vOut msbs
+      // 1 lsbs of upper2 as vOut msb
+      v = (lo >> 3) | (hi << 2) | (hi2 << 7);
+      break;
+    case 2:
+      // 4 msbs of lower as vOut lsbs
+      // 4 lsbs of upper as vOut msbs
+      v = (lo >> 1) | (hi << 4);
+      break;
+    case 3:
+      // 1 msbs of lower as vOut lsbs
+      // 5 lsbs of upper as vOut msbs
+      // 2 lsbs of upper2 as vOut msb
+      v = (lo >> 4) | (hi << 1) | (hi2 << 6);
+      break;
+    case 4:
+      // 3 msbs of lower as vOut lsbs
+      // 5 lsbs of upper as vOut msbs
+      v = (lo >> 2) | (hi << 3);
+      break;
+  }
+
+  return v;
+}
+
+inline uint8_t pack6(int i, uint8_t lo, uint8_t hi) {
+  FAISS_ASSERT((lo & 0x3f) == lo);
+  FAISS_ASSERT((hi & 0x3f) == hi);
+
+  uint8_t v = 0;
+
+  // lsb     ...    msb
+  // 0: 0 0 0 0 0 0 1 1
+  // 1: 1 1 1 1 2 2 2 2
+  // 2: 2 2 3 3 3 3 3 3
+  switch (i % 3) {
+    case 0:
+      // 6 msbs of lower as vOut lsbs
+      // 2 lsbs of upper as vOut msbs
+      v = (lo & 0x3f) | (hi << 6);
+      break;
+    case 1:
+      // 4 msbs of lower as vOut lsbs
+      // 4 lsbs of upper as vOut msbs
+      v = (lo >> 2) | (hi << 4);
+      break;
+    case 2:
+      // 2 msbs of lower as vOut lsbs
+      // 6 lsbs of upper as vOut msbs
+      v = (lo >> 4) | (hi << 2);
+      break;
+  }
+
+  return v;
+}
+
+
+std::vector<uint8_t>
+packNonInterleaved(std::vector<uint8_t> data,
+                   int numVecs,
+                   int dims,
+                   int bitsPerCode) {
+  // bit codes padded to whole bytes
+  FAISS_ASSERT(data.size() == numVecs * dims * utils::divUp(bitsPerCode, 8));
+
+  if (bitsPerCode == 8 ||
+      bitsPerCode == 16 ||
+      bitsPerCode == 32) {
+    // nothing to do, whole words are already where they need to be
+    return data;
+  }
+
+  // bits packed into a whole number of bytes
+  int bytesPerVec = utils::divUp(dims * bitsPerCode, 8);
+
+  std::vector<uint8_t> out(numVecs * bytesPerVec);
+
+  if (bitsPerCode == 4) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      for (int j = 0; j < bytesPerVec; ++j) {
+        int dimLo = j * 2;
+        int dimHi = dimLo + 1;
+        FAISS_ASSERT(dimLo < dims);
+        FAISS_ASSERT(dimHi <= dims);
+
+        uint8_t lo = data[i * dims + dimLo];
+        uint8_t hi = dimHi < dims ? data[i * dims + dimHi] : 0;
+
+        out[i * bytesPerVec + j] = (hi << 4) | (lo & 0xf);
+      }
+    }
+  } else if (bitsPerCode == 5) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      for (int j = 0; j < bytesPerVec; ++j) {
+        int dimLo = (j * 8) / 5;
+        int dimHi = dimLo + 1;
+        int dimHi2 = dimHi + 1;
+        FAISS_ASSERT(dimLo < dims);
+        FAISS_ASSERT(dimHi <= dims);
+        FAISS_ASSERT(dimHi <= dims + 1);
+
+        uint8_t lo = data[i * dims + dimLo];
+        uint8_t hi = dimHi < dims ? data[i * dims + dimHi] : 0;
+        uint8_t hi2 = dimHi2 < dims ? data[i * dims + dimHi2] : 0;
+
+        out[i * bytesPerVec + j] = pack5(j, lo, hi, hi2);
+      }
+    }
+  } else if (bitsPerCode == 6) {
+#pragma omp parallel for
+    for (int i = 0; i < numVecs; ++i) {
+      for (int j = 0; j < bytesPerVec; ++j) {
+        int dimLo = (j * 8) / 6;
+        int dimHi = dimLo + 1;
+        FAISS_ASSERT(dimLo < dims);
+        FAISS_ASSERT(dimHi <= dims);
+
+        uint8_t lo = data[i * dims + dimLo];
+        uint8_t hi = dimHi < dims ? data[i * dims + dimHi] : 0;
+
+        out[i * bytesPerVec + j] = pack6(j, lo, hi);
+      }
+    }
+  } else {
+    // unhandled
+    FAISS_ASSERT(false);
+  }
+
+  return out;
+}
+
+template <typename T>
+void
+packInterleavedWord(const T* in,
+                    T* out,
+                    int numVecs,
+                    int dims,
+                    int bitsPerCode) {
+  int wordsPerDimBlock = 32 * bitsPerCode / (8 * sizeof(T));
+  int wordsPerBlock = wordsPerDimBlock * dims;
+  int numBlocks = utils::divUp(numVecs, 32);
+
+  // We're guaranteed that all other slots not filled by the vectors present are
+  // initialized to zero (from the vector constructor in packInterleaved)
+#pragma omp parallel for
+  for (int i = 0; i < numVecs; ++i) {
+    int block = i / 32;
+    FAISS_ASSERT(block < numBlocks);
+    int lane = i % 32;
+
+    for (int j = 0; j < dims; ++j) {
+      int dstOffset = block * wordsPerBlock + j * wordsPerDimBlock + lane;
+      out[dstOffset] = in[i * dims + j];
+    }
+  }
+}
+
+std::vector<uint8_t>
+packInterleaved(std::vector<uint8_t> data,
+                int numVecs,
+                int dims,
+                int bitsPerCode) {
+  int bytesPerDimBlock = 32 * bitsPerCode / 8;
+  int bytesPerBlock = bytesPerDimBlock * dims;
+  int numBlocks = utils::divUp(numVecs, 32);
+  size_t totalSize = (size_t) bytesPerBlock * numBlocks;
+
+  // bit codes padded to whole bytes
+  FAISS_ASSERT(data.size() == numVecs * dims * utils::divUp(bitsPerCode, 8));
+
+  // packs based on blocks
+  std::vector<uint8_t> out(totalSize, 0);
+
+  if (bitsPerCode == 8) {
+    packInterleavedWord<uint8_t>(data.data(), out.data(),
+                                 numVecs, dims, bitsPerCode);
+  } else if (bitsPerCode == 16) {
+    packInterleavedWord<uint16_t>((uint16_t*) data.data(),
+                                  (uint16_t*) out.data(),
+                                  numVecs, dims, bitsPerCode);
+  } else if (bitsPerCode == 32) {
+    packInterleavedWord<uint32_t>((uint32_t*) data.data(),
+                                  (uint32_t*) out.data(),
+                                  numVecs, dims, bitsPerCode);
+  } else if (bitsPerCode == 4) {
+#pragma omp parallel for
+    for (int i = 0; i < numBlocks; ++i) {
+      for (int j = 0; j < dims; ++j) {
+        for (int k = 0; k < bytesPerDimBlock; ++k) {
+          int loVec = i * 32 + k * 2;
+          int hiVec = loVec + 1;
+
+          uint8_t lo = loVec < numVecs ? data[loVec * dims + j] : 0;
+          uint8_t hi = hiVec < numVecs ? data[hiVec * dims + j] : 0;
+
+          out[i * bytesPerBlock + j * bytesPerDimBlock + k] =
+            (hi << 4) | (lo & 0xf);
+        }
+      }
+    }
+  } else if (bitsPerCode == 5) {
+#pragma omp parallel for
+    for (int i = 0; i < numBlocks; ++i) {
+      for (int j = 0; j < dims; ++j) {
+        for (int k = 0; k < bytesPerDimBlock; ++k) {
+          // What input vectors we are pulling from
+          int loVec = i * 32 + (k * 8) / 5;
+          int hiVec = loVec + 1;
+          int hiVec2 = hiVec + 1;
+
+          uint8_t lo = loVec < numVecs ? data[loVec * dims + j] : 0;
+          uint8_t hi = hiVec < numVecs ? data[hiVec * dims + j] : 0;
+          uint8_t hi2 = hiVec2 < numVecs ? data[hiVec2 * dims + j] : 0;
+
+          out[i * bytesPerBlock + j * bytesPerDimBlock + k] = pack5(k, lo, hi, hi2);
+        }
+      }
+    }
+  } else if (bitsPerCode == 6) {
+#pragma omp parallel for
+    for (int i = 0; i < numBlocks; ++i) {
+      for (int j = 0; j < dims; ++j) {
+        for (int k = 0; k < bytesPerDimBlock; ++k) {
+          // What input vectors we are pulling from
+          int loVec = i * 32 + (k * 8) / 6;
+          int hiVec = loVec + 1;
+
+          uint8_t lo = loVec < numVecs ? data[loVec * dims + j] : 0;
+          uint8_t hi = hiVec < numVecs ? data[hiVec * dims + j] : 0;
+
+          out[i * bytesPerBlock + j * bytesPerDimBlock + k] = pack6(k, lo, hi);
+        }
+      }
+    }
+  } else {
+    // unimplemented
+    FAISS_ASSERT(false);
+  }
+
+  return out;
+}
+
+} } // namespace