Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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