Constraint-queue reconciler for quantization annotation

src/coreai_opt/quantization/_graph/_annotation_config.py

@@ -5,10 +5,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping, Set
+from collections.abc import Mapping
 from dataclasses import dataclass
 
-import torch.fx
 from torchao.quantization.pt2e.quantizer import (
     QuantizationSpec as TorchAOQuantizationSpec,
 )
@@ -20,27 +19,17 @@
 
 @dataclass(frozen=True)
 class AnnotationContext:
-    """Pass-invariant inputs an annotator may need.
-
-    Held constant across all matches in a single annotation pass. Constructed
-    once when ``_AnnotationHandler.annotate`` begins and shared by every
-    annotator invocation during that pass.
-
-    Distinct from :class:`AnnotationConfig`, which carries per-op specs that
-    vary per match.
+    """Pass-invariant inputs for legacy kv-cache annotation overrides.
 
     Attributes:
         module_name_to_state_names_map (Mapping[str, Mapping[str, list[str]]]):
             For each module name, a mapping from each state target (FQN) to the
             list of local names the module uses for that state. Used during
             state-input annotation to translate a state node's target into the
             consumer module's local name(s).
-        shared_observer_nodes (Set[torch.fx.Node]): Nodes whose output annotations
-            are shared with their input annotations if any.
     """
 
     module_name_to_state_names_map: Mapping[str, Mapping[str, list[str]]]
-    shared_observer_nodes: Set[torch.fx.Node]
 
 
 class AnnotationConfig:

src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py

@@ -4,10 +4,9 @@
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, Generic, TypeAlias, TypeVar
+from typing import Generic, TypeVar
 
 import torch
 from torch.fx.passes.utils.matcher_utils import InternalMatch
@@ -16,42 +15,37 @@
 from coreai_opt._utils.registry_utils import ClassRegistryMixin
 
 from . import _annotation_utils
-from ._annotation_config import (
-    AnnotationConfig as _AnnotationConfig,
-    AnnotationContext as _AnnotationContext,
-)
 from ._annotation_utils import (
     OpsListPattern as _OpsListPattern,
 )
+from ._qspec_constraints import (
+    Constraint,
+    PER_CHANNEL_SCHEMES,
+    ShareFields,
+    ShareObserverInstance,
+)
+from ._qspec_types import (
+    FieldName,
+    NodeSlot,
+    ProvisionalQSpec,
+    ProvisionalQSpecMap,
+    SlotKind,
+)
 
 # Generic type variable for match results
 MatchType = TypeVar("MatchType")
 
-# Generic annotator function type.
-# The function is expected to take exactly 3 inputs:
-# 1. Matched nodes to annotate. The type of this entity is flexible depending
-#    on the implementation of the AnnotationPattern subclass. Whatever entity
-#    is returned in the subclass's match_single_pattern dictionary values will
-#    be passed into this function as the first input.
-# 2. Quantization Config to use when annotating the matched nodes (per-match,
-#    derived from OpQuantizerConfig).
-# 3. Annotation pass context. Holds pass-invariant inputs the annotator may
-#    need (the model's module-name-to-state-names map and the set of shared
-#    observer nodes computed at the start of this annotation pass).
-AnnotatorFunc: TypeAlias = Callable[[MatchType, _AnnotationConfig, _AnnotationContext], Any]
-
 
 @dataclass(frozen=True)
 class AnnotatorMatchInfo(Generic[MatchType]):
-    """
-    Holds info related to nodes matched with a particular annotator.
+    """Info about a single annotator's match on the graph.
 
-    annotator_func: A function used to annotate nodes in annotator_match
-    annotator_match: Nodes in the model matched with the annotator
-    pattern_length: Length of the annotation pattern
+    Attributes:
+        annotator_match (MatchType): Nodes in the model matched with the
+            annotator.
+        pattern_length (int): Length of the annotation pattern.
     """
 
-    annotator_func: AnnotatorFunc[MatchType]
     annotator_match: MatchType
     pattern_length: int
 
@@ -88,8 +82,6 @@ class BaseAnnotationPattern(Generic[MatchType], ABC):
     Base class for annotation patterns.
 
     Each pattern class should implement:
-    - get_annotator_func(): Returns a function used to annotate nodes in accordance with
-                            the Annotation pattern.
     - generate_patterns(): Returns list of graph modules representing
                            the patterns to match
     - match_single_pattern(): Matches graph for a single pattern
@@ -98,18 +90,6 @@ class BaseAnnotationPattern(Generic[MatchType], ABC):
     # Class-level cache for patterns
     _patterns: list[torch.fx.GraphModule | _OpsListPattern] | None = None
 
