Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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