teeny_kernels/nn/pad/
constant_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 constant padding forward pass.
29///
30/// Grid: `pid = (b * C + c) * num_ol_tiles + ol_tile`
31///
32/// Output length `OL = PAD_LEFT + L + PAD_RIGHT`.
33/// Positions outside `[PAD_LEFT, PAD_LEFT + L)` are filled with `value`.
34#[kernel]
35pub fn constant_pad1d_forward<
36    T: Triton,
37    D: Num,
38    const PAD_LEFT: i32,
39    const PAD_RIGHT: i32,
40    const BLOCK_OL: i32,
41>(
42    input_ptr: T::Pointer<D>,
43    output_ptr: T::Pointer<D>,
44    _B: i32,
45    C: i32,
46    L: i32,
47    OL: i32,
48    value: f32,
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_ol_tiles = T::cdiv(OL, BLOCK_OL);
58
59    let ol_tile = pid % num_ol_tiles;
60    let bc = pid / num_ol_tiles;
61    let c = bc % C;
62    let b = bc / C;
63
64    let ol_start = ol_tile * BLOCK_OL;
65    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
66    let ol_mask = ol_range.lt(OL);
67
68    let in_bc_base = (b * C + c) * L;
69    let out_bc_base = (b * C + c) * OL;
70
71    let ip_range = ol_range - PAD_LEFT;
72    let in_bounds = ip_range.ge(0) & ip_range.lt(L);
73    let combined_mask = ol_mask & in_bounds;
74
75    let value_vec = T::cast::<f32, D>(T::full::<f32>(&[BLOCK_OL], value), None, false);
76    let tile = T::load(
77        input_ptr.add_offsets(ip_range + in_bc_base),
78        Some(combined_mask),
79        Some(value_vec),
80        &[],
81        None,
82        None,
83        None,
84        false,
85    );
86    let result = T::where_(ol_mask & in_bounds, tile, value_vec);
87
88    let out_offsets = ol_range + out_bc_base;
89    T::store(
90        output_ptr.add_offsets(out_offsets),
91        result,
92        Some(ol_mask),
93        &[],
94        None,
95        None,
96    );
97}
98
99/// 1-D constant padding backward pass.
100///
101/// Gradient flows only through input positions; padded positions have zero grad.
102/// Tile over output positions and store `dy` back to `dx` at input offset.
103/// `dx` must be zero-initialised before launch.
104#[kernel]
105pub fn constant_pad1d_backward<
106    T: Triton,
107    D: Num,
108    const PAD_LEFT: i32,
109    const PAD_RIGHT: i32,
110    const BLOCK_OL: i32,
111>(
112    dy_ptr: T::Pointer<D>,
113    dx_ptr: T::Pointer<D>,
114    _B: i32,
115    C: i32,
116    L: i32,
117    OL: i32,
118) where
119    T::I32Tensor: Tensor<i32, 1>,
120    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
121    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
122    T::BoolTensor: BitOr<Output = T::BoolTensor>,
123    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
124{
125    let pid = T::program_id(Axis::X);
126    let num_ol_tiles = T::cdiv(OL, BLOCK_OL);
127
128    let ol_tile = pid % num_ol_tiles;
129    let bc = pid / num_ol_tiles;
130    let c = bc % C;
131    let b = bc / C;
132
133    let ol_start = ol_tile * BLOCK_OL;
134    let ol_range = T::arange(0, BLOCK_OL) + ol_start;
135    let ol_mask = ol_range.lt(OL);
136
137    let dy_bc_base = (b * C + c) * OL;
138    let dx_bc_base = (b * C + c) * L;
139
140    let ip_range = ol_range - PAD_LEFT;
141    let in_bounds = ip_range.ge(0) & ip_range.lt(L);
142    let load_mask = ol_mask & in_bounds;
143
144    let dy_offsets = ol_range + dy_bc_base;
145    let dy_tile = T::load(
146        dy_ptr.add_offsets(dy_offsets),
147        Some(ol_mask),
148        Some(T::zeros::<D>(&[BLOCK_OL])),
149        &[],
150        None,
151        None,
152        None,
153        false,
154    );
155
156    let dx_offsets = ip_range + dx_bc_base;
157    T::store(
158        dx_ptr.add_offsets(dx_offsets),
159        dy_tile,
160        Some(load_mask),
161        &[],
162        None,
163        None,
164    );
165}
166
167pub struct ConstantPad1dOp<'a, T: Num> {
168    pub forward: ConstantPad1dForward<T>,
169    pub backward: ConstantPad1dBackward<T>,
170    _marker: core::marker::PhantomData<&'a ()>,
171}
teeny_kernels/nn/pad/constant_pad1d.rs

teeny_kernels/nn/pad/
constant_pad1d.rs
