render/sensor_pipeline_runner.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include "chrono_batch_render/render/sensor_pipeline_runner.h" | ||
| 3 | |||
| 4 | #include <algorithm> | ||
| 5 | #include <atomic> | ||
| 6 | |||
| 7 | #include "chrono_batch_render/device/queue.h" | ||
| 8 | #include "chrono_batch_render/render/sensor_pipelines.h" | ||
| 9 | #include "chrono_batch_render/scene/scene_impl.h" | ||
| 10 | #include "chrono_batch_render/sensor_pipeline/node.h" | ||
| 11 | #include "chrono_batch_render/sensor_pipeline/resolve.h" | ||
| 12 | #include "chrono_batch_render/util/error.h" | ||
| 13 | |||
| 14 | namespace chrono::batch_render { | ||
| 15 | |||
| 16 | namespace { | ||
| 17 | |||
| 18 | 30 | uint32_t PatternChannels(DetectorPattern pattern) { | |
| 19 | 30 | switch (pattern) { | |
| 20 | case DetectorPattern::Mono: | ||
| 21 | return 1; | ||
| 22 | ✗ | case DetectorPattern::Rccc: | |
| 23 | ✗ | return 2; | |
| 24 | ✗ | case DetectorPattern::Bands: | |
| 25 | ✗ | return kDetectorBands; | |
| 26 | 29 | default: | |
| 27 | 29 | return 3; | |
| 28 | } | ||
| 29 | } | ||
| 30 | |||
| 31 | 30 | bool PatternMosaics(DetectorPattern pattern) { | |
| 32 | 30 | switch (pattern) { | |
| 33 | case DetectorPattern::Rggb: | ||
| 34 | case DetectorPattern::Bggr: | ||
| 35 | case DetectorPattern::Grbg: | ||
| 36 | case DetectorPattern::Gbrg: | ||
| 37 | case DetectorPattern::Rccb: | ||
| 38 | case DetectorPattern::Rccc: | ||
| 39 | return true; | ||
| 40 | 30 | default: | |
| 41 | 30 | return false; | |
| 42 | } | ||
| 43 | } | ||
| 44 | |||
| 45 | 30 | DetectorRecord Record(const DetectorDesc& desc) { | |
| 46 | 30 | DetectorRecord record{}; | |
| 47 | 30 | record.pattern = uint32_t(desc.pattern); | |
| 48 | 30 | record.channels = PatternChannels(desc.pattern); | |
| 49 | 30 | record.fullWell = desc.fullWell; | |
| 50 | 30 | record.readNoise = desc.readNoise; | |
| 51 | 30 | record.darkCurrent = desc.darkCurrent; | |
| 52 | 30 | record.darkCurrentDoubling = desc.darkCurrentDoubling; | |
| 53 | 30 | record.temperature = desc.temperature; | |
| 54 | 30 | record.prnu = desc.prnu; | |
| 55 | 30 | record.dsnu = desc.dsnu; | |
| 56 | 30 | record.systemGain = desc.systemGain; | |
| 57 | 30 | record.bitDepth = desc.bitDepth; | |
| 58 | 30 | record.blackLevel = desc.blackLevel; | |
| 59 | 30 | record.seed = desc.seed; | |
| 60 | 30 | return record; | |
| 61 | } | ||
| 62 | |||
| 63 | std::atomic<uint32_t> g_standaloneDefinitions{1u << 24}; | ||
| 64 | |||
| 65 | } | ||
| 66 | |||
| 67 | struct SensorPipelineRunner::State { | ||
| 68 | device::Device* device = nullptr; | ||
| 69 | std::shared_ptr<const sensor_pipeline::NodeRegistry> registry; | ||
| 70 | std::string path; | ||
| 71 | sensor_pipeline::ResolvedSensorPipeline instance; | ||
| 72 | std::shared_ptr<const CompiledSensorPipeline> compiled; | ||
| 73 | std::optional<std::vector<uint32_t>> words; | ||
| 74 | |||
| 75 | 9 | State(device::Device& on, std::shared_ptr<const sensor_pipeline::NodeRegistry> nodes, | |
| 76 | std::string name, sensor_pipeline::ResolvedSensorPipeline resolved) | ||
| 77 | 9 | : device(&on), registry(std::move(nodes)), path(std::move(name)), | |
| 78 | 9 | instance(std::move(resolved)) { | |
| 79 | 9 | auto definition = std::make_shared<sensor_pipeline::SensorPipelineDefinition>( | |
| 80 | 18 | sensor_pipeline::BuildDefinition(g_standaloneDefinitions++, instance, *registry)); | |
| 81 | 9 | compiled = CompileSensorPipeline(on, std::move(definition), *registry); | |
| 82 | 9 | } | |
| 83 | }; | ||
| 84 | |||
| 85 | 35 | std::optional<SensorPipelineNodeInfo> FindSensorPipelineNode(std::string_view id) { | |
| 86 | 35 | const std::shared_ptr<const sensor_pipeline::NodeRegistry> nodes = | |
| 87 | 35 | sensor_pipeline::SensorPipelineNodes(); | |
| 88 | 35 | const sensor_pipeline::NodeRegistry& registry = *nodes; | |
| 89 | 35 | const uint32_t found = registry.Find(id); | |
| 90 | 35 | if (found == sensor_pipeline::kNoDef) { | |
| 91 | 1 | return std::nullopt; | |
| 92 | } | ||
| 93 | 34 | const sensor_pipeline::NodeDef& def = registry[found]; | |
| 94 | 34 | SensorPipelineNodeInfo info{ | |
| 95 | 34 | .id = def.id, .kind = std::string(sensor_pipeline::kNodeKindNames[uint32_t(def.kind)])}; | |
| 96 | 65 | for (const sensor_pipeline::PortDef& port : def.inputs) { | |
| 97 | 31 | info.inputs.push_back( | |
| 98 | 31 | {port.name, std::string(sensor_pipeline::Signal(port.type).name), port.optional}); | |
| 99 | } | ||
| 100 | 68 | for (const sensor_pipeline::PortDef& port : def.outputs) { | |
| 101 | 34 | info.outputs.push_back( | |
| 102 | 34 | {port.name, std::string(sensor_pipeline::Signal(port.type).name), port.optional}); | |
| 103 | } | ||
| 104 | 34 | return info; | |
| 105 | 199 | } | |
| 106 | |||
| 107 | 4 | SensorPipelineRunner::SensorPipelineRunner(const Ref<Scene>& scene, std::string_view path) { | |
| 108 | 4 | if (scene->Pending()) { | |
| 109 | 2 | scene->Publish().Wait(); | |
| 110 | } | ||
| 111 | 4 | const std::optional<Prim> prim = scene->Find(path); | |
| 112 | 4 | if (!prim || !prim->IsA<SensorPipeline>()) { | |
| 113 | ✗ | throw Error("{} isn't a SensorPipeline", path); | |
| 114 | } | ||
| 115 | 4 | const SensorPipelineTable& table = SceneImpl::Get(*scene).GetSensorPipelines(); | |
| 116 | 4 | const std::shared_ptr<const SensorPipelineInstance> instance = table.Find(prim->Index()); | |
| 117 | 4 | if (!instance) { | |
| 118 | ✗ | throw Error("{} isn't resolved yet", path); | |
| 119 | } | ||
| 120 | 4 | if (!instance->error.empty()) { | |
| 121 | ✗ | throw Error("{}", instance->error); | |
| 122 | } | ||
| 123 | 8 | m_state = std::make_unique<State>(scene->GetDevice(), table.Registry(), std::string(path), | |
| 124 | 8 | instance->resolved); | |
| 125 | 4 | m_state->words = table.ReadWords(*instance); | |
| 126 | 4 | } | |
| 127 | |||
| 128 | 5 | SensorPipelineRunner::SensorPipelineRunner(device::Device& device, | |
| 129 | 5 | const sensor_pipeline::SensorPipelineDesc& desc) | |
| 130 | 10 | : SensorPipelineRunner(device, desc, *sensor_pipeline::SensorPipelineNodes()) {} | |
| 131 | |||
| 132 | 5 | SensorPipelineRunner::SensorPipelineRunner(device::Device& device, | |
| 133 | const sensor_pipeline::SensorPipelineDesc& desc, | ||
| 134 | 5 | const sensor_pipeline::NodeRegistry& registry) | |
| 135 | 5 | : m_state(std::make_unique<State>( | |
| 136 | 5 | device, std::make_shared<sensor_pipeline::NodeRegistry>(registry), desc.path, | |
| 137 | 10 | sensor_pipeline::ResolveSensorPipeline(desc, registry))) {} | |
| 138 | |||
| 139 | 9 | SensorPipelineRunner::~SensorPipelineRunner() = default; | |
| 140 | ✗ | SensorPipelineRunner::SensorPipelineRunner(SensorPipelineRunner&&) noexcept = default; | |
| 141 | ✗ | SensorPipelineRunner& SensorPipelineRunner::operator=(SensorPipelineRunner&&) noexcept = default; | |
| 142 | |||
| 143 | 3 | device::Device& SensorPipelineRunner::GetDevice() const { | |
| 144 | 3 | return *m_state->device; | |
| 145 | } | ||
| 146 | |||
| 147 | 2 | std::vector<std::string> SensorPipelineRunner::Sources() const { | |
| 148 | 2 | std::vector<std::string> sources; | |
| 149 | 10 | for (const sensor_pipeline::ResolvedNode& node : m_state->instance.nodes) { | |
| 150 | 8 | if ((*m_state->registry)[node.def].kind == sensor_pipeline::NodeKind::Source) { | |
| 151 | 2 | sources.push_back(node.path); | |
| 152 | } | ||
| 153 | } | ||
| 154 | 2 | return sources; | |
| 155 | ✗ | } | |
| 156 | |||
| 157 | 2 | std::vector<std::string> SensorPipelineRunner::Outputs() const { | |
| 158 | 2 | std::vector<std::string> outputs; | |
| 159 | 4 | for (const sensor_pipeline::ResolvedOutput& output : m_state->instance.outputs) { | |
| 160 | 2 | outputs.push_back(output.name); | |
| 161 | } | ||
| 162 | 2 | return outputs; | |
| 163 | ✗ | } | |
| 164 | |||
| 165 | 20 | SensorPipelineOutputs SensorPipelineRunner::Run(const SensorPipelineRunDesc& desc) const { | |
| 166 | 20 | device::Tensor state; | |
| 167 | 37 | return Run(desc, state); | |
| 168 | 20 | } | |
| 169 | |||
| 170 | 30 | SensorPipelineOutputs SensorPipelineRunner::Run(const SensorPipelineRunDesc& desc, | |
| 171 | device::Tensor& state) const { | ||
| 172 | 30 | const sensor_pipeline::NodeRegistry& registry = *m_state->registry; | |
| 173 | 30 | const sensor_pipeline::ResolvedSensorPipeline& instance = m_state->instance; | |
| 174 | 30 | const sensor_pipeline::SensorPipelineDefinition& definition = *m_state->compiled->definition; | |
| 175 | 30 | const sensor_pipeline::GeneratedSensorPipeline& generated = m_state->compiled->generated; | |
| 176 | 30 | const DetectorRecord detector = Record(desc.detector); | |
| 177 | 30 | const std::vector<sensor_pipeline::SignalLayout> layouts = | |
| 178 | sensor_pipeline::PropagateLayouts(instance, registry, | ||
| 179 | 30 | {.channels = detector.channels, | |
| 180 | 30 | .pattern = detector.pattern, | |
| 181 | 30 | .mosaics = PatternMosaics(desc.detector.pattern)}); | |
| 182 | |||
| 183 | 30 | int64_t captures = -1; | |
| 184 | 30 | int64_t height = 0; | |
| 185 | 30 | int64_t width = 0; | |
| 186 | 33 | std::vector<const device::Tensor*> sources(generated.slots.size(), nullptr); | |
| 187 | 87 | for (uint32_t slot = 0; slot < definition.segments.stored.size(); ++slot) { | |
| 188 | 60 | const sensor_pipeline::Value& value = instance.values[generated.slots[slot]]; | |
| 189 | 60 | const sensor_pipeline::ResolvedNode& node = instance.nodes[value.node]; | |
| 190 | 60 | if (registry[node.def].kind != sensor_pipeline::NodeKind::Source) { | |
| 191 | 30 | continue; | |
| 192 | } | ||
| 193 | 30 | auto found = desc.inputs.find(node.path); | |
| 194 | 41 | if (found == desc.inputs.end() && node.path.starts_with(m_state->path + "/")) { | |
| 195 | 22 | found = desc.inputs.find(std::string_view(node.path).substr(m_state->path.size() + 1)); | |
| 196 | } | ||
| 197 | 30 | if (found == desc.inputs.end()) { | |
| 198 | 2 | throw Error("The run has no input for {}", node.path); | |
| 199 | } | ||
| 200 | 28 | const device::Tensor& tensor = found->second; | |
| 201 | 28 | const Span<const int64_t> shape = tensor.Shape(); | |
| 202 | 28 | const uint32_t channels = layouts[generated.slots[slot]].channels; | |
| 203 | 28 | const int64_t given = shape.size() == 4 ? shape[3] : 1; | |
| 204 | 28 | if (tensor.Type() != device::DType::Float32 || (shape.size() != 3 && shape.size() != 4) || | |
| 205 | 28 | given != channels) { | |
| 206 | 1 | throw Error("{}'s input must be float32 shaped [sensors, height, width{}]", node.path, | |
| 207 | 2 | channels == 1 ? std::string() : ", " + std::to_string(channels)); | |
| 208 | } | ||
| 209 | 27 | if (captures >= 0 && (shape[0] != captures || shape[1] != height || shape[2] != width)) { | |
| 210 | ✗ | throw Error("{}'s input is shaped unlike the run's other inputs", node.path); | |
| 211 | } | ||
| 212 | 27 | captures = shape[0]; | |
| 213 | 27 | height = shape[1]; | |
| 214 | 27 | width = shape[2]; | |
| 215 | 27 | sources[slot] = &tensor; | |
| 216 | } | ||
| 217 | 27 | if (captures <= 0 || height <= 0 || width <= 0) { | |
| 218 | ✗ | throw Error("A standalone run needs at least one nonempty input"); | |
| 219 | } | ||
| 220 | |||
| 221 | 27 | device::Device& device = *m_state->device; | |
| 222 | 27 | const uint64_t perCapture = uint64_t(height) * width; | |
| 223 | 27 | const uint64_t elements = perCapture * captures; | |
| 224 | 30 | std::vector<device::Tensor> tensors(generated.slots.size()); | |
| 225 | 27 | std::vector<SensorPipelineValue> values; | |
| 226 | 27 | const auto storedCount = uint32_t(definition.segments.stored.size()); | |
| 227 | 112 | for (uint32_t slot = 0; slot < generated.slots.size(); ++slot) { | |
| 228 | 85 | const sensor_pipeline::SignalLayout& layout = layouts[generated.slots[slot]]; | |
| 229 | 85 | if (sources[slot] != nullptr) { | |
| 230 | 27 | tensors[slot] = *sources[slot]; | |
| 231 | } else { | ||
| 232 | 58 | std::vector<int64_t> shape = {captures}; | |
| 233 | 58 | if (!layout.perCapture) { | |
| 234 | 46 | shape.insert(shape.end(), {height, width}); | |
| 235 | } | ||
| 236 | 58 | if (layout.channels > 1) { | |
| 237 | 58 | shape.push_back(layout.channels); | |
| 238 | } | ||
| 239 | 116 | const uint64_t count = | |
| 240 | 58 | uint64_t(captures) * (layout.perCapture ? 1 : perCapture) * layout.channels; | |
| 241 | 116 | tensors[slot] = device::Tensor( | |
| 242 | 116 | device.CreateBuffer({.bytes = count * sizeof(float), .label = "pipeline value"}), 0, | |
| 243 | 116 | device::DType::Float32, std::move(shape)); | |
| 244 | 58 | } | |
| 245 | 85 | values.push_back( | |
| 246 | 85 | {tensors[slot].Address(), layout.channels, layout.mosaic ? kSignalMosaic : 0u}); | |
| 247 | } | ||
| 248 | |||
| 249 | 27 | std::vector<uint32_t> words = m_state->words | |
| 250 | 27 | ? *m_state->words | |
| 251 | 27 | : sensor_pipeline::PackWords(definition, instance, registry); | |
| 252 | 27 | words.push_back(0); | |
| 253 | 27 | CameraRecord camera{}; | |
| 254 | 27 | camera.exposureTime = desc.exposureTime; | |
| 255 | 27 | camera.iso = desc.iso; | |
| 256 | 27 | std::vector<SensorPipelineValue> slots; | |
| 257 | 27 | std::vector<SensorPipelineCapture> records; | |
| 258 | 27 | const Ref<device::Buffer> wordBuffer = device.CreateBuffer({.label = "pipeline words"}, words); | |
| 259 | 54 | const uint64_t stateBytes = std::max<uint64_t>( | |
| 260 | 54 | uint64_t(definition.initialState.size()) * captures * sizeof(float), sizeof(float)); | |
| 261 | 27 | if (state.Empty() || state.GetBuffer().Bytes() != stateBytes) { | |
| 262 | 20 | std::vector<float> states; | |
| 263 | 40 | for (int64_t capture = 0; capture < captures; ++capture) { | |
| 264 | 20 | states.insert(states.end(), definition.initialState.begin(), | |
| 265 | definition.initialState.end()); | ||
| 266 | } | ||
| 267 | 20 | states.resize(stateBytes / sizeof(float)); | |
| 268 | 60 | state = device::Tensor(device.CreateBuffer({.label = "pipeline state"}, states), 0, | |
| 269 | 40 | device::DType::Float32, {int64_t(states.size())}); | |
| 270 | 20 | } | |
| 271 | 54 | for (int64_t capture = 0; capture < captures; ++capture) { | |
| 272 | 112 | for (uint32_t slot = 0; slot < values.size(); ++slot) { | |
| 273 | 85 | SensorPipelineValue value = values[slot]; | |
| 274 | 85 | const bool stat = layouts[generated.slots[slot]].perCapture; | |
| 275 | 170 | value.address += | |
| 276 | 85 | uint64_t(capture) * (stat ? 1 : perCapture) * value.channels * sizeof(float); | |
| 277 | 85 | slots.push_back(value); | |
| 278 | } | ||
| 279 | 27 | records.push_back( | |
| 280 | 27 | {.firstElement = uint64_t(capture) * perCapture, | |
| 281 | 27 | .words = {wordBuffer->Address()}, | |
| 282 | .width = uint32_t(width), | ||
| 283 | .height = uint32_t(height), | ||
| 284 | 27 | .capture = desc.capture, | |
| 285 | 27 | .identity = desc.identity + uint32_t(capture), | |
| 286 | 27 | .firstValue = uint32_t(capture * values.size()), | |
| 287 | 27 | .state = {state.Address() + | |
| 288 | 27 | uint64_t(capture) * definition.initialState.size() * sizeof(float)}}); | |
| 289 | } | ||
| 290 | 27 | const Ref<device::Buffer> cameras = | |
| 291 | 54 | device.CreateBuffer({.label = "pipeline camera"}, std::vector<CameraRecord>{camera}); | |
| 292 | 27 | const Ref<device::Buffer> detectors = | |
| 293 | 54 | device.CreateBuffer({.label = "pipeline detector"}, std::vector<DetectorRecord>{detector}); | |
| 294 | 27 | RenderTables tables{}; | |
| 295 | 27 | tables.cameras.address = cameras->Address(); | |
| 296 | 27 | tables.detectors.address = detectors->Address(); | |
| 297 | 27 | const Ref<device::Buffer> tableBuffer = | |
| 298 | 54 | device.CreateBuffer({.label = "pipeline tables"}, std::vector<RenderTables>{tables}); | |
| 299 | 27 | const Ref<device::Buffer> valueBuffer = | |
| 300 | 27 | device.CreateBuffer({.label = "pipeline values"}, slots); | |
| 301 | 27 | const Ref<device::Buffer> captureBuffer = | |
| 302 | 27 | device.CreateBuffer({.label = "pipeline captures"}, records); | |
| 303 | 27 | const SensorPipelineParams params{.tables = {tableBuffer->Address()}, | |
| 304 | 27 | .captures = {captureBuffer->Address()}, | |
| 305 | 27 | .values = {valueBuffer->Address()}, | |
| 306 | .elementCount = elements, | ||
| 307 | 27 | .captureCount = uint32_t(captures)}; | |
| 308 | 27 | device::CommandEncoder encoder = device.GetQueue().CreateEncoder(); | |
| 309 | 27 | RecordSensorPipeline(encoder, *m_state->compiled, params, uint32_t(captures), elements); | |
| 310 | 27 | device.GetQueue().Submit(encoder.Finish()).Wait(); | |
| 311 | |||
| 312 | 27 | SensorPipelineOutputs outputs; | |
| 313 | 55 | for (uint32_t output = 0; output < instance.outputs.size(); ++output) { | |
| 314 | 28 | outputs.emplace(instance.outputs[output].name, tensors[storedCount + output]); | |
| 315 | } | ||
| 316 | 27 | return outputs; | |
| 317 | 618 | } | |
| 318 | |||
| 319 | 10 | SensorPipelineOutputs SensorPipelineSession::Run(SensorPipelineRunDesc desc) { | |
| 320 | 10 | desc.capture = m_captures++; | |
| 321 | 10 | return m_pipeline.Run(desc, m_state); | |
| 322 | } | ||
| 323 | |||
| 324 | } | ||
| 325 |