GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 97.4% 75 / 0 / 77
Functions: 100.0% 10 / 0 / 10
Branches: 31.8% 14 / 0 / 44
src/cpu/geglu.cpp
Line Branch Exec Source
1 // ─── CPU GEGLU ops (CHUNK 2) ───────────────────────────────────────────────
2 //
3 // FP32 scalar host implementations. Ports the GEGLU kernels in
4 // src/cuda/elementwise.cu — kernel math reproduced verbatim, FP32 path only.
5 //
6 // X is (B, 2D); split along the last dim into A = X[:, :D] (the value half)
7 // and B_half = X[:, D:] (the gate half).
8 // forward: Y(B, D) = A * gelu(B_half)
9 // backward: dX(B,2D): dX[:, :D] = dY * gelu(B_half)
10 // dX[:, D:] = dY * A * gelu'(B_half)
11 //
12 // Two variants: tanh-approximation GELU and exact (erf-based) GELU. Backward
13 // overwrites dX (the GPU kernel writes both halves directly).
14
15 #include <brotensor/tensor.h>
16
17 #include <cmath>
18 #include <stdexcept>
19
20 namespace brotensor::detail::cpu {
21
22 namespace {
23
24 583 inline float gelu_tanh_scalar(float v) {
25 583 constexpr float kSqrt2OverPi = 0.7978845608f;
26 583 const float u = kSqrt2OverPi * (v + 0.044715f * v * v * v);
27 583 return 0.5f * v * (1.0f + std::tanh(u));
28 }
29
30 // Backward needs both gelu(v) and gelu'(v) for the same v; both are
31 // algebraically derivable from a single std::tanh evaluation, so compute it
32 // once here instead of calling gelu_tanh_scalar plus a separate gradient
33 // function (which would redo the tanh).
34 513 inline void gelu_tanh_value_grad(float v, float& g, float& gprime) {
35 513 constexpr float kSqrt2OverPi = 0.7978845608f;
36 513 const float u = kSqrt2OverPi * (v + 0.044715f * v * v * v);
37 513 const float t = std::tanh(u);
38 513 const float dudx = kSqrt2OverPi * (1.0f + 3.0f * 0.044715f * v * v);
39 513 g = 0.5f * v * (1.0f + t);
40 513 gprime = 0.5f * (1.0f + t) + 0.5f * v * (1.0f - t * t) * dudx;
41 513 }
42
43 548 inline float gelu_exact_scalar(float v) {
44 548 constexpr float kInvSqrt2 = 0.70710678118654752440f;
45 548 return 0.5f * v * (1.0f + std::erf(v * kInvSqrt2));
46 }
47
48 // Backward needs both gelu(v) and gelu'(v) for the same v; both share the
49 // same std::erf evaluation (cdf_term), so compute it once here instead of
50 // calling gelu_exact_scalar plus a separate gradient function.
51 513 inline void gelu_exact_value_grad(float v, float& g, float& gprime) {
52 513 constexpr float kInvSqrt2 = 0.70710678118654752440f;
53 513 constexpr float kInvSqrt2Pi = 0.39894228040143267794f;
54 513 const float cdf_term = 0.5f * (1.0f + std::erf(v * kInvSqrt2));
55 513 g = v * cdf_term;
56 513 const float pdf = kInvSqrt2Pi * std::exp(-0.5f * v * v);
57 513 gprime = cdf_term + v * pdf;
58 513 }
59
60 template <typename GeluFn>
61 9 void geglu_forward_impl(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y,
62 const char* op, GeluFn gelu) {
63
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (X.cols % 2 != 0) {
64 ✗ throw std::runtime_error(std::string(op) + ": X.cols must be even (2*D)");
65 }
66 9 const int B = X.rows;
67 9 const int D = X.cols / 2;
68
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 9 if (Y.rows != B || Y.cols != D) Y.resize(B, D);
69 9 const int total = B * D;
70
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (total == 0) return;
71 9 const float* Xp = X.host_f32();
72 9 float* Yp = Y.host_f32_mut();
73 9 const int two_d = 2 * D;
74
2/2
✓ Branch 0 taken 49 times.
✓ Branch 1 taken 9 times.
 58 for (int b = 0; b < B; ++b) {
75
2/2
✓ Branch 0 taken 1131 times.
✓ Branch 1 taken 49 times.
 1180 for (int d = 0; d < D; ++d) {
76 1131 const float a = Xp[b * two_d + d];
77 1131 const float gv_raw = Xp[b * two_d + D + d];
78 1131 Yp[b * D + d] = a * gelu(gv_raw);
79 1131 }
80 49 }
81 9 }
82
83 template <typename GeluValueGradFn>
84 6 void geglu_backward_impl(const ::brotensor::Tensor& X,
85 const ::brotensor::Tensor& dY,
86 ::brotensor::Tensor& dX,
87 const char* op, GeluValueGradFn gelu_value_grad) {
88
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (X.cols % 2 != 0) {
89 ✗ throw std::runtime_error(std::string(op) + ": X.cols must be even (2*D)");
90 }
91 6 const int B = X.rows;
92 6 const int D = X.cols / 2;
93
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 6 if (dX.rows != B || dX.cols != 2 * D) dX.resize(B, 2 * D);
94 6 const int total = B * D;
95
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (total == 0) return;
96 6 const float* Xp = X.host_f32();
97 6 const float* dYp = dY.host_f32();
98 6 float* dXp = dX.host_f32_mut();
99 6 const int two_d = 2 * D;
100
2/2
✓ Branch 0 taken 34 times.
✓ Branch 1 taken 6 times.
 40 for (int b = 0; b < B; ++b) {
101
2/2
✓ Branch 0 taken 1026 times.
✓ Branch 1 taken 34 times.
 1060 for (int d = 0; d < D; ++d) {
102 1026 const float a = Xp[b * two_d + d];
103 1026 const float bh = Xp[b * two_d + D + d];
104 1026 const float dy = dYp[b * D + d];
105 float g, gprime;
106 1026 gelu_value_grad(bh, g, gprime);
107 1026 dXp[b * two_d + d] = dy * g;
108 1026 dXp[b * two_d + D + d] = dy * a * gprime;
109 1026 }
110 34 }
111 6 }
112
113 } // namespace
114
115 5 void geglu_forward(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y) {
116 5 geglu_forward_impl(X, Y, "geglu_forward", gelu_tanh_scalar);
117 5 }
118
119 3 void geglu_backward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& dY,
120 ::brotensor::Tensor& dX) {
121 3 geglu_backward_impl(X, dY, dX, "geglu_backward", gelu_tanh_value_grad);
122 3 }
123
124 4 void geglu_exact_forward(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y) {
125 4 geglu_forward_impl(X, Y, "geglu_exact_forward", gelu_exact_scalar);
126 4 }
127
128 3 void geglu_exact_backward(const ::brotensor::Tensor& X,
129 const ::brotensor::Tensor& dY,
130 ::brotensor::Tensor& dX) {
131 3 geglu_backward_impl(X, dY, dX, "geglu_exact_backward", gelu_exact_value_grad);
132 3 }
133
134 } // namespace brotensor::detail::cpu
135