burn-jit 0.16.1

Generic backend that can be compiled just-in-time to any shader language target
Documentation 
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

use cubecl::{calculate_cube_count_elemwise, prelude::*};

use crate::{tensor::JitTensor, FloatElement, JitRuntime};

#[cube(launch_unchecked)]
fn interpolate_nearest_kernel<F: Float>(input: &Tensor<F>, output: &mut Tensor<F>) {
    if ABSOLUTE_POS >= output.len() {
        return;
    }

    let batch = ABSOLUTE_POS / output.stride(0) % output.shape(0);
    let channel = ABSOLUTE_POS / output.stride(1) % output.shape(1);
    let y = ABSOLUTE_POS / output.stride(2) % output.shape(2);
    let x = ABSOLUTE_POS / output.stride(3) % output.shape(3);

    let factor = F::cast_from(y);
    let numerator = F::cast_from(input.shape(2));
    let denominator = F::cast_from(output.shape(2));
    let y: F = Floor::floor(factor * numerator / denominator);

    let factor = F::cast_from(x);
    let numerator = F::cast_from(input.shape(3));
    let denominator = F::cast_from(output.shape(3));
    let x: F = Floor::floor(factor * numerator / denominator);

    let index = batch * input.stride(0)
        + channel * input.stride(1)
        + u32::cast_from(y) * input.stride(2)
        + u32::cast_from(x) * input.stride(3);

    output[ABSOLUTE_POS] = input[index];
}

pub(crate) fn interpolate_nearest_launch<R: JitRuntime, E: FloatElement>(
    input: JitTensor<R>,
    output: JitTensor<R>,
) -> JitTensor<R> {
    let cube_dim = CubeDim::default();
    let cube_count = calculate_cube_count_elemwise(output.shape.num_elements(), cube_dim);

    unsafe {
        interpolate_nearest_kernel::launch_unchecked::<E, R>(
            &input.client,
            cube_count,
            cube_dim,
            input.as_tensor_arg::<E>(1),
            output.as_tensor_arg::<E>(1),
        )
    };

    output
}