90 lines
2.4 KiB
C++
90 lines
2.4 KiB
C++
/**
|
|
* 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/utils/MemorySpace.h>
|
|
#include <faiss/impl/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;
|
|
}
|
|
}
|
|
|
|
} }
|