onnx-runtime-ep-cpu 0.1.0-dev.2

//! Sequence construction and scans: `Tile`, `Range`, and `CumSum`.

use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, TensorView};
use onnx_runtime_ir::{DataType, Node};

use super::{
    check_arity, elem_size, to_dense_bytes, to_dense_f32, to_dense_i64, write_dense_bytes,
    write_dense_f32,
};
use crate::strided::numel;

pub struct TileKernel;
pub struct TileFactory;
impl KernelFactory for TileFactory {
    fn create(&self, _: &Node, _: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(TileKernel))
    }
}
impl Kernel for TileKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        check_arity("Tile", inputs, outputs, 2, 2, 1)?;
        let x = &inputs[0];
        let repeats = to_dense_i64(&inputs[1])?;
        if repeats.len() != x.shape.len() || repeats.iter().any(|&r| r < 0) {
            return Err(EpError::KernelFailed(
                "Tile: repeats must be non-negative and match input rank".into(),
            ));
        }
        if outputs[0].dtype != x.dtype {
            return Err(EpError::KernelFailed(
                "Tile: output dtype must match input".into(),
            ));
        }
        let esize = elem_size(x.dtype)?;
        let src = to_dense_bytes(x)?;
        let out_shape = outputs[0].shape.to_vec();
        if out_shape.len() != x.shape.len()
            || out_shape
                .iter()
                .zip(x.shape)
                .zip(&repeats)
                .any(|((&o, &d), &r)| o != d * r as usize)
        {
            return Err(EpError::KernelFailed(
                "Tile: output shape does not match repeats".into(),
            ));
        }
        let in_strides = strides(x.shape);
        let out_strides = strides(&out_shape);
        let mut out = vec![0; numel(&out_shape) * esize];
        for linear in 0..numel(&out_shape) {
            let mut rem = linear;
            let mut source = 0;
            for d in 0..out_shape.len() {
                let coordinate = rem / out_strides[d];
                rem %= out_strides[d];
                source += (coordinate % x.shape[d]) * in_strides[d];
            }
            out[linear * esize..(linear + 1) * esize]
                .copy_from_slice(&src[source * esize..(source + 1) * esize]);
        }
        write_dense_bytes(&mut outputs[0], &out)
    }
    fn supports_strided_input(&self, _: usize) -> bool {
        true
    }
}

pub struct RangeKernel;
pub struct RangeFactory;
impl KernelFactory for RangeFactory {
    fn create(&self, _: &Node, _: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(RangeKernel))
    }
}
impl Kernel for RangeKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        check_arity("Range", inputs, outputs, 3, 3, 1)?;
        if inputs.iter().any(|v| !v.shape.is_empty()) {
            return Err(EpError::KernelFailed(
                "Range: inputs must be scalars".into(),
            ));
        }
        if inputs.iter().any(|v| v.dtype != inputs[0].dtype) || outputs[0].dtype != inputs[0].dtype
        {
            return Err(EpError::KernelFailed(
                "Range: input and output dtypes must match".into(),
            ));
        }
        match inputs[0].dtype {
            DataType::Float32 => {
                let (start, limit, delta) = (
                    to_dense_f32(&inputs[0])?[0],
                    to_dense_f32(&inputs[1])?[0],
                    to_dense_f32(&inputs[2])?[0],
                );
                if delta == 0. {
                    return Err(EpError::KernelFailed(
                        "Range: delta must not be zero".into(),
                    ));
                }
                let count = float_range_count(start, limit, delta)?;
                let mut out = alloc_range_output::<f32>(count)?;
                out.extend((0..count).map(|i| start + i as f32 * delta));
                write_dense_f32(&mut outputs[0], &out)
            }
            DataType::Int64 => {
                let (start, limit, delta) = (
                    to_dense_i64(&inputs[0])?[0],
                    to_dense_i64(&inputs[1])?[0],
                    to_dense_i64(&inputs[2])?[0],
                );
                if delta == 0 {
                    return Err(EpError::KernelFailed(
                        "Range: delta must not be zero".into(),
                    ));
                }
                let count = int_range_count(start, limit, delta)?;
                let mut out = alloc_range_output::<i64>(count)?;
                for i in 0..count {
                    let v = i64::try_from(start as i128 + i as i128 * delta as i128)
                        .map_err(|_| EpError::KernelFailed("Range: value overflow".into()))?;
                    out.push(v);
                }
                let bytes: Vec<u8> = out.iter().flat_map(|v| v.to_le_bytes()).collect();
                write_dense_bytes(&mut outputs[0], &bytes)
            }
            dtype => Err(EpError::KernelFailed(format!(
                "Range: unsupported dtype {dtype:?}"
            ))),
        }
    }
    fn supports_strided_input(&self, _: usize) -> bool {
        true
    }
}

