Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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