tract_cuda/kernels/array/
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 Gather;
12
13impl fmt::Display for Gather {
14    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
15        write!(f, "{self:?}")
16    }
17}
18
19impl Gather {
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 gather op", dt);
26        let tname = DeviceTensor::tname(dt)?;
27        Ok(format!("gather_{tname}"))
28    }
29
30    pub fn eval(
31        &self,
32        stream: &TractCudaStream,
33        data: &DeviceTensor,
34        indices: &DeviceTensor,
35        axis: usize,
36    ) -> TractResult<DeviceTensor> {
37        ensure!(data.rank() > axis);
38        let mut out_shape: TVec<usize> = data.shape()[..axis].into();
39        out_shape.extend(indices.shape().iter().copied());
40        out_shape.extend(data.shape()[axis + 1..].iter().copied());
41        let output = unsafe { DeviceTensor::uninitialized_dt(data.datum_type(), &out_shape)? };
42        self.dispatch_eval(stream, data, indices, axis, &output)?;
43        stream.synchronize()?;
44        Ok(output)
45    }
46
47    pub fn dispatch_eval(
48        &self,
49        stream: &TractCudaStream,
50        data: &DeviceTensor,
51        indices: &DeviceTensor,
52        axis: usize,
53        output: &DeviceTensor,
54    ) -> TractResult<()> {
55        ensure!(data.rank() > axis);
56        ensure!(indices.datum_type() == i64::datum_type());
57        ensure!(output.datum_type() == data.datum_type());
58
59        let data_shape = data.shape();
60        let pre: usize = data_shape[..axis].iter().product();
61        let a_size: usize = data_shape[axis];
62        let post: usize = data_shape[axis + 1..].iter().product();
63        let n_indices: usize = indices.shape().iter().product();
64
65        // Output volume must match (pre, n_indices, post) under natural strides;
66        // the GpuGather::output_facts code computes this same shape.
67        let expected: usize = pre * n_indices * post;
68        ensure!(
69            output.shape().iter().product::<usize>() == expected,
70            "Gather output shape mismatch: data={:?} axis={} indices={:?} output={:?}",
71            data_shape,
72            axis,
73            indices.shape(),
74            output.shape()
75        );
76
77        let d_view = get_cuda_view(data);
78        let i_view = get_cuda_view(indices);
79        let o_view = get_cuda_view(output);
80
81        let func = cuda_context()
82            .load_pipeline(LibraryName::Array, self.kernel_name(data.datum_type())?)?;
83
84        let mut launch_args = TractLaunchArgs::new(stream, &func);
85        launch_args.push_view(&d_view);
86        launch_args.push_view(&i_view);
87        launch_args.push_view(&o_view);
88        launch_args.push::<i32>(pre as i32);
89        launch_args.push::<i32>(a_size as i32);
90        launch_args.push::<i32>(post as i32);
91        launch_args.push::<i32>(n_indices as i32);
92
93        let block_x = post.clamp(32, MAX_THREADS);
94        let grid_x = post.div_ceil(block_x);
95        let cfg = LaunchConfig {
96            grid_dim: (grid_x as _, n_indices as _, pre as _),
97            block_dim: (block_x as _, 1, 1),
98            shared_mem_bytes: 0,
99        };
100        launch_args.launch(cfg)
101    }
102}
103
104pub fn cuda_gather_dispatch(
105    data: &DeviceTensor,
106    indices: &DeviceTensor,
107    axis: usize,
108    output: &DeviceTensor,
109) -> TractResult<()> {
110    crate::with_cuda_stream(|stream| Gather.dispatch_eval(stream, data, indices, axis, output))
111}
112
113crate::register_cuda_op!(tract_core::ops::array::Gather, |source, node, op| {
114    let facts = source.node_input_facts(node.id)?;
115    // Plain-tensor path only.  The CPU op also handles block-quant and packed
116    // matrix storage; those decompose into a dequantization step that the GPU
117    // path doesn't (yet) cover.
118    rule_if!(facts[0].is_plain());
119    rule_if!(Gather::is_supported_dt(facts[0].datum_type));
120    rule_if!(facts[1].datum_type == i64::datum_type());
121    rule_if!(op.output_type.is_none() || op.output_type == Some(facts[0].datum_type));
122    Ok(Some(Box::new(tract_gpu::ops::gather::GpuGather::new(
123        op.axis,
124        "Cuda",
125        cuda_gather_dispatch,
126    ))))
127});
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use tract_core::internal::Tensor;
133    use tract_core::ops::array::Gather as CpuGather;
134    use tract_gpu::tensor::IntoDevice;
135
136    fn run_against_cpu(
137        data_shape: &[usize],
138        indices_shape: &[usize],
139        indices_data: &[i64],
140        axis: usize,
141    ) -> TractResult<()> {
142        crate::with_cuda_stream(|stream| {
143            let n: usize = data_shape.iter().product();
144            let data = Tensor::from_shape(
145                data_shape,
146                &(0..n).map(|i| i as f32 / 10.0).collect::<Vec<_>>(),
147            )?;
148            let indices = Tensor::from_shape(indices_shape, indices_data)?;
149            let cuda_data = data.clone().into_device()?;
150            let cuda_indices = indices.clone().into_device()?;
151
152            let cpu_op = CpuGather::new(axis);
153            let cpu_out = cpu_op.eval(
154                &EvalContext::out_of_plan(),
155                tvec![data.into_tvalue(), indices.into_tvalue()],
156            )?[0]
157                .clone()
158                .into_tensor();
159            let cuda_out = Gather.eval(stream, &cuda_data, &cuda_indices, axis)?;
160            cpu_out
161                .close_enough(&cuda_out.to_host()?.into_tensor(), Approximation::Exact)
162                .with_context(|| {
163                    format!(
164                        "data={data_shape:?} indices={indices_shape:?} axis={axis} \
165                         indices_data={indices_data:?}"
166                    )
167                })
168        })
169    }
170
171    /// Embedding lookup: rank-2 table, rank-2 index — the nemotron decoder shape.
172    #[test]
173    fn test_gather_embedding() -> TractResult<()> {
174        run_against_cpu(&[1025, 640], &[1, 1], &[42], 0)
175    }
176
177    /// Multi-batch embedding lookup.
178    #[test]
179    fn test_gather_embedding_multi() -> TractResult<()> {
180        run_against_cpu(&[100, 16], &[2, 3], &[0, 1, 99, 50, 25, 7], 0)
181    }
182
183    /// Non-zero axis: pre-batch axes flatten correctly.
184    #[test]
185    fn test_gather_axis_1() -> TractResult<()> {
186        run_against_cpu(&[3, 10, 4], &[2], &[0, 9], 1)
187    }
188
189    /// Negative indices wrap (axis size = 100, so -1 → 99, -100 → 0).
190    #[test]
191    fn test_gather_negative_indices() -> TractResult<()> {
192        run_against_cpu(&[100, 4], &[3], &[-1, -100, -50], 0)
193    }
194
195    /// Scalar index input (rank-0).
196    #[test]
197    fn test_gather_scalar_index() -> TractResult<()> {
198        run_against_cpu(&[5, 8], &[], &[3], 0)
199    }
200}
tract_cuda/kernels/array/gather.rs

tract_cuda/kernels/array/
gather.rs
