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