Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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