Skip to main content

cubek_interpolate/
lib.rs

1use core::result::Result;
2
3use cubecl::{Runtime, client::ComputeClient, prelude::TensorBinding, prelude::*};
4
5use crate::definition::{InterpolateMode, InterpolateOptions};
6
7pub mod definition;
8mod error;
9mod modes;
10pub use error::InterpolateError;
11
12use crate::modes::bicubic::interpolate_bicubic_launch;
13use crate::modes::bilinear::interpolate_bilinear_launch;
14use crate::modes::lanczos3::interpolate_lanczos3_launch;
15use crate::modes::nearest::interpolate_nearest_launch;
16
17#[cfg(feature = "cpu-reference")]
18pub mod cpu_reference;
19
20/// Interpolate operation
21///
22/// Supports nearest, bilinear, bicubic and lanczos3 modes.
23///
24/// Expects input in NHWC layout.
25pub fn interpolate<R: Runtime>(
26    client: &ComputeClient<R>,
27    input: TensorBinding<R>,
28    output: TensorBinding<R>,
29    options: InterpolateOptions,
30    dtype: StorageType,
31) -> Result<(), InterpolateError> {
32    let _align_corners = options.align_corners;
33
34    match options.mode {
35        InterpolateMode::Nearest => interpolate_nearest_launch(client, input, output, dtype),
36        InterpolateMode::Bilinear => {
37            interpolate_bilinear_launch(client, input, output, _align_corners, dtype)
38        }
39        InterpolateMode::Bicubic => {
40            interpolate_bicubic_launch(client, input, output, _align_corners, dtype)
41        }
42        InterpolateMode::Lanczos3 => {
43            interpolate_lanczos3_launch(client, input, output, _align_corners, dtype)
44        }
45    }
46}