teeny_kernels/nn/norm/
instancenorm.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//! InstanceNorm Triton kernels.
18//!
19//! InstanceNorm normalises over the spatial dimensions (L) independently per
20//! sample (n) and per channel (c):
21//!
22//!   y[n,c,l] = (x[n,c,l] - mean[n,c]) / sqrt(var[n,c] + eps) * γ[c] + β[c]
23//!
24//! Input shape: `[N, C, L]` — N batch, C channels, L spatial elements.
25//! Grid: `[N * C]` — one CTA per (sample, channel) pair.
26//! The CTA index encodes the pair as `cta = n * C + c`.
27
28#![allow(non_snake_case)]
29
30use teeny_core::dtype::Float;
31use teeny_macros::kernel;
32use teeny_triton::triton::{
33    types::{AddOffsets, Comparison},
34    *,
35};
36
37// ─── Inference ───────────────────────────────────────────────────────────────
38
39/// InstanceNorm forward (inference — no running stats).
40///
41/// Grid: `[N * C]` — one CTA per (sample, channel).
42#[kernel]
43pub fn instance_norm_forward_inference<T: Triton, D: Float, const BLOCK_L: i32>(
44    x_ptr: T::Pointer<D>,
45    y_ptr: T::Pointer<D>,
46    weight_ptr: T::Pointer<D>,
47    bias_ptr: T::Pointer<D>,
48    _N: i32,
49    C: i32,
50    L: i32,
51    eps: f32,
52) where
53    T::I32Tensor: types::Tensor<i32, 1>,
54    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
55    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
56{
57    let pid = T::program_id(Axis::X);
58    let n = pid / C;
59    let c = pid - n * C;
60    let row_start = (n * C + c) * L;
61
62    let c_idx = T::arange(0, 1) + c;
63    let zeros = T::zeros::<D>(&[BLOCK_L]);
64    let zero_1 = T::zeros::<D>(&[1]);
65    let l_inv = T::cast::<f32, D>(T::full::<f32>(&[1], 1.0f32 / (L as f32)), None, false);
66
67    // ── Pass 1: mean ─────────────────────────────────────────────────────────
68    let mut sum = zero_1;
69    let mut l_start: i32 = 0;
70    while l_start < L {
71        let col_offs = T::arange(0, BLOCK_L) + l_start;
72        let mask = col_offs.lt(L);
73        let x_tile = T::load(
74            x_ptr.add_offsets(col_offs + row_start),
75            Some(mask),
76            Some(zeros),
77            &[],
78            None,
79            None,
80            None,
81            false,
82        );
83        sum = sum + T::sum(x_tile, None, true);
84        l_start += BLOCK_L;
85    }
86    let mean_1 = sum * l_inv;
87    let mean = T::broadcast_to(mean_1, &[BLOCK_L]);
88
89    // ── Pass 2: variance ─────────────────────────────────────────────────────
90    let mut var_sum = zero_1;
91    l_start = 0;
92    while l_start < L {
93        let col_offs = T::arange(0, BLOCK_L) + l_start;
94        let mask = col_offs.lt(L);
95        let x_tile = T::load(
96            x_ptr.add_offsets(col_offs + row_start),
97            Some(mask),
98            Some(zeros),
99            &[],
100            None,
101            None,
102            None,
103            false,
104        );
105        // Mask the diff so out-of-bounds positions don't contribute mean^2 to variance.
106        let diff = T::where_::<D>(mask, x_tile - mean, zeros);
107        var_sum = var_sum + T::sum(diff * diff, None, true);
108        l_start += BLOCK_L;
109    }
110    let eps_t = T::cast::<f32, D>(T::full::<f32>(&[1], eps), None, false);
111    let rstd = T::broadcast_to(T::rsqrt(var_sum * l_inv + eps_t), &[BLOCK_L]);
112
113    let gamma = T::broadcast_to(
114        T::load(
115            weight_ptr.add_offsets(c_idx),
116            None,
117            None,
118            &[],
119            None,
120            None,
121            None,
122            false,
123        ),
124        &[BLOCK_L],
125    );
126    let beta = T::broadcast_to(
127        T::load(
128            bias_ptr.add_offsets(c_idx),
129            None,
130            None,
131            &[],
132            None,
133            None,
134            None,
135            false,
136        ),
137        &[BLOCK_L],
138    );
139
140    // ── Pass 3: normalise ─────────────────────────────────────────────────────
141    l_start = 0;
142    while l_start < L {
143        let col_offs = T::arange(0, BLOCK_L) + l_start;
144        let mask = col_offs.lt(L);
145        let x_tile = T::load(
146            x_ptr.add_offsets(col_offs + row_start),
147            Some(mask),
148            Some(zeros),
149            &[],
150            None,
151            None,
152            None,
153            false,
154        );
155        let y_tile = (x_tile - mean) * rstd * gamma + beta;
156        T::store(
157            y_ptr.add_offsets(col_offs + row_start),
158            y_tile,
159            Some(mask),
160            &[],
161            None,
162            None,
163        );
164        l_start += BLOCK_L;
165    }
166}
167
168// ─── Training forward ─────────────────────────────────────────────────────────
169
170/// InstanceNorm training forward — saves per-(n,c) mean and rstd.
171///
172/// Grid: `[N * C]` — one CTA per (sample, channel).
173#[cfg(feature = "training")]
174#[kernel]
175pub fn instance_norm_forward<T: Triton, D: Float, const BLOCK_L: i32>(
176    x_ptr: T::Pointer<D>,
177    y_ptr: T::Pointer<D>,
178    weight_ptr: T::Pointer<D>,
179    bias_ptr: T::Pointer<D>,
180    mean_ptr: T::Pointer<D>,
181    rstd_ptr: T::Pointer<D>,
182    _N: i32,
183    C: i32,
184    L: i32,
185    eps: f32,
186) where
187    T::I32Tensor: types::Tensor<i32, 1>,
188    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
189    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
190{
191    let pid = T::program_id(Axis::X);
192    let n = pid / C;
193    let c = pid - n * C;
194    let row_start = (n * C + c) * L;
195    let stat_idx = T::arange(0, 1) + pid;
196    let c_idx = T::arange(0, 1) + c;
197
198    let zeros = T::zeros::<D>(&[BLOCK_L]);
199    let zero_1 = T::zeros::<D>(&[1]);
200    let l_inv = T::cast::<f32, D>(T::full::<f32>(&[1], 1.0f32 / (L as f32)), None, false);
201
202    // ── Pass 1: mean ─────────────────────────────────────────────────────────
203    let mut sum = zero_1;
204    let mut l_start: i32 = 0;
205    while l_start < L {
206        let col_offs = T::arange(0, BLOCK_L) + l_start;
207        let mask = col_offs.lt(L);
208        let x_tile = T::load(
209            x_ptr.add_offsets(col_offs + row_start),
210            Some(mask),
211            Some(zeros),
212            &[],
213            None,
214            None,
215            None,
216            false,
217        );
218        sum = sum + T::sum(x_tile, None, true);
219        l_start += BLOCK_L;
220    }
221    let mean_1 = sum * l_inv;
222    let mean = T::broadcast_to(mean_1, &[BLOCK_L]);
223
224    // ── Pass 2: variance ─────────────────────────────────────────────────────
225    let mut var_sum = zero_1;
226    l_start = 0;
227    while l_start < L {
228        let col_offs = T::arange(0, BLOCK_L) + l_start;
229        let mask = col_offs.lt(L);
230        let x_tile = T::load(
231            x_ptr.add_offsets(col_offs + row_start),
232            Some(mask),
233            Some(zeros),
234            &[],
235            None,
236            None,
237            None,
238            false,
239        );
240        // Mask the diff so out-of-bounds positions don't contribute mean^2 to variance.
241        let diff = T::where_::<D>(mask, x_tile - mean, zeros);
242        var_sum = var_sum + T::sum(diff * diff, None, true);
243        l_start += BLOCK_L;
244    }
245    let eps_t = T::cast::<f32, D>(T::full::<f32>(&[1], eps), None, false);
246    let rstd_1 = T::rsqrt(var_sum * l_inv + eps_t);
247    let rstd = T::broadcast_to(rstd_1, &[BLOCK_L]);
248
249    T::store(
250        mean_ptr.add_offsets(stat_idx),
251        mean_1,
252        None,
253        &[],
254        None,
255        None,
256    );
257    T::store(
258        rstd_ptr.add_offsets(stat_idx),
259        rstd_1,
260        None,
261        &[],
262        None,
263        None,
264    );
265
266    let gamma = T::broadcast_to(
267        T::load(
268            weight_ptr.add_offsets(c_idx),
269            None,
270            None,
271            &[],
272            None,
273            None,
274            None,
275            false,
276        ),
277        &[BLOCK_L],
278    );
279    let beta = T::broadcast_to(
280        T::load(
281            bias_ptr.add_offsets(c_idx),
282            None,
283            None,
284            &[],
285            None,
286            None,
287            None,
288            false,
289        ),
290        &[BLOCK_L],
291    );
292
293    // ── Pass 3: normalise ─────────────────────────────────────────────────────
294    l_start = 0;
295    while l_start < L {
296        let col_offs = T::arange(0, BLOCK_L) + l_start;
297        let mask = col_offs.lt(L);
298        let x_tile = T::load(
299            x_ptr.add_offsets(col_offs + row_start),
300            Some(mask),
301            Some(zeros),
302            &[],
303            None,
304            None,
305            None,
306            false,
307        );
308        let y_tile = (x_tile - mean) * rstd * gamma + beta;
309        T::store(
310            y_ptr.add_offsets(col_offs + row_start),
311            y_tile,
312            Some(mask),
313            &[],
314            None,
315            None,
316        );
317        l_start += BLOCK_L;
318    }
319}
320
321// ─── Training backward ───────────────────────────────────────────────────────
322
323/// InstanceNorm backward pass.
324///
325/// Grid: `[N * C]` — one CTA per (sample, channel).
326#[cfg(feature = "training")]
327#[kernel]
328pub fn instance_norm_backward<T: Triton, D: Float, const BLOCK_L: i32>(
329    dy_ptr: T::Pointer<D>,
330    x_ptr: T::Pointer<D>,
331    dx_ptr: T::Pointer<D>,
332    weight_ptr: T::Pointer<D>,
333    dweight_ptr: T::Pointer<D>,
334    dbias_ptr: T::Pointer<D>,
335    mean_ptr: T::Pointer<D>,
336    rstd_ptr: T::Pointer<D>,
337    _N: i32,
338    C: i32,
339    L: i32,
340) where
341    T::I32Tensor: types::Tensor<i32, 1>,
342    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
343    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
344{
345    let pid = T::program_id(Axis::X);
346    let n = pid / C;
347    let c = pid - n * C;
348    let row_start = (n * C + c) * L;
349    let stat_idx = T::arange(0, 1) + pid;
350    let c_idx = T::arange(0, 1) + c;
351
352    let zeros = T::zeros::<D>(&[BLOCK_L]);
353    let zero_1 = T::zeros::<D>(&[1]);
354    let l_inv = T::cast::<f32, D>(T::full::<f32>(&[1], 1.0f32 / (L as f32)), None, false);
355
356    let rstd_1 = T::load(
357        rstd_ptr.add_offsets(stat_idx),
358        None,
359        None,
360        &[],
361        None,
362        None,
363        None,
364        false,
365    );
366    let mean_1 = T::load(
367        mean_ptr.add_offsets(stat_idx),
368        None,
369        None,
370        &[],
371        None,
372        None,
373        None,
374        false,
375    );
376    let rstd = T::broadcast_to(rstd_1, &[BLOCK_L]);
377    let mean = T::broadcast_to(mean_1, &[BLOCK_L]);
378
379    let gamma = T::broadcast_to(
380        T::load(
381            weight_ptr.add_offsets(c_idx),
382            None,
383            None,
384            &[],
385            None,
386            None,
387            None,
388            false,
389        ),
390        &[BLOCK_L],
391    );
392
393    // ── Pass 1: accumulate row dot products ───────────────────────────────────
394    let mut sum_dy_gamma = zero_1;
395    let mut sum_dy_gamma_xhat = zero_1;
396    let mut l_start: i32 = 0;
397    while l_start < L {
398        let col_offs = T::arange(0, BLOCK_L) + l_start;
399        let mask = col_offs.lt(L);
400        let x_tile = T::load(
401            x_ptr.add_offsets(col_offs + row_start),
402            Some(mask),
403            Some(zeros),
404            &[],
405            None,
406            None,
407            None,
408            false,
409        );
410        let dy_tile = T::load(
411            dy_ptr.add_offsets(col_offs + row_start),
412            Some(mask),
413            Some(zeros),
414            &[],
415            None,
416            None,
417            None,
418            false,
419        );
420        let xhat = (x_tile - mean) * rstd;
421        sum_dy_gamma = sum_dy_gamma + T::sum(dy_tile * gamma, None, true);
422        sum_dy_gamma_xhat = sum_dy_gamma_xhat + T::sum(dy_tile * gamma * xhat, None, true);
423        l_start += BLOCK_L;
424    }
425    let c1 = T::broadcast_to(sum_dy_gamma * l_inv, &[BLOCK_L]);
426    let c2 = T::broadcast_to(sum_dy_gamma_xhat * l_inv, &[BLOCK_L]);
427
428    // ── Pass 2: dx and dweight / dbias ───────────────────────────────────────
429    l_start = 0;
430    while l_start < L {
431        let col_offs = T::arange(0, BLOCK_L) + l_start;
432        let mask = col_offs.lt(L);
433        let x_tile = T::load(
434            x_ptr.add_offsets(col_offs + row_start),
435            Some(mask),
436            Some(zeros),
437            &[],
438            None,
439            None,
440            None,
441            false,
442        );
443        let dy_tile = T::load(
444            dy_ptr.add_offsets(col_offs + row_start),
445            Some(mask),
446            Some(zeros),
447            &[],
448            None,
449            None,
450            None,
451            false,
452        );
453        let dw_old = T::load(
454            dweight_ptr.add_offsets(c_idx),
455            None,
456            None,
457            &[],
458            None,
459            None,
460            None,
461            false,
462        );
463        let db_old = T::load(
464            dbias_ptr.add_offsets(c_idx),
465            None,
466            None,
467            &[],
468            None,
469            None,
470            None,
471            false,
472        );
473
474        let xhat = (x_tile - mean) * rstd;
475        let dx_tile = rstd * gamma * (dy_tile - c1 - xhat * c2);
476
477        T::store(
478            dx_ptr.add_offsets(col_offs + row_start),
479            dx_tile,
480            Some(mask),
481            &[],
482            None,
483            None,
484        );
485        T::store(
486            dweight_ptr.add_offsets(c_idx),
487            dw_old + T::sum(dy_tile * xhat, None, true),
488            None,
489            &[],
490            None,
491            None,
492        );
493        T::store(
494            dbias_ptr.add_offsets(c_idx),
495            db_old + T::sum(dy_tile, None, true),
496            None,
497            &[],
498            None,
499            None,
500        );
501        l_start += BLOCK_L;
502    }
503}
teeny_kernels/nn/norm/instancenorm.rs

teeny_kernels/nn/norm/
instancenorm.rs
