Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
ThrustAllocator.cuh
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 "MemorySpace.h"
12 #include <cuda.h>
13 #include <unordered_set>
14
15 namespace faiss { namespace gpu {
16
17 /// Allocator for Thrust that comes out of a specified memory space
18 class GpuResourcesThrustAllocator {
19 public:
20 typedef char value_type;
21
22 GpuResourcesThrustAllocator(void* mem, size_t size)
23 : start_((char*) mem),
24 cur_((char*) mem),
25 end_((char*) mem + size) {
26 }
27
28 ~GpuResourcesThrustAllocator() {
29 // In the case of an exception being thrown, we may not have called
30 // deallocate on all of our sub-allocations. Free them here
31 for (auto p : mallocAllocs_) {
32 freeMemorySpace(MemorySpace::Device, p);
33 }
34 }
35
36 char* allocate(std::ptrdiff_t size) {
37 if (size <= (end_ - cur_)) {
38 char* p = cur_;
39 cur_ += size;
40 FAISS_ASSERT(cur_ <= end_);
41
42 return p;
43 } else {
44 char* p = nullptr;
45 allocMemorySpace(MemorySpace::Device, &p, size);
46 mallocAllocs_.insert(p);
47 return p;
48 }
49 }
50
51 void deallocate(char* p, size_t size) {
52 // Allocations could be returned out-of-order; ignore those we
53 // didn't cudaMalloc
54 auto it = mallocAllocs_.find(p);
55 if (it != mallocAllocs_.end()) {
56 freeMemorySpace(MemorySpace::Device, p);
57 mallocAllocs_.erase(it);
58 }
59 }
60
61 private:
62 char* start_;
63 char* cur_;
64 char* end_;
65 std::unordered_set<char*> mallocAllocs_;
66 };
67
68
69 } } // namespace
faiss::gpu::GpuResourcesThrustAllocator
Allocator for Thrust that comes out of a specified memory space.
Definition: ThrustAllocator.cuh:18