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

teeny_kernels/nn/pad/
replication_pad2d.rs
