Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
ThrustAllocator.cuh
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.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 }
30
31 char* allocate(std::ptrdiff_t size) {
32 if (size <= (end_ - cur_)) {
33 char* p = cur_;
34 cur_ += size;
35 FAISS_ASSERT(cur_ <= end_);
36
37 return p;
38 } else {
39 char* p = nullptr;
40 CUDA_VERIFY(cudaMalloc(&p, size));
41 mallocAllocs_.insert(p);
42 return p;
43 }
44 }
45
46 void deallocate(char* p, size_t size) {
47 // Allocations could be returned out-of-order; ignore those we
48 // didn't cudaMalloc
49 auto it = mallocAllocs_.find(p);
50 if (it != mallocAllocs_.end()) {
51 CUDA_VERIFY(cudaFree(p));
52 mallocAllocs_.erase(it);
53 }
54 }
55
56 private:
57 char* start_;
58 char* cur_;
59 char* end_;
60 std::unordered_set<char*> mallocAllocs_;
61 };
62
63
64 } } // namespace
faiss::gpu::GpuResourcesThrustAllocator
Allocator for Thrust that comes out of a specified memory space.
Definition: ThrustAllocator.cuh:18