teeny_kernels/nn/pad/
reflection_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 reflection padding forward pass.
29///
30/// Grid: `pid = ((b*C+c)*OH+oh) * num_ow_tiles + ow_tile`
31#[kernel]
32pub fn reflection_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    // Reflect height index
71    let ih_raw = oh - PT;
72    let ih = if ih_raw < 0 {
73        -ih_raw
74    } else if ih_raw >= H {
75        2 * (H - 1) - ih_raw
76    } else {
77        ih_raw
78    };
79
80    let in_bc_base = (b * C + c) * H * W + ih * W;
81    let out_bc_base = ((b * C + c) * OH + oh) * OW;
82
83    let iw_raw = ow_range - PL;
84    let left_cond = iw_raw.lt(0);
85    let right_cond = iw_raw.ge(W);
86    let in_bounds = iw_raw.ge(0) & iw_raw.lt(W);
87
88    let iw_left = iw_raw * (-1);
89    let iw_right = iw_raw * (-1) + (2 * (W - 1));
90
91    let zeros = T::zeros::<D>(&[BLOCK_OW]);
92    let val_center = T::load(
93        input_ptr.add_offsets(iw_raw + in_bc_base),
94        Some(ow_mask & in_bounds),
95        Some(zeros),
96        &[],
97        None,
98        None,
99        None,
100        false,
101    );
102    let val_left = T::load(
103        input_ptr.add_offsets(iw_left + in_bc_base),
104        Some(ow_mask & left_cond),
105        Some(zeros),
106        &[],
107        None,
108        None,
109        None,
110        false,
111    );
112    let val_right = T::load(
113        input_ptr.add_offsets(iw_right + in_bc_base),
114        Some(ow_mask & right_cond),
115        Some(zeros),
116        &[],
117        None,
118        None,
119        None,
120        false,
121    );
122
123    let result = T::where_(
124        left_cond,
125        val_left,
126        T::where_(right_cond, val_right, val_center),
127    );
128
129    let out_offsets = ow_range + out_bc_base;
130    T::store(
131        output_ptr.add_offsets(out_offsets),
132        result,
133        Some(ow_mask),
134        &[],
135        None,
136        None,
137    );
138}
139
140/// 2-D reflection padding backward pass.
141#[kernel]
142pub fn reflection_pad2d_backward<
143    T: Triton,
144    D: Num,
145    const PT: i32,
146    const PB: i32,
147    const PL: i32,
148    const PR: i32,
149    const BLOCK_OW: i32,
150>(
151    dy_ptr: T::Pointer<D>,
152    dx_ptr: T::Pointer<D>,
153    _B: i32,
154    C: i32,
155    H: i32,
156    W: i32,
157    OH: i32,
158    OW: i32,
159) where
160    T::I32Tensor: Tensor<i32, 1>,
161    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
162    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
163    T::BoolTensor: BitOr<Output = T::BoolTensor>,
164    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
165{
166    let pid = T::program_id(Axis::X);
167    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
168
169    let ow_tile = pid % num_ow_tiles;
170    let rest = pid / num_ow_tiles;
171    let oh = rest % OH;
172    let bc = rest / OH;
173    let c = bc % C;
174    let b = bc / C;
175
176    let ow_start = ow_tile * BLOCK_OW;
177    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
178    let ow_mask = ow_range.lt(OW);
179
180    let ih_raw = oh - PT;
181    let ih = if ih_raw < 0 {
182        -ih_raw
183    } else if ih_raw >= H {
184        2 * (H - 1) - ih_raw
185    } else {
186        ih_raw
187    };
188
189    let dy_bc_base = ((b * C + c) * OH + oh) * OW;
190    let dx_bc_base = (b * C + c) * H * W + ih * W;
191
192    let dy_offsets = ow_range + dy_bc_base;
193    let dy_tile = T::load(
194        dy_ptr.add_offsets(dy_offsets),
195        Some(ow_mask),
196        Some(T::zeros::<D>(&[BLOCK_OW])),
197        &[],
198        None,
199        None,
200        None,
201        false,
202    );
203
204    let iw_raw = ow_range - PL;
205    let left_cond = iw_raw.lt(0);
206    let right_cond = iw_raw.ge(W);
207    let in_bounds = iw_raw.ge(0) & iw_raw.lt(W);
208
209    let iw_left = iw_raw * (-1);
210    let iw_right = iw_raw * (-1) + (2 * (W - 1));
211
212    T::atomic_add(
213        dx_ptr.add_offsets(iw_raw + dx_bc_base),
214        dy_tile,
215        Some(ow_mask & in_bounds),
216        None,
217        None,
218    );
219    T::atomic_add(
220        dx_ptr.add_offsets(iw_left + dx_bc_base),
221        dy_tile,
222        Some(ow_mask & left_cond),
223        None,
224        None,
225    );
226    T::atomic_add(
227        dx_ptr.add_offsets(iw_right + dx_bc_base),
228        dy_tile,
229        Some(ow_mask & right_cond),
230        None,
231        None,
232    );
233}
234
235pub struct ReflectionPad2dOp<'a, T: Num> {
236    pub forward: ReflectionPad2dForward<T>,
237    pub backward: ReflectionPad2dBackward<T>,
238    _marker: core::marker::PhantomData<&'a ()>,
239}
teeny_kernels/nn/pad/reflection_pad2d.rs

teeny_kernels/nn/pad/
reflection_pad2d.rs
