gpu/kernel.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "terraingen/gpu/kernel.h" | ||
| 2 | |||
| 3 | #include "terraingen/gpu/cuda.h" | ||
| 4 | #include "terraingen/gpu/metal.h" | ||
| 5 | #include "terraingen/gpu/vulkan.h" | ||
| 6 | |||
| 7 | #include <algorithm> | ||
| 8 | #include <array> | ||
| 9 | #include <cmath> | ||
| 10 | #include <cstdint> | ||
| 11 | #include <vector> | ||
| 12 | |||
| 13 | namespace terraingen { | ||
| 14 | |||
| 15 | namespace { | ||
| 16 | |||
| 17 | /// Most cells one dispatch covers, keeping each buffer under the smallest storage buffer range | ||
| 18 | /// Vulkan allows (128 MiB). | ||
| 19 | constexpr std::size_t kMaxStripCells = std::size_t{1} << 24; | ||
| 20 | |||
| 21 | /// Cells per workgroup side in kernels. | ||
| 22 | constexpr std::size_t kGroupSide = 16; | ||
| 23 | |||
| 24 | /// Most workgroups Vulkan guarantees in one dimension of a dispatch; CUDA and Metal allow more. | ||
| 25 | constexpr std::size_t kMaxGroups = 65535; | ||
| 26 | |||
| 27 | /// The push constants of kernels (`terraingen/gpu/kernel.slang`). | ||
| 28 | struct KernelParams { | ||
| 29 | std::uint32_t width = 0; | ||
| 30 | std::uint32_t rows = 0; | ||
| 31 | std::uint32_t noData = 0; ///< 0 for none, 1 for NaN, 2 for `noDataValue`. | ||
| 32 | float noDataValue = 0.0f; | ||
| 33 | std::uint32_t firstRow = 0; ///< The strip's first row in the grid. | ||
| 34 | }; | ||
| 35 | |||
| 36 | } // namespace | ||
| 37 | |||
| 38 | 17249 | bool RunKernel(Backend backend, const KernelProgram& program, std::span<const std::byte> kernel, | |
| 39 | Raster<float>& grid, const Vec2d& origin, std::span<const KernelBuffer> extra) { | ||
| 40 | 17249 | const GeoTransform& transform = grid.Transform(); | |
| 41 | 17249 | const std::size_t width = grid.Width(); | |
| 42 | 34496 | if (transform.rowRotation != 0.0 || transform.colRotation != 0.0 || width == 0 || | |
| 43 | 51745 | grid.Height() == 0 || (width + kGroupSide - 1) / kGroupSide > kMaxGroups) { | |
| 44 | 2 | return false; | |
| 45 | } | ||
| 46 | 17247 | VulkanCompute* vulkan = nullptr; | |
| 47 | 17247 | CudaCompute* cuda = nullptr; | |
| 48 | 17247 | MetalCompute* metal = nullptr; | |
| 49 | 17247 | switch (backend) { | |
| 50 | 31 | case Backend::Vulkan: | |
| 51 | 31 | vulkan = program.spirv.empty() ? nullptr : VulkanCompute::Get(); | |
| 52 | 31 | break; | |
| 53 | 1 | case Backend::Cuda: | |
| 54 | 1 | cuda = program.ptx.empty() ? nullptr : CudaCompute::Get(); | |
| 55 | 1 | break; | |
| 56 | 1 | case Backend::Metal: | |
| 57 | 1 | metal = program.metal.empty() ? nullptr : MetalCompute::Get(); | |
| 58 | 1 | break; | |
| 59 | 17214 | default: | |
| 60 | 17214 | break; | |
| 61 | } | ||
| 62 | 17247 | if (vulkan == nullptr && cuda == nullptr && metal == nullptr) { | |
| 63 | 17217 | return false; | |
| 64 | } | ||
| 65 | 60 | std::vector<float> xs(width); | |
| 66 | 7302 | for (std::size_t col = 0; col < width; ++col) { | |
| 67 | 7272 | xs[col] = static_cast<float>(grid.CellCenter(col, 0).x - origin.x); | |
| 68 | } | ||
| 69 | 60 | std::vector<float> ys(grid.Height()); | |
| 70 | 8381 | for (std::size_t row = 0; row < ys.size(); ++row) { | |
| 71 | 8351 | ys[row] = static_cast<float>(grid.CellCenter(0, row).y - origin.y); | |
| 72 | } | ||
| 73 | 60 | std::vector<std::byte> block(kernel.begin(), kernel.end()); | |
| 74 | ✗ | std::vector<std::vector<std::byte>> extraCopies; | |
| 75 | 110 | for (const KernelBuffer& buffer : extra) { | |
| 76 | 80 | extraCopies.emplace_back(buffer.bytes.begin(), buffer.bytes.end()); | |
| 77 | } | ||
| 78 | 30 | KernelParams push; | |
| 79 | 30 | push.width = static_cast<std::uint32_t>(width); | |
| 80 | 30 | if (const auto& noData = grid.NoData()) { | |
| 81 | 2 | push.noData = std::isnan(*noData) ? 1u : 2u; | |
| 82 | 2 | push.noDataValue = *noData; | |
| 83 | } | ||
| 84 | const std::size_t stripRows = | ||
| 85 | 30 | std::clamp<std::size_t>(kMaxStripCells / width, 1, kMaxGroups * kGroupSide); | |
| 86 | 30 | const std::array<std::uint32_t, 3> groupSize{kGroupSide, kGroupSide, 1u}; | |
| 87 | 30 | std::span<float> values(grid.Values()); | |
| 88 | 61 | for (std::size_t row0 = 0; row0 < ys.size(); row0 += stripRows) { | |
| 89 | 31 | const std::size_t rows = std::min(stripRows, ys.size() - row0); | |
| 90 | 31 | push.rows = static_cast<std::uint32_t>(rows); | |
| 91 | 31 | push.firstRow = static_cast<std::uint32_t>(row0); | |
| 92 | std::vector<std::span<std::byte>> buffers{ | ||
| 93 | 31 | std::as_writable_bytes(values.subspan(row0 * width, rows * width)), | |
| 94 | 31 | std::as_writable_bytes(std::span(xs)), | |
| 95 | 62 | std::as_writable_bytes(std::span(ys).subspan(row0, rows)), std::span(block)}; | |
| 96 | 111 | for (std::vector<std::byte>& copy : extraCopies) { | |
| 97 | 80 | buffers.push_back(std::span(copy)); | |
| 98 | } | ||
| 99 | 31 | const std::span<const std::byte> pushBytes = std::as_bytes(std::span(&push, 1)); | |
| 100 | 31 | const std::array<std::uint32_t, 3> groups{ | |
| 101 | 31 | static_cast<std::uint32_t>((width + kGroupSide - 1) / kGroupSide), | |
| 102 | 31 | static_cast<std::uint32_t>((rows + kGroupSide - 1) / kGroupSide), 1u}; | |
| 103 | 31 | if (vulkan != nullptr) { | |
| 104 | 31 | vulkan->Run({program.spirv.data(), program.spirv.size()}, buffers, pushBytes, groups); | |
| 105 | ✗ | } else if (cuda != nullptr) { | |
| 106 | ✗ | std::vector<CudaBuffer> cudaBuffers{{buffers[0], rows * width}, | |
| 107 | ✗ | {buffers[1], width}, | |
| 108 | ✗ | {buffers[2], rows}, | |
| 109 | ✗ | {buffers[3], 1}}; | |
| 110 | ✗ | for (std::size_t n = 0; n < extra.size(); ++n) { | |
| 111 | ✗ | cudaBuffers.push_back({buffers[4 + n], extra[n].count}); | |
| 112 | } | ||
| 113 | ✗ | cuda->Run(program.ptx, cudaBuffers, pushBytes, groups, groupSize); | |
| 114 | ✗ | } else { | |
| 115 | ✗ | metal->Run(program.metal, buffers, pushBytes, groups, groupSize); | |
| 116 | } | ||
| 117 | 31 | } | |
| 118 | 30 | return true; | |
| 119 | } | ||
| 120 | |||
| 121 | } // namespace terraingen | ||
| 122 |