Update implementation of find_sequential_partitions

src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py

@@ -261,7 +261,7 @@ def match_single_pattern(
         return _annotation_utils.match_pattern_with_subgraph_matcher(model, pattern)
 
 
-class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition]]):
+class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition, ...]]):
     """
     Base class for annotation patterns that handle an op of one or more inputs along
     with an optional activation.
@@ -285,7 +285,7 @@ class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition]]):
     @classmethod
     def match_single_pattern(
         cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern
-    ) -> dict[torch.fx.Node, tuple[SourcePartition]]:
+    ) -> dict[torch.fx.Node, tuple[SourcePartition, ...]]:
         """
         Match nodes in model to the provided pattern and return a dictionary mapping
         nodes to matches.
@@ -300,7 +300,7 @@ def match_single_pattern(
         return _annotation_utils.match_pattern_with_sequential_partitions(model, pattern)
 
 
-class SharedObserverModulePattern(BaseAnnotationPattern[tuple[SourcePartition]]):
+class SharedObserverModulePattern(BaseAnnotationPattern[tuple[SourcePartition, ...]]):
     """
     Base class for annotation patterns that handle shared observer modules.
 
@@ -353,7 +353,7 @@ def _validate_patterns_length(cls):
     @classmethod
     def match_single_pattern(
         cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | OpsListPattern
-    ) -> dict[torch.fx.Node, tuple[SourcePartition]]:
+    ) -> dict[torch.fx.Node, tuple[SourcePartition, ...]]:
         """
         Match nodes in model to the provided pattern and return a dictionary mapping
         nodes to matches.

src/coreai_opt/quantization/_graph/_annotation_utils.py

@@ -19,7 +19,7 @@
     SubgraphMatcherWithNameNodeMap,
 )
 from torch.fx.passes.utils.source_matcher_utils import SourcePartition
