diff --git a/faiss/gpu/GpuResources.h b/faiss/gpu/GpuResources.h index 0a0304e7f4..d4106b9acc 100644 --- a/faiss/gpu/GpuResources.h +++ b/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 diff --git a/faiss/gpu/StandardGpuResources.cpp b/faiss/gpu/StandardGpuResources.cpp index f34979c358..d537b7ebb4 100644 --- a/faiss/gpu/StandardGpuResources.cpp +++ b/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 waitingStreams{stream}; + std::vector 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 waitingStreams{newStream}; + std::vector waitOnStreams{prevStream}; + this->streamWait(waitingStreams, waitOnStreams); #if defined USE_NVIDIA_CUVS // update the stream on the raft handle for this device diff --git a/faiss/gpu/StandardGpuResources.h b/faiss/gpu/StandardGpuResources.h index 82758ad8bd..2ab2bec983 100644 --- a/faiss/gpu/StandardGpuResources.h +++ b/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 diff --git a/faiss/gpu/test/CMakeLists.txt b/faiss/gpu/test/CMakeLists.txt index bbd789dddd..1504c43527 100644 --- a/faiss/gpu/test/CMakeLists.txt +++ b/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) diff --git a/faiss/gpu/test/TestGpuMemoryPool.cpp b/faiss/gpu/test/TestGpuMemoryPool.cpp new file mode 100644 index 0000000000..7b50a1e420 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include + +// 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(); +} diff --git a/faiss/gpu/utils/DeviceUtils.h b/faiss/gpu/utils/DeviceUtils.h index 2a560afa1e..691df023da 100644 --- a/faiss/gpu/utils/DeviceUtils.h +++ b/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 void streamWaitBase(const L1& listWaiting, const L2& listWaitOn) { // For all the streams we are waiting on, create an event