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