Skip to main content

upsample_nearest2d_forward

teeny_kernels::nn::tensor::upsample_nearest2d

Function upsample_nearest2d_forward 

Source 
pub fn upsample_nearest2d_forward<T: Triton, D: Num, const SCALE_H: i32, const SCALE_W: i32, const BLOCK_OW: i32>(
    x_ptr: T::Pointer<D>,
    y_ptr: T::Pointer<D>,
    _B: i32,
    C: i32,
    H: i32,
    W: i32,
    OH: i32,
    OW: i32,
)
where
    T::I32Tensor: Tensor<i32, 1> + Comparison<i32, BoolTensor = T::BoolTensor>,
    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
Expand description

Nearest-neighbour 2-D upsample forward pass — NCHW layout.

Grid: one flat pid per (b, c, oh, ow-tile): pid = ((b * C + c) * OH + oh) * num_ow_tiles + ow_tile

Each CTA writes a BLOCK_OW-wide strip of output columns by reading from the nearest input position via floor division: ih = oh / SCALE_H, iw = ow / SCALE_W

Output shape: [B, C, OH=H*SCALE_H, OW=W*SCALE_W].