cubek-interpolate 0.3.0

CubeK: Interpolation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

mod nearest;

pub(crate) use nearest::reference_nearest_backward;

use super::{f32_elem_type, make_random_f32_host, make_zero_handle};
use crate::definition::{InterpolateBackwardProblem, InterpolateMode, InterpolateOptions};
use cubecl::client::Client;
use cubek_test_utils::{
    ExecutionOutcome, HostData, HostDataType, Progress, launch_and_capture_outcome,
};

use crate::interpolate_backward;

pub fn kernel_result(
    client: Client,
    problem: InterpolateBackwardProblem,
    seed: u64,
) -> Result<HostData, String> {
    let dtype = f32_elem_type();

    let out_grad_shape = problem.out_grad_shape.to_vec();

    let input_shape = vec![
        out_grad_shape[0],
        problem.input_size[0],
        problem.input_size[1],
        out_grad_shape[3],
    ];
    let (input_handle, _input_host) = make_random_f32_host(&client, input_shape.clone(), seed);

    let (out_grad_handle, _out_grad_host) =
        make_random_f32_host(&client, out_grad_shape.clone(), seed);

    let input_grad_shape = input_shape;
    let input_grad_handle = make_zero_handle(&client, input_grad_shape, dtype);

    let outcome = launch_and_capture_outcome(&client, &[&input_grad_handle.handle], |c| {
        interpolate_backward(
            c,
            input_handle.clone().binding(),
            out_grad_handle.clone().binding(),
            input_grad_handle.clone().binding(),
            problem.options,
            dtype,
        )
        .into()
    });

    match outcome {
        ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")),
        ExecutionOutcome::Executed => Ok(HostData::from_tensor_handle(
            &client,
            input_grad_handle,
            HostDataType::F32,
        )),
    }
}

pub fn cpu_reference_result(
    client: Client,
    problem: InterpolateBackwardProblem,
    seed: u64,
    progress: Option<&Progress>,
) -> Result<HostData, String> {
    let out_grad_shape = problem.out_grad_shape.to_vec();

    if let Some(p) = progress {
        let total: usize = out_grad_shape.iter().product();
        p.set_total(total as u64);
    }

    let (_out_grad_handle, out_grad_host) = make_random_f32_host(&client, out_grad_shape, seed);

    let input_grad_shape = vec![
        problem.out_grad_shape[0],
        problem.input_size[0],
        problem.input_size[1],
        problem.out_grad_shape[3],
    ];

    Ok(reference_for_backward_interpolation_mode(
        &out_grad_host,
        &input_grad_shape,
        &problem.options,
        progress,
    ))
}

pub fn reference_for_backward_interpolation_mode(
    out_grad: &HostData,
    output_shape: &[usize],
    options: &InterpolateOptions,
    progress: Option<&Progress>,
) -> HostData {
    match options.mode {
        InterpolateMode::Nearest(nearest_mode) => {
            reference_nearest_backward(out_grad, output_shape, nearest_mode, progress)
        }
        InterpolateMode::Bilinear => {
            panic!("Bilinear interpolation backward is not supported by CPU reference")
        }
        InterpolateMode::Bicubic => {
            panic!("Bicubic interpolation backward is not supported by CPU reference")
        }
        InterpolateMode::Lanczos3 => {
            panic!("Lanczos3 interpolation backward is not supported by CPU reference")
        }
    }
}