src/backends/memory_pool.cpp - with_cpu - vuron/adept

src/backends/memory_pool.cpp
126 строк4 KB
kolkir
Fix banch of errors and add TBB parallelization
22 янв 2025, 19:17
22 янв 2025, 19:17514572a
#include <adept/backends/memory_pool.h>
#include <adept/exception.hpp>

#include <iostream>

// #define ALLOC_TRACE

namespace adept {

PoolBuffer::PoolBuffer(buffer_ptr_t buffer, std::weak_ptr<MemoryPool> parent_pool)
    : original_buffer_(std::move(buffer)), parent_pool_(std::move(parent_pool)) {}

PoolBuffer::~PoolBuffer() {
  auto pool = parent_pool_.lock();
  if (pool) {
    pool->free(std::move(original_buffer_));
  } else {
    std::cerr << "PoolBuffer can't free memory because parent pool was destroyed";
  }
}

MemoryPool::MemoryPool(allocator_ptr_t allocator, size_t max_size, float32_t size_diff_ratio)
    : allocator_(allocator), max_size_(max_size), size_diff_ratio_(size_diff_ratio) {}

MemoryPool::~MemoryPool() {
  for (auto buf : buffers_cache_) {
    allocator_->free(std::move(buf));
  }
  buffers_cache_.clear();

  if (!buffers_in_use_.empty()) {
    std::cerr << "There are " << buffers_in_use_.size() << " leaked buffers in MemoryPool"
              << std::endl;
  }
  buffers_in_use_.clear();
}

pool_buffer_ptr_t MemoryPool::allocate(size_t size) {
  auto mem_size = allocator_->to_mem_size(size);
  {
    std::scoped_lock lock(guard_);
    auto lower_it = cache_lookup_.lower_bound(mem_size);

    // search exsisting buffer with suitable size
    if (lower_it != cache_lookup_.end() &&
        lower_it->first <= mem_size + static_cast<size_t>(mem_size * size_diff_ratio_)) {
#ifdef ALLOC_TRACE
      std::cout << this << " cache hit: " << mem_size << "(" << size
                << "), low: " << (*lower_it->second)->size() << std::endl;
#endif
      auto buffer = *lower_it->second;
      buffers_cache_.erase(lower_it->second);
      cache_lookup_.erase(lower_it);
      buffers_in_use_.push_back(buffer);
      return std::make_shared<PoolBuffer>(std::move(buffer), weak_from_this());
    }

#ifdef ALLOC_TRACE
    std::cout << this << " cache miss: " << mem_size << "(" << size << ")" << std::endl;
#endif

    // appropriate buffer was not found
    // 1. drop cache
    while (buffers_cache_.size() >= max_size_) {
      auto it_remove = std::prev(buffers_cache_.end());
      auto [lower_it, upper_it] = cache_lookup_.equal_range((*it_remove)->size());
      for (; lower_it != upper_it; ++lower_it) {
        if (lower_it->second == it_remove) {
          allocator_->free(std::move(*it_remove));
          buffers_cache_.erase(it_remove);
          cache_lookup_.erase(lower_it);
          break;
        }
      }
#ifdef ALLOC_TRACE
      std::cout << this << " cache erase , num in use: " << buffers_in_use()
                << ", num in cache: " << buffers_in_cache() << std::endl;
#endif
    }
  }

  // 2. allocate new one
  auto new_buffer = allocator_->allocate(size);
  std::scoped_lock lock(guard_);

#ifdef ALLOC_TRACE
  std::cout << this << " alloc: " << mem_size << "(" << size
            << "), num in use: " << buffers_in_use() << ", num in cache: " << buffers_in_cache()
            << std::endl;
#endif

  buffers_in_use_.push_back(new_buffer);
  return std::make_shared<PoolBuffer>(std::move(new_buffer), weak_from_this());
}

void MemoryPool::free(buffer_ptr_t&& buffer) {
  std::scoped_lock lock(guard_);
  buffer_ptr_t ret_buffer = std::move(buffer);
  auto it = buffers_in_use_.begin();
  for (; it != buffers_in_use_.end(); ++it) {
    if (*it == ret_buffer) {
      buffers_cache_.push_front(ret_buffer);
      cache_lookup_.emplace(ret_buffer->size(), buffers_cache_.begin());
      buffers_in_use_.erase(it);

#ifdef ALLOC_TRACE
      std::cout << this << " back to cache: " << ret_buffer->size()
                << ", num in use: " << buffers_in_use() << ", num in cache: " << buffers_in_cache()
                << std::endl;
#endif

      return;
    }
  }
  THROW_ERROR("MemoryPool::free failed: undeclared buffer");
}

size_t MemoryPool::buffers_in_use() const {
  return buffers_in_use_.size();
}

size_t MemoryPool::buffers_in_cache() const {
  return buffers_cache_.size();
}

}  // namespace adept