tract_cuda/kernels/array/
mod.rs

1mod cast;
2mod copy;
3mod diag_gather;
4mod dispatch;
5mod gather;
6mod rotate_half;
7
8pub use cast::Cast;
9pub use cast::cuda_cast_dispatch;
10pub use copy::Memcpy;
11pub use diag_gather::DiagGather;
12pub use diag_gather::cuda_diag_gather_dispatch;
13pub use dispatch::cuda_copy_nd_dispatch;
14pub use gather::Gather;
15pub use gather::cuda_gather_dispatch;
16pub use rotate_half::RotateHalf;
17pub use rotate_half::cuda_rotate_half_dispatch;
18
19pub fn all_functions() -> Vec<String> {
20    use std::collections::HashSet;
21    use tract_gpu::utils::BroadcastKind;
22    let mut functions = HashSet::<String>::new();
23
24    functions.extend(BroadcastKind::all_copy_kernel_names(""));
25
26    functions.extend(
27        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
28            .into_iter()
29            .flat_map(|dt1| {
30                tract_gpu::tensor::DeviceTensor::SUPPORTED_DT.into_iter().map(move |dt2| (dt1, dt2))
31            })
32            .flat_map(|(dt1, dt2)| Cast.kernel_name(dt1, dt2).into_iter()),
33    );
34
35    functions.extend(
36        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
37            .into_iter()
38            .flat_map(|dt| DiagGather.kernel_name(dt).into_iter()),
39    );
40
41    functions.extend(
42        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
43            .into_iter()
44            .flat_map(|dt| Gather.kernel_name(dt).into_iter()),
45    );
46
47    functions.into_iter().collect()
48}
tract_cuda/kernels/array/mod.rs

tract_cuda/kernels/array/
mod.rs
