burn-cubecl 0.22.0

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

use crate::kernel::memory_order::in_memory_order;
use crate::{
    kernel::utils::{address_type, broadcast_shape},
    ops::{max_vector_size, numeric::empty_device_dtype},
    tensor::CubeTensor,
};
use burn_backend::cubecl::dtype_to_storage_type;
use cubecl::{
    calculate_cube_count_elemwise,
    prelude::*,
    std::tensor::layout::linear::{LinearView, LinearViewMut},
};

pub(crate) trait BinaryOpFloatFamily: Send + Sync + 'static {
    type BinaryOp<C: Float, N: Size>: BinaryOpFloat<C, N>;
}

#[cube]
pub(crate) trait BinaryOpFloat<C: Float, N: Size>: 'static + Send + Sync {
    /// Execute a binary operation.
    fn execute(lhs: Vector<C, N>, rhs: Vector<C, N>) -> Vector<C, N>;
}

pub(crate) struct ArcTan2Op;

impl BinaryOpFloatFamily for ArcTan2Op {
    type BinaryOp<C: Float, N: Size> = Self;
}

#[cube]
impl<T: Float, N: Size> BinaryOpFloat<T, N> for ArcTan2Op {
    fn execute(lhs: Vector<T, N>, rhs: Vector<T, N>) -> Vector<T, N> {
        Vector::atan2(lhs, rhs)
    }
}

#[cube(launch_unchecked, address_type = "dynamic")]
pub(crate) fn kernel_binop<C: Float, N: Size, O: BinaryOpFloatFamily>(
    lhs: LinearView<'_, Vector<C, N>>,
    rhs: LinearView<'_, Vector<C, N>>,
    mut out: LinearViewMut<'_, Vector<C, N>>,
    #[define(C)] _dtype: ElemType,
) {
    if !out.is_in_bounds(ABSOLUTE_POS) {
        terminate!();
    }

    let res = O::BinaryOp::<C, N>::execute(lhs.read(ABSOLUTE_POS), rhs.read(ABSOLUTE_POS));

    out.write(ABSOLUTE_POS, res)
}

pub(crate) fn launch_binop_float<O: BinaryOpFloatFamily>(
    lhs: CubeTensor,
    rhs: CubeTensor,
) -> CubeTensor {
    let output_shape = broadcast_shape(&[&lhs, &rhs]);
    in_memory_order([lhs, rhs], output_shape, |[lhs, rhs], shape_out| {
        let vector_size_lhs = max_vector_size(&lhs);
        let vector_size_rhs = max_vector_size(&rhs);
        let vector_size = Ord::min(vector_size_lhs, vector_size_rhs);

        let dtype = lhs.dtype;

        // A zero-sized broadcast output has no elements to compute, and the in-place/kernel paths
        // below assume a non-empty output. Return the empty output directly.
        if shape_out.num_elements() == 0 {
            return empty_device_dtype(lhs.client.clone(), lhs.device.clone(), shape_out, dtype);
        }

        let client = lhs.client.clone();
        let num_elems = shape_out.num_elements();
        let working_units = num_elems / vector_size as usize;

        let cube_dim = CubeDim::new(&lhs.client, working_units);
        let cube_count = calculate_cube_count_elemwise(&lhs.client, working_units, cube_dim);

        unsafe {
            if lhs.can_mut_broadcast(&rhs) {
                kernel_binop::launch_unchecked::<O>(
                    &client,
                    cube_count,
                    cube_dim,
                    address_type!(lhs, rhs),
                    vector_size,
                    lhs.clone().into_linear_view(),
                    rhs.clone().into_linear_view_like(&lhs),
                    lhs.as_linear_view_alias(0),
                    dtype_to_storage_type(dtype),
                );

                lhs
            } else if rhs.can_mut_broadcast(&lhs) {
                kernel_binop::launch_unchecked::<O>(
                    &client,
                    cube_count,
                    cube_dim,
                    address_type!(lhs, rhs),
                    vector_size,
                    lhs.into_linear_view_like(&rhs),
                    rhs.clone().into_linear_view(),
                    rhs.as_linear_view_alias(1),
                    dtype_to_storage_type(dtype),
                );

                rhs
            } else {
                let output =
                    empty_device_dtype(lhs.client.clone(), lhs.device.clone(), shape_out, dtype);

                kernel_binop::launch_unchecked::<O>(
                    &client,
                    cube_count,
                    cube_dim,
                    address_type!(lhs, rhs, output),
                    vector_size,
                    lhs.into_linear_view_like(&output),
                    rhs.into_linear_view_like(&output),
                    output.clone().into_linear_view(),
                    dtype_to_storage_type(dtype),
                );

                output
            }
        }
    })
}

/// Calculate the four-quadrant inverse tangent of `lhs / rhs`.
pub fn atan2(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor {
    launch_binop_float::<ArcTan2Op>(lhs, rhs)
}