burn_cubecl/kernel/interpolate/
base.rs

1#[cfg(feature = "autotune")]
2use crate::kernel::interpolate::interpolate_autotune;
3use crate::{
4    kernel::into_contiguous,
5    ops::{numeric::empty_device_dtype, permute_nchw_to_nhwc, permute_nhwc_to_nchw},
6    tensor::CubeTensor,
7};
8use burn_backend::cubecl::dtype_to_storage_type;
9use burn_backend::{Shape, TensorMetadata, ops::InterpolateMode, ops::InterpolateOptions};
10use cubek::interpolate::{
11    InterpolateStrategy as CubekInterpolateStrategy,
12    definition::{
13        InterpolateError, InterpolateMode as CubekInterpolateMode,
14        InterpolateOptions as CubekInterpolateOptions, NearestMode as CubekNearestMode,
15    },
16    interpolate as cubek_interpolate, interpolate_backward as cubek_interpolate_backward,
17};
18
19#[derive(Debug)]
20/// Strategy used to select how interpolation runs.
21pub enum InterpolateStrategy {
22    /// Run the strategy given rather than searching for one. cubek resolves it against the device
23    /// and the problem, so an intent states what the launch optimizes for and leaves the geometry
24    /// to cubek, while [`Forced`](CubekInterpolateStrategy::Forced) pins the geometry outright.
25    Specific(CubekInterpolateStrategy),
26
27    /// Automatically benchmark and select the best strategy at runtime.
28    #[cfg(feature = "autotune")]
29    Autotune,
30}
31
32impl Default for InterpolateStrategy {
33    fn default() -> Self {
34        // if autotune is enabled, default to autotune
35        #[cfg(feature = "autotune")]
36        return InterpolateStrategy::Autotune;
37
38        // Interpolation reads one tensor and writes another, so a build that measures nothing
39        // runs the intent that takes memory for the limit.
40        #[cfg(not(feature = "autotune"))]
41        InterpolateStrategy::Specific(CubekInterpolateStrategy::MaximizeThroughput)
42    }
43}
44
45/// Interpolate operation
46///
47/// Supports nearest, bilinear, bicubic and lanczos3 modes
48pub fn interpolate(
49    input: CubeTensor,
50    output_size: [usize; 2],
51    options: InterpolateOptions,
52    strategy: InterpolateStrategy,
53) -> Result<CubeTensor, InterpolateError> {
54    match strategy {
55        InterpolateStrategy::Specific(strategy) => {
56            execute_interpolate(input, output_size, options, strategy)
57        }
58        #[cfg(feature = "autotune")]
59        InterpolateStrategy::Autotune => Ok(interpolate_autotune(input, output_size, options)),
60    }
61}
62
63/// Execute interpolation with the given strategy, without autotuning. This is used by the
64/// autotune implementation to run each candidate strategy.
65pub fn execute_interpolate(
66    input: CubeTensor,
67    output_size: [usize; 2],
68    options: InterpolateOptions,
69    strategy: CubekInterpolateStrategy,
70) -> Result<CubeTensor, InterpolateError> {
71    let [batch_size, channels, _, _] = input.meta.shape().dims();
72    let [out_height, out_width] = output_size;
73
74    let input = into_contiguous(permute_nchw_to_nhwc(input));
75
76    let shape_out = Shape::new([batch_size, out_height, out_width, channels]);
77    let output = empty_device_dtype(
78        input.client.clone(),
79        input.device.clone(),
80        shape_out,
81        input.dtype,
82    );
83
84    cubek_interpolate(
85        &input.client.clone(),
86        input.clone().binding(),
87        output.clone().binding(),
88        map_options(options.clone()),
89        strategy,
90        dtype_to_storage_type(input.dtype),
91    )?;
92
93    Ok(permute_nhwc_to_nchw(output))
94}
95
96/// Backward interpolate operation
97///
98/// Note: only nearest mode is supported
99pub fn interpolate_backward(
100    input: CubeTensor,
101    out_grad: CubeTensor,
102    _output_size: [usize; 2],
103    options: InterpolateOptions,
104) -> CubeTensor {
105    let input = permute_nchw_to_nhwc(input);
106    let out_grad = permute_nchw_to_nhwc(out_grad);
107
108    let output_shape = input.shape();
109    let output = empty_device_dtype(
110        input.client.clone(),
111        input.device.clone(),
112        output_shape,
113        input.dtype,
114    );
115
116    cubek_interpolate_backward(
117        &input.client.clone(),
118        input.clone().binding(),
119        out_grad.binding(),
120        output.clone().binding(),
121        map_options(options.clone()),
122        dtype_to_storage_type(input.dtype),
123    )
124    .unwrap_or_else(|e| {
125        panic!(
126            "interpolate_backward kernel failed (device={0:?}, dtype={1:?}, options={2:?}): {3}",
127            input.device, input.dtype, options, e
128        )
129    });
130
131    permute_nhwc_to_nchw(output)
132}
133
134pub(crate) fn map_mode(mode: InterpolateMode) -> CubekInterpolateMode {
135    match mode {
136        InterpolateMode::Nearest => CubekInterpolateMode::Nearest(CubekNearestMode::Floor),
137        InterpolateMode::NearestExact => CubekInterpolateMode::Nearest(CubekNearestMode::Exact),
138        InterpolateMode::Bilinear => CubekInterpolateMode::Bilinear,
139        InterpolateMode::Bicubic => CubekInterpolateMode::Bicubic,
140        InterpolateMode::Lanczos3 => CubekInterpolateMode::Lanczos3,
141    }
142}
143
144pub(crate) fn map_options(options: InterpolateOptions) -> CubekInterpolateOptions {
145    CubekInterpolateOptions {
146        mode: map_mode(options.mode),
147        align_corners: options.align_corners,
148    }
149}
burn_cubecl/kernel/interpolate/base.rs

burn_cubecl/kernel/interpolate/
base.rs
