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 <stdexcept> | ||
| 25 | #include <string> | ||
| 26 | |||
| 27 | namespace brotensor::detail::cpu { | ||
| 28 | |||
| 29 | namespace { | ||
| 30 | |||
| 31 | 2 | [[noreturn]] inline void fail(const char* op, const std::string& reason) { | |
| 32 |
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); |
| 33 | 2 | } | |
| 34 | |||
| 35 | 12 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 36 | const char* op, const char* name) { | ||
| 37 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | if (t.dtype != Dtype::FP32) { |
| 38 | ✗ | fail(op, std::string(name) + | |
| 39 | " must be FP32 (CPU backend is FP32-only)"); | ||
| 40 | } | ||
| 41 | 12 | } | |
| 42 | |||
| 43 | // "(a, idx_a) precedes (b, idx_b)" in our descending-value, ascending-index | ||
| 44 | // ordering. Returns true iff a is strictly preferable to b (i.e. should | ||
| 45 | // appear earlier in the output). | ||
| 46 | 9952 | inline bool prefers(float a, int idx_a, float b, int idx_b) { | |
| 47 |
2/2✓ Branch 0 taken 9942 times.
✓ Branch 1 taken 10 times.
|
9952 | if (a != b) return a > b; |
| 48 | 10 | return idx_a < idx_b; | |
| 49 | 9952 | } | |
| 50 | |||
| 51 | } // namespace | ||
| 52 | |||
| 53 | 12 | void top_k_rows(const ::brotensor::Tensor& X, int k, | |
| 54 | ::brotensor::Tensor& Vals, ::brotensor::Tensor& Idx) { | ||
| 55 | 12 | const char* op = "top_k_rows"; | |
| 56 | 12 | check_fp32(X, op, "X"); | |
| 57 | 12 | const int R = X.rows, C = X.cols; | |
| 58 |
3/4✓ Branch 0 taken 11 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
13 | if (k < 1) fail(op, "k must be >= 1"); |
| 59 |
3/4✓ Branch 0 taken 10 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
11 | if (k > C) fail(op, "k must be <= C (per-row length)"); |
| 60 | |||
| 61 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
10 | if (Vals.rows != R || Vals.cols != k || Vals.dtype != Dtype::FP32) { |
| 62 | 10 | Vals.resize(R, k, Dtype::FP32); | |
| 63 | 10 | } | |
| 64 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
10 | if (Idx.rows != R || Idx.cols != k || Idx.dtype != Dtype::INT32) { |
| 65 | 10 | Idx.resize(R, k, Dtype::INT32); | |
| 66 | 10 | } | |
| 67 |
2/4✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
|
10 | if (R == 0 || k == 0) return; |
| 68 | |||
| 69 | 10 | const float* Xp = X.host_f32(); | |
| 70 | 10 | float* Vp = Vals.host_f32_mut(); | |
| 71 | 10 | int32_t* Ip = static_cast<int32_t*>(Idx.data); | |
| 72 | |||
| 73 | // Per-row working arrays (size k). Keeping them on the stack via VLA | ||
| 74 | // isn't portable on MSVC; allocate on the heap once and reuse. | ||
| 75 | // k is small in practice but C can be huge — so we allocate per row. | ||
| 76 | // The inner top-k is O(k) per candidate scanned. | ||
| 77 |
2/2✓ Branch 0 taken 10 times.
✓ Branch 1 taken 95 times.
|
105 | for (int r = 0; r < R; ++r) { |
| 78 | 95 | const float* row = Xp + static_cast<long>(r) * C; | |
| 79 | 95 | float* out_v = Vp + static_cast<long>(r) * k; | |
| 80 | 95 | int32_t* out_i = Ip + static_cast<long>(r) * k; | |
| 81 | |||
| 82 | // Step 1: seed the working set with the first k elements (verbatim | ||
| 83 | // order — ties get resolved on later replacements). | ||
| 84 |
2/2✓ Branch 0 taken 435 times.
✓ Branch 1 taken 95 times.
|
530 | for (int j = 0; j < k; ++j) { |
| 85 | 435 | out_v[j] = row[j]; | |
| 86 | 435 | out_i[j] = j; | |
| 87 | 435 | } | |
| 88 | // Track the current weakest entry — anything that beats it gets | ||
| 89 | // swapped in. Scan the working set for the weakest each replace. | ||
| 90 | // (For modest k this is faster + cleaner than a heap.) | ||
| 91 | 95 | int weakest = 0; | |
| 92 |
2/2✓ Branch 0 taken 340 times.
✓ Branch 1 taken 95 times.
|
435 | for (int j = 1; j < k; ++j) { |
| 93 |
2/2✓ Branch 0 taken 245 times.
✓ Branch 1 taken 95 times.
|
340 | if (prefers(out_v[weakest], out_i[weakest], out_v[j], out_i[j])) { |
| 94 | 95 | weakest = j; | |
| 95 | 95 | } | |
| 96 | 340 | } | |
| 97 | |||
| 98 | // Step 2: scan the remainder. Replace the weakest if beaten. | ||
| 99 |
2/2✓ Branch 0 taken 4509 times.
✓ Branch 1 taken 95 times.
|
4604 | for (int c = k; c < C; ++c) { |
| 100 | 4509 | const float v = row[c]; | |
| 101 |
2/2✓ Branch 0 taken 3681 times.
✓ Branch 1 taken 828 times.
|
4509 | if (prefers(v, c, out_v[weakest], out_i[weakest])) { |
| 102 | 828 | out_v[weakest] = v; | |
| 103 | 828 | out_i[weakest] = c; | |
| 104 | // Re-find the weakest after replacement. | ||
| 105 | 828 | weakest = 0; | |
| 106 |
2/2✓ Branch 0 taken 4338 times.
✓ Branch 1 taken 828 times.
|
5166 | for (int j = 1; j < k; ++j) { |
| 107 |
4/4✓ Branch 0 taken 3353 times.
✓ Branch 1 taken 985 times.
✓ Branch 2 taken 3353 times.
✓ Branch 3 taken 985 times.
|
8676 | if (prefers(out_v[weakest], out_i[weakest], |
| 108 | 4338 | out_v[j], out_i[j])) { | |
| 109 | 985 | weakest = j; | |
| 110 | 985 | } | |
| 111 | 4338 | } | |
| 112 | 828 | } | |
| 113 | 4509 | } | |
| 114 | |||
| 115 | // Step 3: sort the k survivors into descending-value / ascending- | ||
| 116 | // index order (insertion sort — small k). | ||
| 117 |
2/2✓ Branch 0 taken 340 times.
✓ Branch 1 taken 95 times.
|
435 | for (int i = 1; i < k; ++i) { |
| 118 | 340 | const float v = out_v[i]; | |
| 119 | 340 | const int32_t idx = out_i[i]; | |
| 120 | 340 | int j = i; | |
| 121 |
4/4✓ Branch 0 taken 100 times.
✓ Branch 1 taken 765 times.
✓ Branch 2 taken 525 times.
✓ Branch 3 taken 340 times.
|
865 | while (j > 0 && prefers(v, idx, out_v[j - 1], out_i[j - 1])) { |
| 122 | 525 | out_v[j] = out_v[j - 1]; | |
| 123 | 525 | out_i[j] = out_i[j - 1]; | |
| 124 | 525 | --j; | |
| 125 | } | ||
| 126 | 340 | out_v[j] = v; | |
| 127 | 340 | out_i[j] = idx; | |
| 128 | 340 | } | |
| 129 | 95 | } | |
| 130 | 12 | } | |
| 131 | |||
| 132 | } // namespace brotensor::detail::cpu | ||
| 133 |