-    @classmethod
-    @abstractmethod
-    def get_annotator_func(cls) -> AnnotatorFunc[MatchType]:
-        """
-        Return a function which is used to annotate nodes in accordance with a
-        particular Annotation pattern.
-
-        Returns:
-            Function used to annotate nodes.
-        """
-        raise NotImplementedError
-
     @classmethod
     def get_patterns(cls) -> list[torch.fx.GraphModule | _OpsListPattern]:
         """
@@ -239,7 +219,6 @@ def _match_all_patterns(
                 {
                     node: AnnotatorMatchInfo(
                         pattern_length=cls.get_pattern_length(),
-                        annotator_func=cls.get_annotator_func(),
                         annotator_match=match,
                     )
                     for node, match in node_to_match_dict.items()
@@ -264,14 +243,6 @@ class WeightedModulePattern(BaseAnnotationPattern[InternalMatch]):
     be present in the pattern.
     """
 
-    @classmethod
-    def get_annotator_func(cls) -> AnnotatorFunc[InternalMatch]:
-        """
-        Return a function which is used to annotate nodes in accordance with
-        WeightedModulePattern.
-        """
-        return _annotation_utils.annotate_weighted_mod_match
-
     @classmethod
     def match_single_pattern(
         cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | _OpsListPattern
@@ -310,14 +281,6 @@ class NAryActPattern(BaseAnnotationPattern[tuple[SourcePartition]]):
     Add, Add+activation, etc.).
     """
 
-    @classmethod
-    def get_annotator_func(cls) -> AnnotatorFunc[tuple[SourcePartition]]:
-        """
-        Return a function which is used to annotate nodes in accordance with
-        NAryActPattern.
-        """
-        return _annotation_utils.annotate_n_ary_act_match
-
     @classmethod
     def match_single_pattern(
         cls, model: torch.fx.GraphModule, pattern: torch.fx.GraphModule | _OpsListPattern
@@ -349,15 +312,39 @@ class SharedObserverModulePattern(BaseAnnotationPattern[tuple[SourcePartition]])
 
     Subclasses must implement the generate_patterns() method to define the
     specific patterns they want to match (e.g., maxpool, avgpool, flatten, etc.).
+
+    Subclasses must also implement :meth:`generate_qspec_sharing_constraints`
+    to declare how quantization specs are shared across the op's input and
+    output slots. Missing this override raises ``TypeError`` at instantiation.
     """
 
     @classmethod
-    def get_annotator_func(cls) -> AnnotatorFunc[tuple[SourcePartition]]:
-        """
-        Return a function which is used to annotate nodes in accordance with
-        SharedObserverModulePattern.
+    @abstractmethod
+    def generate_qspec_sharing_constraints(
+        cls,
+        node: torch.fx.Node,
+        qspecs: ProvisionalQSpecMap,
+    ) -> list[Constraint]:
+        """Return the constraints that express this op's qspec-sharing semantics.
+
+        Called by the reconciler once per matched shared-observer node with
+        the current provisional-qspec state. Each subclass owns the decision
+        of which of its input / output slots must share dtype, is_dynamic,
+        or a full observer instance.
+
+        Args:
+            node (torch.fx.Node): The matched op node whose slots the
+                returned constraints reference.
+            qspecs (ProvisionalQSpecMap): Current provisional-qspec state,
+                available for dynamic peeking (e.g. concat's axis-aware
+                observer-sharing decision reads the current output-slot
+                qscheme and ch_axis).
+
+        Returns:
+            list[Constraint]: Constraints to enqueue. Empty list means
+            "no sharing this op wants to enforce."
         """
-        return _annotation_utils.annotate_shared_observer_match
+        raise NotImplementedError
 
     @classmethod
     def _validate_patterns_length(cls):
@@ -388,6 +375,32 @@ def match_single_pattern(
         return _annotation_utils.match_pattern_with_sequential_partitions(model, pattern)
 
 
+def _shared_op_boundary_slots(node: torch.fx.Node) -> tuple[NodeSlot, list[NodeSlot]]:
+    """Return the output slot and list of input slots on a shared-observer node.
+
+    Utility for :class:`SharedObserverModulePattern` subclass implementations
+    of :meth:`generate_qspec_sharing_constraints`.
+    """
+    output_slot = NodeSlot(node=node, kind=SlotKind.OUTPUT, arg_index=0)
+    input_slots = [
+        NodeSlot(node=node, kind=SlotKind.INPUT, arg_index=arg_index)
+        for arg_index in range(len(node.all_input_nodes))
+    ]
+    return output_slot, input_slots
+
+
+def _any_input_slot_populated(
+    input_slots: list[NodeSlot], qspecs: ProvisionalQSpecMap
+) -> bool:
+    """Return True iff any of the input slots already has a proposal in ``qspecs``.
+
+    The shared-observer sharing rule for cat / maxpool / etc. only fires when
+    at least one input side wants observation — otherwise there's nothing
+    for the sharing to enforce.
+    """
+    return any(input_slot in qspecs for input_slot in input_slots)
+
+
 class _AnnotationPatternRegistry(ClassRegistryMixin):
     """
     A registry of quantization annotation pattern classes.
