1/* Layer normalization of a row (mesh-ref gpt2.rs `row_reduce`, `normalize_row`). */
2#include "mesh/ref/gpt2_norm.h"
3
4#include "mesh/ftz.h"
5#include "mesh/ref/block.h"
6#include "mesh/ref/gpt2_scalar.h"
7#include "mesh/ref/gpt2_shape.h"
8
9/* Thread t's sequential sum (from +0) over i = t, t + NT, .., then the block tree. sq: of (x - mean)^2. */
10static float row_reduce(const float *x, int sq, float mean) {
11    float parts[MESH_GPT2_NT];
12    for (int t = 0; t < MESH_GPT2_NT; t++) {
13        float s = 0.0f;
14        for (int i = t; i < MESH_GPT2_WIDTH; i += MESH_GPT2_NT) {
15            float f = x[i];
16            if (sq) {
17                float d = mesh_ftz_sub(x[i], mean);
18                f = mesh_ftz_mul(d, d);
19            }
20            s = mesh_ftz_add(s, f);
21        }
22        parts[t] = s;
23    }
24    return mesh_block_reduce_f32(parts, MESH_GPT2_NT, mesh_block_sum_f32_id(), mesh_block_add_f32);
25}
26
27void mesh_gpt2_normalize_row(const float *x, const float *g, const float *b, float *out) {
28    float mean = mesh_ftz_div(row_reduce(x, 0, 0.0f), (float)MESH_GPT2_WIDTH);
29    float var = mesh_ftz_div(row_reduce(x, 1, mean), (float)MESH_GPT2_WIDTH);
30    float inv = mesh_ftz_div(1.0f, mesh_gpt2_sqrt_newton(mesh_ftz_add(var, mesh_ftz_from_bits(MESH_GPT2_LN_EPS_BITS))));
31    for (int i = 0; i < MESH_GPT2_WIDTH; i++)
32        out[i] = mesh_ftz_add(mesh_ftz_mul(mesh_ftz_mul(mesh_ftz_sub(x[i], mean), inv), g[i]), b[i]);
33}