teeny_kernels/nn/pad/
replication_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 replication padding forward pass.
29///
30/// Grid: `pid = (((b*C+c)*OD+od)*OH+oh)*num_ow_tiles + ow_tile`
31#[kernel]
32pub fn replication_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 = if id_raw < 0 {
78        0
79    } else if id_raw >= Dv {
80        Dv - 1
81    } else {
82        id_raw
83    };
84
85    let ih_raw = oh - PH1;
86    let ih = if ih_raw < 0 {
87        0
88    } else if ih_raw >= H {
89        H - 1
90    } else {
91        ih_raw
92    };
93
94    let in_bc_base = ((b * C + c) * Dv + id) * H * W + ih * W;
95    let out_bc_base = (((b * C + c) * OD + od) * OH + oh) * OW;
96
97    let iw_raw = ow_range - PW1;
98    let cond_left = iw_raw.lt(0);
99    let cond_right = iw_raw.ge(W);
100    let in_bounds = iw_raw.ge(0) & iw_raw.lt(W);
101
102    #[allow(clippy::erasing_op)]
103    let zero_iw = iw_raw * 0;
104    #[allow(clippy::erasing_op)]
105    let wm1_iw = iw_raw * 0 + (W - 1);
106
107    let zeros = T::zeros::<D>(&[BLOCK_OW]);
108    let val_center = T::load(
109        input_ptr.add_offsets(iw_raw + in_bc_base),
110        Some(ow_mask & in_bounds),
111        Some(zeros),
112        &[],
113        None,
114        None,
115        None,
116        false,
117    );
118    let val_left = T::load(
119        input_ptr.add_offsets(zero_iw + in_bc_base),
120        Some(ow_mask & cond_left),
121        Some(zeros),
122        &[],
123        None,
124        None,
125        None,
126        false,
127    );
128    let val_right = T::load(
129        input_ptr.add_offsets(wm1_iw + in_bc_base),
130        Some(ow_mask & cond_right),
131        Some(zeros),
132        &[],
133        None,
134        None,
135        None,
136        false,
137    );
138
139    let result = T::where_(
140        cond_left,
141        val_left,
142        T::where_(cond_right, val_right, val_center),
143    );
144
145    let out_offsets = ow_range + out_bc_base;
146    T::store(
147        output_ptr.add_offsets(out_offsets),
148        result,
149        Some(ow_mask),
150        &[],
151        None,
152        None,
153    );
154}
155
156/// 3-D replication padding backward pass.
157#[kernel]
158pub fn replication_pad3d_backward<
159    T: Triton,
160    D: Num,
161    const PD1: i32,
162    const PD2: i32,
163    const PH1: i32,
164    const PH2: i32,
165    const PW1: i32,
166    const PW2: i32,
167    const BLOCK_OW: i32,
168>(
169    dy_ptr: T::Pointer<D>,
170    dx_ptr: T::Pointer<D>,
171    _B: i32,
172    C: i32,
173    Dv: i32,
174    H: i32,
175    W: i32,
176    OD: i32,
177    OH: i32,
178    OW: i32,
179) where
180    T::I32Tensor: Tensor<i32, 1>,
181    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
182    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
183    T::BoolTensor: BitOr<Output = T::BoolTensor>,
184    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
185{
186    let pid = T::program_id(Axis::X);
187    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
188
189    let ow_tile = pid % num_ow_tiles;
190    let rest = pid / num_ow_tiles;
191    let oh = rest % OH;
192    let rest2 = rest / OH;
193    let od = rest2 % OD;
194    let bco = rest2 / OD;
195    let c = bco % C;
196    let b = bco / C;
197
198    let ow_start = ow_tile * BLOCK_OW;
199    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
200    let ow_mask = ow_range.lt(OW);
201
202    let id_raw = od - PD1;
203    let id = if id_raw < 0 {
204        0
205    } else if id_raw >= Dv {
206        Dv - 1
207    } else {
208        id_raw
209    };
210
211    let ih_raw = oh - PH1;
212    let ih = if ih_raw < 0 {
213        0
214    } else if ih_raw >= H {
215        H - 1
216    } else {
217        ih_raw
218    };
219
220    let dy_bc_base = (((b * C + c) * OD + od) * OH + oh) * OW;
221    let dx_bc_base = ((b * C + c) * Dv + id) * H * W + ih * W;
222
223    let dy_offsets = ow_range + dy_bc_base;
224    let dy_tile = T::load(
225        dy_ptr.add_offsets(dy_offsets),
226        Some(ow_mask),
227        Some(T::zeros::<D>(&[BLOCK_OW])),
228        &[],
229        None,
230        None,
231        None,
232        false,
233    );
234
235    let iw_raw = ow_range - PW1;
236    let cond_left = iw_raw.lt(0);
237    let cond_right = iw_raw.ge(W);
238    let in_bounds = iw_raw.ge(0) & iw_raw.lt(W);
239
240    #[allow(clippy::erasing_op)]
241    let zero_iw = iw_raw * 0;
242    #[allow(clippy::erasing_op)]
243    let wm1_iw = iw_raw * 0 + (W - 1);
244
245    T::atomic_add(
246        dx_ptr.add_offsets(iw_raw + dx_bc_base),
247        dy_tile,
248        Some(ow_mask & in_bounds),
249        None,
250        None,
251    );
252    T::atomic_add(
253        dx_ptr.add_offsets(zero_iw + dx_bc_base),
254        dy_tile,
255        Some(ow_mask & cond_left),
256        None,
257        None,
258    );
259    T::atomic_add(
260        dx_ptr.add_offsets(wm1_iw + dx_bc_base),
261        dy_tile,
262        Some(ow_mask & cond_right),
263        None,
264        None,
265    );
266}
267
268pub struct ReplicationPad3dOp<'a, T: Num> {
269    pub forward: ReplicationPad3dForward<T>,
270    pub backward: ReplicationPad3dBackward<T>,
271    _marker: core::marker::PhantomData<&'a ()>,
272}
teeny_kernels/nn/pad/replication_pad3d.rs

teeny_kernels/nn/pad/
replication_pad3d.rs
