singe-kernel 0.1.0-alpha.4

Reusable CPU and GPU kernels.
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

// @generated by cargo xtask gen-cutile-dtype-kernels
// Do not edit generated entries by hand.

#[cutile::module]
mod kernels {
    use cutile::core::*;
    pub fn softmax_impl<
        const BM: i32,
        const BN: i32,
        const LOG_OUTPUT: bool,
        const NEGATE_INPUT: bool,
    >(
        out: &mut Tensor<f16, { [BM, BN] }>,
        input: &Tensor<f16, { [-1, -1] }>,
    ) {
        let input_tile_raw: Tile<f16, { [BM, BN] }> = load_tile_like(input, out);
        let input_tile_raw: Tile<f32, { [BM, BN] }> = convert_tile(input_tile_raw);
        let zero = constant(0.0f32, out.shape());
        let input_tile_raw = if NEGATE_INPUT {
            zero - input_tile_raw
        } else {
            input_tile_raw
        };
        let col_iota: Tile<i32, { [BN] }> = iota(const_shape![BN]);
        let col_iota = col_iota
            .reshape(const_shape![1i32, BN])
            .broadcast(out.shape());
        let input_shape: [i32; 2] = get_tensor_shape(input);
        let cols_tile = broadcast_scalar(input_shape[1], out.shape());
        let valid_cols = cmpi(col_iota, cols_tile, predicate::LessThan);
        let negative_infinity = constant(f32::NEG_INFINITY, out.shape());
        let input_tile = select(valid_cols, input_tile_raw, negative_infinity);
        let input_max: Tile<f32, { [BM] }> = reduce_max(input_tile, 1i32);
        let input_max = input_max
            .reshape(const_shape![BM, 1i32])
            .broadcast(out.shape());
        let numerator = exp(input_tile - input_max);
        let denominator: Tile<f32, { [BM] }> = reduce_sum(numerator, 1i32);
        let denominator = denominator
            .reshape(const_shape![BM, 1i32])
            .broadcast(out.shape());
        let value = if LOG_OUTPUT {
            input_tile - input_max - log(denominator)
        } else {
            numerator / denominator
        };
        out.store(convert_tile(value));
    }
    #[cutile::entry()]
    pub fn softmax_f16<const BM: i32, const BN: i32>(
        out: &mut Tensor<f16, { [BM, BN] }>,
        input: &Tensor<f16, { [-1, -1] }>,
    ) {
        softmax_impl::<BM, BN, false, false>(out, input);
    }
    #[cutile::entry()]
    pub fn log_softmax_f16<const BM: i32, const BN: i32>(
        out: &mut Tensor<f16, { [BM, BN] }>,
        input: &Tensor<f16, { [-1, -1] }>,
    ) {
        softmax_impl::<BM, BN, true, false>(out, input);
    }
    #[cutile::entry()]
    pub fn softmin_f16<const BM: i32, const BN: i32>(
        out: &mut Tensor<f16, { [BM, BN] }>,
        input: &Tensor<f16, { [-1, -1] }>,
    ) {
        softmax_impl::<BM, BN, false, true>(out, input);
    }
}
pub use crate::cuda::cutile::kernel::f16::utility::softmax_lse_f16;
pub use kernels::*;