teeny_kernels/nn/pool/
avgpool2d.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/// 2-D average-pooling forward pass.
27///
28/// Grid: one flat 1-D pid per (b, c, oh, ow-tile) combination:
29///   pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile
30///
31/// Each CTA accumulates a BLOCK_OW-wide strip of output columns via a flat
32/// loop over all KH×KW kernel positions (avoids nested scf.for).
33///
34/// **Constraints**: no padding; `OH = (H - KH) / STRIDE_H + 1`, `OW = (W - KW) / STRIDE_W + 1`.
35#[kernel]
36pub fn avgpool2d_forward<
37    T: Triton,
38    D: Num,
39    const KH: i32,
40    const KW: i32,
41    const STRIDE_H: i32,
42    const STRIDE_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    // Decode flat pid → (b, c, oh, ow_tile).
62    let ow_tile = pid % num_ow_tiles;
63    let bco = pid / num_ow_tiles;
64    let oh = bco % OH;
65    let bc = bco / OH;
66    let c = bc % C;
67    let b = bc / C;
68
69    let ow_start = ow_tile * BLOCK_OW;
70    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
71    let ow_mask = ow_range.lt(OW);
72
73    let in_bc_base = (b * C + c) * H * W;
74    let out_bc_base = (b * C + c) * OH * OW;
75
76    let mut acc = T::zeros::<D>(&[BLOCK_OW]);
77
78    // Flat loop over KH * KW kernel positions to avoid nested scf.for.
79    let loop_bound = KH * KW;
80    for idx in 0..loop_bound {
81        let kw = idx % KW;
82        let kh = idx / KW;
83        let ih = oh * STRIDE_H + kh;
84        let iw_range = ow_range * STRIDE_W + kw;
85        let in_offsets = iw_range + (in_bc_base + ih * W);
86        let tile = T::load(
87            input_ptr.add_offsets(in_offsets),
88            Some(ow_mask),
89            Some(T::zeros::<D>(&[BLOCK_OW])),
90            &[],
91            None,
92            None,
93            None,
94            false,
95        );
96        acc = acc + tile;
97    }
98
99    // Scale by 1/(KH*KW): build [1] i32 tensor from arange then broadcast.
100    let ksize_1 = T::full::<i32>(&[1], KH * KW);
101    let ksize_f_1 = T::cast::<i32, D>(ksize_1, None, false);
102    let ksize = T::broadcast_to(ksize_f_1, &[BLOCK_OW]);
103    let result = acc / ksize;
104
105    let out_offsets = ow_range + (out_bc_base + oh * OW);
106    T::store(
107        output_ptr.add_offsets(out_offsets),
108        result,
109        Some(ow_mask),
110        &[],
111        None,
112        None,
113    );
114}
115
116/// 2-D average-pooling backward pass.
117///
118/// Uses the same flat 1-D grid as the forward pass, iterating over output
119/// positions. For each output element dy[b, c, oh, ow], the gradient is
120/// spread uniformly across the KH×KW input window via `atomic_add` so that
121/// overlapping pooling windows (stride < kernel size) are handled correctly.
122///
123/// **Constraints**: `dx` must be zero-initialised before launch; same spatial
124/// constraints as `avgpool2d_forward`.
125#[kernel]
126pub fn avgpool2d_backward<
127    T: Triton,
128    D: Num,
129    const KH: i32,
130    const KW: i32,
131    const STRIDE_H: i32,
132    const STRIDE_W: i32,
133    const BLOCK_OW: i32,
134>(
135    dy_ptr: T::Pointer<D>,
136    dx_ptr: T::Pointer<D>,
137    _B: i32,
138    C: i32,
139    H: i32,
140    W: i32,
141    OH: i32,
142    OW: i32,
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 dy_bc_base = (b * C + c) * OH * OW;
163    let dx_bc_base = (b * C + c) * H * W;
164
165    // Load upstream gradient tile and scale by 1/(KH*KW).
166    let dy_offsets = ow_range + (dy_bc_base + oh * OW);
167    let dy_tile = T::load(
168        dy_ptr.add_offsets(dy_offsets),
169        Some(ow_mask),
170        Some(T::zeros::<D>(&[BLOCK_OW])),
171        &[],
172        None,
173        None,
174        None,
175        false,
176    );
177    let ksize_1 = T::full::<i32>(&[1], KH * KW);
178    let ksize_f_1 = T::cast::<i32, D>(ksize_1, None, false);
179    let ksize = T::broadcast_to(ksize_f_1, &[BLOCK_OW]);
180    let grad = dy_tile / ksize;
181
182    // Scatter scaled gradient to the KH×KW input window via flat loop.
183    let loop_bound = KH * KW;
184    for idx in 0..loop_bound {
185        let kw = idx % KW;
186        let kh = idx / KW;
187        let ih = oh * STRIDE_H + kh;
188        let iw_range = ow_range * STRIDE_W + kw;
189        let dx_offsets = iw_range + (dx_bc_base + ih * W);
190        T::atomic_add(
191            dx_ptr.add_offsets(dx_offsets),
192            grad,
193            Some(ow_mask),
194            None,
195            None,
196        );
197    }
198}
199
200impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for Avgpool2dForward<D> {
201    fn n_activation_inputs(&self) -> usize {
202        1
203    }
204
205    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
206        Vec::new()
207    }
208
209    fn pack_args(
210        &self,
211        inputs: &[(teeny_core::model::RawPtr, &[usize])],
212        _params: &[teeny_core::model::RawPtr],
213        output: teeny_core::model::RawPtr,
214        output_shape: &[usize],
215        _output_row_stride: i32,
216        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
217    ) {
218        // kernel args: input_ptr, output_ptr, B, C, H, W, OH, OW
219        // input_shape = [B, C, H, W], output_shape = [B, C, OH, OW]
220        let input_shape = inputs[0].1;
221        visitor.visit_ptr(inputs[0].0);
222        visitor.visit_ptr(output);
223        visitor.visit_i32(input_shape[0] as i32); // B
224        visitor.visit_i32(input_shape[1] as i32); // C
225        visitor.visit_i32(input_shape[2] as i32); // H
226        visitor.visit_i32(input_shape[3] as i32); // W
227        visitor.visit_i32(output_shape[2] as i32); // OH
228        visitor.visit_i32(output_shape[3] as i32); // OW
229    }
230
231    fn block(&self) -> [u32; 3] {
232        [128, 1, 1]
233    }
234
235    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
236        // pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile
237        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
238        [
239            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
240            1,
241            1,
242        ]
243    }
244}
245
246pub struct Avgpool2dOp<'a, T: Num> {
247    pub forward: Avgpool2dForward<T>,
248    pub backward: Avgpool2dBackward<T>,
249    _marker: core::marker::PhantomData<&'a ()>,
250}
teeny_kernels/nn/pool/avgpool2d.rs

teeny_kernels/nn/pool/
avgpool2d.rs
