chrono_batch_render/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
63 / 63
100.0%
Branches:
44 / 50
88.0%
Functions:
6 / 6
100.0%

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
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 948 times.
949 if (capacity > RadixSort::kMaxItems) {
26
1/1
✓ Branch 2 taken 1 time.
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
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 947 times.
948 if (keyBytes != 4 && keyBytes != 8) {
33
1/1
✓ Branch 2 taken 1 time.
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
2/2
✓ Branch 1 taken 947 times.
✓ Branch 5 taken 947 times.
947 m_program(
43
1/1
✓ Branch 2 taken 947 times.
947 device.CreateProgram({.modules = {"chrono_batch_render.algorithms.kernels.radix_sort"},
44 .library = &kAlgorithmsKernels,
45 .entries = {"CountDigits", "Scatter"},
46
1/1
✓ Branch 1 taken 947 times.
947 .label = "radix sort"})),
47
1/1
✓ Branch 1 taken 947 times.
947 m_countDigits(&m_program->Entry("CountDigits")), m_scatter(&m_program->Entry("Scatter")),
48
2/3
✓ Branch 1 taken 947 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 947 times.
947 m_scan(device, kDigitBins * GroupCount(capacity, kSortTileItems)) {
49
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 947 times.
947 const uint64_t items = std::max(capacity, 1u);
50
1/1
✓ Branch 1 taken 947 times.
947 const uint64_t tiles = GroupCount(items, kSortTileItems);
51
2/3
✓ Branch 1 taken 947 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 947 times.
947 m_keys = device.CreateBuffer({.bytes = items * m_keyWords * 4, .label = "radix sort keys"});
52
2/3
✓ Branch 1 taken 947 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 947 times.
947 m_values = device.CreateBuffer({.bytes = items * 4, .label = "radix sort values"});
53
2/3
✓ Branch 1 taken 947 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 947 times.
947 m_counts = device.CreateBuffer({.bytes = kDigitBins * tiles * 4, .label = "radix sort counts"});
54 947 m_unlimitedCount =
55
3/4
✓ Branch 1 taken 947 times.
✓ Branch 4 taken 947 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 947 times.
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
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1383 times.
1385 if (count > m_capacity) {
73
1/1
✓ Branch 2 taken 2 times.
2 throw Error("Sort of {} items; the RadixSort was made for at most {}", count, m_capacity);
74 }
75
2/2
✓ Branch 0 taken 1368 times.
✓ Branch 1 taken 15 times.
1383 const uint32_t bits = keyBits.value_or(m_keyWords * 32);
76
4/4
✓ Branch 0 taken 1382 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✓ Branch 3 taken 1381 times.
1383 if (bits == 0 || bits > m_keyWords * 32) {
77
1/1
✓ Branch 2 taken 2 times.
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
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1376 times.
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
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1376 times.
1378 if (count == 0) {
93 2 return;
94 }
95
96 1376 const uint32_t passes = (bits + kDigitBits - 1) / kDigitBits;
97
2/2
✓ Branch 0 taken 197 times.
✓ Branch 1 taken 1179 times.
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
2/2
✓ Branch 0 taken 5293 times.
✓ Branch 1 taken 1376 times.
6669 for (uint32_t pass = 0; pass < passes; ++pass) {
106 5293 params.shift = pass * kDigitBits;
107
2/2
✓ Branch 0 taken 1281 times.
✓ Branch 1 taken 4012 times.
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