67 lines
1.4 KiB
Plaintext
67 lines
1.4 KiB
Plaintext
|
|
/**
|
|
* Copyright (c) 2015-present, Facebook, Inc.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the CC-by-NC license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
// Copyright 2004-present Facebook. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#include <cuda.h>
|
|
#include <unordered_set>
|
|
|
|
namespace faiss { namespace gpu {
|
|
|
|
/// Allocator for Thrust that comes out of a specified memory space
|
|
class GpuResourcesThrustAllocator {
|
|
public:
|
|
typedef char value_type;
|
|
|
|
GpuResourcesThrustAllocator(void* mem, size_t size)
|
|
: start_((char*) mem),
|
|
cur_((char*) mem),
|
|
end_((char*) mem + size) {
|
|
}
|
|
|
|
~GpuResourcesThrustAllocator() {
|
|
}
|
|
|
|
char* allocate(std::ptrdiff_t size) {
|
|
if (size <= (end_ - cur_)) {
|
|
char* p = cur_;
|
|
cur_ += size;
|
|
FAISS_ASSERT(cur_ <= end_);
|
|
|
|
return p;
|
|
} else {
|
|
char* p = nullptr;
|
|
CUDA_VERIFY(cudaMalloc(&p, size));
|
|
mallocAllocs_.insert(p);
|
|
return p;
|
|
}
|
|
}
|
|
|
|
void deallocate(char* p, size_t size) {
|
|
// Allocations could be returned out-of-order; ignore those we
|
|
// didn't cudaMalloc
|
|
auto it = mallocAllocs_.find(p);
|
|
if (it != mallocAllocs_.end()) {
|
|
CUDA_VERIFY(cudaFree(p));
|
|
mallocAllocs_.erase(it);
|
|
}
|
|
}
|
|
|
|
private:
|
|
char* start_;
|
|
char* cur_;
|
|
char* end_;
|
|
std::unordered_set<char*> mallocAllocs_;
|
|
};
|
|
|
|
|
|
} } // namespace
|