device/program.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include "chrono_batch_render/device/program.h" | ||
| 3 | |||
| 4 | #include <algorithm> | ||
| 5 | #include <chrono> | ||
| 6 | #include <condition_variable> | ||
| 7 | #include <exception> | ||
| 8 | #include <functional> | ||
| 9 | #include <future> | ||
| 10 | #include <map> | ||
| 11 | #include <mutex> | ||
| 12 | #include <set> | ||
| 13 | #include <thread> | ||
| 14 | #include <tuple> | ||
| 15 | #include <utility> | ||
| 16 | #include <vector> | ||
| 17 | |||
| 18 | #include <fmt/format.h> | ||
| 19 | |||
| 20 | #include "chrono_batch_render/device/device.h" | ||
| 21 | #include "chrono_batch_render/device/slang/compiler.h" | ||
| 22 | #include "chrono_batch_render/device/slang/debug.h" | ||
| 23 | #include "chrono_batch_render/util/disk_cache.h" | ||
| 24 | #include "chrono_batch_render/util/error.h" | ||
| 25 | #include "chrono_batch_render/util/log.h" | ||
| 26 | #include "chrono_batch_render/util/thread_pool.h" | ||
| 27 | |||
| 28 | namespace chrono::batch_render::device { | ||
| 29 | namespace { | ||
| 30 | |||
| 31 | struct Registry { | ||
| 32 | std::mutex mutex; | ||
| 33 | std::shared_ptr<const SourceFiles> files = std::make_shared<SourceFiles>(); | ||
| 34 | }; | ||
| 35 | |||
| 36 | 9209 | Registry& Files() { | |
| 37 | 9209 | static Registry registry; | |
| 38 | 9209 | return registry; | |
| 39 | } | ||
| 40 | |||
| 41 | 9203 | std::shared_ptr<const SourceFiles> RegisteredFiles() { | |
| 42 | 9203 | Registry& registry = Files(); | |
| 43 | 9203 | const std::lock_guard lock(registry.mutex); | |
| 44 | 18406 | return registry.files; | |
| 45 | 9203 | } | |
| 46 | |||
| 47 | 2454 | ThreadPool& Pool() { | |
| 48 | 2454 | static ThreadPool& pool = | |
| 49 | 2465 | *new ThreadPool(std::clamp(std::thread::hardware_concurrency() / 2, 1u, 4u)); | |
| 50 | 2454 | return pool; | |
| 51 | } | ||
| 52 | |||
| 53 | struct Builds { | ||
| 54 | std::mutex mutex; | ||
| 55 | std::condition_variable done; | ||
| 56 | uint64_t pending = 0; | ||
| 57 | uint64_t queued = 0; | ||
| 58 | uint64_t finished = 0; | ||
| 59 | uint64_t cached = 0; | ||
| 60 | std::multiset<std::string> building; | ||
| 61 | }; | ||
| 62 | |||
| 63 | 3712 | Builds& PendingBuilds() { | |
| 64 | 3712 | static Builds& builds = *new Builds; | |
| 65 | 3712 | return builds; | |
| 66 | } | ||
| 67 | |||
| 68 | 1119 | void FinishBuild(const std::string& name = {}, bool cached = false) { | |
| 69 | 1119 | Builds& builds = PendingBuilds(); | |
| 70 | 1119 | { | |
| 71 | 1119 | const std::lock_guard lock(builds.mutex); | |
| 72 | 1119 | --builds.pending; | |
| 73 | 1119 | ++builds.finished; | |
| 74 | 1119 | builds.cached += cached ? 1 : 0; | |
| 75 | 1119 | if (const auto found = builds.building.find(name); found != builds.building.end()) { | |
| 76 | 1119 | builds.building.erase(found); | |
| 77 | } | ||
| 78 | 1119 | } | |
| 79 | 1119 | builds.done.notify_all(); | |
| 80 | 1119 | } | |
| 81 | |||
| 82 | 1117 | std::string ResourceUsage(const KernelStats& stats) { | |
| 83 | 1117 | std::string usage; | |
| 84 | 1117 | if (stats.registers) { | |
| 85 | ✗ | usage += fmt::format(", {} registers", *stats.registers); | |
| 86 | } | ||
| 87 | 1117 | if (stats.spillBytes) { | |
| 88 | ✗ | usage += fmt::format(", {} bytes of scratch memory", *stats.spillBytes); | |
| 89 | } | ||
| 90 | 1117 | return usage; | |
| 91 | ✗ | } | |
| 92 | |||
| 93 | } | ||
| 94 | |||
| 95 | 6 | void RegisterFile(std::string path, std::string source) { | |
| 96 | 6 | Registry& registry = Files(); | |
| 97 | 6 | const std::lock_guard lock(registry.mutex); | |
| 98 | 6 | auto files = std::make_shared<SourceFiles>(*registry.files); | |
| 99 | 6 | (*files)[std::move(path)] = std::move(source); | |
| 100 | 6 | registry.files = std::move(files); | |
| 101 | 6 | } | |
| 102 | |||
| 103 | namespace { | ||
| 104 | |||
| 105 | 1119 | struct BuildJob { | |
| 106 | enum class Stage { Compiling, Compiled, Loading, Done, Abandoned }; | ||
| 107 | |||
| 108 | std::mutex mutex; | ||
| 109 | std::condition_variable loaded; | ||
| 110 | Stage stage = Stage::Compiling; | ||
| 111 | bool wanted = false; | ||
| 112 | std::optional<KernelCode> code; | ||
| 113 | }; | ||
| 114 | |||
| 115 | } | ||
| 116 | |||
| 117 | struct Kernel::State { | ||
| 118 | std::shared_ptr<BuildJob> job = std::make_shared<BuildJob>(); | ||
| 119 | std::promise<KernelShape> promise; | ||
| 120 | std::shared_future<KernelShape> compiled = promise.get_future().share(); | ||
| 121 | KernelStats stats; | ||
| 122 | |||
| 123 | 232830 | const KernelShape& Get(const Kernel& kernel) const { | |
| 124 | 232830 | if (compiled.wait_for(std::chrono::seconds(0)) != std::future_status::ready) { | |
| 125 | 216 | if (const std::function<void()> queued = Pool().Take(&kernel)) { | |
| 126 | 132 | queued(); | |
| 127 | 216 | } | |
| 128 | } | ||
| 129 | 232830 | return compiled.get(); | |
| 130 | } | ||
| 131 | }; | ||
| 132 | |||
| 133 | 1119 | Kernel::Kernel() : m_state(std::make_unique<State>()) {} | |
| 134 | |||
| 135 | 1119 | Kernel::~Kernel() = default; | |
| 136 | |||
| 137 | 6811 | bool Kernel::Ready() const { | |
| 138 | 6811 | return m_state->compiled.wait_for(std::chrono::seconds(0)) == std::future_status::ready; | |
| 139 | } | ||
| 140 | |||
| 141 | 247 | void Kernel::Wait() const { | |
| 142 | 247 | m_state->Get(*this); | |
| 143 | 245 | } | |
| 144 | |||
| 145 | ✗ | Size3 Kernel::GroupSize() const { | |
| 146 | ✗ | return m_state->Get(*this).groupSize; | |
| 147 | } | ||
| 148 | |||
| 149 | 232582 | uint32_t Kernel::ParamBytes() const { | |
| 150 | 232582 | return m_state->Get(*this).paramBytes; | |
| 151 | } | ||
| 152 | |||
| 153 | ✗ | bool Kernel::Foldable() const { | |
| 154 | ✗ | return m_state->Get(*this).foldable; | |
| 155 | } | ||
| 156 | |||
| 157 | 1 | const KernelStats& Kernel::Stats() const { | |
| 158 | 1 | m_state->Get(*this); | |
| 159 | 1 | return m_state->stats; | |
| 160 | } | ||
| 161 | |||
| 162 | 17 | struct Device::ProgramCache { | |
| 163 | using Key = std::tuple<const KernelLibrary*, std::vector<std::string>, Traversal, | ||
| 164 | const SourceFiles*, SourceFiles>; | ||
| 165 | |||
| 166 | std::mutex mutex; | ||
| 167 | std::map<Key, Ref<Program>> programs; | ||
| 168 | }; | ||
| 169 | |||
| 170 | struct Program::State { | ||
| 171 | const KernelLibrary* library = nullptr; | ||
| 172 | std::shared_ptr<const SourceFiles> sources; | ||
| 173 | std::vector<std::string> modules; | ||
| 174 | Traversal traversal = Traversal::Software; | ||
| 175 | std::mutex mutex; | ||
| 176 | std::map<std::string, std::unique_ptr<Kernel>, std::less<>> kernels; | ||
| 177 | }; | ||
| 178 | |||
| 179 | 310 | Program::Program(Device& device, const ProgramDesc& desc, | |
| 180 | 310 | std::shared_ptr<const SourceFiles> sources) | |
| 181 | 620 | : Resource(device, desc.label), m_state(std::make_unique<State>()) { | |
| 182 | 310 | m_state->library = desc.library; | |
| 183 | 310 | m_state->sources = std::move(sources); | |
| 184 | 650 | for (const std::string& name : desc.modules) { | |
| 185 | 341 | if (!HasModule(m_state->library, *m_state->sources, name)) { | |
| 186 | 1 | throw Error("program \"{}\" names module {}, which is neither registered nor in " | |
| 187 | "Chrono Batch Render's kernel library", | ||
| 188 | 1 | desc.label, name); | |
| 189 | } | ||
| 190 | } | ||
| 191 | 309 | m_state->modules = desc.modules; | |
| 192 | |||
| 193 | 309 | m_state->traversal = desc.traversal == Traversal::Auto ? device.GetTraversal() : desc.traversal; | |
| 194 | 309 | if (m_state->traversal == Traversal::Hardware && device.GetTraversal() != Traversal::Hardware) { | |
| 195 | ✗ | throw UnsupportedError("program \"{}\" traverses rays in hardware, which {} doesn't", | |
| 196 | ✗ | desc.label, device.Adapter().alias); | |
| 197 | } | ||
| 198 | 311 | } | |
| 199 | |||
| 200 | 618 | Program::~Program() { | |
| 201 | 1428 | for (const auto& [name, kernel] : m_state->kernels) { | |
| 202 | 1119 | if (Pool().Cancel(kernel.get())) { | |
| 203 | ✗ | FinishBuild(); | |
| 204 | ✗ | continue; | |
| 205 | } | ||
| 206 | 1119 | BuildJob& job = *kernel->m_state->job; | |
| 207 | 1119 | std::unique_lock lock(job.mutex); | |
| 208 | 1119 | if (job.stage == BuildJob::Stage::Compiling) { | |
| 209 | ✗ | job.stage = BuildJob::Stage::Abandoned; | |
| 210 | ✗ | continue; | |
| 211 | } | ||
| 212 | 1119 | if (job.stage == BuildJob::Stage::Compiled) { | |
| 213 | ✗ | continue; | |
| 214 | } | ||
| 215 | 1120 | job.loaded.wait(lock, [&] { return job.stage == BuildJob::Stage::Done; }); | |
| 216 | 1119 | } | |
| 217 | 618 | } | |
| 218 | |||
| 219 | 70076 | Kernel& Program::Entry(std::string_view name) { | |
| 220 | 70076 | return Request(name, true); | |
| 221 | } | ||
| 222 | |||
| 223 | 107910 | Kernel& Program::Request(std::string_view name, bool wanted) { | |
| 224 | 107910 | const std::lock_guard lock(m_state->mutex); | |
| 225 | 107910 | if (const auto found = m_state->kernels.find(name); found != m_state->kernels.end()) { | |
| 226 | 106791 | Kernel& kernel = *found->second; | |
| 227 | 106791 | if (wanted) { | |
| 228 | 106791 | BuildJob& job = *kernel.m_state->job; | |
| 229 | 106791 | const std::lock_guard staging(job.mutex); | |
| 230 | 106791 | job.wanted = true; | |
| 231 | 106791 | if (job.stage == BuildJob::Stage::Compiled) { | |
| 232 | ✗ | job.stage = BuildJob::Stage::Loading; | |
| 233 | ✗ | QueueLoad(kernel); | |
| 234 | } | ||
| 235 | 106791 | } | |
| 236 | 106791 | return kernel; | |
| 237 | } | ||
| 238 | |||
| 239 | 1119 | std::unique_ptr<Kernel> made = GetDevice().NewKernel(); | |
| 240 | 1119 | made->m_name = name; | |
| 241 | 1119 | made->m_program = this; | |
| 242 | 1119 | made->m_state->job->wanted = wanted; | |
| 243 | 1119 | Kernel& kernel = *m_state->kernels.emplace(std::string(name), std::move(made)).first->second; | |
| 244 | 1119 | { | |
| 245 | 1119 | Builds& builds = PendingBuilds(); | |
| 246 | 1119 | const std::lock_guard pending(builds.mutex); | |
| 247 | 1119 | ++builds.pending; | |
| 248 | 1119 | ++builds.queued; | |
| 249 | 1119 | } | |
| 250 | 1119 | const AdapterInfo& adapter = GetDevice().Adapter(); | |
| 251 | 1119 | const std::string label = Label() + ": " + kernel.Name(); | |
| 252 | 2238 | Pool().Add(&kernel, | |
| 253 | 3357 | [this, &kernel, label, job = kernel.m_state->job, backend = adapter.backend, | |
| 254 | 1119 | traversal = m_state->traversal, sampling = GetDevice().GetSampling(), | |
| 255 | 1119 | debugInfo = GetDevice().m_kernelDebugInfo, library = m_state->library, | |
| 256 | 2238 | sources = m_state->sources, modules = m_state->modules, entry = kernel.Name(), | |
| 257 | 1119 | program = Label(), directory = GetDevice().KernelCacheDirectory(), | |
| 258 | 1119 | narrow = adapter.features.waveSizeMin < 8] { | |
| 259 | 1119 | { | |
| 260 | 1119 | Builds& builds = PendingBuilds(); | |
| 261 | 1119 | const std::lock_guard building(builds.mutex); | |
| 262 | 1119 | builds.building.insert(label); | |
| 263 | ✗ | } | |
| 264 | 1119 | const auto start = std::chrono::steady_clock::now(); | |
| 265 | 3357 | const DiskCache slang(directory / "slang"); | |
| 266 | 3357 | const DiskCache native(directory / BackendName(backend)); | |
| 267 | 1119 | std::optional<KernelCode> code; | |
| 268 | 1119 | std::exception_ptr failure; | |
| 269 | 1119 | try { | |
| 270 | 2236 | code = Compile(backend, traversal, sampling, debugInfo, library, *sources, | |
| 271 | 1119 | modules, entry, program, slang, narrow); | |
| 272 | 2 | } catch (...) { | |
| 273 | 2 | failure = std::current_exception(); | |
| 274 | 2 | } | |
| 275 | 1119 | { | |
| 276 | 1119 | const std::lock_guard staging(job->mutex); | |
| 277 | 1119 | if (job->stage == BuildJob::Stage::Abandoned) { | |
| 278 | ✗ | FinishBuild(label, false); | |
| 279 | return; | ||
| 280 | } | ||
| 281 | 1119 | if (!failure && !job->wanted) { | |
| 282 | ✗ | const bool cached = code->cached; | |
| 283 | ✗ | job->code = std::move(code); | |
| 284 | ✗ | job->stage = BuildJob::Stage::Compiled; | |
| 285 | ✗ | FinishBuild(label, cached); | |
| 286 | return; | ||
| 287 | } | ||
| 288 | 1119 | job->stage = BuildJob::Stage::Loading; | |
| 289 | ✗ | } | |
| 290 | 1119 | if (failure) { | |
| 291 | 6 | kernel.m_state->promise.set_exception(failure); | |
| 292 | } else { | ||
| 293 | 1117 | Load(kernel, *code, native, start); | |
| 294 | } | ||
| 295 | 1119 | Finish(kernel, label); | |
| 296 | 3357 | }); | |
| 297 | 1119 | return kernel; | |
| 298 | 109029 | } | |
| 299 | |||
| 300 | ✗ | void Program::QueueLoad(Kernel& kernel) { | |
| 301 | ✗ | { | |
| 302 | ✗ | Builds& builds = PendingBuilds(); | |
| 303 | ✗ | const std::lock_guard pending(builds.mutex); | |
| 304 | ✗ | ++builds.pending; | |
| 305 | ✗ | ++builds.queued; | |
| 306 | ✗ | } | |
| 307 | ✗ | Pool().Add(&kernel, [this, &kernel, label = Label() + ": " + kernel.Name(), | |
| 308 | ✗ | directory = GetDevice().KernelCacheDirectory() / | |
| 309 | ✗ | BackendName(GetDevice().Adapter().backend)] { | |
| 310 | ✗ | { | |
| 311 | ✗ | Builds& builds = PendingBuilds(); | |
| 312 | ✗ | const std::lock_guard building(builds.mutex); | |
| 313 | ✗ | builds.building.insert(label); | |
| 314 | ✗ | } | |
| 315 | ✗ | BuildJob& job = *kernel.m_state->job; | |
| 316 | ✗ | Load(kernel, *job.code, DiskCache(directory), std::chrono::steady_clock::now()); | |
| 317 | ✗ | job.code.reset(); | |
| 318 | ✗ | Finish(kernel, label); | |
| 319 | ✗ | }); | |
| 320 | ✗ | } | |
| 321 | |||
| 322 | 1119 | void Program::Finish(Kernel& kernel, const std::string& label) { | |
| 323 | 1119 | const bool cached = kernel.Ready() && kernel.m_state->stats.cached; | |
| 324 | 1119 | BuildJob& job = *kernel.m_state->job; | |
| 325 | 1119 | { | |
| 326 | 1119 | const std::lock_guard staging(job.mutex); | |
| 327 | 1119 | job.stage = BuildJob::Stage::Done; | |
| 328 | 1119 | } | |
| 329 | 1119 | job.loaded.notify_all(); | |
| 330 | 1119 | FinishBuild(label, cached); | |
| 331 | 1119 | } | |
| 332 | |||
| 333 | 1117 | void Program::Load(Kernel& kernel, const KernelCode& code, const DiskCache& cache, | |
| 334 | std::chrono::steady_clock::time_point start) { | ||
| 335 | 1117 | try { | |
| 336 | 1117 | const AdapterInfo& adapter = GetDevice().Adapter(); | |
| 337 | 1117 | if (code.shape.paramBytes > adapter.limits.maxParamBytes) { | |
| 338 | ✗ | throw Error("kernel {} of program \"{}\" takes {} bytes of parameters; {} takes at " | |
| 339 | "most {}", | ||
| 340 | ✗ | kernel.Name(), Label(), code.shape.paramBytes, adapter.alias, | |
| 341 | ✗ | adapter.limits.maxParamBytes); | |
| 342 | } | ||
| 343 | |||
| 344 | 1117 | GetDevice().Debug().AddKernel(kernel.Name(), code.strings); | |
| 345 | 1117 | const auto loading = std::chrono::steady_clock::now(); | |
| 346 | 1117 | KernelStats stats = kernel.Load(code, cache, GetDevice().Globals()); | |
| 347 | 1117 | const auto loaded = std::chrono::steady_clock::now(); | |
| 348 | 1117 | stats.compileMs = std::chrono::duration<double, std::milli>(loaded - start).count(); | |
| 349 | 1117 | stats.cached = stats.cached && code.cached; | |
| 350 | 1117 | Log(LogLevel::Debug, | |
| 351 | "kernel {} of program \"{}\" for {} {} in {:.1f} ms, {:.1f} ms of it in {}{}", | ||
| 352 | 1117 | kernel.Name(), Label(), adapter.alias, | |
| 353 | 1119 | stats.cached ? "loaded from the cache" : "compiled", stats.compileMs, | |
| 354 | 1117 | std::chrono::duration<double, std::milli>(loaded - loading).count(), | |
| 355 | 2234 | BackendName(adapter.backend), ResourceUsage(stats)); | |
| 356 | |||
| 357 | 1117 | kernel.m_state->stats = std::move(stats); | |
| 358 | 1117 | kernel.m_state->promise.set_value(code.shape); | |
| 359 | 1117 | } catch (...) { | |
| 360 | ✗ | kernel.m_state->promise.set_exception(std::current_exception()); | |
| 361 | ✗ | } | |
| 362 | 1117 | } | |
| 363 | |||
| 364 | 3 | BuildProgress CurrentBuilds() { | |
| 365 | 3 | Builds& builds = PendingBuilds(); | |
| 366 | 3 | const std::lock_guard lock(builds.mutex); | |
| 367 | 3 | return {builds.queued, builds.finished, builds.cached, | |
| 368 | 3 | std::vector<std::string>(builds.building.begin(), builds.building.end())}; | |
| 369 | 3 | } | |
| 370 | |||
| 371 | 352 | void WaitForBuilds() { | |
| 372 | 352 | Builds& builds = PendingBuilds(); | |
| 373 | 352 | std::unique_lock lock(builds.mutex); | |
| 374 | 1088 | builds.done.wait(lock, [&] { return builds.pending == 0; }); | |
| 375 | 352 | } | |
| 376 | |||
| 377 | 9203 | Ref<Program> Device::CreateProgram(const ProgramDesc& desc) { | |
| 378 | 9203 | static std::mutex creating; | |
| 379 | 9203 | std::shared_ptr<ProgramCache> cache; | |
| 380 | 9203 | { | |
| 381 | 9203 | const std::lock_guard lock(creating); | |
| 382 | 9203 | if (!m_programs) { | |
| 383 | 17 | m_programs = std::make_shared<ProgramCache>(); | |
| 384 | } | ||
| 385 | 9203 | cache = m_programs; | |
| 386 | ✗ | } | |
| 387 | 9203 | const std::shared_ptr<const SourceFiles> registered = RegisteredFiles(); | |
| 388 | 9203 | const Traversal traversal = desc.traversal == Traversal::Auto ? m_traversal : desc.traversal; | |
| 389 | 9203 | const ProgramCache::Key key{desc.library, desc.modules, traversal, registered.get(), | |
| 390 | 18407 | SourceFiles(desc.sources.begin(), desc.sources.end())}; | |
| 391 | 9203 | Ref<Program> program; | |
| 392 | 9203 | { | |
| 393 | 9203 | const std::lock_guard lock(cache->mutex); | |
| 394 | 9203 | if (const auto found = cache->programs.find(key); found != cache->programs.end()) { | |
| 395 | 8893 | program = found->second; | |
| 396 | } else { | ||
| 397 | 310 | std::shared_ptr<const SourceFiles> sources = registered; | |
| 398 | 310 | if (!desc.sources.empty()) { | |
| 399 | 82 | auto merged = std::make_shared<SourceFiles>(*registered); | |
| 400 | 245 | for (const auto& [path, source] : desc.sources) { | |
| 401 | 326 | (*merged)[path] = source; | |
| 402 | } | ||
| 403 | 82 | sources = std::move(merged); | |
| 404 | 82 | } | |
| 405 | 312 | program = Ref<Program>(new Program(*this, desc, std::move(sources))); | |
| 406 | 309 | Adopt(*program); | |
| 407 | 309 | cache->programs.emplace(key, program); | |
| 408 | 310 | } | |
| 409 | 1 | } | |
| 410 | 47036 | for (const std::string& entry : desc.entries) { | |
| 411 | 37834 | program->Request(entry, desc.loadEntries); | |
| 412 | } | ||
| 413 | 9202 | return program; | |
| 414 | 27609 | } | |
| 415 | |||
| 416 | } | ||
| 417 |