mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
afe0fdc161
Facebook sync (Mar 2019) - MatrixStats object - option to round coordinates during k-means optimization - alternative option for search in HNSW - moved stats and imbalance_factor of IndexIVF to InvertedLists object - range search for IVFScalarQuantizer - direct unit8 codec in ScalarQuantizer - renamed IndexProxy to IndexReplicas and moved to main Faiss - better support for PQ code assignment with external index - support for IMI2x16 (4B virtual centroids!) - support for k = 2048 search on GPU (instead of 1024) - most CUDA mem alloc failures throw exceptions instead of terminating on an assertion - support for renaming an ondisk invertedlists - interrupt computations with ctrl-C in python
91 lines
2.5 KiB
C++
91 lines
2.5 KiB
C++
/**
|
|
* Copyright (c) 2015-present, Facebook, Inc.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the BSD+Patents license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
|
|
#include "MemorySpace.h"
|
|
#include "../../FaissAssert.h"
|
|
#include <cuda_runtime.h>
|
|
|
|
namespace faiss { namespace gpu {
|
|
|
|
/// Allocates CUDA memory for a given memory space
|
|
void allocMemorySpaceV(MemorySpace space, void** p, size_t size) {
|
|
switch (space) {
|
|
case MemorySpace::Device:
|
|
{
|
|
auto err = cudaMalloc(p, size);
|
|
|
|
// Throw if we fail to allocate
|
|
FAISS_THROW_IF_NOT_FMT(
|
|
err == cudaSuccess,
|
|
"failed to cudaMalloc %zu bytes (error %d %s)",
|
|
size, (int) err, cudaGetErrorString(err));
|
|
}
|
|
break;
|
|
case MemorySpace::Unified:
|
|
{
|
|
#ifdef FAISS_UNIFIED_MEM
|
|
auto err = cudaMallocManaged(p, size);
|
|
|
|
// Throw if we fail to allocate
|
|
FAISS_THROW_IF_NOT_FMT(
|
|
err == cudaSuccess,
|
|
"failed to cudaMallocManaged %zu bytes (error %d %s)",
|
|
size, (int) err, cudaGetErrorString(err));
|
|
#else
|
|
FAISS_THROW_MSG("Attempting to allocate via cudaMallocManaged "
|
|
"without CUDA 8+ support");
|
|
#endif
|
|
}
|
|
break;
|
|
case MemorySpace::HostPinned:
|
|
{
|
|
auto err = cudaHostAlloc(p, size, cudaHostAllocDefault);
|
|
|
|
// Throw if we fail to allocate
|
|
FAISS_THROW_IF_NOT_FMT(
|
|
err == cudaSuccess,
|
|
"failed to cudaHostAlloc %zu bytes (error %d %s)",
|
|
size, (int) err, cudaGetErrorString(err));
|
|
}
|
|
break;
|
|
default:
|
|
FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) space);
|
|
break;
|
|
}
|
|
}
|
|
|
|
// We'll allow allocation to fail, but free should always succeed and be a
|
|
// fatal error if it doesn't free
|
|
void freeMemorySpace(MemorySpace space, void* p) {
|
|
switch (space) {
|
|
case MemorySpace::Device:
|
|
case MemorySpace::Unified:
|
|
{
|
|
auto err = cudaFree(p);
|
|
FAISS_ASSERT_FMT(err == cudaSuccess,
|
|
"Failed to cudaFree pointer %p (error %d %s)",
|
|
p, (int) err, cudaGetErrorString(err));
|
|
}
|
|
break;
|
|
case MemorySpace::HostPinned:
|
|
{
|
|
auto err = cudaFreeHost(p);
|
|
FAISS_ASSERT_FMT(err == cudaSuccess,
|
|
"Failed to cudaFreeHost pointer %p (error %d %s)",
|
|
p, (int) err, cudaGetErrorString(err));
|
|
}
|
|
break;
|
|
default:
|
|
FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) space);
|
|
break;
|
|
}
|
|
}
|
|
|
|
} }
|