src/cpu/loss.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU loss ops (CHUNK 1) ──────────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations. Ports the GPU loss kernels in | ||
| 4 | // src/cuda/loss.cu — formulas reproduced verbatim, FP32 path only. | ||
| 5 | // | ||
| 6 | // mse_vec_forward — mean of squared error over all elements. | ||
| 7 | // mse_vec_backward — dPred = (2/n) * (pred - target), overwrite. | ||
| 8 | // softmax_xent_fused — stable softmax + cross-entropy over the flat | ||
| 9 | // tensor; writes probs, dLogits = p - t, returns loss. | ||
| 10 | |||
| 11 | #include <brotensor/tensor.h> | ||
| 12 | |||
| 13 | #include <cmath> | ||
| 14 | |||
| 15 | namespace brotensor::detail::cpu { | ||
| 16 | |||
| 17 | 3 | float mse_vec_forward(const ::brotensor::Tensor& pred, | |
| 18 | const ::brotensor::Tensor& target) { | ||
| 19 | 3 | const int n = pred.size(); | |
| 20 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (n == 0) return 0.0f; |
| 21 | 3 | const float* pp = pred.host_f32(); | |
| 22 | 3 | const float* tp = target.host_f32(); | |
| 23 | 3 | float sum = 0.0f; | |
| 24 |
2/2✓ Branch 0 taken 265 times.
✓ Branch 1 taken 3 times.
|
268 | for (int i = 0; i < n; ++i) { |
| 25 | 265 | const float d = pp[i] - tp[i]; | |
| 26 | 265 | sum += d * d; | |
| 27 | 265 | } | |
| 28 | 3 | return sum / static_cast<float>(n); | |
| 29 | 3 | } | |
| 30 | |||
| 31 | 3 | void mse_vec_backward(const ::brotensor::Tensor& pred, | |
| 32 | const ::brotensor::Tensor& target, | ||
| 33 | ::brotensor::Tensor& dPred) { | ||
| 34 | 3 | const int n = pred.size(); | |
| 35 |
3/6✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
|
3 | if (dPred.rows != pred.rows || dPred.cols != pred.cols || |
| 36 | 3 | dPred.dtype != ::brotensor::Dtype::FP32) { | |
| 37 | ✗ | dPred.resize(pred.rows, pred.cols, ::brotensor::Dtype::FP32); | |
| 38 | ✗ | } | |
| 39 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (n == 0) return; |
| 40 | 3 | const float scale = 2.0f / static_cast<float>(n); | |
| 41 | 3 | const float* pp = pred.host_f32(); | |
| 42 | 3 | const float* tp = target.host_f32(); | |
| 43 | 3 | float* dp = dPred.host_f32_mut(); | |
| 44 |
2/2✓ Branch 0 taken 265 times.
✓ Branch 1 taken 3 times.
|
268 | for (int i = 0; i < n; ++i) dp[i] = scale * (pp[i] - tp[i]); |
| 45 | 3 | } | |
| 46 | |||
| 47 | 5 | float softmax_xent_fused(const ::brotensor::Tensor& logits, | |
| 48 | const ::brotensor::Tensor& target, | ||
| 49 | const float* d_mask, | ||
| 50 | ::brotensor::Tensor& probs, | ||
| 51 | ::brotensor::Tensor& dLogits) { | ||
| 52 | 5 | const int n = logits.size(); | |
| 53 |
3/6✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 5 times.
|
5 | if (probs.rows != logits.rows || probs.cols != logits.cols || |
| 54 | 5 | probs.dtype != ::brotensor::Dtype::FP32) { | |
| 55 | ✗ | probs.resize(logits.rows, logits.cols, ::brotensor::Dtype::FP32); | |
| 56 | ✗ | } | |
| 57 |
3/6✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 5 times.
|
5 | if (dLogits.rows != logits.rows || dLogits.cols != logits.cols || |
| 58 | 5 | dLogits.dtype != ::brotensor::Dtype::FP32) { | |
| 59 | ✗ | dLogits.resize(logits.rows, logits.cols, ::brotensor::Dtype::FP32); | |
| 60 | ✗ | } | |
| 61 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (n == 0) return 0.0f; |
| 62 | |||
| 63 | 5 | const float* lp = logits.host_f32(); | |
| 64 | 5 | const float* tp = target.host_f32(); | |
| 65 | 5 | float* pp = probs.host_f32_mut(); | |
| 66 | 5 | float* dp = dLogits.host_f32_mut(); | |
| 67 | |||
| 68 | // Stable softmax over the flat tensor (single segment). | ||
| 69 | 5 | float m = -1e30f; | |
| 70 |
2/2✓ Branch 0 taken 369 times.
✓ Branch 1 taken 5 times.
|
374 | for (int i = 0; i < n; ++i) { |
| 71 |
4/4✓ Branch 0 taken 96 times.
✓ Branch 1 taken 273 times.
✓ Branch 2 taken 64 times.
✓ Branch 3 taken 32 times.
|
369 | if (d_mask && d_mask[i] < 0.5f) continue; |
| 72 |
2/2✓ Branch 0 taken 315 times.
✓ Branch 1 taken 22 times.
|
337 | if (lp[i] > m) m = lp[i]; |
| 73 | 337 | } | |
| 74 | 5 | float sum = 0.0f; | |
| 75 |
2/2✓ Branch 0 taken 369 times.
✓ Branch 1 taken 5 times.
|
374 | for (int i = 0; i < n; ++i) { |
| 76 |
4/4✓ Branch 0 taken 96 times.
✓ Branch 1 taken 273 times.
✓ Branch 2 taken 64 times.
✓ Branch 3 taken 32 times.
|
369 | if (d_mask && d_mask[i] < 0.5f) { pp[i] = 0.0f; continue; } |
| 77 | 337 | const float e = std::exp(lp[i] - m); | |
| 78 | 337 | pp[i] = e; | |
| 79 | 337 | sum += e; | |
| 80 | 337 | } | |
| 81 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | const float inv = sum > 0.0f ? 1.0f / sum : 0.0f; |
| 82 | |||
| 83 | 5 | float loss = 0.0f; | |
| 84 |
2/2✓ Branch 0 taken 369 times.
✓ Branch 1 taken 5 times.
|
374 | for (int i = 0; i < n; ++i) { |
| 85 |
4/4✓ Branch 0 taken 96 times.
✓ Branch 1 taken 273 times.
✓ Branch 2 taken 64 times.
✓ Branch 3 taken 32 times.
|
369 | if (d_mask && d_mask[i] < 0.5f) { dp[i] = 0.0f; continue; } |
| 86 | 337 | const float p = pp[i] * inv; | |
| 87 | 337 | pp[i] = p; | |
| 88 | 337 | const float t = tp[i]; | |
| 89 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 337 times.
|
337 | if (t > 0.0f) { |
| 90 |
1/2✓ Branch 0 taken 337 times.
✗ Branch 1 not taken.
|
337 | const float pc = p > 1e-12f ? p : 1e-12f; |
| 91 | 337 | loss -= t * std::log(pc); | |
| 92 | 337 | } | |
| 93 | 337 | dp[i] = p - t; | |
| 94 | 337 | } | |
| 95 | 5 | return loss; | |
| 96 | 5 | } | |
| 97 | |||
| 98 | } // namespace brotensor::detail::cpu | ||
| 99 |