teeny_kernels/nn/pool/
maxpool1d.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::Num;
20use teeny_macros::kernel;
21use teeny_triton::triton::{
22    types::{AddOffsets, Comparison, Tensor},
23    *,
24};
25
26/// 1-D max-pooling forward pass.
27///
28/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
29///
30/// Initialises the accumulator to `-inf` and reduces over KL positions.
31///
32/// **Constraints**: no padding; `OL = (L - KL) / STRIDE + 1`.
33#[kernel]
34pub fn maxpool1d_forward<T: Triton, D: Num, const KL: i32, const STRIDE: i32, const BLOCK_OL: i32>(
35    input_ptr: T::Pointer<D>,
36    output_ptr: T::Pointer<D>,
37    _B: i32,
38    C: i32,
39    L: i32,
40    OL: i32,
41) where
42    T::I32Tensor: Tensor<i32, 1>,
43    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
44    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
45{
46    let pid = T::program_id(Axis::X);
47    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
48
49    let ol_tile = pid % num_ol_tiles;
50    let bc = pid / num_ol_tiles;
51    let c = bc % C;
52    let b = bc / C;
53
54    let ol_start = ol_tile * BLOCK_OL;
55    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
56    let ol_mask = ol_range.lt(OL);
57
58    let in_bc_base = (b * C + c) * L;
59    let out_bc_base = (b * C + c) * OL;
60
61    let mut acc = T::cast::<f32, D>(T::full::<f32>(&[BLOCK_OL], -3.4028235e38_f32), None, false);
62
63    let loop_bound = KL;
64    for kl in 0..loop_bound {
65        let il_range = ol_range * STRIDE + kl;
66        let in_offsets = il_range + in_bc_base;
67        let tile = T::load(
68            input_ptr.add_offsets(in_offsets),
69            Some(ol_mask),
70            Some(T::cast::<f32, D>(
71                T::full::<f32>(&[BLOCK_OL], -3.4028235e38_f32),
72                None,
73                false,
74            )),
75            &[],
76            None,
77            None,
78            None,
79            false,
80        );
81        acc = T::maximum(acc, tile);
82    }
83
84    let out_offsets = ol_range + out_bc_base;
85    T::store(
86        output_ptr.add_offsets(out_offsets),
87        acc,
88        Some(ol_mask),
89        &[],
90        None,
91        None,
92    );
93}
94
95/// 1-D max-pooling backward pass.
96///
97/// Re-scans the input window to find the max position, then scatters `dy` to
98/// all input elements that equal the stored output maximum. `dx` must be
99/// zero-initialised before launch.
100#[kernel]
101pub fn maxpool1d_backward<
102    T: Triton,
103    D: Num,
104    const KL: i32,
105    const STRIDE: i32,
106    const BLOCK_OL: i32,
107>(
108    dy_ptr: T::Pointer<D>,
109    x_ptr: T::Pointer<D>,
110    y_ptr: T::Pointer<D>,
111    dx_ptr: T::Pointer<D>,
112    _B: i32,
113    C: i32,
114    L: i32,
115    OL: i32,
116) where
117    T::I32Tensor: Tensor<i32, 1>,
118    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
119    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
120{
121    let pid = T::program_id(Axis::X);
122    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
123
124    let ol_tile = pid % num_ol_tiles;
125    let bc = pid / num_ol_tiles;
126    let c = bc % C;
127    let b = bc / C;
128
129    let ol_start = ol_tile * BLOCK_OL;
130    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
131    let ol_mask = ol_range.lt(OL);
132
133    let in_bc_base = (b * C + c) * L;
134    let out_bc_base = (b * C + c) * OL;
135
136    let dy_offsets = ol_range + out_bc_base;
137    let dy_tile = T::load(
138        dy_ptr.add_offsets(dy_offsets),
139        Some(ol_mask),
140        Some(T::zeros::<D>(&[BLOCK_OL])),
141        &[],
142        None,
143        None,
144        None,
145        false,
146    );
147
148    let y_tile = T::load(
149        y_ptr.add_offsets(dy_offsets),
150        Some(ol_mask),
151        Some(T::zeros::<D>(&[BLOCK_OL])),
152        &[],
153        None,
154        None,
155        None,
156        false,
157    );
158
159    let loop_bound = KL;
160    for kl in 0..loop_bound {
161        let il_range = ol_range * STRIDE + kl;
162        let in_offsets = il_range + in_bc_base;
163        let x_tile = T::load(
164            x_ptr.add_offsets(in_offsets),
165            Some(ol_mask),
166            Some(T::cast::<f32, D>(
167                T::full::<f32>(&[BLOCK_OL], -3.4028235e38_f32),
168                None,
169                false,
170            )),
171            &[],
172            None,
173            None,
174            None,
175            false,
176        );
177        let is_max = T::eq(x_tile, y_tile);
178        let grad = T::where_(is_max, dy_tile, T::zeros::<D>(&[BLOCK_OL]));
179        T::atomic_add(
180            dx_ptr.add_offsets(in_offsets),
181            grad,
182            Some(ol_mask),
183            None,
184            None,
185        );
186    }
187}
188
189pub struct Maxpool1dOp<'a, T: Num> {
190    pub forward: Maxpool1dForward<T>,
191    pub backward: Maxpool1dBackward<T>,
192    _marker: core::marker::PhantomData<&'a ()>,
193}
teeny_kernels/nn/pool/maxpool1d.rs

teeny_kernels/nn/pool/
maxpool1d.rs
