teeny_kernels/nn/pad/
constant_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 constant padding forward pass.
29///
30/// Grid: `pid = (((b*C+c)*OD+od)*OH+oh)*num_ow_tiles + ow_tile`
31///
32/// `OD = PD1+D+PD2`, `OH = PH1+H+PH2`, `OW = PW1+W+PW2`.
33#[kernel]
34pub fn constant_pad3d_forward<
35    T: Triton,
36    D: Num,
37    const PD1: i32,
38    const PD2: i32,
39    const PH1: i32,
40    const PH2: i32,
41    const PW1: i32,
42    const PW2: i32,
43    const BLOCK_OW: i32,
44>(
45    input_ptr: T::Pointer<D>,
46    output_ptr: T::Pointer<D>,
47    _B: i32,
48    C: i32,
49    Dv: i32,
50    H: i32,
51    W: i32,
52    OD: i32,
53    OH: i32,
54    OW: i32,
55    value: f32,
56) where
57    T::I32Tensor: Tensor<i32, 1>,
58    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
59    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
60    T::BoolTensor: BitOr<Output = T::BoolTensor>,
61    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
62{
63    let pid = T::program_id(Axis::X);
64    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
65
66    let ow_tile = pid % num_ow_tiles;
67    let rest = pid / num_ow_tiles;
68    let oh = rest % OH;
69    let rest2 = rest / OH;
70    let od = rest2 % OD;
71    let bco = rest2 / OD;
72    let c = bco % C;
73    let b = bco / C;
74
75    let ow_start = ow_tile * BLOCK_OW;
76    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
77    let ow_mask = ow_range.lt(OW);
78
79    let id = od - PD1;
80    let ih = oh - PH1;
81    let iw_range = ow_range - PW1;
82
83    let value_vec = T::cast::<f32, D>(T::full::<f32>(&[BLOCK_OW], value), None, false);
84
85    // `ow_range * 0` is the only way to splat a scalar into an I32Tensor (no broadcast API).
86    // A compound `if id < 0 || ...` with 4 conditions triggers a compiler phi-node bug
87    // (cond_br: phi local not in ssa_values), so keep the branchless mask approach here.
88    #[allow(clippy::erasing_op)]
89    let id_t = ow_range * 0 + id;
90    #[allow(clippy::erasing_op)]
91    let ih_t = ow_range * 0 + ih;
92    let d_in_bounds = id_t.ge(0) & id_t.lt(Dv);
93    let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
94    let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
95
96    let in_bc_base = ((b * C + c) * Dv + id) * H * W + ih * W;
97    let out_bc_base = (((b * C + c) * OD + od) * OH + oh) * OW;
98
99    let combined_mask = ow_mask & d_in_bounds & h_in_bounds & w_in_bounds;
100
101    let tile = T::load(
102        input_ptr.add_offsets(iw_range + in_bc_base),
103        Some(combined_mask),
104        Some(value_vec),
105        &[],
106        None,
107        None,
108        None,
109        false,
110    );
111    let result = T::where_(combined_mask, tile, value_vec);
112
113    let out_offsets = ow_range + out_bc_base;
114    T::store(
115        output_ptr.add_offsets(out_offsets),
116        result,
117        Some(ow_mask),
118        &[],
119        None,
120        None,
121    );
122}
123
124/// 3-D constant padding backward pass.
125#[kernel]
126pub fn constant_pad3d_backward<
127    T: Triton,
128    D: Num,
129    const PD1: i32,
130    const PD2: i32,
131    const PH1: i32,
132    const PH2: i32,
133    const PW1: i32,
134    const PW2: i32,
135    const BLOCK_OW: i32,
136>(
137    dy_ptr: T::Pointer<D>,
138    dx_ptr: T::Pointer<D>,
139    _B: i32,
140    C: i32,
141    Dv: i32,
142    H: i32,
143    W: i32,
144    OD: i32,
145    OH: i32,
146    OW: 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_ow_tiles = T::cdiv(OW, BLOCK_OW);
156
157    let ow_tile = pid % num_ow_tiles;
158    let rest = pid / num_ow_tiles;
159    let oh = rest % OH;
160    let rest2 = rest / OH;
161    let od = rest2 % OD;
162    let bco = rest2 / OD;
163    let c = bco % C;
164    let b = bco / C;
165
166    let ow_start = ow_tile * BLOCK_OW;
167    let ow_range = T::arange(0, BLOCK_OW) + ow_start;
168    let ow_mask = ow_range.lt(OW);
169
170    let id = od - PD1;
171    let ih = oh - PH1;
172    let iw_range = ow_range - PW1;
173
174    #[allow(clippy::erasing_op)]
175    let id_t = ow_range * 0 + id;
176    #[allow(clippy::erasing_op)]
177    let ih_t = ow_range * 0 + ih;
178    let d_in_bounds = id_t.ge(0) & id_t.lt(Dv);
179    let h_in_bounds = ih_t.ge(0) & ih_t.lt(H);
180    let w_in_bounds = iw_range.ge(0) & iw_range.lt(W);
181
182    let dy_bc_base = (((b * C + c) * OD + od) * OH + oh) * OW;
183    let dx_bc_base = ((b * C + c) * Dv + id) * H * W + ih * W;
184
185    let store_mask = ow_mask & d_in_bounds & h_in_bounds & w_in_bounds;
186
187    let dy_offsets = ow_range + dy_bc_base;
188    let dy_tile = T::load(
189        dy_ptr.add_offsets(dy_offsets),
190        Some(ow_mask),
191        Some(T::zeros::<D>(&[BLOCK_OW])),
192        &[],
193        None,
194        None,
195        None,
196        false,
197    );
198
199    let dx_offsets = iw_range + dx_bc_base;
200    T::store(
201        dx_ptr.add_offsets(dx_offsets),
202        dy_tile,
203        Some(store_mask),
204        &[],
205        None,
206        None,
207    );
208}
209
210pub struct ConstantPad3dOp<'a, T: Num> {
211    pub forward: ConstantPad3dForward<T>,
212    pub backward: ConstantPad3dBackward<T>,
213    _marker: core::marker::PhantomData<&'a ()>,
214}
teeny_kernels/nn/pad/constant_pad3d.rs

teeny_kernels/nn/pad/
constant_pad3d.rs
