reimplement pool lowering

add pool validation
align PIM ops/codegen/parser with the ISA
move constant materialization to MLIR
rename the PIM verification/materialization passes
better folded-constant handling

validation/operations/README.md

@@ -23,6 +23,18 @@ python3 validation/operations/gen_tests.py
 | With bias 3x3 | `conv/with_bias_3x3` | [1,3,5,5] | [1,2,3,3] | 3x3 | 1 | none | yes | Multi-channel with bias |
 | Large spatial | `conv/large_spatial` | [1,1,8,8] | [1,1,6,6] | 3x3 | 1 | none | no | Larger spatial input |
 
+## Pool
+
+| Test | Directory | Input | Output | Kernel | Stride | Padding | Notes |
+|------|-----------|-------|--------|--------|--------|---------|-------|
+| Max basic | `pool/max_basic` | [1,1,4,4] | [1,1,3,3] | 2x2 | 1 | none | Basic max pooling |
+| Max stride 2 multi-channel | `pool/max_stride2_multichannel` | [1,5,6,6] | [1,5,3,3] | 2x2 | 2 | none | Channel-preserving max pool |
+| Max SAME_UPPER | `pool/max_same_upper` | [1,1,5,5] | [1,1,3,3] | 3x3 | 2 | SAME_UPPER | Deprecated auto_pad path |
+| Avg basic | `pool/avg_basic` | [1,3,4,4] | [1,3,3,3] | 2x2 | 1 | none | Basic average pooling |
+| Avg explicit padding | `pool/avg_explicit_padding` | [1,2,4,4] | [1,2,2,2] | 3x3 | 2 | [1,1,1,1] | `count_include_pad=0` |
+| Avg include pad | `pool/avg_include_pad` | [1,2,4,4] | [1,2,2,2] | 3x3 | 2 | [1,1,1,1] | `count_include_pad=1` |
+| Max after Conv | `pool/max_after_conv` | [1,3,6,6] | [1,4,2,2] | Conv 3x3 then Pool 2x2 | 2 | none | Regression for `pool(conv(...))` |
+
 ## Gemm
 
 | Test | Directory | A (input) | W (weight) | Output | transB | alpha | beta | Bias | Notes |

validation/operations/gen_tests.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Generate ONNX test models for validating GEMM and Conv implementations."""
+"""Generate ONNX test models for validating GEMM, Conv, and Pooling implementations."""
 
 import numpy as np
 import onnx
@@ -248,6 +248,85 @@ def conv_large_spatial():
     save_model(model, "conv/large_spatial", "conv_large_spatial.onnx")
 
 
+# ---------------------------------------------------------------------------
+# Pooling tests
+# ---------------------------------------------------------------------------
+
+def maxpool_basic():
+    """MaxPool 2x2 with stride 1."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 4, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 3, 3])
+    node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
+    graph = helper.make_graph([node], "maxpool_basic", [X], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/max_basic", "maxpool_basic.onnx")
+
+
+def maxpool_stride2_multichannel():
+    """MaxPool 2x2 with stride 2 on multiple channels."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 5, 6, 6])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 5, 3, 3])
+    node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0])
+    graph = helper.make_graph([node], "maxpool_stride2_multichannel", [X], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/max_stride2_multichannel", "maxpool_stride2_multichannel.onnx")
+
+
+def maxpool_same_upper():
+    """MaxPool 3x3 with SAME_UPPER padding."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 5, 5])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 3, 3])
+    node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[3, 3], strides=[2, 2], auto_pad="SAME_UPPER")
+    graph = helper.make_graph([node], "maxpool_same_upper", [X], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/max_same_upper", "maxpool_same_upper.onnx")
+
+
+def avgpool_basic():
+    """AveragePool 2x2 with stride 1."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 4, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 3, 3])
+    node = helper.make_node("AveragePool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
+    graph = helper.make_graph([node], "avgpool_basic", [X], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/avg_basic", "avgpool_basic.onnx")
+
+
+def avgpool_explicit_padding():
+    """AveragePool 3x3 with explicit padding, excluding pad from the divisor."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2, 4, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 2, 2])
+    node = helper.make_node("AveragePool", ["X"], ["Y"],
+                            kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 1, 1], count_include_pad=0)
+    graph = helper.make_graph([node], "avgpool_explicit_padding", [X], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/avg_explicit_padding", "avgpool_explicit_padding.onnx")
+
+
+def avgpool_include_pad():
+    """AveragePool 3x3 with explicit padding, including pad in the divisor."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2, 4, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 2, 2])
+    node = helper.make_node("AveragePool", ["X"], ["Y"],
+                            kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 1, 1], count_include_pad=1)
+    graph = helper.make_graph([node], "avgpool_include_pad", [X], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/avg_include_pad", "avgpool_include_pad.onnx")
+
+
+def maxpool_after_conv():
+    """Conv followed by MaxPool to validate pooling on lowered conv results."""
+    rng = np.random.default_rng(59)
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 6, 6])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4, 2, 2])
+    W = numpy_helper.from_array(rng.uniform(-1, 1, (4, 3, 3, 3)).astype(np.float32), name="W")
+    conv = helper.make_node("Conv", ["X", "W"], ["C"], kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
+    pool = helper.make_node("MaxPool", ["C"], ["Y"], kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0])
+    graph = helper.make_graph([conv, pool], "maxpool_after_conv", [X], [Y], initializer=[W])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "pool/max_after_conv", "maxpool_after_conv.onnx")
+
+
 # ---------------------------------------------------------------------------
 # Main
 # ---------------------------------------------------------------------------
@@ -273,4 +352,13 @@ if __name__ == "__main__":
     conv_batch_2()
     conv_large_spatial()
 
+    print("\nGenerating Pooling tests:")
+    maxpool_basic()
+    maxpool_stride2_multichannel()
+    maxpool_same_upper()
+    avgpool_basic()
+    avgpool_explicit_padding()
+    avgpool_include_pad()
+    maxpool_after_conv()
+
     print("\nDone.")