@@ -787,6 +800,13 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         """Returns flatten pattern."""
         return _get_all_patterns_from_base_ops({"flatten"})
 
+    @classmethod
+    def generate_qspec_sharing_constraints(
+        cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap
+    ) -> list[Constraint]:
+        """All slots share one observer instance — flatten preserves values."""
+        return _all_slots_share_observer(node, qspecs)
+
 
 @_AnnotationPatternRegistry.register("maxpool")
 class MaxPoolPattern(SharedObserverModulePattern):
@@ -805,6 +825,13 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
             }
         )
 
+    @classmethod
+    def generate_qspec_sharing_constraints(
+        cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap
+    ) -> list[Constraint]:
+        """All slots share one observer instance — maxpool preserves values."""
+        return _all_slots_share_observer(node, qspecs)
+
 
 @_AnnotationPatternRegistry.register("avgpool")
 class AvgPoolPattern(SharedObserverModulePattern):
@@ -827,6 +854,13 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]:
             }
         )
 
+    @classmethod
+    def generate_qspec_sharing_constraints(
+        cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap
+    ) -> list[Constraint]:
+        """All slots share one observer instance — averaging preserves range roughly."""
+        return _all_slots_share_observer(node, qspecs)
+
 
 @_AnnotationPatternRegistry.register("concat")
 class ConcatPattern(SharedObserverModulePattern):
@@ -838,3 +872,65 @@ class ConcatPattern(SharedObserverModulePattern):
     def generate_patterns(cls) -> list[torch.fx.GraphModule]:
         """Returns concat pattern."""
         return _get_all_patterns_from_base_ops({"cat", "concat"})
+
+    @classmethod
+    def generate_qspec_sharing_constraints(
+        cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap
+    ) -> list[Constraint]:
+        """Axis-aware sharing:
+
+        * All slots always share ``DTYPE`` and ``IS_DYNAMIC``.
+        * If the reconciled output qscheme is per-channel along the concat
+          dimension, each input keeps its own scale — no observer-instance
+          sharing needed. Per-channel along a *different* axis (or per-tensor)
+          requires all inputs to observe together.
+        """
+        output_slot, input_slots = _shared_op_boundary_slots(node)
+        if not _any_input_slot_populated(input_slots, qspecs):
+            return []
+
+        all_slots = frozenset([output_slot, *input_slots])
+        constraints: list[Constraint] = [
+            ShareFields(
+                _slots=all_slots,
+                fields=frozenset({FieldName.DTYPE, FieldName.IS_DYNAMIC}),
+            )
+        ]
+
+        concat_dim = _concat_dim(node)
+        output_qspec = qspecs.get(output_slot, ProvisionalQSpec()).fields
+        output_qscheme = output_qspec.get(FieldName.QSCHEME)
+        output_ch_axis = output_qspec.get(FieldName.CH_AXIS)
+
+        per_channel_along_concat_axis = (
+            output_qscheme is not None
+            and output_qscheme.value in PER_CHANNEL_SCHEMES
+            and output_ch_axis is not None
+            and output_ch_axis.value == concat_dim
+        )
+        if not per_channel_along_concat_axis:
+            constraints.append(ShareObserverInstance(_slots=all_slots))
+        return constraints
+
+
+def _all_slots_share_observer(
+    node: torch.fx.Node, qspecs: ProvisionalQSpecMap
+) -> list[Constraint]:
+    """Emit ``ShareObserverInstance`` across the op's input + output slots.
+
+    Used by shared-observer ops (flatten, maxpool, avgpool) whose sharing
+    semantics don't depend on any spec field — they always tie every input
+    and output slot into one observer instance when at least one input
+    side wants observation.
+    """
+    output_slot, input_slots = _shared_op_boundary_slots(node)
+    if not _any_input_slot_populated(input_slots, qspecs):
+        return []
+    return [ShareObserverInstance(_slots=frozenset([output_slot, *input_slots]))]
+
+
+def _concat_dim(node: torch.fx.Node) -> int:
+    """Concat dim from ``aten.cat(tensors, dim=?)`` (defaults to 0)."""
+    if len(node.args) >= 2:
+        return int(node.args[1])
+    return int(node.kwargs.get("dim", 0))