tract_cuda/kernels/array/
diag_gather.rs

1use crate::context::{TractCudaStream, cuda_context};
2use crate::kernels::launch_args::TractLaunchArgs;
3use crate::kernels::{LibraryName, MAX_THREADS, get_cuda_view};
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 DiagGather;
12
13impl fmt::Display for DiagGather {
14    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
15        write!(f, "{self:?}")
16    }
17}
18
19impl DiagGather {
20    pub fn is_supported_dt(dt: DatumType) -> bool {
21        matches!(dt, DatumType::F32 | DatumType::F16)
22    }
23
24    pub fn kernel_name(&self, dt: DatumType) -> TractResult<String> {
25        ensure!(Self::is_supported_dt(dt), "Unsupported dt {:?} for cuda diag_gather op", dt);
26        let tname = DeviceTensor::tname(dt)?;
27        Ok(format!("diag_gather_{tname}"))
28    }
29
30    pub fn eval(
31        &self,
32        stream: &TractCudaStream,
33        input: &DeviceTensor,
34        offset: i64,
35        out_len: usize,
36    ) -> TractResult<DeviceTensor> {
37        let rank = input.rank();
38        ensure!(rank >= 2);
39        let mut out_shape: TVec<usize> = input.shape().into();
40        out_shape[rank - 1] = out_len;
41        let output = unsafe { DeviceTensor::uninitialized_dt(input.datum_type(), &out_shape)? };
42        self.dispatch_eval(stream, input, offset, out_len, &output)?;
43        stream.synchronize()?;
44        Ok(output)
45    }
46
47    pub fn dispatch_eval(
48        &self,
49        stream: &TractCudaStream,
50        input: &DeviceTensor,
51        offset: i64,
52        out_len: usize,
53        output: &DeviceTensor,
54    ) -> TractResult<()> {
55        let rank = input.rank();
56        ensure!(rank >= 2);
57        ensure!(output.rank() == rank);
58        ensure!(output.datum_type() == input.datum_type());
59        let in_shape = input.shape();
60        let out_shape = output.shape();
61        // Output shares all leading axes with input, only the last differs.
62        ensure!(in_shape[..rank - 2] == out_shape[..rank - 2]);
63        ensure!(in_shape[rank - 2] == out_shape[rank - 2]);
64        ensure!(out_shape[rank - 1] == out_len);
65        // i64 -> i32 down-cast: model widths are well under i32::MAX in practice.
66        let offset_i32: i32 = offset.try_into().context("DiagGather offset overflows i32")?;
67        let out_len_i32: i32 = out_len.try_into().context("DiagGather out_len overflows i32")?;
68
69        // Flatten the (rank-2) leading axes into one batch axis.  Assumes the
70        // leading block is plain row-major (encoder use: rank-4 BxHxTxR with
71        // natural strides), so the batch stride is `t_q * (R or out_len)`.
72        let in_strides = input.strides();
73        let out_strides = output.strides();
74        let batch: usize = in_shape[..rank - 2].iter().product();
75        let t_q = in_shape[rank - 2];
76        let r_in = in_shape[rank - 1];
77        let in_stride_b: i32 = if rank >= 3 { (t_q * r_in) as i32 } else { 0 };
78        let in_stride_i = in_strides[rank - 2] as i32;
79        let in_stride_r = in_strides[rank - 1] as i32;
80        let out_stride_b: i32 = if rank >= 3 { (t_q * out_len) as i32 } else { 0 };
81        let out_stride_i = out_strides[rank - 2] as i32;
82        let out_stride_k = out_strides[rank - 1] as i32;
83
84        let i_view = get_cuda_view(input);
85        let o_view = get_cuda_view(output);
86
87        let func = cuda_context()
88            .load_pipeline(LibraryName::Array, self.kernel_name(input.datum_type())?)?;
89
90        let mut launch_args = TractLaunchArgs::new(stream, &func);
91        launch_args.push_view(&i_view);
92        launch_args.push_view(&o_view);
93        launch_args.push::<i32>(offset_i32);
94        launch_args.push::<i32>(batch as i32);
95        launch_args.push::<i32>(t_q as i32);
96        launch_args.push::<i32>(r_in as i32);
97        launch_args.push::<i32>(out_len_i32);
98        launch_args.push::<i32>(in_stride_b);
99        launch_args.push::<i32>(in_stride_i);
100        launch_args.push::<i32>(in_stride_r);
101        launch_args.push::<i32>(out_stride_b);
102        launch_args.push::<i32>(out_stride_i);
103        launch_args.push::<i32>(out_stride_k);
104
105        // Grid: x = out_len cols, y = T_q rows, z = batch.  Threads per block
106        // along x = min(out_len, 256) rounded down to multiple of 32.
107        let block_x = out_len.clamp(32, MAX_THREADS);
108        let grid_x = out_len.div_ceil(block_x);
109        let cfg = LaunchConfig {
110            grid_dim: (grid_x as _, t_q as _, batch as _),
111            block_dim: (block_x as _, 1, 1),
112            shared_mem_bytes: 0,
113        };
114        launch_args.launch(cfg)
115    }
116}
117
118pub fn cuda_diag_gather_dispatch(
119    input: &DeviceTensor,
120    offset: i64,
121    out_len: usize,
122    output: &DeviceTensor,
123) -> TractResult<()> {
124    crate::with_cuda_stream(|stream| {
125        DiagGather.dispatch_eval(stream, input, offset, out_len, output)
126    })
127}
128
129crate::register_cuda_op!(tract_transformers::ops::diag_gather::DiagGather, |source, node, op| {
130    rule_if!(DiagGather::is_supported_dt(source.node_input_facts(node.id)?[0].datum_type));
131    Ok(Some(Box::new(tract_gpu::ops::diag_gather::GpuDiagGather::new(
132        op.offset.clone(),
133        op.out_len.clone(),
134        "Cuda",
135        cuda_diag_gather_dispatch,
136    ))))
137});
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use tract_core::internal::Tensor;
143    use tract_gpu::tensor::IntoDevice;
144    use tract_transformers::ops::diag_gather as cpu_dg;
145
146    fn run_against_cpu(shape: &[usize], offset: i64, out_len: usize) -> TractResult<()> {
147        use tract_core::plan::TurnState;
148        crate::with_cuda_stream(|stream| {
149            let len: usize = shape.iter().product();
150            let data: Vec<f32> = (0..len).map(|i| i as f32).collect();
151            let cpu_in = Tensor::from_shape(shape, &data)?;
152            let cuda_in = cpu_in.clone().into_device()?;
153
154            // CPU DiagGather only implements eval_with_turn (it resolves
155            // TDims against the turn's resolved_symbols); pass an empty
156            // TurnState since the TDims here are already concrete.
157            let cpu_op = cpu_dg::DiagGather { offset: offset.to_dim(), out_len: out_len.to_dim() };
158            let cpu_out = cpu_op.eval(&EvalContext::out_of_plan(), tvec![cpu_in.into_tvalue()])?[0]
159                .clone()
160                .into_tensor();
161            let cuda_out = DiagGather.eval(stream, &cuda_in, offset, out_len)?;
162            cpu_out
163                .close_enough(&cuda_out.to_host()?.into_tensor(), Approximation::Exact)
164                .with_context(|| format!("shape={shape:?} offset={offset} out_len={out_len}"))
165        })
166    }
167
168    #[test]
169    fn test_diag_gather_skew_basic() -> TractResult<()> {
170        // Classic skew-trick shape: [B*H, T, 2T-1] -> [B*H, T, T] with offset = T-1.
171        let t = 4;
172        run_against_cpu(&[2, t, 2 * t - 1], (t - 1) as i64, t)
173    }
174
175    #[test]
176    fn test_diag_gather_rank4_encoder_like() -> TractResult<()> {
177        // Mirrors the encoder shape: [B, H, T_q, R].
178        let t = 14;
179        run_against_cpu(&[1, 8, t, 2 * t - 1], (t - 1) as i64, t)
180    }
181
182    #[test]
183    fn test_diag_gather_out_of_bounds_zero_fill() -> TractResult<()> {
184        // out_len > R so some `r = offset + k - i` fall outside [0, R) and
185        // must be zero-filled (matches CPU contract).
186        let r = 5;
187        let t = 4;
188        run_against_cpu(&[1, t, r], 1, 8)
189    }
190
191    #[test]
192    fn test_diag_gather_partial_overlap() -> TractResult<()> {
193        // offset = 0: row i reads input[..., i, -i..-i+out_len], so the
194        // first `i` columns of each row are out of bounds and zeroed.
195        let t = 4;
196        let r = 6;
197        run_against_cpu(&[1, t, r], 0, t)
198    }
199
200    #[test]
201    fn test_diag_gather_rank2() -> TractResult<()> {
202        // Smallest valid rank: no leading batch axes.
203        run_against_cpu(&[5, 9], 4, 5)
204    }
205}
tract_cuda/kernels/array/diag_gather.rs

tract_cuda/kernels/array/
diag_gather.rs
