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

teeny_kernels/nn/pad/
circular_pad3d.rs
