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
110

mod backward;
mod forward;

use crate::{
    InterpolateStrategy,
    definition::{InterpolateOptions, InterpolateProblem},
};
use cubecl::std::tensor::TensorHandle;
use cubecl::{client::Client, prelude::*, zspace::Strides};
use cubek_test_utils::{HostData, Progress, TestInput};

pub(crate) fn f32_elem_type() -> ElemType {
    f32::elem_type_native()
}

pub(crate) fn make_random_f32_host(
    client: &Client,
    shape: Vec<usize>,
    seed: u64,
) -> (TensorHandle, HostData) {
    TestInput::builder(client.clone(), shape)
        .uniform(seed, -1., 1.)
        .generate_with_f32_host_data()
}

pub(crate) fn make_zero_handle(
    client: &Client,
    shape: Vec<usize>,
    dtype: ElemType,
) -> TensorHandle {
    TestInput::builder(client.clone(), shape)
        .dtype(dtype)
        .zeros()
        .generate_without_host_data()
}

pub fn kernel_result(
    client: Client,
    problem: InterpolateProblem,
    strategy: InterpolateStrategy,
    seed: u64,
) -> Result<HostData, String> {
    match problem {
        InterpolateProblem::Forward(prob) => forward::kernel_result(client, prob, strategy, seed),
        InterpolateProblem::Backward(prob) => backward::kernel_result(client, prob, seed),
    }
}

pub fn cpu_reference_result(
    client: Client,
    problem: InterpolateProblem,
    seed: u64,
    progress: Option<&Progress>,
) -> Result<HostData, String> {
    match problem {
        InterpolateProblem::Forward(prob) => {
            forward::cpu_reference_result(client, prob, seed, progress)
        }
        InterpolateProblem::Backward(prob) => {
            backward::cpu_reference_result(client, prob, seed, progress)
        }
    }
}

pub fn cpu_reference_interpolate_from_host(
    input: &HostData,
    output_shape: &[usize],
    options: &InterpolateOptions,
) -> HostData {
    forward::reference_for_interpolation_mode(input, output_shape, options, None)
}

pub fn cpu_reference_interpolate_backward_from_host(
    out_grad: &HostData,
    output_shape: &[usize],
    options: &InterpolateOptions,
) -> HostData {
    backward::reference_for_backward_interpolation_mode(out_grad, output_shape, options, None)
}

pub(crate) fn for_each_output_coord(output_shape: &[usize], mut f: impl FnMut(usize, &[usize])) {
    let rank = output_shape.len();
    if rank == 0 {
        f(0, &[]);
        return;
    }
    let num: usize = output_shape.iter().product();
    let mut coord = vec![0usize; rank];
    for linear in 0..num {
        let mut rem = linear;
        for d in (0..rank).rev() {
            coord[d] = rem % output_shape[d];
            rem /= output_shape[d];
        }
        f(linear, &coord);
    }
}

pub(crate) fn contiguous_strides(shape: &[usize]) -> Strides {
    let n = shape.len();
    if n == 0 {
        return Strides::new(&[] as &[usize]);
    }
    let mut s = vec![0usize; n];
    s[n - 1] = 1;
    for i in (0..n - 1).rev() {
        s[i] = s[i + 1] * shape[i + 1];
    }
    Strides::new(&s)
}