Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
MemorySpace.cpp
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8
9 #include "MemorySpace.h"
10 #include "../../FaissAssert.h"
11 #include <cuda_runtime.h>
12
13 namespace faiss { namespace gpu {
14
15 /// Allocates CUDA memory for a given memory space
16 void allocMemorySpaceV(MemorySpace space, void** p, size_t size) {
17 switch (space) {
18 case MemorySpace::Device:
19 {
20 auto err = cudaMalloc(p, size);
21
22 // Throw if we fail to allocate
23 FAISS_THROW_IF_NOT_FMT(
24 err == cudaSuccess,
25 "failed to cudaMalloc %zu bytes (error %d %s)",
26 size, (int) err, cudaGetErrorString(err));
27 }
28 break;
29 case MemorySpace::Unified:
30 {
31 #ifdef FAISS_UNIFIED_MEM
32 auto err = cudaMallocManaged(p, size);
33
34 // Throw if we fail to allocate
35 FAISS_THROW_IF_NOT_FMT(
36 err == cudaSuccess,
37 "failed to cudaMallocManaged %zu bytes (error %d %s)",
38 size, (int) err, cudaGetErrorString(err));
39 #else
40 FAISS_THROW_MSG("Attempting to allocate via cudaMallocManaged "
41 "without CUDA 8+ support");
42 #endif
43 }
44 break;
45 case MemorySpace::HostPinned:
46 {
47 auto err = cudaHostAlloc(p, size, cudaHostAllocDefault);
48
49 // Throw if we fail to allocate
50 FAISS_THROW_IF_NOT_FMT(
51 err == cudaSuccess,
52 "failed to cudaHostAlloc %zu bytes (error %d %s)",
53 size, (int) err, cudaGetErrorString(err));
54 }
55 break;
56 default:
57 FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) space);
58 break;
59 }
60 }
61
62 // We'll allow allocation to fail, but free should always succeed and be a
63 // fatal error if it doesn't free
64 void freeMemorySpace(MemorySpace space, void* p) {
65 switch (space) {
66 case MemorySpace::Device:
67 case MemorySpace::Unified:
68 {
69 auto err = cudaFree(p);
70 FAISS_ASSERT_FMT(err == cudaSuccess,
71 "Failed to cudaFree pointer %p (error %d %s)",
72 p, (int) err, cudaGetErrorString(err));
73 }
74 break;
75 case MemorySpace::HostPinned:
76 {
77 auto err = cudaFreeHost(p);
78 FAISS_ASSERT_FMT(err == cudaSuccess,
79 "Failed to cudaFreeHost pointer %p (error %d %s)",
80 p, (int) err, cudaGetErrorString(err));
81 }
82 break;
83 default:
84 FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) space);
85 break;
86 }
87 }
88
89 } }