burn-cubecl 0.22.0-pre.4

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

use cubecl::prelude::*;

use crate::{
    kernel::{
        FloatUnaryOp, FloatUnaryOpFamily, NumericUnaryOp, NumericUnaryOpFamily, launch_unary_float,
        launch_unary_numeric,
    },
    tensor::CubeTensor,
};

#[derive(CubeLaunch, CubeType)]
struct Options {
    min_value: InputScalar,
    max_value: InputScalar,
}

/// Clamp a float tensor, keeping NaN inputs.
pub(crate) fn clamp_float(
    input: CubeTensor,
    min_value: InputScalar,
    max_value: InputScalar,
) -> CubeTensor {
    struct ClampOp;

    #[cube]
    impl<F: Float, N: Size> FloatUnaryOp<F, N> for ClampOp {
        type Options = Options;

        fn execute(input: Vector<F, N>, options: &Self::Options) -> Vector<F, N> {
            // clamp lowers to max(min(x, max), min), which returns the non-NaN operand and so
            // mapped NaN to a bound. Comparisons against NaN are false, so it survives here, in
            // the same order as the clamp_min/clamp_max trait default.
            let min_value = Vector::new(options.min_value.get::<F>());
            let max_value = Vector::new(options.max_value.get::<F>());

            let clamped = select(input > max_value, max_value, input);
            select(clamped < min_value, min_value, clamped)
        }
    }

    impl FloatUnaryOpFamily for ClampOp {
        type Options = Options;
        type Unary<F: Float, N: Size> = Self;
    }

    launch_unary_float::<ClampOp, _>(input, |_| OptionsLaunch::new(min_value, max_value))
}

/// Clamp an int tensor.
pub(crate) fn clamp_int(
    input: CubeTensor,
    min_value: InputScalar,
    max_value: InputScalar,
) -> CubeTensor {
    struct ClampOp;

    #[cube]
    impl<T: Numeric, N: Size> NumericUnaryOp<T, N> for ClampOp {
        type Options = Options;

        fn execute(input: Vector<T, N>, options: &Self::Options) -> Vector<T, N> {
            cubecl::prelude::clamp(
                input,
                Vector::new(options.min_value.get::<T>()),
                Vector::new(options.max_value.get::<T>()),
            )
        }
    }

    impl NumericUnaryOpFamily for ClampOp {
        type Options = Options;
        type Unary<T: Numeric, N: Size> = Self;
    }

    launch_unary_numeric::<ClampOp, _>(input, |_| OptionsLaunch::new(min_value, max_value))
}