use derive_new::new;
use onnx_ir_derive::NodeBuilder;
use crate::ir::{ArgType, Argument, Node, RawNode, RuntimeInputRef, TensorDataExt};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
#[derive(Debug, Clone, new)]
pub struct SliceConfig {
pub starts: SliceInput,
pub ends: SliceInput,
pub axes: Option<SliceInput>,
pub steps: Option<SliceInput>,
}
#[derive(Debug, Clone, NodeBuilder)]
pub struct SliceNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: SliceConfig,
}
#[derive(Debug, Clone)]
pub enum SliceInput {
Static(Vec<i64>),
Runtime(RuntimeInputRef),
}
impl Default for SliceInput {
fn default() -> Self {
Self::Static(Vec::new())
}
}
fn normalize_axes(axes: &mut [i64], rank: usize, _node_name: &str) {
for axis in axes.iter_mut() {
if *axis < 0 {
*axis += rank as i64;
}
}
}
fn calculate_shape_slice_output_len(
start: i64,
end: i64,
step: i64,
shape_rank: usize,
_node_name: &str,
) -> usize {
let shape_len = shape_rank as i64;
let norm_start = if start < 0 {
(shape_len + start).max(0)
} else {
start.min(shape_len)
};
let norm_end = if end == i64::MAX || end >= shape_len {
shape_len
} else if end < 0 {
(shape_len + end).max(0)
} else {
end.min(shape_len)
};
let range_len = (norm_end - norm_start).max(0);
if step.abs() == 1 {
range_len as usize
} else {
((range_len + step.abs() - 1) / step.abs()) as usize
}
}
pub(crate) struct SliceProcessor;
impl NodeProcessor for SliceProcessor {
type Config = SliceConfig;
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 10,
max_opset: None,
inputs: InputSpec::AtLeast(3),
outputs: OutputSpec::Exact(1),
}
}
fn input_preferences(
&self,
node: &RawNode,
_opset: usize,
) -> Result<Option<crate::processor::InputPreferences>, ProcessError> {
use crate::processor::{ArgPreference, InputPreferences};
let mut prefs = InputPreferences::new();
for input in node.inputs.iter().skip(1) {
prefs = prefs.add(&input.name, ArgPreference::Shape);
}
Ok(Some(prefs))
}
fn lift_constants(&self, node: &mut RawNode, _opset: usize) -> Result<(), ProcessError> {
if node.inputs.len() > 1 && node.inputs[1].is_constant() {
node.inputs[1].to_static()?;
}
if node.inputs.len() > 2 && node.inputs[2].is_constant() {
node.inputs[2].to_static()?;
}
if node.inputs.len() > 3 && node.inputs[3].is_constant() {
node.inputs[3].to_static()?;
}
if node.inputs.len() > 4 && node.inputs[4].is_constant() {
node.inputs[4].to_static()?;
}
Ok(())
}
fn infer_types(
&self,
node: &mut RawNode,
opset: usize,
_output_preferences: &OutputPreferences,
) -> Result<(), ProcessError> {
let config = self.extract_config(node, opset)?;
let input_ty = node.inputs[0].ty.clone();
match input_ty {
ArgType::Tensor(_) => {
node.outputs[0].ty = input_ty;
}
ArgType::Shape(shape_rank) => {
let (starts, ends, steps) = match (&config.starts, &config.ends, &config.steps) {
(SliceInput::Static(s), SliceInput::Static(e), steps_opt) => {
let step_values = match steps_opt {
Some(SliceInput::Static(st)) => st.clone(),
_ => vec![1], };
(s, e, step_values)
}
_ => {
return Err(ProcessError::Custom(format!(
"Runtime slice on Shape input is not supported for node {}",
node.name
)));
}
};
if starts.len() != 1 || ends.len() != 1 {
return Err(ProcessError::Custom(format!(
"Slice on Shape input requires exactly one dimension slice config for node {}",
node.name
)));
}
let step = if steps.is_empty() { 1 } else { steps[0] };
let output_len = calculate_shape_slice_output_len(
starts[0], ends[0], step, shape_rank, &node.name,
);
node.outputs[0].ty = ArgType::Shape(output_len);
}
unsupported_type => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor or Shape".to_string(),
actual: format!("{:?}", unsupported_type),
});
}
}
Ok(())
}
fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
fn get_slice_input(
node: &RawNode,
index: usize,
) -> Result<Option<SliceInput>, ProcessError> {
let input = match node.inputs.get(index) {
Some(i) => i,
None => return Ok(None),
};
if matches!(input.ty, ArgType::Shape(_)) {
match input.value() {
Some(tensor_data) => match tensor_data.to_i64_vec() {
Ok(vec) => return Ok(Some(SliceInput::Static(vec))),
Err(_) => {
return Err(ProcessError::Custom(format!(
"Slice Shape input at index {} must be int32 or int64",
index
)));
}
},
None => {
return Ok(Some(SliceInput::Runtime(RuntimeInputRef::new(
input.name.clone(),
index,
))));
}
}
}
match input.value() {
None => {
Ok(Some(SliceInput::Runtime(RuntimeInputRef::new(
input.name.clone(),
index,
))))
}
Some(tensor_data) => match tensor_data.to_i64_vec() {
Ok(vec) => Ok(Some(SliceInput::Static(vec))),
Err(_) => Err(ProcessError::Custom(format!(
"Slice input at index {} must be int32 or int64",
index
))),
},
}
}
let starts = get_slice_input(node, 1)?
.ok_or_else(|| ProcessError::MissingInput("starts".to_string()))?;
let ends = get_slice_input(node, 2)?
.ok_or_else(|| ProcessError::MissingInput("ends".to_string()))?;
let axes = get_slice_input(node, 3)?;
let steps = get_slice_input(node, 4)?;
let (mut axes, steps) = if let SliceInput::Static(ref starts_vec) = starts {
let num_slices = starts_vec.len();
let steps = if steps.is_none() && num_slices > 0 {
Some(SliceInput::Static(vec![1; num_slices]))
} else {
steps
};
let axes = if axes.is_none() && num_slices > 0 {
Some(SliceInput::Static((0..num_slices as i64).collect()))
} else {
axes
};
(axes, steps)
} else {
(axes, steps)
};
if let Some(SliceInput::Static(ref step_values)) = steps
&& step_values.contains(&0)
{
return Err(ProcessError::Custom(
"Slice: step values cannot be zero".to_string(),
));
}
if let Some(SliceInput::Static(ref mut axes_values)) = axes
&& let ArgType::Tensor(ref tensor_type) = node.inputs[0].ty
{
normalize_axes(axes_values, tensor_type.rank, &node.name);
}
let config = SliceConfig {
starts,
ends,
axes,
steps,
};
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
Node::Slice(SliceNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
config,
})
}
}
#[cfg(test)]
mod tests {
use crate::ir::{DType, NodeType};
use crate::node::test_utils::TestNodeBuilder;
use super::*;
fn create_test_node(
starts: Vec<i64>,
ends: Vec<i64>,
axes: Option<Vec<i64>>,
) -> TestNodeBuilder {
let mut builder = TestNodeBuilder::new(NodeType::Slice, "test_slice")
.input_tensor_f32("data", 3, None)
.output_default("output");
builder = builder.input_tensor_i64_data("starts", starts.clone(), vec![starts.len()]);
builder = builder.input_tensor_i64_data("ends", ends.clone(), vec![ends.len()]);
if let Some(axes_vec) = axes.clone() {
builder = builder.input_tensor_i64_data("axes", axes_vec.clone(), vec![axes_vec.len()]);
}
builder
}
fn create_shape_input_node(start: i64, end: i64) -> TestNodeBuilder {
TestNodeBuilder::new(NodeType::Slice, "test_slice_shape")
.input_shape("data", 5)
.input_tensor_i64_data("starts", vec![start], vec![1])
.input_tensor_i64_data("ends", vec![end], vec![1])
.input_tensor_i64_data("axes", vec![0], vec![1])
.output_default("output")
}
fn create_runtime_slice_node() -> TestNodeBuilder {
TestNodeBuilder::new(NodeType::Slice, "test_runtime_slice")
.input_tensor_f32("data", 2, None)
.input_tensor_i64("starts", 0, None) .input_tensor_i64("ends", 0, None) .input_tensor_i64_data("axes", vec![0], vec![1])
.input_tensor_i64_data("steps", vec![1], vec![1])
.output_default("output")
}
fn create_mixed_slice_node_runtime_start() -> TestNodeBuilder {
TestNodeBuilder::new(NodeType::Slice, "test_mixed_slice")
.input_tensor_f32("data", 2, None)
.input_tensor_i64("starts", 0, None) .input_tensor_i64_data("ends", vec![3], vec![1]) .input_tensor_i64_data("axes", vec![0], vec![1])
.input_tensor_i64_data("steps", vec![1], vec![1])
.output_default("output")
}
fn create_mixed_slice_node_runtime_end() -> TestNodeBuilder {
TestNodeBuilder::new(NodeType::Slice, "test_mixed_slice")
.input_tensor_f32("data", 2, None)
.input_tensor_i64_data("starts", vec![1], vec![1]) .input_tensor_i64("ends", 0, None) .input_tensor_i64_data("axes", vec![0], vec![1])
.input_tensor_i64_data("steps", vec![1], vec![1])
.output_default("output")
}
#[test]
fn test_slice_config_basic() {
let node =
create_test_node(vec![1, 0], vec![3, 2], Some(vec![0, 2])).build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends) {
(SliceInput::Static(starts), SliceInput::Static(ends)) => {
assert_eq!(starts, &vec![1, 0]);
assert_eq!(ends, &vec![3, 2]);
if let Some(SliceInput::Static(axes)) = &result.axes {
assert_eq!(axes, &vec![0, 2]);
}
if let Some(SliceInput::Static(steps)) = &result.steps {
assert_eq!(steps, &vec![1, 1]);
} else {
panic!("Expected steps to have ONNX spec default");
}
}
_ => panic!("Expected static config"),
}
}
#[test]
fn test_slice_config_negative_axes() {
let node = create_test_node(vec![1], vec![3], Some(vec![-3])).build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends) {
(SliceInput::Static(starts), SliceInput::Static(ends)) => {
assert_eq!(starts, &vec![1]);
assert_eq!(ends, &vec![3]);
if let Some(SliceInput::Static(axes)) = &result.axes {
assert_eq!(axes, &vec![0]); }
}
_ => panic!("Expected static config"),
}
}
#[test]
fn test_slice_config_default_axes() {
let node = create_test_node(vec![1, 2], vec![3, 4], None).build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends) {
(SliceInput::Static(starts), SliceInput::Static(ends)) => {
assert_eq!(starts, &vec![1, 2]);
assert_eq!(ends, &vec![3, 4]);
if let Some(SliceInput::Static(axes)) = &result.axes {
assert_eq!(axes, &vec![0, 1]);
} else {
panic!("Expected axes to have ONNX spec default");
}
}
_ => panic!("Expected static config"),
}
}
#[test]
fn test_slice_config_runtime() {
let node = create_runtime_slice_node().build();
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends) {
(SliceInput::Runtime(starts), SliceInput::Runtime(ends)) => {
assert_eq!(starts.name, "starts");
assert_eq!(ends.name, "ends");
if let Some(SliceInput::Static(axes)) = &result.axes {
assert_eq!(axes, &vec![0]);
}
if let Some(SliceInput::Static(steps)) = &result.steps {
assert_eq!(steps, &vec![1]);
}
}
_ => panic!("Expected runtime config"),
}
}
#[test]
fn test_slice_update_output_rank_tensor_input() {
let mut node = create_test_node(vec![1, 2], vec![3, 4], None).build_with_graph_data(16);
assert!(matches!(node.inputs[0].ty, ArgType::Tensor(_)));
assert!(matches!(node.outputs[0].ty, ArgType::Tensor(_)));
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert!(
matches!(&node.outputs[0].ty, ArgType::Tensor(tensor_type) if tensor_type.dtype == DType::F32 && tensor_type.rank == 3)
);
}
#[test]
fn test_slice_update_output_rank_shape_input() {
let mut node = create_shape_input_node(1, 3).build_with_graph_data(16);
assert!(matches!(node.inputs[0].ty, ArgType::Shape(5)));
assert!(matches!(node.outputs[0].ty, ArgType::Tensor(ref t) if t.rank == 0));
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert!(matches!(&node.outputs[0].ty, ArgType::Shape(2)));
}
#[test]
fn test_slice_config_mixed_runtime_start() {
let node = create_mixed_slice_node_runtime_start().build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends) {
(SliceInput::Runtime(starts), SliceInput::Static(ends)) => {
assert_eq!(starts.name, "starts");
assert_eq!(ends, &vec![3]);
}
_ => panic!("Expected mixed config with runtime start and static end"),
}
}
#[test]
fn test_slice_config_mixed_runtime_end() {
let node = create_mixed_slice_node_runtime_end().build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends) {
(SliceInput::Static(starts), SliceInput::Runtime(ends)) => {
assert_eq!(starts, &vec![1]);
assert_eq!(ends.name, "ends");
}
_ => panic!("Expected mixed config with static start and runtime end"),
}
}
#[test]
fn test_slice_config_with_steps() {
let builder = TestNodeBuilder::new(NodeType::Slice, "test_slice_with_steps")
.input_tensor_f32("data", 3, None)
.input_tensor_i64_data("starts", vec![0, 0], vec![2])
.input_tensor_i64_data("ends", vec![10, 10], vec![2])
.input_tensor_i64_data("axes", vec![0, 1], vec![2])
.input_tensor_i64_data("steps", vec![2, 3], vec![2])
.output_default("output");
let node = builder.build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match (&result.starts, &result.ends, &result.steps) {
(
SliceInput::Static(starts),
SliceInput::Static(ends),
Some(SliceInput::Static(steps)),
) => {
assert_eq!(starts, &vec![0, 0]);
assert_eq!(ends, &vec![10, 10]);
assert_eq!(steps, &vec![2, 3]);
}
_ => panic!("Expected static config with steps"),
}
}
#[test]
fn test_slice_config_zero_step() {
let builder = TestNodeBuilder::new(NodeType::Slice, "test_zero_step")
.input_tensor_f32("data", 2, None)
.input_tensor_i64_data("starts", vec![0], vec![1])
.input_tensor_i64_data("ends", vec![10], vec![1])
.input_tensor_i64_data("axes", vec![0], vec![1])
.input_tensor_i64_data("steps", vec![0], vec![1])
.output_default("output");
let node = builder.build_with_graph_data(16);
let node = node;
let processor = SliceProcessor;
let result = processor.extract_config(&node, 16);
assert!(matches!(result, Err(ProcessError::Custom(_))));
}
#[test]
fn test_slice_config_negative_steps() {
let builder = TestNodeBuilder::new(NodeType::Slice, "test_negative_steps")
.input_tensor_f32("data", 2, None)
.input_tensor_i64_data("starts", vec![0, 2], vec![2])
.input_tensor_i64_data("ends", vec![10, 8], vec![2])
.input_tensor_i64_data("axes", vec![0, 1], vec![2])
.input_tensor_i64_data("steps", vec![-1, -2], vec![2])
.output_default("output");
let node = builder.build_with_graph_data(16);
let mut node = node;
let processor = SliceProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
let result = processor.extract_config(&node, 16).unwrap();
match &result.steps {
Some(SliceInput::Static(steps)) => {
assert_eq!(steps, &vec![-1, -2]);
}
_ => panic!("Expected static steps with negative values"),
}
}
}