teeny_kernels/nn/pool/
maxpool2d.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 core::marker::PhantomData;
20
21use teeny_core::dtype::Num;
22use teeny_macros::kernel;
23use teeny_triton::triton::{
24    types::{AddOffsets, Comparison, Tensor},
25    *,
26};
27
28/// 2-D max-pooling forward pass with optional symmetric padding.
29///
30/// Grid: `pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile`
31///
32/// `OH = (H + 2*PAD_H - KH) / STRIDE_H + 1`, `OW = (W + 2*PAD_W - KW) / STRIDE_W + 1`.
33#[kernel]
34pub fn maxpool2d_forward<
35    T: Triton,
36    D: Num,
37    const KH: i32,
38    const KW: i32,
39    const STRIDE_H: i32,
40    const STRIDE_W: i32,
41    const PAD_H: i32,
42    const PAD_W: i32,
43    const BLOCK_OW: i32,
44>(
45    input_ptr: T::Pointer<D>,
46    output_ptr: T::Pointer<D>,
47    _B: i32,
48    C: i32,
49    H: i32,
50    W: i32,
51    OH: i32,
52    OW: i32,
53) where
54    T::I32Tensor: Tensor<i32, 1>,
55    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
56    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
57{
58    let pid = T::program_id(Axis::X);
59    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
60
61    let ow_tile = pid % num_ow_tiles;
62    let bco = pid / num_ow_tiles;
63    let oh = bco % OH;
64    let bc = bco / OH;
65    let c = bc % C;
66    let b = bc / C;
67
68    let ow_start = ow_tile * BLOCK_OW;
69    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
70    let ow_mask = ow_range.lt(OW);
71
72    let in_bc_base = (b * C + c) * H * W;
73    let out_bc_base = (b * C + c) * OH * OW;
74
75    let mut acc = T::cast::<f32, D>(T::full::<f32>(&[BLOCK_OW], -3.4028235e38_f32), None, false);
76
77    // The Triton MLIR frontend only generates correct scf.while (with body) for
78    // single-variable, single-condition loops with no nested control flow.
79    // A nested `for kw` inside `while kh` creates extra basic blocks that break
80    // the outer loop's structural conversion — the body gets dropped silently.
81    //
82    // Solution: flatten (kh, kw) into a single linear loop `iter < n_valid_kh * KW`,
83    // recover kh/kw per-iteration via div/rem (KW is a compile-time constant).
84    //
85    // Phase 1: skip-loop — advance kh until ih = oh*STRIDE_H + kh - PAD_H >= 0.
86    let mut kh: i32 = 0;
87    let mut ih: i32 = oh * STRIDE_H - PAD_H;
88    while ih < 0 {
89        kh += 1;
90        ih += 1; // each kh step advances ih by exactly 1
91    }
92    // Phase 2: clamp kh_hi = min(kh + (H - ih), KH) via countdown.
93    let mut kh_hi: i32 = kh + (H - ih);
94    while kh_hi > KH {
95        kh_hi -= 1;
96    }
97    // Phase 3: flat loop over all (kh, kw) pairs — matches BN kernel structure
98    // (single variable `iter`, single condition, vector body, no nested loops).
99    // If kh_hi <= kh (no valid rows), total_iters <= 0 → loop runs 0 times.
100    let total_iters = (kh_hi - kh) * KW;
101    let ih_lo = ih;
102    let mut iter: i32 = 0;
103    while iter < total_iters {
104        let kw_idx = iter % KW;
105        let ih_local = ih_lo + iter / KW;
106        let iw_range = ow_range * STRIDE_W + kw_idx - PAD_W;
107        let iw_valid = iw_range.ge(0) & iw_range.lt(W);
108        let valid_mask = ow_mask & iw_valid;
109        let in_offsets = iw_range + (in_bc_base + ih_local * W);
110        let tile = T::load(
111            input_ptr.add_offsets(in_offsets),
112            Some(valid_mask),
113            Some(T::cast::<f32, D>(
114                T::full::<f32>(&[BLOCK_OW], -3.4028235e38_f32),
115                None,
116                false,
117            )),
118            &[],
119            None,
120            None,
121            None,
122            false,
123        );
124        acc = T::maximum(acc, tile);
125        iter += 1;
126    }
127
128    let out_offsets = ow_range + (out_bc_base + oh * OW);
129    T::store(
130        output_ptr.add_offsets(out_offsets),
131        acc,
132        Some(ow_mask),
133        &[],
134        None,
135        None,
136    );
137}
138
139/// 2-D max-pooling backward pass with optional symmetric padding.
140///
141/// Re-scans the input window and scatters `dy` to positions where
142/// `input == output_max`. `dx` must be zero-initialised before launch.
143#[kernel]
144pub fn maxpool2d_backward<
145    T: Triton,
146    D: Num,
147    const KH: i32,
148    const KW: i32,
149    const STRIDE_H: i32,
150    const STRIDE_W: i32,
151    const PAD_H: i32,
152    const PAD_W: i32,
153    const BLOCK_OW: i32,
154>(
155    dy_ptr: T::Pointer<D>,
156    x_ptr: T::Pointer<D>,
157    y_ptr: T::Pointer<D>,
158    dx_ptr: T::Pointer<D>,
159    _B: i32,
160    C: i32,
161    H: i32,
162    W: i32,
163    OH: i32,
164    OW: i32,
165) where
166    T::I32Tensor: Tensor<i32, 1>,
167    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
168    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
169{
170    let pid = T::program_id(Axis::X);
171    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
172
173    let ow_tile = pid % num_ow_tiles;
174    let bco = pid / num_ow_tiles;
175    let oh = bco % OH;
176    let bc = bco / OH;
177    let c = bc % C;
178    let b = bc / C;
179
180    let ow_start = ow_tile * BLOCK_OW;
181    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
182    let ow_mask = ow_range.lt(OW);
183
184    let in_bc_base = (b * C + c) * H * W;
185    let out_bc_base = (b * C + c) * OH * OW;
186
187    let out_offsets = ow_range + (out_bc_base + oh * OW);
188    let dy_tile = T::load(
189        dy_ptr.add_offsets(out_offsets),
190        Some(ow_mask),
191        Some(T::zeros::<D>(&[BLOCK_OW])),
192        &[],
193        None,
194        None,
195        None,
196        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
209    let mut kh: i32 = 0;
210    let mut ih: i32 = oh * STRIDE_H - PAD_H;
211    while ih < 0 {
212        kh += 1;
213        ih += 1;
214    }
215    let mut kh_hi: i32 = kh + (H - ih);
216    while kh_hi > KH {
217        kh_hi -= 1;
218    }
219    let total_iters = (kh_hi - kh) * KW;
220    let ih_lo = ih;
221    let mut iter: i32 = 0;
222    while iter < total_iters {
223        let kw_idx = iter % KW;
224        let ih_local = ih_lo + iter / KW;
225        let iw_range = ow_range * STRIDE_W + kw_idx - PAD_W;
226        let iw_valid = iw_range.ge(0) & iw_range.lt(W);
227        let valid_mask = ow_mask & iw_valid;
228        let in_offsets = iw_range + (in_bc_base + ih_local * W);
229        let x_tile = T::load(
230            x_ptr.add_offsets(in_offsets),
231            Some(valid_mask),
232            Some(T::cast::<f32, D>(
233                T::full::<f32>(&[BLOCK_OW], -3.4028235e38_f32),
234                None,
235                false,
236            )),
237            &[],
238            None,
239            None,
240            None,
241            false,
242        );
243        let is_max = T::eq(x_tile, y_tile);
244        let grad = T::where_(is_max, dy_tile, T::zeros::<D>(&[BLOCK_OW]));
245        T::atomic_add(
246            dx_ptr.add_offsets(in_offsets),
247            grad,
248            Some(valid_mask),
249            None,
250            None,
251        );
252        iter += 1;
253    }
254}
255
256impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for Maxpool2dForward<D> {
257    fn n_activation_inputs(&self) -> usize {
258        1
259    }
260
261    fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
262        Vec::new()
263    }
264
265    fn pack_args(
266        &self,
267        inputs: &[(teeny_core::model::RawPtr, &[usize])],
268        _params: &[teeny_core::model::RawPtr],
269        output: teeny_core::model::RawPtr,
270        output_shape: &[usize],
271        _output_row_stride: i32,
272        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
273    ) {
274        let input_shape = inputs[0].1;
275        visitor.visit_ptr(inputs[0].0);
276        visitor.visit_ptr(output);
277        visitor.visit_i32(input_shape[0] as i32); // B
278        visitor.visit_i32(input_shape[1] as i32); // C
279        visitor.visit_i32(input_shape[2] as i32); // H
280        visitor.visit_i32(input_shape[3] as i32); // W
281        visitor.visit_i32(output_shape[2] as i32); // OH
282        visitor.visit_i32(output_shape[3] as i32); // OW
283    }
284
285    fn block(&self) -> [u32; 3] {
286        [128, 1, 1]
287    }
288
289    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
290        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
291        [
292            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
293            1,
294            1,
295        ]
296    }
297
298    #[cfg(feature = "training")]
299    fn has_backward(&self) -> bool {
300        true
301    }
302
303    /// kernel args: dy, x, y, dx, B, C, H, W, OH, OW
304    #[cfg(feature = "training")]
305    fn pack_backward_args(
306        &self,
307        inputs: &[(teeny_core::model::RawPtr, &[usize])],
308        _params: &[teeny_core::model::RawPtr],
309        output: teeny_core::model::RawPtr,
310        output_shape: &[usize],
311        grad_output: teeny_core::model::RawPtr,
312        _grad_output_row_stride: i32,
313        grad_inputs: &[teeny_core::model::RawPtr],
314        _grad_params: &[teeny_core::model::RawPtr],
315        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
316    ) {
317        let in_shape = inputs[0].1; // [B, C, H, W]
318        visitor.visit_ptr(grad_output); // dy_ptr
319        visitor.visit_ptr(inputs[0].0); // x_ptr (saved activation)
320        visitor.visit_ptr(output); // y_ptr (forward output = max values)
321        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
322        visitor.visit_i32(in_shape[0] as i32); // B
323        visitor.visit_i32(in_shape[1] as i32); // C
324        visitor.visit_i32(in_shape[2] as i32); // H
325        visitor.visit_i32(in_shape[3] as i32); // W
326        visitor.visit_i32(output_shape[2] as i32); // OH
327        visitor.visit_i32(output_shape[3] as i32); // OW
328    }
329
330    #[cfg(feature = "training")]
331    fn backward_block(&self) -> [u32; 3] {
332        [128, 1, 1]
333    }
334
335    /// Grid over output positions (same formula as forward).
336    #[cfg(feature = "training")]
337    fn backward_grid(&self, _input_shapes: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
338        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
339        [
340            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
341            1,
342            1,
343        ]
344    }
345}
346
347pub struct Maxpool2dOp<'a, T: Num> {
348    pub forward: Maxpool2dForward<T>,
349    pub backward: Maxpool2dBackward<T>,
350    _marker: PhantomData<&'a ()>,
351}
teeny_kernels/nn/pool/maxpool2d.rs

teeny_kernels/nn/pool/
maxpool2d.rs
