teeny_kernels/nn/pad/
replication_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 replication padding forward pass.
29///
30/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
31///
32/// Clamps the input index to `[0, L-1]` (replication of the boundary value).
33#[kernel]
34pub fn replication_pad1d_forward<
35    T: Triton,
36    D: Num,
37    const PAD_LEFT: i32,
38    const PAD_RIGHT: i32,
39    const BLOCK_OL: i32,
40>(
41    input_ptr: T::Pointer<D>,
42    output_ptr: T::Pointer<D>,
43    _B: i32,
44    C: i32,
45    L: i32,
46    OL: i32,
47) where
48    T::I32Tensor: Tensor<i32, 1>,
49    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
50    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
51    T::BoolTensor: BitOr<Output = T::BoolTensor>,
52    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
53{
54    let pid = T::program_id(Axis::X);
55    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
56
57    let ol_tile = pid % num_ol_tiles;
58    let bc = pid / num_ol_tiles;
59    let c = bc % C;
60    let b = bc / C;
61
62    let ol_start = ol_tile * BLOCK_OL;
63    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
64    let ol_mask = ol_range.lt(OL);
65
66    let in_bc_base = (b * C + c) * L;
67    let out_bc_base = (b * C + c) * OL;
68
69    let ip_raw = ol_range - PAD_LEFT;
70    let cond_left = ip_raw.lt(0);
71    let cond_right = ip_raw.ge(L);
72    let in_bounds = ip_raw.ge(0) & ip_raw.lt(L);
73
74    // `ip_raw * 0` is the only way to splat a scalar into an I32Tensor (no broadcast API).
75    #[allow(clippy::erasing_op)]
76    let zero_ip = ip_raw * 0;
77    #[allow(clippy::erasing_op)]
78    let lm1_ip = ip_raw * 0 + (L - 1);
79
80    let zeros = T::zeros::<D>(&[BLOCK_OL]);
81    let val_center = T::load(
82        input_ptr.add_offsets(ip_raw + in_bc_base),
83        Some(ol_mask & in_bounds),
84        Some(zeros),
85        &[],
86        None,
87        None,
88        None,
89        false,
90    );
91    let val_left = T::load(
92        input_ptr.add_offsets(zero_ip + in_bc_base),
93        Some(ol_mask & cond_left),
94        Some(zeros),
95        &[],
96        None,
97        None,
98        None,
99        false,
100    );
101    let val_right = T::load(
102        input_ptr.add_offsets(lm1_ip + in_bc_base),
103        Some(ol_mask & cond_right),
104        Some(zeros),
105        &[],
106        None,
107        None,
108        None,
109        false,
110    );
111
112    let result = T::where_(
113        cond_left,
114        val_left,
115        T::where_(cond_right, val_right, val_center),
116    );
117
118    let out_offsets = ol_range + out_bc_base;
119    T::store(
120        output_ptr.add_offsets(out_offsets),
121        result,
122        Some(ol_mask),
123        &[],
124        None,
125        None,
126    );
127}
128
129/// 1-D replication padding backward pass.
130///
131/// Gradients accumulate to boundary elements for padded lanes.
132/// `dx` must be zero-initialised before launch.
133#[kernel]
134pub fn replication_pad1d_backward<
135    T: Triton,
136    D: Num,
137    const PAD_LEFT: i32,
138    const PAD_RIGHT: i32,
139    const BLOCK_OL: i32,
140>(
141    dy_ptr: T::Pointer<D>,
142    dx_ptr: T::Pointer<D>,
143    _B: i32,
144    C: i32,
145    L: i32,
146    OL: i32,
147) where
148    T::I32Tensor: Tensor<i32, 1>,
149    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
150    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
151    T::BoolTensor: BitOr<Output = T::BoolTensor>,
152    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
153{
154    let pid = T::program_id(Axis::X);
155    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
156
157    let ol_tile = pid % num_ol_tiles;
158    let bc = pid / num_ol_tiles;
159    let c = bc % C;
160    let b = bc / C;
161
162    let ol_start = ol_tile * BLOCK_OL;
163    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
164    let ol_mask = ol_range.lt(OL);
165
166    let dy_bc_base = (b * C + c) * OL;
167    let dx_bc_base = (b * C + c) * L;
168
169    let dy_offsets = ol_range + dy_bc_base;
170    let dy_tile = T::load(
171        dy_ptr.add_offsets(dy_offsets),
172        Some(ol_mask),
173        Some(T::zeros::<D>(&[BLOCK_OL])),
174        &[],
175        None,
176        None,
177        None,
178        false,
179    );
180
181    let ip_raw = ol_range - PAD_LEFT;
182    let cond_left = ip_raw.lt(0);
183    let cond_right = ip_raw.ge(L);
184    let in_bounds = ip_raw.ge(0) & ip_raw.lt(L);
185
186    #[allow(clippy::erasing_op)]
187    let zero_ip = ip_raw * 0;
188    #[allow(clippy::erasing_op)]
189    let lm1_ip = ip_raw * 0 + (L - 1);
190
191    T::atomic_add(
192        dx_ptr.add_offsets(ip_raw + dx_bc_base),
193        dy_tile,
194        Some(ol_mask & in_bounds),
195        None,
196        None,
197    );
198    T::atomic_add(
199        dx_ptr.add_offsets(zero_ip + dx_bc_base),
200        dy_tile,
201        Some(ol_mask & cond_left),
202        None,
203        None,
204    );
205    T::atomic_add(
206        dx_ptr.add_offsets(lm1_ip + dx_bc_base),
207        dy_tile,
208        Some(ol_mask & cond_right),
209        None,
210        None,
211    );
212}
213
214pub struct ReplicationPad1dOp<'a, T: Num> {
215    pub forward: ReplicationPad1dForward<T>,
216    pub backward: ReplicationPad1dBackward<T>,
217    _marker: core::marker::PhantomData<&'a ()>,
218}
teeny_kernels/nn/pad/replication_pad1d.rs

teeny_kernels/nn/pad/
replication_pad1d.rs
