Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
DeviceMemory.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 <cuda_runtime.h>
12 #include <string>
13
14 namespace faiss { namespace gpu {
15
16 class DeviceMemory;
17
18 class DeviceMemoryReservation {
19 public:
20 DeviceMemoryReservation();
21 DeviceMemoryReservation(DeviceMemory* state,
22 int device, void* p, size_t size,
23 cudaStream_t stream);
24 DeviceMemoryReservation(DeviceMemoryReservation&& m) noexcept;
25 ~DeviceMemoryReservation();
26
27 DeviceMemoryReservation& operator=(DeviceMemoryReservation&& m);
28
29 int device() { return device_; }
30 void* get() { return data_; }
31 size_t size() { return size_; }
32 cudaStream_t stream() { return stream_; }
33
34 private:
35 DeviceMemory* state_;
36
37 int device_;
38 void* data_;
39 size_t size_;
40 cudaStream_t stream_;
41 };
42
43 /// Manages temporary memory allocations on a GPU device
44 class DeviceMemory {
45 public:
46 virtual ~DeviceMemory();
47
48 /// Returns the device we are managing memory for
49 virtual int getDevice() const = 0;
50
51 /// Obtains a temporary memory allocation for our device,
52 /// whose usage is ordered with respect to the given stream.
53 virtual DeviceMemoryReservation getMemory(cudaStream_t stream,
54 size_t size) = 0;
55
56 /// Returns the current size available without calling cudaMalloc
57 virtual size_t getSizeAvailable() const = 0;
58
59 /// Returns a string containing our current memory manager state
60 virtual std::string toString() const = 0;
61
62 /// Returns the high-water mark of cudaMalloc allocations for our
63 /// device
64 virtual size_t getHighWaterCudaMalloc() const = 0;
65
66 protected:
67 friend class DeviceMemoryReservation;
68 virtual void returnAllocation(DeviceMemoryReservation& m) = 0;
69 };
70
71 } } // namespace
faiss::gpu::DeviceMemory::getHighWaterCudaMalloc
virtual size_t getHighWaterCudaMalloc() const =0
faiss::gpu::DeviceMemory::toString
virtual std::string toString() const =0
Returns a string containing our current memory manager state.
faiss::gpu::DeviceMemory::getSizeAvailable
virtual size_t getSizeAvailable() const =0
Returns the current size available without calling cudaMalloc.
faiss::gpu::DeviceMemory::getMemory
virtual DeviceMemoryReservation getMemory(cudaStream_t stream, size_t size)=0
faiss::gpu::DeviceMemory::getDevice
virtual int getDevice() const =0
Returns the device we are managing memory for.
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:44
faiss::gpu::DeviceMemoryReservation
Definition: DeviceMemory.h:18