teeny_kernels/nn/pool/
lppool2d.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#![allow(non_snake_case)]
18
19use teeny_core::dtype::Float;
20use teeny_macros::kernel;
21use teeny_triton::triton::{
22    types::{AddOffsets, Comparison, Tensor},
23    *,
24};
25
26/// 2-D Lp-norm pooling forward pass.
27///
28/// Grid: `pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile`
29///
30/// **Constraints**: no padding; `OH = (H - KH) / STRIDE_H + 1`, `OW = (W - KW) / STRIDE_W + 1`.
31#[kernel]
32pub fn lppool2d_forward<
33    T: Triton,
34    D: Float,
35    const KH: i32,
36    const KW: i32,
37    const STRIDE_H: i32,
38    const STRIDE_W: i32,
39    const BLOCK_OW: i32,
40>(
41    input_ptr: T::Pointer<D>,
42    output_ptr: T::Pointer<D>,
43    _B: i32,
44    C: i32,
45    H: i32,
46    W: i32,
47    OH: i32,
48    OW: i32,
49    p: f32,
50) where
51    T::I32Tensor: Tensor<i32, 1>,
52    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
53    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
54{
55    let pid = T::program_id(Axis::X);
56    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
57
58    let ow_tile = pid % num_ow_tiles;
59    let bco = pid / num_ow_tiles;
60    let oh = bco % OH;
61    let bc = bco / OH;
62    let c = bc % C;
63    let b = bc / C;
64
65    let ow_start = ow_tile * BLOCK_OW;
66    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
67    let ow_mask = ow_range.lt(OW);
68
69    let in_bc_base = (b * C + c) * H * W;
70    let out_bc_base = (b * C + c) * OH * OW;
71
72    let p_vec = T::full::<f32>(&[BLOCK_OW], p);
73    let inv_p_vec = T::full::<f32>(&[BLOCK_OW], 1.0_f32 / p);
74    let eps_vec = T::full::<f32>(&[BLOCK_OW], 1e-12_f32);
75
76    let mut acc = T::zeros::<f32>(&[BLOCK_OW]);
77
78    let loop_bound = KH * KW;
79    for idx in 0..loop_bound {
80        let kw = idx % KW;
81        let kh = idx / KW;
82        let ih = oh * STRIDE_H + kh;
83        let iw_range = ow_range * STRIDE_W + kw;
84        let in_offsets = iw_range + (in_bc_base + ih * W);
85        let tile = T::load(
86            input_ptr.add_offsets(in_offsets),
87            Some(ow_mask),
88            Some(T::zeros::<D>(&[BLOCK_OW])),
89            &[],
90            None,
91            None,
92            None,
93            false,
94        );
95        let tile_f32 = T::cast::<D, f32>(tile, None, false);
96        let abs_tile = T::abs(tile_f32);
97        let safe_abs = T::maximum(abs_tile, eps_vec);
98        let pow_tile = T::exp(p_vec * T::log(safe_abs));
99        acc = acc + pow_tile;
100    }
101
102    let safe_acc = T::maximum(acc, eps_vec);
103    let result_f32 = T::exp(T::log(safe_acc) * inv_p_vec);
104    let result = T::cast::<f32, D>(result_f32, None, false);
105
106    let out_offsets = ow_range + (out_bc_base + oh * OW);
107    T::store(
108        output_ptr.add_offsets(out_offsets),
109        result,
110        Some(ow_mask),
111        &[],
112        None,
113        None,
114    );
115}
116
117/// 2-D Lp-norm pooling backward pass.
118///
119/// `dx_i = dy * sign(x_i) * (|x_i| / max(y, ε))^(p-1) / max(y, ε)`.
120///
121/// `dx` must be zero-initialised before launch.
122#[kernel]
123pub fn lppool2d_backward<
124    T: Triton,
125    D: Float,
126    const KH: i32,
127    const KW: i32,
128    const STRIDE_H: i32,
129    const STRIDE_W: i32,
130    const BLOCK_OW: i32,
131>(
132    dy_ptr: T::Pointer<D>,
133    x_ptr: T::Pointer<D>,
134    y_ptr: T::Pointer<D>,
135    dx_ptr: T::Pointer<D>,
136    _B: i32,
137    C: i32,
138    H: i32,
139    W: i32,
140    OH: i32,
141    OW: i32,
142    p: f32,
143) where
144    T::I32Tensor: Tensor<i32, 1>,
145    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
146    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
147{
148    let pid = T::program_id(Axis::X);
149    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
150
151    let ow_tile = pid % num_ow_tiles;
152    let bco = pid / num_ow_tiles;
153    let oh = bco % OH;
154    let bc = bco / OH;
155    let c = bc % C;
156    let b = bc / C;
157
158    let ow_start = ow_tile * BLOCK_OW;
159    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
160    let ow_mask = ow_range.lt(OW);
161
162    let in_bc_base = (b * C + c) * H * W;
163    let out_bc_base = (b * C + c) * OH * OW;
164
165    let pm1_vec = T::full::<f32>(&[BLOCK_OW], p - 1.0_f32);
166    let eps_vec = T::full::<f32>(&[BLOCK_OW], 1e-12_f32);
167    let zeros_f32 = T::zeros::<f32>(&[BLOCK_OW]);
168
169    let out_offsets = ow_range + (out_bc_base + oh * OW);
170    let dy_tile = T::load(
171        dy_ptr.add_offsets(out_offsets),
172        Some(ow_mask),
173        Some(T::zeros::<D>(&[BLOCK_OW])),
174        &[],
175        None,
176        None,
177        None,
178        false,
179    );
180    let dy_f32 = T::cast::<D, f32>(dy_tile, None, false);
181
182    let y_tile = T::load(
183        y_ptr.add_offsets(out_offsets),
184        Some(ow_mask),
185        Some(T::zeros::<D>(&[BLOCK_OW])),
186        &[],
187        None,
188        None,
189        None,
190        false,
191    );
192    let y_f32 = T::cast::<D, f32>(y_tile, None, false);
193    let safe_y = T::maximum(y_f32, eps_vec);
194
195    let loop_bound = KH * KW;
196    for idx in 0..loop_bound {
197        let kw = idx % KW;
198        let kh = idx / KW;
199        let ih = oh * STRIDE_H + kh;
200        let iw_range = ow_range * STRIDE_W + kw;
201        let in_offsets = iw_range + (in_bc_base + ih * W);
202        let x_tile = T::load(
203            x_ptr.add_offsets(in_offsets),
204            Some(ow_mask),
205            Some(T::zeros::<D>(&[BLOCK_OW])),
206            &[],
207            None,
208            None,
209            None,
210            false,
211        );
212        let x_f32 = T::cast::<D, f32>(x_tile, None, false);
213        let abs_x = T::abs(x_f32);
214        let safe_abs = T::maximum(abs_x, eps_vec);
215
216        let pos = T::where_(
217            T::gt(x_f32, zeros_f32),
218            T::full(&[BLOCK_OW], 1.0_f32),
219            zeros_f32,
220        );
221        let neg = T::where_(
222            T::gt(zeros_f32, x_f32),
223            T::full(&[BLOCK_OW], 1.0_f32),
224            zeros_f32,
225        );
226        let sign_x = pos - neg;
227
228        let ratio = safe_abs / safe_y;
229        let safe_ratio = T::maximum(ratio, eps_vec);
230        let pow_ratio = T::exp(pm1_vec * T::log(safe_ratio));
231
232        let dx_f32 = dy_f32 * sign_x * pow_ratio;
233        let dx_tile = T::cast::<f32, D>(dx_f32, None, false);
234
235        T::atomic_add(
236            dx_ptr.add_offsets(in_offsets),
237            dx_tile,
238            Some(ow_mask),
239            None,
240            None,
241        );
242    }
243}
244
245pub struct Lppool2dOp<'a, T: Float> {
246    pub forward: Lppool2dForward<T>,
247    pub backward: Lppool2dBackward<T>,
248    _marker: core::marker::PhantomData<&'a ()>,
249}
teeny_kernels/nn/pool/lppool2d.rs

teeny_kernels/nn/pool/
lppool2d.rs
