render/scene_columns.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file scene_columns.h | ||
| 3 | /// @brief Scene columns whose source of truth is the device: per-step values, such as the poses | ||
| 4 | /// of thousands of worlds, written in place by kernels, CUDA or Warp, and bound by passes as they | ||
| 5 | /// are. The host scene keeps the structure and the authored values; per-step data never passes | ||
| 6 | /// through it. | ||
| 7 | #pragma once | ||
| 8 | |||
| 9 | #include "ernie/device/column.h" | ||
| 10 | #include "ernie/device/staging.h" | ||
| 11 | #include "ernie/graph/graph.h" | ||
| 12 | #include "ernie/scene/scene.h" | ||
| 13 | |||
| 14 | #include <cstdint> | ||
| 15 | #include <functional> | ||
| 16 | #include <memory> | ||
| 17 | #include <optional> | ||
| 18 | #include <span> | ||
| 19 | #include <string_view> | ||
| 20 | #include <unordered_map> | ||
| 21 | #include <vector> | ||
| 22 | |||
| 23 | namespace ernie { | ||
| 24 | |||
| 25 | /// @brief How a host column's rows look on the device: floats for the scene's doubles. | ||
| 26 | 114 | struct DeviceElement { | |
| 27 | DType type = DType::Float32; | ||
| 28 | std::vector<std::int64_t> shape; | ||
| 29 | }; | ||
| 30 | |||
| 31 | /// @return The device element of a column type: Double becomes one float, Vec2 to Vec4 that many, | ||
| 32 | /// Quat four (x, y, z, w), Matrix sixteen in glm's column-major order, Int an int32 and Bool a | ||
| 33 | /// byte. Variable-length types have none. | ||
| 34 | std::optional<DeviceElement> ToDeviceElement(ColumnType type); | ||
| 35 | |||
| 36 | /// @brief Device columns that follow their scene buckets' rows. Each starts from the scene's | ||
| 37 | /// values, then the device owns them: when prims are removed, the bucket's swap-removes are | ||
| 38 | /// replayed as one batch of row copies on the device, rows added on the host start from the | ||
| 39 | /// host's values, and values authored on the host overwrite those rows. Rows the device wrote and | ||
| 40 | /// the host only moved keep the device's values. | ||
| 41 | class SceneColumns { | ||
| 42 | public: | ||
| 43 | /// @param arenas Where the columns live: `MemoryUse::Shared` arenas to share them with CUDA. | ||
| 44 | SceneColumns(DeviceContext& context, Arenas& arenas, const Scene& scene); | ||
| 45 | ~SceneColumns(); | ||
| 46 | SceneColumns(const SceneColumns&) = delete; | ||
| 47 | SceneColumns& operator=(const SceneColumns&) = delete; | ||
| 48 | |||
| 49 | /// @brief Puts a bucket's column of `attribute` on the device, from the scene's values. | ||
| 50 | /// @throws std::invalid_argument when the bucket has no such column, or its type has no | ||
| 51 | /// device element. | ||
| 52 | const DeviceColumn& Bind(BucketId bucket, std::string_view attribute); | ||
| 53 | /// @brief Puts a bucket's column of `attribute` on the device in `column`, such as a CUDA, | ||
| 54 | /// PyTorch or Warp tensor wrapped in place (`CudaInterop::Wrap`): its values are what the | ||
| 55 | /// device holds, not the scene's, and values the host authors later still overwrite them. | ||
| 56 | /// @throws std::invalid_argument unless `column` has the bucket's rows and the attribute's | ||
| 57 | /// device element, or when the attribute is bound already. | ||
| 58 | const DeviceColumn& Bind(BucketId bucket, std::string_view attribute, | ||
| 59 | std::unique_ptr<DeviceColumn> column); | ||
| 60 | /// @return The bucket's device column of `attribute`, or null when it is not bound. | ||
| 61 | const DeviceColumn* Find(BucketId bucket, std::string_view attribute) const; | ||
| 62 | /// @brief Follows the scene's changes since the last sync, after `waits`. | ||
| 63 | /// @return The ticket passes wait on before reading the columns. | ||
| 64 | Ticket Sync(std::span<const Ticket> waits = {}); | ||
| 65 | /// @brief Follows the scene's changes as a pass of a graph run on the main queue: rows move | ||
| 66 | /// and host values land before later passes that import the columns read them. | ||
| 67 | void AddSync(graph::Graph& graph); | ||
| 68 | |||
| 69 | private: | ||
| 70 | 109 | struct Binding { | |
| 71 | BucketId bucket = 0; | ||
| 72 | AttributeId attribute = 0; | ||
| 73 | std::uint64_t cursor = 0; | ||
| 74 | std::unique_ptr<DeviceColumn> column; | ||
| 75 | /// The device bytes of what the host last gave each row, to tell host edits apart. | ||
| 76 | std::vector<std::byte> shadow; | ||
| 77 | }; | ||
| 78 | |||
| 79 | /// A bucket's host column of an attribute, and its device element. | ||
| 80 | 36 | struct Host { | |
| 81 | AttributeId attribute = 0; | ||
| 82 | const Column* column = nullptr; | ||
| 83 | DeviceElement element; | ||
| 84 | std::uint32_t rows = 0; | ||
| 85 | }; | ||
| 86 | |||
| 87 | /// @throws std::invalid_argument when the bucket has no such column, or it has no device | ||
| 88 | /// element. | ||
| 89 | Host HostColumn(BucketId bucket, std::string_view attribute) const; | ||
| 90 | /// Where each row changed since a cursor has its values from: a row before, or the host. | ||
| 91 | using Origins = std::unordered_map<std::uint32_t, std::uint32_t>; | ||
| 92 | |||
| 93 | static Origins RowOrigins(const ChangeSet& changes, BucketId bucket); | ||
| 94 | /// @return The row `row`'s values come from; rows that did not change keep their own. | ||
| 95 | static std::uint32_t Origin(const Origins& origins, std::uint32_t row); | ||
| 96 | /// @return The device copies that move the kept rows to their new places, moving the shadow | ||
| 97 | /// alike; rows new from the host are added to `hostRows`. | ||
| 98 | std::vector<CopyRegion> MoveRows(Binding& binding, const Origins& origins, std::uint32_t rows, | ||
| 99 | std::vector<std::uint32_t>& hostRows); | ||
| 100 | /// Adds the rows whose host values were authored anew to `hostRows`, taking their values | ||
| 101 | /// into the shadow. | ||
| 102 | void FindAuthoredRows(const ChangeSet& changes, Binding& binding, const Origins& origins, | ||
| 103 | std::uint32_t rows, std::vector<std::uint32_t>& hostRows); | ||
| 104 | /// Uploads the binding's shadow at sorted `rows`. | ||
| 105 | void Upload(const Binding& binding, std::span<const std::uint32_t> rows); | ||
| 106 | /// Follows the changes on the host: resizes the columns, plans the row moves and stages the | ||
| 107 | /// host's rows. | ||
| 108 | /// @return What the moves and uploads wait on. | ||
| 109 | std::vector<Ticket> Prepare(std::span<const Ticket> waits); | ||
| 110 | /// Records the planned moves, then the staged uploads. | ||
| 111 | void Record(CommandList commands, const std::function<void(std::shared_ptr<const void>)>& keep); | ||
| 112 | const Binding* FindBinding(BucketId bucket, AttributeId attribute) const; | ||
| 113 | |||
| 114 | DeviceContext& context_; | ||
| 115 | Arenas& arenas_; | ||
| 116 | const Scene& scene_; | ||
| 117 | Uploader uploader_; | ||
| 118 | std::vector<Binding> bindings_; | ||
| 119 | /// Authored ranges' device bytes, compared with the shadow. | ||
| 120 | std::vector<std::byte> scratch_; | ||
| 121 | /// Each column's row copies, from the rows before the sync to the rows after it. | ||
| 122 | std::vector<std::pair<const DeviceColumn*, std::vector<CopyRegion>>> moves_; | ||
| 123 | }; | ||
| 124 | |||
| 125 | } // namespace ernie | ||
| 126 |