algorithms/radix_sort.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include "chrono_batch_render/algorithms/radix_sort.h" | ||
| 3 | |||
| 4 | #include <algorithm> | ||
| 5 | #include <utility> | ||
| 6 | #include <vector> | ||
| 7 | |||
| 8 | #include "chrono_batch_render/algorithms/kernels.h" | ||
| 9 | #include "chrono_batch_render/algorithms/kernels/radix_sort.h" | ||
| 10 | #include "chrono_batch_render/device/device.h" | ||
| 11 | #include "chrono_batch_render/device/slang/dispatch.h" | ||
| 12 | #include "chrono_batch_render/util/error.h" | ||
| 13 | |||
| 14 | namespace chrono::batch_render::algorithms { | ||
| 15 | |||
| 16 | using device::BufferRange; | ||
| 17 | using device::CommandEncoder; | ||
| 18 | using device::Device; | ||
| 19 | using device::GroupCount; | ||
| 20 | using device::RequireRange; | ||
| 21 | |||
| 22 | namespace { | ||
| 23 | |||
| 24 | 949 | uint32_t RequireCapacity(uint32_t capacity) { | |
| 25 | 949 | if (capacity > RadixSort::kMaxItems) { | |
| 26 | 1 | throw Error("RadixSort of {} items; it sorts at most {}", capacity, RadixSort::kMaxItems); | |
| 27 | } | ||
| 28 | 948 | return capacity; | |
| 29 | } | ||
| 30 | |||
| 31 | 948 | uint32_t RequireKeyWords(uint32_t keyBytes) { | |
| 32 | 948 | if (keyBytes != 4 && keyBytes != 8) { | |
| 33 | 1 | throw Error("RadixSort of keys of {} bytes; it sorts keys of 4 or 8", keyBytes); | |
| 34 | } | ||
| 35 | 947 | return keyBytes / 4; | |
| 36 | } | ||
| 37 | |||
| 38 | } | ||
| 39 | |||
| 40 | 949 | RadixSort::RadixSort(Device& device, uint32_t capacity, uint32_t keyBytes) | |
| 41 | 949 | : m_capacity(RequireCapacity(capacity)), m_keyWords(RequireKeyWords(keyBytes)), | |
| 42 | 947 | m_program( | |
| 43 | 947 | device.CreateProgram({.modules = {"chrono_batch_render.algorithms.kernels.radix_sort"}, | |
| 44 | .library = &kAlgorithmsKernels, | ||
| 45 | .entries = {"CountDigits", "Scatter"}, | ||
| 46 | 947 | .label = "radix sort"})), | |
| 47 | 947 | m_countDigits(&m_program->Entry("CountDigits")), m_scatter(&m_program->Entry("Scatter")), | |
| 48 | 947 | m_scan(device, kDigitBins * GroupCount(capacity, kSortTileItems)) { | |
| 49 | 947 | const uint64_t items = std::max(capacity, 1u); | |
| 50 | 947 | const uint64_t tiles = GroupCount(items, kSortTileItems); | |
| 51 | 947 | m_keys = device.CreateBuffer({.bytes = items * m_keyWords * 4, .label = "radix sort keys"}); | |
| 52 | 947 | m_values = device.CreateBuffer({.bytes = items * 4, .label = "radix sort values"}); | |
| 53 | 947 | m_counts = device.CreateBuffer({.bytes = kDigitBins * tiles * 4, .label = "radix sort counts"}); | |
| 54 | 947 | m_unlimitedCount = | |
| 55 | 1894 | device.CreateBuffer({.label = "radix sort unlimited count"}, std::vector{~0u}); | |
| 56 | 6629 | } | |
| 57 | |||
| 58 | 1360 | void RadixSort::Sort(CommandEncoder& encoder, BufferRange keys, BufferRange values, uint32_t count, | |
| 59 | std::optional<uint32_t> keyBits) const { | ||
| 60 | 1360 | Record(encoder, keys, values, count, {m_unlimitedCount->Address()}, keyBits); | |
| 61 | 1355 | } | |
| 62 | |||
| 63 | 25 | void RadixSort::Sort(CommandEncoder& encoder, BufferRange keys, BufferRange values, | |
| 64 | BufferRange count, uint32_t capacity, std::optional<uint32_t> keyBits) const { | ||
| 65 | 25 | Record(encoder, keys, values, capacity, RequireRange<uint32_t>(*m_program, "Sort", count, 1), | |
| 66 | keyBits); | ||
| 67 | 23 | } | |
| 68 | |||
| 69 | 1385 | void RadixSort::Record(CommandEncoder& encoder, BufferRange keys, BufferRange values, | |
| 70 | uint32_t count, Ptr<uint32_t> liveCount, | ||
| 71 | std::optional<uint32_t> keyBits) const { | ||
| 72 | 1385 | if (count > m_capacity) { | |
| 73 | 2 | throw Error("Sort of {} items; the RadixSort was made for at most {}", count, m_capacity); | |
| 74 | } | ||
| 75 | 1383 | const uint32_t bits = keyBits.value_or(m_keyWords * 32); | |
| 76 | 1383 | if (bits == 0 || bits > m_keyWords * 32) { | |
| 77 | 2 | throw Error("Sort by {} bits of its keys; they have 1 to {}", bits, m_keyWords * 32); | |
| 78 | } | ||
| 79 | 1381 | const uint64_t keyBytes = uint64_t(count) * m_keyWords * 4; | |
| 80 | 1381 | const uint64_t valueBytes = uint64_t(count) * 4; | |
| 81 | 1381 | const uint32_t tiles = GroupCount(count, kSortTileItems); | |
| 82 | 1381 | SortParams params{.keys = RequireRange<uint32_t>(*m_program, "Sort", keys, keyBytes / 4), | |
| 83 | 1379 | .values = RequireRange<uint32_t>(*m_program, "Sort", values, count), | |
| 84 | 1378 | .keysOut = {m_keys->Address()}, | |
| 85 | 1378 | .valuesOut = {m_values->Address()}, | |
| 86 | 1378 | .counts = {m_counts->Address()}, | |
| 87 | .deviceCount = liveCount, | ||
| 88 | .count = count, | ||
| 89 | .tiles = tiles, | ||
| 90 | 1378 | .keyWords = m_keyWords}; | |
| 91 | |||
| 92 | 1378 | if (count == 0) { | |
| 93 | 2 | return; | |
| 94 | } | ||
| 95 | |||
| 96 | 1376 | const uint32_t passes = (bits + kDigitBits - 1) / kDigitBits; | |
| 97 | 1376 | if (passes % 2 != 0) { | |
| 98 | 197 | encoder.Barrier(); | |
| 99 | 197 | encoder.Copy(keys.Slice(0, keyBytes), m_keys->Slice(0, keyBytes)); | |
| 100 | 197 | encoder.Copy(values.Slice(0, valueBytes), m_values->Slice(0, valueBytes)); | |
| 101 | 197 | std::swap(params.keys, params.keysOut); | |
| 102 | 197 | std::swap(params.values, params.valuesOut); | |
| 103 | } | ||
| 104 | |||
| 105 | 6669 | for (uint32_t pass = 0; pass < passes; ++pass) { | |
| 106 | 5293 | params.shift = pass * kDigitBits; | |
| 107 | 5293 | params.mask = (1u << std::min(kDigitBits, bits - params.shift)) - 1; | |
| 108 | 5293 | encoder.Barrier(); | |
| 109 | 5293 | encoder.Dispatch(*m_countDigits, params, {.x = tiles}); | |
| 110 | 5293 | m_scan.ExclusiveSum(encoder, *m_counts, *m_counts, kDigitBins * tiles); | |
| 111 | 5293 | encoder.Barrier(); | |
| 112 | 5293 | encoder.Dispatch(*m_scatter, params, {.x = tiles}); | |
| 113 | 5293 | std::swap(params.keys, params.keysOut); | |
| 114 | 5293 | std::swap(params.values, params.valuesOut); | |
| 115 | } | ||
| 116 | } | ||
| 117 | |||
| 118 | } | ||
| 119 |