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
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

use crate::{element::JitElement, ops::numeric::empty_device, tensor::JitTensor, JitRuntime};
use burn_tensor::Shape;
use cubecl::{calculate_cube_count_elemwise, prelude::*};

#[cube(launch)]
fn adaptive_avg_pool2d_direct<E: Numeric>(input: &Tensor<E>, output: &mut Tensor<E>) {
    let (output_stride_0, output_stride_1, output_stride_2, output_stride_3) = (
        output.stride(0),
        output.stride(1),
        output.stride(2),
        output.stride(3),
    );
    let (output_shape_0, output_shape_1, output_shape_2, output_shape_3) = (
        output.shape(0),
        output.shape(1),
        output.shape(2),
        output.shape(3),
    );
    let (input_stride_0, input_stride_1, input_stride_2, input_stride_3) = (
        input.stride(0),
        input.stride(1),
        input.stride(2),
        input.stride(3),
    );
    let (input_shape_2, input_shape_3) = (input.shape(2), input.shape(3));

    let b = (ABSOLUTE_POS / output_stride_0) % output_shape_0;
    let c = (ABSOLUTE_POS / output_stride_1) % output_shape_1;
    let oh = (ABSOLUTE_POS / output_stride_2) % output_shape_2;
    let ow = (ABSOLUTE_POS / output_stride_3) % output_shape_3;

    let ih_start = start_index(oh, output_shape_2, input_shape_2);
    let ih_end = end_index(oh, output_shape_2, input_shape_2);

    let iw_start = start_index(ow, output_shape_3, input_shape_3);
    let iw_end = end_index(ow, output_shape_3, input_shape_3);

    let mut sum = E::from_int(0);

    let index_input_0 = b * input_stride_0;
    let index_input_1 = c * input_stride_1;

    for ih in ih_start..ih_end {
        let index_input_2 = ih * input_stride_2;

        for iw in iw_start..iw_end {
            let index_input_3 = iw * input_stride_3;

            let index_input = index_input_0 + index_input_1 + index_input_2 + index_input_3;
            sum += input[index_input];
        }
    }

    let num_ih = ih_end - ih_start;
    let num_iw = iw_end - iw_start;

    output[ABSOLUTE_POS] = sum / E::cast_from(num_ih * num_iw);
}

#[cube]
fn start_index(output_size_index: u32, output_size: u32, input_size: u32) -> u32 {
    (output_size_index * input_size) / output_size
}

#[allow(unknown_lints)] // `manual_div_ceil` only appeared in 1.83
#[allow(clippy::manual_div_ceil)]
#[cube]
fn end_index(output_size_index: u32, output_size: u32, input_size: u32) -> u32 {
    let index = (output_size_index + 1) * input_size;
    let index = (index + output_size - 1) / output_size;

    if input_size < index {
        input_size
    } else {
        index
    }
}

pub(crate) fn adaptive_avg_pool2d<R: JitRuntime, E: JitElement>(
    input: JitTensor<R>,
    output_size: [usize; 2],
) -> JitTensor<R> {
    let [batch_size, channels, _, _] = input.shape.dims();

    let output_shape = Shape::new([batch_size, channels, output_size[0], output_size[1]]);
    let num_elems: usize = output_shape.num_elements();
    let output = empty_device::<R, E>(input.client.clone(), input.device.clone(), output_shape);

    let cube_dim = CubeDim::default();
    let cube_count = calculate_cube_count_elemwise(num_elems, cube_dim);

    adaptive_avg_pool2d_direct::launch::<E, R>(
        &input.client,
        cube_count,
        cube_dim,
        input.as_tensor_arg::<E>(1),
        output.as_tensor_arg::<E>(1),
    );

    output
}