use derive_new::new;
use onnx_ir_derive::NodeBuilder;
use crate::ir::{Argument, Node, RawNode};
use crate::node::padding::{AutoPad, PaddingConfig2d, padding_config_2d};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
#[derive(Debug, Clone, new)]
pub struct LpPool2dConfig {
pub kernel_size: [usize; 2],
pub strides: [usize; 2],
pub padding: PaddingConfig2d,
pub dilation: [usize; 2],
pub ceil_mode: bool,
pub auto_pad: AutoPad,
pub p: i64,
}
#[derive(Debug, Clone, NodeBuilder)]
pub struct LpPool2dNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: LpPool2dConfig,
}
pub(crate) struct LpPool2dProcessor;
impl NodeProcessor for LpPool2dProcessor {
type Config = LpPool2dConfig;
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> {
let mut has_kernel_shape = false;
for (key, value) in node.attrs.iter() {
match key.as_str() {
"kernel_shape" => {
has_kernel_shape = true;
let kernel_shape = value.clone().into_i64s();
if kernel_shape.len() != 2 {
return Err(ProcessError::Custom(format!(
"LpPool2d: kernel_shape must have length 2, got {:?}",
kernel_shape
)));
}
if kernel_shape[0] <= 0 || kernel_shape[1] <= 0 {
return Err(ProcessError::Custom(format!(
"LpPool2d: kernel_shape values must be > 0, got {:?}",
kernel_shape
)));
}
}
"strides" => {
let strides = value.clone().into_i64s();
if strides.len() != 2 {
return Err(ProcessError::Custom(format!(
"LpPool2d: strides must have length 2, got {:?}",
strides
)));
}
if strides[0] <= 0 || strides[1] <= 0 {
return Err(ProcessError::Custom(format!(
"LpPool2d: strides values must be > 0, got {:?}",
strides
)));
}
}
"pads" => {
let pads = value.clone().into_i64s();
if pads.len() != 4 {
return Err(ProcessError::Custom(format!(
"LpPool2d: pads must have length 4, got {:?}",
pads
)));
}
}
"p" => {
let p = value.clone().into_i64();
if p <= 0 {
return Err(ProcessError::Custom(format!(
"LpPool2d: p must be > 0, got {}",
p
)));
}
}
"ceil_mode" => {
let ceil_mode = value.clone().into_i64();
if ceil_mode != 0 && opset < 18 {
return Err(ProcessError::Custom(format!(
"LpPool: ceil_mode requires opset 18+, got opset {}",
opset
)));
}
}
"dilations" => {
let dilations = value.clone().into_i64s();
if dilations.len() != 2 {
return Err(ProcessError::Custom(format!(
"LpPool2d: dilations must have length 2, got {:?}",
dilations
)));
}
if dilations[0] <= 0 || dilations[1] <= 0 {
return Err(ProcessError::Custom(format!(
"LpPool2d: dilations values must be > 0, got {:?}",
dilations
)));
}
if dilations.iter().any(|&d| d != 1) {
if opset < 11 {
return Err(ProcessError::Custom(format!(
"LpPool2d: dilations requires opset 11+, got opset {}",
opset
)));
}
return Err(ProcessError::Custom(
"LpPool2d: dilations != 1 is not supported in burn-onnx yet"
.to_string(),
));
}
}
"auto_pad" => {
AutoPad::parse(&value.clone().into_string())?;
}
_ => {
return Err(ProcessError::InvalidAttribute {
name: key.clone(),
reason: format!("Unexpected attribute for LpPool2d: {key}"),
});
}
}
}
if !has_kernel_shape {
return Err(ProcessError::Custom(
"LpPool2d: missing required attribute kernel_shape".to_string(),
));
}
crate::processor::same_as_input(node);
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];
let mut pads = vec![0, 0, 0, 0];
let mut dilations = vec![1, 1];
let mut ceil_mode: i64 = 0;
let mut auto_pad = AutoPad::NotSet;
let mut p: i64 = 2;
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(),
"dilations" => dilations = value.clone().into_i64s(),
"ceil_mode" => ceil_mode = value.clone().into_i64(),
"auto_pad" => auto_pad = AutoPad::parse(&value.clone().into_string())?,
"p" => p = value.clone().into_i64(),
_ => {}
}
}
let padding = padding_config_2d(&pads);
let config = LpPool2dConfig::new(
[kernel_shape[0] as usize, kernel_shape[1] as usize],
[strides[0] as usize, strides[1] as usize],
padding,
[dilations[0] as usize, dilations[1] as usize],
ceil_mode != 0,
auto_pad,
p,
);
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
Node::LpPool2d(LpPool2dNode {
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>,
dilations: Option<Vec<i64>>,
ceil_mode: i64,
p: Option<i64>,
) -> RawNode {
let mut builder = TestNodeBuilder::new(NodeType::LpPool2d, "test_lppool2d")
.input_tensor_f32("data", 4, None)
.output_tensor_f32("output", 4, None)
.attr_ints("kernel_shape", kernel_shape)
.attr_ints("strides", strides)
.attr_ints("pads", pads)
.attr_int("ceil_mode", ceil_mode);
if let Some(dilations) = dilations {
builder = builder.attr_ints("dilations", dilations);
}
if let Some(p) = p {
builder = builder.attr_int("p", p);
}
builder.build()
}
#[test]
fn test_lppool2d_config_basic() {
let node = create_test_node(vec![2, 3], vec![1, 2], vec![0, 1, 1, 0], None, 0, None);
let mut node = node;
let processor = LpPool2dProcessor;
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, 3]);
assert_eq!(config.strides, [1, 2]);
assert_eq!(config.dilation, [1, 1]);
assert_eq!(config.p, 2);
assert!(!config.ceil_mode);
assert!(matches!(
config.padding,
PaddingConfig2d::Explicit(0, 1, 1, 0)
));
}
#[test]
fn test_lppool2d_config_with_custom_p() {
let node = create_test_node(vec![3, 3], vec![2, 2], vec![1, 1, 1, 1], None, 0, Some(4));
let mut node = node;
let processor = LpPool2dProcessor;
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]);
assert_eq!(config.strides, [2, 2]);
assert_eq!(config.p, 4);
}
#[test]
fn test_lppool2d_ceil_mode_opset_validation() {
let mut node = create_test_node(vec![2, 2], vec![1, 1], vec![0, 0, 0, 0], None, 1, None);
let processor = LpPool2dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(err_msg.contains("ceil_mode requires opset 18+"));
}
#[test]
fn test_lppool2d_missing_kernel_shape_validation() {
let mut node = TestNodeBuilder::new(NodeType::LpPool2d, "test_lppool2d_missing_kernel")
.input_tensor_f32("data", 4, None)
.output_tensor_f32("output", 4, None)
.build();
let processor = LpPool2dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(err_msg.contains("missing required attribute kernel_shape"));
}
#[test]
fn test_lppool2d_dilation_unsupported_validation() {
let mut node = create_test_node(
vec![2, 2],
vec![1, 1],
vec![0, 0, 0, 0],
Some(vec![2, 1]),
0,
None,
);
let processor = LpPool2dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(err_msg.contains("dilations != 1 is not supported"));
}
#[test]
fn test_lppool2d_dilation_opset_validation() {
let mut node = create_test_node(
vec![2, 2],
vec![1, 1],
vec![0, 0, 0, 0],
Some(vec![2, 1]),
0,
None,
);
let processor = LpPool2dProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 10, &prefs);
assert!(result.is_err());
let err_msg = format!("{}", result.unwrap_err());
assert!(err_msg.contains("dilations requires opset 11+"));
}
#[test]
fn test_lppool2d_non_positive_values_validation() {
let processor = LpPool2dProcessor;
let prefs = OutputPreferences::new();
let mut kernel_zero =
create_test_node(vec![0, 2], vec![1, 1], vec![0, 0, 0, 0], None, 0, None);
let err = processor
.infer_types(&mut kernel_zero, 16, &prefs)
.expect_err("Expected non-positive kernel_shape to fail");
assert!(format!("{}", err).contains("kernel_shape values must be > 0"));
let mut stride_zero =
create_test_node(vec![2, 2], vec![0, 1], vec![0, 0, 0, 0], None, 0, None);
let err = processor
.infer_types(&mut stride_zero, 16, &prefs)
.expect_err("Expected non-positive stride to fail");
assert!(format!("{}", err).contains("strides values must be > 0"));
let mut dilation_zero = create_test_node(
vec![2, 2],
vec![1, 1],
vec![0, 0, 0, 0],
Some(vec![0, 1]),
0,
None,
);
let err = processor
.infer_types(&mut dilation_zero, 16, &prefs)
.expect_err("Expected non-positive dilation to fail");
assert!(format!("{}", err).contains("dilations values must be > 0"));
let mut p_zero =
create_test_node(vec![2, 2], vec![1, 1], vec![0, 0, 0, 0], None, 0, Some(0));
let err = processor
.infer_types(&mut p_zero, 16, &prefs)
.expect_err("Expected non-positive p to fail");
assert!(format!("{}", err).contains("LpPool2d: p must be > 0"));
}
}