tract_cuda/kernels/array/
cast.rs

1use cudarc::driver::{CudaStream, LaunchConfig, PushKernelArg};
2use derive_new::new;
3use std::fmt;
4use tract_core::internal::*;
5use tract_gpu::tensor::DeviceTensor;
6
7use crate::context::{TractCudaStream, cuda_context};
8use crate::kernels::launch_args::TractLaunchArgs;
9use crate::kernels::{LibraryName, get_cuda_view, launch_args};
10
11#[derive(Debug, Clone, new, PartialEq, Eq, Hash)]
12pub struct Cast;
13
14impl fmt::Display for Cast {
15    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
16        write!(f, "{self:?}")
17    }
18}
19
20impl Cast {
21    pub fn is_supported_dt(dt: DatumType) -> bool {
22        matches!(
23            dt,
24            DatumType::F32
25                | DatumType::F16
26                | DatumType::U8
27                | DatumType::U16
28                | DatumType::U32
29                | DatumType::U64
30                | DatumType::I8
31                | DatumType::I16
32                | DatumType::I32
33                | DatumType::I64
34                | DatumType::Bool
35        )
36    }
37
38    pub fn kernel_name(&self, from_dt: DatumType, to_dt: DatumType) -> TractResult<String> {
39        ensure!(
40            Self::is_supported_dt(from_dt),
41            "Unsupported from_dt {:?} for cuda castop",
42            from_dt
43        );
44        ensure!(Self::is_supported_dt(to_dt), "Unsupported to_dt {:?} for cuda castop", to_dt);
45        let from_tname = DeviceTensor::tname(from_dt)?;
46        let to_tname = DeviceTensor::tname(to_dt)?;
47        Ok(format!("cast_{from_tname}_{to_tname}"))
48    }
49
50    pub fn eval(
51        &self,
52        stream: &TractCudaStream,
53        input: &DeviceTensor,
54        to_dt: DatumType,
55    ) -> TractResult<DeviceTensor> {
56        let output = unsafe { DeviceTensor::uninitialized_dt(to_dt, input.shape())? };
57        self.dispatch_eval(stream, input, &output)?;
58        stream.synchronize()?;
59        Ok(output)
60    }
61
62    pub fn dispatch_eval(
63        &self,
64        stream: &TractCudaStream,
65        input: &DeviceTensor,
66        output: &DeviceTensor,
67    ) -> TractResult<()> {
68        ensure!(
69            input.shape() == output.shape(),
70            "Cast I/O don't have the same shape in: {:?}, out: {:?}",
71            input.shape(),
72            output.shape()
73        );
74
75        let kernel_name = self.kernel_name(input.datum_type(), output.datum_type())?;
76
77        let i_view = get_cuda_view(input);
78        let o_view = get_cuda_view(output);
79        let len = output.len();
80        let func = cuda_context().load_pipeline(LibraryName::Array, kernel_name)?;
81
82        let mut launch_args = TractLaunchArgs::new(stream, &func);
83        launch_args.push_view(&i_view);
84        launch_args.push_view(&o_view);
85        launch_args.push_i32(len);
86        let cfg = LaunchConfig::for_num_elems(len as _);
87
88        launch_args.launch(cfg)
89    }
90}
91
92pub fn cuda_cast_dispatch(input: &DeviceTensor, output: &DeviceTensor) -> TractResult<()> {
93    crate::with_cuda_stream(|stream| Cast.dispatch_eval(stream, input, output))
94}
95
96crate::register_cuda_op!(tract_core::ops::cast::Cast, |_source, _node, op| {
97    Ok(crate::transform::cuda_cast_new(op.to).map(|c| Box::new(c) as _))
98});
99
100#[cfg(test)]
101mod tests {
102
103    use super::*;
104    use tract_gpu::tensor::IntoDevice;
105    use tract_itertools::Itertools;
106
107    use num_traits::{FromPrimitive, Zero};
108
109    use tract_core::internal::Tensor;
110
111    fn run_test_case<T0: Datum + Copy + FromPrimitive, T1: Datum>(
112        shape: &[usize],
113    ) -> TractResult<()> {
114        crate::with_cuda_stream(|stream| {
115            let len = shape.iter().product::<usize>();
116            let data = (0..len).map(|f| T0::from_f32(f as f32 / 2.).unwrap()).collect::<Vec<_>>();
117            let input = Tensor::from_shape(shape, &data)?;
118
119            let output = Cast {}.eval(stream, &input.clone().into_device()?, T1::datum_type())?;
120
121            assert_eq!(
122                output.to_host()?.into_tensor(),
123                input.cast_to_dt(T1::datum_type())?.into_owned()
124            );
125            Ok(())
126        })
127    }
128
129    #[test]
130    fn test_cast() -> TractResult<()> {
131        run_test_case::<f16, f32>(&[3, 4])?;
132        run_test_case::<u8, f32>(&[2, 5])?;
133        run_test_case::<f16, u32>(&[3, 2, 2])?;
134        Ok(())
135    }
136}
tract_cuda/kernels/array/cast.rs

tract_cuda/kernels/array/
cast.rs
