[Rewriter] improve docstring and rename objects (#2301)

Author: Johansmm

onnxscript/rewriter/__init__.py

@@ -19,9 +19,9 @@
     broadcast_to_matmul,
     cast_constant_of_shape,
     collapse_slices,
+    fuse_pad_into_conv,
     no_op,
     pattern,
-    fuse_pad_into_conv,
 )
 
 _ModelProtoOrIr = TypeVar("_ModelProtoOrIr", onnx.ModelProto, ir.Model)

onnxscript/rewriter/fuse_pad_into_conv.py

@@ -18,6 +18,20 @@
 
 
 def fill_pads_with_axes(pads: Sequence[int], axes: Sequence[int], rank: int) -> List[int]:
+    """Converts the parameters of the ONNX Pad operator into an explicit list of values.
+
+    A filled list of pads will be returned following the format:
+    [x1_begin, x2_begin, ..., x{rank}_begin, x1_end, x2_end, ..., x{rank}_end]
+
+    Args:
+        pads: list of integers indicating the number of padding elements to add at
+            the beginning and end of each axis.
+        axes: list of axes that pads apply to.
+        rank: value to compute the size of the filled list (2 * rank).
+
+    Returns:
+        The filled list of pads.
+    """
     new_pads = []
     for axis in range(rank):
         if axis not in axes:
@@ -47,11 +61,13 @@ def read_conv_attributes(ir_conv: ir.Node) -> dict[str, Sequence[int] | str]:
     return attributes
 
 
-class _FusePadConvBase(orp.RewriteRuleClassBase):
+class _FuseConvPadBase(orp.RewriteRuleClassBase):
     """Interface for PadConv nodes fusion."""
 
     def __init__(self, as_function: bool = False):
-        # Remove nodes is set to False to remove unused nodes after the rewrite.
+        # Remove nodes is set to False to remove unused nodes after the rewrite, since
+        # Pad or Conv inputs can come from constant nodes.
+        # With remove_nodes=False these nodes are removed if these nodes are no longer needed.
         super().__init__(remove_nodes=False, as_function=as_function)
 
     def rewrite(
@@ -89,6 +105,22 @@ def rewrite(
         )
 
     def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.MatchResult:
+        """Condition to check if we need to replace the pattern.
+
+        If Pad inputs can be added in 'pads' attribute of the Conv operator.
+
+        To validate this, we need to check the following:
+        1. `Pad<mode>` attribute has 'constant' as value
+        2. `Pad` operator inputs are constants ('pads', 'constant_value', 'axes')
+        3. 'constant_value' is equal to 0.0.
+        4. `Pad` operator is only used for the spatial dimensions (batch dimension and channels
+           remain unchanged).
+
+        If the above are true, then we don't need the reshapes.
+
+        Returns:
+            True if we need to replace the pattern, False otherwise.
+        """
         del context  # Unused
         check_result = orp.MatchResult()
         pad_node = pad.producer()
@@ -123,8 +155,8 @@ def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.Matc
         return check_result
 
 
-class FusePadConv(_FusePadConvBase):
-    """Replaces ``Pad(Conv(x))`` with ``Conv(x)``."""
+class FuseConvPad(_FuseConvPadBase):
+    """Replaces ``Conv(Pad(x))`` with ``Conv(x)``."""
 
     def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
         return op.Conv(
@@ -147,8 +179,8 @@ def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.Matc
         return check_result
 
 
-class FusePadConvInteger(FusePadConv):
-    """Replaces ``Pad(ConvInteger(x))`` with ``ConvInteger(x)``."""
+class FuseConvIntegerPad(FuseConvPad):
+    """Replaces ``ConvInteger(Pad(x))`` with ``ConvInteger(x)``."""
 
     def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
         return op.ConvInteger(
@@ -195,6 +227,24 @@ def rewrite(self, op: ir.tape.Tape, conv: ir.Value, **__) -> ir.Value:
         )
 
     def check(self, context, conv: ir.Value, **__) -> orp.MatchResult:
+        """Condition to check if we need to replace the pattern.
+
+        If it is possible to deduce 'pads'.
+
+        To validate this, we need to check the following:
+        1. `Conv<auto_pad != "NOTSET">` (nothing to do in this case, since 'pads' are
+           already explicit)
+        2. it is possible to deduce the input rank when `Conv<auto_pad == "VALID">`
+        3. When `Conv<auto_pad != "VALID">`:
+            * spatial input/output shapes are static
+            * it is possible to infer `kernel_shape` either from the `Conv` operator attribute
+              or from the kernel input
+
+        If the above are true, then we don't need the reshapes.
+
+        Returns:
+            True if we need to replace the pattern, False otherwise.
+        """
         del context
         check_result = orp.MatchResult()
 
@@ -274,8 +324,8 @@ def pattern(self, op: ir.tape.Tape, x: ir.Value) -> ir.Value:
 
 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()
+fuse_pad_into_conv = FuseConvPad.rule()
+fuse_pad_into_conv_integer = FuseConvIntegerPad.rule()
 
 
 def fuse_pad_into_conv_rule_set() -> orp.RewriteRuleSet:

onnxscript/rewriter/fuse_pad_into_conv_test.py

@@ -23,7 +23,7 @@ def _clone_model(model: ir.Model) -> ir.Model:
     return ir.from_proto(ir.to_proto(model))
 
 
-class FusePadConvBaseTest(unittest.TestCase):
+class FuseConvPadBaseTest(unittest.TestCase):
     @property
     def rng(self):
         return np.random.default_rng(20250522)
@@ -89,7 +89,7 @@ def build_model(
         return ir_model
 
 
-class FusePadConvTest(FusePadConvBaseTest):
+class FuseConvPadTest(FuseConvPadBaseTest):
     @parameterized.parameterized.expand(
         [
             (pad_pads, const_value, axes, conv_pads, conv_auto_pad)
@@ -218,7 +218,7 @@ def test_unsupported_fuse_pad_into_conv(
         self.assertRegex(tracer_match.match_result.reason, err_msg)
 
 
-class FusePadConvIntegerTest(FusePadConvBaseTest):
+class FuseConvIntegerPadTest(FuseConvPadBaseTest):
     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:
@@ -267,7 +267,7 @@ def test_fuse_pad_into_conv_integer(
         testing.assert_numerically_equal(base_model, updated_model, (inputs,), atol=0, rtol=0)
 
 
-class NormalizePadFormatTest(FusePadConvBaseTest):
+class NormalizePadFormatTest(FuseConvPadBaseTest):
     def build_model(
         self,
         input_shape: ir.Shape,