Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
StackDeviceMemory.h
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 #pragma once
10
11 #include "DeviceMemory.h"
12 #include <list>
13 #include <memory>
14 #include <unordered_map>
15
16 namespace faiss { namespace gpu {
17
18 /// Device memory manager that provides temporary memory allocations
19 /// out of a region of memory
20 class StackDeviceMemory : public DeviceMemory {
21 public:
22 /// Allocate a new region of memory that we manage
23 explicit StackDeviceMemory(int device, size_t allocPerDevice);
24
25 /// Manage a region of memory for a particular device, with or
26 /// without ownership
27 StackDeviceMemory(int device, void* p, size_t size, bool isOwner);
28
29 ~StackDeviceMemory() override;
30
31 /// Enable or disable the warning about not having enough temporary memory
32 /// when cudaMalloc gets called
33 void setCudaMallocWarning(bool b);
34
35 int getDevice() const override;
36
37 DeviceMemoryReservation getMemory(cudaStream_t stream,
38 size_t size) override;
39
40 size_t getSizeAvailable() const override;
41 std::string toString() const override;
42 size_t getHighWaterCudaMalloc() const override;
43
44 protected:
45 void returnAllocation(DeviceMemoryReservation& m) override;
46
47 protected:
48 /// Previous allocation ranges and the streams for which
49 /// synchronization is required
50 struct Range {
51 inline Range(char* s, char* e, cudaStream_t str) :
52 start_(s), end_(e), stream_(str) {
53 }
54
55 // References a memory range [start, end)
56 char* start_;
57 char* end_;
58 cudaStream_t stream_;
59 };
60
61 struct Stack {
62 /// Constructor that allocates memory via cudaMalloc
63 Stack(int device, size_t size);
64
65 /// Constructor that references a pre-allocated region of memory
66 Stack(int device, void* p, size_t size, bool isOwner);
67 ~Stack();
68
69 /// Returns how much size is available for an allocation without
70 /// calling cudaMalloc
71 size_t getSizeAvailable() const;
72
73 /// Obtains an allocation; all allocations are guaranteed to be 16
74 /// byte aligned
75 char* getAlloc(size_t size, cudaStream_t stream);
76
77 /// Returns an allocation
78 void returnAlloc(char* p, size_t size, cudaStream_t stream);
79
80 /// Returns the stack state
81 std::string toString() const;
82
83 /// Returns the high-water mark of cudaMalloc activity
84 size_t getHighWaterCudaMalloc() const;
85
86 /// Device this allocation is on
87 int device_;
88
89 /// Do we own our region of memory?
90 bool isOwner_;
91
92 /// Where our allocation begins and ends
93 /// [start_, end_) is valid
94 char* start_;
95 char* end_;
96
97 /// Total size end_ - start_
98 size_t size_;
99
100 /// Stack head within [start, end)
101 char* head_;
102
103 /// List of previous last users of allocations on our stack, for
104 /// possible synchronization purposes
105 std::list<Range> lastUsers_;
106
107 /// How much cudaMalloc memory is currently outstanding?
108 size_t mallocCurrent_;
109
110 /// What's the high water mark in terms of memory used from the
111 /// temporary buffer?
112 size_t highWaterMemoryUsed_;
113
114 /// What's the high water mark in terms of memory allocated via
115 /// cudaMalloc?
116 size_t highWaterMalloc_;
117
118 /// Whether or not a warning upon cudaMalloc is generated
119 bool cudaMallocWarning_;
120 };
121
122 /// Our device
123 int device_;
124
125 /// Memory stack
126 Stack stack_;
127 };
128
129 } } // namespace
faiss::gpu::StackDeviceMemory::Stack::isOwner_
bool isOwner_
Do we own our region of memory?
Definition: StackDeviceMemory.h:90
faiss::gpu::StackDeviceMemory::getMemory
DeviceMemoryReservation getMemory(cudaStream_t stream, size_t size) override
Definition: StackDeviceMemory.cpp:204
faiss::gpu::StackDeviceMemory::Stack::getHighWaterCudaMalloc
size_t getHighWaterCudaMalloc() const
Returns the high-water mark of cudaMalloc activity.
Definition: StackDeviceMemory.cpp:175
faiss::gpu::StackDeviceMemory::Stack::cudaMallocWarning_
bool cudaMallocWarning_
Whether or not a warning upon cudaMalloc is generated.
Definition: StackDeviceMemory.h:119
faiss::gpu::StackDeviceMemory::Stack::size_
size_t size_
Total size end_ - start_.
Definition: StackDeviceMemory.h:98
faiss::gpu::StackDeviceMemory::Stack::lastUsers_
std::list< Range > lastUsers_
Definition: StackDeviceMemory.h:105
faiss::gpu::StackDeviceMemory::Stack::Stack
Stack(int device, size_t size)
Constructor that allocates memory via cudaMalloc.
Definition: StackDeviceMemory.cpp:19
faiss::gpu::StackDeviceMemory::Stack::mallocCurrent_
size_t mallocCurrent_
How much cudaMalloc memory is currently outstanding?
Definition: StackDeviceMemory.h:108
faiss::gpu::StackDeviceMemory::setCudaMallocWarning
void setCudaMallocWarning(bool b)
Definition: StackDeviceMemory.cpp:194
faiss::gpu::StackDeviceMemory::stack_
Stack stack_
Memory stack.
Definition: StackDeviceMemory.h:126
faiss::gpu::StackDeviceMemory::Stack::returnAlloc
void returnAlloc(char *p, size_t size, cudaStream_t stream)
Returns an allocation.
Definition: StackDeviceMemory.cpp:132
faiss::gpu::StackDeviceMemory::Stack::head_
char * head_
Stack head within [start, end)
Definition: StackDeviceMemory.h:101
faiss::gpu::StackDeviceMemory::getSizeAvailable
size_t getSizeAvailable() const override
Returns the current size available without calling cudaMalloc.
Definition: StackDeviceMemory.cpp:217
faiss::gpu::StackDeviceMemory::Stack
Definition: StackDeviceMemory.h:61
faiss::gpu::StackDeviceMemory::Stack::getSizeAvailable
size_t getSizeAvailable() const
Definition: StackDeviceMemory.cpp:60
faiss::gpu::StackDeviceMemory::Stack::highWaterMemoryUsed_
size_t highWaterMemoryUsed_
Definition: StackDeviceMemory.h:112
faiss::gpu::StackDeviceMemory::Stack::device_
int device_
Device this allocation is on.
Definition: StackDeviceMemory.h:87
faiss::gpu::StackDeviceMemory
Definition: StackDeviceMemory.h:20
faiss::gpu::StackDeviceMemory::Stack::highWaterMalloc_
size_t highWaterMalloc_
Definition: StackDeviceMemory.h:116
faiss::gpu::StackDeviceMemory::toString
std::string toString() const override
Returns a string containing our current memory manager state.
Definition: StackDeviceMemory.cpp:222
faiss::gpu::StackDeviceMemory::Range
Definition: StackDeviceMemory.h:50
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:44
faiss::gpu::StackDeviceMemory::Stack::toString
std::string toString() const
Returns the stack state.
Definition: StackDeviceMemory.cpp:154
faiss::gpu::StackDeviceMemory::getHighWaterCudaMalloc
size_t getHighWaterCudaMalloc() const override
Definition: StackDeviceMemory.cpp:227
faiss::gpu::StackDeviceMemory::Stack::getAlloc
char * getAlloc(size_t size, cudaStream_t stream)
Definition: StackDeviceMemory.cpp:65
faiss::gpu::StackDeviceMemory::Stack::start_
char * start_
Definition: StackDeviceMemory.h:94
faiss::gpu::StackDeviceMemory::device_
int device_
Our device.
Definition: StackDeviceMemory.h:123
faiss::gpu::DeviceMemoryReservation
Definition: DeviceMemory.h:18
faiss::gpu::StackDeviceMemory::getDevice
int getDevice() const override
Returns the device we are managing memory for.
Definition: StackDeviceMemory.cpp:199
faiss::gpu::StackDeviceMemory::StackDeviceMemory
StackDeviceMemory(int device, size_t allocPerDevice)
Allocate a new region of memory that we manage.
Definition: StackDeviceMemory.cpp:179