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