Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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