Fixes in Graph mode Quantizer

src/coreai_opt/_utils/torch_utils.py

@@ -39,6 +39,9 @@
     torch.ops.aten.conv1d.default: torch.nn.Conv1d,
     torch.ops.aten.conv2d.default: torch.nn.Conv2d,
     torch.ops.aten.conv3d.default: torch.nn.Conv3d,
+    torch.ops.aten.conv1d.padding: torch.nn.Conv1d,
+    torch.ops.aten.conv2d.padding: torch.nn.Conv2d,
+    torch.ops.aten.conv3d.padding: torch.nn.Conv3d,
     torch.ops.aten.conv_transpose1d.default: torch.nn.ConvTranspose1d,
     torch.ops.aten.conv_transpose2d.input: torch.nn.ConvTranspose2d,
     torch.ops.aten.conv_transpose3d.input: torch.nn.ConvTranspose3d,

src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py

@@ -405,20 +405,20 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         for conv_dim in [1, 2, 3]:
             for act_fn in _annotation_utils._supported_activations:
                 for act_in_place in [True, False]:
-                    pattern_gm = _annotation_utils.get_conv_bn_pattern(
+                    pattern_gms = _annotation_utils.get_conv_bn_pattern(
                         conv_dim=conv_dim,
                         is_transpose=False,
                         act_fn=act_fn,
                         act_in_place=act_in_place,
                     )
-                    patterns.append(pattern_gm)
+                    patterns.extend(pattern_gms)
 
             # Add patterns for activations without inplace support
             for act_fn in _annotation_utils._supported_activations_no_inplace:
-                pattern_gm = _annotation_utils.get_conv_bn_pattern(
+                pattern_gms = _annotation_utils.get_conv_bn_pattern(
                     conv_dim=conv_dim, is_transpose=False, act_fn=act_fn, act_in_place=False
                 )
-                patterns.append(pattern_gm)
+                patterns.extend(pattern_gms)
         return patterns
 
 
@@ -435,20 +435,20 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         for conv_dim in [1, 2, 3]:
             for act_fn in _annotation_utils._supported_activations:
                 for act_in_place in [True, False]:
-                    pattern_gm = _annotation_utils.get_conv_bn_pattern(
+                    pattern_gms = _annotation_utils.get_conv_bn_pattern(
                         conv_dim=conv_dim,
                         is_transpose=True,
                         act_fn=act_fn,
                         act_in_place=act_in_place,
                     )
-                    patterns.append(pattern_gm)
+                    patterns.extend(pattern_gms)
 
             # Add patterns for activations without inplace support
             for act_fn in _annotation_utils._supported_activations_no_inplace:
-                pattern_gm = _annotation_utils.get_conv_bn_pattern(
+                pattern_gms = _annotation_utils.get_conv_bn_pattern(
                     conv_dim=conv_dim, is_transpose=True, act_fn=act_fn, act_in_place=False
                 )
-                patterns.append(pattern_gm)
+                patterns.extend(pattern_gms)
 
         return patterns
 
@@ -466,17 +466,17 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         for conv_dim in [1, 2, 3]:
             for act_fn in _annotation_utils._supported_activations:
                 for act_in_place in [True, False]:
-                    pattern_gm = _annotation_utils.get_conv_pattern(
+                    pattern_gms = _annotation_utils.get_conv_pattern(
                         conv_dim=conv_dim, act_fn=act_fn, act_in_place=act_in_place
                     )
-                    patterns.append(pattern_gm)
+                    patterns.extend(pattern_gms)
 
             # Add patterns for activations without inplace support
             for act_fn in _annotation_utils._supported_activations_no_inplace:
-                pattern_gm = _annotation_utils.get_conv_pattern(
+                pattern_gms = _annotation_utils.get_conv_pattern(
                     conv_dim=conv_dim, act_fn=act_fn, act_in_place=False
                 )
-                patterns.append(pattern_gm)
+                patterns.extend(pattern_gms)
 
         return patterns
 
@@ -494,20 +494,20 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         for conv_dim in [1, 2, 3]:
             for act_fn in _annotation_utils._supported_activations:
                 for act_in_place in [True, False]:
-                    pattern_gm = _annotation_utils.get_conv_pattern(
+                    pattern_gms = _annotation_utils.get_conv_pattern(
                         conv_dim=conv_dim,
                         is_transpose=True,
                         act_fn=act_fn,
                         act_in_place=act_in_place,
                     )
-                    patterns.append(pattern_gm)
+                    patterns.extend(pattern_gms)
 
             # Add patterns for activations without inplace support
             for act_fn in _annotation_utils._supported_activations_no_inplace:
-                pattern_gm = _annotation_utils.get_conv_pattern(
+                pattern_gms = _annotation_utils.get_conv_pattern(
                     conv_dim=conv_dim, is_transpose=True, act_fn=act_fn, act_in_place=False
                 )
-                patterns.append(pattern_gm)
+                patterns.extend(pattern_gms)
 
         return patterns
 
@@ -523,10 +523,10 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         """Returns all conv pattern variations."""
         patterns = []
         for conv_dim in [1, 2, 3]:
-            pattern_gm = _annotation_utils.get_conv_bn_pattern(
+            pattern_gms = _annotation_utils.get_conv_bn_pattern(
                 conv_dim=conv_dim, is_transpose=False, act_fn=None
             )
-            patterns.append(pattern_gm)
+            patterns.extend(pattern_gms)
         return patterns
 
 
@@ -541,10 +541,10 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         """Returns all conv_transpose pattern variations."""
         patterns = []
         for conv_dim in [1, 2, 3]:
-            pattern_gm = _annotation_utils.get_conv_bn_pattern(
+            pattern_gms = _annotation_utils.get_conv_bn_pattern(
                 conv_dim=conv_dim, is_transpose=True, act_fn=None
             )
-            patterns.append(pattern_gm)
+            patterns.extend(pattern_gms)
         return patterns
 
 
@@ -559,8 +559,8 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         """Returns all conv pattern variations."""
         patterns = []
         for conv_dim in [1, 2, 3]:
-            pattern_gm = _annotation_utils.get_conv_pattern(conv_dim=conv_dim, act_fn=None)
-            patterns.append(pattern_gm)
+            pattern_gms = _annotation_utils.get_conv_pattern(conv_dim=conv_dim, act_fn=None)
+            patterns.extend(pattern_gms)
         return patterns
 
 
@@ -575,10 +575,10 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         """Returns all conv_transpose pattern variations."""
         patterns = []
         for conv_dim in [1, 2, 3]:
-            pattern_gm = _annotation_utils.get_conv_pattern(
+            pattern_gms = _annotation_utils.get_conv_pattern(
                 conv_dim=conv_dim, is_transpose=True, act_fn=None
             )
-            patterns.append(pattern_gm)
+            patterns.extend(pattern_gms)
         return patterns
 
 

src/coreai_opt/quantization/_graph/_annotation_utils.py

@@ -6,6 +6,7 @@
 import logging
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Literal
 
 import torch
@@ -181,7 +182,7 @@ def _get_weighted_mod_pattern(
     example_inputs: tuple[torch.Tensor, ...],
     act_fn: Callable | None = None,
     act_in_place: bool = False,
-) -> torch.nn.Module:
+) -> torch.fx.GraphModule:
     """
     Returns an aten graph corresponding to a sequence of these ops:
     input -> weighted_mod -> activation -> output
@@ -199,7 +200,6 @@ def _weighted_mod(input, weight, bias):
         output = mod_out
         node_dict = {
             "input": input,
-            "mod": mod_out,
             "weight": weight,
             "bias": bias,
         }
@@ -217,7 +217,7 @@ def _get_weighted_mod_bn_pattern(
     example_inputs: tuple[torch.Tensor, ...],
     act_fn: Callable | None = None,
     act_in_place: bool = False,
-) -> torch.nn.Module:
+) -> torch.fx.GraphModule:
     """
     Returns an aten graph corresponding to a sequence of these ops:
     input -> weighted_mod -> batch_norm -> activation -> output
@@ -240,7 +240,6 @@ def _weight_mod_bn(input, weight, bias, bn_weight, bn_bias, bn_run_mean, bn_run_
 
         node_dict = {
             "input": input,
-            "mod": mod_out,
             "bn": bn_out,
             "weight": weight,
             "bias": bias,
@@ -342,61 +341,66 @@ def forward(self, *args, **kwargs):
     )
 
 
-def get_conv_pattern(
-    conv_dim: int,
-    is_transpose: bool = False,
-    act_fn: Callable | None = None,
-    act_in_place: bool = False,
-) -> torch.nn.Module:
-    """
-    Returns an aten graph corresponding to a sequence of these ops:
-    input -> conv/conv_transpose -> activation -> output
+def _get_conv_fns(conv_dim: int, is_transpose: bool) -> list[Callable]:
+    """Return the conv functions to trace patterns from, one per aten overload of the conv.
 
-    No activation is used if ``act_fn`` is ``None``.
-    ``act_fn`` is an activation function from _supported_activations or
-    _supported_activations_no_inplace
+    A conv with integer padding exports as ``aten.convNd.default``, while one with string
+    padding ("same" / "valid") exports as ``aten.convNd.padding``.
+    Transposed convs only accept integer padding.
     """
     if is_transpose:
         assert conv_dim in _conv_transpose_fn_map, (
             f"Dimension {conv_dim} is not supported for ConvTranspose layers."
         )
+        conv_fns = [_conv_transpose_fn_map[conv_dim]]
     else:
         assert conv_dim in _conv_fn_map, (
             f"Dimension {conv_dim} is not supported for Convolution layers."
         )
+        conv_fn = _conv_fn_map[conv_dim]
+        conv_fns = [conv_fn, partial(conv_fn, padding="valid")]
+    return conv_fns
+
 
-    conv_func = _conv_transpose_fn_map[conv_dim] if is_transpose else _conv_fn_map[conv_dim]
+def get_conv_pattern(
+    conv_dim: int,
+    is_transpose: bool = False,
+    act_fn: Callable | None = None,
+    act_in_place: bool = False,
+) -> list[torch.fx.GraphModule]:
+    """
+    Returns aten graphs corresponding to a sequence of these ops:
+    input -> conv/conv_transpose -> activation -> output
 
+    One graph is returned per conv aten overload (see ``_get_conv_fns``).
+    No activation is used if ``act_fn`` is ``None``.
+    ``act_fn`` is an activation function from _supported_activations or
+    _supported_activations_no_inplace
+    """
     example_inputs = (
         torch.randn(1, 1, *[3] * conv_dim),  # input
         torch.randn(1, 1, *[1] * conv_dim),  # conv weight
         torch.randn(1),  # conv bias
     )
-    return _get_weighted_mod_pattern(conv_func, example_inputs, act_fn, act_in_place)
+    patterns = []
+    for conv_fn in _get_conv_fns(conv_dim, is_transpose):
+        pattern_gm = _get_weighted_mod_pattern(conv_fn, example_inputs, act_fn, act_in_place)
+        patterns.append(pattern_gm)
+    return patterns
 
 
 def get_conv_bn_pattern(
     conv_dim: int, is_transpose: bool, act_fn: Callable | None = None, act_in_place: bool = False
-) -> torch.nn.Module:
+) -> list[torch.fx.GraphModule]:
     """
-    Returns an aten graph corresponding to a sequence of these ops:
+    Returns aten graphs corresponding to a sequence of these ops:
     input -> conv -> batch_norm -> activation -> output
 
+    One graph is returned per conv aten overload (see ``_get_conv_fns``).
     No activation is used if ``act_fn`` is ``None``.
     ``act_fn`` is an activation function from _supported_activations or
     _supported_activations_no_inplace
     """
-    if is_transpose:
-        assert conv_dim in _conv_transpose_fn_map, (
-            f"Dimension {conv_dim} is not supported for ConvTranspose layers."
-        )
-    else:
-        assert conv_dim in _conv_fn_map, (
-            f"Dimension {conv_dim} is not supported for Convolution layers."
-        )
-
-    conv_func = _conv_transpose_fn_map[conv_dim] if is_transpose else _conv_fn_map[conv_dim]
-
     example_inputs = (
         torch.randn(1, 1, *[3] * conv_dim),  # input
         torch.randn(1, 1, *[1] * conv_dim),  # conv weight
@@ -406,7 +410,11 @@ def get_conv_bn_pattern(
         torch.randn(1),  # bn_run_mean
         torch.randn(1),  # bn_run_var
     )
-    return _get_weighted_mod_bn_pattern(conv_func, example_inputs, act_fn, act_in_place)
+    patterns = []
+    for conv_fn in _get_conv_fns(conv_dim, is_transpose):
+        pattern_gm = _get_weighted_mod_bn_pattern(conv_fn, example_inputs, act_fn, act_in_place)
+        patterns.append(pattern_gm)
+    return patterns
 
 
 def get_linear_pattern(
@@ -466,7 +474,6 @@ def _embedding(input, weight):
         output = mod_out
         node_dict = {
             "input": input,
-            "mod": mod_out,
             "weight": weight,
         }
         return output, node_dict
@@ -751,19 +758,42 @@ def match_pattern_with_subgraph_matcher(
     Helper function for matching model nodes with pattern using
     SubgraphMatcherWithNameNodeMap.
 
-    Returns a dictionary mapping nodes to matches.
+    Used for ``WeightedModulePattern`` patterns, whose name node map must have a
+    "weight" entry. The primary node of each match is the op in the match that
+    consumes the weight (the weight may also feed ops outside the match).
+
+    Returns a dictionary mapping each match's primary node to the match.
+
+    Raises:
+        RuntimeError: If the pattern's name node map has no "weight" entry.
     """
     node_to_match_dict: dict[torch.fx.Node, InternalMatch] = {}
     matcher = SubgraphMatcherWithNameNodeMap(pattern_gm, ignore_literals=True)
+    if "weight" not in matcher.name_node_map:
+        raise RuntimeError(
+            'Expected the pattern\'s name node map to have a "weight" entry, got keys '
+            f"{list(matcher.name_node_map)}."
+        )
     matches = matcher.match(model.graph)
 
     for subgraph_match in matches:
-        name_node_map = subgraph_match.name_node_map
-        mod_node = name_node_map["mod"]
-        assert mod_node not in node_to_match_dict
-
-        # Use the node marked with "mod" as the key
-        node_to_match_dict[mod_node] = subgraph_match
+        # "subgraph_match.nodes_map" maps each node in "pattern_gm" to its node in "model".
+        # To find "weight_users", keep only the model nodes whose pattern node is a
+        # call_function (not a placeholder).
+        # When weights are shared there maybe multiple users within the model, but only
+        # one should be present in the subgraph corresponding to the pattern.
+        matched_ops = {
+            model_node
+            for pattern_node, model_node in subgraph_match.nodes_map.items()
+            if pattern_node.op == "call_function"
+        }
+        weight_users = [
+            user for user in subgraph_match.name_node_map["weight"].users if user in matched_ops
+        ]
+        assert len(weight_users) == 1, f"Expected one weighted op in the match, got {weight_users}"
+        primary_node = weight_users[0]
+        assert primary_node not in node_to_match_dict
+        node_to_match_dict[primary_node] = subgraph_match
     return node_to_match_dict