Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuResources.h
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 #pragma once
13
14 #include "utils/DeviceMemory.h"
15 #include <cuda_runtime.h>
16 #include <cublas_v2.h>
17 #include <utility>
18 #include <vector>
19
20 namespace faiss { namespace gpu {
21
22 /// Base class of GPU-side resource provider; hides provision of
23 /// cuBLAS handles, CUDA streams and a temporary memory manager
24 class GpuResources {
25 public:
26 virtual ~GpuResources();
27
28 /// Call to pre-allocate resources for a particular device. If this is
29 /// not called, then resources will be allocated at the first time
30 /// of demand
31 virtual void initializeForDevice(int device) = 0;
32
33 virtual cublasHandle_t getBlasHandle(int device) = 0;
34
35 virtual cudaStream_t getDefaultStream(int device) = 0;
36
37 virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0;
38
39 virtual DeviceMemory& getMemoryManager(int device) = 0;
40
41 virtual std::pair<void*, size_t> getPinnedMemory() = 0;
42
43 virtual cudaStream_t getAsyncCopyStream(int device) = 0;
44
45 cublasHandle_t getBlasHandleCurrentDevice();
46
47 cudaStream_t getDefaultStreamCurrentDevice();
48
49 std::vector<cudaStream_t> getAlternateStreamsCurrentDevice();
50
51 DeviceMemory& getMemoryManagerCurrentDevice();
52
53 cudaStream_t getAsyncCopyStreamCurrentDevice();
54 };
55
56 } } // namespace
faiss::gpu::GpuResources::initializeForDevice
virtual void initializeForDevice(int device)=0
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:47