Complete GPU stream-event reuse and document async allocation.

faiss/gpu/GpuResources.h

@@ -76,8 +76,9 @@ enum AllocType {
     /// When using StandardGpuResources, any MemorySpace::Temporary allocations
     /// that cannot be satisfied within the TemporaryMemoryBuffer region fall
     /// back
-    /// to calling cudaMalloc which are sized to just the request at hand. These
-    /// "overflow" temporary allocations are marked with this AllocType.
+    /// to a GpuMemoryPool if configured, or to on-demand Device or Unified
+    /// allocation sized to just the request at hand. These "overflow" temporary
+    /// allocations are marked with this AllocType.
     TemporaryMemoryOverflow = 11,
 };
 
@@ -88,10 +89,10 @@ std::string allocTypeToString(AllocType t);
 enum MemorySpace {
     /// Temporary device memory (guaranteed to no longer be used upon exit of a
     /// top-level index call, and where the streams using it have completed GPU
-    /// work). Typically backed by Device memory (cudaMalloc/cudaFree).
+    /// work). Typically backed by Device memory (cudaMallocAsync/cudaFreeAsync).
     Temporary = 0,
 
-    /// Managed using cudaMalloc/cudaFree (typical GPU device memory)
+    /// Managed using cudaMallocAsync/cudaFreeAsync (typical GPU device memory)
     Device = 1,
 
     /// Managed using cudaMallocManaged/cudaFree (typical Unified CPU/GPU

faiss/gpu/StandardGpuResources.cpp

@@ -238,8 +238,8 @@ void StandardGpuResourcesImpl::setTempMemory(size_t size) {
 
         // We need to re-initialize memory resources for all current devices
         // that have been initialized. This should be safe to do, even if we are
-        // currently running work, because the cudaFree call that this implies
-        // will force-synchronize all GPUs with the CPU
+        // currently running work, because tearing down the temp memory stack
+        // implies deallocation that will force-synchronize all GPUs with the CPU
         for (auto& p : tempMemory_) {
             int device = p.first;
             // Free the existing memory first
@@ -296,7 +296,11 @@ void StandardGpuResourcesImpl::setDefaultStream(
         }
 
         if (prevStream != stream) {
-            streamWait({stream}, {prevStream});
+            addEventForStream(prevStream);
+            addEventForStream(stream);
+            std::vector<cudaStream_t> waitingStreams{stream};
+            std::vector<cudaStream_t> waitOnStreams{prevStream};
+            this->streamWait(waitingStreams, waitOnStreams);
         }
 #if defined USE_NVIDIA_CUVS
         // delete the raft handle for this device, which will be initialized
@@ -323,7 +327,11 @@ void StandardGpuResourcesImpl::revertDefaultStream(int device) {
             FAISS_ASSERT(defaultStreams_.count(device));
             cudaStream_t newStream = defaultStreams_[device];
 
-            streamWait({newStream}, {prevStream});
+            addEventForStream(prevStream);
+            addEventForStream(newStream);
+            std::vector<cudaStream_t> waitingStreams{newStream};
+            std::vector<cudaStream_t> waitOnStreams{prevStream};
+            this->streamWait(waitingStreams, waitOnStreams);
 
 #if defined USE_NVIDIA_CUVS
             // update the stream on the raft handle for this device

faiss/gpu/StandardGpuResources.h

@@ -53,7 +53,8 @@ class StandardGpuResourcesImpl : public GpuResources {
     bool supportsBFloat16(int device) override;
 
     /// Disable allocation of temporary memory; all temporary memory
-    /// requests will call cudaMalloc / cudaFree at the point of use
+    /// requests will allocate directly via the Device or Unified memory path
+    /// at the point of use
     void noTempMemory();
 
     /// Specify that we wish to use a certain fixed size of memory on
@@ -64,14 +65,14 @@ class StandardGpuResourcesImpl : public GpuResources {
     void setTempMemory(size_t size);
 
     /// Set memory space for all temporary memory allocations (both the
-    /// pool and overflow). Default is MemorySpace::Device (cudaMalloc).
+    /// pool and overflow). Default is MemorySpace::Device (cudaMallocAsync).
     /// Set to MemorySpace::Unified for cudaMallocManaged.
     /// Must be called before any device is initialized.
     void setTempMemorySpace(MemorySpace space);
 
     /// Set a memory pool to use for temporary memory overflow allocations.
-    /// If not set, we will fallback to on-demand cudaMalloc for overflow
-    /// allocations.
+    /// If not set, we will fallback to on-demand Device or Unified allocation
+    /// for overflow allocations.
     void setTempMemoryOverflowPool(GpuMemoryPool* pool);
 
     /// Set amount of pinned memory to allocate, for async GPU <-> CPU
@@ -239,7 +240,8 @@ class StandardGpuResources : public GpuResourcesProvider {
     bool supportsBFloat16CurrentDevice();
 
     /// Disable allocation of temporary memory; all temporary memory
-    /// requests will call cudaMalloc / cudaFree at the point of use
+    /// requests will allocate directly via the Device or Unified memory path
+    /// at the point of use
     void noTempMemory();
 
     /// Specify that we wish to use a certain fixed size of memory on
@@ -250,14 +252,14 @@ class StandardGpuResources : public GpuResourcesProvider {
     void setTempMemory(size_t size);
 
     /// Set memory space for all temporary memory allocations (both the
-    /// pool and overflow). Default is MemorySpace::Device (cudaMalloc).
+    /// pool and overflow). Default is MemorySpace::Device (cudaMallocAsync).
     /// Set to MemorySpace::Unified for cudaMallocManaged.
     /// Must be called before any device is initialized.
     void setTempMemorySpace(MemorySpace space);
 
     /// Set a memory pool to use for temporary memory overflow allocations.
-    /// If not set, we will fallback to on-demand cudaMalloc for overflow
-    /// allocations.
+    /// If not set, we will fallback to on-demand Device or Unified allocation
+    /// for overflow allocations.
     void setTempMemoryOverflowPool(GpuMemoryPool* pool);
 
     /// Set amount of pinned memory to allocate, for async GPU <-> CPU

faiss/gpu/test/CMakeLists.txt

@@ -43,6 +43,7 @@ faiss_gpu_test(TestGpuIndexFlat.cpp)
 faiss_gpu_test(TestGpuIndexIVFFlat.cpp)
 faiss_gpu_test(TestGpuIndexBinaryFlat.cpp)
 faiss_gpu_test(TestGpuMemoryException.cpp)
+faiss_gpu_test(TestGpuMemoryPool.cpp)
 faiss_gpu_test(TestGpuIcmEncoder.cpp)
 faiss_gpu_test(TestGpuIndexIVFPQ.cpp)
 faiss_gpu_test(TestGpuIndexIVFScalarQuantizer.cpp)

faiss/gpu/test/TestGpuMemoryPool.cpp

@@ -0,0 +1,64 @@
+/*
+ * 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/IndexFlat.h>
+#include <faiss/gpu/GpuIndexFlat.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/test/TestUtils.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/gpu/utils/MemoryPool.h>
+#include <gtest/gtest.h>
+
+// Exercises temporary-memory overflow through an externally managed GpuMemoryPool.
+TEST(TestGpuMemoryPool, TempOverflowViaPool) {
+    int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
+    const size_t kTempMem = 64 * 1024;
+    const size_t kPoolCap = 64 * 1024 * 1024;
+
+    faiss::gpu::GpuMemoryPool pool(device, kPoolCap);
+
+    faiss::gpu::StandardGpuResources res;
+    res.setTempMemory(kTempMem);
+    res.setTempMemoryOverflowPool(&pool);
+
+    const int dim = 128;
+    const int numVecs = 10000;
+    const int numQuery = 10;
+    const int k = 10;
+
+    faiss::gpu::GpuIndexFlatConfig config;
+    config.device = device;
+    config.use_cuvs = false;
+
+    faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
+    faiss::IndexFlatL2 cpuIndex(dim);
+
+    auto vecs = faiss::gpu::randVecs(numVecs, dim);
+    gpuIndex.add(numVecs, vecs.data());
+    cpuIndex.add(numVecs, vecs.data());
+
+    auto queries = faiss::gpu::randVecs(numQuery, dim);
+    compareIndices(
+            queries,
+            cpuIndex,
+            gpuIndex,
+            numQuery,
+            dim,
+            k,
+            "",
+            6e-3f,
+            0.1f,
+            0.015f);
+}
+
+int main(int argc, char** argv) {
+    testing::InitGoogleTest(&argc, argv);
+
+    faiss::gpu::setTestSeed(100);
+
+    return RUN_ALL_TESTS();
+}

faiss/gpu/utils/DeviceUtils.h

@@ -169,7 +169,9 @@ class CudaEvent {
     } while (0)
 #endif
 
-/// Call for a collection of streams to wait on
+/// Call for a collection of streams to wait on.
+/// Prefer GpuResources::streamWait when a GpuResources object is available;
+/// it reuses per-stream CUDA events instead of creating/destroying them here.
 template <typename L1, typename L2>
 void streamWaitBase(const L1& listWaiting, const L2& listWaitOn) {
     // For all the streams we are waiting on, create an event