src/cpu/top_k.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU per-row top-k ────────────────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementation. For each row of X(R, C), selects the k | ||
| 4 | // largest values and returns them in descending order in `Vals(R, k)` with | ||
| 5 | // their original column indices in `Idx(R, k)`. Ties are broken by smaller | ||
| 6 | // column index — deterministic and stable across runs. | ||
| 7 | // | ||
| 8 | // Used for: classification heads (top-5), NMS pre-filter (keep top-K box | ||
| 9 | // scores per class), beam-search candidates, retrieval rerankers. Not | ||
| 10 | // differentiable — no backward op. | ||
| 11 | // | ||
| 12 | // Algorithm: partial sort via a max-heap of (value, -index) tuples isn't a | ||
| 13 | // great fit for "tie -> smaller index wins" (we'd want a min-key composite), | ||
| 14 | // so for clarity we use a streaming-replacement strategy: maintain a working | ||
| 15 | // array of (value, index) for the current top-k, scan the row, and replace | ||
| 16 | // the current minimum if a candidate beats it. O(C * k) per row — k is | ||
| 17 | // always small (5, 100, etc.) in real use, so the simpler code wins. | ||
| 18 | // | ||
| 19 | // ── ACCUMULATION ──────────────────────────────────────────────────────────── | ||
| 20 | // top_k_rows — Vals and Idx OVERWRITTEN. | ||
| 21 | |||
| 22 | #include <brotensor/tensor.h> | ||
| 23 | |||
| 24 | #include <algorithm> | ||
| 25 | #include <stdexcept> | ||
| 26 | #include <string> | ||
| 27 | #include <vector> | ||
| 28 | |||
| 29 | namespace brotensor::detail::cpu { | ||
| 30 | |||
| 31 | namespace { | ||
| 32 | |||
| 33 | 2 | [[noreturn]] inline void fail(const char* op, const std::string& reason) { | |
| 34 |
7/14✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
✓ Branch 12 taken 2 times.
✗ Branch 13 not taken.
|
2 | throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason); |
| 35 | 2 | } | |
| 36 | |||
| 37 | 32 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 38 | const char* op, const char* name) { | ||
| 39 |
1/2✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
|
32 | if (t.dtype != Dtype::FP32) { |
| 40 | ✗ | fail(op, std::string(name) + | |
| 41 | " must be FP32 (CPU backend is FP32-only)"); | ||
| 42 | } | ||
| 43 | 32 | } | |
| 44 | |||
| 45 | // "(a, idx_a) precedes (b, idx_b)" in our descending-value, ascending-index | ||
| 46 | // ordering. Returns true iff a is strictly preferable to b (i.e. should | ||
| 47 | // appear earlier in the output). | ||
| 48 | 3166996 | inline bool prefers(float a, int idx_a, float b, int idx_b) { | |
| 49 |
2/2✓ Branch 0 taken 2867456 times.
✓ Branch 1 taken 299540 times.
|
3166996 | if (a != b) return a > b; |
| 50 | 299540 | return idx_a < idx_b; | |
| 51 | 3166996 | } | |
| 52 | |||
| 53 | } // namespace | ||
| 54 | |||
| 55 | 32 | void top_k_rows(const ::brotensor::Tensor& X, int k, | |
| 56 | ::brotensor::Tensor& Vals, ::brotensor::Tensor& Idx) { | ||
| 57 | 32 | const char* op = "top_k_rows"; | |
| 58 | 32 | check_fp32(X, op, "X"); | |
| 59 | 32 | const int R = X.rows, C = X.cols; | |
| 60 |
3/4✓ Branch 0 taken 31 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
33 | if (k < 1) fail(op, "k must be >= 1"); |
| 61 |
3/4✓ Branch 0 taken 30 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
31 | if (k > C) fail(op, "k must be <= C (per-row length)"); |
| 62 | |||
| 63 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
30 | if (Vals.rows != R || Vals.cols != k || Vals.dtype != Dtype::FP32) { |
| 64 | 30 | Vals.resize(R, k, Dtype::FP32); | |
| 65 | 30 | } | |
| 66 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
30 | if (Idx.rows != R || Idx.cols != k || Idx.dtype != Dtype::INT32) { |
| 67 | 30 | Idx.resize(R, k, Dtype::INT32); | |
| 68 | 30 | } | |
| 69 |
2/4✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 30 times.
|
30 | if (R == 0 || k == 0) return; |
| 70 | |||
| 71 | 30 | const float* Xp = X.host_f32(); | |
| 72 | 30 | float* Vp = Vals.host_f32_mut(); | |
| 73 | 30 | int32_t* Ip = static_cast<int32_t*>(Idx.data); | |
| 74 | |||
| 75 | // Large k (the masked-diffusion host loop asks for k up to the full row | ||
| 76 | // length): the streaming replacement below is O(C * k) per row, so past | ||
| 77 | // a modest k a full sort of the row's indices under the same total order | ||
| 78 | // — O(C log C), independent of k — is the cheaper route. The order is | ||
| 79 | // total (descending value, ascending index), so both paths produce the | ||
| 80 | // identical Vals / Idx; only the cost model differs. | ||
| 81 | 30 | constexpr int kSortThreshold = 64; | |
| 82 | 30 | std::vector<int32_t> order; | |
| 83 |
3/4✓ Branch 0 taken 14 times.
✓ Branch 1 taken 16 times.
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
|
30 | if (k > kSortThreshold) order.resize(C); |
| 84 | |||
| 85 | // Per-row working arrays (size k). Keeping them on the stack via VLA | ||
| 86 | // isn't portable on MSVC; allocate on the heap once and reuse. | ||
| 87 | // k is small in practice but C can be huge — so we allocate per row. | ||
| 88 | // The inner top-k is O(k) per candidate scanned. | ||
| 89 |
2/2✓ Branch 0 taken 117 times.
✓ Branch 1 taken 30 times.
|
147 | for (int r = 0; r < R; ++r) { |
| 90 | 117 | const float* row = Xp + static_cast<long>(r) * C; | |
| 91 | 117 | float* out_v = Vp + static_cast<long>(r) * k; | |
| 92 | 117 | int32_t* out_i = Ip + static_cast<long>(r) * k; | |
| 93 | |||
| 94 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 101 times.
|
117 | if (k > kSortThreshold) { |
| 95 |
2/2✓ Branch 0 taken 375848 times.
✓ Branch 1 taken 16 times.
|
375864 | for (int c = 0; c < C; ++c) order[c] = c; |
| 96 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
32 | std::partial_sort(order.begin(), order.begin() + k, order.end(), |
| 97 | 2941626 | [row](int32_t a, int32_t b) { | |
| 98 | 2941610 | return prefers(row[a], a, row[b], b); | |
| 99 | }); | ||
| 100 |
2/2✓ Branch 0 taken 110203 times.
✓ Branch 1 taken 16 times.
|
110219 | for (int j = 0; j < k; ++j) { |
| 101 | 110203 | out_i[j] = order[j]; | |
| 102 | 110203 | out_v[j] = row[order[j]]; | |
| 103 | 110203 | } | |
| 104 | 16 | continue; | |
| 105 | } | ||
| 106 | |||
| 107 | // Step 1: seed the working set with the first k elements (verbatim | ||
| 108 | // order — ties get resolved on later replacements). | ||
| 109 |
2/2✓ Branch 0 taken 569 times.
✓ Branch 1 taken 101 times.
|
670 | for (int j = 0; j < k; ++j) { |
| 110 | 569 | out_v[j] = row[j]; | |
| 111 | 569 | out_i[j] = j; | |
| 112 | 569 | } | |
| 113 | // Track the current weakest entry — anything that beats it gets | ||
| 114 | // swapped in. Scan the working set for the weakest each replace. | ||
| 115 | // (For modest k this is faster + cleaner than a heap.) | ||
| 116 | 101 | int weakest = 0; | |
| 117 |
2/2✓ Branch 0 taken 101 times.
✓ Branch 1 taken 468 times.
|
569 | for (int j = 1; j < k; ++j) { |
| 118 |
3/4✓ Branch 0 taken 468 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 366 times.
✓ Branch 3 taken 102 times.
|
468 | if (prefers(out_v[weakest], out_i[weakest], out_v[j], out_i[j])) { |
| 119 | 102 | weakest = j; | |
| 120 | 102 | } | |
| 121 | 468 | } | |
| 122 | |||
| 123 | // Step 2: scan the remainder. Replace the weakest if beaten. | ||
| 124 |
2/2✓ Branch 0 taken 101 times.
✓ Branch 1 taken 171415 times.
|
171516 | for (int c = k; c < C; ++c) { |
| 125 | 171415 | const float v = row[c]; | |
| 126 |
3/4✓ Branch 0 taken 171415 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 169818 times.
✓ Branch 3 taken 1597 times.
|
171415 | if (prefers(v, c, out_v[weakest], out_i[weakest])) { |
| 127 | 1597 | out_v[weakest] = v; | |
| 128 | 1597 | out_i[weakest] = c; | |
| 129 | // Re-find the weakest after replacement. | ||
| 130 | 1597 | weakest = 0; | |
| 131 |
2/2✓ Branch 0 taken 1597 times.
✓ Branch 1 taken 50446 times.
|
52043 | for (int j = 1; j < k; ++j) { |
| 132 |
4/6✓ Branch 0 taken 50446 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 50446 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 47116 times.
✓ Branch 5 taken 3330 times.
|
100892 | if (prefers(out_v[weakest], out_i[weakest], |
| 133 | 50446 | out_v[j], out_i[j])) { | |
| 134 | 3330 | weakest = j; | |
| 135 | 3330 | } | |
| 136 | 50446 | } | |
| 137 | 1597 | } | |
| 138 | 171415 | } | |
| 139 | |||
| 140 | // Step 3: sort the k survivors into descending-value / ascending- | ||
| 141 | // index order (insertion sort — small k). | ||
| 142 |
2/2✓ Branch 0 taken 468 times.
✓ Branch 1 taken 101 times.
|
569 | for (int i = 1; i < k; ++i) { |
| 143 | 468 | const float v = out_v[i]; | |
| 144 | 468 | const int32_t idx = out_i[i]; | |
| 145 | 468 | int j = i; | |
| 146 |
5/6✓ Branch 0 taken 109 times.
✓ Branch 1 taken 3057 times.
✓ Branch 2 taken 3057 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 468 times.
✓ Branch 5 taken 2698 times.
|
3634 | while (j > 0 && prefers(v, idx, out_v[j - 1], out_i[j - 1])) { |
| 147 | 2698 | out_v[j] = out_v[j - 1]; | |
| 148 | 2698 | out_i[j] = out_i[j - 1]; | |
| 149 | 2698 | --j; | |
| 150 | } | ||
| 151 | 468 | out_v[j] = v; | |
| 152 | 468 | out_i[j] = idx; | |
| 153 | 468 | } | |
| 154 | 101 | } | |
| 155 | 32 | } | |
| 156 | |||
| 157 | } // namespace brotensor::detail::cpu | ||
| 158 |