[Rewriter]: introduce normalize_pad_format_conv (#2301)

Convert 'auto_pad' attribute into a list of explicit pads.

Author: Johansmm

onnxscript/rewriter/fuse_pad_into_conv.py

@@ -3,6 +3,8 @@
 """Fuses Pad nodes into preceding nodes. Supported fusion patterns:
 - Pad ∘ Conv          -> Conv
 - Pad ∘ ConvInteger   -> ConvInteger
+
+To make some rules possible, we implicitly transform `auto_pad` attribute into its explicit list.
 """
 
 import typing
@@ -29,6 +31,26 @@ def fill_pads_with_axes(
     return new_pads
 
 
+def read_conv_attributes(ir_conv: ir.Node) -> dict[str, typing.Sequence[int] | str]:
+    # Read attributes
+    attributes = {}
+    if (kernel_shape := ir_conv.attributes.get("kernel_shape", None)) is not None:
+        attributes["kernel_shape"] = kernel_shape.as_ints()
+    else:
+        attributes["kernel_shape"] = ir_conv.inputs[1].shape[2:]
+    if (strides := ir_conv.attributes.get("strides", None)) is not None:
+        attributes["strides"] = strides.as_ints()
+    else:
+        attributes["strides"] = [1] * len(ir_conv.inputs[0].shape[2:])
+    if (auto_pad := ir_conv.attributes.get("auto_pad", None)) is not None:
+        attributes["auto_pad"] = auto_pad.as_string()
+    else:
+        attributes["auto_pad"] = "NOTSET"
+    if (pads := ir_conv.attributes.get("pads", None)) is not None:
+        attributes["pads"] = pads.as_ints()
+    return attributes
+
+
 class _FusePadConvBase(orp.RewriteRuleClassBase):
     """Interface for PadConv nodes fusion."""
 
@@ -145,6 +167,103 @@ def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
         )
 
 
+class _NormalizePadFormatBase(orp.RewriteRuleClassBase):
+    """Interface to normalize pad attributes in conv nodes."""
+
+    def compute_pads(
+        self,
+        input_shape: typing.Sequence[int],
+        output_shape: typing.Sequence[int],
+        attributes: dict[str, typing.Sequence[int] | str],
+    ) -> typing.Sequence[int]:
+        raise NotImplementedError("Child have to implement this function")
+
+    def rewrite(self, op: ir.tape.Tape, conv: ir.Value, **__) -> ir.Value:
+        cnode = conv.producer()
+
+        # Read spatial dimensions and attributes
+        input_shape = cnode.inputs[0].shape[2:]
+        output_shape = cnode.outputs[0].shape[2:]
+        attributes = read_conv_attributes(cnode)
+
+        # Convert auto_pad mode into an explicit list
+        pads = self.compute_pads(input_shape, output_shape, attributes)
+
+        # Replace auto_pad, forcing to the explicit list
+        conv_attr: typing.Mapping[str, ir.Attr] = cnode.attributes.copy()
+        conv_attr["auto_pad"] = ir.convenience.convert_attribute("auto_pad", "NOTSET")
+        if any(x != 0 for x in pads):
+            conv_attr["pads"] = ir.convenience.convert_attribute("pads", pads)
+
+        return op.op(
+            cnode.op_type,
+            inputs=cnode.inputs,
+            attributes=conv_attr,
+            domain=cnode.domain,
+            name=cnode.name,
+        )
+
+    def check(self, context, conv: ir.Value, **__) -> orp.MatchResult:
+        del context
+        check_result = orp.MatchResult()
+
+        # Conv constraints: attributes
+        cnode = conv.producer()
+        auto_pad = cnode.attributes.get("auto_pad", None)
+        if auto_pad is None or auto_pad.as_string() == "NOTSET":
+            return check_result.fail(f"{cnode.name} auto_pad must be different to 'NOTSET'.")
+
+        # Conv constraints: inputs/outputs
+        if cnode.inputs[0].shape is None:
+            return check_result.fail(f"Input shapes are not defined on {cnode.name}.")
+        if cnode.outputs[0].shape is None:
+            return check_result.fail(f"Output shapes are not defined on {cnode.name}.")
+        return check_result
+
+
+class NormalizePadFormatConv(_NormalizePadFormatBase):
+    """Convert auto_pad attribute into 'NOTSET' in Conv nodes ."""
+
+    def compute_pads(
+        self,
+        input_shape: typing.Sequence[int],
+        output_shape: typing.Sequence[int],
+        attributes: dict[str, typing.Sequence[int] | str],
+    ) -> typing.Sequence[int]:
+        # Compute pads, following auto_pad/pads attributes
+        if attributes["auto_pad"] in ["NOTSET", "VALID"]:
+            return attributes.get("pads", [0] * len(input_shape) * 2)
+
+        bottom_pads, top_pads = [], []
+        kernel_shape, strides = attributes["kernel_shape"], attributes["strides"]
+        for x, y, k, s in zip(input_shape, output_shape, kernel_shape, strides):
+            # Compute the output shape and the total padding to apply
+            total_pads = max(0, (y - 1) * s + k - x)
+
+            # Depending of mode, apply the padding to the upper or lower part
+            pad1 = total_pads // 2
+            pad2 = total_pads - pad1
+            if attributes["auto_pad"] == "SAME_UPPER":
+                bottom_pads.append(pad1)
+                top_pads.append(pad2)
+            else:
+                top_pads.append(pad1)
+                bottom_pads.append(pad2)
+        return bottom_pads + top_pads
+
+    def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
+        return op.Conv(x, _allow_other_inputs=True, _outputs=["conv"])
+
+
+class NormalizePadFormatConvInteger(NormalizePadFormatConv):
+    """Convert auto_pad attribute into 'NOTSET' in ConvInteger nodes ."""
+
+    def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
+        return op.ConvInteger(x, _allow_other_inputs=True, _outputs=["conv"])
+
+
+normalize_pad_format_conv = NormalizePadFormatConv.rule()
+normalize_pad_format_conv_integer = NormalizePadFormatConvInteger.rule()
 fuse_pad_into_conv = FusePadConv.rule()
 fuse_pad_into_conv_integer = FusePadConvInteger.rule()
 
@@ -159,6 +278,8 @@ def fuse_pad_into_conv_rule_set() -> orp.RewriteRuleSet:
     """
     return orp.RewriteRuleSet(
         [
+            normalize_pad_format_conv,
+            normalize_pad_format_conv_integer,
             fuse_pad_into_conv,
             fuse_pad_into_conv_integer,
         ]

onnxscript/rewriter/fuse_pad_into_conv_test.py

@@ -6,13 +6,14 @@
 import numpy as np
 import onnx_ir as ir
 import parameterized
-from onnx_ir.passes.common import onnx_checker
+from onnx_ir.passes.common import onnx_checker, shape_inference
 
 from onnxscript.rewriter import pattern as orp
 from onnxscript.rewriter import testing
 from onnxscript.rewriter.fuse_pad_into_conv import (
     fuse_pad_into_conv,
     fuse_pad_into_conv_rule_set,
+    normalize_pad_format_conv,
 )
 
 
@@ -83,22 +84,24 @@ def build_model(
             ir_version=9,
         )
         onnx_checker.CheckerPass(True)(ir_model)
+        ir_model = shape_inference.infer_shapes(ir_model)
         return ir_model
 
 
 class FusePadConvTest(FusePadConvBaseTest):
     @parameterized.parameterized.expand(
         [
-            (pad_pads, const_value, axes, conv_pads)
-            for pad_pads, axes, conv_pads in [
-                ([0, 0, 2, 2, 0, 0, 2, 2], None, None),
-                ([0, 2, 2, 0, 2, 2], ir.tensor([1, -2, -1], name="axes"), [2, 0, 2, 0]),
-                ([1, 1, 1, 1], ir.tensor([-2, 3], name="axes"), [0, 1, 0, 1]),
+            (pad_pads, const_value, axes, conv_pads, conv_auto_pad)
+            for pad_pads, axes, conv_pads, conv_auto_pad in [
+                ([0, 0, 2, 2, 0, 0, 2, 2], None, None, None),
+                ([0, 2, 2, 0, 2, 2], ir.tensor([1, -2, -1], name="axes"), [2, 0, 2, 0], None),
+                ([1, 1, 1, 1], ir.tensor([-2, 3], name="axes"), [0, 1, 0, 1], None),
+                ([1, 3, 1, 3], ir.tensor([3, 2], name="axes"), None, "VALID"),
             ]
             for const_value in [None, 0.0]
         ]
     )
-    def test_fuse_pad_into_conv(self, pad_pads, const_value, axes, conv_pads):
+    def test_fuse_pad_into_conv(self, pad_pads, const_value, axes, conv_pads, conv_auto_pad):
         pad_inputs = [ir.tensor(pad_pads, name="pads")]
         if const_value is not None or axes is not None:
             pad_inputs.append(const_value)
@@ -109,15 +112,15 @@ def test_fuse_pad_into_conv(self, pad_pads, const_value, axes, conv_pads):
             input_shape=ir.Shape(("N", 32, 14, 16)),
             weight_shape=(10, 32, 3, 3),
             pad_inputs=pad_inputs,
-            conv_attributes={"pads": conv_pads},
+            conv_attributes={"pads": conv_pads, "auto_pad": conv_auto_pad},
         )
         updated_model = _clone_model(base_model)
 
         # Apply rule
         count = fuse_pad_into_conv_rule_set().apply_to_model(updated_model)
 
         # Check that Pad was fused
-        self.assertEqual(count, 1)
+        self.assertEqual(count, 1 if conv_auto_pad is None else 2)
         self.assertEqual(updated_model.graph.num_nodes(), 1)
         onnx_checker.CheckerPass(True)(updated_model)
 
@@ -223,16 +226,19 @@ def get_conv_weights(self, shape: typing.Sequence[int], tape: ir.tape.Tape = Non
 
     @parameterized.parameterized.expand(
         [
-            (pad_pads, const_value, axes, conv_pads)
-            for pad_pads, axes, conv_pads in [
-                ([0, 0, 3, 2, 0, 0, 1, 4], None, [1, 1, 1, 1]),
-                ([2, 2, 0, 2, 2, 0], ir.tensor([-2, -1, 1], name="axes"), None),
-                ([1, 2, 2, 1], ir.tensor([-1, 2], name="axes"), [0, 1, 0, 1]),
+            (pad_pads, const_value, axes, conv_pads, conv_auto_pad)
+            for pad_pads, axes, conv_pads, conv_auto_pad in [
+                ([0, 0, 3, 2, 0, 0, 1, 4], None, [1, 1, 1, 1], None),
+                ([2, 2, 0, 2, 2, 0], ir.tensor([-2, -1, 1], name="axes"), None, None),
+                ([1, 2, 2, 1], ir.tensor([-1, 2], name="axes"), [0, 1, 0, 1], None),
+                ([3, 3], ir.tensor([2], name="axes"), None, "SAME_UPPER"),
             ]
             for const_value in [None, ir.tensor(np.array([0], "uint8"), name="const_value")]
         ]
     )
-    def test_fuse_pad_into_conv_integer(self, pad_pads, const_value, axes, conv_pads):
+    def test_fuse_pad_into_conv_integer(
+        self, pad_pads, const_value, axes, conv_pads, conv_auto_pad
+    ):
         pad_inputs = [ir.tensor(pad_pads, name="pads")]
         if const_value is not None or axes is not None:
             pad_inputs.append(const_value)
@@ -243,15 +249,15 @@ def test_fuse_pad_into_conv_integer(self, pad_pads, const_value, axes, conv_pads
             input_shape=ir.Shape(("N", 24, 19, 23)),
             weight_shape=(8, 24, 3, 3),
             pad_inputs=pad_inputs,
-            conv_attributes={"pads": conv_pads},
+            conv_attributes={"pads": conv_pads, "auto_pad": conv_auto_pad},
         )
         updated_model = _clone_model(base_model)
 
         # Apply rule
         count = fuse_pad_into_conv_rule_set().apply_to_model(updated_model)
 
         # Check that Pad was fused
-        self.assertEqual(count, 1)
+        self.assertEqual(count, 1 if conv_auto_pad is None else 2)
         self.assertEqual(updated_model.graph.num_nodes(), 1)
         onnx_checker.CheckerPass(True)(updated_model)
 
@@ -260,5 +266,67 @@ def test_fuse_pad_into_conv_integer(self, pad_pads, const_value, axes, conv_pads
         testing.assert_numerically_equal(base_model, updated_model, (inputs,), atol=0, rtol=0)
 
 
+class NormalizePadFormatTest(FusePadConvBaseTest):
+    @parameterized.parameterized.expand(
+        [
+            (strides, kernel_shape, auto_pad)
+            for strides, kernel_shape in [((2, 3), (1, 4)), ((2, 1), (5, 2))]
+            for auto_pad in ["SAME_UPPER", "SAME_LOWER", "VALID"]
+        ]
+    )
+    def test_normalize_pad_format(self, strides, kernel_shape, auto_pad):
+        pad_inputs = [
+            ir.tensor([1, 1, 1, 1], name="pads"),
+            None,
+            ir.tensor([2, 3], name="axes"),
+        ]
+        base_model = self.build_model(
+            op_type="Conv",
+            input_shape=ir.Shape(("N", 32, 22, 27)),
+            weight_shape=(32, 32, *kernel_shape),
+            pad_inputs=pad_inputs,
+            conv_attributes={
+                "strides": strides,
+                "auto_pad": auto_pad,
+                "kernel_shape": kernel_shape,
+            },
+        )
+        updated_model = _clone_model(base_model)
+
+        # Apply rule
+        count = fuse_pad_into_conv_rule_set().apply_to_model(updated_model)
+
+        # Check that Pad was fused
+        self.assertEqual(count, 2)
+        self.assertEqual(updated_model.graph.num_nodes(), 1)
+        onnx_checker.CheckerPass(True)(updated_model)
+
+        # Check inference
+        inputs = self.rng.random((1, 32, 22, 27), dtype="float32")
+        testing.assert_numerically_equal(base_model, updated_model, (inputs,), atol=0, rtol=0)
+
+    def test_unsupported_normalize_pad_format(self):
+        base_model = self.build_model(
+            op_type="Conv",
+            input_shape=ir.Shape(("N", 32, 14)),
+            weight_shape=(32, 11, 4),
+            pad_inputs=[ir.tensor([0, 0, 0, 0, 0, 0], name="pads")],
+            conv_attributes={"auto_pad": "VALID"},
+        )
+        # Drop convolutional input shape
+        base_model.graph[0].outputs[0].shape = None
+        onnx_checker.CheckerPass(True)(base_model)
+
+        # Apply rule and check it was not applied
+        tracer = orp.MatchingTracer()
+        count = normalize_pad_format_conv.apply_to_model(base_model, tracer=tracer)
+        self.assertEqual(count, 0)
+
+        # Check that the error message is the expected one
+        tracer_match = tracer.best_matches_map[normalize_pad_format_conv][0]
+        self.assertEqual(tracer_match.status.value, orp.MatchStatus.CONDITION_FAILED)
+        self.assertRegex(tracer_match.match_result.reason, "Input shapes are not defined")
+
+
 if __name__ == "__main__":
     unittest.main()