teeny_kernels/nn/pad/
reflection_pad1d.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/// 1-D reflection padding forward pass.
29///
30/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
31///
32/// For output position `op`:
33/// - `ip = op - PAD_LEFT`
34/// - if `ip < 0`: source = input[-ip]
35/// - if `ip >= L`: source = input[2*(L-1) - ip]
36/// - else: source = input[ip]
37///
38/// **Constraints**: `PAD_LEFT < L`, `PAD_RIGHT < L`.
39#[kernel]
40pub fn reflection_pad1d_forward<
41    T: Triton,
42    D: Num,
43    const PAD_LEFT: i32,
44    const PAD_RIGHT: i32,
45    const BLOCK_OL: i32,
46>(
47    input_ptr: T::Pointer<D>,
48    output_ptr: T::Pointer<D>,
49    _B: i32,
50    C: i32,
51    L: i32,
52    OL: i32,
53) where
54    T::I32Tensor: Tensor<i32, 1>,
55    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
56    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
57    T::BoolTensor: BitOr<Output = T::BoolTensor>,
58    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
59{
60    let pid = T::program_id(Axis::X);
61    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
62
63    let ol_tile = pid % num_ol_tiles;
64    let bc = pid / num_ol_tiles;
65    let c = bc % C;
66    let b = bc / C;
67
68    let ol_start = ol_tile * BLOCK_OL;
69    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
70    let ol_mask = ol_range.lt(OL);
71
72    let in_bc_base = (b * C + c) * L;
73    let out_bc_base = (b * C + c) * OL;
74
75    let ip_raw = ol_range - PAD_LEFT;
76    let left_cond = ip_raw.lt(0);
77    let right_cond = ip_raw.ge(L);
78
79    // Reflected index for each case (safe to compute for all lanes)
80    let ip_left = ip_raw * (-1); // -ip: always in [1, PAD_LEFT] for left pad
81    let ip_right = ip_raw * (-1) + (2 * (L - 1)); // 2*(L-1) - ip for right pad
82
83    // Load all three candidate positions. Masks keep each load safe:
84    // - left load: mask = left_cond & ol_mask (only left-pad lanes)
85    // - right load: mask = right_cond & ol_mask
86    // - center load: mask = in_bounds & ol_mask
87    let in_bounds = ip_raw.ge(0) & ip_raw.lt(L);
88
89    let zeros = T::zeros::<D>(&[BLOCK_OL]);
90    let val_center = T::load(
91        input_ptr.add_offsets(ip_raw + in_bc_base),
92        Some(ol_mask & in_bounds),
93        Some(zeros),
94        &[],
95        None,
96        None,
97        None,
98        false,
99    );
100    let val_left = T::load(
101        input_ptr.add_offsets(ip_left + in_bc_base),
102        Some(ol_mask & left_cond),
103        Some(zeros),
104        &[],
105        None,
106        None,
107        None,
108        false,
109    );
110    let val_right = T::load(
111        input_ptr.add_offsets(ip_right + in_bc_base),
112        Some(ol_mask & right_cond),
113        Some(zeros),
114        &[],
115        None,
116        None,
117        None,
118        false,
119    );
120
121    let result = T::where_(
122        left_cond,
123        val_left,
124        T::where_(right_cond, val_right, val_center),
125    );
126
127    let out_offsets = ol_range + out_bc_base;
128    T::store(
129        output_ptr.add_offsets(out_offsets),
130        result,
131        Some(ol_mask),
132        &[],
133        None,
134        None,
135    );
136}
137
138/// 1-D reflection padding backward pass.
139///
140/// Each output gradient position maps back to one input position via the same
141/// reflection rule. Multiple output positions may map to the same input
142/// position (the boundary elements reflect), so `atomic_add` is used.
143/// `dx` must be zero-initialised before launch.
144#[kernel]
145pub fn reflection_pad1d_backward<
146    T: Triton,
147    D: Num,
148    const PAD_LEFT: i32,
149    const PAD_RIGHT: i32,
150    const BLOCK_OL: i32,
151>(
152    dy_ptr: T::Pointer<D>,
153    dx_ptr: T::Pointer<D>,
154    _B: i32,
155    C: i32,
156    L: i32,
157    OL: i32,
158) where
159    T::I32Tensor: Tensor<i32, 1>,
160    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
161    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
162    T::BoolTensor: BitOr<Output = T::BoolTensor>,
163    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
164{
165    let pid = T::program_id(Axis::X);
166    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
167
168    let ol_tile = pid % num_ol_tiles;
169    let bc = pid / num_ol_tiles;
170    let c = bc % C;
171    let b = bc / C;
172
173    let ol_start = ol_tile * BLOCK_OL;
174    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
175    let ol_mask = ol_range.lt(OL);
176
177    let dy_bc_base = (b * C + c) * OL;
178    let dx_bc_base = (b * C + c) * L;
179
180    let dy_offsets = ol_range + dy_bc_base;
181    let dy_tile = T::load(
182        dy_ptr.add_offsets(dy_offsets),
183        Some(ol_mask),
184        Some(T::zeros::<D>(&[BLOCK_OL])),
185        &[],
186        None,
187        None,
188        None,
189        false,
190    );
191
192    let ip_raw = ol_range - PAD_LEFT;
193    let left_cond = ip_raw.lt(0);
194    let right_cond = ip_raw.ge(L);
195    let in_bounds = ip_raw.ge(0) & ip_raw.lt(L);
196
197    let ip_left = ip_raw * (-1);
198    let ip_right = ip_raw * (-1) + (2 * (L - 1));
199
200    // Scatter to center (input) positions
201    T::atomic_add(
202        dx_ptr.add_offsets(ip_raw + dx_bc_base),
203        dy_tile,
204        Some(ol_mask & in_bounds),
205        None,
206        None,
207    );
208    // Scatter to reflected left positions
209    T::atomic_add(
210        dx_ptr.add_offsets(ip_left + dx_bc_base),
211        dy_tile,
212        Some(ol_mask & left_cond),
213        None,
214        None,
215    );
216    // Scatter to reflected right positions
217    T::atomic_add(
218        dx_ptr.add_offsets(ip_right + dx_bc_base),
219        dy_tile,
220        Some(ol_mask & right_cond),
221        None,
222        None,
223    );
224}
225
226pub struct ReflectionPad1dOp<'a, T: Num> {
227    pub forward: ReflectionPad1dForward<T>,
228    pub backward: ReflectionPad1dBackward<T>,
229    _marker: core::marker::PhantomData<&'a ()>,
230}
teeny_kernels/nn/pad/reflection_pad1d.rs

teeny_kernels/nn/pad/
reflection_pad1d.rs
