Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuResources.h
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9
10 #pragma once
11
12 #include "utils/DeviceMemory.h"
13 #include <cuda_runtime.h>
14 #include <cublas_v2.h>
15 #include <utility>
16 #include <vector>
17
18 namespace faiss { namespace gpu {
19
20 /// Base class of GPU-side resource provider; hides provision of
21 /// cuBLAS handles, CUDA streams and a temporary memory manager
22 class GpuResources {
23 public:
24 virtual ~GpuResources();
25
26 /// Call to pre-allocate resources for a particular device. If this is
27 /// not called, then resources will be allocated at the first time
28 /// of demand
29 virtual void initializeForDevice(int device) = 0;
30
31 /// Returns the cuBLAS handle that we use for the given device
32 virtual cublasHandle_t getBlasHandle(int device) = 0;
33
34 /// Returns the stream that we order all computation on for the
35 /// given device
36 virtual cudaStream_t getDefaultStream(int device) = 0;
37
38 /// Returns the set of alternative streams that we use for the given device
39 virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0;
40
41 /// Returns the temporary memory manager for the given device
42 virtual DeviceMemory& getMemoryManager(int device) = 0;
43
44 /// Returns the available CPU pinned memory buffer
45 virtual std::pair<void*, size_t> getPinnedMemory() = 0;
46
47 /// Returns the stream on which we perform async CPU <-> GPU copies
48 virtual cudaStream_t getAsyncCopyStream(int device) = 0;
49
50 /// Calls getBlasHandle with the current device
51 cublasHandle_t getBlasHandleCurrentDevice();
52
53 /// Calls getDefaultStream with the current device
54 cudaStream_t getDefaultStreamCurrentDevice();
55
56 /// Synchronizes the CPU with respect to the default stream for the
57 /// given device
58 // equivalent to cudaDeviceSynchronize(getDefaultStream(device))
59 void syncDefaultStream(int device);
60
61 /// Calls syncDefaultStream for the current device
62 void syncDefaultStreamCurrentDevice();
63
64 /// Calls getAlternateStreams for the current device
65 std::vector<cudaStream_t> getAlternateStreamsCurrentDevice();
66
67 /// Calls getMemoryManager for the current device
68 DeviceMemory& getMemoryManagerCurrentDevice();
69
70 /// Calls getAsyncCopyStream for the current device
71 cudaStream_t getAsyncCopyStreamCurrentDevice();
72 };
73
74 } } // namespace
faiss::gpu::GpuResources::initializeForDevice
virtual void initializeForDevice(int device)=0
faiss::gpu::GpuResources::getDefaultStreamCurrentDevice
cudaStream_t getDefaultStreamCurrentDevice()
Calls getDefaultStream with the current device.
Definition: GpuResources.cpp:24
faiss::gpu::GpuResources::getBlasHandle
virtual cublasHandle_t getBlasHandle(int device)=0
Returns the cuBLAS handle that we use for the given device.
faiss::gpu::GpuResources
Definition: GpuResources.h:22
faiss::gpu::GpuResources::getDefaultStream
virtual cudaStream_t getDefaultStream(int device)=0
faiss::gpu::GpuResources::getBlasHandleCurrentDevice
cublasHandle_t getBlasHandleCurrentDevice()
Calls getBlasHandle with the current device.
Definition: GpuResources.cpp:19
faiss::gpu::GpuResources::getMemoryManagerCurrentDevice
DeviceMemory & getMemoryManagerCurrentDevice()
Calls getMemoryManager for the current device.
Definition: GpuResources.cpp:34
faiss::gpu::GpuResources::syncDefaultStreamCurrentDevice
void syncDefaultStreamCurrentDevice()
Calls syncDefaultStream for the current device.
Definition: GpuResources.cpp:49
faiss::gpu::GpuResources::getAsyncCopyStreamCurrentDevice
cudaStream_t getAsyncCopyStreamCurrentDevice()
Calls getAsyncCopyStream for the current device.
Definition: GpuResources.cpp:39
faiss::gpu::GpuResources::getAlternateStreamsCurrentDevice
std::vector< cudaStream_t > getAlternateStreamsCurrentDevice()
Calls getAlternateStreams for the current device.
Definition: GpuResources.cpp:29
faiss::gpu::GpuResources::getAsyncCopyStream
virtual cudaStream_t getAsyncCopyStream(int device)=0
Returns the stream on which we perform async CPU &lt;-&gt; GPU copies.
faiss::gpu::GpuResources::syncDefaultStream
void syncDefaultStream(int device)
Definition: GpuResources.cpp:44
faiss::gpu::GpuResources::getMemoryManager
virtual DeviceMemory & getMemoryManager(int device)=0
Returns the temporary memory manager for the given device.
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:45
faiss::gpu::GpuResources::getPinnedMemory
virtual std::pair< void *, size_t > getPinnedMemory()=0
Returns the available CPU pinned memory buffer.
faiss::gpu::GpuResources::getAlternateStreams
virtual std::vector< cudaStream_t > getAlternateStreams(int device)=0
Returns the set of alternative streams that we use for the given device.