teeny_kernels/nn/pool/
avgpool1d.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 average-pooling forward pass.
27///
28/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
29///
30/// Each CTA sums a BLOCK_OL-wide strip over the KL kernel positions then
31/// divides by KL.
32///
33/// **Constraints**: no padding; `OL = (L - KL) / STRIDE + 1`.
34#[kernel]
35pub fn avgpool1d_forward<T: Triton, D: Num, const KL: i32, const STRIDE: i32, const BLOCK_OL: i32>(
36    input_ptr: T::Pointer<D>,
37    output_ptr: T::Pointer<D>,
38    _B: i32,
39    C: i32,
40    L: i32,
41    OL: i32,
42) where
43    T::I32Tensor: Tensor<i32, 1>,
44    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
45    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
46{
47    let pid = T::program_id(Axis::X);
48    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
49
50    let ol_tile = pid % num_ol_tiles;
51    let bc = pid / num_ol_tiles;
52    let c = bc % C;
53    let b = bc / C;
54
55    let ol_start = ol_tile * BLOCK_OL;
56    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
57    let ol_mask = ol_range.lt(OL);
58
59    let in_bc_base = (b * C + c) * L;
60    let out_bc_base = (b * C + c) * OL;
61
62    let mut acc = T::zeros::<D>(&[BLOCK_OL]);
63
64    let loop_bound = KL;
65    for kl in 0..loop_bound {
66        let il_range = ol_range * STRIDE + kl;
67        let in_offsets = il_range + in_bc_base;
68        let tile = T::load(
69            input_ptr.add_offsets(in_offsets),
70            Some(ol_mask),
71            Some(T::zeros::<D>(&[BLOCK_OL])),
72            &[],
73            None,
74            None,
75            None,
76            false,
77        );
78        acc = acc + tile;
79    }
80
81    let ksize_1 = T::full::<i32>(&[1], KL);
82    let ksize_f_1 = T::cast::<i32, D>(ksize_1, None, false);
83    let ksize = T::broadcast_to(ksize_f_1, &[BLOCK_OL]);
84    let result = acc / ksize;
85
86    let out_offsets = ol_range + out_bc_base;
87    T::store(
88        output_ptr.add_offsets(out_offsets),
89        result,
90        Some(ol_mask),
91        &[],
92        None,
93        None,
94    );
95}
96
97/// 1-D average-pooling backward pass.
98///
99/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
100///
101/// Spreads each output gradient uniformly across its KL input positions via
102/// `atomic_add`. `dx` must be zero-initialised before launch.
103#[kernel]
104pub fn avgpool1d_backward<
105    T: Triton,
106    D: Num,
107    const KL: i32,
108    const STRIDE: i32,
109    const BLOCK_OL: i32,
110>(
111    dy_ptr: T::Pointer<D>,
112    dx_ptr: T::Pointer<D>,
113    _B: i32,
114    C: i32,
115    L: i32,
116    OL: i32,
117) where
118    T::I32Tensor: Tensor<i32, 1>,
119    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
120    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
121{
122    let pid = T::program_id(Axis::X);
123    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
124
125    let ol_tile = pid % num_ol_tiles;
126    let bc = pid / num_ol_tiles;
127    let c = bc % C;
128    let b = bc / C;
129
130    let ol_start = ol_tile * BLOCK_OL;
131    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
132    let ol_mask = ol_range.lt(OL);
133
134    let dy_bc_base = (b * C + c) * OL;
135    let dx_bc_base = (b * C + c) * L;
136
137    let dy_offsets = ol_range + dy_bc_base;
138    let dy_tile = T::load(
139        dy_ptr.add_offsets(dy_offsets),
140        Some(ol_mask),
141        Some(T::zeros::<D>(&[BLOCK_OL])),
142        &[],
143        None,
144        None,
145        None,
146        false,
147    );
148    let ksize_1 = T::full::<i32>(&[1], KL);
149    let ksize_f_1 = T::cast::<i32, D>(ksize_1, None, false);
150    let ksize = T::broadcast_to(ksize_f_1, &[BLOCK_OL]);
151    let grad = dy_tile / ksize;
152
153    let loop_bound = KL;
154    for kl in 0..loop_bound {
155        let il_range = ol_range * STRIDE + kl;
156        let dx_offsets = il_range + dx_bc_base;
157        T::atomic_add(
158            dx_ptr.add_offsets(dx_offsets),
159            grad,
160            Some(ol_mask),
161            None,
162            None,
163        );
164    }
165}
166
167pub struct Avgpool1dOp<'a, T: Num> {
168    pub forward: Avgpool1dForward<T>,
169    pub backward: Avgpool1dBackward<T>,
170    _marker: core::marker::PhantomData<&'a ()>,
171}
teeny_kernels/nn/pool/avgpool1d.rs

teeny_kernels/nn/pool/
avgpool1d.rs
