device/column.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file column.h | ||
| 3 | /// @brief Columns whose source of truth is device memory: per-step data such as poses, written | ||
| 4 | /// and read in place on the GPU, which the host sees only through an explicit readback. | ||
| 5 | #pragma once | ||
| 6 | |||
| 7 | #include "ernie/device/arena.h" | ||
| 8 | #include "ernie/device/staging.h" | ||
| 9 | #include "ernie/shaders/shared.h" | ||
| 10 | #include "ernie/util/algorithm.h" | ||
| 11 | #include "ernie/util/tensor.h" | ||
| 12 | |||
| 13 | #include <cstdint> | ||
| 14 | #include <memory> | ||
| 15 | #include <span> | ||
| 16 | #include <string> | ||
| 17 | #include <string_view> | ||
| 18 | #include <vector> | ||
| 19 | |||
| 20 | namespace ernie { | ||
| 21 | |||
| 22 | /// @brief Rows of one element shape in device memory: a range of Ernie's arenas, or memory | ||
| 23 | /// another API owns and shares in place. Kernels, and through export CUDA, Warp or PyTorch, | ||
| 24 | /// write rows in place; nothing is mirrored on the host. | ||
| 25 | class DeviceColumn { | ||
| 26 | public: | ||
| 27 | /// @param element The shape of one row, such as {4, 3} for a transform; empty for a scalar. | ||
| 28 | DeviceColumn(Arenas& arenas, std::string_view name, DType type, | ||
| 29 | std::vector<std::int64_t> element, std::uint64_t rows); | ||
| 30 | /// @brief A column over memory it does not own, such as a CUDA tensor imported in place, | ||
| 31 | /// from `offset` in `storage`. It cannot grow; it keeps `owner` alive, and both until the | ||
| 32 | /// device is done with them. | ||
| 33 | DeviceColumn(DeviceContext& context, std::shared_ptr<const Buffer> storage, | ||
| 34 | std::uint64_t offset, std::string_view name, DType type, | ||
| 35 | std::vector<std::int64_t> element, std::uint64_t rows, | ||
| 36 | std::shared_ptr<const void> owner = {}); | ||
| 37 | ~DeviceColumn(); | ||
| 38 | DeviceColumn(const DeviceColumn&) = delete; | ||
| 39 | DeviceColumn& operator=(const DeviceColumn&) = delete; | ||
| 40 | |||
| 41 | /// @brief Grows or shrinks to `rows`; new rows are undefined. Within the capacity the rows | ||
| 42 | /// stay where they are. Past it the capacity at least doubles, so growing a row at a time | ||
| 43 | /// costs amortized constant time: the rows are copied to a new range on `queue` after | ||
| 44 | /// `waits`, and the old one is freed once the copy is done. | ||
| 45 | /// @return The copy's ticket; a default one when nothing was copied. | ||
| 46 | /// @throws std::logic_error when a column over another API's memory would grow. | ||
| 47 | Ticket Resize(std::uint64_t rows, Queue& queue, std::span<const Ticket> waits = {}); | ||
| 48 | |||
| 49 | ✗ | const std::string& Name() const { | |
| 50 | 480 | return name_; | |
| 51 | } | ||
| 52 | 6 | DType Type() const { | |
| 53 | 6 | return type_; | |
| 54 | } | ||
| 55 | ✗ | const std::vector<std::int64_t>& Element() const { | |
| 56 | ✗ | return element_; | |
| 57 | } | ||
| 58 | 262 | std::uint64_t Rows() const { | |
| 59 | 261 | return rows_; | |
| 60 | } | ||
| 61 | /// @return The rows, then the element's shape: the column as a tensor sees it. | ||
| 62 | ✗ | std::vector<std::int64_t> Shape() const { | |
| 63 | ✗ | std::vector<std::int64_t> shape{static_cast<std::int64_t>(rows_)}; | |
| 64 | ✗ | shape.insert(shape.end(), element_.begin(), element_.end()); | |
| 65 | ✗ | return shape; | |
| 66 | ✗ | } | |
| 67 | /// @return The rows' bytes; the range may hold more, for growth. | ||
| 68 | 514 | std::uint64_t Bytes() const { | |
| 69 | 514 | return rows_ * rowBytes_; | |
| 70 | } | ||
| 71 | 2 | std::uint64_t Capacity() const { | |
| 72 | 2 | return capacity_; | |
| 73 | } | ||
| 74 | 3772 | std::uint64_t RowBytes() const { | |
| 75 | 2160 | return rowBytes_; | |
| 76 | } | ||
| 77 | /// @return Whether the memory is another API's, shared in place. | ||
| 78 | ✗ | bool Adopted() const { | |
| 79 | ✗ | return arenas_ == nullptr; | |
| 80 | } | ||
| 81 | /// @return The rows' range: (arena, offset) on every backend; for adopted memory, the offset | ||
| 82 | /// in `Storage`. | ||
| 83 | const BufferRange& Range() const { | ||
| 84 | return range_; | ||
| 85 | } | ||
| 86 | 1611 | const Buffer& Storage() const { | |
| 87 | 1611 | return arenas_ != nullptr ? arenas_->Storage(range_) : *storage_; | |
| 88 | } | ||
| 89 | /// @return The first row's device address, as Vulkan shaders read it. | ||
| 90 | 422 | VkDeviceAddress Address() const { | |
| 91 | 148 | return Storage().Address() + range_.offset; | |
| 92 | } | ||
| 93 | |||
| 94 | private: | ||
| 95 | BufferRange Allocate(std::uint64_t rows); | ||
| 96 | |||
| 97 | Arenas* arenas_ = nullptr; | ||
| 98 | DeviceContext* context_ = nullptr; | ||
| 99 | std::shared_ptr<const Buffer> storage_; | ||
| 100 | std::shared_ptr<const void> owner_; | ||
| 101 | std::string name_; | ||
| 102 | DType type_; | ||
| 103 | std::vector<std::int64_t> element_; | ||
| 104 | std::uint64_t rowBytes_; | ||
| 105 | std::uint64_t rows_; | ||
| 106 | std::uint64_t capacity_; | ||
| 107 | BufferRange range_; | ||
| 108 | }; | ||
| 109 | |||
| 110 | /// @return A shader's handle of a column's rows as `T`s. | ||
| 111 | 394 | template <typename T> shaders::BufferRef<T> ToBufferRef(const DeviceColumn& column) { | |
| 112 | 394 | return {column.Address()}; | |
| 113 | } | ||
| 114 | |||
| 115 | /// @brief Stages `rows` of a host mirror into the same rows of `column`. The rows are sorted and | ||
| 116 | /// made unique in place, and those past the mirror's end dropped; the caller clears them. | ||
| 117 | template <typename T> | ||
| 118 | 437 | void StageRows(Uploader& uploader, const DeviceColumn& column, const std::vector<T>& host, | |
| 119 | std::vector<std::uint32_t>& rows) { | ||
| 120 | 437 | SortUnique(rows); | |
| 121 | 491 | std::erase_if(rows, [&](std::uint32_t row) { return row >= host.size(); }); | |
| 122 | 437 | uploader.WriteRows(column.Storage(), column.Range().offset, column.RowBytes(), | |
| 123 | std::as_bytes(std::span(host)), rows); | ||
| 124 | 437 | } | |
| 125 | |||
| 126 | } // namespace ernie | ||
| 127 |