sensor_pipeline/codegen.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include "chrono_batch_render/sensor_pipeline/codegen.h" | ||
| 3 | |||
| 4 | #include <algorithm> | ||
| 5 | #include <cstddef> | ||
| 6 | #include <set> | ||
| 7 | #include <string_view> | ||
| 8 | |||
| 9 | #include <fmt/format.h> | ||
| 10 | |||
| 11 | #include "chrono_batch_render/util/error.h" | ||
| 12 | #include "chrono_batch_render/util/text.h" | ||
| 13 | |||
| 14 | extern "C" { | ||
| 15 | extern const unsigned char chrono_batch_render_sensor_pipeline_module[]; | ||
| 16 | extern const std::size_t chrono_batch_render_sensor_pipeline_module_size; | ||
| 17 | extern const unsigned char chrono_batch_render_sensor_pipeline_elements[]; | ||
| 18 | extern const std::size_t chrono_batch_render_sensor_pipeline_elements_size; | ||
| 19 | extern const unsigned char chrono_batch_render_sensor_pipeline_captures[]; | ||
| 20 | extern const std::size_t chrono_batch_render_sensor_pipeline_captures_size; | ||
| 21 | extern const unsigned char chrono_batch_render_sensor_pipeline_reduce_groups[]; | ||
| 22 | extern const std::size_t chrono_batch_render_sensor_pipeline_reduce_groups_size; | ||
| 23 | extern const unsigned char chrono_batch_render_sensor_pipeline_reduce_partial[]; | ||
| 24 | extern const std::size_t chrono_batch_render_sensor_pipeline_reduce_partial_size; | ||
| 25 | } | ||
| 26 | |||
| 27 | namespace chrono::batch_render::sensor_pipeline { | ||
| 28 | |||
| 29 | namespace { | ||
| 30 | |||
| 31 | 197 | std::string NodeType(const NodeDef& def, const ResolvedNode& node) { | |
| 32 | 197 | if (node.uniforms.empty()) { | |
| 33 | 158 | return def.code.type; | |
| 34 | } | ||
| 35 | 39 | std::vector<std::string> choices; | |
| 36 | 126 | for (const uint32_t choice : node.uniforms) { | |
| 37 | 87 | choices.push_back(std::to_string(choice)); | |
| 38 | } | ||
| 39 | 39 | return fmt::format("{}<{}>", def.code.type, Join(choices)); | |
| 40 | 39 | } | |
| 41 | |||
| 42 | 184 | std::string_view Embedded(const unsigned char* bytes, std::size_t size) { | |
| 43 | 184 | return {reinterpret_cast<const char*>(bytes), size}; | |
| 44 | } | ||
| 45 | |||
| 46 | 110 | std::string_view KernelTemplate(DispatchShape shape) { | |
| 47 | 110 | switch (shape) { | |
| 48 | 99 | case DispatchShape::Elements: | |
| 49 | 99 | return Embedded(chrono_batch_render_sensor_pipeline_elements, | |
| 50 | 99 | chrono_batch_render_sensor_pipeline_elements_size); | |
| 51 | 5 | case DispatchShape::ReduceGroups: | |
| 52 | 5 | return Embedded(chrono_batch_render_sensor_pipeline_reduce_groups, | |
| 53 | 5 | chrono_batch_render_sensor_pipeline_reduce_groups_size); | |
| 54 | 6 | case DispatchShape::Captures: | |
| 55 | 6 | return Embedded(chrono_batch_render_sensor_pipeline_captures, | |
| 56 | 6 | chrono_batch_render_sensor_pipeline_captures_size); | |
| 57 | } | ||
| 58 | ✗ | return {}; | |
| 59 | } | ||
| 60 | |||
| 61 | class Generator { | ||
| 62 | public: | ||
| 63 | 69 | Generator(const SensorPipelineDefinition& definition, const NodeRegistry& registry) | |
| 64 | 69 | : m_definition(definition), m_pipeline(definition.pipeline), m_registry(registry), | |
| 65 | 69 | m_storedSlot(definition.pipeline.values.size(), kNoValue) { | |
| 66 | 138 | m_generated.module = fmt::format( | |
| 67 | 69 | "chrono_batch_render.sensor_pipeline.generated.definition{}", definition.id); | |
| 68 | 69 | m_generated.path = fmt::format( | |
| 69 | 69 | "chrono_batch_render/sensor_pipeline/generated/definition{}.slang", definition.id); | |
| 70 | 169 | for (const StoredValue& stored : definition.segments.stored) { | |
| 71 | 100 | m_storedSlot[stored.value] = uint32_t(m_generated.slots.size()); | |
| 72 | 100 | m_generated.slots.push_back(stored.value); | |
| 73 | } | ||
| 74 | 155 | for (const ResolvedOutput& output : m_pipeline.outputs) { | |
| 75 | 86 | m_generated.slots.push_back(output.value); | |
| 76 | } | ||
| 77 | 69 | } | |
| 78 | |||
| 79 | 69 | GeneratedSensorPipeline Generate() && { | |
| 80 | 69 | AddImports(); | |
| 81 | 162 | for (uint32_t index = 0; index < m_definition.segments.segments.size(); ++index) { | |
| 82 | 93 | switch (m_definition.segments.segments[index].kind) { | |
| 83 | 5 | case NodeKind::Reduce: | |
| 84 | 5 | AddReduce(index); | |
| 85 | 5 | break; | |
| 86 | 6 | case NodeKind::Global: | |
| 87 | 6 | AddGlobal(index); | |
| 88 | 6 | break; | |
| 89 | 82 | case NodeKind::Point: | |
| 90 | 82 | case NodeKind::Stencil: | |
| 91 | 82 | AddFused(index); | |
| 92 | 82 | AddUpdates(index); | |
| 93 | 82 | break; | |
| 94 | ✗ | case NodeKind::Source: | |
| 95 | ✗ | case NodeKind::Compact: | |
| 96 | ✗ | throw Error("Pipeline passes of compact nodes can't be generated yet"); | |
| 97 | } | ||
| 98 | } | ||
| 99 | 69 | AddSourceCopies(); | |
| 100 | 69 | return std::move(m_generated); | |
| 101 | } | ||
| 102 | |||
| 103 | private: | ||
| 104 | 929 | const ResolvedNode& Node(uint32_t node) const { | |
| 105 | 177 | return m_pipeline.nodes[node]; | |
| 106 | } | ||
| 107 | |||
| 108 | 1209 | const NodeDef& Def(uint32_t node) const { | |
| 109 | 1209 | return m_registry[m_pipeline.nodes[node].def]; | |
| 110 | } | ||
| 111 | |||
| 112 | 288 | std::string_view SlangType(uint32_t value) const { | |
| 113 | 17 | return Signal(m_pipeline.values[value].type).slangType; | |
| 114 | } | ||
| 115 | |||
| 116 | 210 | std::string Context(uint32_t node) const { | |
| 117 | 210 | return fmt::format("PointContext(p, element, {}, {}, {})", node, | |
| 118 | 210 | m_definition.parameterOffsets[node], m_definition.stateOffsets[node]); | |
| 119 | } | ||
| 120 | |||
| 121 | 204 | std::vector<uint32_t> OutputSlots(uint32_t value) const { | |
| 122 | 204 | std::vector<uint32_t> slots; | |
| 123 | 204 | if (m_storedSlot[value] != kNoValue) { | |
| 124 | 32 | slots.push_back(m_storedSlot[value]); | |
| 125 | } | ||
| 126 | 204 | const auto stored = uint32_t(m_definition.segments.stored.size()); | |
| 127 | 479 | for (uint32_t output = 0; output < m_pipeline.outputs.size(); ++output) { | |
| 128 | 275 | if (m_pipeline.outputs[output].value == value) { | |
| 129 | 83 | slots.push_back(stored + output); | |
| 130 | } | ||
| 131 | } | ||
| 132 | 204 | return slots; | |
| 133 | ✗ | } | |
| 134 | |||
| 135 | 618 | void Emit(std::string_view text) { | |
| 136 | 3 | m_body += text; | |
| 137 | 618 | } | |
| 138 | |||
| 139 | 110 | void BeginKernel(std::string entry, DispatchShape shape) { | |
| 140 | 110 | m_body.clear(); | |
| 141 | 110 | m_generated.kernels.push_back({std::move(entry), shape}); | |
| 142 | 220 | } | |
| 143 | |||
| 144 | 110 | void EndKernel() { | |
| 145 | 110 | const GeneratedKernel& kernel = m_generated.kernels.back(); | |
| 146 | 110 | m_generated.source += | |
| 147 | 220 | FillTemplate(KernelTemplate(kernel.shape), {{"entry", kernel.entry}, {"body", m_body}}); | |
| 148 | 110 | } | |
| 149 | |||
| 150 | 199 | void Write(uint32_t value, std::string_view expression) { | |
| 151 | 309 | for (const uint32_t slot : OutputSlots(value)) { | |
| 152 | 220 | Emit(fmt::format(" WriteSignal(p, element, {}, {});\n", slot, expression)); | |
| 153 | 199 | } | |
| 154 | 199 | } | |
| 155 | |||
| 156 | 69 | void AddImports() { | |
| 157 | 69 | std::set<std::string> modules; | |
| 158 | 333 | for (const ResolvedNode& node : m_pipeline.nodes) { | |
| 159 | 264 | const NodeDef& def = m_registry[node.def]; | |
| 160 | 264 | if (def.kind != NodeKind::Source && !def.code.module.empty()) { | |
| 161 | 196 | modules.insert(def.code.module); | |
| 162 | } | ||
| 163 | } | ||
| 164 | 69 | std::string imports; | |
| 165 | 161 | for (const std::string& module : modules) { | |
| 166 | 184 | imports += fmt::format("import {};\n", module); | |
| 167 | } | ||
| 168 | 138 | m_generated.source = FillTemplate(Embedded(chrono_batch_render_sensor_pipeline_module, | |
| 169 | chrono_batch_render_sensor_pipeline_module_size), | ||
| 170 | 69 | {{"imports", imports}}); | |
| 171 | 69 | } | |
| 172 | |||
| 173 | 5 | void AddReduce(uint32_t index) { | |
| 174 | 5 | const uint32_t node = m_definition.segments.segments[index].nodes.front(); | |
| 175 | 5 | const ResolvedNode& resolved = Node(node); | |
| 176 | 5 | const NodeDef& def = Def(node); | |
| 177 | 5 | const uint32_t input = resolved.inputs[0]; | |
| 178 | 5 | if (m_storedSlot[input] == kNoValue) { | |
| 179 | ✗ | throw Error("{} is a reduction, but its {} isn't stored", resolved.path, | |
| 180 | ✗ | def.inputs[0].name); | |
| 181 | } | ||
| 182 | 5 | const std::string type = NodeType(def, resolved); | |
| 183 | 5 | const std::vector<uint32_t> slots = OutputSlots(resolved.outputs[0]); | |
| 184 | |||
| 185 | 5 | BeginKernel(fmt::format("Segment{}Partial", index), DispatchShape::ReduceGroups); | |
| 186 | 15 | Emit(FillTemplate(Embedded(chrono_batch_render_sensor_pipeline_reduce_partial, | |
| 187 | chrono_batch_render_sensor_pipeline_reduce_partial_size), | ||
| 188 | {{"type", type}, | ||
| 189 | 5 | {"context", Context(node)}, | |
| 190 | 5 | {"signal", SlangType(input)}, | |
| 191 | 10 | {"input", std::to_string(m_storedSlot[input])}, | |
| 192 | 10 | {"slot", std::to_string(slots.front())}})); | |
| 193 | 5 | EndKernel(); | |
| 194 | |||
| 195 | 5 | BeginKernel(fmt::format("Segment{}Finish", index), DispatchShape::Captures); | |
| 196 | 10 | Emit(fmt::format( | |
| 197 | " const CaptureStat total = {}.Finish({}, SumStatPartials(p, capture, {}));\n", type, | ||
| 198 | 5 | Context(node), slots.front())); | |
| 199 | 10 | for (const uint32_t slot : slots) { | |
| 200 | 10 | Emit(fmt::format(" WriteStat(p, capture, {}, total);\n", slot)); | |
| 201 | } | ||
| 202 | 5 | EndKernel(); | |
| 203 | 5 | } | |
| 204 | |||
| 205 | 6 | void AddGlobal(uint32_t index) { | |
| 206 | 6 | const uint32_t node = m_definition.segments.segments[index].nodes.front(); | |
| 207 | 6 | const ResolvedNode& resolved = Node(node); | |
| 208 | 6 | const NodeDef& def = Def(node); | |
| 209 | 6 | const uint32_t input = resolved.inputs[0]; | |
| 210 | 6 | if (input == kNoValue || m_storedSlot[input] == kNoValue) { | |
| 211 | ✗ | throw Error("{} is global, but its {} isn't stored", resolved.path, def.inputs[0].name); | |
| 212 | } | ||
| 213 | 6 | const uint32_t output = resolved.outputs[0]; | |
| 214 | 20 | for (uint32_t pass = 0; pass < def.passes; ++pass) { | |
| 215 | 14 | const uint32_t previous = pass == 0 ? input : resolved.scratch[pass - 1]; | |
| 216 | 14 | const bool last = pass + 1 == def.passes; | |
| 217 | 14 | BeginKernel(fmt::format("Segment{}Pass{}", index, pass), DispatchShape::Elements); | |
| 218 | 28 | Emit(fmt::format(" {0} result;\n" | |
| 219 | " {1}.Evaluate({2}, {3}, Neighborhood<{0}>(p, element, {4}),\n" | ||
| 220 | " Neighborhood<{0}>(p, element, {5}), result);\n", | ||
| 221 | 28 | SlangType(output), NodeType(def, resolved), Context(node), pass, | |
| 222 | 14 | m_storedSlot[input], m_storedSlot[previous])); | |
| 223 | 14 | Write(last ? output : resolved.scratch[pass], "result"); | |
| 224 | 14 | EndKernel(); | |
| 225 | } | ||
| 226 | 6 | } | |
| 227 | |||
| 228 | 82 | void AddFused(uint32_t index) { | |
| 229 | 82 | BeginKernel(fmt::format("Segment{}", index), DispatchShape::Elements); | |
| 230 | 82 | std::set<uint32_t> available; | |
| 231 | 267 | for (const uint32_t node : m_definition.segments.segments[index].nodes) { | |
| 232 | 185 | std::vector<std::string> arguments; | |
| 233 | 371 | for (uint32_t port = 0; port < Def(node).inputs.size(); ++port) { | |
| 234 | 372 | arguments.push_back(Argument(node, port, available)); | |
| 235 | } | ||
| 236 | 185 | if (Def(node).code.applies) { | |
| 237 | 8 | Apply(node, std::move(arguments)); | |
| 238 | } else { | ||
| 239 | 177 | Evaluate(node, std::move(arguments)); | |
| 240 | } | ||
| 241 | 370 | for (const uint32_t value : Node(node).outputs) { | |
| 242 | 185 | available.insert(value); | |
| 243 | 370 | Write(value, fmt::format("v{}", value)); | |
| 244 | } | ||
| 245 | 185 | } | |
| 246 | 82 | EndKernel(); | |
| 247 | 82 | } | |
| 248 | |||
| 249 | 186 | std::string Argument(uint32_t node, uint32_t port, std::set<uint32_t>& available) { | |
| 250 | 186 | const ResolvedNode& resolved = Node(node); | |
| 251 | 186 | const NodeDef& def = Def(node); | |
| 252 | 186 | const uint32_t value = resolved.inputs[port]; | |
| 253 | 186 | const SignalType taken = def.inputs[port].type; | |
| 254 | 186 | if (value == kNoValue) { | |
| 255 | 1 | return taken == SignalType::CaptureStat | |
| 256 | 1 | ? std::string("CaptureStat.Zero()") | |
| 257 | 1 | : fmt::format("{}.Make(Channels(0))", Signal(taken).slangType); | |
| 258 | } | ||
| 259 | 185 | if (def.kind == NodeKind::Stencil && port == 0) { | |
| 260 | 14 | if (m_storedSlot[value] == kNoValue) { | |
| 261 | ✗ | throw Error("{} is a stencil, but its {} isn't stored", resolved.path, | |
| 262 | ✗ | def.inputs[port].name); | |
| 263 | } | ||
| 264 | 28 | Emit(fmt::format( | |
| 265 | " const Neighborhood<{0}> n{1} = Neighborhood<{0}>(p, element, {2});\n", | ||
| 266 | 14 | SlangType(value), value, m_storedSlot[value])); | |
| 267 | 14 | return fmt::format("n{}", value); | |
| 268 | } | ||
| 269 | 171 | if (!available.contains(value)) { | |
| 270 | 69 | if (m_storedSlot[value] == kNoValue) { | |
| 271 | ✗ | throw Error("{}'s {} reads a value no earlier pass stored", resolved.path, | |
| 272 | ✗ | def.inputs[port].name); | |
| 273 | } | ||
| 274 | 138 | Emit( | |
| 275 | taken == SignalType::CaptureStat | ||
| 276 | 69 | ? fmt::format(" const CaptureStat v{} = ReadStat(p, element.capture, {});\n", | |
| 277 | 2 | value, m_storedSlot[value]) | |
| 278 | : fmt::format(" const {0} v{1} = ReadSignal<{0}>(p, element, {2});\n", | ||
| 279 | 67 | SlangType(value), value, m_storedSlot[value])); | |
| 280 | 69 | available.insert(value); | |
| 281 | } | ||
| 282 | 171 | return fmt::format("v{}", value); | |
| 283 | } | ||
| 284 | |||
| 285 | 8 | void Apply(uint32_t node, std::vector<std::string> arguments) { | |
| 286 | 8 | const NodeDef& def = Def(node); | |
| 287 | 8 | const std::string context = fmt::format("context{}", node); | |
| 288 | 16 | Emit(fmt::format(" const PointContext {} = {};\n", context, Context(node))); | |
| 289 | 16 | Emit(fmt::format(" {} node{};\n", def.code.type, node)); | |
| 290 | 8 | uint32_t word = 0; | |
| 291 | 22 | for (const ParameterDef& parameter : def.parameters) { | |
| 292 | 28 | Emit(fmt::format(" node{}.{} = {};\n", node, parameter.name, | |
| 293 | 14 | ReadParameter(parameter, context, word))); | |
| 294 | 14 | word += parameter.words; | |
| 295 | } | ||
| 296 | 8 | if (def.code.context != kNoContextArgument) { | |
| 297 | 8 | const auto at = std::min(size_t(def.code.context), arguments.size()); | |
| 298 | 8 | arguments.insert(arguments.begin() + ptrdiff_t(at), context); | |
| 299 | } | ||
| 300 | 8 | const uint32_t value = Node(node).outputs.front(); | |
| 301 | 16 | Emit(fmt::format(" const {} v{} = node{}.Apply({});\n", SlangType(value), value, node, | |
| 302 | 16 | Join(arguments))); | |
| 303 | 8 | } | |
| 304 | |||
| 305 | 177 | void Evaluate(uint32_t node, std::vector<std::string> arguments) { | |
| 306 | 177 | arguments.insert(arguments.begin(), Context(node)); | |
| 307 | 354 | for (const uint32_t value : Node(node).outputs) { | |
| 308 | 354 | Emit(fmt::format(" {} v{};\n", SlangType(value), value)); | |
| 309 | 354 | arguments.push_back(fmt::format("v{}", value)); | |
| 310 | } | ||
| 311 | 354 | Emit(fmt::format(" {}.Evaluate({});\n", NodeType(Def(node), Node(node)), | |
| 312 | 177 | Join(arguments))); | |
| 313 | 177 | } | |
| 314 | |||
| 315 | 82 | void AddUpdates(uint32_t index) { | |
| 316 | 267 | for (const uint32_t node : m_definition.segments.segments[index].nodes) { | |
| 317 | 185 | const ResolvedNode& resolved = Node(node); | |
| 318 | 185 | const NodeDef& def = Def(node); | |
| 319 | 185 | if (def.state.empty()) { | |
| 320 | 184 | continue; | |
| 321 | } | ||
| 322 | 2 | std::vector<std::string> arguments = {Context(node)}; | |
| 323 | 3 | for (uint32_t port = 0; port < def.inputs.size(); ++port) { | |
| 324 | 2 | if (def.inputs[port].type != SignalType::CaptureStat) { | |
| 325 | 1 | continue; | |
| 326 | } | ||
| 327 | 1 | const uint32_t value = resolved.inputs[port]; | |
| 328 | 2 | arguments.push_back(value == kNoValue ? std::string("CaptureStat.Zero()") | |
| 329 | : fmt::format("ReadStat(p, capture, {})", | ||
| 330 | 1 | m_storedSlot[value])); | |
| 331 | } | ||
| 332 | 1 | BeginKernel(fmt::format("Segment{}Update{}", index, node), DispatchShape::Captures); | |
| 333 | 2 | Emit(fmt::format(" {}.Update({});\n", NodeType(def, resolved), Join(arguments))); | |
| 334 | 1 | EndKernel(); | |
| 335 | 1 | } | |
| 336 | 83 | } | |
| 337 | |||
| 338 | 69 | void AddSourceCopies() { | |
| 339 | 69 | const auto stored = uint32_t(m_definition.segments.stored.size()); | |
| 340 | 69 | std::string copies; | |
| 341 | 155 | for (uint32_t output = 0; output < m_pipeline.outputs.size(); ++output) { | |
| 342 | 86 | const uint32_t value = m_pipeline.outputs[output].value; | |
| 343 | 86 | if (Def(m_pipeline.values[value].node).kind == NodeKind::Source) { | |
| 344 | ✗ | copies += fmt::format( | |
| 345 | " WriteSignal(p, element, {}, ReadSignal<{}>(p, element, {}));\n", | ||
| 346 | 6 | stored + output, SlangType(value), m_storedSlot[value]); | |
| 347 | } | ||
| 348 | } | ||
| 349 | 69 | if (!copies.empty()) { | |
| 350 | 3 | BeginKernel("CopySources", DispatchShape::Elements); | |
| 351 | 3 | Emit(copies); | |
| 352 | 3 | EndKernel(); | |
| 353 | } | ||
| 354 | 69 | } | |
| 355 | |||
| 356 | const SensorPipelineDefinition& m_definition; | ||
| 357 | const ResolvedSensorPipeline& m_pipeline; | ||
| 358 | const NodeRegistry& m_registry; | ||
| 359 | std::vector<uint32_t> m_storedSlot; | ||
| 360 | GeneratedSensorPipeline m_generated; | ||
| 361 | std::string m_body; | ||
| 362 | }; | ||
| 363 | |||
| 364 | } | ||
| 365 | |||
| 366 | 70 | std::vector<std::string> GeneratedSensorPipeline::Entries() const { | |
| 367 | 70 | std::vector<std::string> entries; | |
| 368 | 181 | for (const GeneratedKernel& kernel : kernels) { | |
| 369 | 111 | entries.push_back(kernel.entry); | |
| 370 | } | ||
| 371 | 70 | return entries; | |
| 372 | ✗ | } | |
| 373 | |||
| 374 | 69 | GeneratedSensorPipeline GenerateSensorPipeline(const SensorPipelineDefinition& definition, | |
| 375 | const NodeRegistry& registry) { | ||
| 376 | 69 | return Generator(definition, registry).Generate(); | |
| 377 | } | ||
| 378 | |||
| 379 | } | ||
| 380 |