teeny_kernels/nn/pad/
circular_pad2d.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::ops::{BitAnd, BitOr};
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 circular padding forward pass.
29///
30/// Grid: `pid = ((b*C+c)*OH+oh) * num_ow_tiles + ow_tile`
31#[kernel]
32pub fn circular_pad2d_forward<
33    T: Triton,
34    D: Num,
35    const PT: i32,
36    const PB: i32,
37    const PL: i32,
38    const PR: i32,
39    const BLOCK_OW: i32,
40>(
41    input_ptr: T::Pointer<D>,
42    output_ptr: T::Pointer<D>,
43    _B: i32,
44    C: i32,
45    H: i32,
46    W: i32,
47    OH: i32,
48    OW: i32,
49) where
50    T::I32Tensor: Tensor<i32, 1>,
51    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
52    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
53    T::BoolTensor: BitOr<Output = 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_ow_tiles = T::cdiv(OW, BLOCK_OW);
58
59    let ow_tile = pid % num_ow_tiles;
60    let rest = pid / num_ow_tiles;
61    let oh = rest % OH;
62    let bc = rest / OH;
63    let c = bc % C;
64    let b = bc / C;
65
66    let ow_start = ow_tile * BLOCK_OW;
67    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
68    let ow_mask = ow_range.lt(OW);
69
70    // Wrap height index (branchless: works when |ih_raw| < H)
71    let ih_raw = oh - PT;
72    let ih = (ih_raw + H) % H;
73
74    let in_bc_base = (b * C + c) * H * W + ih * W;
75    let out_bc_base = ((b * C + c) * OH + oh) * OW;
76
77    let iw_raw = ow_range - PL;
78    let cond_left = iw_raw.lt(0);
79    let cond_right = iw_raw.ge(W);
80    let in_bounds = iw_raw.ge(0) & iw_raw.lt(W);
81
82    let iw_wrap_left = iw_raw + W;
83    let iw_wrap_right = iw_raw - W;
84
85    let zeros = T::zeros::<D>(&[BLOCK_OW]);
86    let val_center = T::load(
87        input_ptr.add_offsets(iw_raw + in_bc_base),
88        Some(ow_mask & in_bounds),
89        Some(zeros),
90        &[],
91        None,
92        None,
93        None,
94        false,
95    );
96    let val_left = T::load(
97        input_ptr.add_offsets(iw_wrap_left + in_bc_base),
98        Some(ow_mask & cond_left),
99        Some(zeros),
100        &[],
101        None,
102        None,
103        None,
104        false,
105    );
106    let val_right = T::load(
107        input_ptr.add_offsets(iw_wrap_right + in_bc_base),
108        Some(ow_mask & cond_right),
109        Some(zeros),
110        &[],
111        None,
112        None,
113        None,
114        false,
115    );
116
117    let result = T::where_(
118        cond_left,
119        val_left,
120        T::where_(cond_right, val_right, val_center),
121    );
122
123    let out_offsets = ow_range + out_bc_base;
124    T::store(
125        output_ptr.add_offsets(out_offsets),
126        result,
127        Some(ow_mask),
128        &[],
129        None,
130        None,
131    );
132}
133
134/// 2-D circular padding backward pass.
135#[kernel]
136pub fn circular_pad2d_backward<
137    T: Triton,
138    D: Num,
139    const PT: i32,
140    const PB: i32,
141    const PL: i32,
142    const PR: i32,
143    const BLOCK_OW: i32,
144>(
145    dy_ptr: T::Pointer<D>,
146    dx_ptr: T::Pointer<D>,
147    _B: i32,
148    C: i32,
149    H: i32,
150    W: i32,
151    OH: i32,
152    OW: i32,
153) where
154    T::I32Tensor: Tensor<i32, 1>,
155    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
156    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
157    T::BoolTensor: BitOr<Output = T::BoolTensor>,
158    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
159{
160    let pid = T::program_id(Axis::X);
161    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
162
163    let ow_tile = pid % num_ow_tiles;
164    let rest = pid / num_ow_tiles;
165    let oh = rest % OH;
166    let bc = rest / OH;
167    let c = bc % C;
168    let b = bc / C;
169
170    let ow_start = ow_tile * BLOCK_OW;
171    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
172    let ow_mask = ow_range.lt(OW);
173
174    let ih_raw = oh - PT;
175    let ih = (ih_raw + H) % H;
176
177    let dy_bc_base = ((b * C + c) * OH + oh) * OW;
178    let dx_bc_base = (b * C + c) * H * W + ih * W;
179
180    let dy_offsets = ow_range + dy_bc_base;
181    let dy_tile = T::load(
182        dy_ptr.add_offsets(dy_offsets),
183        Some(ow_mask),
184        Some(T::zeros::<D>(&[BLOCK_OW])),
185        &[],
186        None,
187        None,
188        None,
189        false,
190    );
191
192    let iw_raw = ow_range - PL;
193    let cond_left = iw_raw.lt(0);
194    let cond_right = iw_raw.ge(W);
195    let in_bounds = iw_raw.ge(0) & iw_raw.lt(W);
196
197    let iw_wrap_left = iw_raw + W;
198    let iw_wrap_right = iw_raw - W;
199
200    T::atomic_add(
201        dx_ptr.add_offsets(iw_raw + dx_bc_base),
202        dy_tile,
203        Some(ow_mask & in_bounds),
204        None,
205        None,
206    );
207    T::atomic_add(
208        dx_ptr.add_offsets(iw_wrap_left + dx_bc_base),
209        dy_tile,
210        Some(ow_mask & cond_left),
211        None,
212        None,
213    );
214    T::atomic_add(
215        dx_ptr.add_offsets(iw_wrap_right + dx_bc_base),
216        dy_tile,
217        Some(ow_mask & cond_right),
218        None,
219        None,
220    );
221}
222
223pub struct CircularPad2dOp<'a, T: Num> {
224    pub forward: CircularPad2dForward<T>,
225    pub backward: CircularPad2dBackward<T>,
226    _marker: core::marker::PhantomData<&'a ()>,
227}
teeny_kernels/nn/pad/circular_pad2d.rs

teeny_kernels/nn/pad/
circular_pad2d.rs