-from torchao.quantization.pt2e import WrapperModule, find_sequential_partitions
+from torchao.quantization.pt2e import WrapperModule
 from torchao.quantization.pt2e.quantizer import (
     QuantizationAnnotation,
     QuantizationSpec as TorchAOQuantizationSpec,
@@ -35,6 +35,7 @@
 from coreai_opt._utils.fx_utils import (
     get_local_state_name,
     get_module_boundary_nodes,
+    get_node_type,
     is_coreai_compressed_state_node,
 )
 from coreai_opt._utils.python_utils import get_fn_arg_names
@@ -704,15 +705,6 @@ def _get_call_function_node_from_partition(partition: SourcePartition) -> torch.
     """
     call_function_nodes = [node for node in partition.nodes if node.op == "call_function"]
     if len(call_function_nodes) != 1:
-        # torch.export's insert_deferred_runtime_asserts synthesizes one SymInt mul per
-        # shape-runtime assertion, all sharing one torch_fn tag, so several can collapse
-        # into a single partition. They carry no tensor value to annotate, so picking any
-        # one of them is safe here; downstream floating-point filtering no-ops on SymInt.
-        if call_function_nodes and all(
-            isinstance(node.meta.get("val"), torch.SymInt) for node in call_function_nodes
-        ):
-            return call_function_nodes[0]
-
         module_names = {
             name
             for node in call_function_nodes
@@ -733,18 +725,89 @@ def _get_call_function_node_from_partition(partition: SourcePartition) -> torch.
     return call_function_nodes[0]
 
 
+def _single_op_partition(node: Node, op_type: str) -> SourcePartition:
+    """Wrap one op node in a ``SourcePartition`` of its own."""
+    input_nodes = [arg for arg in node.all_input_nodes if arg.op != "get_attr"]
+    output_nodes = [node] if node.users else []
+    partition = SourcePartition(
+        nodes=[node], source=op_type, input_nodes=input_nodes, output_nodes=output_nodes
+    )
+    return partition
+
+
+def _find_sequential_partitions(
+    model: torch.fx.GraphModule, op_types: list[str]
+) -> list[tuple[SourcePartition, ...]]:
+    """Find every chain of ops in ``model`` whose op types are ``op_types``, in order.
+
+    Each chain is returned as a tuple with one ``SourcePartition`` per element of
+    ``op_types``, each holding exactly one call_function node.
+
+    This is a local replacement for torchao's ``find_sequential_partitions``: it returns
+    results in the same format, but does not merge unrelated ops into one partition.
+
+    torchao's ``find_sequential_partitions`` gets the ops from torch's
+    ``get_source_partitions``, which can put separate nodes into a single partition: it
+    builds partitions based on the first element of ``node.meta["torch_fn"]``, a name that
+    can be the same for two separate nodes. For example::
+
+        def forward(self, x1, x2):
+            s = x1 + x2           # add:   torch_fn ("add_1", "method_descriptor.add")
+            t = torch.add(s, x2)  # add_1: torch_fn ("add_1", "builtin_function_or_method.add")
+            out = torch.relu(t)   # relu:  only add_1 feeds it
+            return out
+
+    Both adds get the ``torch_fn`` name ``"add_1"``. Writing each partition as the list of
+    its nodes, torchao's ``find_sequential_partitions`` returns:
+
+    - for ``["add"]``: one chain, ``([add, add_1],)``: a single partition holding both adds.
+    - for ``["add", "relu"]``: one chain, ``([add, add_1], [relu])``.
+
+    This function returns:
+
+    - for ``["add"]``: two chains, ``([add],)`` and ``([add_1],)``.
+    - for ``["add", "relu"]``: one chain, ``([add_1], [relu])``.
+
+    So each add gets its own match, and the inputs of both adds are quantized.
+    """
+    chains = []
+    for node in model.graph.nodes:
+        if node.op == "call_function" and get_node_type(node, warn_on_failure=False) == op_types[0]:
+            chains.append([node])
+
+    for op_type in op_types[1:]:
+        extended_chains = []
+        for chain in chains:
+            for user in chain[-1].users:
+                if (
+                    user.op == "call_function"
+                    and get_node_type(user, warn_on_failure=False) == op_type
+                ):
+                    extended_chains.append([*chain, user])
+        chains = extended_chains
+
+    sequences = []
+    for chain in chains:
+        sequence = tuple(
+            _single_op_partition(node, op_type)
+            for node, op_type in zip(chain, op_types, strict=True)
+        )
+        sequences.append(sequence)
+    return sequences
+
+
 def match_pattern_with_sequential_partitions(
     model: torch.fx.GraphModule, pattern: OpsListPattern
-) -> dict[torch.fx.Node, tuple[SourcePartition]]:
+) -> dict[torch.fx.Node, tuple[SourcePartition, ...]]:
     """
     Helper function for matching model nodes with pattern using
-    find_sequential_partitions.
+    ``_find_sequential_partitions``.
 
     Returns a dictionary mapping nodes to matches.
     """
-    node_to_match_dict: dict[torch.fx.Node, tuple[SourcePartition]] = {}
+    node_to_match_dict: dict[torch.fx.Node, tuple[SourcePartition, ...]] = {}
 
-    partitions = find_sequential_partitions(model, pattern.pattern)
+    partitions = _find_sequential_partitions(model, pattern.pattern)
     # Each sequential partition is a tuple of source partitions, one per pattern element.
     for sequential_partition in partitions:
         if _is_branching_partition(sequential_partition):
@@ -757,7 +820,7 @@ def match_pattern_with_sequential_partitions(
     return node_to_match_dict
 
 
-def _is_branching_partition(sequential_partition: tuple[SourcePartition]) -> bool:
+def _is_branching_partition(sequential_partition: tuple[SourcePartition, ...]) -> bool:
     """
     Return True if the partition is a branching partition, False otherwise.
 

tests/quantization/test_annotation_pattern_registry.py

@@ -2088,6 +2088,30 @@ def test_cat_inputs_per_channel_along_other_axis(self):
         )
         assert token1_quantizer is token2_quantizer
 
