luma_tensor/device/cpu/kernels/
shape.rs

1//! Shape-movement kernels that materialize data: concatenation.
2
3use super::iter::gather;
4use crate::{Error, Layout, Result, Shape};
5
6/// Concatenate contiguous logical views along `dim`. Each source is first
7/// materialized in logical order, then copied block-by-block into the output.
8pub fn cat<T: Copy + Default>(srcs: &[(&[T], &Layout)], dim: usize) -> Result<(Vec<T>, Shape)> {
9    if srcs.is_empty() {
10        return Err(Error::OpRequiresAtLeastOneTensor { op: "cat" });
11    }
12    let first = srcs[0].1;
13    let rank = first.shape().rank();
14    // output shape: sum sizes along `dim`, all others must match.
15    let mut out_dims = first.dims().to_vec();
16    let mut cat_size = 0usize;
17    for (n, (_, l)) in srcs.iter().enumerate() {
18        let d = l.dims();
19        if d.len() != rank {
20            return Err(Error::ShapeMismatchCat { dim, first_shape: first.shape().clone(), n, nth_shape: l.shape().clone() });
21        }
22        for (i, (&a, &b)) in first.dims().iter().zip(d.iter()).enumerate() {
23            if i != dim && a != b {
24                return Err(Error::ShapeMismatchCat { dim, first_shape: first.shape().clone(), n, nth_shape: l.shape().clone() });
25            }
26        }
27        cat_size += d[dim];
28    }
29    out_dims[dim] = cat_size;
30    let out_shape = Shape::from(out_dims.clone());
31
32    let outer: usize = out_dims[..dim].iter().product();
33    let right: usize = out_dims[dim + 1..].iter().product();
34
35    let mut out = vec![T::default(); out_shape.element_count()];
36    // materialize each source contiguously, then interleave along `dim`.
37    let materialized: Vec<Vec<T>> = srcs.iter().map(|(d, l)| gather(d, l)).collect();
38
39    for o in 0..outer {
40        let mut dst_dim_base = 0usize;
41        for (src_idx, (_, l)) in srcs.iter().enumerate() {
42            let src_dim = l.dims()[dim];
43            let block = src_dim * right;
44            let src = &materialized[src_idx][o * block..o * block + block];
45            let dst_start = (o * cat_size + dst_dim_base) * right;
46            out[dst_start..dst_start + block].copy_from_slice(src);
47            dst_dim_base += src_dim;
48        }
49    }
50    Ok((out, out_shape))
51}
luma_tensor/device/cpu/kernels/shape.rs

luma_tensor/device/cpu/kernels/
shape.rs
