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

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

#[cutile::module]
mod kernels {
    use cutile::core::*;
    #[cutile::entry()]
    pub fn quantize_f16_u8<const B: i32>(
        out: &mut Tensor<u8, { [B] }>,
        input: &Tensor<f16, { [-1] }>,
        scale: f16,
        zero_point: f16,
    ) {
        let values: Tile<f16, { [B] }> = load_tile_like(input, out);
        let raw = values / broadcast_scalar(scale, out.shape())
            + broadcast_scalar(zero_point, out.shape());
        let low = constant(f16::from_f32(0.0), out.shape());
        let high = constant(f16::from_f32(255.0), out.shape());
        out.store(ftoi(
            min_tile(max_tile(raw, low), high),
            rounding::NearestEven,
        ));
    }
    #[cutile::entry()]
    pub fn quantize_f16_i8<const B: i32>(
        out: &mut Tensor<i8, { [B] }>,
        input: &Tensor<f16, { [-1] }>,
        scale: f16,
        zero_point: f16,
    ) {
        let values: Tile<f16, { [B] }> = load_tile_like(input, out);
        let raw = values / broadcast_scalar(scale, out.shape())
            + broadcast_scalar(zero_point, out.shape());
        let low = constant(f16::from_f32(-128.0), out.shape());
        let high = constant(f16::from_f32(127.0), out.shape());
        out.store(ftoi(
            min_tile(max_tile(raw, low), high),
            rounding::NearestEven,
        ));
    }
}
pub use crate::cuda::cutile::kernel::f16::utility::dequantize_i8_f16;
pub use crate::cuda::cutile::kernel::f16::utility::dequantize_u8_f16;
pub use kernels::*;