+    def test_add_operator_and_add_function(self):
+        """
+        Test that both adds are quantized
+        """
+
+        class AddOperatorAndFunctionModel(nn.Module):
+            def forward(self, x1, x2):
+                out = x1 + x2
+                out = torch.add(out, x2)
+                return out
+
+        model = AddOperatorAndFunctionModel().eval()
+        quantizer = Quantizer(model, QuantizerConfig())
+        prepared_model = quantizer.prepare(example_inputs=(torch.randn(2, 4), torch.randn(2, 4)))
+
+        patterns = analyze_graph_structure(prepared_model)
+        check_fake_quantize_placement(
+            patterns,
+            {
+                "add": {"input_fq": [True, True], "output_fq": True},
+                "add_1": {"input_fq": [True, True], "output_fq": True},
+            },
+        )
+
     @pytest.mark.parametrize(
         "granularity",
         [

tests/quantization/test_annotation_utils.py

@@ -10,10 +10,14 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch._ops import OpOverloadPacket
-from torchao.quantization.pt2e import WrapperModule
+from torchao.quantization.pt2e import (
+    WrapperModule,
+    find_sequential_partitions as torchao_find_sequential_partitions,
+)
 
 from coreai_opt._utils.torch_utils import export_model
 from coreai_opt.quantization._graph._annotation_utils import (
+    _find_sequential_partitions,
     _get_aten_graph_module_for_pattern,
     match_pattern_with_subgraph_matcher,
 )
@@ -160,3 +164,69 @@ def _relu_pattern(input):
     model = export_model(nn.ReLU(), (torch.randn(2),), None, False)
     with pytest.raises(RuntimeError, match='"weight" entry'):
         match_pattern_with_subgraph_matcher(model, pattern_gm)
+
+
+class TestFindSequentialPartitions:
+    """Tests for ``_find_sequential_partitions``."""
+
+    class AddReluMul(nn.Module):
+        def forward(self, x1, x2):
+            out = torch.relu(x1 + x2) * x2
+            return out
+
+    class AddOperatorAndFunction(nn.Module):
+        """One add written as ``+`` and one as ``torch.add``: both get ``torch_fn`` "add_1"."""
+
+        def forward(self, x1, x2):
+            out = x1 + x2
+            out = torch.add(out, x2)
+            return out
+
+    @staticmethod
+    def _partition_node_names(sequences) -> list[list[list[str]]]:
+        """Return the node names of each partition, for each sequence."""
+        names = []
+        for sequence in sequences:
+            sequence_names = []
+            for partition in sequence:
+                sequence_names.append([node.name for node in partition.nodes])
+            names.append(sequence_names)
+        return names
+
+    @pytest.mark.parametrize(
+        "op_types",
+        [["add"], ["add", "relu"], ["relu", "mul"], ["mul"]],
+        ids=["add", "add_relu", "relu_mul", "mul"],
+    )
+    def test_matches_torchao(self, op_types):
+        """``_find_sequential_partitions`` mostly returns what torchao's version does.
+
+        They differ only when torch merges unrelated ops into one partition (see
+        ``test_torch_fn_name_collision``), and in torchao also matching equivalent types,
+        e.g. "add_" for "add", which the op-list patterns list explicitly instead.
+        """
+        model = export_model(
+            self.AddReluMul().eval(), (torch.randn(2, 4), torch.randn(2, 4)), None, False
+        )
+
+        expected = self._partition_node_names(torchao_find_sequential_partitions(model, op_types))
+        actual = self._partition_node_names(_find_sequential_partitions(model, op_types))
+
+        assert len(actual) == 1
+        assert actual == expected
+
+    def test_torch_fn_name_collision(self):
+        """Each op gets its own partition, even when two ops share a ``torch_fn`` name."""
+        model = export_model(
+            self.AddOperatorAndFunction().eval(),
+            (torch.randn(2, 4), torch.randn(2, 4)),
+            None,
+            False,
+        )
+
+        sequences = _find_sequential_partitions(model, ["add"])
+
+        # torchao's find_sequential_partitions returns [[["add", "add_1"]]] here: one
+        # sequence whose single partition holds both adds, since torch's
+        # get_source_partitions groups nodes by torch_fn name.
+        assert self._partition_node_names(sequences) == [[["add"]], [["add_1"]]]

tests/quantization/test_graph_mode_quantizer.py

@@ -555,9 +555,9 @@ def test_prepare_with_symint_mul_partition_collision(self):
 
         torch.export's ``insert_deferred_runtime_asserts`` synthesizes one SymInt
         ``mul`` per shape-runtime assertion, all sharing one ``torch_fn`` tag, so
-        they collapse into a single ``SourcePartition``. The annotator picks the
-        first node and the downstream ``_is_fx_node_floating_point`` filter no-ops
-        on SymInt inputs.
+        they collapse into a single ``SourcePartition``. The annotator splits it
+        into one partition per node, and the downstream ``_is_fx_node_floating_point``
+        filter no-ops on the SymInt inputs.
         """
 
         H, W, B, embed_dim = 4, 4, 1, 8