src/cpu/group_norm.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU GroupNorm ops (CHUNK 3) ─────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations. Ports src/cuda/group_norm.cu — FP32 path | ||
| 4 | // only. NCHW activations; one tile per (sample, group). | ||
| 5 | // | ||
| 6 | // A tile spans channels_per_group * spatial elements, where | ||
| 7 | // channels_per_group = C / num_groups and spatial = H * W. | ||
| 8 | // | ||
| 9 | // Forward: | ||
| 10 | // mean = (1/M) Σ x | ||
| 11 | // var = (1/M) Σ x² - mean² (GPU's biased estimator) | ||
| 12 | // rstd = 1 / sqrt(var + eps) | ||
| 13 | // y[c,s] = (x - mean) * rstd * gamma[c] + beta[c] | ||
| 14 | // | ||
| 15 | // Backward (per tile; dx̂ = dY*γ_c, x̂ = (x-mean)*rstd): | ||
| 16 | // sum1 = Σ dx̂ | ||
| 17 | // sum2 = Σ dx̂ · x̂ | ||
| 18 | // dX = rstd * (dx̂ - (sum1 + x̂*sum2) / M) | ||
| 19 | // dGamma_c += Σ_{s in tile,batch} dY · x̂ | ||
| 20 | // dBeta_c += Σ_{s in tile,batch} dY | ||
| 21 | // | ||
| 22 | // ACCUMULATION (matches the GPU kernels): | ||
| 23 | // group_norm_forward — Y OVERWRITTEN. | ||
| 24 | // group_norm_backward — dX OVERWRITTEN; dGamma / dBeta ACCUMULATE (+=). | ||
| 25 | // The GPU atomicAdds into FP32 scratch then folds | ||
| 26 | // into the caller's dGamma/dBeta — caller zeros them. | ||
| 27 | // | ||
| 28 | // FP32 accumulation mirrors the GPU (single-precision sums); a wider double | ||
| 29 | // accumulator would diverge from the GPU and break parity. | ||
| 30 | |||
| 31 | #include <brotensor/tensor.h> | ||
| 32 | #include <brotensor/detail/cpu/thread_pool.h> | ||
| 33 | |||
| 34 | #include <cmath> | ||
| 35 | #include <cstddef> | ||
| 36 | #include <stdexcept> | ||
| 37 | #include <vector> | ||
| 38 | |||
| 39 | namespace brotensor::detail::cpu { | ||
| 40 | |||
| 41 | 41 | void group_norm_forward(const ::brotensor::Tensor& X, | |
| 42 | const ::brotensor::Tensor& gamma, | ||
| 43 | const ::brotensor::Tensor& beta, | ||
| 44 | int N, int C, int H, int W, | ||
| 45 | int num_groups, | ||
| 46 | float eps, | ||
| 47 | ::brotensor::Tensor& Y) { | ||
| 48 |
1/2✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
|
41 | if (gamma.dtype != X.dtype || beta.dtype != X.dtype) { |
| 49 | ✗ | throw std::runtime_error("group_norm_forward: gamma/beta dtype must match X"); | |
| 50 | } | ||
| 51 |
1/2✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
|
41 | if (num_groups <= 0 || C % num_groups != 0) { |
| 52 | ✗ | throw std::runtime_error("group_norm_forward: num_groups must divide C"); | |
| 53 | } | ||
| 54 | 41 | const int spatial = H * W; | |
| 55 | 41 | const int cols = C * spatial; | |
| 56 |
1/2✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
|
41 | if (gamma.size() != C || beta.size() != C) { |
| 57 | ✗ | throw std::runtime_error("group_norm_forward: gamma/beta must have C elements"); | |
| 58 | } | ||
| 59 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 41 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
41 | if (Y.rows != N || Y.cols != cols || Y.dtype != X.dtype) { |
| 60 | 41 | Y.resize(N, cols, X.dtype); | |
| 61 | 41 | } | |
| 62 |
2/4✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 41 times.
|
41 | if (N == 0 || cols == 0) return; |
| 63 | |||
| 64 | 41 | const float* Xp = X.host_f32(); | |
| 65 | 41 | const float* gp = gamma.host_f32(); | |
| 66 | 41 | const float* bp = beta.host_f32(); | |
| 67 | 41 | float* Yp = Y.host_f32_mut(); | |
| 68 | |||
| 69 | 41 | const int channels_per_group = C / num_groups; | |
| 70 | 41 | const int tile_size = channels_per_group * spatial; | |
| 71 | 41 | const int sample_stride = C * spatial; | |
| 72 | 41 | const float inv_M = 1.0f / static_cast<float>(tile_size); | |
| 73 | |||
| 74 | // Each sample n owns its own slice of Y exclusively (X/gamma/beta are | ||
| 75 | // read-only shared), so this parallelizes across n with no cross-thread | ||
| 76 | // writes. | ||
| 77 |
1/2✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
|
98 | parallel_for(static_cast<std::size_t>(N), [&](std::size_t ni) { |
| 78 | 57 | const int n = static_cast<int>(ni); | |
| 79 |
2/2✓ Branch 0 taken 204 times.
✓ Branch 1 taken 57 times.
|
261 | for (int g = 0; g < num_groups; ++g) { |
| 80 | 204 | const int chan_base = g * channels_per_group; | |
| 81 | 204 | const float* x_tile = Xp + n * sample_stride + chan_base * spatial; | |
| 82 | 204 | float* y_tile = Yp + n * sample_stride + chan_base * spatial; | |
| 83 | |||
| 84 | // Two passes: mean, then the sum of squared deviations from it. | ||
| 85 | // E[x^2] - E[x]^2 is one pass but cancels catastrophically once a | ||
| 86 | // tile's mean dwarfs its spread — a near-constant channel puts both | ||
| 87 | // terms on the same large value, so their FP32 difference is noise | ||
| 88 | // and can come out negative, making rstd NaN. Deviations are | ||
| 89 | // non-negative by construction, so var cannot go negative. | ||
| 90 | 204 | float sum = 0.0f; | |
| 91 |
2/2✓ Branch 0 taken 51768 times.
✓ Branch 1 taken 204 times.
|
51972 | for (int i = 0; i < tile_size; ++i) sum += x_tile[i]; |
| 92 | 204 | const float mean = sum * inv_M; | |
| 93 | |||
| 94 | 204 | float sumsq = 0.0f; | |
| 95 |
2/2✓ Branch 0 taken 51768 times.
✓ Branch 1 taken 204 times.
|
51972 | for (int i = 0; i < tile_size; ++i) { |
| 96 | 51768 | const float d = x_tile[i] - mean; | |
| 97 | 51768 | sumsq += d * d; | |
| 98 | 51768 | } | |
| 99 | 204 | const float var = sumsq * inv_M; | |
| 100 | 204 | const float rstd = 1.0f / std::sqrt(var + eps); | |
| 101 | |||
| 102 |
2/2✓ Branch 0 taken 51768 times.
✓ Branch 1 taken 204 times.
|
51972 | for (int i = 0; i < tile_size; ++i) { |
| 103 | 51768 | const int local_c = i / spatial; | |
| 104 | 51768 | const int channel = chan_base + local_c; | |
| 105 | 51768 | const float yn = (x_tile[i] - mean) * rstd; | |
| 106 | 51768 | y_tile[i] = yn * gp[channel] + bp[channel]; // overwrite | |
| 107 | 51768 | } | |
| 108 | 204 | } | |
| 109 | 57 | }); | |
| 110 | 41 | } | |
| 111 | |||
| 112 | 25 | void group_norm_backward(const ::brotensor::Tensor& X, | |
| 113 | const ::brotensor::Tensor& gamma, | ||
| 114 | const ::brotensor::Tensor& dY, | ||
| 115 | int N, int C, int H, int W, | ||
| 116 | int num_groups, | ||
| 117 | float eps, | ||
| 118 | ::brotensor::Tensor& dX, | ||
| 119 | ::brotensor::Tensor& dGamma, | ||
| 120 | ::brotensor::Tensor& dBeta) { | ||
| 121 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | if (gamma.dtype != X.dtype || dY.dtype != X.dtype) { |
| 122 | ✗ | throw std::runtime_error("group_norm_backward: gamma/dY dtype must match X"); | |
| 123 | } | ||
| 124 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | if (dGamma.dtype != X.dtype || dBeta.dtype != X.dtype) { |
| 125 | ✗ | throw std::runtime_error("group_norm_backward: dGamma/dBeta dtype must match X"); | |
| 126 | } | ||
| 127 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | if (num_groups <= 0 || C % num_groups != 0) { |
| 128 | ✗ | throw std::runtime_error("group_norm_backward: num_groups must divide C"); | |
| 129 | } | ||
| 130 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
50 | if (dGamma.rows != C || dGamma.cols != 1 || |
| 131 | 25 | dBeta.rows != C || dBeta.cols != 1) { | |
| 132 | ✗ | throw std::runtime_error("group_norm_backward: dGamma/dBeta must be (C,1)"); | |
| 133 | } | ||
| 134 | 25 | const int spatial = H * W; | |
| 135 | 25 | const int cols = C * spatial; | |
| 136 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | if (gamma.size() != C) { |
| 137 | ✗ | throw std::runtime_error("group_norm_backward: gamma must have C elements"); | |
| 138 | } | ||
| 139 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | if (dY.rows != N || dY.cols != cols) { |
| 140 | ✗ | throw std::runtime_error("group_norm_backward: dY shape mismatch"); | |
| 141 | } | ||
| 142 |
4/6✓ Branch 0 taken 8 times.
✓ Branch 1 taken 17 times.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
|
25 | if (dX.rows != N || dX.cols != cols || dX.dtype != X.dtype) { |
| 143 | 17 | dX.resize(N, cols, X.dtype); | |
| 144 | 17 | } | |
| 145 |
2/4✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 25 times.
|
25 | if (N == 0 || cols == 0) return; |
| 146 | |||
| 147 | 25 | const float* Xp = X.host_f32(); | |
| 148 | 25 | const float* gp = gamma.host_f32(); | |
| 149 | 25 | const float* dYp = dY.host_f32(); | |
| 150 | 25 | float* dXp = dX.host_f32_mut(); | |
| 151 | 25 | float* dGp = dGamma.host_f32_mut(); | |
| 152 | 25 | float* dBp = dBeta.host_f32_mut(); | |
| 153 | |||
| 154 | 25 | const int channels_per_group = C / num_groups; | |
| 155 | 25 | const int tile_size = channels_per_group * spatial; | |
| 156 | 25 | const int sample_stride = C * spatial; | |
| 157 | 25 | const float inv_M = 1.0f / static_cast<float>(tile_size); | |
| 158 | |||
| 159 | // dGamma_c / dBeta_c accumulate over EVERY n (a fixed channel appears in | ||
| 160 | // exactly one group but in all N samples' tiles), so a naive | ||
| 161 | // parallel-for over n would race on dGp[channel]/dBp[channel] — that's a | ||
| 162 | // shared reduction across the batch axis, not disjoint per n. | ||
| 163 | // | ||
| 164 | // Fix: each n instead writes its per-channel partial dg/db into row n of | ||
| 165 | // a (N, C) scratch buffer (disjoint per n — every channel is touched by | ||
| 166 | // exactly one group/local_c per n, so each element of the row is | ||
| 167 | // written exactly once, no init needed). A final single-threaded pass | ||
| 168 | // sums those partials down into dGp/dBp. Pass 1/2/3 (stats, partial | ||
| 169 | // dg/db + x̂ caching, dX write) are otherwise fully independent per n and | ||
| 170 | // run inside the parallel_for; x̂'s scratch buffer is declared LOCAL to | ||
| 171 | // the lambda (one per n, sized tile_size) instead of shared across n, so | ||
| 172 | // concurrent n's never touch the same buffer. | ||
| 173 | 25 | std::vector<float> dg_partial(static_cast<std::size_t>(N) * C); | |
| 174 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | std::vector<float> db_partial(static_cast<std::size_t>(N) * C); |
| 175 | |||
| 176 |
2/4✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 25 times.
✗ Branch 3 not taken.
|
58 | parallel_for(static_cast<std::size_t>(N), [&](std::size_t ni) { |
| 177 | 33 | const int n = static_cast<int>(ni); | |
| 178 | 33 | float* dg_row = dg_partial.data() + static_cast<std::size_t>(n) * C; | |
| 179 | 33 | float* db_row = db_partial.data() + static_cast<std::size_t>(n) * C; | |
| 180 | // Local per-n scratch for x̂ — safe even though this lambda runs on | ||
| 181 | // different threads, since each invocation gets its own stack/heap- | ||
| 182 | // local buffer. | ||
| 183 | 33 | std::vector<float> xhat_buf(static_cast<std::size_t>(tile_size)); | |
| 184 | |||
| 185 |
2/2✓ Branch 0 taken 124 times.
✓ Branch 1 taken 33 times.
|
157 | for (int g = 0; g < num_groups; ++g) { |
| 186 | 124 | const int chan_base = g * channels_per_group; | |
| 187 | 124 | const float* x_tile = Xp + n * sample_stride + chan_base * spatial; | |
| 188 | 124 | const float* dy_tile = dYp + n * sample_stride + chan_base * spatial; | |
| 189 | 124 | float* dx_tile = dXp + n * sample_stride + chan_base * spatial; | |
| 190 | |||
| 191 | // Pass 1: mean, var, rstd over the tile. Variance from squared | ||
| 192 | // deviations, for the cancellation reason spelled out in | ||
| 193 | // group_norm_forward above. | ||
| 194 | 124 | float sum = 0.0f; | |
| 195 |
2/2✓ Branch 0 taken 11832 times.
✓ Branch 1 taken 124 times.
|
11956 | for (int i = 0; i < tile_size; ++i) sum += x_tile[i]; |
| 196 | 124 | const float mean = sum * inv_M; | |
| 197 | |||
| 198 | 124 | float sumsq = 0.0f; | |
| 199 |
2/2✓ Branch 0 taken 11832 times.
✓ Branch 1 taken 124 times.
|
11956 | for (int i = 0; i < tile_size; ++i) { |
| 200 | 11832 | const float d = x_tile[i] - mean; | |
| 201 | 11832 | sumsq += d * d; | |
| 202 | 11832 | } | |
| 203 | 124 | const float var = sumsq * inv_M; | |
| 204 | 124 | const float rstd = 1.0f / std::sqrt(var + eps); | |
| 205 | |||
| 206 | // Pass 2: sum1 = Σ dx̂, sum2 = Σ dx̂·x̂; stash this n's per-channel | ||
| 207 | // dg/db into the partial buffers; cache x̂ into xhat_buf for | ||
| 208 | // reuse in pass 3. Nested (channel, spatial) loops keep | ||
| 209 | // local_c/channel loop-invariant per channel instead of | ||
| 210 | // deriving it via division per element. | ||
| 211 | 124 | float sum1 = 0.0f, sum2 = 0.0f; | |
| 212 |
2/2✓ Branch 0 taken 392 times.
✓ Branch 1 taken 124 times.
|
516 | for (int local_c = 0; local_c < channels_per_group; ++local_c) { |
| 213 | 392 | const int channel = chan_base + local_c; | |
| 214 | 392 | const float gv = gp[channel]; | |
| 215 | 392 | const int base = local_c * spatial; | |
| 216 | 392 | float dg = 0.0f, db = 0.0f; | |
| 217 |
2/2✓ Branch 0 taken 11798 times.
✓ Branch 1 taken 392 times.
|
12190 | for (int s = 0; s < spatial; ++s) { |
| 218 | 11798 | const int i = base + s; | |
| 219 | 11798 | const float dyv = dy_tile[i]; | |
| 220 | 11798 | const float xh = (x_tile[i] - mean) * rstd; | |
| 221 | 11798 | xhat_buf[i] = xh; | |
| 222 | 11798 | const float dxh = dyv * gv; | |
| 223 | 11798 | sum1 += dxh; | |
| 224 | 11798 | sum2 += dxh * xh; | |
| 225 | 11798 | dg += dyv * xh; | |
| 226 | 11798 | db += dyv; | |
| 227 | 11798 | } | |
| 228 | 392 | dg_row[channel] = dg; // disjoint per-n write | |
| 229 | 392 | db_row[channel] = db; // disjoint per-n write | |
| 230 | 392 | } | |
| 231 | |||
| 232 | // Pass 3: dX = rstd * (dx̂ - (sum1 + x̂*sum2) / M). | ||
| 233 |
2/2✓ Branch 0 taken 392 times.
✓ Branch 1 taken 124 times.
|
516 | for (int local_c = 0; local_c < channels_per_group; ++local_c) { |
| 234 | 392 | const int channel = chan_base + local_c; | |
| 235 | 392 | const float gv = gp[channel]; | |
| 236 | 392 | const int base = local_c * spatial; | |
| 237 |
2/2✓ Branch 0 taken 11829 times.
✓ Branch 1 taken 392 times.
|
12221 | for (int s = 0; s < spatial; ++s) { |
| 238 | 11829 | const int i = base + s; | |
| 239 | 11829 | const float dyv = dy_tile[i]; | |
| 240 | 11829 | const float xh = xhat_buf[i]; | |
| 241 | 11829 | const float dxh = dyv * gv; | |
| 242 | 11829 | dx_tile[i] = rstd * (dxh - (sum1 + xh * sum2) * inv_M); // overwrite | |
| 243 | 11829 | } | |
| 244 | 392 | } | |
| 245 | 124 | } | |
| 246 | 33 | }); | |
| 247 | |||
| 248 | // Single-threaded: sum the per-n partials down into dGamma/dBeta. | ||
| 249 |
2/2✓ Branch 0 taken 264 times.
✓ Branch 1 taken 25 times.
|
289 | for (int c = 0; c < C; ++c) { |
| 250 | 264 | float dg = 0.0f, db = 0.0f; | |
| 251 |
2/2✓ Branch 0 taken 392 times.
✓ Branch 1 taken 264 times.
|
656 | for (int n = 0; n < N; ++n) { |
| 252 | 392 | dg += dg_partial[static_cast<std::size_t>(n) * C + c]; | |
| 253 | 392 | db += db_partial[static_cast<std::size_t>(n) * C + c]; | |
| 254 | 392 | } | |
| 255 | 264 | dGp[c] += dg; // accumulate | |
| 256 | 264 | dBp[c] += db; // accumulate | |
| 257 | 264 | } | |
| 258 | 25 | } | |
| 259 | |||
| 260 | } // namespace brotensor::detail::cpu | ||
| 261 |