pub struct CumSumKernel {
    exclusive: bool,
    reverse: bool,
}
pub struct CumSumFactory;
impl KernelFactory for CumSumFactory {
    fn create(&self, node: &Node, _: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(CumSumKernel {
            exclusive: node.attr("exclusive").and_then(|a| a.as_int()).unwrap_or(0) != 0,
            reverse: node.attr("reverse").and_then(|a| a.as_int()).unwrap_or(0) != 0,
        }))
    }
}
impl Kernel for CumSumKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        check_arity("CumSum", inputs, outputs, 2, 2, 1)?;
        let axis = to_dense_i64(&inputs[1])?;
        if axis.len() != 1 {
            return Err(EpError::KernelFailed(
                "CumSum: axis must be a scalar".into(),
            ));
        }
        let rank = inputs[0].shape.len();
        let raw = axis[0];
        let axis = if raw < 0 { raw + rank as i64 } else { raw };
        if axis < 0 || axis as usize >= rank {
            return Err(EpError::KernelFailed("CumSum: axis out of range".into()));
        }
        let axis = axis as usize;
        if outputs[0].dtype != inputs[0].dtype {
            return Err(EpError::KernelFailed(
                "CumSum: output dtype must match input".into(),
            ));
        }
        match inputs[0].dtype {
            DataType::Float32 => {
                let mut out = to_dense_f32(&inputs[0])?;
                scan(
                    &mut out,
                    inputs[0].shape,
                    axis,
                    self.exclusive,
                    self.reverse,
                    |a, b| a + b,
                );
                write_dense_f32(&mut outputs[0], &out)
            }
            DataType::Int64 => {
                let mut out = to_dense_i64(&inputs[0])?;
                scan(
                    &mut out,
                    inputs[0].shape,
                    axis,
                    self.exclusive,
                    self.reverse,
                    i64::wrapping_add,
                );
                let bytes: Vec<u8> = out.iter().flat_map(|v| v.to_le_bytes()).collect();
                write_dense_bytes(&mut outputs[0], &bytes)
            }
            dtype => Err(EpError::KernelFailed(format!(
                "CumSum: unsupported dtype {dtype:?}"
            ))),
        }
    }
    fn supports_strided_input(&self, _: usize) -> bool {
        true
    }
}

fn float_range_count(start: f32, limit: f32, delta: f32) -> Result<usize> {
    let count = ((limit - start) / delta).ceil().max(0.0);
    if !count.is_finite() || count >= usize::MAX as f32 {
        return Err(EpError::KernelFailed(
            "Range: element count exceeds addressable memory".into(),
        ));
    }
    Ok(count as usize)
}

fn int_range_count(start: i64, limit: i64, delta: i64) -> Result<usize> {
    let count = if delta > 0 && start < limit {
        let distance = limit as i128 - start as i128;
        (distance + delta as i128 - 1) / delta as i128
    } else if delta < 0 && start > limit {
        let distance = start as i128 - limit as i128;
        let step = -(delta as i128);
        (distance + step - 1) / step
    } else {
        0
    };
    usize::try_from(count).map_err(|_| {
        EpError::KernelFailed("Range: element count exceeds addressable memory".into())
    })
}

/// Validate a Range output's element count against byte addressability and
/// pre-allocate its backing buffer without ever panicking.
///
/// `count * size_of::<T>()` must not exceed `isize::MAX` (the largest size any
/// Rust allocation can address); otherwise `Vec` allocation/index math would
/// panic on user-controlled inputs. We check the bound up front and fall back
/// to `try_reserve` so an over-large request fails as a clean kernel error.
fn alloc_range_output<T>(count: usize) -> Result<Vec<T>> {
    let elem = std::mem::size_of::<T>();
    let bytes = count.checked_mul(elem);
    if bytes.is_none_or(|b| b > isize::MAX as usize) {
        return Err(EpError::KernelFailed(format!(
            "Range output too large: {count} elements ({} bytes) exceeds addressable limit",
            bytes.map_or_else(|| "overflow".to_string(), |b| b.to_string()),
        )));
    }
    let mut out = Vec::new();
    out.try_reserve(count).map_err(|_| {
        EpError::KernelFailed(format!(
            "Range output too large: failed to allocate {count} elements ({} bytes)",
            count * elem,
        ))
    })?;
    Ok(out)
}

