meganeura 0.2.0

E-graph optimized neural network training on Blade
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

// Two entry points:
//   rms_norm_grad_w: gradient wrt weight, dispatch [ceil(cols/256), 1, 1]
//   rms_norm_grad_x: gradient wrt input, dispatch [rows, 1, 1]
// Params field names: m=rows, n=cols, k=eps_bits

struct Params {
    m: u32,
    n: u32,
    k: u32,
    _pad: u32,
}

var<storage> src_a: array<f32>;  // dy
var<storage> src_b: array<f32>;  // x
var<storage> bias: array<f32>;   // w (weight)
var<storage, read_write> dst: array<f32>;
var<uniform> params: Params;
var<workgroup> wg_data: array<f32, 256>;

// grad_w[col] = sum_i( dy[i*n+col] * x[i*n+col] * rsqrt_i )
//
// Dispatch: [ceil(cols/256), 1, 1]
// Each thread handles one column. Workgroup cooperatively precomputes
// rsqrt for each row via shared-memory reduction (O(rows × cols/256) per
// thread), then accumulates grad_w in O(rows) per thread.
// Total work: O(rows × cols) — down from the naive O(rows × cols²).
@compute @workgroup_size(256)
fn rms_norm_grad_w(@builtin(workgroup_id) wgid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let col = wgid.x * 256u + lid.x;
    let tid = lid.x;
    let rows = params.m;
    let cols = params.n;
    let eps = bitcast<f32>(params.k);

    // Process one row at a time: cooperative rsqrt then per-column accumulation
    var acc = 0.0;
    for (var i = 0u; i < rows; i++) {
        let offset = i * cols;

        // Cooperative rsqrt: each thread sums a stride of the row
        var ss = 0.0;
        var j = tid;
        loop {
            if j >= cols { break; }
            let v = src_b[offset + j];
            ss += v * v;
            j += 256u;
        }
        wg_data[tid] = ss;
        workgroupBarrier();

        var stride = 128u;
        loop {
            if stride == 0u { break; }
            if tid < stride {
                wg_data[tid] += wg_data[tid + stride];
            }
            workgroupBarrier();
            stride >>= 1u;
        }
        let rsqrt_i = inverseSqrt(wg_data[0] / f32(cols) + eps);

        // Accumulate this row's contribution
        if col < cols {
            acc += src_a[offset + col] * src_b[offset + col] * rsqrt_i;
        }
        workgroupBarrier();
    }

    if col < cols {
        dst[col] = acc;
    }
}

// grad_x[i,j] = rsqrt_i * (dy[i,j]*w[j] - x[i,j] * s_i)
// where s_i = (rsqrt_i^2 / cols) * sum_j(dy[i,j]*w[j]*x[i,j])
@compute @workgroup_size(256)
fn rms_norm_grad_x(@builtin(workgroup_id) wgid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let row = wgid.x;
    let tid = lid.x;
    let rows = params.m;
    let cols = params.n;
    let eps = bitcast<f32>(params.k);
    if row >= rows { return; }
    let offset = row * cols;

    // Phase 1: Compute rsqrt via shared memory reduction
    var ss = 0.0;
    var j = tid;
    loop {
        if j >= cols { break; }
        let v = src_b[offset + j];
        ss += v * v;
        j += 256u;
    }
    wg_data[tid] = ss;
    workgroupBarrier();

    var stride = 128u;
    loop {
        if stride == 0u { break; }
        if tid < stride {
            wg_data[tid] += wg_data[tid + stride];
        }
        workgroupBarrier();
        stride >>= 1u;
    }
    let rsqrt_val = inverseSqrt(wg_data[0] / f32(cols) + eps);

    // Phase 2: Compute s_i = (rsqrt^2 / cols) * sum_j(dy[j]*w[j]*x[j])
    var dot = 0.0;
    j = tid;
    loop {
        if j >= cols { break; }
        dot += src_a[offset + j] * bias[j] * src_b[offset + j];
        j += 256u;
    }
    wg_data[tid] = dot;
    workgroupBarrier();

    stride = 128u;
    loop {
        if stride == 0u { break; }
        if tid < stride {
            wg_data[tid] += wg_data[tid + stride];
        }
        workgroupBarrier();
        stride >>= 1u;
    }
    let s_i = (rsqrt_val * rsqrt_val / f32(cols)) * wg_data[0];

    // Phase 3: Write output
    j = tid;
    loop {
        if j >= cols { break; }
        let dy = src_a[offset + j];
        let w = bias[j];
        let x = src_b[offset + j];
        dst[offset + j] = rsqrt_val * (dy * w - x * s_i);
        j += 256u;
    }
}