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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use crate::{
    element::JitElement,
    kernel::{self},
    tensor::JitTensor,
    IntElement, JitRuntime,
};
use cubecl::prelude::*;
use cubecl::{calculate_cube_count_elemwise, CubeDim};

#[cube(launch_unchecked)]
fn scatter_kernel<T: Numeric, I: Int>(
    input: &mut Tensor<T>,
    indices: &Tensor<I>,
    value: &Tensor<T>,
    dim: &u32,
) {
    let stride_input = input.stride(*dim);
    let shape_value = value.shape(*dim);

    let mut offset_input = 0;
    let mut offset_value = 0;
    let mut num_elems = 1;

    for i in 0..value.rank() {
        let shouldnt_skip = i != *dim;
        if shouldnt_skip {
            let shape_input_loop = input.shape(i);
            let shape_value_loop = value.shape(i);

            let stride_value_loop = value.stride(i);
            let stride_input_loop = input.stride(i);
            let stride_tmp = indices.stride(i);

            let mut num_blocks = ABSOLUTE_POS / stride_tmp;
            num_blocks %= shape_input_loop;

            let mut offset_tmp = num_blocks * stride_input_loop;
            offset_input += offset_tmp;

            offset_tmp = num_blocks * stride_value_loop;
            offset_value += offset_tmp;

            num_elems *= shape_value_loop;
        }
    }

    let should_stop = ABSOLUTE_POS >= num_elems;
    if should_stop {
        return;
    }

    for i in 0..shape_value {
        let mut idx = stride_input * i;
        idx += offset_value;

        let result_value = value[idx];
        let result_indices = u32::cast_from(indices[idx]);

        let mut index_input = stride_input * result_indices;
        index_input += offset_input;

        let mut result_input = input[index_input];
        result_input += result_value;
        input[index_input] = result_input;
    }
}

pub(crate) fn scatter<R: JitRuntime, E: JitElement, I: IntElement>(
    dim: usize,
    tensor: JitTensor<R>,
    indices: JitTensor<R>,
    value: JitTensor<R>,
) -> JitTensor<R> {
    let ndims = tensor.shape.num_dims();
    let mut indices = kernel::into_contiguous(indices);
    let tensor = kernel::into_contiguous(tensor);
    let value = kernel::into_contiguous(value);

    let tensor = match tensor.can_mut() {
        true => tensor,
        false => tensor.copy(),
    };

    let mut strides = vec![0; ndims];
    let mut current = 1;
    let mut num_elems = 1;

    tensor
        .shape
        .dims
        .iter()
        .enumerate()
        .rev()
        .filter(|(index, _val)| *index != dim)
        .for_each(|(index, val)| {
            strides[index] = current;
            current *= val;
            num_elems *= tensor.shape.dims[index];
        });

    // Fake strides of the virtual output where the strides of dim is hardcoded to one.
    indices.strides = strides;

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

    unsafe {
        scatter_kernel::launch_unchecked::<E, I, R>(
            &indices.client.clone(),
            cube_count,
            cube_dim,
            tensor.as_tensor_arg::<E>(1),
            indices.as_tensor_arg::<I>(1),
            value.as_tensor_arg::<E>(1),
            ScalarArg::new(dim as u32),
        )
    }
    tensor
}