device/arena.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file arena.h | ||
| 3 | /// @brief Device memory as a few large buffers sub-allocated by offset. Shaders reach data as | ||
| 4 | /// (arena, offset): a buffer device address on Vulkan, a bound storage buffer and offset on | ||
| 5 | /// WebGPU, which has no pointers. | ||
| 6 | #pragma once | ||
| 7 | |||
| 8 | #include "ernie/device/buffer.h" | ||
| 9 | |||
| 10 | #include <cstdint> | ||
| 11 | #include <memory> | ||
| 12 | #include <mutex> | ||
| 13 | #include <unordered_map> | ||
| 14 | #include <vector> | ||
| 15 | |||
| 16 | namespace ernie { | ||
| 17 | |||
| 18 | /// @brief Where data lives on the device: a range of one arena. | ||
| 19 | struct BufferRange { | ||
| 20 | std::uint32_t arena = 0; | ||
| 21 | std::uint64_t offset = 0; | ||
| 22 | std::uint64_t size = 0; | ||
| 23 | bool operator==(const BufferRange&) const = default; | ||
| 24 | }; | ||
| 25 | |||
| 26 | /// @brief The arenas of one memory use. Ranges are handed out by offset; a full arena is joined | ||
| 27 | /// by another, and a range larger than the arena size gets an arena of its own. No arena is | ||
| 28 | /// larger than the device's largest addressable buffer. | ||
| 29 | class Arenas { | ||
| 30 | public: | ||
| 31 | /// @param arenaSize Bytes per arena, capped at `DeviceContext::MaxBufferSize`. | ||
| 32 | Arenas(DeviceContext& context, MemoryUse use, std::string_view name, | ||
| 33 | std::uint64_t arenaSize = std::uint64_t{256} << 20); | ||
| 34 | /// Waits for the device, so ranges freed late are back before the arenas go. | ||
| 35 | ~Arenas(); | ||
| 36 | Arenas(const Arenas&) = delete; | ||
| 37 | Arenas& operator=(const Arenas&) = delete; | ||
| 38 | |||
| 39 | /// @brief Takes a range of `size` bytes at a multiple of `alignment`. | ||
| 40 | /// @throws DeviceError when `size` exceeds the largest buffer the device can address, or the | ||
| 41 | /// memory runs out. | ||
| 42 | BufferRange Allocate(std::uint64_t size, std::uint64_t alignment = 16); | ||
| 43 | /// @brief Gives a range back once every queue's work submitted so far is done. | ||
| 44 | void Free(const BufferRange& range); | ||
| 45 | |||
| 46 | /// @return The buffer of a range's arena: its handle, and its mapping for host-visible uses. | ||
| 47 | const Buffer& Storage(const BufferRange& range) const; | ||
| 48 | /// @return The range's device address, as shaders read it on Vulkan. | ||
| 49 | 17 | VkDeviceAddress Address(const BufferRange& range) const { | |
| 50 | 17 | return Storage(range).Address() + range.offset; | |
| 51 | } | ||
| 52 | 77 | std::size_t Count() const { | |
| 53 | 77 | return arenas_.size(); | |
| 54 | } | ||
| 55 | /// @return The bytes handed out and not yet given back. | ||
| 56 | std::uint64_t Used() const; | ||
| 57 | |||
| 58 | private: | ||
| 59 | struct Arena { | ||
| 60 | Buffer buffer; | ||
| 61 | VmaVirtualBlock block = VK_NULL_HANDLE; | ||
| 62 | /// Each range's allocation in `block`, by offset. | ||
| 63 | std::unordered_map<std::uint64_t, VmaVirtualAllocation> allocations; | ||
| 64 | std::uint64_t used = 0; | ||
| 65 | ~Arena(); | ||
| 66 | }; | ||
| 67 | |||
| 68 | std::uint32_t AddArena(std::uint64_t size); | ||
| 69 | void Release(const BufferRange& range); | ||
| 70 | |||
| 71 | DeviceContext& context_; | ||
| 72 | MemoryUse use_; | ||
| 73 | std::string name_; | ||
| 74 | std::uint64_t arenaSize_; | ||
| 75 | mutable std::mutex mutex_; | ||
| 76 | std::vector<std::unique_ptr<Arena>> arenas_; | ||
| 77 | }; | ||
| 78 | |||
| 79 | } // namespace ernie | ||
| 80 |