tract-cuda 0.23.8

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
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

mod apply_rope;
mod gelu_approximate;
mod leaky_relu;
mod reduce;
mod rms_norm;
mod scaled_masked_softmax;
mod softmax;

pub use apply_rope::{ApplyRope, cuda_apply_rope_dispatch};
pub use gelu_approximate::GeluApproximate;
pub use gelu_approximate::cuda_gelu_approximate_dispatch;
pub use leaky_relu::LeakyRelu;
pub use leaky_relu::cuda_leaky_relu_dispatch;
pub use reduce::{Reducer, cuda_reduce_launch};
pub use rms_norm::RmsNorm;
pub use rms_norm::cuda_rms_norm_dispatch;
pub use scaled_masked_softmax::{ScaledMaskedSoftmax, cuda_scaled_masked_softmax_dispatch};
pub use softmax::Softmax;
pub use softmax::cuda_softmax_dispatch;

use crate::kernels::{BroadcastKind, MAX_THREADS};

fn sms_block_sizes() -> Vec<i32> {
    let mut range = vec![0i32];

    for i in 5..=15 {
        range.push(2i32.pow(i));
    }

    range
}

pub fn all_functions() -> Vec<String> {
    use std::collections::HashSet;
    let mut functions = HashSet::<String>::new();

    functions.extend(
        Reducer::ALL
            .into_iter()
            .flat_map(|op| {
                tract_gpu::tensor::DeviceTensor::SUPPORTED_DT.into_iter().map(move |dt| (op, dt))
            })
            .flat_map(|(op, dt)| [0, MAX_THREADS].into_iter().map(move |n_cols| (op, dt, n_cols)))
            .flat_map(|(op, dt, n_cols)| reduce::kernel_name(&op, dt, n_cols).into_iter()),
    );
    functions.extend(
        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
            .into_iter()
            .flat_map(|dt| [0, MAX_THREADS].into_iter().map(move |n_cols| (dt, n_cols)))
            .flat_map(|(dt, n_cols)| Softmax.kernel_name(dt, n_cols).into_iter()),
    );

    functions.extend(
        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
            .into_iter()
            .flat_map(|dt| sms_block_sizes().into_iter().map(move |bs| (dt, bs as usize)))
            .flat_map(|(dt, bs)| {
                [false, true]
                    .into_iter()
                    .flat_map(move |mb| ScaledMaskedSoftmax.kernel_name(dt, mb, bs).into_iter())
            }),
    );

    functions.extend(
        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
            .into_iter()
            .flat_map(|in_dt| {
                tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
                    .into_iter()
                    .map(move |out_dt| (in_dt, out_dt))
            })
            .flat_map(|(in_dt, out_dt)| {
                [0, MAX_THREADS].into_iter().map(move |n_cols| (in_dt, out_dt, n_cols))
            })
            .flat_map(|(in_dt, out_dt, n_cols)| {
                [false, true].into_iter().flat_map(move |scaled| {
                    [false, true]
                        .into_iter()
                        .map(move |with_residual| (in_dt, out_dt, n_cols, scaled, with_residual))
                })
            })
            .flat_map(|(in_dt, out_dt, n_cols, scaled, with_residual)| {
                [false, true]
                    .into_iter()
                    .map(move |round| (in_dt, out_dt, n_cols, scaled, with_residual, round))
            })
            .flat_map(|(in_dt, out_dt, n_cols, scaled, with_residual, round)| {
                RmsNorm.kernel_name(in_dt, out_dt, n_cols, scaled, with_residual, round).into_iter()
            }),
    );

    functions.extend(
        BroadcastKind::ALL
            .into_iter()
            .flat_map(|brdcast| {
                tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
                    .into_iter()
                    .map(move |dt| (dt, brdcast))
            })
            .flat_map(|(dt, brdcast)| ApplyRope.kernel_name(dt, brdcast).into_iter()),
    );

    functions.extend(
        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
            .into_iter()
            .flat_map(|dt| [true, false].into_iter().map(move |fast_impl| (dt, fast_impl)))
            .flat_map(|(dt, fast_impl)| GeluApproximate { fast_impl }.kernel_name(dt).into_iter()),
    );

    functions.extend(
        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
            .into_iter()
            .flat_map(|dt| LeakyRelu.kernel_name(dt).into_iter()),
    );

    functions.into_iter().collect()
}