tract_cuda/kernels/array/
rotate_half.rs

1use crate::context::{TractCudaStream, cuda_context};
2use crate::kernels::launch_args::TractLaunchArgs;
3use crate::kernels::{LibraryName, get_cuda_view, utils};
4use anyhow::ensure;
5use cudarc::driver::{CudaStream, LaunchConfig, PushKernelArg};
6use std::fmt;
7use tract_core::internal::*;
8use tract_gpu::tensor::DeviceTensor;
9
10#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub struct RotateHalf;
12
13impl fmt::Display for RotateHalf {
14    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
15        write!(f, "{self:?}")
16    }
17}
18
19impl RotateHalf {
20    pub fn is_supported_dt(dt: DatumType) -> bool {
21        matches!(
22            dt,
23            DatumType::F32
24                | DatumType::F16
25                | DatumType::I8
26                | DatumType::I16
27                | DatumType::I32
28                | DatumType::I64
29        )
30    }
31
32    pub fn kernel_name(&self, dt: DatumType) -> TractResult<String> {
33        ensure!(Self::is_supported_dt(dt), "Unsupported dt {:?} for cuda rotate halfop", dt);
34        let tname = DeviceTensor::tname(dt)?;
35        Ok(format!("rotate_half_nd2_{tname}"))
36    }
37
38    pub fn eval(
39        &self,
40        stream: &TractCudaStream,
41        input: &DeviceTensor,
42    ) -> TractResult<DeviceTensor> {
43        let output = unsafe { DeviceTensor::uninitialized_dt(input.datum_type(), input.shape())? };
44        self.dispatch_eval(stream, input, &output)?;
45        stream.synchronize()?;
46        Ok(output)
47    }
48
49    pub fn dispatch_eval(
50        &self,
51        stream: &TractCudaStream,
52        input: &DeviceTensor,
53        output: &DeviceTensor,
54    ) -> TractResult<()> {
55        let shape_nd2 = utils::reshape_to_rank_2(input.shape(), input.rank() - 1);
56        ensure!(
57            shape_nd2[1].is_multiple_of(2),
58            "Rotate half required most inner dimension to be a multiple of 2: {:?}",
59            input.shape()
60        );
61        let strides_nd2 = Tensor::natural_strides(&shape_nd2);
62
63        let kernel_name = self.kernel_name(input.datum_type())?;
64
65        let func = cuda_context().load_pipeline(LibraryName::Array, kernel_name)?;
66
67        let i_view = get_cuda_view(input);
68        let o_view = get_cuda_view(output);
69
70        let mut launch_args = TractLaunchArgs::new(stream, &func);
71        launch_args.push_view(&i_view);
72        launch_args.push_view(&o_view);
73        launch_args.push_slice_i32(&shape_nd2);
74        launch_args.push_slice_i32(&strides_nd2);
75
76        let cfg = LaunchConfig {
77            grid_dim: ((shape_nd2[1] / 2) as _, shape_nd2[0] as _, 1),
78            block_dim: (1, 1, 1),
79            shared_mem_bytes: 0,
80        };
81        launch_args.launch(cfg)
82    }
83}
84
85pub fn cuda_rotate_half_dispatch(input: &DeviceTensor, output: &DeviceTensor) -> TractResult<()> {
86    crate::with_cuda_stream(|stream| RotateHalf.dispatch_eval(stream, input, output))
87}
88
89crate::register_cuda_op!(tract_transformers::ops::apply_rope::RotateHalf, |source, node, _op| {
90    rule_if!(RotateHalf::is_supported_dt(source.node_input_facts(node.id)?[0].datum_type));
91    Ok(Some(Box::new(tract_gpu::ops::rotate_half::GpuRotateHalf::new(
92        "Cuda",
93        cuda_rotate_half_dispatch,
94    ))))
95});
96
97#[cfg(test)]
98mod tests {
99
100    use super::*;
101    use num_traits::AsPrimitive;
102    use tract_core::internal::Tensor;
103    use tract_gpu::tensor::IntoDevice;
104    use tract_transformers::ops::apply_rope;
105
106    fn run_test_case<F>(shape: &[usize]) -> TractResult<()>
107    where
108        F: Copy + 'static + Datum,
109        usize: AsPrimitive<F>,
110    {
111        crate::with_cuda_stream(|stream| {
112            let len = shape.iter().product::<usize>();
113
114            let a =
115                Tensor::from_shape(shape, &(0..len).map(|f| -> F { f.as_() }).collect::<Vec<_>>())?;
116
117            let cuda_a = a.clone().into_device()?;
118
119            let cpu_output = apply_rope::RotateHalf
120                .eval(&EvalContext::out_of_plan(), tvec![a.clone().into()])?[0]
121                .clone()
122                .into_tensor();
123            let cuda_output = RotateHalf.eval(stream, &cuda_a)?;
124
125            cpu_output
126                .close_enough(&cuda_output.to_host()?.into_tensor(), Approximation::Exact)
127                .with_context(|| {
128                    format!(
129                        "Input: {:?} Cpu: {:?}, Cuda: {:?}",
130                        a.dump(true),
131                        cpu_output.dump(true),
132                        cuda_output.to_host().and_then(|it| it.dump(true))
133                    )
134                })?;
135            Ok(())
136        })
137    }
138
139    #[test]
140    fn test_rotate_half() -> TractResult<()> {
141        run_test_case::<f32>(&[2, 2])?;
142        run_test_case::<f32>(&[512, 512])?;
143        run_test_case::<f32>(&[10, 8, 8])?;
144        run_test_case::<f32>(&[10, 512, 1024])?;
145        run_test_case::<f32>(&[10, 512, 1024])?;
146        run_test_case::<f16>(&[10, 256, 4])?;
147        Ok(())
148    }
149}
tract_cuda/kernels/array/rotate_half.rs

tract_cuda/kernels/array/
rotate_half.rs
