burn_cubecl/kernel/grid_sample/
base.rs

1use cubecl::prelude::*;
2
3use crate::tensor::CubeTensor;
4use burn_backend::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode};
5
6use super::bilinear::grid_sample_bilinear_launch;
7
8/// Grid sample operation supporting bilinear interpolation
9pub fn grid_sample(input: CubeTensor, grid: CubeTensor, options: GridSampleOptions) -> CubeTensor {
10    match options.mode {
11        InterpolateMode::Bilinear => grid_sample_bilinear_launch(input, grid, options),
12        _ => panic!(
13            "Unsupported grid_sample interpolation mode: {:?}",
14            options.mode
15        ),
16    }
17}
18
19/// Compile-time padding mode for kernel specialization
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
21pub enum PaddingMode {
22    /// Fill with zeros for out-of-bounds coordinates.
23    Zeros,
24    /// Clamp coordinates to the border (use nearest edge value).
25    Border,
26    /// Reflect coordinates at the boundary.
27    Reflection,
28}
29
30impl From<GridSamplePaddingMode> for PaddingMode {
31    fn from(mode: GridSamplePaddingMode) -> Self {
32        match mode {
33            GridSamplePaddingMode::Zeros => PaddingMode::Zeros,
34            GridSamplePaddingMode::Border => PaddingMode::Border,
35            GridSamplePaddingMode::Reflection => PaddingMode::Reflection,
36        }
37    }
38}
39
40/// Fetch value based on padding mode (dispatch to appropriate handler)
41#[cube]
42pub(crate) fn fetch_value<F: Float>(
43    input: &Tensor<F>,
44    base: usize,
45    stride_h: usize,
46    stride_w: usize,
47    y: i32,
48    x: i32,
49    h: i32,
50    w: i32,
51    #[comptime] padding_mode: PaddingMode,
52) -> F {
53    match padding_mode {
54        PaddingMode::Zeros => fetch_with_zeros(input, base, stride_h, stride_w, y, x, h, w),
55        PaddingMode::Border => fetch_with_border(input, base, stride_h, stride_w, y, x, h, w),
56        PaddingMode::Reflection => {
57            fetch_with_reflection(input, base, stride_h, stride_w, y, x, h, w)
58        }
59    }
60}
61
62/// Fetch value with zeros padding (return 0 for out-of-bounds).
63#[cube]
64pub(crate) fn fetch_with_zeros<F: Float>(
65    input: &Tensor<F>,
66    base: usize,
67    stride_h: usize,
68    stride_w: usize,
69    y: i32,
70    x: i32,
71    h: i32,
72    w: i32,
73) -> F {
74    let in_bounds = x >= 0 && x < w && y >= 0 && y < h;
75    let x_clamped = clamp(x, 0, w - 1) as usize;
76    let y_clamped = clamp(y, 0, h - 1) as usize;
77    let idx = base + y_clamped * stride_h + x_clamped * stride_w;
78    select(in_bounds, input[idx], F::new(0.0_f32))
79}
80
81/// Fetch value with border padding (clamp to edge).
82#[cube]
83pub(crate) fn fetch_with_border<F: Float>(
84    input: &Tensor<F>,
85    base: usize,
86    stride_h: usize,
87    stride_w: usize,
88    y: i32,
89    x: i32,
90    h: i32,
91    w: i32,
92) -> F {
93    let x_clamped = clamp(x, 0, w - 1) as usize;
94    let y_clamped = clamp(y, 0, h - 1) as usize;
95    let idx = base + y_clamped * stride_h + x_clamped * stride_w;
96    input[idx]
97}
98
99/// Fetch value with reflection padding.
100/// Assumes float reflection was applied to center, so indices are at most 2 steps out of bounds.
101#[cube]
102pub(crate) fn fetch_with_reflection<F: Float>(
103    input: &Tensor<F>,
104    base: usize,
105    stride_h: usize,
106    stride_w: usize,
107    y: i32,
108    x: i32,
109    h: i32,
110    w: i32,
111) -> F {
112    let x_reflected = reflect_coord_bounded(x, w);
113    let y_reflected = reflect_coord_bounded(y, h);
114    let idx = base + y_reflected * stride_h + x_reflected * stride_w;
115    input[idx]
116}
117
118/// Reflect an integer index that may be out of bounds.
119/// After float reflection, indices can be up to 2 steps out for bicubic (1 step for bilinear).
120#[cube]
121fn reflect_coord_bounded(idx: i32, size: i32) -> usize {
122    let max_idx = size - 1;
123    let neg_reflected = -idx - 1;
124    let pos_reflected = 2 * max_idx + 1 - idx;
125    let result = select(
126        idx < 0,
127        neg_reflected,
128        select(idx > max_idx, pos_reflected, idx),
129    );
130    clamp(result, 0, max_idx) as usize
131}
132
133/// Reflect a float coordinate into the valid sampling range.
134#[cube]
135pub(crate) fn reflect_coord<F: Float>(coord: F, size: u32, #[comptime] align_corners: bool) -> F {
136    let size_f = F::cast_from(size);
137    if align_corners {
138        reflect_float_impl::<F>(coord, F::new(0.0_f32), size_f - F::new(1.0_f32))
139    } else {
140        reflect_float_impl::<F>(coord, F::new(-0.5_f32), size_f - F::new(0.5_f32))
141    }
142}
143
144/// Reflect a float coordinate into [min_val, max_val] using a triangle wave pattern.
145#[cube]
146fn reflect_float_impl<F: Float>(coord: F, min_val: F, max_val: F) -> F {
147    let span = max_val - min_val;
148
149    let is_valid = span > F::new(0.0_f32);
150    let safe_span = select(is_valid, span, F::new(1.0_f32));
151
152    // Triangle wave formula: span - |((x mod 2*span) - span)| + min_val
153    let period = safe_span * F::new(2.0_f32);
154    let x = (coord - min_val).abs();
155    let x_mod = x - (x / period).floor() * period;
156    let reflected = safe_span - (x_mod - safe_span).abs() + min_val;
157
158    select(is_valid, reflected, min_val)
159}
burn_cubecl/kernel/grid_sample/base.rs

burn_cubecl/kernel/grid_sample/
base.rs
