src/cpu/gather_rows.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU row gather + scatter-add ─────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementation. General "look up these rows" — the same | ||
| 4 | // compute as embedding_lookup_forward but taking Idx as a Tensor (so it | ||
| 5 | // routes through the dispatcher and works on any backend). Used by SAM's | ||
| 6 | // prompt encoder (point indices -> positional embeddings), DETR object | ||
| 7 | // queries, per-row attention reorderings, etc. | ||
| 8 | // | ||
| 9 | // gather_rows : Y[m, :] = X[Idx[m], :] | ||
| 10 | // scatter_rows_add : dX[Idx[m], :] += dY[m, :] (with dX zeroed first; | ||
| 11 | // duplicate indices in Idx accumulate — adjoint of | ||
| 12 | // gather_rows). | ||
| 13 | // | ||
| 14 | // Idx layout: (M, 1) INT32, matching argmax_rows / top_k_rows' INT32 idx | ||
| 15 | // conventions. M is read from Idx.rows (cols must be 1). | ||
| 16 | // | ||
| 17 | // ── ACCUMULATION ──────────────────────────────────────────────────────────── | ||
| 18 | // gather_rows — Y OVERWRITTEN (pure copy from X). | ||
| 19 | // scatter_rows_add — dX OVERWRITTEN (zeroed, then scatter-add). | ||
| 20 | // | ||
| 21 | // Bounds: an OOB Idx value triggers undefined behavior on all backends. | ||
| 22 | // The caller owns the precondition (same as embedding_lookup). | ||
| 23 | |||
| 24 | #include <brotensor/tensor.h> | ||
| 25 | |||
| 26 | #include <cstring> | ||
| 27 | #include <stdexcept> | ||
| 28 | #include <string> | ||
| 29 | |||
| 30 | namespace brotensor::detail::cpu { | ||
| 31 | |||
| 32 | namespace { | ||
| 33 | |||
| 34 | 2 | [[noreturn]] inline void fail(const char* op, const std::string& reason) { | |
| 35 |
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); |
| 36 | 2 | } | |
| 37 | |||
| 38 | 17 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 39 | const char* op, const char* name) { | ||
| 40 |
1/2✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
|
17 | if (t.dtype != Dtype::FP32) { |
| 41 | ✗ | fail(op, std::string(name) + | |
| 42 | " must be FP32 (CPU backend is FP32-only)"); | ||
| 43 | } | ||
| 44 | 17 | } | |
| 45 | |||
| 46 | 17 | inline void check_idx(const ::brotensor::Tensor& Idx, const char* op) { | |
| 47 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 1 time.
|
17 | if (Idx.dtype != Dtype::INT32) { |
| 48 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
2 | fail(op, "Idx must be INT32"); |
| 49 | } | ||
| 50 |
2/2✓ Branch 0 taken 15 times.
✓ Branch 1 taken 1 time.
|
16 | if (Idx.cols != 1) { |
| 51 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | fail(op, "Idx must be shaped (M, 1)"); |
| 52 | } | ||
| 53 | 17 | } | |
| 54 | |||
| 55 | } // namespace | ||
| 56 | |||
| 57 | 7 | void gather_rows(const ::brotensor::Tensor& X, | |
| 58 | const ::brotensor::Tensor& Idx, | ||
| 59 | ::brotensor::Tensor& Y) { | ||
| 60 | 7 | const char* op = "gather_rows"; | |
| 61 | 7 | check_fp32(X, op, "X"); | |
| 62 | 7 | check_idx(Idx, op); | |
| 63 | 7 | const int M = Idx.rows; | |
| 64 | 7 | const int C = X.cols; | |
| 65 |
3/6✓ Branch 0 taken 1 time.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
7 | if (Y.rows != M || Y.cols != C || Y.dtype != Dtype::FP32) { |
| 66 | 7 | Y.resize(M, C, Dtype::FP32); | |
| 67 | 7 | } | |
| 68 |
3/4✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
|
7 | if (M == 0 || C == 0) return; |
| 69 | |||
| 70 | 6 | const float* Xp = X.host_f32(); | |
| 71 | 6 | const int32_t* Ip = static_cast<const int32_t*>(Idx.data); | |
| 72 | 6 | float* Yp = Y.host_f32_mut(); | |
| 73 | |||
| 74 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 91 times.
|
97 | for (int m = 0; m < M; ++m) { |
| 75 | 91 | const int r = Ip[m]; | |
| 76 | 182 | std::memcpy(Yp + static_cast<long>(m) * C, | |
| 77 | 91 | Xp + static_cast<long>(r) * C, | |
| 78 | 91 | static_cast<size_t>(C) * sizeof(float)); | |
| 79 | 91 | } | |
| 80 | 7 | } | |
| 81 | |||
| 82 | 6 | void scatter_rows_add(const ::brotensor::Tensor& dY, | |
| 83 | const ::brotensor::Tensor& Idx, int R, | ||
| 84 | ::brotensor::Tensor& dX) { | ||
| 85 | 6 | const char* op = "scatter_rows_add"; | |
| 86 | 6 | check_fp32(dY, op, "dY"); | |
| 87 | 6 | check_idx(Idx, op); | |
| 88 |
1/4✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
6 | if (R < 0) fail(op, "R must be >= 0"); |
| 89 | 6 | const int M = Idx.rows; | |
| 90 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (dY.rows != M) { |
| 91 | ✗ | fail(op, "dY.rows must equal Idx.rows"); | |
| 92 | } | ||
| 93 | 6 | const int C = dY.cols; | |
| 94 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
6 | if (dX.rows != R || dX.cols != C || dX.dtype != Dtype::FP32) { |
| 95 | 6 | dX.resize(R, C, Dtype::FP32); | |
| 96 | 6 | } | |
| 97 |
2/4✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
|
6 | if (R == 0 || C == 0) return; |
| 98 | |||
| 99 | 6 | const float* dYp = dY.host_f32(); | |
| 100 | 6 | const int32_t* Ip = static_cast<const int32_t*>(Idx.data); | |
| 101 | 6 | float* dXp = dX.host_f32_mut(); | |
| 102 | |||
| 103 | 6 | const long total = static_cast<long>(R) * C; | |
| 104 |
2/2✓ Branch 0 taken 8464 times.
✓ Branch 1 taken 6 times.
|
8470 | for (long i = 0; i < total; ++i) dXp[i] = 0.0f; |
| 105 | |||
| 106 |
2/2✓ Branch 0 taken 91 times.
✓ Branch 1 taken 6 times.
|
97 | for (int m = 0; m < M; ++m) { |
| 107 | 91 | const int r = Ip[m]; | |
| 108 | 91 | const float* src = dYp + static_cast<long>(m) * C; | |
| 109 | 91 | float* dst = dXp + static_cast<long>(r) * C; | |
| 110 |
2/2✓ Branch 0 taken 16589 times.
✓ Branch 1 taken 91 times.
|
16680 | for (int c = 0; c < C; ++c) dst[c] += src[c]; |
| 111 | 91 | } | |
| 112 | 6 | } | |
| 113 | |||
| 114 | 2 | void scatter_rows(const ::brotensor::Tensor& Y, | |
| 115 | const ::brotensor::Tensor& Idx, | ||
| 116 | ::brotensor::Tensor& X) { | ||
| 117 | 2 | const char* op = "scatter_rows"; | |
| 118 | 2 | check_fp32(Y, op, "Y"); | |
| 119 | 2 | check_idx(Idx, op); | |
| 120 | 2 | const int M = Idx.rows; | |
| 121 |
1/4✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
2 | if (Y.rows != M) fail(op, "Y.rows must equal Idx.rows"); |
| 122 |
1/4✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
2 | if (X.dtype != Dtype::FP32) fail(op, "X must be FP32"); |
| 123 |
1/4✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
2 | if (X.cols != Y.cols) fail(op, "X.cols must equal Y.cols"); |
| 124 | 2 | const int C = Y.cols; | |
| 125 |
2/4✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
|
2 | if (M == 0 || C == 0) return; |
| 126 | |||
| 127 | 2 | const float* Yp = Y.host_f32(); | |
| 128 | 2 | const int32_t* Ip = static_cast<const int32_t*>(Idx.data); | |
| 129 | 2 | float* Xp = X.host_f32_mut(); | |
| 130 | |||
| 131 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 7 times.
|
9 | for (int m = 0; m < M; ++m) { |
| 132 | 7 | const int r = Ip[m]; | |
| 133 | 14 | std::memcpy(Xp + static_cast<long>(r) * C, | |
| 134 | 7 | Yp + static_cast<long>(m) * C, | |
| 135 | 7 | static_cast<size_t>(C) * sizeof(float)); | |
| 136 | 7 | } | |
| 137 | 2 | } | |
| 138 | |||
| 139 | } // namespace brotensor::detail::cpu | ||
| 140 |