graph/pool.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file pool.h | ||
| 3 | /// @brief Memory for the frame graph's transient resources: images pooled across frames by | ||
| 4 | /// description, and buffers as ranges of arenas. | ||
| 5 | #pragma once | ||
| 6 | |||
| 7 | #include "ernie/device/arena.h" | ||
| 8 | #include "ernie/device/image.h" | ||
| 9 | |||
| 10 | #include <cstdint> | ||
| 11 | #include <memory> | ||
| 12 | #include <vector> | ||
| 13 | |||
| 14 | namespace ernie::graph { | ||
| 15 | |||
| 16 | /// @brief Lends graphs images and buffer ranges for one execution. An image a graph returns may | ||
| 17 | /// be lent to a later graph on the same queue, whose barriers order it after the earlier | ||
| 18 | /// graph's work; images are freed only once the device is done with them. | ||
| 19 | class TransientPool { | ||
| 20 | public: | ||
| 21 | /// @param arenaSize Bytes per arena of transient buffer ranges. | ||
| 22 | explicit TransientPool(DeviceContext& context, | ||
| 23 | std::uint64_t arenaSize = std::uint64_t{64} << 20); | ||
| 24 | ~TransientPool(); | ||
| 25 | TransientPool(const TransientPool&) = delete; | ||
| 26 | TransientPool& operator=(const TransientPool&) = delete; | ||
| 27 | |||
| 28 | struct ImageLease { | ||
| 29 | DeviceImage* image = nullptr; | ||
| 30 | /// Whether an earlier graph used it. | ||
| 31 | bool recycled = false; | ||
| 32 | }; | ||
| 33 | /// @brief Lends an idle image of `description`, or makes one. | ||
| 34 | ImageLease AcquireImage(const ImageDescription& description); | ||
| 35 | void Release(const DeviceImage& image); | ||
| 36 | |||
| 37 | /// @brief Lends a range of `bytes`, given back once every queue's work so far is done. | ||
| 38 | BufferRange AcquireBuffer(std::uint64_t bytes); | ||
| 39 | void Release(const BufferRange& range); | ||
| 40 | 86 | const Arenas& Buffers() const { | |
| 41 | 86 | return buffers_; | |
| 42 | } | ||
| 43 | |||
| 44 | /// @brief Frees idle images no graph leased in the last `frames` trims, once the device is | ||
| 45 | /// done with them. | ||
| 46 | void Trim(std::uint32_t frames = 0); | ||
| 47 | /// @return The images held, lent or idle. | ||
| 48 | 2 | std::size_t Images() const { | |
| 49 | 2 | return images_.size(); | |
| 50 | } | ||
| 51 | |||
| 52 | private: | ||
| 53 | 6 | struct Entry { | |
| 54 | std::shared_ptr<DeviceImage> image; | ||
| 55 | bool leased = false; | ||
| 56 | /// Trims since it was last leased. | ||
| 57 | std::uint32_t idle = 0; | ||
| 58 | }; | ||
| 59 | |||
| 60 | DeviceContext& context_; | ||
| 61 | Arenas buffers_; | ||
| 62 | std::vector<Entry> images_; | ||
| 63 | std::uint64_t made_ = 0; | ||
| 64 | }; | ||
| 65 | |||
| 66 | } // namespace ernie::graph | ||
| 67 |