teeny_kernels/nn/pad/
circular_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 circular padding forward pass.
29///
30/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
31///
32/// For output position `op`:
33/// - `ip_raw = op - PAD_LEFT`
34/// - if `ip_raw < 0`: wraps to `ip_raw + L`
35/// - if `ip_raw >= L`: wraps to `ip_raw - L`
36/// - else: `ip_raw`
37///
38/// **Constraints**: `PAD_LEFT <= L`, `PAD_RIGHT <= L`.
39#[kernel]
40pub fn circular_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 cond_left = ip_raw.lt(0);
77    let cond_right = ip_raw.ge(L);
78    let in_bounds = ip_raw.ge(0) & ip_raw.lt(L);
79
80    let ip_wrapped_left = ip_raw + L; // ip_raw + L for left pad
81    let ip_wrapped_right = ip_raw - L; // ip_raw - L for right pad
82
83    let zeros = T::zeros::<D>(&[BLOCK_OL]);
84    let val_center = T::load(
85        input_ptr.add_offsets(ip_raw + in_bc_base),
86        Some(ol_mask & in_bounds),
87        Some(zeros),
88        &[],
89        None,
90        None,
91        None,
92        false,
93    );
94    let val_left = T::load(
95        input_ptr.add_offsets(ip_wrapped_left + in_bc_base),
96        Some(ol_mask & cond_left),
97        Some(zeros),
98        &[],
99        None,
100        None,
101        None,
102        false,
103    );
104    let val_right = T::load(
105        input_ptr.add_offsets(ip_wrapped_right + in_bc_base),
106        Some(ol_mask & cond_right),
107        Some(zeros),
108        &[],
109        None,
110        None,
111        None,
112        false,
113    );
114
115    let result = T::where_(
116        cond_left,
117        val_left,
118        T::where_(cond_right, val_right, val_center),
119    );
120
121    let out_offsets = ol_range + out_bc_base;
122    T::store(
123        output_ptr.add_offsets(out_offsets),
124        result,
125        Some(ol_mask),
126        &[],
127        None,
128        None,
129    );
130}
131
132/// 1-D circular padding backward pass.
133///
134/// Multiple output positions may wrap to the same input position, so
135/// `atomic_add` is used. `dx` must be zero-initialised before launch.
136#[kernel]
137pub fn circular_pad1d_backward<
138    T: Triton,
139    D: Num,
140    const PAD_LEFT: i32,
141    const PAD_RIGHT: i32,
142    const BLOCK_OL: i32,
143>(
144    dy_ptr: T::Pointer<D>,
145    dx_ptr: T::Pointer<D>,
146    _B: i32,
147    C: i32,
148    L: i32,
149    OL: i32,
150) where
151    T::I32Tensor: Tensor<i32, 1>,
152    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
153    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
154    T::BoolTensor: BitOr<Output = T::BoolTensor>,
155    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
156{
157    let pid = T::program_id(Axis::X);
158    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
159
160    let ol_tile = pid % num_ol_tiles;
161    let bc = pid / num_ol_tiles;
162    let c = bc % C;
163    let b = bc / C;
164
165    let ol_start = ol_tile * BLOCK_OL;
166    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
167    let ol_mask = ol_range.lt(OL);
168
169    let dy_bc_base = (b * C + c) * OL;
170    let dx_bc_base = (b * C + c) * L;
171
172    let dy_offsets = ol_range + dy_bc_base;
173    let dy_tile = T::load(
174        dy_ptr.add_offsets(dy_offsets),
175        Some(ol_mask),
176        Some(T::zeros::<D>(&[BLOCK_OL])),
177        &[],
178        None,
179        None,
180        None,
181        false,
182    );
183
184    let ip_raw = ol_range - PAD_LEFT;
185    let cond_left = ip_raw.lt(0);
186    let cond_right = ip_raw.ge(L);
187    let in_bounds = ip_raw.ge(0) & ip_raw.lt(L);
188
189    let ip_wrapped_left = ip_raw + L;
190    let ip_wrapped_right = ip_raw - L;
191
192    T::atomic_add(
193        dx_ptr.add_offsets(ip_raw + dx_bc_base),
194        dy_tile,
195        Some(ol_mask & in_bounds),
196        None,
197        None,
198    );
199    T::atomic_add(
200        dx_ptr.add_offsets(ip_wrapped_left + dx_bc_base),
201        dy_tile,
202        Some(ol_mask & cond_left),
203        None,
204        None,
205    );
206    T::atomic_add(
207        dx_ptr.add_offsets(ip_wrapped_right + dx_bc_base),
208        dy_tile,
209        Some(ol_mask & cond_right),
210        None,
211        None,
212    );
213}
214
215pub struct CircularPad1dOp<'a, T: Num> {
216    pub forward: CircularPad1dForward<T>,
217    pub backward: CircularPad1dBackward<T>,
218    _marker: core::marker::PhantomData<&'a ()>,
219}
teeny_kernels/nn/pad/circular_pad1d.rs

teeny_kernels/nn/pad/
circular_pad1d.rs
