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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use cubecl::{calculate_cube_count_elemwise, linalg::tensor::index_offset_with_layout, prelude::*};

use crate::{
    element::JitElement,
    ops::{max_vectorization, numeric::empty_device},
    tensor::JitTensor,
    BoolElement, JitRuntime,
};

#[cube(launch)]
fn mask_where_readonly_kernel<T: CubePrimitive, B: Int>(
    input: &Tensor<Line<T>>,
    mask: &Tensor<Line<B>>,
    value: &Tensor<Line<T>>,
    output: &mut Tensor<Line<T>>,
    #[comptime] rank: u32,
) {
    if ABSOLUTE_POS >= output.len() {
        return;
    }

    let index_input = index_offset_with_layout(input, output, ABSOLUTE_POS, 0, rank, true);
    let index_mask = index_offset_with_layout(mask, output, ABSOLUTE_POS, 0, rank, true);
    let index_value = index_offset_with_layout(value, output, ABSOLUTE_POS, 0, rank, true);
    let mask = Line::cast_from(mask[index_mask]);

    output[ABSOLUTE_POS] = select_many(mask, value[index_value], input[index_input]);
}

#[cube(launch)]
fn mask_where_inplace_kernel<T: CubePrimitive, B: Int>(
    input: &mut Tensor<Line<T>>,
    mask: &Tensor<Line<B>>,
    value: &Tensor<Line<T>>,
    reverse: B,
    #[comptime] rank: u32,
) {
    if ABSOLUTE_POS >= input.len() {
        return;
    }

    let index_mask = index_offset_with_layout(mask, input, ABSOLUTE_POS, 0, rank, true);
    let index_value = index_offset_with_layout(value, input, ABSOLUTE_POS, 0, rank, true);

    input[ABSOLUTE_POS] = select(
        mask[index_mask] != Line::new(reverse),
        value[index_value],
        input[ABSOLUTE_POS],
    );
}

#[derive(Clone, Copy, Debug)]
/// Define how to run the mask where kernel.
///
/// # Notes
///
/// All assertions should be done before choosing the strategy.
pub enum MaskWhereStrategy {
    /// Don't mutate any input.
    Readonly,
    /// Reuse the lhs tensor inplace.
    InplaceLhs,
    /// Reuse the rhs tensor inplace.
    InplaceRhs,
}

/// Execute the mask where kernel with the given strategy.
pub fn mask_where<R: JitRuntime, E: JitElement, BT: BoolElement>(
    input: JitTensor<R>,
    mask: JitTensor<R>,
    value: JitTensor<R>,
    strategy: MaskWhereStrategy,
) -> JitTensor<R> {
    match strategy {
        MaskWhereStrategy::Readonly => mask_where_readonly::<R, E, BT>(input, mask, value),
        MaskWhereStrategy::InplaceLhs => mask_where_inplace::<R, E, BT>(input, mask, value, false),
        MaskWhereStrategy::InplaceRhs => mask_where_inplace::<R, E, BT>(value, mask, input, true),
    }
}

fn mask_where_readonly<R: JitRuntime, EI: JitElement, EM: BoolElement>(
    input: JitTensor<R>,
    mask: JitTensor<R>,
    value: JitTensor<R>,
) -> JitTensor<R> {
    let ndims = input.shape.num_dims();
    let output = empty_device::<R, EI>(
        input.client.clone(),
        input.device.clone(),
        input.shape.clone(),
    );

    let cube_dim = CubeDim::default();
    let cube_count = calculate_cube_count_elemwise(input.shape.num_elements(), cube_dim);
    let vectorization = max_vectorization(&input);

    mask_where_readonly_kernel::launch::<EI, EM, R>(
        &input.client,
        cube_count,
        cube_dim,
        input.as_tensor_arg::<EI>(vectorization),
        mask.as_tensor_arg::<EM>(vectorization),
        value.as_tensor_arg::<EI>(vectorization),
        output.as_tensor_arg::<EI>(vectorization),
        ndims as u32,
    );

    output
}

fn mask_where_inplace<R: JitRuntime, EI: JitElement, EM: BoolElement>(
    input: JitTensor<R>,
    mask: JitTensor<R>,
    value: JitTensor<R>,
    reverse: bool,
) -> JitTensor<R> {
    let ndims = input.shape.num_dims();
    let cube_dim = CubeDim::default();
    let cube_count = calculate_cube_count_elemwise(input.shape.num_elements(), cube_dim);
    let vectorization = max_vectorization(&input);

    mask_where_inplace_kernel::launch::<EI, EM, R>(
        &input.client,
        cube_count,
        cube_dim,
        input.as_tensor_arg::<EI>(vectorization),
        mask.as_tensor_arg::<EM>(vectorization),
        value.as_tensor_arg::<EI>(vectorization),
        ScalarArg::new(EM::new_bool(reverse)),
        ndims as u32,
    );

    input
}