GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 92.6% 112 / 0 / 121
Functions: 100.0% 9 / 0 / 9
Branches: 69.7% 83 / 0 / 119
src/cpu/matmul.cpp
Line Branch Exec Source
1 // ─── CPU matmul ops (CHUNK 2) ──────────────────────────────────────────────
2 //
3 // FP32 scalar host implementations. Ports src/cuda/matmul.cu and
4 // src/cuda/matmul_backward.cu — FP32 path only, row-major throughout.
5 //
6 // forward: C(M, N) = A(M, K) @ B(K, N)
7 // backward: dA(M, K) += dC(M, N) @ B^T(N, K)
8 // dB(K, N) += A^T(K, M) @ dC(M, N)
9 //
10 // ACCUMULATION: the GPU backward kernels atomicAdd partial products into the
11 // caller's dA / dB buffers, so they ACCUMULATE (+=). The caller is responsible
12 // for zeroing dA / dB before the call if a fresh gradient is wanted; the GPU
13 // kernel also requires dA / dB to be pre-sized to (M, K) / (K, N).
14
15 #include <brotensor/tensor.h>
16 #include <brotensor/detail/cpu/thread_pool.h>
17
18 #include <cmath>
19 #include <cstdint>
20 #include <stdexcept>
21
22 namespace brotensor::detail::cpu {
23
24 14 void matmul(const ::brotensor::Tensor& A, const ::brotensor::Tensor& B,
25 ::brotensor::Tensor& C) {
26
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (A.dtype != B.dtype) {
27 ✗ throw std::runtime_error("matmul: A and B must share dtype");
28 }
29 14 const int M = A.rows;
30 14 const int K = A.cols;
31
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (B.rows != K) {
32 ✗ throw std::runtime_error("matmul: shape mismatch (A.cols != B.rows)");
33 }
34 14 const int N = B.cols;
35
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 14 if (C.rows != M || C.cols != N) C.resize(M, N);
36
2/4
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 14 times.
 14 if (M == 0 || N == 0) return;
37
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (K == 0) {
38 ✗ C.zero();
39 ✗ return;
40 }
41 14 const float* Ap = A.host_f32();
42 14 const float* Bp = B.host_f32();
43 14 float* Cp = C.host_f32_mut();
44 // m-k-n order: broadcast A[m,k], walk B's row k and C's row m contiguously
45 // in the innermost loop (both stride-1) instead of striding through B by N
46 // floats per k step. Each m exclusively owns C's row m (B is read-only),
47 // so the outer loop parallelizes across m with no cross-thread writes.
48
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 299 parallel_for(static_cast<std::size_t>(M), [&](std::size_t mi) {
49 285 const int m = static_cast<int>(mi);
50 285 float* Crow = Cp + static_cast<size_t>(m) * N;
51
2/2
✓ Branch 0 taken 20076 times.
✓ Branch 1 taken 285 times.
 20361 for (int n = 0; n < N; ++n) Crow[n] = 0.0f;
52
2/2
✓ Branch 0 taken 13166 times.
✓ Branch 1 taken 285 times.
 13451 for (int k = 0; k < K; ++k) {
53 13166 const float a_mk = Ap[m * K + k];
54 13166 const float* Brow = Bp + static_cast<size_t>(k) * N;
55
2/2
✓ Branch 0 taken 814466 times.
✓ Branch 1 taken 13166 times.
 827632 for (int n = 0; n < N; ++n) {
56 814466 Crow[n] += a_mk * Brow[n];
57 814466 }
58 13166 }
59 285 });
60 14 }
61
62 // Batched A @ B^T, 16-bit (FP16/BF16), FP32 accumulation. Reference triple
63 // loop mirroring fp16_internal::launch_matmul_ABT_batched_impl semantics:
64 // C[b][m,n] = sum_k A[b][m,k] * B[b][n,k] (+ bias[n], then activation).
65 // bias is per-N (broadcast over rows), length N. C is caller-sized.
66 51 void matmul_abt(const ::brotensor::Tensor& A, const ::brotensor::Tensor& B,
67 ::brotensor::Tensor& C,
68 int batch, int M, int N, int K,
69 long long strideA, long long strideB, long long strideC,
70 const ::brotensor::Tensor* bias, int act) {
71
2/2
✓ Branch 0 taken 49 times.
✓ Branch 1 taken 2 times.
 51 if (A.dtype != B.dtype || A.dtype != C.dtype) {
72
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 throw std::runtime_error("matmul_abt: A, B, C must share dtype");
73 }
74 49 const bool is_fp16 = (A.dtype == ::brotensor::Dtype::FP16);
75 49 const bool is_bf16 = (A.dtype == ::brotensor::Dtype::BF16);
76
4/4
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 27 times.
✓ Branch 2 taken 21 times.
✓ Branch 3 taken 1 time.
 49 if (!is_fp16 && !is_bf16) {
77
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 throw std::runtime_error("matmul_abt: dtype must be FP16 or BF16");
78 }
79
4/4
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 28 times.
✓ Branch 2 taken 19 times.
✓ Branch 3 taken 1 time.
 48 if (bias && bias->dtype != A.dtype) {
80
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 throw std::runtime_error("matmul_abt: bias dtype must match operands");
81 }
82
6/6
✓ Branch 0 taken 45 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 44 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 1 time.
✓ Branch 5 taken 43 times.
 47 if (batch <= 0 || M == 0 || N == 0) return;
83
84
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 21 times.
 43 const uint16_t* Ap = is_fp16 ? A.host_fp16() : A.host_bf16();
85
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 21 times.
 43 const uint16_t* Bp = is_fp16 ? B.host_fp16() : B.host_bf16();
86
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 21 times.
 43 uint16_t* Cp = is_fp16 ? C.host_fp16_mut() : C.host_bf16_mut();
87 43 const uint16_t* bp =
88
4/4
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 19 times.
✓ Branch 2 taken 10 times.
✓ Branch 3 taken 9 times.
 43 bias ? (is_fp16 ? bias->host_fp16() : bias->host_bf16()) : nullptr;
89
90 6723 auto to_f32 = [&](uint16_t b) {
91
2/2
✓ Branch 0 taken 3350 times.
✓ Branch 1 taken 3330 times.
 6680 return is_fp16 ? ::brotensor::fp16_bits_to_fp32(b)
92 3330 : ::brotensor::bf16_bits_to_fp32(b);
93 };
94 683 auto from_f32 = [&](float v) {
95
2/2
✓ Branch 0 taken 322 times.
✓ Branch 1 taken 318 times.
 640 return is_fp16 ? ::brotensor::fp32_to_fp16_bits(v)
96 318 : ::brotensor::fp32_to_bf16_bits(v);
97 };
98 683 auto apply_act = [&](float v) -> float {
99
7/7
✓ Branch 0 taken 100 times.
✓ Branch 1 taken 96 times.
✓ Branch 2 taken 96 times.
✓ Branch 3 taken 96 times.
✓ Branch 4 taken 24 times.
✓ Branch 5 taken 132 times.
✓ Branch 6 taken 96 times.
 640 switch (act) {
100 100 case 0: return v; // none
101
2/2
✓ Branch 0 taken 69 times.
✓ Branch 1 taken 63 times.
 132 case 1: return v > 0.0f ? v : 0.0f; // relu
102 case 2: { // gelu (tanh)
103 96 const float c = 0.7978845608028654f; // sqrt(2/pi)
104 96 const float t = c * (v + 0.044715f * v * v * v);
105 96 return 0.5f * v * (1.0f + std::tanh(t));
106 }
107 case 3: // gelu (exact)
108 96 return 0.5f * v * (1.0f + std::erf(v * 0.7071067811865476f));
109 case 4: // silu
110 96 return v / (1.0f + std::exp(-v));
111 case 5: // quick gelu
112 96 return v / (1.0f + std::exp(-1.702f * v));
113 24 default: return v;
114 }
115 640 };
116
117
2/2
✓ Branch 0 taken 43 times.
✓ Branch 1 taken 69 times.
 112 for (int b = 0; b < batch; ++b) {
118 69 const uint16_t* Ab = Ap + static_cast<long long>(b) * strideA;
119 69 const uint16_t* Bb = Bp + static_cast<long long>(b) * strideB;
120 69 uint16_t* Cb = Cp + static_cast<long long>(b) * strideC;
121
2/2
✓ Branch 0 taken 154 times.
✓ Branch 1 taken 69 times.
 223 for (int m = 0; m < M; ++m) {
122
2/2
✓ Branch 0 taken 640 times.
✓ Branch 1 taken 154 times.
 794 for (int n = 0; n < N; ++n) {
123 640 float acc = 0.0f;
124
2/2
✓ Branch 0 taken 3128 times.
✓ Branch 1 taken 640 times.
 3768 for (int k = 0; k < K; ++k) {
125 3128 acc += to_f32(Ab[m * K + k]) * to_f32(Bb[n * K + k]);
126 3128 }
127
2/2
✓ Branch 0 taken 216 times.
✓ Branch 1 taken 424 times.
 640 if (bp) acc += to_f32(bp[n]);
128 640 Cb[m * N + n] = from_f32(apply_act(acc));
129 640 }
130 154 }
131 69 }
132 47 }
133
134 8 void matmul_backward(const ::brotensor::Tensor& A, const ::brotensor::Tensor& B,
135 const ::brotensor::Tensor& dC,
136 ::brotensor::Tensor& dA, ::brotensor::Tensor& dB) {
137
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 16 if (A.dtype != B.dtype || A.dtype != dC.dtype ||
138 8 A.dtype != dA.dtype || A.dtype != dB.dtype) {
139 ✗ throw std::runtime_error("matmul_backward: dtype mismatch");
140 }
141 8 const int M = A.rows;
142 8 const int K = A.cols;
143
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (B.rows != K) {
144 ✗ throw std::runtime_error("matmul_backward: shape mismatch (A.cols != B.rows)");
145 }
146 8 const int N = B.cols;
147
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (dC.rows != M || dC.cols != N) {
148 ✗ throw std::runtime_error("matmul_backward: dC shape mismatch");
149 }
150
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (dA.rows != M || dA.cols != K) {
151 ✗ throw std::runtime_error("matmul_backward: dA must be pre-sized to (M, K)");
152 }
153
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (dB.rows != K || dB.cols != N) {
154 ✗ throw std::runtime_error("matmul_backward: dB must be pre-sized to (K, N)");
155 }
156
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (M == 0 || N == 0 || K == 0) return;
157
158 8 const float* Ap = A.host_f32();
159 8 const float* Bp = B.host_f32();
160 8 const float* dCp = dC.host_f32();
161 8 float* dAp = dA.host_f32_mut();
162 8 float* dBp = dB.host_f32_mut();
163
164 // dA[m, k] += sum_n dC[m, n] * B[k, n] (accumulate — matches GPU atomicAdd)
165 // Each m exclusively owns dA's row m (dC/B are read-only), so this
166 // parallelizes across m with no cross-thread writes.
167
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 87 parallel_for(static_cast<std::size_t>(M), [&](std::size_t mi) {
168 79 const int m = static_cast<int>(mi);
169
2/2
✓ Branch 0 taken 1285 times.
✓ Branch 1 taken 79 times.
 1364 for (int k = 0; k < K; ++k) {
170 1285 float acc = 0.0f;
171
2/2
✓ Branch 0 taken 17263 times.
✓ Branch 1 taken 1285 times.
 18548 for (int n = 0; n < N; ++n) {
172 17263 acc += dCp[m * N + n] * Bp[k * N + n];
173 17263 }
174 1285 dAp[m * K + k] += acc;
175 1285 }
176 79 });
177 // dB[k, n] += sum_m A[m, k] * dC[m, n] (accumulate — matches GPU atomicAdd)
178 // k-m-n order: each k exclusively owns dB's row k (accumulated over all
179 // m), so this axis parallelizes across k with no cross-thread writes —
180 // unlike an m-outer order, which would need every thread to accumulate
181 // into the same shared dB rows across m. dC's row m and dB's row k are
182 // still walked contiguously (both stride-1) in the innermost loop; only
183 // the A[m,k] broadcast load is strided, which doesn't block
184 // vectorization of the n loop.
185
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 126 parallel_for(static_cast<std::size_t>(K), [&](std::size_t ki) {
186 118 const int k = static_cast<int>(ki);
187 118 float* dBrow = dBp + static_cast<size_t>(k) * N;
188
2/2
✓ Branch 0 taken 1237 times.
✓ Branch 1 taken 118 times.
 1355 for (int m = 0; m < M; ++m) {
189 1237 const float a_mk = Ap[m * K + k];
190 1237 const float* dCrow = dCp + static_cast<size_t>(m) * N;
191
2/2
✓ Branch 0 taken 16022 times.
✓ Branch 1 taken 1237 times.
 17259 for (int n = 0; n < N; ++n) {
192 16022 dBrow[n] += a_mk * dCrow[n];
193 16022 }
194 1237 }
195 118 });
196 8 }
197
198 } // namespace brotensor::detail::cpu
199