fn scan<T: Copy + Default>(
    values: &mut [T],
    shape: &[usize],
    axis: usize,
    exclusive: bool,
    reverse: bool,
    add: impl Fn(T, T) -> T,
) {
    let inner = numel(&shape[axis + 1..]);
    let width = shape[axis];
    for outer in 0..numel(&shape[..axis]) {
        for i in 0..inner {
            let mut total = T::default();
            for n in 0..width {
                let d = if reverse { width - 1 - n } else { n };
                let offset = (outer * width + d) * inner + i;
                let v = values[offset];
                if exclusive {
                    values[offset] = total;
                    total = add(total, v);
                } else {
                    total = add(total, v);
                    values[offset] = total;
                }
            }
        }
    }
}
fn strides(shape: &[usize]) -> Vec<usize> {
    let mut result = vec![1; shape.len()];
    for d in (0..shape.len()).rev().skip(1) {
        result[d] = result[d + 1] * shape[d + 1];
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::kernels::testutil::Owned;
    #[test]
    fn tile_repeats_each_axis() {
        let x = Owned::f32(&[1, 2], &[1., 2.]);
        let r = Owned::i64(&[2], &[2, 3]);
        let mut y = Owned::zeros_f32(&[2, 6]);
        TileKernel
            .execute(&[x.view(), r.view()], &mut [y.view_mut()])
            .unwrap();
        assert_eq!(
            y.to_f32(),
            vec![1., 2., 1., 2., 1., 2., 1., 2., 1., 2., 1., 2.]
        );
    }
    #[test]
    fn range_descending_int64() {
        let a = Owned::i64(&[], &[5]);
        let b = Owned::i64(&[], &[-1]);
        let d = Owned::i64(&[], &[-2]);
        let mut y = Owned::zeros(DataType::Int64, &[3]);
        RangeKernel
            .execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()])
            .unwrap();
        assert_eq!(y.to_i64(), vec![5, 3, 1]);
    }
    #[test]
    fn range_float32_generates_by_index() {
        let a = Owned::f32(&[], &[0.]);
        let b = Owned::f32(&[], &[1.]);
        let d = Owned::f32(&[], &[0.25]);
        let mut y = Owned::zeros_f32(&[4]);
        RangeKernel
            .execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()])
            .unwrap();
        assert_eq!(y.to_f32(), vec![0., 0.25, 0.5, 0.75]);
    }
    #[test]
    fn range_float32_descending() {
        let a = Owned::f32(&[], &[1.]);
        let b = Owned::f32(&[], &[-0.1]);
        let d = Owned::f32(&[], &[-0.25]);
        let mut y = Owned::zeros_f32(&[5]);
        RangeKernel
            .execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()])
            .unwrap();
        assert_eq!(y.to_f32(), vec![1., 0.75, 0.5, 0.25, 0.]);
    }
    #[test]
    fn range_float32_no_progress_terminates() {
        let a = Owned::f32(&[], &[16_777_216.]);
        let b = Owned::f32(&[], &[16_777_218.]);
        let d = Owned::f32(&[], &[1.]);
        let mut y = Owned::zeros_f32(&[2]);
        RangeKernel
            .execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()])
            .unwrap();
        assert_eq!(y.to_f32().len(), 2);
        assert_eq!(y.to_f32().first(), Some(&16_777_216.));
        assert_eq!(y.to_f32().last(), Some(&16_777_216.));
    }
    #[test]
    fn range_int64_overflow_returns_error() {
        // count ~= i64::MAX elements; count * size_of::<i64>() overflows the
        // addressable-byte limit, so the kernel must return an Err before any
        // allocation is attempted instead of panicking.
        let a = Owned::i64(&[], &[0]);
        let b = Owned::i64(&[], &[i64::MAX]);
        let d = Owned::i64(&[], &[1]);
        let mut y = Owned::zeros(DataType::Int64, &[1]);
        let result = RangeKernel.execute(&[a.view(), b.view(), d.view()], &mut [y.view_mut()]);
        assert!(
            matches!(result, Err(EpError::KernelFailed(_))),
            "expected KernelFailed, got {result:?}"
        );
    }
    #[test]
    fn cumsum_reverse_exclusive() {
        let x = Owned::f32(&[3], &[1., 2., 3.]);
        let a = Owned::i64(&[], &[0]);
        let mut y = Owned::zeros_f32(&[3]);
        CumSumKernel {
            exclusive: true,
            reverse: true,
        }
        .execute(&[x.view(), a.view()], &mut [y.view_mut()])
        .unwrap();
        assert_eq!(y.to_f32(), vec![5., 3., 0.]);
    }
}