teeny_kernels/nn/tensor/
upsample_nearest2d.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::marker::PhantomData;
20use teeny_core::dtype::Num;
21use teeny_macros::kernel;
22use teeny_triton::triton::{
23    types::{AddOffsets, Comparison, Tensor},
24    *,
25};
26
27/// Nearest-neighbour 2-D upsample forward pass — NCHW layout.
28///
29/// Grid: one flat pid per (b, c, oh, ow-tile):
30///   `pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile`
31///
32/// Each CTA writes a BLOCK_OW-wide strip of output columns by reading from
33/// the nearest input position via floor division:
34///   `ih = oh / SCALE_H`,  `iw = ow / SCALE_W`
35///
36/// Output shape: `[B, C, OH=H*SCALE_H, OW=W*SCALE_W]`.
37#[kernel]
38pub fn upsample_nearest2d_forward<
39    T: Triton,
40    D: Num,
41    const SCALE_H: i32,
42    const SCALE_W: i32,
43    const BLOCK_OW: i32,
44>(
45    x_ptr: T::Pointer<D>, // input  [B, C, H, W]
46    y_ptr: T::Pointer<D>, // output [B, C, OH, OW]
47    _B: i32,
48    C: i32,
49    H: i32,
50    W: 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::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
57{
58    let pid = T::program_id(Axis::X);
59    let num_ow_tiles = T::cdiv(OW, BLOCK_OW);
60
61    // Decode flat pid → (b, c, oh, ow_tile).
62    let ow_tile = pid % num_ow_tiles;
63    let bco = pid / num_ow_tiles;
64    let oh = bco % OH;
65    let bc = bco / OH;
66    let c = bc % C;
67    let b = bc / C;
68
69    // Nearest-neighbour source row (scalar).
70    let ih = oh / SCALE_H;
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    // Nearest-neighbour source column per output lane: iw = ow / SCALE_W.
77    let iw_range = ow_range / SCALE_W;
78
79    let in_bc_base = (b * C + c) * H * W;
80    let out_bc_base = (b * C + c) * OH * OW;
81
82    let in_offsets = iw_range + (in_bc_base + ih * W);
83    let out_offsets = ow_range + (out_bc_base + oh * OW);
84
85    let x = T::load(
86        x_ptr.add_offsets(in_offsets),
87        Some(ow_mask),
88        None,
89        &[],
90        None,
91        None,
92        None,
93        false,
94    );
95    T::store(
96        y_ptr.add_offsets(out_offsets),
97        x,
98        Some(ow_mask),
99        &[],
100        None,
101        None,
102    );
103}
104
105/// Nearest-neighbour 2-D upsample backward pass — NCHW layout.
106///
107/// Grid: one flat pid per (b, c, ih, iw-tile) — same spatial extent as the
108/// *input* tensor:
109///   `pid = ((b * C + c) * H + ih) * num_iw_tiles + iw_tile`
110///
111/// Each CTA computes the gradient for a BLOCK_IW-wide strip of input columns
112/// by summing the SCALE_H × SCALE_W upstream gradients that each input pixel
113/// received during the forward pass.  No atomic operations are needed because
114/// each input element is the sole accumulator for exactly SCALE_H×SCALE_W
115/// output lanes (their ranges are disjoint).
116///
117/// `dx` does NOT need to be zero-initialised (every element is written once).
118#[kernel]
119pub fn upsample_nearest2d_backward<
120    T: Triton,
121    D: Num,
122    const SCALE_H: i32,
123    const SCALE_W: i32,
124    const BLOCK_IW: i32,
125>(
126    dy_ptr: T::Pointer<D>, // upstream grad [B, C, OH, OW]
127    dx_ptr: T::Pointer<D>, // input    grad [B, C, H,  W]
128    _B: i32,
129    C: i32,
130    H: i32,
131    W: i32,
132    OH: i32,
133    OW: i32,
134) where
135    T::I32Tensor: Tensor<i32, 1>,
136    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
137    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
138{
139    let pid = T::program_id(Axis::X);
140    let num_iw_tiles = T::cdiv(W, BLOCK_IW);
141
142    // Decode flat pid → (b, c, ih, iw_tile).
143    let iw_tile = pid % num_iw_tiles;
144    let bch = pid / num_iw_tiles;
145    let ih = bch % H;
146    let bc = bch / H;
147    let c = bc % C;
148    let b = bc / C;
149
150    let iw_start = iw_tile * BLOCK_IW;
151    let iw_range = T::arange(0, BLOCK_IW) + iw_start;
152    let iw_mask = iw_range.lt(W);
153
154    let dy_bc_base = (b * C + c) * OH * OW;
155    let dx_bc_base = (b * C + c) * H * W;
156
157    let mut acc = T::zeros::<D>(&[BLOCK_IW]);
158
159    // Flat loop over SCALE_H × SCALE_W upstream gradient positions.
160    let loop_bound = SCALE_H * SCALE_W;
161    for idx in 0..loop_bound {
162        let sw = idx % SCALE_W;
163        let sh = idx / SCALE_W;
164        let oh = ih * SCALE_H + sh; // scalar row
165        let ow_range = iw_range * SCALE_W + sw; // tensor cols
166        let dy_offsets = ow_range + (dy_bc_base + oh * OW);
167        let dy_mask = ow_range.lt(OW);
168        let tile = T::load(
169            dy_ptr.add_offsets(dy_offsets),
170            Some(dy_mask),
171            Some(T::zeros::<D>(&[BLOCK_IW])),
172            &[],
173            None,
174            None,
175            None,
176            false,
177        );
178        acc = acc + tile;
179    }
180
181    let dx_offsets = iw_range + (dx_bc_base + ih * W);
182    T::store(
183        dx_ptr.add_offsets(dx_offsets),
184        acc,
185        Some(iw_mask),
186        &[],
187        None,
188        None,
189    );
190}
191
192impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for UpsampleNearest2dForward<D> {
193    fn n_activation_inputs(&self) -> usize {
194        1
195    }
196
197    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
198        Vec::new()
199    }
200
201    fn pack_args(
202        &self,
203        inputs: &[(teeny_core::model::RawPtr, &[usize])],
204        _params: &[teeny_core::model::RawPtr],
205        output: teeny_core::model::RawPtr,
206        output_shape: &[usize],
207        _output_row_stride: i32,
208        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
209    ) {
210        // kernel args: x_ptr, y_ptr, B, C, H, W, OH, OW
211        let input_shape = inputs[0].1;
212        visitor.visit_ptr(inputs[0].0);
213        visitor.visit_ptr(output);
214        visitor.visit_i32(input_shape[0] as i32); // B
215        visitor.visit_i32(input_shape[1] as i32); // C
216        visitor.visit_i32(input_shape[2] as i32); // H
217        visitor.visit_i32(input_shape[3] as i32); // W
218        visitor.visit_i32(output_shape[2] as i32); // OH
219        visitor.visit_i32(output_shape[3] as i32); // OW
220    }
221
222    fn block(&self) -> [u32; 3] {
223        [128, 1, 1]
224    }
225
226    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
227        // pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile
228        let num_ow_tiles = output_shape[3].div_ceil(self.block_ow as usize);
229        [
230            (output_shape[0] * output_shape[1] * output_shape[2] * num_ow_tiles) as u32,
231            1,
232            1,
233        ]
234    }
235
236    #[cfg(feature = "training")]
237    fn has_backward(&self) -> bool {
238        true
239    }
240
241    /// kernel args: dy, dx, B, C, H, W, OH, OW
242    #[cfg(feature = "training")]
243    fn pack_backward_args(
244        &self,
245        inputs: &[(teeny_core::model::RawPtr, &[usize])],
246        _params: &[teeny_core::model::RawPtr],
247        _output: teeny_core::model::RawPtr,
248        output_shape: &[usize],
249        grad_output: teeny_core::model::RawPtr,
250        _grad_output_row_stride: i32,
251        grad_inputs: &[teeny_core::model::RawPtr],
252        _grad_params: &[teeny_core::model::RawPtr],
253        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
254    ) {
255        let in_shape = inputs[0].1; // [B, C, H, W]
256        visitor.visit_ptr(grad_output); // dy_ptr
257        visitor.visit_ptr(grad_inputs[0]); // dx_ptr
258        visitor.visit_i32(in_shape[0] as i32); // B
259        visitor.visit_i32(in_shape[1] as i32); // C
260        visitor.visit_i32(in_shape[2] as i32); // H
261        visitor.visit_i32(in_shape[3] as i32); // W
262        visitor.visit_i32(output_shape[2] as i32); // OH
263        visitor.visit_i32(output_shape[3] as i32); // OW
264    }
265
266    #[cfg(feature = "training")]
267    fn backward_block(&self) -> [u32; 3] {
268        [128, 1, 1]
269    }
270
271    /// Grid over input spatial positions: `pid = ((b * C + c) * H + ih) * num_iw_tiles + iw_tile`
272    #[cfg(feature = "training")]
273    fn backward_grid(&self, input_shapes: &[&[usize]], _output_shape: &[usize]) -> [u32; 3] {
274        let in_shape = input_shapes[0]; // [B, C, H, W]
275        let num_iw_tiles = in_shape[3].div_ceil(self.block_ow as usize);
276        [
277            (in_shape[0] * in_shape[1] * in_shape[2] * num_iw_tiles) as u32,
278            1,
279            1,
280        ]
281    }
282}
283
284pub struct UpsampleNearest2dOp<'a, D: Num> {
285    pub forward: UpsampleNearest2dForward<D>,
286    pub backward: UpsampleNearest2dBackward<D>,
287    _marker: PhantomData<&'a ()>,
288}
teeny_kernels/nn/tensor/upsample_nearest2d.rs

teeny_kernels/nn/tensor/
upsample_nearest2d.rs
