1
#include <c10/core/impl/alloc_cpu.h>
2
#include <c10/mobile/CPUCachingAllocator.h>
3
#include <c10/util/Exception.h>
8
thread_local CPUCachingAllocator* caching_allocator_ptr{nullptr};
11
std::mutex CPUCachingAllocator::mutex_;
12
ska::flat_hash_map<void*, size_t> CPUCachingAllocator::allocation_map_;
14
inline void* CPUCachingAllocator::allocate_and_cache(const size_t bytes) {
18
ptr = c10::alloc_cpu(bytes);
19
} catch (c10::Error&) {
27
ptr = c10::alloc_cpu(bytes);
29
allocation_map_[ptr] = bytes;
33
void* CPUCachingAllocator::allocate(const size_t bytes) {
34
std::lock_guard<std::mutex> guard(mutex_);
35
const auto& it = available_map_.find(bytes);
36
if (it == available_map_.end() || it->second.empty()) {
37
return allocate_and_cache(bytes);
39
return it->second.pop_back_val();
42
void CPUCachingAllocator::free(void* ptr) {
50
std::lock_guard<std::mutex> guard(mutex_);
53
const auto& it = allocation_map_.find(ptr);
54
if (it == allocation_map_.end()) {
58
const size_t alloc_size = it->second;
59
available_map_[alloc_size].push_back(ptr);
62
void CPUCachingAllocator::record_free(void* ptr) {
72
std::lock_guard<std::mutex> guard(mutex_);
73
const auto& it = allocation_map_.find(ptr);
74
if (it != allocation_map_.end()) {
75
allocation_map_.erase(it);
79
void CPUCachingAllocator::free_cached() {
80
for (const auto& it : available_map_) {
81
for (const auto ptr : it.second) {
85
allocation_map_.erase(ptr);
88
available_map_.clear();
91
CPUCachingAllocator::~CPUCachingAllocator() {
95
CPUCachingAllocator* GetThreadLocalCachingAllocator() {
96
return caching_allocator_ptr;
99
WithCPUCachingAllocatorGuard::WithCPUCachingAllocatorGuard(
100
CPUCachingAllocator* allocator)
101
: prev_caching_allocator_ptr_(GetThreadLocalCachingAllocator()) {
102
caching_allocator_ptr = allocator;
105
WithCPUCachingAllocatorGuard::~WithCPUCachingAllocatorGuard() {
106
caching_allocator_ptr = prev_caching_allocator_ptr_;