Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
DeviceUtils.cpp
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved.
11
12 #include "DeviceUtils.h"
13 #include "../../FaissAssert.h"
14 #include <mutex>
15 #include <unordered_map>
16
17 namespace faiss { namespace gpu {
18
19 int getCurrentDevice() {
20 int dev = -1;
21 CUDA_VERIFY(cudaGetDevice(&dev));
22 FAISS_ASSERT(dev != -1);
23
24 return dev;
25 }
26
27 void setCurrentDevice(int device) {
28 CUDA_VERIFY(cudaSetDevice(device));
29 }
30
31 int getNumDevices() {
32 int numDev = -1;
33 CUDA_VERIFY(cudaGetDeviceCount(&numDev));
34 FAISS_ASSERT(numDev != -1);
35
36 return numDev;
37 }
38
39 void synchronizeAllDevices() {
40 for (int i = 0; i < getNumDevices(); ++i) {
41 DeviceScope scope(i);
42
43 CUDA_VERIFY(cudaDeviceSynchronize());
44 }
45 }
46
47 cudaDeviceProp& getDeviceProperties(int device) {
48 static std::mutex mutex;
49 static std::unordered_map<int, cudaDeviceProp> properties;
50
51 std::lock_guard<std::mutex> guard(mutex);
52
53 auto it = properties.find(device);
54 if (it == properties.end()) {
55 cudaDeviceProp prop;
56 CUDA_VERIFY(cudaGetDeviceProperties(&prop, device));
57
58 properties[device] = prop;
59 it = properties.find(device);
60 }
61
62 return it->second;
63 }
64
65 int getMaxThreads(int device) {
66 return getDeviceProperties(device).maxThreadsPerBlock;
67 }
68
69 int getMaxThreadsCurrentDevice() {
70 return getMaxThreads(getCurrentDevice());
71 }
72
73 size_t getMaxSharedMemPerBlock(int device) {
74 return getDeviceProperties(device).sharedMemPerBlock;
75 }
76
77 size_t getMaxSharedMemPerBlockCurrentDevice() {
78 return getMaxSharedMemPerBlock(getCurrentDevice());
79 }
80
81 int getDeviceForAddress(const void* p) {
82 if (!p) {
83 return -1;
84 }
85
86 cudaPointerAttributes att;
87 cudaError_t err = cudaPointerGetAttributes(&att, p);
88 FAISS_ASSERT(err == cudaSuccess ||
89 err == cudaErrorInvalidValue);
90
91 if (err == cudaErrorInvalidValue) {
92 // Make sure the current thread error status has been reset
93 err = cudaGetLastError();
94 FAISS_ASSERT(err == cudaErrorInvalidValue);
95 return -1;
96 } else if (att.memoryType == cudaMemoryTypeHost) {
97 return -1;
98 } else {
99 return att.device;
100 }
101 }
102
103 DeviceScope::DeviceScope(int device) {
104 prevDevice_ = getCurrentDevice();
105
106 if (prevDevice_ != device) {
107 setCurrentDevice(device);
108 } else {
109 prevDevice_ = -1;
110 }
111 }
112
113 DeviceScope::~DeviceScope() {
114 if (prevDevice_ != -1) {
115 setCurrentDevice(prevDevice_);
116 }
117 }
118
119 CublasHandleScope::CublasHandleScope() {
120 auto blasStatus = cublasCreate(&blasHandle_);
121 FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
122 }
123
124 CublasHandleScope::~CublasHandleScope() {
125 auto blasStatus = cublasDestroy(blasHandle_);
126 FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
127 }
128
129 CudaEvent::CudaEvent(cudaStream_t stream)
130 : event_(0) {
131 CUDA_VERIFY(cudaEventCreateWithFlags(&event_, cudaEventDisableTiming));
132 CUDA_VERIFY(cudaEventRecord(event_, stream));
133 }
134
135 CudaEvent::CudaEvent(CudaEvent&& event) noexcept
136 : event_(std::move(event.event_)) {
137 event.event_ = 0;
138 }
139
140 CudaEvent::~CudaEvent() {
141 if (event_) {
142 CUDA_VERIFY(cudaEventDestroy(event_));
143 }
144 }
145
146 CudaEvent&
147 CudaEvent::operator=(CudaEvent&& event) noexcept {
148 event_ = std::move(event.event_);
149 event.event_ = 0;
150
151 return *this;
152 }
153
154 void
155 CudaEvent::streamWaitOnEvent(cudaStream_t stream) {
156 CUDA_VERIFY(cudaStreamWaitEvent(stream, event_, 0));
157 }
158
159 void
160 CudaEvent::cpuWaitOnEvent() {
161 CUDA_VERIFY(cudaEventSynchronize(event_));
162 }
163
164 } } // namespace
faiss::gpu::CudaEvent::cpuWaitOnEvent
void cpuWaitOnEvent()
Have the CPU wait for the completion of this event.
Definition: DeviceUtils.cpp:160
faiss::gpu::CudaEvent::streamWaitOnEvent
void streamWaitOnEvent(cudaStream_t stream)
Wait on this event in this stream.
Definition: DeviceUtils.cpp:155
faiss::gpu::CudaEvent::CudaEvent
CudaEvent(cudaStream_t stream)
Creates an event and records it in this stream.
Definition: DeviceUtils.cpp:129
faiss::gpu::CudaEvent
Definition: DeviceUtils.h:77