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

teeny_kernels/nn/pool/
lppool3d.rs
