GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 95.1% 214 / 0 / 225
Functions: 100.0% 5 / 0 / 5
Branches: 61.0% 111 / 0 / 182
src/cpu/lstm.cpp
Line Branch Exec Source
1 // CPU implementation of the trainable LSTM cell (ops/lstm.h) — training
2 // forward with per-step caching, and full back-propagation-through-time.
3 //
4 // PyTorch nn.LSTM weight layout: W_ih (4H,I), W_hh (4H,H), gate-row blocks in
5 // order [input, forget, cell, output]; gate g uses tanh, the rest sigmoid.
6 // CPU-resident, FP32 throughout.
7
8 #include <brotensor/tensor.h>
9
10 #include <cmath>
11 #include <stdexcept>
12 #include <string>
13 #include <vector>
14
15 namespace brotensor::detail::cpu {
16
17 namespace {
18
19 108162 inline float sigmoidf(float x) { return 1.0f / (1.0f + std::exp(-x)); }
20
21 11578 void require_cpu_f32(const ::brotensor::Tensor& t, const char* what) {
22
1/2
✓ Branch 0 taken 11578 times.
✗ Branch 1 not taken.
 11578 if (t.device != ::brotensor::Device::CPU ||
23 11578 t.dtype != ::brotensor::Dtype::FP32) {
24 ✗ throw std::runtime_error(std::string("lstm: ") + what +
25 " must be CPU FP32");
26 }
27 11578 }
28
29 // An optional bias/state tensor: present iff non-null and non-empty.
30 4418 const float* opt_ptr(const ::brotensor::Tensor* t, const char* what,
31 int rows, int cols) {
32
3/4
✓ Branch 0 taken 4398 times.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4398 times.
 4418 if (!t || t->empty()) return nullptr;
33 4398 require_cpu_f32(*t, what);
34
1/2
✓ Branch 0 taken 4398 times.
✗ Branch 1 not taken.
 4398 if (t->rows != rows || t->cols != cols) {
35 ✗ throw std::runtime_error(std::string("lstm: ") + what + " has wrong shape");
36 }
37 4398 return t->host_f32();
38 4418 }
39
40 } // namespace
41
42 901 void lstm_forward_train(const ::brotensor::Tensor& X, const ::brotensor::Tensor& W_ih,
43 const ::brotensor::Tensor& W_hh,
44 const ::brotensor::Tensor* b_ih, const ::brotensor::Tensor* b_hh,
45 const ::brotensor::Tensor* h0, const ::brotensor::Tensor* c0,
46 int T, int B,
47 ::brotensor::Tensor& Y, ::brotensor::Tensor& gates,
48 ::brotensor::Tensor& C,
49 ::brotensor::Tensor* hT, ::brotensor::Tensor* cT) {
50 901 require_cpu_f32(X, "X");
51 901 require_cpu_f32(W_ih, "W_ih");
52 901 require_cpu_f32(W_hh, "W_hh");
53 901 const int H = W_hh.cols;
54 901 const int I = W_ih.cols;
55 901 const int G = 4 * H;
56
1/2
✓ Branch 0 taken 901 times.
✗ Branch 1 not taken.
 901 if (T <= 0 || B <= 0)
57 ✗ throw std::runtime_error("lstm_forward_train: T and B must be > 0");
58
1/2
✓ Branch 0 taken 901 times.
✗ Branch 1 not taken.
 901 if (W_ih.rows != G || W_hh.rows != G)
59 ✗ throw std::runtime_error("lstm_forward_train: W_ih/W_hh must have 4*H rows");
60
1/2
✓ Branch 0 taken 901 times.
✗ Branch 1 not taken.
 901 if (X.rows != T * B || X.cols != I)
61 ✗ throw std::runtime_error("lstm_forward_train: X must be (T*B, I)");
62
63 901 const float* bih = opt_ptr(b_ih, "b_ih", G, 1);
64 901 const float* bhh = opt_ptr(b_hh, "b_hh", G, 1);
65 901 const float* h0p = opt_ptr(h0, "h0", B, H);
66 901 const float* c0p = opt_ptr(c0, "c0", B, H);
67
68 901 Y.resize(T * B, H);
69 901 gates.resize(T * B, G);
70 901 C.resize(T * B, H);
71
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 6 times.
 901 if (hT) hT->resize(B, H);
72
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 6 times.
 901 if (cT) cT->resize(B, H);
73
74 901 const float* xp = X.host_f32();
75 901 const float* wih = W_ih.host_f32();
76 901 const float* whh = W_hh.host_f32();
77 901 float* yp = Y.host_f32_mut();
78 901 float* gp = gates.host_f32_mut();
79 901 float* cp = C.host_f32_mut();
80
81
2/2
✓ Branch 0 taken 3598 times.
✓ Branch 1 taken 901 times.
 4499 for (int t = 0; t < T; ++t) {
82
2/2
✓ Branch 0 taken 7200 times.
✓ Branch 1 taken 3598 times.
 10798 for (int b = 0; b < B; ++b) {
83 7200 const std::size_t row = static_cast<std::size_t>(t) * B + b;
84 7200 const float* x = xp + row * I;
85
4/4
✓ Branch 0 taken 5398 times.
✓ Branch 1 taken 1802 times.
✓ Branch 2 taken 1796 times.
✓ Branch 3 taken 6 times.
 7200 const float* hprev = (t == 0) ? (h0p ? h0p + static_cast<std::size_t>(b) * H : nullptr)
86 5398 : yp + (static_cast<std::size_t>(t - 1) * B + b) * H;
87
4/4
✓ Branch 0 taken 5398 times.
✓ Branch 1 taken 1802 times.
✓ Branch 2 taken 1796 times.
✓ Branch 3 taken 6 times.
 7200 const float* cprev = (t == 0) ? (c0p ? c0p + static_cast<std::size_t>(b) * H : nullptr)
88 5398 : cp + (static_cast<std::size_t>(t - 1) * B + b) * H;
89 7200 float* grow = gp + row * G; // post-activation [i|f|g|o]
90 7200 float* crow = cp + row * H;
91 7200 float* hrow = yp + row * H;
92
93
2/2
✓ Branch 0 taken 36054 times.
✓ Branch 1 taken 7200 times.
 43254 for (int n = 0; n < H; ++n) {
94
2/2
✓ Branch 0 taken 35950 times.
✓ Branch 1 taken 104 times.
 36054 float zi = bih ? bih[n] : 0.0f;
95
2/2
✓ Branch 0 taken 35950 times.
✓ Branch 1 taken 104 times.
 36054 float zf = bih ? bih[H + n] : 0.0f;
96
2/2
✓ Branch 0 taken 35950 times.
✓ Branch 1 taken 104 times.
 36054 float zg = bih ? bih[2 * H + n] : 0.0f;
97
2/2
✓ Branch 0 taken 35950 times.
✓ Branch 1 taken 104 times.
 36054 float zo = bih ? bih[3 * H + n] : 0.0f;
98
2/2
✓ Branch 0 taken 104 times.
✓ Branch 1 taken 35950 times.
 36054 if (bhh) { zi += bhh[n]; zf += bhh[H + n]; zg += bhh[2 * H + n]; zo += bhh[3 * H + n]; }
99
100 36054 const float* wi = wih + static_cast<std::size_t>(n) * I;
101 36054 const float* wf = wih + static_cast<std::size_t>(H + n) * I;
102 36054 const float* wg = wih + static_cast<std::size_t>(2 * H + n) * I;
103 36054 const float* wo = wih + static_cast<std::size_t>(3 * H + n) * I;
104
2/2
✓ Branch 0 taken 109236 times.
✓ Branch 1 taken 36054 times.
 145290 for (int j = 0; j < I; ++j) {
105 109236 const float xj = x[j];
106 109236 zi += wi[j] * xj; zf += wf[j] * xj; zg += wg[j] * xj; zo += wo[j] * xj;
107 109236 }
108
2/2
✓ Branch 0 taken 28 times.
✓ Branch 1 taken 36026 times.
 36054 if (hprev) {
109 36026 const float* ui = whh + static_cast<std::size_t>(n) * H;
110 36026 const float* uf = whh + static_cast<std::size_t>(H + n) * H;
111 36026 const float* ug = whh + static_cast<std::size_t>(2 * H + n) * H;
112 36026 const float* uo = whh + static_cast<std::size_t>(3 * H + n) * H;
113
2/2
✓ Branch 0 taken 180534 times.
✓ Branch 1 taken 36026 times.
 216560 for (int m = 0; m < H; ++m) {
114 180534 const float hm = hprev[m];
115 180534 zi += ui[m] * hm; zf += uf[m] * hm; zg += ug[m] * hm; zo += uo[m] * hm;
116 180534 }
117 36026 }
118
119 36054 const float ig = sigmoidf(zi);
120 36054 const float fg = sigmoidf(zf);
121 36054 const float gg = std::tanh(zg);
122 36054 const float og = sigmoidf(zo);
123
2/2
✓ Branch 0 taken 36026 times.
✓ Branch 1 taken 28 times.
 36054 const float cprev_n = cprev ? cprev[n] : 0.0f;
124 36054 const float cn = fg * cprev_n + ig * gg;
125
126 36054 grow[n] = ig; grow[H + n] = fg; grow[2 * H + n] = gg; grow[3 * H + n] = og;
127 36054 crow[n] = cn;
128 36054 hrow[n] = og * std::tanh(cn);
129 36054 }
130 7200 }
131 3598 }
132
133
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 6 times.
 901 if (hT) {
134 6 float* p = hT->host_f32_mut();
135
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 for (int b = 0; b < B; ++b)
136
2/2
✓ Branch 0 taken 65 times.
✓ Branch 1 taken 12 times.
 77 for (int n = 0; n < H; ++n)
137 65 p[static_cast<std::size_t>(b) * H + n] =
138 77 yp[(static_cast<std::size_t>(T - 1) * B + b) * H + n];
139 6 }
140
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 6 times.
 901 if (cT) {
141 6 float* p = cT->host_f32_mut();
142
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 for (int b = 0; b < B; ++b)
143
2/2
✓ Branch 0 taken 65 times.
✓ Branch 1 taken 12 times.
 77 for (int n = 0; n < H; ++n)
144 65 p[static_cast<std::size_t>(b) * H + n] =
145 77 cp[(static_cast<std::size_t>(T - 1) * B + b) * H + n];
146 6 }
147 901 }
148
149 407 void lstm_backward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& W_ih,
150 const ::brotensor::Tensor& W_hh,
151 const ::brotensor::Tensor* h0, const ::brotensor::Tensor* c0,
152 const ::brotensor::Tensor& Y, const ::brotensor::Tensor& gates,
153 const ::brotensor::Tensor& C,
154 const ::brotensor::Tensor& dY, int T, int B,
155 ::brotensor::Tensor& dX, ::brotensor::Tensor& dW_ih,
156 ::brotensor::Tensor& dW_hh,
157 ::brotensor::Tensor* db_ih, ::brotensor::Tensor* db_hh,
158 ::brotensor::Tensor* dh0, ::brotensor::Tensor* dc0) {
159 407 require_cpu_f32(X, "X");
160 407 require_cpu_f32(W_ih, "W_ih");
161 407 require_cpu_f32(W_hh, "W_hh");
162 407 require_cpu_f32(Y, "Y");
163 407 require_cpu_f32(gates, "gates");
164 407 require_cpu_f32(C, "C");
165 407 require_cpu_f32(dY, "dY");
166 407 const int H = W_hh.cols;
167 407 const int I = W_ih.cols;
168 407 const int G = 4 * H;
169
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 if (T <= 0 || B <= 0)
170 ✗ throw std::runtime_error("lstm_backward: T and B must be > 0");
171
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 if (W_ih.rows != G || W_hh.rows != G)
172 ✗ throw std::runtime_error("lstm_backward: W_ih/W_hh must have 4*H rows");
173
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 if (X.rows != T * B || X.cols != I)
174 ✗ throw std::runtime_error("lstm_backward: X must be (T*B, I)");
175
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 if (dY.rows != T * B || dY.cols != H)
176 ✗ throw std::runtime_error("lstm_backward: dY must be (T*B, H)");
177
178 // Parameter grads accumulate — require pre-sized buffers, do NOT zero them.
179 407 require_cpu_f32(dW_ih, "dW_ih");
180 407 require_cpu_f32(dW_hh, "dW_hh");
181
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 if (dW_ih.rows != G || dW_ih.cols != I)
182 ✗ throw std::runtime_error("lstm_backward: dW_ih must be (4H, I), zeroed");
183
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 if (dW_hh.rows != G || dW_hh.cols != H)
184 ✗ throw std::runtime_error("lstm_backward: dW_hh must be (4H, H), zeroed");
185
2/6
✗ Branch 0 not taken.
✓ Branch 1 taken 407 times.
✓ Branch 2 taken 407 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 407 if (db_ih) { require_cpu_f32(*db_ih, "db_ih"); if (db_ih->rows != G || db_ih->cols != 1) throw std::runtime_error("lstm_backward: db_ih must be (4H, 1), zeroed"); }
186
2/6
✗ Branch 0 not taken.
✓ Branch 1 taken 407 times.
✓ Branch 2 taken 407 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 407 if (db_hh) { require_cpu_f32(*db_hh, "db_hh"); if (db_hh->rows != G || db_hh->cols != 1) throw std::runtime_error("lstm_backward: db_hh must be (4H, 1), zeroed"); }
187
188 407 const float* h0p = opt_ptr(h0, "h0", B, H);
189 407 const float* c0p = opt_ptr(c0, "c0", B, H);
190
191 407 dX.resize(T * B, I);
192
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 7 times.
 407 if (dh0) dh0->resize(B, H);
193
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 7 times.
 407 if (dc0) dc0->resize(B, H);
194
195 407 const float* xp = X.host_f32();
196 407 const float* yp = Y.host_f32();
197 407 const float* gp = gates.host_f32();
198 407 const float* cp = C.host_f32();
199 407 const float* dyp = dY.host_f32();
200 407 const float* wih = W_ih.host_f32();
201 407 const float* whh = W_hh.host_f32();
202 407 float* dxp = dX.host_f32_mut();
203 407 float* dwih = dW_ih.host_f32_mut();
204 407 float* dwhh = dW_hh.host_f32_mut();
205
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 float* dbih = db_ih ? db_ih->host_f32_mut() : nullptr;
206
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 float* dbhh = db_hh ? db_hh->host_f32_mut() : nullptr;
207
208 407 std::fill(dxp, dxp + static_cast<std::size_t>(T) * B * I, 0.0f);
209
210 // Recurrent grad carried backward in time: dL/dh_{t} and dL/dc_{t} arriving
211 // from step t+1. After the loop these hold dL/dh0 and dL/dc0.
212 407 std::vector<float> dh_next(static_cast<std::size_t>(B) * H, 0.0f);
213
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 std::vector<float> dc_next(static_cast<std::size_t>(B) * H, 0.0f);
214
1/2
✓ Branch 0 taken 407 times.
✗ Branch 1 not taken.
 407 std::vector<float> dh_new(static_cast<std::size_t>(H));
215
216
2/2
✓ Branch 0 taken 1625 times.
✓ Branch 1 taken 407 times.
 2032 for (int t = T - 1; t >= 0; --t) {
217
2/2
✓ Branch 0 taken 1625 times.
✓ Branch 1 taken 3252 times.
 4877 for (int b = 0; b < B; ++b) {
218 3252 const std::size_t row = static_cast<std::size_t>(t) * B + b;
219 3252 const float* grow = gp + row * G;
220 3252 const float* crow = cp + row * H;
221 3252 const float* x = xp + row * I;
222 3252 const float* dyrow = dyp + row * H;
223
4/4
✓ Branch 0 taken 2438 times.
✓ Branch 1 taken 814 times.
✓ Branch 2 taken 811 times.
✓ Branch 3 taken 3 times.
 3252 const float* hprev = (t == 0) ? (h0p ? h0p + static_cast<std::size_t>(b) * H : nullptr)
224 2438 : yp + (static_cast<std::size_t>(t - 1) * B + b) * H;
225
4/4
✓ Branch 0 taken 2438 times.
✓ Branch 1 taken 814 times.
✓ Branch 2 taken 811 times.
✓ Branch 3 taken 3 times.
 3252 const float* cprev = (t == 0) ? (c0p ? c0p + static_cast<std::size_t>(b) * H : nullptr)
226 2438 : cp + (static_cast<std::size_t>(t - 1) * B + b) * H;
227 3252 float* dxrow = dxp + row * I;
228 3252 float* dh_b = dh_next.data() + static_cast<std::size_t>(b) * H;
229 3252 float* dc_b = dc_next.data() + static_cast<std::size_t>(b) * H;
230
231
1/2
✓ Branch 0 taken 3252 times.
✗ Branch 1 not taken.
 3252 std::fill(dh_new.begin(), dh_new.end(), 0.0f);
232
233
2/2
✓ Branch 0 taken 16287 times.
✓ Branch 1 taken 3252 times.
 19539 for (int n = 0; n < H; ++n) {
234 16287 const float ig = grow[n], fg = grow[H + n], gg = grow[2 * H + n], og = grow[3 * H + n];
235 16287 const float tc = std::tanh(crow[n]);
236 16287 const float dh = dyrow[n] + dh_b[n];
237 16287 const float do_ = dh * tc;
238 16287 const float dc = dh * og * (1.0f - tc * tc) + dc_b[n];
239
2/2
✓ Branch 0 taken 16273 times.
✓ Branch 1 taken 14 times.
 16287 const float cprev_n = cprev ? cprev[n] : 0.0f;
240 16287 const float df = dc * cprev_n;
241 16287 const float di = dc * gg;
242 16287 const float dg = dc * ig;
243 16287 const float dcprev = dc * fg;
244
245 // Gate pre-activation grads (sigmoid'/tanh' on cached outputs).
246 16287 const float dzi = di * ig * (1.0f - ig);
247 16287 const float dzf = df * fg * (1.0f - fg);
248 16287 const float dzg = dg * (1.0f - gg * gg);
249 16287 const float dzo = do_ * og * (1.0f - og);
250
251
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16287 times.
 16287 if (dbih) { dbih[n] += dzi; dbih[H + n] += dzf; dbih[2 * H + n] += dzg; dbih[3 * H + n] += dzo; }
252
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16287 times.
 16287 if (dbhh) { dbhh[n] += dzi; dbhh[H + n] += dzf; dbhh[2 * H + n] += dzg; dbhh[3 * H + n] += dzo; }
253
254 // Input path: dX = W_ih^T·dz ; dW_ih += dz·x^T
255 16287 const float* wi = wih + static_cast<std::size_t>(n) * I;
256 16287 const float* wf = wih + static_cast<std::size_t>(H + n) * I;
257 16287 const float* wg = wih + static_cast<std::size_t>(2 * H + n) * I;
258 16287 const float* wo = wih + static_cast<std::size_t>(3 * H + n) * I;
259 16287 float* dwi = dwih + static_cast<std::size_t>(n) * I;
260 16287 float* dwf = dwih + static_cast<std::size_t>(H + n) * I;
261 16287 float* dwg = dwih + static_cast<std::size_t>(2 * H + n) * I;
262 16287 float* dwo = dwih + static_cast<std::size_t>(3 * H + n) * I;
263
2/2
✓ Branch 0 taken 49398 times.
✓ Branch 1 taken 16287 times.
 65685 for (int j = 0; j < I; ++j) {
264 49398 dxrow[j] += wi[j] * dzi + wf[j] * dzf + wg[j] * dzg + wo[j] * dzo;
265 49398 const float xj = x[j];
266 49398 dwi[j] += dzi * xj; dwf[j] += dzf * xj; dwg[j] += dzg * xj; dwo[j] += dzo * xj;
267 49398 }
268
269 // Recurrent path: dh_prev = W_hh^T·dz ; dW_hh += dz·h_prev^T
270 16287 const float* ui = whh + static_cast<std::size_t>(n) * H;
271 16287 const float* uf = whh + static_cast<std::size_t>(H + n) * H;
272 16287 const float* ug = whh + static_cast<std::size_t>(2 * H + n) * H;
273 16287 const float* uo = whh + static_cast<std::size_t>(3 * H + n) * H;
274 16287 float* dui = dwhh + static_cast<std::size_t>(n) * H;
275 16287 float* duf = dwhh + static_cast<std::size_t>(H + n) * H;
276 16287 float* dug = dwhh + static_cast<std::size_t>(2 * H + n) * H;
277 16287 float* duo = dwhh + static_cast<std::size_t>(3 * H + n) * H;
278
2/2
✓ Branch 0 taken 81633 times.
✓ Branch 1 taken 16287 times.
 97920 for (int m = 0; m < H; ++m) {
279 81633 dh_new[m] += ui[m] * dzi + uf[m] * dzf + ug[m] * dzg + uo[m] * dzo;
280
2/2
✓ Branch 0 taken 66 times.
✓ Branch 1 taken 81567 times.
 81633 if (hprev) {
281 81567 const float hm = hprev[m];
282 81567 dui[m] += dzi * hm; duf[m] += dzf * hm; dug[m] += dzg * hm; duo[m] += dzo * hm;
283 81567 }
284 81633 }
285
286 // dL/dc_{t-1} for this unit (cell carry is per-unit, no coupling).
287 16287 dc_b[n] = dcprev;
288 16287 }
289
290 // Swap in the freshly accumulated dL/dh_{t-1} for the earlier step.
291
2/2
✓ Branch 0 taken 16287 times.
✓ Branch 1 taken 3252 times.
 19539 for (int m = 0; m < H; ++m) dh_b[m] = dh_new[m];
292 3252 }
293 1625 }
294
295
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 400 times.
 407 if (dh0) {
296
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 float* p = dh0->host_f32_mut();
297
2/2
✓ Branch 0 taken 75 times.
✓ Branch 1 taken 7 times.
 82 for (std::size_t i = 0; i < dh_next.size(); ++i) p[i] = dh_next[i];
298 7 }
299
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 400 times.
 407 if (dc0) {
300
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 float* p = dc0->host_f32_mut();
301
2/2
✓ Branch 0 taken 75 times.
✓ Branch 1 taken 7 times.
 82 for (std::size_t i = 0; i < dc_next.size(); ++i) p[i] = dc_next[i];
302 7 }
303 407 }
304
305 } // namespace brotensor::detail::cpu
306