vuron/adept: Платформа для обучения и вывода нейронных сетей, с API приближенным к PyTorch.

tests/mempool_tests.cpp
125 строк4 KB
kolkir
Cpp serialization implementation and tests
25 янв 2025, 14:06
25 янв 2025, 14:06f3f352c
#include <adept/backends/memory_pool.h>

#include "catch.hpp"

using namespace adept;

namespace {
struct MockBuffer : public Buffer {
  MockBuffer(void* ptr, size_t size) : ptr_(ptr), size_(size) {}
  void* ptr() const override { return ptr_; }
  size_t size() const override { return size_; }

  void* ptr_{nullptr};
  size_t size_{0};
};

struct MockAllocator : public Allocator {
  buffer_ptr_t allocate(size_t size) override {
    ++count;
    ++alloc_count;
    return std::make_shared<MockBuffer>(::malloc(size), size);
  }
  void free(buffer_ptr_t&& buffer) override {
    ::free(std::move(buffer)->ptr());
    --count;
    ++free_count;
  }
  size_t to_mem_size(size_t numel) override { return numel; }
  size_t count{0};
  size_t alloc_count{0};
  size_t free_count{0};
};
}  // namespace

TEST_CASE("MemoryPool allocate buffer", "[memory pool]") {
  auto allocator = std::make_shared<MockAllocator>();
  {
    auto mem_pool = std::make_shared<MemoryPool>(allocator, /*max_size=*/5);
    auto buf = mem_pool->allocate(32);
    REQUIRE(buf->size() == 32);
    REQUIRE(buf->ptr() != nullptr);
    REQUIRE(mem_pool->buffers_in_use() == 1);
    REQUIRE(mem_pool->buffers_in_cache() == 0);
    buf.reset();
    REQUIRE(mem_pool->buffers_in_cache() == 1);
  }
  REQUIRE(allocator->count == 0);
  REQUIRE(allocator->alloc_count == 1);
  REQUIRE(allocator->free_count == 1);
}

TEST_CASE("MemoryPool reuse buffers", "[memory pool]") {
  auto allocator = std::make_shared<MockAllocator>();
  {
    auto mem_pool = std::make_shared<MemoryPool>(allocator, /*max_size=*/5);
    std::vector<pool_buffer_ptr_t> buffers;
    for (int i = 0; i < 5; ++i) {
      auto buf = mem_pool->allocate(32);
      REQUIRE(buf->size() == 32);
      REQUIRE(buf->ptr() != nullptr);
      buffers.emplace_back(std::move(buf));
    }
    REQUIRE(mem_pool->buffers_in_use() == 5);
    REQUIRE(mem_pool->buffers_in_cache() == 0);
    buffers[0].reset();
    buffers[1].reset();
    buffers[2].reset();
    REQUIRE(mem_pool->buffers_in_cache() == 3);
    REQUIRE(mem_pool->buffers_in_use() == 2);
    buffers[3].reset();
    buffers[4].reset();
    REQUIRE(mem_pool->buffers_in_cache() == 5);
    REQUIRE(mem_pool->buffers_in_use() == 0);
    REQUIRE(allocator->count == 5);
    // reuse cache
    buffers.clear();
    for (int i = 0; i < 5; ++i) {
      auto buf = mem_pool->allocate(26);
      REQUIRE(buf->size() == 32);
      REQUIRE(buf->ptr() != nullptr);
      buffers.emplace_back(std::move(buf));
    }
    REQUIRE(allocator->alloc_count == 5);  // no new allocations
    REQUIRE(mem_pool->buffers_in_cache() == 0);
    REQUIRE(mem_pool->buffers_in_use() == 5);
    buffers.clear();
    REQUIRE(mem_pool->buffers_in_cache() == 5);
    REQUIRE(mem_pool->buffers_in_use() == 0);
  }
  REQUIRE(allocator->count == 0);
  REQUIRE(allocator->alloc_count == 5);
  REQUIRE(allocator->free_count == 5);
}

TEST_CASE("MemoryPool drop buffers", "[memory pool]") {
  auto allocator = std::make_shared<MockAllocator>();
  {
    auto mem_pool = std::make_shared<MemoryPool>(allocator, /*max_size=*/5);
    std::vector<pool_buffer_ptr_t> buffers;
    for (int i = 0; i < 5; ++i) {
      auto buf = mem_pool->allocate(32);
      REQUIRE(buf->size() == 32);
      REQUIRE(buf->ptr() != nullptr);
      buffers.emplace_back(std::move(buf));
    }
    REQUIRE(allocator->alloc_count == 5);
    REQUIRE(mem_pool->buffers_in_use() == 5);
    REQUIRE(mem_pool->buffers_in_cache() == 0);
    buffers.clear();  // return to cache
    REQUIRE(mem_pool->buffers_in_cache() == 5);
    REQUIRE(mem_pool->buffers_in_use() == 0);
    auto buf = mem_pool->allocate(64);
    REQUIRE(buf->size() == 64);
    REQUIRE(buf->ptr() != nullptr);
    REQUIRE(allocator->alloc_count == 6);
    REQUIRE(mem_pool->buffers_in_use() == 1);
    REQUIRE(mem_pool->buffers_in_cache() == 4);
    buf.reset();
    REQUIRE(mem_pool->buffers_in_use() == 0);
    REQUIRE(mem_pool->buffers_in_cache() == 5);
  }
  REQUIRE(allocator->count == 0);
  REQUIRE(allocator->alloc_count == 6);
  REQUIRE(allocator->free_count == 6);
}