teeny_kernels/nn/norm/
groupnorm.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! GroupNorm Triton kernels.
18//!
19//! GroupNorm partitions C channels into G groups and normalises over each
20//! group independently, per sample:
21//!
22//!   y[n,c,l] = (x[n,c,l] - mean[n,g]) / sqrt(var[n,g] + eps) * γ[c] + β[c]
23//!
24//! where g = c / (C / G) is the group index.
25//!
26//! Input shape: `[N, C, L]`. Grid: `[N * G]` — one CTA per (sample, group).
27//! Each CTA covers channels `[g*(C/G), (g+1)*(C/G))` × all L elements,
28//! so the normalised tile has `(C/G) * L` elements.
29//!
30//! BLOCK_NL must be >= (C/G) * L and a power of two.
31
32#![allow(non_snake_case)]
33
34use teeny_core::dtype::Float;
35use teeny_macros::kernel;
36use teeny_triton::triton::{
37    types::{AddOffsets, Comparison},
38    *,
39};
40
41// ─── Inference ───────────────────────────────────────────────────────────────
42
43/// GroupNorm forward (inference).
44///
45/// Grid: `[N * G]` — one CTA per (sample, group).
46#[kernel]
47pub fn group_norm_forward_inference<T: Triton, D: Float, const BLOCK_NL: i32>(
48    x_ptr: T::Pointer<D>,
49    y_ptr: T::Pointer<D>,
50    weight_ptr: T::Pointer<D>,
51    bias_ptr: T::Pointer<D>,
52    _N: i32,
53    C: i32,
54    L: i32,
55    G: i32,
56    eps: f32,
57) where
58    T::I32Tensor: types::Tensor<i32, 1>,
59    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
60    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
61{
62    let pid = T::program_id(Axis::X);
63    let n = pid / G;
64    let g = pid - n * G;
65    let channels_per_group = C / G;
66
67    // Flat element count for this (n, g) group: channels_per_group * L.
68    let group_size = channels_per_group * L;
69    // Offset to the first element of this group in the flat [N, C, L] layout.
70    let group_start = n * C * L + g * channels_per_group * L;
71
72    let zeros = T::zeros::<D>(&[BLOCK_NL]);
73    let zero_1 = T::zeros::<D>(&[1]);
74    let gs_inv = T::cast::<f32, D>(
75        T::full::<f32>(&[1], 1.0f32 / (group_size as f32)),
76        None,
77        false,
78    );
79
80    // ── Pass 1: mean ─────────────────────────────────────────────────────────
81    let mut sum = zero_1;
82    let mut t: i32 = 0;
83    while t < group_size {
84        let offs = T::arange(0, BLOCK_NL) + t;
85        let mask = offs.lt(group_size);
86        let x_tile = T::load(
87            x_ptr.add_offsets(offs + group_start),
88            Some(mask),
89            Some(zeros),
90            &[],
91            None,
92            None,
93            None,
94            false,
95        );
96        sum = sum + T::sum(x_tile, None, true);
97        t += BLOCK_NL;
98    }
99    let mean_1 = sum * gs_inv;
100    let mean = T::broadcast_to(mean_1, &[BLOCK_NL]);
101
102    // ── Pass 2: variance ─────────────────────────────────────────────────────
103    let mut var_sum = zero_1;
104    t = 0;
105    while t < group_size {
106        let offs = T::arange(0, BLOCK_NL) + t;
107        let mask = offs.lt(group_size);
108        let x_tile = T::load(
109            x_ptr.add_offsets(offs + group_start),
110            Some(mask),
111            Some(zeros),
112            &[],
113            None,
114            None,
115            None,
116            false,
117        );
118        // Mask the diff so out-of-bounds positions (x_tile=0, mean=true_mean)
119        // don't contribute (0-mean)^2 = mean^2 to the variance sum.
120        let diff = T::where_::<D>(mask, x_tile - mean, zeros);
121        var_sum = var_sum + T::sum(diff * diff, None, true);
122        t += BLOCK_NL;
123    }
124    let eps_t = T::cast::<f32, D>(T::full::<f32>(&[1], eps), None, false);
125    let rstd = T::broadcast_to(T::rsqrt(var_sum * gs_inv + eps_t), &[BLOCK_NL]);
126
127    // ── Pass 3: normalise with per-channel affine ─────────────────────────────
128    // Element at group-flat index (t+i) maps to channel g*cpg + (t+i)/L.
129    t = 0;
130    while t < group_size {
131        let offs = T::arange(0, BLOCK_NL) + t;
132        let mask = offs.lt(group_size);
133        let chan_offs = (T::arange(0, BLOCK_NL) + t) / L + g * channels_per_group;
134        let x_tile = T::load(
135            x_ptr.add_offsets(offs + group_start),
136            Some(mask),
137            Some(zeros),
138            &[],
139            None,
140            None,
141            None,
142            false,
143        );
144        let gamma = T::load(
145            weight_ptr.add_offsets(chan_offs),
146            Some(mask),
147            Some(zeros),
148            &[],
149            None,
150            None,
151            None,
152            false,
153        );
154        let beta = T::load(
155            bias_ptr.add_offsets(chan_offs),
156            Some(mask),
157            Some(zeros),
158            &[],
159            None,
160            None,
161            None,
162            false,
163        );
164        let y_tile = (x_tile - mean) * rstd * gamma + beta;
165        T::store(
166            y_ptr.add_offsets(offs + group_start),
167            y_tile,
168            Some(mask),
169            &[],
170            None,
171            None,
172        );
173        t += BLOCK_NL;
174    }
175}
176
177// ─── Training forward ─────────────────────────────────────────────────────────
178
179/// GroupNorm training forward — saves per-(n,g) mean and rstd.
180///
181/// Grid: `[N * G]` — one CTA per (sample, group).
182#[cfg(feature = "training")]
183#[kernel]
184pub fn group_norm_forward<T: Triton, D: Float, const BLOCK_NL: i32>(
185    x_ptr: T::Pointer<D>,
186    y_ptr: T::Pointer<D>,
187    weight_ptr: T::Pointer<D>,
188    bias_ptr: T::Pointer<D>,
189    mean_ptr: T::Pointer<D>,
190    rstd_ptr: T::Pointer<D>,
191    _N: i32,
192    C: i32,
193    L: i32,
194    G: i32,
195    eps: f32,
196) where
197    T::I32Tensor: types::Tensor<i32, 1>,
198    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
199    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
200{
201    let pid = T::program_id(Axis::X);
202    let n = pid / G;
203    let g = pid - n * G;
204    let channels_per_group = C / G;
205    let group_size = channels_per_group * L;
206    let group_start = n * C * L + g * channels_per_group * L;
207    let stat_idx = T::arange(0, 1) + pid;
208
209    let zeros = T::zeros::<D>(&[BLOCK_NL]);
210    let zero_1 = T::zeros::<D>(&[1]);
211    let gs_inv = T::cast::<f32, D>(
212        T::full::<f32>(&[1], 1.0f32 / (group_size as f32)),
213        None,
214        false,
215    );
216
217    // ── Pass 1: mean ─────────────────────────────────────────────────────────
218    let mut sum = zero_1;
219    let mut t: i32 = 0;
220    while t < group_size {
221        let offs = T::arange(0, BLOCK_NL) + t;
222        let mask = offs.lt(group_size);
223        let x_tile = T::load(
224            x_ptr.add_offsets(offs + group_start),
225            Some(mask),
226            Some(zeros),
227            &[],
228            None,
229            None,
230            None,
231            false,
232        );
233        sum = sum + T::sum(x_tile, None, true);
234        t += BLOCK_NL;
235    }
236    let mean_1 = sum * gs_inv;
237    let mean = T::broadcast_to(mean_1, &[BLOCK_NL]);
238
239    // ── Pass 2: variance ─────────────────────────────────────────────────────
240    let mut var_sum = zero_1;
241    t = 0;
242    while t < group_size {
243        let offs = T::arange(0, BLOCK_NL) + t;
244        let mask = offs.lt(group_size);
245        let x_tile = T::load(
246            x_ptr.add_offsets(offs + group_start),
247            Some(mask),
248            Some(zeros),
249            &[],
250            None,
251            None,
252            None,
253            false,
254        );
255        // Mask the diff so out-of-bounds positions (x_tile=0, mean=true_mean)
256        // don't contribute (0-mean)^2 = mean^2 to the variance sum.
257        let diff = T::where_::<D>(mask, x_tile - mean, zeros);
258        var_sum = var_sum + T::sum(diff * diff, None, true);
259        t += BLOCK_NL;
260    }
261    let eps_t = T::cast::<f32, D>(T::full::<f32>(&[1], eps), None, false);
262    let rstd_1 = T::rsqrt(var_sum * gs_inv + eps_t);
263    let rstd = T::broadcast_to(rstd_1, &[BLOCK_NL]);
264
265    T::store(
266        mean_ptr.add_offsets(stat_idx),
267        mean_1,
268        None,
269        &[],
270        None,
271        None,
272    );
273    T::store(
274        rstd_ptr.add_offsets(stat_idx),
275        rstd_1,
276        None,
277        &[],
278        None,
279        None,
280    );
281
282    // ── Pass 3: normalise ─────────────────────────────────────────────────────
283    t = 0;
284    while t < group_size {
285        let offs = T::arange(0, BLOCK_NL) + t;
286        let mask = offs.lt(group_size);
287        let chan_offs = (T::arange(0, BLOCK_NL) + t) / L + g * channels_per_group;
288        let x_tile = T::load(
289            x_ptr.add_offsets(offs + group_start),
290            Some(mask),
291            Some(zeros),
292            &[],
293            None,
294            None,
295            None,
296            false,
297        );
298        let gamma = T::load(
299            weight_ptr.add_offsets(chan_offs),
300            Some(mask),
301            Some(zeros),
302            &[],
303            None,
304            None,
305            None,
306            false,
307        );
308        let beta = T::load(
309            bias_ptr.add_offsets(chan_offs),
310            Some(mask),
311            Some(zeros),
312            &[],
313            None,
314            None,
315            None,
316            false,
317        );
318        let y_tile = (x_tile - mean) * rstd * gamma + beta;
319        T::store(
320            y_ptr.add_offsets(offs + group_start),
321            y_tile,
322            Some(mask),
323            &[],
324            None,
325            None,
326        );
327        t += BLOCK_NL;
328    }
329}
330
331// ─── Training backward ───────────────────────────────────────────────────────
332
333/// GroupNorm backward pass.
334///
335/// Grid: `[N * G]` — one CTA per (sample, group).
336#[cfg(feature = "training")]
337#[kernel]
338pub fn group_norm_backward<T: Triton, D: Float, const BLOCK_NL: i32>(
339    dy_ptr: T::Pointer<D>,
340    x_ptr: T::Pointer<D>,
341    dx_ptr: T::Pointer<D>,
342    weight_ptr: T::Pointer<D>,
343    dweight_ptr: T::Pointer<D>,
344    dbias_ptr: T::Pointer<D>,
345    mean_ptr: T::Pointer<D>,
346    rstd_ptr: T::Pointer<D>,
347    _N: i32,
348    C: i32,
349    L: i32,
350    G: i32,
351) where
352    T::I32Tensor: types::Tensor<i32, 1>,
353    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
354    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
355{
356    let pid = T::program_id(Axis::X);
357    let n = pid / G;
358    let g = pid - n * G;
359    let channels_per_group = C / G;
360    let group_size = channels_per_group * L;
361    let group_start = n * C * L + g * channels_per_group * L;
362    let stat_idx = T::arange(0, 1) + pid;
363
364    let zeros = T::zeros::<D>(&[BLOCK_NL]);
365    let zero_1 = T::zeros::<D>(&[1]);
366    let gs_inv = T::cast::<f32, D>(
367        T::full::<f32>(&[1], 1.0f32 / (group_size as f32)),
368        None,
369        false,
370    );
371
372    let rstd_1 = T::load(
373        rstd_ptr.add_offsets(stat_idx),
374        None,
375        None,
376        &[],
377        None,
378        None,
379        None,
380        false,
381    );
382    let mean_1 = T::load(
383        mean_ptr.add_offsets(stat_idx),
384        None,
385        None,
386        &[],
387        None,
388        None,
389        None,
390        false,
391    );
392    let rstd = T::broadcast_to(rstd_1, &[BLOCK_NL]);
393    let mean = T::broadcast_to(mean_1, &[BLOCK_NL]);
394
395    // ── Pass 1: row-level dot products ────────────────────────────────────────
396    let mut sum_dy_gamma = zero_1;
397    let mut sum_dy_gamma_xhat = zero_1;
398    let mut t: i32 = 0;
399    while t < group_size {
400        let offs = T::arange(0, BLOCK_NL) + t;
401        let mask = offs.lt(group_size);
402        let chan_offs = (T::arange(0, BLOCK_NL) + t) / L + g * channels_per_group;
403        let x_tile = T::load(
404            x_ptr.add_offsets(offs + group_start),
405            Some(mask),
406            Some(zeros),
407            &[],
408            None,
409            None,
410            None,
411            false,
412        );
413        let dy_tile = T::load(
414            dy_ptr.add_offsets(offs + group_start),
415            Some(mask),
416            Some(zeros),
417            &[],
418            None,
419            None,
420            None,
421            false,
422        );
423        let gamma = T::load(
424            weight_ptr.add_offsets(chan_offs),
425            Some(mask),
426            Some(zeros),
427            &[],
428            None,
429            None,
430            None,
431            false,
432        );
433        let xhat = (x_tile - mean) * rstd;
434        sum_dy_gamma = sum_dy_gamma + T::sum(dy_tile * gamma, None, true);
435        sum_dy_gamma_xhat = sum_dy_gamma_xhat + T::sum(dy_tile * gamma * xhat, None, true);
436        t += BLOCK_NL;
437    }
438    let c1 = T::broadcast_to(sum_dy_gamma * gs_inv, &[BLOCK_NL]);
439    let c2 = T::broadcast_to(sum_dy_gamma_xhat * gs_inv, &[BLOCK_NL]);
440
441    // ── Pass 2: dx and dweight / dbias ───────────────────────────────────────
442    t = 0;
443    while t < group_size {
444        let offs = T::arange(0, BLOCK_NL) + t;
445        let mask = offs.lt(group_size);
446        let chan_offs = (T::arange(0, BLOCK_NL) + t) / L + g * channels_per_group;
447        let x_tile = T::load(
448            x_ptr.add_offsets(offs + group_start),
449            Some(mask),
450            Some(zeros),
451            &[],
452            None,
453            None,
454            None,
455            false,
456        );
457        let dy_tile = T::load(
458            dy_ptr.add_offsets(offs + group_start),
459            Some(mask),
460            Some(zeros),
461            &[],
462            None,
463            None,
464            None,
465            false,
466        );
467        let gamma = T::load(
468            weight_ptr.add_offsets(chan_offs),
469            Some(mask),
470            Some(zeros),
471            &[],
472            None,
473            None,
474            None,
475            false,
476        );
477        let dw_old = T::load(
478            dweight_ptr.add_offsets(chan_offs),
479            Some(mask),
480            Some(zeros),
481            &[],
482            None,
483            None,
484            None,
485            false,
486        );
487        let db_old = T::load(
488            dbias_ptr.add_offsets(chan_offs),
489            Some(mask),
490            Some(zeros),
491            &[],
492            None,
493            None,
494            None,
495            false,
496        );
497
498        let xhat = (x_tile - mean) * rstd;
499        let dx_tile = rstd * gamma * (dy_tile - c1 - xhat * c2);
500
501        T::store(
502            dx_ptr.add_offsets(offs + group_start),
503            dx_tile,
504            Some(mask),
505            &[],
506            None,
507            None,
508        );
509        T::store(
510            dweight_ptr.add_offsets(chan_offs),
511            dw_old + dy_tile * xhat,
512            Some(mask),
513            &[],
514            None,
515            None,
516        );
517        T::store(
518            dbias_ptr.add_offsets(chan_offs),
519            db_old + dy_tile,
520            Some(mask),
521            &[],
522            None,
523            None,
524        );
525        t += BLOCK_NL;
526    }
527}
teeny_kernels/nn/norm/groupnorm.rs

teeny_kernels/nn/norm/
groupnorm.rs
