src/cpu/public_reductions.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU public reductions (CHUNK 1) ─────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations. Ports the GPU kernels in | ||
| 4 | // src/cuda/public_reductions.cu — FP32 path only. | ||
| 5 | // | ||
| 6 | // sum_rows — per-row sum; out is (M, 1). | ||
| 7 | // sum_cols — per-column sum; out is (1, N). | ||
| 8 | // argmax_rows — per-row argmax index, written as FP32; out is (M, 1). | ||
| 9 | // rows_count_above — per-row strict-> counts at two thresholds; counts is | ||
| 10 | // (R, 2) INT32. Accepts FP32 or FP16 X (matches CUDA). | ||
| 11 | // | ||
| 12 | // GPU argmax: ties keep the lowest index (strict `v > best_v`); the scalar | ||
| 13 | // loop here matches that. | ||
| 14 | |||
| 15 | #include <brotensor/tensor.h> | ||
| 16 | |||
| 17 | #include <cstddef> | ||
| 18 | #include <cstdint> | ||
| 19 | #include <stdexcept> | ||
| 20 | |||
| 21 | namespace brotensor::detail::cpu { | ||
| 22 | |||
| 23 | 6 | void sum_rows(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y) { | |
| 24 | 6 | const int M = X.rows; | |
| 25 | 6 | const int N = X.cols; | |
| 26 |
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 (Y.rows != M || Y.cols != 1 || Y.dtype != ::brotensor::Dtype::FP32) { |
| 27 | 6 | Y.resize(M, 1, ::brotensor::Dtype::FP32); | |
| 28 | 6 | } | |
| 29 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (M == 0) return; |
| 30 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (N == 0) { Y.zero(); return; } |
| 31 | 6 | const float* xp = X.host_f32(); | |
| 32 | 6 | float* yp = Y.host_f32_mut(); | |
| 33 |
2/2✓ Branch 0 taken 24 times.
✓ Branch 1 taken 6 times.
|
30 | for (int m = 0; m < M; ++m) { |
| 34 | 24 | const float* row = xp + static_cast<std::size_t>(m) * N; | |
| 35 | 24 | float acc = 0.0f; | |
| 36 |
2/2✓ Branch 0 taken 2056 times.
✓ Branch 1 taken 24 times.
|
2080 | for (int n = 0; n < N; ++n) acc += row[n]; |
| 37 | 24 | yp[m] = acc; | |
| 38 | 24 | } | |
| 39 | 6 | } | |
| 40 | |||
| 41 | 6 | void sum_cols(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y) { | |
| 42 | 6 | const int M = X.rows; | |
| 43 | 6 | const int N = X.cols; | |
| 44 |
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 (Y.rows != 1 || Y.cols != N || Y.dtype != ::brotensor::Dtype::FP32) { |
| 45 | 6 | Y.resize(1, N, ::brotensor::Dtype::FP32); | |
| 46 | 6 | } | |
| 47 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (N == 0) return; |
| 48 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (M == 0) { Y.zero(); return; } |
| 49 | 6 | const float* xp = X.host_f32(); | |
| 50 | 6 | float* yp = Y.host_f32_mut(); | |
| 51 |
2/2✓ Branch 0 taken 72 times.
✓ Branch 1 taken 6 times.
|
78 | for (int n = 0; n < N; ++n) { |
| 52 | 72 | float acc = 0.0f; | |
| 53 |
2/2✓ Branch 0 taken 2056 times.
✓ Branch 1 taken 72 times.
|
2128 | for (int m = 0; m < M; ++m) acc += xp[static_cast<std::size_t>(m) * N + n]; |
| 54 | 72 | yp[n] = acc; | |
| 55 | 72 | } | |
| 56 | 6 | } | |
| 57 | |||
| 58 | // Output dtype is opt-in (mirrors the CUDA contract): an INT32-typed `Idx` | ||
| 59 | // receives the index as int32; otherwise it is written FP32 (the default). | ||
| 60 | 7 | void argmax_rows(const ::brotensor::Tensor& X, ::brotensor::Tensor& Idx) { | |
| 61 | using ::brotensor::Dtype; | ||
| 62 | 7 | const int M = X.rows; | |
| 63 | 7 | const int N = X.cols; | |
| 64 | 7 | const Dtype out_dt = (Idx.dtype == Dtype::INT32) ? Dtype::INT32 : Dtype::FP32; | |
| 65 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
7 | if (Idx.rows != M || Idx.cols != 1 || Idx.dtype != out_dt) { |
| 66 | 7 | Idx.resize(M, 1, out_dt); | |
| 67 | 7 | } | |
| 68 |
1/2✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
|
7 | if (M == 0) return; |
| 69 |
1/2✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
|
7 | if (N == 0) { Idx.zero(); return; } |
| 70 | 7 | const float* xp = X.host_f32(); | |
| 71 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
|
7 | float* fp = (out_dt == Dtype::FP32) ? Idx.host_f32_mut() : nullptr; |
| 72 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
|
7 | int32_t* ip = (out_dt == Dtype::INT32) |
| 73 | ✗ | ? static_cast<int32_t*>(Idx.host_raw_mut()) : nullptr; | |
| 74 |
2/2✓ Branch 0 taken 7 times.
✓ Branch 1 taken 32 times.
|
39 | for (int m = 0; m < M; ++m) { |
| 75 | 32 | const float* row = xp + static_cast<std::size_t>(m) * N; | |
| 76 | 32 | float best_v = -3.4028235e38f; | |
| 77 | 32 | int best_i = 0; | |
| 78 |
2/2✓ Branch 0 taken 3558 times.
✓ Branch 1 taken 32 times.
|
3590 | for (int n = 0; n < N; ++n) { |
| 79 |
2/2✓ Branch 0 taken 3443 times.
✓ Branch 1 taken 115 times.
|
3558 | if (row[n] > best_v) { best_v = row[n]; best_i = n; } |
| 80 | 3558 | } | |
| 81 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
|
32 | if (ip) ip[m] = best_i; |
| 82 | 32 | else fp[m] = static_cast<float>(best_i); | |
| 83 | 32 | } | |
| 84 | 7 | } | |
| 85 | |||
| 86 | // One pass per row over both thresholds. Strict >: a value exactly at a | ||
| 87 | // threshold is not counted. | ||
| 88 | 5 | void rows_count_above(const ::brotensor::Tensor& X, float t_lo, float t_hi, | |
| 89 | ::brotensor::Tensor& counts) { | ||
| 90 | using ::brotensor::Dtype; | ||
| 91 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
5 | if (X.dtype != Dtype::FP32 && X.dtype != Dtype::FP16) { |
| 92 | ✗ | throw std::runtime_error("rows_count_above: X must be FP32 or FP16"); | |
| 93 | } | ||
| 94 | 5 | const int R = X.rows; | |
| 95 | 5 | const int C = X.cols; | |
| 96 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
5 | if (counts.rows != R || counts.cols != 2 || counts.dtype != Dtype::INT32) { |
| 97 | 5 | counts.resize(R, 2, Dtype::INT32); | |
| 98 | 5 | } | |
| 99 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (R == 0) return; |
| 100 | 5 | int32_t* cp = static_cast<int32_t*>(counts.host_raw_mut()); | |
| 101 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
|
5 | const float* xf = (X.dtype == Dtype::FP32) ? X.host_f32() : nullptr; |
| 102 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
|
5 | const uint16_t* xh = (X.dtype == Dtype::FP16) ? X.host_fp16() : nullptr; |
| 103 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 5 times.
|
21 | for (int r = 0; r < R; ++r) { |
| 104 | 16 | int32_t n_lo = 0, n_hi = 0; | |
| 105 | 16 | const std::size_t base = static_cast<std::size_t>(r) * C; | |
| 106 |
2/2✓ Branch 0 taken 6001 times.
✓ Branch 1 taken 16 times.
|
6017 | for (int c = 0; c < C; ++c) { |
| 107 |
1/2✓ Branch 0 taken 6001 times.
✗ Branch 1 not taken.
|
6001 | const float v = xf ? xf[base + c] |
| 108 | ✗ | : ::brotensor::fp16_bits_to_fp32(xh[base + c]); | |
| 109 | 6001 | n_lo += (v > t_lo) ? 1 : 0; | |
| 110 | 6001 | n_hi += (v > t_hi) ? 1 : 0; | |
| 111 | 6001 | } | |
| 112 | 16 | cp[2 * r + 0] = n_lo; | |
| 113 | 16 | cp[2 * r + 1] = n_hi; | |
| 114 | 16 | } | |
| 115 | 5 | } | |
| 116 | |||
| 117 | } // namespace brotensor::detail::cpu | ||
| 118 |