graph/pool.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file pool.cpp | ||
| 3 | /// @brief Lending pooled images and arena ranges to graphs. | ||
| 4 | #include "ernie/graph/pool.h" | ||
| 5 | |||
| 6 | #include <algorithm> | ||
| 7 | #include <string> | ||
| 8 | |||
| 9 | namespace ernie::graph { | ||
| 10 | |||
| 11 | 26 | TransientPool::TransientPool(DeviceContext& context, std::uint64_t arenaSize) | |
| 12 | 26 | : context_(context), buffers_(context, MemoryUse::Device, "graph transients", arenaSize) {} | |
| 13 | |||
| 14 | 26 | TransientPool::~TransientPool() { | |
| 15 | 27 | for (Entry& entry : images_) { | |
| 16 | 1 | context_.Retire(std::move(entry.image)); | |
| 17 | } | ||
| 18 | 26 | } | |
| 19 | |||
| 20 | 4 | TransientPool::ImageLease TransientPool::AcquireImage(const ImageDescription& description) { | |
| 21 | 4 | for (Entry& entry : images_) { | |
| 22 | 2 | if (!entry.leased && entry.image->Description() == description) { | |
| 23 | 2 | entry.leased = true; | |
| 24 | 2 | entry.idle = 0; | |
| 25 | 2 | return {entry.image.get(), true}; | |
| 26 | } | ||
| 27 | } | ||
| 28 | 2 | Entry& entry = images_.emplace_back( | |
| 29 | 4 | Entry{std::make_shared<DeviceImage>(context_, description, | |
| 30 | 4 | "graph transient " + std::to_string(made_++)), | |
| 31 | true, 0}); | ||
| 32 | 2 | return {entry.image.get(), false}; | |
| 33 | } | ||
| 34 | |||
| 35 | 4 | void TransientPool::Release(const DeviceImage& image) { | |
| 36 | 4 | const auto found = std::ranges::find_if( | |
| 37 | 8 | images_, [&](const Entry& entry) { return entry.image.get() == ℑ }); | |
| 38 | 4 | if (found != images_.end()) { | |
| 39 | 4 | found->leased = false; | |
| 40 | } | ||
| 41 | 4 | } | |
| 42 | |||
| 43 | 86 | BufferRange TransientPool::AcquireBuffer(std::uint64_t bytes) { | |
| 44 | 86 | return buffers_.Allocate(bytes); | |
| 45 | } | ||
| 46 | |||
| 47 | 86 | void TransientPool::Release(const BufferRange& range) { | |
| 48 | 86 | buffers_.Free(range); | |
| 49 | 86 | } | |
| 50 | |||
| 51 | 1 | void TransientPool::Trim(std::uint32_t frames) { | |
| 52 | 1 | std::erase_if(images_, [&](Entry& entry) { | |
| 53 | 1 | if (entry.leased || entry.idle++ < frames) { | |
| 54 | ✗ | return false; | |
| 55 | } | ||
| 56 | 1 | context_.Retire(std::move(entry.image)); | |
| 57 | 1 | return true; | |
| 58 | }); | ||
| 59 | 1 | } | |
| 60 | |||
| 61 | } // namespace ernie::graph | ||
| 62 |