faiss 0.3.1 → 0.3.3

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -26,6 +26,12 @@ struct GpuIndexIVFConfig : public GpuIndexConfig {
 
     /// Configuration for the coarse quantizer object
     GpuIndexFlatConfig flatConfig;
+
+    /// This flag controls the CPU fallback logic for coarse quantizer
+    /// component of the index. When set to false (default), the cloner will
+    /// throw an exception for indices not implemented on GPU. When set to
+    /// true, it will fallback to a CPU implementation.
+    bool allowCpuCoarseQuantizer = false;
 };
 
 /// Base class of all GPU IVF index types. This (for now) deliberately does not

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -87,6 +87,8 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
     /// Trains the coarse quantizer based on the given vector data
     void train(idx_t n, const float* x) override;
 
+    void reconstruct_n(idx_t i0, idx_t n, float* out) const override;
+
    protected:
     /// Initialize appropriate index
     void setIndex_(

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -34,7 +34,7 @@ struct GpuIndexIVFPQConfig : public GpuIndexIVFConfig {
 
     /// Use the alternative memory layout for the IVF lists
     /// WARNING: this is a feature under development, and is only supported with
-    /// RAFT enabled for the index. Do not use if RAFT is not enabled.
+    /// cuVS enabled for the index. Do not use if cuVS is not enabled.
     bool interleavedLayout = false;
 
     /// Use GEMM-backed computation of PQ code distances for the no precomputed

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,3 +1,4 @@
+// @lint-ignore-every LICENSELINT
 /**
  * Copyright (c) Facebook, Inc. and its affiliates.
  *
@@ -5,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -160,6 +161,10 @@ GpuMemoryReservation::~GpuMemoryReservation() {
 
 GpuResources::~GpuResources() = default;
 
+bool GpuResources::supportsBFloat16CurrentDevice() {
+    return supportsBFloat16(getCurrentDevice());
+}
+
 cublasHandle_t GpuResources::getBlasHandleCurrentDevice() {
     return getBlasHandle(getCurrentDevice());
 }
@@ -168,7 +173,7 @@ cudaStream_t GpuResources::getDefaultStreamCurrentDevice() {
     return getDefaultStream(getCurrentDevice());
 }
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
 raft::device_resources& GpuResources::getRaftHandleCurrentDevice() {
     return getRaftHandle(getCurrentDevice());
 }

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

@@ -1,3 +1,4 @@
+// @lint-ignore-every LICENSELINT
 /**
  * Copyright (c) Facebook, Inc. and its affiliates.
  *
@@ -5,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -30,7 +31,7 @@
 #include <utility>
 #include <vector>
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 #endif
@@ -161,7 +162,7 @@ struct AllocRequest : public AllocInfo {
     /// The size in bytes of the allocation
     size_t size = 0;
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
     rmm::mr::device_memory_resource* mr = nullptr;
 #endif
 };
@@ -204,6 +205,9 @@ class GpuResources {
     /// of demand
     virtual void initializeForDevice(int device) = 0;
 
+    /// Does the given GPU support bfloat16?
+    virtual bool supportsBFloat16(int device) = 0;
+
     /// Returns the cuBLAS handle that we use for the given device
     virtual cublasHandle_t getBlasHandle(int device) = 0;
 
@@ -211,7 +215,7 @@ class GpuResources {
     /// given device
     virtual cudaStream_t getDefaultStream(int device) = 0;
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
     /// Returns the raft handle for the given device which can be used to
     /// make calls to other raft primitives.
     virtual raft::device_resources& getRaftHandle(int device) = 0;
@@ -251,6 +255,9 @@ class GpuResources {
     /// Functions provided by default
     ///
 
+    /// Does the current GPU support bfloat16?
+    bool supportsBFloat16CurrentDevice();
+
     /// Calls getBlasHandle with the current device
     cublasHandle_t getBlasHandleCurrentDevice();
 

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

@@ -1,3 +1,4 @@
+// @lint-ignore-every LICENSELINT
 /**
  * Copyright (c) Facebook, Inc. and its affiliates.
  *
@@ -5,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,7 +21,7 @@
  * limitations under the License.
  */
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
 #include <rmm/mr/device/managed_memory_resource.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
@@ -90,7 +91,7 @@ std::string allocsToString(const std::unordered_map<void*, AllocRequest>& map) {
 
 StandardGpuResourcesImpl::StandardGpuResourcesImpl()
         :
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
           mmr_(new rmm::mr::managed_memory_resource),
           pmr_(new rmm::mr::pinned_memory_resource),
 #endif
@@ -129,6 +130,10 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
     FAISS_ASSERT_MSG(
             !allocError, "GPU memory allocations not properly cleaned up");
 
+#if defined USE_NVIDIA_CUVS
+    raftHandles_.clear();
+#endif
+
     for (auto& entry : defaultStreams_) {
         DeviceScope scope(entry.first);
 
@@ -158,7 +163,7 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
     }
 
     if (pinnedMemAlloc_) {
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
         pmr_->deallocate(pinnedMemAlloc_, pinnedMemAllocSize_);
 #else
         auto err = cudaFreeHost(pinnedMemAlloc_);
@@ -201,6 +206,13 @@ size_t StandardGpuResourcesImpl::getDefaultTempMemForGPU(
     return requested;
 }
 
+/// Does the given GPU support bfloat16?
+bool StandardGpuResourcesImpl::supportsBFloat16(int device) {
+    initializeForDevice(device);
+    auto& prop = getDeviceProperties(device);
+    return prop.major >= 8;
+}
+
 void StandardGpuResourcesImpl::noTempMemory() {
     setTempMemory(0);
 }
@@ -257,6 +269,14 @@ void StandardGpuResourcesImpl::setDefaultStream(
         if (prevStream != stream) {
             streamWait({stream}, {prevStream});
         }
+#if defined USE_NVIDIA_CUVS
+        // delete the raft handle for this device, which will be initialized
+        // with the updated stream during any subsequent calls to getRaftHandle
+        auto it2 = raftHandles_.find(device);
+        if (it2 != raftHandles_.end()) {
+            raft::resource::set_cuda_stream(it2->second, stream);
+        }
+#endif
     }
 
     userDefaultStreams_[device] = stream;
@@ -274,6 +294,24 @@ void StandardGpuResourcesImpl::revertDefaultStream(int device) {
             cudaStream_t newStream = defaultStreams_[device];
 
             streamWait({newStream}, {prevStream});
+
+#if defined USE_NVIDIA_CUVS
+            // update the stream on the raft handle for this device
+            auto it2 = raftHandles_.find(device);
+            if (it2 != raftHandles_.end()) {
+                raft::resource::set_cuda_stream(it2->second, newStream);
+            }
+#endif
+        } else {
+#if defined USE_NVIDIA_CUVS
+            // delete the raft handle for this device, which will be initialized
+            // with the updated stream during any subsequent calls to
+            // getRaftHandle
+            auto it2 = raftHandles_.find(device);
+            if (it2 != raftHandles_.end()) {
+                raftHandles_.erase(it2);
+            }
+#endif
         }
     }
 
@@ -307,7 +345,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
     // If this is the first device that we're initializing, create our
     // pinned memory allocation
     if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
         // If this is the first device that we're initializing, create our
         // pinned memory allocation
         if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
@@ -347,11 +385,20 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
             prop.major,
             prop.minor);
 
+#if USE_AMD_ROCM
+    // Our code is pre-built with and expects warpSize == 32 or 64, validate
+    // that
+    FAISS_ASSERT_FMT(
+            prop.warpSize == 32 || prop.warpSize == 64,
+            "Device id %d does not have expected warpSize of 32 or 64",
+            device);
+#else
     // Our code is pre-built with and expects warpSize == 32, validate that
     FAISS_ASSERT_FMT(
             prop.warpSize == 32,
             "Device id %d does not have expected warpSize of 32",
             device);
+#endif
 
     // Create streams
     cudaStream_t defaultStream = nullptr;
@@ -360,7 +407,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
 
     defaultStreams_[device] = defaultStream;
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
     raftHandles_.emplace(std::make_pair(device, defaultStream));
 #endif
 
@@ -426,7 +473,7 @@ cudaStream_t StandardGpuResourcesImpl::getDefaultStream(int device) {
     return defaultStreams_[device];
 }
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
 raft::device_resources& StandardGpuResourcesImpl::getRaftHandle(int device) {
     initializeForDevice(device);
 
@@ -497,7 +544,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
         // Otherwise, we can handle this locally
         p = tempMemory_[adjReq.device]->allocMemory(adjReq.stream, adjReq.size);
     } else if (adjReq.space == MemorySpace::Device) {
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
         try {
             rmm::mr::device_memory_resource* current_mr =
                     rmm::mr::get_per_device_resource(
@@ -531,7 +578,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
         }
 #endif
     } else if (adjReq.space == MemorySpace::Unified) {
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
         try {
             // for now, use our own managed MR to do Unified Memory allocations.
             // TODO: change this to use the current device resource once RMM has
@@ -600,7 +647,7 @@ void StandardGpuResourcesImpl::deallocMemory(int device, void* p) {
     } else if (
             req.space == MemorySpace::Device ||
             req.space == MemorySpace::Unified) {
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
         req.mr->deallocate_async(p, req.size, req.stream);
 #else
         auto err = cudaFree(p);
@@ -661,6 +708,14 @@ std::shared_ptr<GpuResources> StandardGpuResources::getResources() {
     return res_;
 }
 
+bool StandardGpuResources::supportsBFloat16(int device) {
+    return res_->supportsBFloat16(device);
+}
+
+bool StandardGpuResources::supportsBFloat16CurrentDevice() {
+    return res_->supportsBFloat16CurrentDevice();
+}
+
 void StandardGpuResources::noTempMemory() {
     res_->noTempMemory();
 }
@@ -694,7 +749,7 @@ cudaStream_t StandardGpuResources::getDefaultStream(int device) {
     return res_->getDefaultStream(device);
 }
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
 raft::device_resources& StandardGpuResources::getRaftHandle(int device) {
     return res_->getRaftHandle(device);
 }

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

@@ -1,3 +1,4 @@
+// @lint-ignore-every LICENSELINT
 /**
  * Copyright (c) Facebook, Inc. and its affiliates.
  *
@@ -5,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,7 +23,7 @@
 
 #pragma once
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
 #include <rmm/mr/host/pinned_memory_resource.hpp>
 #endif
@@ -47,6 +48,9 @@ class StandardGpuResourcesImpl : public GpuResources {
 
     ~StandardGpuResourcesImpl() override;
 
+    /// Does the given GPU support bfloat16?
+    bool supportsBFloat16(int device) override;
+
     /// Disable allocation of temporary memory; all temporary memory
     /// requests will call cudaMalloc / cudaFree at the point of use
     void noTempMemory();
@@ -79,7 +83,7 @@ class StandardGpuResourcesImpl : public GpuResources {
     /// this stream upon exit from an index or other Faiss GPU call.
     cudaStream_t getDefaultStream(int device) override;
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
     /// Returns the raft handle for the given device which can be used to
     /// make calls to other raft primitives.
     raft::device_resources& getRaftHandle(int device) override;
@@ -151,7 +155,7 @@ class StandardGpuResourcesImpl : public GpuResources {
     /// cuBLAS handle for each device
     std::unordered_map<int, cublasHandle_t> blasHandles_;
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
     /// raft handle for each device
     std::unordered_map<int, raft::device_resources> raftHandles_;
 
@@ -198,6 +202,12 @@ class StandardGpuResources : public GpuResourcesProvider {
 
     std::shared_ptr<GpuResources> getResources() override;
 
+    /// Whether or not the given device supports native bfloat16 arithmetic
+    bool supportsBFloat16(int device);
+
+    /// Whether or not the current device supports native bfloat16 arithmetic
+    bool supportsBFloat16CurrentDevice();
+
     /// Disable allocation of temporary memory; all temporary memory
     /// requests will call cudaMalloc / cudaFree at the point of use
     void noTempMemory();
@@ -234,7 +244,7 @@ class StandardGpuResources : public GpuResourcesProvider {
     /// Returns the current default stream
     cudaStream_t getDefaultStream(int device);
 
-#if defined USE_NVIDIA_RAFT
+#if defined USE_NVIDIA_CUVS
     /// Returns the raft handle for the given device which can be used to
     /// make calls to other raft primitives.
     raft::device_resources& getRaftHandle(int device);

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,11 +1,12 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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/gpu/utils/DeviceUtils.h>
 #include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/impl/FaissAssert.h>
 
@@ -166,15 +167,16 @@ void unpackInterleavedWord(
         int numVecs,
         int dims,
         int bitsPerCode) {
-    int wordsPerDimBlock = 32 * bitsPerCode / (8 * sizeof(T));
+    int warpSize = getWarpSizeCurrentDevice();
+    int wordsPerDimBlock = (size_t)warpSize * bitsPerCode / (8 * sizeof(T));
     int wordsPerBlock = wordsPerDimBlock * dims;
-    int numBlocks = utils::divUp(numVecs, 32);
+    int numBlocks = utils::divUp(numVecs, warpSize);
 
 #pragma omp parallel for
     for (int i = 0; i < numVecs; ++i) {
-        int block = i / 32;
+        int block = i / warpSize;
         FAISS_ASSERT(block < numBlocks);
-        int lane = i % 32;
+        int lane = i % warpSize;
 
         for (int j = 0; j < dims; ++j) {
             int srcOffset = block * wordsPerBlock + j * wordsPerDimBlock + lane;
@@ -188,9 +190,10 @@ std::vector<uint8_t> unpackInterleaved(
         int numVecs,
         int dims,
         int bitsPerCode) {
-    int bytesPerDimBlock = 32 * bitsPerCode / 8;
+    int warpSize = getWarpSizeCurrentDevice();
+    int bytesPerDimBlock = warpSize * bitsPerCode / 8;
     int bytesPerBlock = bytesPerDimBlock * dims;
-    int numBlocks = utils::divUp(numVecs, 32);
+    int numBlocks = utils::divUp(numVecs, warpSize);
     size_t totalSize = (size_t)bytesPerBlock * numBlocks;
     FAISS_ASSERT(data.size() == totalSize);
 
@@ -217,8 +220,8 @@ std::vector<uint8_t> unpackInterleaved(
     } else if (bitsPerCode == 4) {
 #pragma omp parallel for
         for (int i = 0; i < numVecs; ++i) {
-            int block = i / 32;
-            int lane = i % 32;
+            int block = i / warpSize;
+            int lane = i % warpSize;
 
             int word = lane / 2;
             int subWord = lane % 2;
@@ -235,8 +238,8 @@ std::vector<uint8_t> unpackInterleaved(
     } else if (bitsPerCode == 5) {
 #pragma omp parallel for
         for (int i = 0; i < numVecs; ++i) {
-            int block = i / 32;
-            int blockVector = i % 32;
+            int block = i / warpSize;
+            int blockVector = i % warpSize;
 
             for (int j = 0; j < dims; ++j) {
                 uint8_t* dimBlock =
@@ -257,8 +260,8 @@ std::vector<uint8_t> unpackInterleaved(
     } else if (bitsPerCode == 6) {
 #pragma omp parallel for
         for (int i = 0; i < numVecs; ++i) {
-            int block = i / 32;
-            int blockVector = i % 32;
+            int block = i / warpSize;
+            int blockVector = i % warpSize;
 
             for (int j = 0; j < dims; ++j) {
                 uint8_t* dimBlock =
@@ -442,17 +445,18 @@ void packInterleavedWord(
         int numVecs,
         int dims,
         int bitsPerCode) {
-    int wordsPerDimBlock = 32 * bitsPerCode / (8 * sizeof(T));
+    int warpSize = getWarpSizeCurrentDevice();
+    int wordsPerDimBlock = (size_t)warpSize * bitsPerCode / (8 * sizeof(T));
     int wordsPerBlock = wordsPerDimBlock * dims;
-    int numBlocks = utils::divUp(numVecs, 32);
+    int numBlocks = utils::divUp(numVecs, warpSize);
 
     // 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;
+        int block = i / warpSize;
         FAISS_ASSERT(block < numBlocks);
-        int lane = i % 32;
+        int lane = i % warpSize;
 
         for (int j = 0; j < dims; ++j) {
             int dstOffset = block * wordsPerBlock + j * wordsPerDimBlock + lane;
@@ -466,9 +470,10 @@ std::vector<uint8_t> packInterleaved(
         int numVecs,
         int dims,
         int bitsPerCode) {
-    int bytesPerDimBlock = 32 * bitsPerCode / 8;
+    int warpSize = getWarpSizeCurrentDevice();
+    int bytesPerDimBlock = warpSize * bitsPerCode / 8;
     int bytesPerBlock = bytesPerDimBlock * dims;
-    int numBlocks = utils::divUp(numVecs, 32);
+    int numBlocks = utils::divUp(numVecs, warpSize);
     size_t totalSize = (size_t)bytesPerBlock * numBlocks;
 
     // bit codes padded to whole bytes
@@ -499,7 +504,7 @@ std::vector<uint8_t> packInterleaved(
         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 loVec = i * warpSize + k * 2;
                     int hiVec = loVec + 1;
 
                     uint8_t lo = loVec < numVecs ? data[loVec * dims + j] : 0;
@@ -516,7 +521,7 @@ std::vector<uint8_t> packInterleaved(
             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 loVec = i * warpSize + (k * 8) / 5;
                     int hiVec = loVec + 1;
                     int hiVec2 = hiVec + 1;
 
@@ -536,7 +541,7 @@ std::vector<uint8_t> packInterleaved(
             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 loVec = i * warpSize + (k * 8) / 6;
                     int hiVec = loVec + 1;
 
                     uint8_t lo = loVec < numVecs ? data[loVec * dims + j] : 0;

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -17,6 +17,7 @@
 #include <vector>
 
 #include <cuda_profiler_api.h>
+#include <faiss/impl/AuxIndexStructures.h>
 
 DEFINE_int32(num, 10000, "# of vecs");
 DEFINE_int32(k, 100, "# of clusters");
@@ -34,6 +35,7 @@ DEFINE_int64(
         "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");
+DEFINE_double(timeout, 0, "timeout in seconds");
 
 using namespace faiss::gpu;
 
@@ -99,10 +101,14 @@ int main(int argc, char** argv) {
         cp.max_points_per_centroid = FLAGS_max_points;
     }
 
+    auto tc = new faiss::TimeoutCallback();
+    faiss::InterruptCallback::instance.reset(tc);
+
     faiss::Clustering kmeans(FLAGS_dim, FLAGS_k, cp);
 
     // Time k-means
     {
+        tc->set_timeout(FLAGS_timeout);
         CpuTimer timer;
 
         kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));