teeny_kernels/nn/pool/
lppool1d.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/// 1-D Lp-norm pooling forward pass.
27///
28/// `y = (Σ |x_i|^p)^(1/p)` over the kernel window.
29///
30/// `pow(|x|, p)` is computed as `exp(p * log(max(|x|, ε)))` to avoid
31/// `log(0)`. `p` is a runtime float parameter.
32///
33/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
34///
35/// **Constraints**: no padding; `OL = (L - KL) / STRIDE + 1`.
36#[kernel]
37pub fn lppool1d_forward<
38    T: Triton,
39    D: Float,
40    const KL: i32,
41    const STRIDE: i32,
42    const BLOCK_OL: i32,
43>(
44    input_ptr: T::Pointer<D>,
45    output_ptr: T::Pointer<D>,
46    _B: i32,
47    C: i32,
48    L: i32,
49    OL: i32,
50    p: f32,
51) where
52    T::I32Tensor: Tensor<i32, 1>,
53    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
54    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
55{
56    let pid = T::program_id(Axis::X);
57    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
58
59    let ol_tile = pid % num_ol_tiles;
60    let bc = pid / num_ol_tiles;
61    let c = bc % C;
62    let b = bc / C;
63
64    let ol_start = ol_tile * BLOCK_OL;
65    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
66    let ol_mask = ol_range.lt(OL);
67
68    let in_bc_base = (b * C + c) * L;
69    let out_bc_base = (b * C + c) * OL;
70
71    let p_vec = T::full::<f32>(&[BLOCK_OL], p);
72    let inv_p_vec = T::full::<f32>(&[BLOCK_OL], 1.0_f32 / p);
73    let eps_vec = T::full::<f32>(&[BLOCK_OL], 1e-12_f32);
74
75    let mut acc = T::zeros::<f32>(&[BLOCK_OL]);
76
77    let loop_bound = KL;
78    for kl in 0..loop_bound {
79        let il_range = ol_range * STRIDE + kl;
80        let in_offsets = il_range + in_bc_base;
81        let tile = T::load(
82            input_ptr.add_offsets(in_offsets),
83            Some(ol_mask),
84            Some(T::zeros::<D>(&[BLOCK_OL])),
85            &[],
86            None,
87            None,
88            None,
89            false,
90        );
91        let tile_f32 = T::cast::<D, f32>(tile, None, false);
92        let abs_tile = T::abs(tile_f32);
93        let safe_abs = T::maximum(abs_tile, eps_vec);
94        // |x|^p = exp(p * log(|x|))
95        let pow_tile = T::exp(p_vec * T::log(safe_abs));
96        acc = acc + pow_tile;
97    }
98
99    // sum^(1/p) = exp(log(sum) / p)
100    let safe_acc = T::maximum(acc, eps_vec);
101    let result_f32 = T::exp(T::log(safe_acc) * inv_p_vec);
102    let result = T::cast::<f32, D>(result_f32, None, false);
103
104    let out_offsets = ol_range + out_bc_base;
105    T::store(
106        output_ptr.add_offsets(out_offsets),
107        result,
108        Some(ol_mask),
109        &[],
110        None,
111        None,
112    );
113}
114
115/// 1-D Lp-norm pooling backward pass.
116///
117/// `dx_i = dy * sign(x_i) * (|x_i| / max(y, ε))^(p-1) / max(y, ε)`.
118///
119/// Requires both the original input `x` and the forward output `y`.
120/// `dx` must be zero-initialised before launch.
121#[kernel]
122pub fn lppool1d_backward<
123    T: Triton,
124    D: Float,
125    const KL: i32,
126    const STRIDE: i32,
127    const BLOCK_OL: i32,
128>(
129    dy_ptr: T::Pointer<D>,
130    x_ptr: T::Pointer<D>,
131    y_ptr: T::Pointer<D>,
132    dx_ptr: T::Pointer<D>,
133    _B: i32,
134    C: i32,
135    L: i32,
136    OL: i32,
137    p: f32,
138) where
139    T::I32Tensor: Tensor<i32, 1>,
140    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
141    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
142{
143    let pid = T::program_id(Axis::X);
144    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
145
146    let ol_tile = pid % num_ol_tiles;
147    let bc = pid / num_ol_tiles;
148    let c = bc % C;
149    let b = bc / C;
150
151    let ol_start = ol_tile * BLOCK_OL;
152    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
153    let ol_mask = ol_range.lt(OL);
154
155    let in_bc_base = (b * C + c) * L;
156    let out_bc_base = (b * C + c) * OL;
157
158    let pm1_vec = T::full::<f32>(&[BLOCK_OL], p - 1.0_f32);
159    let eps_vec = T::full::<f32>(&[BLOCK_OL], 1e-12_f32);
160    let zeros_f32 = T::zeros::<f32>(&[BLOCK_OL]);
161
162    let dy_offsets = ol_range + out_bc_base;
163    let dy_tile = T::load(
164        dy_ptr.add_offsets(dy_offsets),
165        Some(ol_mask),
166        Some(T::zeros::<D>(&[BLOCK_OL])),
167        &[],
168        None,
169        None,
170        None,
171        false,
172    );
173    let dy_f32 = T::cast::<D, f32>(dy_tile, None, false);
174
175    let y_tile = T::load(
176        y_ptr.add_offsets(dy_offsets),
177        Some(ol_mask),
178        Some(T::zeros::<D>(&[BLOCK_OL])),
179        &[],
180        None,
181        None,
182        None,
183        false,
184    );
185    let y_f32 = T::cast::<D, f32>(y_tile, None, false);
186    let safe_y = T::maximum(y_f32, eps_vec);
187
188    let loop_bound = KL;
189    for kl in 0..loop_bound {
190        let il_range = ol_range * STRIDE + kl;
191        let in_offsets = il_range + in_bc_base;
192        let x_tile = T::load(
193            x_ptr.add_offsets(in_offsets),
194            Some(ol_mask),
195            Some(T::zeros::<D>(&[BLOCK_OL])),
196            &[],
197            None,
198            None,
199            None,
200            false,
201        );
202        let x_f32 = T::cast::<D, f32>(x_tile, None, false);
203        let abs_x = T::abs(x_f32);
204        let safe_abs = T::maximum(abs_x, eps_vec);
205
206        // sign(x): 1.0 if x > 0, -1.0 if x < 0, 0.0 if x == 0.
207        let pos = T::where_(
208            T::gt(x_f32, zeros_f32),
209            T::full(&[BLOCK_OL], 1.0_f32),
210            zeros_f32,
211        );
212        let neg = T::where_(
213            T::gt(zeros_f32, x_f32),
214            T::full(&[BLOCK_OL], 1.0_f32),
215            zeros_f32,
216        );
217        let sign_x = pos - neg;
218
219        // (|x| / y)^(p-1) = exp((p-1) * log(|x| / y))
220        let ratio = safe_abs / safe_y;
221        let safe_ratio = T::maximum(ratio, eps_vec);
222        let pow_ratio = T::exp(pm1_vec * T::log(safe_ratio));
223
224        let dx_f32 = dy_f32 * sign_x * pow_ratio;
225        let dx_tile = T::cast::<f32, D>(dx_f32, None, false);
226
227        T::atomic_add(
228            dx_ptr.add_offsets(in_offsets),
229            dx_tile,
230            Some(ol_mask),
231            None,
232            None,
233        );
234    }
235}
236
237pub struct Lppool1dOp<'a, T: Float> {
238    pub forward: Lppool1dForward<T>,
239    pub backward: Lppool1dBackward<T>,
240    _marker: core::marker::PhantomData<&'a ()>,
241}
teeny_kernels/nn/pool/lppool1d.rs

teeny_kernels/nn/pool/
lppool1d.rs
