[Rewriter]: apply suggestions (#2301)

Author: Johansmm

onnxscript/rewriter/fuse_pad_into_conv.py

@@ -1,23 +1,23 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """Fuses Pad nodes into preceding nodes. Supported fusion patterns:
-- Pad ∘ Conv          -> Conv
-- Pad ∘ ConvInteger   -> ConvInteger
+- Conv ∘ Pad          -> Conv
+- ConvInteger ∘ Pad   -> ConvInteger
 
 To make some rules possible, we implicitly transform `auto_pad` attribute into its explicit list.
 """
 
-import typing
+from __future__ import annotations
+
+from typing import List, Sequence
 
 import numpy as np
 import onnx_ir as ir
 
 from onnxscript.rewriter import pattern as orp
 
 
-def fill_pads_with_axes(
-    pads: typing.Sequence[int], axes: typing.Sequence[int], rank: int
-) -> typing.List[int]:
+def fill_pads_with_axes(pads: Sequence[int], axes: Sequence[int], rank: int) -> List[int]:
     new_pads = [0] * 2 * rank
     N = len(axes)
     for start_idx, axis in enumerate(axes):
@@ -26,43 +26,39 @@ def fill_pads_with_axes(
     return new_pads
 
 
-def read_conv_attributes(ir_conv: ir.Node) -> dict[str, typing.Sequence[int] | str]:
+def read_conv_attributes(ir_conv: ir.Node) -> dict[str, 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()
+    ir_attributes = ir_conv.attributes
+    attributes["kernel_shape"] = ir_attributes.get_ints(
+        "kernel_shape", ir_conv.inputs[1].shape[2:]
+    )
+    attributes["strides"] = ir_attributes.get_ints(
+        "strides", [1] * len(ir_conv.inputs[0].shape[2:])
+    )
+    attributes["auto_pad"] = ir_attributes.get_string("auto_pad", "NOTSET")
+    if "pads" in ir_attributes:
+        attributes["pads"] = ir_attributes.get_ints("pads")
     return attributes
 
 
 class _FusePadConvBase(orp.RewriteRuleClassBase):
     """Interface for PadConv nodes fusion."""
 
-    def __init__(self, name: str, as_function: bool = False):
+    def __init__(self, as_function: bool = False):
         # Remove nodes is set to False to remove unused nodes after the rewrite.
-        super().__init__(name=name, remove_nodes=False, as_function=as_function)
+        super().__init__(remove_nodes=False, as_function=as_function)
 
     def rewrite(
         self, op: ir.tape.Tape, x: ir.Value, pad: ir.Value, conv: ir.Value
     ) -> ir.Value:
-        pnode = pad.producer()
-        cnode = conv.producer()
+        pad_node = pad.producer()
+        conv_node = conv.producer()
 
         # Retrieve the padding and axes
         x_rank = len(x.shape)
-        pad_pads = pnode.inputs[1].const_value.numpy().tolist()
-        if len(pnode.inputs) > 3 and (axes := pnode.inputs[3]) is not None:
+        pad_pads = pad_node.inputs[1].const_value.numpy().tolist()
+        if len(pad_node.inputs) > 3 and (axes := pad_node.inputs[3]) is not None:
             axes = [x if x >= 0 else x_rank + x for x in axes.const_value.numpy()]
         else:
             axes = list(range(x_rank))
@@ -74,41 +70,40 @@ def rewrite(
         new_pads = pad_pads[2:x_rank] + pad_pads[x_rank + 2 :]
 
         # Replace conv pads = new + old
-        conv_attr: typing.Mapping[str, ir.Attr] = cnode.attributes.copy()
+        conv_attr = conv_node.attributes.copy()
         if "pads" in conv_attr:
             new_pads = [x + y for x, y in zip(conv_attr["pads"].as_ints(), new_pads)]
-        conv_attr["pads"] = ir.convenience.convert_attribute("pads", new_pads)
+        conv_attr.add(ir.AttrInt64s("pads", new_pads))
 
         return op.op(
-            cnode.op_type,
-            inputs=(x, *cnode.inputs[1:]),
+            conv_node.op_type,
+            inputs=(x, *conv_node.inputs[1:]),
             attributes=conv_attr,
-            domain=cnode.domain,
-            name=cnode.name,
+            domain=conv_node.domain,
+            name=conv_node.name,
         )
 
     def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.MatchResult:
         del context  # Unused
         check_result = orp.MatchResult()
-        pnode = pad.producer()
+        pad_node = pad.producer()
         x_rank = len(x.shape)
 
         # Pad constraints: attributes
-        if (mode := pnode.attributes.get("mode", None)) and mode.as_string() != "constant":
-            return check_result.fail(f"{pnode.name} mode must be 'constant'.")
+        if (mode := pad_node.attributes.get("mode", None)) and mode.as_string() != "constant":
+            return check_result.fail(f"{pad_node.name} mode must be 'constant'.")
 
         # Pad constraints: inputs
-        if (pads := pnode.inputs[1]).const_value is None:
+        if (pads := pad_node.inputs[1]).const_value is None:
             return check_result.fail(f"{pads.name} is not a constant/initializer.")
-        if len(pnode.inputs) > 2 and (constant_value := pnode.inputs[2]) is not None:
+        if len(pad_node.inputs) > 2 and (constant_value := pad_node.inputs[2]) is not None:
             if constant_value.const_value is None:
                 return check_result.fail(
                     f"{constant_value.name} is not a constant/initializer."
                 )
             elif constant_value.const_value.numpy().item() != 0:
                 return check_result.fail(f"{constant_value.name} must be equal to 0.")
-        axes = list(range(x_rank))
-        if len(pnode.inputs) > 3 and (axes := pnode.inputs[3]) is not None:
+        if len(pad_node.inputs) > 3 and (axes := pad_node.inputs[3]) is not None:
             if axes.const_value is None:
                 return check_result.fail(f"{axes.name} is not a constant/initializer.")
             axes_list = [x if x >= 0 else x_rank + x for x in axes.const_value.numpy()]
@@ -126,9 +121,6 @@ def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.Matc
 class FusePadConv(_FusePadConvBase):
     """Replaces ``Pad(Conv(x))`` with ``Conv(x)``."""
 
-    def __init__(self, as_function: bool = False):
-        super().__init__(name="FusePadConv", as_function=as_function)
-
     def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
         return op.Conv(
             op.Pad(x, _allow_other_inputs=True, _outputs=["pad"]),
@@ -142,18 +134,17 @@ def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.Matc
             return check_result
 
         # Conv constraints: attributes
-        cnode = conv.producer()
-        if (apad := cnode.attributes.get("auto_pad", None)) and apad.as_string() != "NOTSET":
-            return check_result.fail(f"{cnode.name} auto_pad must be 'NOTSET'.")
+        conv_node = conv.producer()
+        if (
+            apad := conv_node.attributes.get("auto_pad", None)
+        ) and apad.as_string() != "NOTSET":
+            return check_result.fail(f"{conv_node.name} auto_pad must be 'NOTSET'.")
         return check_result
 
 
 class FusePadConvInteger(FusePadConv):
     """Replaces ``Pad(ConvInteger(x))`` with ``ConvInteger(x)``."""
 
-    def __init__(self, as_function: bool = False):
-        super(FusePadConv, self).__init__(name="FusePadConvInteger", as_function=as_function)
-
     def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
         return op.ConvInteger(
             op.Pad(x, _allow_other_inputs=True, _outputs=["pad"]),
@@ -165,66 +156,68 @@ def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
 class _NormalizePadFormatBase(orp.RewriteRuleClassBase):
     """Interface to normalize pad attributes in conv nodes."""
 
+    @staticmethod
     def compute_pads(
-        self,
-        input_shape: typing.Sequence[int],
-        output_shape: typing.Sequence[int],
-        attributes: dict[str, typing.Sequence[int] | str],
-    ) -> typing.Sequence[int]:
+        input_shape: Sequence[int],
+        output_shape: Sequence[int],
+        attributes: dict[str, Sequence[int] | str],
+    ) -> 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()
+        conv_node = 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)
+        input_shape = conv_node.inputs[0].shape[2:]
+        output_shape = conv_node.outputs[0].shape[2:]
+        attributes = read_conv_attributes(conv_node)
 
         # 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")
+        conv_attr = conv_node.attributes.copy()
+        conv_attr.add(ir.AttrString("auto_pad", "NOTSET"))
         if any(x != 0 for x in pads):
-            conv_attr["pads"] = ir.convenience.convert_attribute("pads", pads)
+            conv_attr.add(ir.AttrInt64s("pads", pads))
 
         return op.op(
-            cnode.op_type,
-            inputs=cnode.inputs,
+            conv_node.op_type,
+            inputs=conv_node.inputs,
             attributes=conv_attr,
-            domain=cnode.domain,
-            name=cnode.name,
+            domain=conv_node.domain,
+            name=conv_node.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)
+        conv_node = conv.producer()
+        auto_pad = conv_node.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'.")
+            return check_result.fail(
+                f"{conv_node.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}.")
+        if conv_node.inputs[0].shape is None:
+            return check_result.fail(f"Input shapes are not defined on {conv_node.name}.")
+        if conv_node.outputs[0].shape is None:
+            return check_result.fail(f"Output shapes are not defined on {conv_node.name}.")
         return check_result
 
 
 class NormalizePadFormatConv(_NormalizePadFormatBase):
     """Convert auto_pad attribute into 'NOTSET' in Conv nodes ."""
 
+    @staticmethod
     def compute_pads(
-        self,
-        input_shape: typing.Sequence[int],
-        output_shape: typing.Sequence[int],
-        attributes: dict[str, typing.Sequence[int] | str],
-    ) -> typing.Sequence[int]:
+        input_shape: Sequence[int],
+        output_shape: Sequence[int],
+        attributes: dict[str, Sequence[int] | str],
+    ) -> 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)

onnxscript/rewriter/fuse_pad_into_conv_test.py

@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-import typing
+from __future__ import annotations
+
 import unittest
+from typing import Mapping, Sequence
 
 import numpy as np
 import onnx_ir as ir
@@ -26,7 +28,7 @@ class FusePadConvBaseTest(unittest.TestCase):
     def rng(self):
         return np.random.default_rng(20250522)
 
-    def get_conv_weights(self, shape: typing.Sequence[int], tape: ir.tape.Tape = None):
+    def get_conv_weights(self, shape: Sequence[int], tape: ir.tape.Tape = None):
         w = ir.tensor(self.rng.uniform(-0.5, 0.5, shape).astype("float32"), name="W")
         if tape is not None:
             w = tape.initializer(w)
@@ -36,11 +38,10 @@ def build_model(
         self,
         op_type: str,
         input_shape: ir.Shape,
-        weight_shape: typing.Sequence[int],
-        pad_inputs: typing.Sequence[ir.TensorProtocol | ir.Value | None],
-        pad_attributes: typing.Mapping[str, ir.Attr] | None = None,
-        conv_attributes: typing.Mapping[str, ir.Attr] | None = None,
-        opset_imports: typing.Mapping[str, int] = {"": 20},
+        weight_shape: Sequence[int],
+        pad_inputs: Sequence[ir.TensorProtocol | ir.Value | None],
+        pad_attributes: Mapping[str, ir.Attr] | None = None,
+        conv_attributes: Mapping[str, ir.Attr] | None = None,
     ) -> ir.Model:
         tape = ir.tape.Tape()
         inputs = []
@@ -78,10 +79,10 @@ def build_model(
                 outputs=[y],
                 nodes=tape.nodes,
                 initializers=tape.initializers,
-                opset_imports=opset_imports,
+                opset_imports={"": 20},
                 name="model",
             ),
-            ir_version=9,
+            ir_version=10,
         )
         onnx_checker.CheckerPass(True)(ir_model)
         ir_model = shape_inference.infer_shapes(ir_model)
@@ -218,7 +219,7 @@ def test_unsupported_fuse_pad_into_conv(
 
 
 class FusePadConvIntegerTest(FusePadConvBaseTest):
-    def get_conv_weights(self, shape: typing.Sequence[int], tape: ir.tape.Tape = None):
+    def get_conv_weights(self, shape: Sequence[int], tape: ir.tape.Tape = None):
         w = ir.tensor(self.rng.integers(0, 256, shape).astype("uint8"), name="W")
         if tape is not None:
             w = tape.initializer(w)