use derive_new::new;
use onnx_ir_derive::NodeBuilder;
use crate::ir::Argument;
use crate::ir::{ArgType, Node, RawNode, TensorType};
use crate::node::padding::{AutoPad, PaddingConfig3d, padding_config_3d};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
#[derive(Debug, Clone, new)]
pub struct AvgPool3dConfig {
pub kernel_size: [usize; 3],
pub strides: [usize; 3],
pub padding: PaddingConfig3d,
pub count_include_pad: bool,
pub dilation: [usize; 3],
pub ceil_mode: bool,
pub auto_pad: AutoPad,
}
#[derive(Debug, Clone, NodeBuilder)]
pub struct AveragePool3dNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: AvgPool3dConfig,
}
pub(crate) struct AvgPool3dProcessor;
impl NodeProcessor for AvgPool3dProcessor {
type Config = AvgPool3dConfig;
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 1,
max_opset: None,
inputs: InputSpec::Exact(1),
outputs: OutputSpec::Exact(1),
}
}
fn infer_types(
&self,
node: &mut RawNode,
opset: usize,
_output_preferences: &OutputPreferences,
) -> Result<(), ProcessError> {
for (key, value) in node.attrs.iter() {
match key.as_str() {
"kernel_shape" | "strides" | "pads" | "count_include_pad" => {}
"ceil_mode" => {
let ceil_mode = value.clone().into_i64();
if ceil_mode != 0 && opset < 19 {
return Err(ProcessError::Custom(format!(
"AveragePool: ceil_mode requires opset 19+, got opset {}",
opset
)));
}
}
"dilations" => {
let dilations = value.clone().into_i64s();
if dilations.iter().any(|&d| d != 1) && opset < 10 {
return Err(ProcessError::Custom(format!(
"AveragePool: dilations requires opset 10+, got opset {}",
opset
)));
}
}
"auto_pad" => {
AutoPad::parse(&value.clone().into_string())?;
}
_ => {
return Err(ProcessError::InvalidAttribute {
name: key.clone(),
reason: format!("Unexpected attribute for AvgPool3d: {}", key),
});
}
}
}
let tensor = match &node.inputs[0].ty {
ArgType::Tensor(tensor) => tensor,
_ => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor".to_string(),
actual: format!("{:?}", node.inputs[0].ty),
});
}
};
if tensor.rank != 5 {
return Err(ProcessError::Custom(format!(
"AvgPool3d expects input tensor of rank 5 (N x C x D x H x W), got rank {}",
tensor.rank
)));
}
node.outputs[0].ty = ArgType::Tensor(TensorType {
dtype: tensor.dtype,
rank: tensor.rank,
static_shape: None,
});
Ok(())
}
fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
let mut kernel_shape = Vec::new();
let mut strides = vec![1, 1, 1];
let mut pads = vec![0, 0, 0, 0, 0, 0];
let mut count_include_pad: i64 = 0;
let mut dilations = vec![1, 1, 1];
let mut ceil_mode: i64 = 0;
let mut auto_pad = AutoPad::NotSet;
for (key, value) in node.attrs.iter() {
match key.as_str() {
"kernel_shape" => kernel_shape = value.clone().into_i64s(),
"strides" => strides = value.clone().into_i64s(),
"pads" => pads = value.clone().into_i64s(),
"count_include_pad" => count_include_pad = value.clone().into_i64(),
"dilations" => dilations = value.clone().into_i64s(),
"ceil_mode" => ceil_mode = value.clone().into_i64(),
"auto_pad" => auto_pad = AutoPad::parse(&value.clone().into_string())?,
_ => {}
}
}
let padding = padding_config_3d(&pads);
let config = AvgPool3dConfig::new(
[
kernel_shape[0] as usize,
kernel_shape[1] as usize,
kernel_shape[2] as usize,
],
[
strides[0] as usize,
strides[1] as usize,
strides[2] as usize,
],
padding,
count_include_pad == 1,
[
dilations[0] as usize,
dilations[1] as usize,
dilations[2] as usize,
],
ceil_mode == 1,
auto_pad,
);
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
Node::AveragePool3d(AveragePool3dNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
config,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::NodeType;
use crate::node::test_utils::TestNodeBuilder;
fn create_test_node(
kernel_shape: Vec<i64>,
strides: Vec<i64>,
pads: Vec<i64>,
count_include_pad: i64,
ceil_mode: i64,
dilations: Option<Vec<i64>>,
) -> RawNode {
let mut builder = TestNodeBuilder::new(NodeType::AveragePool3d, "test_avgpool3d")
.input_tensor_f32("data", 5, None)
.output_tensor_f32("output", 5, None)
.attr_ints("kernel_shape", kernel_shape)
.attr_ints("strides", strides)
.attr_ints("pads", pads)
.attr_int("count_include_pad", count_include_pad)
.attr_int("ceil_mode", ceil_mode);
if let Some(dilations) = dilations {
builder = builder.attr_ints("dilations", dilations);
}
builder.build()
}
#[test]
fn test_avg_pool3d_config_basic() {
let mut node = create_test_node(
vec![3, 3, 3],
vec![1, 1, 1],
vec![0, 0, 0, 0, 0, 0],
0,
0,
None,
);
let processor = AvgPool3dProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.kernel_size, [3, 3, 3]);
assert_eq!(config.strides, [1, 1, 1]);
assert_eq!(config.dilation, [1, 1, 1]);
assert!(!config.count_include_pad);
assert!(!config.ceil_mode);
assert!(matches!(config.padding, PaddingConfig3d::Valid));
}
#[test]
fn test_avg_pool3d_config_with_padding() {
let mut node = create_test_node(
vec![2, 2, 2],
vec![2, 2, 2],
vec![1, 1, 1, 1, 1, 1],
0,
0,
None,
);
let processor = AvgPool3dProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.kernel_size, [2, 2, 2]);
assert_eq!(config.strides, [2, 2, 2]);
assert!(matches!(
config.padding,
PaddingConfig3d::Explicit(1, 1, 1, 1, 1, 1)
));
}
#[test]
fn test_avg_pool3d_config_with_dilation() {
let mut node = create_test_node(
vec![3, 3, 3],
vec![1, 1, 1],
vec![0, 0, 0, 0, 0, 0],
0,
0,
Some(vec![2, 2, 2]),
);
let processor = AvgPool3dProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.dilation, [2, 2, 2]);
}
#[test]
fn test_avg_pool3d_dilation_opset_validation() {
let mut node = create_test_node(
vec![3, 3, 3],
vec![1, 1, 1],
vec![0, 0, 0, 0, 0, 0],
0,
0,
Some(vec![2, 2, 2]),
);
let processor = AvgPool3dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 9, &prefs);
assert!(matches!(result, Err(ProcessError::Custom(_))));
}
#[test]
fn test_avg_pool3d_ceil_mode_opset_validation() {
let mut node = create_test_node(
vec![3, 3, 3],
vec![1, 1, 1],
vec![0, 0, 0, 0, 0, 0],
0,
1,
None,
);
let processor = AvgPool3dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 18, &prefs);
assert!(matches!(result, Err(ProcessError::Custom(_))));
}
#[test]
fn test_avg_pool3d_rejects_non_5d_input() {
let mut node = TestNodeBuilder::new(NodeType::AveragePool3d, "test_avgpool3d_bad")
.input_tensor_f32("data", 4, None)
.output_tensor_f32("output", 4, None)
.attr_ints("kernel_shape", vec![3, 3, 3])
.attr_ints("strides", vec![1, 1, 1])
.attr_ints("pads", vec![0, 0, 0, 0, 0, 0])
.attr_int("count_include_pad", 0)
.attr_int("ceil_mode", 0)
.build();
let processor = AvgPool3dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(matches!(result, Err(ProcessError::Custom(_))));
}
}