/*! * Copyright (c) 2017 by Contributors * \file workspace_pool.h * \brief Workspace pool utility. */ #include "./workspace_pool.h" namespace tvm { namespace runtime { // page size. constexpr size_t kWorkspacePageSize = 4 << 10; class WorkspacePool::Pool { public: // constructor Pool() { // safe guard header on each list. Entry e; e.data = nullptr; e.size = 0; free_list_.push_back(e); allocated_.push_back(e); } // allocate from pool void* Alloc(TVMContext ctx, DeviceAPI* device, size_t nbytes) { // Allocate align to page. nbytes = (nbytes + (kWorkspacePageSize - 1)) / kWorkspacePageSize * kWorkspacePageSize; if (nbytes == 0) nbytes = kWorkspacePageSize; Entry e; TVMType type; type.code = kDLUInt; type.bits = 8; type.lanes = 1; if (free_list_.size() == 2) { e = free_list_.back(); free_list_.pop_back(); if (e.size < nbytes) { // resize the page device->FreeDataSpace(ctx, e.data); e.data = device->AllocDataSpace(ctx, nbytes, kTempAllocaAlignment, type); e.size = nbytes; } } else if (free_list_.size() == 1) { e.data = device->AllocDataSpace(ctx, nbytes, kTempAllocaAlignment, type); e.size = nbytes; } else { if (free_list_.back().size >= nbytes) { // find smallest fit auto it = free_list_.end() - 2; for (; it->size >= nbytes; --it) {} e = *(it + 1); free_list_.erase(it + 1); } else { // resize the page e = free_list_.back(); free_list_.pop_back(); device->FreeDataSpace(ctx, e.data); e.data = device->AllocDataSpace(ctx, nbytes, kTempAllocaAlignment, type); e.size = nbytes; } } allocated_.push_back(e); return e.data; } // free resource back to pool void Free(void* data) { Entry e; if (allocated_.back().data == data) { // quick path, last allocated. e = allocated_.back(); allocated_.pop_back(); } else { int index = static_cast<int>(allocated_.size()) - 2; for (; index > 0 && allocated_[index].data != data; --index) {} CHECK_GT(index, 0) << "trying to free things that has not been allocated"; e = allocated_[index]; allocated_.erase(allocated_.begin() + index); } if (free_list_.back().size < e.size) { free_list_.push_back(e); } else if (free_list_.size() == 2) { free_list_.push_back(free_list_.back()); free_list_[1] = e; } else { size_t i = free_list_.size() - 1; free_list_.resize(free_list_.size() + 1); for (; e.size < free_list_[i].size; --i) { free_list_[i + 1] = free_list_[i]; } free_list_[i + 1] = e; } } // Release all resources void Release(TVMContext ctx, DeviceAPI* device) { CHECK_EQ(allocated_.size(), 1); for (size_t i = 1; i < free_list_.size(); ++i) { device->FreeDataSpace(ctx, free_list_[i].data); } free_list_.clear(); } private: /*! \brief a single entry in the pool */ struct Entry { void* data; size_t size; }; /*! \brief List of free items, sorted from small to big size */ std::vector<Entry> free_list_; /*! \brief List of allocated items */ std::vector<Entry> allocated_; }; WorkspacePool::WorkspacePool(DLDeviceType device_type, std::shared_ptr<DeviceAPI> device) : device_type_(device_type), device_(device) { } WorkspacePool::~WorkspacePool() { for (size_t i = 0; i < array_.size(); ++i) { if (array_[i] != nullptr) { TVMContext ctx; ctx.device_type = device_type_; ctx.device_id = static_cast<int>(i); array_[i]->Release(ctx, device_.get()); delete array_[i]; } } } void* WorkspacePool::AllocWorkspace(TVMContext ctx, size_t size) { if (static_cast<size_t>(ctx.device_id) >= array_.size()) { array_.resize(ctx.device_id + 1, nullptr); } if (array_[ctx.device_id] == nullptr) { array_[ctx.device_id] = new Pool(); } return array_[ctx.device_id]->Alloc(ctx, device_.get(), size); } void WorkspacePool::FreeWorkspace(TVMContext ctx, void* ptr) { CHECK(static_cast<size_t>(ctx.device_id) < array_.size() && array_[ctx.device_id] != nullptr); array_[ctx.device_id]->Free(ptr); } } // namespace runtime } // namespace tvm