diff --git a/src/coreai_opt/quantization/_graph/_annotation_config.py b/src/coreai_opt/quantization/_graph/_annotation_config.py index bec1c2d..2781a35 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_config.py +++ b/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,14 +19,7 @@ @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]]]): @@ -35,12 +27,9 @@ class AnnotationContext: 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: diff --git a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py index e6a1338..fb1acbf 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py +++ b/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)) diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index fe304ff..28885bb 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -4,7 +4,7 @@ # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause import logging -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Any, Literal @@ -182,129 +182,6 @@ def is_node_annotated(node: Node) -> bool: return node and Q_ANNOTATION_KEY in node.meta and node.meta[Q_ANNOTATION_KEY]._annotated -def is_any_annotated(nodes: list[Node]) -> bool: - """ - Given a list of nodes (that represents an operator pattern), - check if any of the node is annotated, return True if any of the node - is annotated, otherwise return False. - """ - return any(is_node_annotated(node) for node in nodes) - - -def is_all_annotated(nodes: list[Node]) -> bool: - """ - Given a list of nodes (that represents an operator pattern), - return True if all of the node is annotated, otherwise return False. - """ - return all(is_node_annotated(node) for node in nodes) - - -def mark_nodes_as_annotated(nodes: Iterable[Node]) -> None: - for node in nodes: - if node is not None: - if Q_ANNOTATION_KEY not in node.meta: - node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation() - node.meta[Q_ANNOTATION_KEY]._annotated = True - - -def _propagate_qscheme_to_child_nodes( - root_node: torch.fx.Node, - qscheme: torch.qscheme, - shared_observer_nodes: set[torch.fx.Node] | None = None, -) -> None: - """ - Given a qscheme, propagate the qscheme to all applicable children. Any input qspecs - which are not shared qspecs will have qschemes updated. The propagation logic - continues downwards through the graph until we encounter a non-shared observer op. - """ - nodes_to_propagate = [(root_node, user) for user in root_node.users.keys()] - while nodes_to_propagate: - parent, curr_node = nodes_to_propagate.pop(0) - if not is_node_annotated(curr_node): - continue - curr_input_qspec = curr_node.meta[Q_ANNOTATION_KEY].input_qspec_map.get(parent) - if curr_input_qspec is None: - continue - if ( - isinstance(curr_input_qspec, _SharedQuantizationSpec) - and root_node not in curr_input_qspec.edge_or_node - ): - # This is a case in which we encounter a multi input op which already has a - # shared input qspec referencing a different edge. Here, we will not update - # the qscheme since it would unnecessarily constrain the range of the other - # edge. Here, there is the possiblity of having back to back quantize - # dequantize ops inserted. - continue - if not isinstance(curr_input_qspec, _SharedQuantizationSpec): - adjusted_qspec = TorchAOQuantizationSpec( - observer_or_fake_quant_ctr=curr_input_qspec.observer_or_fake_quant_ctr, - dtype=curr_input_qspec.dtype, - qscheme=qscheme, - quant_min=curr_input_qspec.quant_min, - quant_max=curr_input_qspec.quant_max, - ) - curr_node.meta[Q_ANNOTATION_KEY].input_qspec_map[parent] = adjusted_qspec - - if curr_node in shared_observer_nodes: - nodes_to_propagate.extend([(curr_node, user) for user in curr_node.users.keys()]) - - -def adjust_output_qspec_for_qscheme_and_propagate( - node: torch.fx.Node, shared_observer_nodes: set[torch.fx.Node] -) -> None: - """ - Adjust output quantization spec for ops which can use fixed qparams - or ops for which we can use affine quantization mode during - symmetric quantization because their output is always positive. - Propagate the updated qscheme to input qspecs of child ops. - - Args: - node: The node whose output qspec should be updated if necessary - shared_observer_nodes: Set of shared observer nodes for which adjusted qschemes - should propagate through - """ - if not is_node_annotated(node): - return - - qspec = node.meta[Q_ANNOTATION_KEY].output_qspec - if qspec is None: - return - - # ReLU6 activation maps to torch.ops.aten.hardtanh.default with - # min_val = 0 and max_val = 6 - is_always_affine_op = node.target in _always_affine_ops or ( - node.target in [torch.ops.aten.hardtanh.default, torch.ops.aten.hardtanh_.default] - and node.args[1] == 0 # min_val, corresponding to ReLU6 - and node.args[2] == 6 # max_val, corresponding to ReLU6 - ) - - adjusted_qspec = None - if node.target in _fixed_q_params_ops: - adjusted_qspec = TorchAOQuantizationSpec( - observer_or_fake_quant_ctr=qspec.observer_or_fake_quant_ctr, - dtype=qspec.dtype, - qscheme=_fixed_q_params_ops[node.target].qscheme, - quant_min=qspec.quant_min, - quant_max=qspec.quant_max, - ) - # FIXME: Because of a bug in PyTorch in function _create_obs_or_fq_from_qspec - # in module torch/ao/quantization/fx/prepare.py which creates a - # FixedQParamsFakeQuantize partial, instead of an instance, we cannot - # actually create FixedQParamsQuantizationSpec - elif is_always_affine_op: - adjusted_qspec = TorchAOQuantizationSpec( - observer_or_fake_quant_ctr=qspec.observer_or_fake_quant_ctr, - dtype=qspec.dtype, - qscheme=torch.per_tensor_affine, - quant_min=qspec.quant_min, - quant_max=qspec.quant_max, - ) - - if adjusted_qspec is not None: - node.meta[Q_ANNOTATION_KEY].output_qspec = adjusted_qspec - _propagate_qscheme_to_child_nodes(node, adjusted_qspec.qscheme, shared_observer_nodes) - - def _get_weighted_mod_pattern( mod_fn: Callable, example_inputs: tuple[torch.Tensor, ...], @@ -603,20 +480,6 @@ def _embedding(input, weight): return _get_aten_graph_module_for_pattern(WrapperModule(_embedding), example_inputs) -def find_consumers( - pattern_nodes: Iterable[torch.fx.Node], target_nodes: list[torch.fx.Node] -) -> set[torch.fx.Node]: - """Find all nodes that take as input any of the target nodes""" - consumers = [] - for node in pattern_nodes: - if node is None: - continue - for arg in node.all_input_nodes: - if arg in target_nodes: - consumers.append(node) - return set(consumers) - - def _is_fx_node_floating_point(node: torch.fx.Node) -> bool: """ Check if a fx node has floating point data. @@ -811,108 +674,6 @@ def _fill_input_qspec_map_for_input( input_qspec_map[input_node] = input_node.meta[Q_ANNOTATION_KEY].output_qspec -def _get_output_qspec( - node: torch.fx.Node, - quantization_config: AnnotationConfig, - shared_observer_nodes: set[torch.fx.Node] | None = None, -) -> TorchAOQuantizationSpec | None: - """ - Get the output qspec which should be associated with the node. Check child nodes - first to see if any input qspec is already set, and if so, reuse the qspec. - If there are multiple child nodes, the qspec will be the first valid one found when - checking all child nodes. - When encountering shared observer nodes as child nodes, continue checking their - children if no valid input qspec is found. - """ - op_output_spec = quantization_config.op_output_spec - - # Early exit if input node is not floating point dtype - if not _is_fx_node_floating_point(node): - # Hardcoding for index 0 for now while we only support single output setting - if 0 in op_output_spec: - _warn_non_quantizable_tensor_setting(node, "output", 0, op_output_spec) - return None - - # First read qspec from config without applying it yet. - qspec_from_config, _ = get_last_matching_spec([0], op_output_spec) - - # Don't set output qspec if it is specified to be None. If the op has multiple child - # ops where a subset of child ops don't have input quantization, we should not - # insert a quantizer for the op's output. - if qspec_from_config is None: - return None - - # Check child ops to see if qspec settings should be taken from them. - # Current logic simply uses the first child qspec found as the qspec to use. - # This logic might need to be updated to be smarter in how the qspec to use is - # determined. - if not shared_observer_nodes: - shared_observer_nodes = set() - nodes_to_check = [user for user in node.users] - while nodes_to_check: - curr_node = nodes_to_check.pop(0) - if ( - is_node_annotated(curr_node) - and curr_node.meta[Q_ANNOTATION_KEY].input_qspec_map.get(node) is not None - ): - return curr_node.meta[Q_ANNOTATION_KEY].input_qspec_map.get(node) - # Check children of shared observer nodes as well in a DFS fashion - if curr_node in shared_observer_nodes: - nodes_to_check[:0] = curr_node.users - - # If no valid qspec was found, refer to quantization config - return qspec_from_config - - -def _propagate_output_qspec( - node: torch.fx.Node, - output_qspec: TorchAOQuantizationSpec, - shared_observer_nodes: set[torch.fx.Node] | None = None, -): - """ - Propagate output qspec to child ops which are shared observer nodes. These node - types should simply echo the qspec which is coming from its input. - Propagated specs will be SharedQuantizationSpecs except for the first input qspec - that is set, since all subsequent SharedQuantizationSpecs will be referring to the - first input qspec set. - This function should be called for each annotation function which is defined. - """ - if output_qspec is None or not shared_observer_nodes or not node.users: - return - if not shared_observer_nodes: - shared_observer_nodes = set() - spec_to_propagate = output_qspec - nodes_to_propagate = [user for user in node.users if user in shared_observer_nodes] - - while nodes_to_propagate: - user = nodes_to_propagate.pop(0) - # Skip propagation for a node which is already annotated - if is_any_annotated([user]): - continue - input_qspec_map = {} - if isinstance(spec_to_propagate, _SharedQuantizationSpec): - # If spec to propagate is already a shared qspec, we can simply update all - # shared observer inputs and outputs with the shared qspec - for input_node in user.all_input_nodes: - input_qspec_map[input_node] = spec_to_propagate - else: - # If spec is not a shared qspec, this is the very first qspec we are - # setting. Set the first input to be this qspec, then update - # spec_to_propagate to be a shared qspec to be used for all other updates - first_user_input = user.all_input_nodes[0] - input_qspec_map[first_user_input] = spec_to_propagate - spec_to_propagate = _SharedQuantizationSpec((first_user_input, user)) - for input_node in user.all_input_nodes[1:]: - input_qspec_map[input_node] = spec_to_propagate - - user.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, output_qspec=spec_to_propagate - ) - mark_nodes_as_annotated([user]) - - nodes_to_propagate.extend([child for child in user.users if child in shared_observer_nodes]) - - def _get_call_function_node_from_partition(partition: SourcePartition) -> torch.fx.Node: """Return the first call_function node in the partition.""" return [node for node in partition.nodes if node.op == "call_function"][0] @@ -980,179 +741,6 @@ def match_pattern_with_subgraph_matcher( return node_to_match_dict -def annotate_weighted_mod_match( - annotator_match: InternalMatch, - quantization_config: AnnotationConfig, - context: AnnotationContext, -) -> None: - """ - Try to annotate specific nodes in the model designated by ``annotator_match`` using - ``quantization_config``. - """ - # Entries in name_node_map will be determined by what is populated in - # node_dict when the pattern was created. - name_node_map = annotator_match.name_node_map - mod_node = name_node_map["mod"] - bn_node = name_node_map.get("bn") - output_node = None - if "output" in name_node_map: - # In this case, an activation is applied to the weighted module output - output_node = name_node_map["output"] - # If the output is same as bn_node or mod_node, it means we have an inplace - # activation, so we need to correct the node. - if bn_node is not None and bn_node == output_node: - bn_node = bn_node.args[0] - elif mod_node == output_node: - mod_node = mod_node.args[0] - - # TODO: skip partition if any intermediate node output is used by an op outside the pattern. - - # Skip partition if already annotated - partition = [mod_node] - partition.extend(filter(None, [bn_node, output_node])) - - if is_any_annotated(partition): - return - - shared_observer_nodes = context.shared_observer_nodes - input_qspec_map = _get_input_qspec_map( - mod_node.all_input_nodes, - quantization_config, - context, - ) - output_qspec = _get_output_qspec( - output_node or mod_node, quantization_config, shared_observer_nodes - ) - # set mod_node - mod_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=None if output_node else output_qspec, - ) - # set output_node if exists - if output_node: - output_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(output_qspec=output_qspec) - - # apply output spec to either final output_node or mod_node - _propagate_output_qspec(output_node or mod_node, output_qspec, shared_observer_nodes) - - # Mark all nodes in pattern as annotated - mark_nodes_as_annotated(partition) - - -def annotate_n_ary_act_match( - annotator_match: tuple[SourcePartition], - quantization_config: AnnotationConfig, - context: AnnotationContext, -) -> None: - """ - Try to annotate specific nodes in the model designated by ``annotator_match`` using - ``quantization_config``. - """ - if len(annotator_match) > 2: - error_msg = ( - "Sequential list of ops is longer than 2. Only lists of up to " - "length 2 (op + [act]_ are supported." - ) - raise RuntimeError(error_msg) - nodes_to_annotate = [ - _get_call_function_node_from_partition(partition) for partition in annotator_match - ] - first_op_node = nodes_to_annotate[0] - last_op_node = nodes_to_annotate[-1] - - if is_any_annotated(nodes_to_annotate): - return - - # TODO: skip partition if any intermediate node output is used by an op outside the pattern. - - shared_observer_nodes = context.shared_observer_nodes - input_qspec_map = _get_input_qspec_map( - first_op_node.all_input_nodes, - quantization_config, - context, - ) - output_qspec = _get_output_qspec(last_op_node, quantization_config, shared_observer_nodes) - if len(nodes_to_annotate) == 1: - first_op_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=output_qspec, - _annotated=True, - ) - else: - first_op_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - _annotated=True, - ) - last_op_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - output_qspec=output_qspec, - _annotated=True, - ) - _propagate_output_qspec(last_op_node, output_qspec, shared_observer_nodes) - - -def _adjust_input_qspec_map_for_shared_observers( - op_node: Node, input_qspec_map: dict[Node, TorchAOQuantizationSpec] -) -> _SharedQuantizationSpec | None: - """ - For shared observer ops, check if any inputs have qspecs. If so, - 1. set the first input to the first valid qspec - 2. set all subsequent inputs to use SharedObserverQspec tied to the first input - - This logic is mainly needed for multi-input shared observer ops (concat). - A SharedObserverQspec is returned for use as output_qspec. - """ - shared_qspec = None - for input_node in op_node.all_input_nodes: - if input_qspec_map[input_node] is not None: - input_qspec_map[op_node.all_input_nodes[0]] = input_qspec_map[input_node] - if input_qspec_map[op_node.all_input_nodes[0]] is not None: - shared_qspec = _SharedQuantizationSpec((op_node.all_input_nodes[0], op_node)) - for input_node in op_node.all_input_nodes[1:]: - input_qspec_map[input_node] = shared_qspec - return shared_qspec - - -def annotate_shared_observer_match( - annotator_match: tuple[SourcePartition], - quantization_config: AnnotationConfig, - context: AnnotationContext, -) -> None: - """ - Try to annotate specific nodes in the model designated by ``annotator_match`` using - ``quantization_config``. - """ - if len(annotator_match) > 1: - error_msg = ( - "Shared observer pattern is expected to be length 1, but got " - f"{len(annotator_match)}. Annotator match: {annotator_match}." - ) - raise RuntimeError(error_msg) - - op_node = _get_call_function_node_from_partition(annotator_match[0]) - if is_node_annotated(op_node): - return - - shared_observer_nodes = context.shared_observer_nodes - input_qspec_map = _get_input_qspec_map( - op_node.all_input_nodes, - quantization_config, - context, - ) - output_qspec = _adjust_input_qspec_map_for_shared_observers(op_node, input_qspec_map) - - if output_qspec is None: - # Only use a different output qspec if it isn't sharing with its input - output_qspec = _get_output_qspec(op_node, quantization_config, shared_observer_nodes) - - # input and output of op will share quantization parameter with input of op - op_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=output_qspec, - _annotated=True, - ) - _propagate_output_qspec(op_node, output_qspec, shared_observer_nodes) - - def annotate_module_level_specs( module_configs: ModuleConfigDict, module_name_to_state_names_map: Mapping[str, Mapping[str, list[str]]], diff --git a/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py new file mode 100644 index 0000000..709253d --- /dev/null +++ b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py @@ -0,0 +1,380 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Phase 2 of the qspec reconciliation pipeline: build the initial +:class:`ProvisionalQSpecMap` from winning configs and op-intrinsic +overrides. + +For each pattern-covered node the pipeline seeds field values on: + + * Each INPUT slot whose producer is outside the node's pattern + group (state producers get ``op_state_spec``; everything else + gets ``op_input_spec``). + * The OUTPUT slot, unless every consumer is inside the pattern + group (internal edges get no annotation — see below). + +For nodes whose target has op-intrinsic quantization semantics +(sigmoid/tanh/hardsigmoid/relu/relu6), the intrinsic overrides the +user's proposal on the corresponding fields. Every override that +actually changes a user-specified value is logged. + +Two pattern-shape assumptions are checked and raised on violation: + + * An op-intrinsic node whose OUTPUT slot would be skipped (all its + consumers live inside its pattern group) means the intrinsic is + about to be silently dropped. Raise instead — a future pattern + author must acknowledge the conflict. + * A pattern-covered node with a mix of in-pattern and external + consumers can't have the pattern's "no observer on internal + edges" contract honored (the observer torchao inserts at the + producer's OUTPUT is visible to both consumer groups). Raise. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import torch +import torch.fx as fx +from torchao.quantization.pt2e.fake_quantize import FixedQParamsFakeQuantize +from torchao.quantization.pt2e.quantizer import QuantizationSpec as TorchAOQuantizationSpec + +from coreai_opt._utils.config_utils import ALL_TENSORS as _ALL_TENSORS +from coreai_opt._utils.fx_utils import ( + get_local_state_name as _get_local_state_name, + is_coreai_compressed_state_node as _is_state_node, +) +from coreai_opt.quantization.config import OpQuantizerConfig + +from ._annotation_config import AnnotationConfig +from ._annotation_utils import ( + _always_affine_ops, + _fixed_q_params_ops, +) +from ._qspec_constraints import _get_or_create +from ._qspec_types import ( + FieldName, + FieldValue, + NodeSlot, + ProvisionalQSpecMap, + ReconciliationError, + SlotKind, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Entry point. +# --------------------------------------------------------------------------- + + +def build_initial_provisional_qspecs( + model: fx.GraphModule, + winning_configs: dict[fx.Node, OpQuantizerConfig], + node_priorities: dict[fx.Node, int], + pattern_groups: dict[fx.Node, frozenset[fx.Node]], +) -> ProvisionalQSpecMap: + """Build the initial :class:`ProvisionalQSpecMap` from + winning-config + intrinsic contributions. + """ + node_to_annotation_config: dict[fx.Node, AnnotationConfig] = { + node: AnnotationConfig.from_quantizer_config(cfg) + for node, cfg in winning_configs.items() + } + + qspecs: ProvisionalQSpecMap = {} + for node, cfg in node_to_annotation_config.items(): + priority = node_priorities[node] + pattern = pattern_groups.get(node, frozenset({node})) + _populate_input_slots(node, cfg, priority, pattern, qspecs) + _populate_output_slot(node, cfg, priority, pattern, qspecs) + + return qspecs + + +# --------------------------------------------------------------------------- +# Input-slot population. +# --------------------------------------------------------------------------- + + +def _populate_input_slots( + node: fx.Node, + cfg: AnnotationConfig, + priority: int, + pattern: frozenset[fx.Node], + qspecs: ProvisionalQSpecMap, +) -> None: + """Populate every non-internal INPUT slot on ``node`` from its config.""" + for arg_index, producer in enumerate(node.all_input_nodes): + if producer in pattern: + continue # internal edge — skip + slot = NodeSlot(node=node, kind=SlotKind.INPUT, arg_index=arg_index) + if _is_state_node(producer): + _validate_state_not_referenced_via_input_spec(producer, arg_index, cfg) + spec = _lookup_state_spec(cfg, producer) + else: + spec = _lookup_by_key(cfg.op_input_spec, arg_index) + if spec is None: + continue + _populate_fields_from_spec(qspecs, slot, spec, priority) + + +def _validate_state_not_referenced_via_input_spec( + state_producer: fx.Node, arg_index: int, cfg: AnnotationConfig +) -> None: + """Raise if the user's ``op_input_spec`` targets an arg index whose + producer is a state tensor. State tensors must be configured via + ``op_state_spec``. + """ + if arg_index in cfg.op_input_spec: + raise RuntimeError( + f"Config is attempting to set op_input_spec idx {arg_index}, " + f"but the input is a state tensor (node: {state_producer.name}). " + f"Use op_state_spec to configure state inputs instead.\n" + f"op_input_spec: {cfg.op_input_spec}" + ) + + +def _lookup_state_spec( + consumer_cfg: AnnotationConfig, state_node: fx.Node +) -> TorchAOQuantizationSpec | None: + """Resolve a state consumer's spec from its ``op_state_spec``.""" + state_name = _get_local_state_name(state_node) + if state_name is None: + return None + return _lookup_by_key(consumer_cfg.op_state_spec, state_name) + + +# --------------------------------------------------------------------------- +# Output-slot population + op-intrinsic override. +# --------------------------------------------------------------------------- + + +def _populate_output_slot( + node: fx.Node, + cfg: AnnotationConfig, + priority: int, + pattern: frozenset[fx.Node], + qspecs: ProvisionalQSpecMap, +) -> None: + """Populate ``node``'s OUTPUT slot from its config, then apply any + op-intrinsic override. + + Enforces two pattern-shape invariants: + + * Op-intrinsic nodes must not be fully internal to a pattern (else + their intrinsic contribution is silently dropped — an unhandled + case for any future pattern that places an intrinsic op mid-chain). + * A covered node must have its consumers either all inside the + pattern or all outside (any mix means the pattern's implicit "no + observer on the internal edge" contract can't be honored). + """ + if not node.users: + return + + consumers_in_pattern = [ + consumer for consumer in node.users if consumer in pattern + ] + consumers_outside_pattern = [ + consumer for consumer in node.users if consumer not in pattern + ] + + has_intrinsic = _op_intrinsic_qscheme(node) is not None + + if consumers_in_pattern and consumers_outside_pattern: + raise ReconciliationError( + f"Pattern-covered node {node.name!r} (target={node.target!r}) has " + f"consumers both inside and outside its pattern group. The " + f"pattern's 'no observer on internal edges' convention can't " + f"be honored — an observer inserted at {node.name}.output for " + f"the external consumer(s) would also be visible to the " + f"in-pattern consumer(s).\n" + f" pattern group: {sorted(covered.name for covered in pattern)}\n" + f" in-pattern consumers: {sorted(consumer.name for consumer in consumers_in_pattern)}\n" + f" external consumers: {sorted(consumer.name for consumer in consumers_outside_pattern)}\n" + f"If this pattern is intentional, update " + f"_provisional_qspec_generation to decide how the internal " + f"and external consumers should share (or not share) the " + f"same observer." + ) + + fully_internal = not consumers_outside_pattern + if fully_internal: + if has_intrinsic: + raise ReconciliationError( + f"Op-intrinsic node {node.name!r} (target={node.target!r}) " + f"has all consumers inside its pattern group " + f"{sorted(covered.name for covered in pattern)}. Its " + f"intrinsic qspec (qscheme / range / observer) would be " + f"silently dropped because internal-edge OUTPUT slots " + f"aren't annotated.\n" + f"If a new pattern legitimately places an intrinsic op " + f"mid-chain, update _provisional_qspec_generation to " + f"decide how the intrinsic's constraints should be applied " + f"(likely by carrying them onto some other slot in the " + f"pattern boundary)." + ) + return + + output_slot = NodeSlot(node=node, kind=SlotKind.OUTPUT, arg_index=0) + out_spec = _lookup_by_key(cfg.op_output_spec, 0) + if out_spec is None: + return + + _populate_fields_from_spec(qspecs, output_slot, out_spec, priority) + if has_intrinsic: + _apply_op_intrinsic_override(node, output_slot, out_spec, priority, qspecs) + + +def _apply_op_intrinsic_override( + node: fx.Node, + output_slot: NodeSlot, + user_spec: TorchAOQuantizationSpec, + priority: int, + qspecs: ProvisionalQSpecMap, +) -> None: + """Overwrite the user's OUTPUT-slot fields with the op-intrinsic's known values. + + For :data:`_fixed_q_params_ops` nodes (sigmoid / tanh / hardsigmoid), + the intrinsic carries a full :class:`FixedQParamsQuantizationSpec` with + known ``qscheme``, ``quant_min``, ``quant_max``, ``scale``, and + ``zero_point``. If the user's dtype matches the intrinsic's, we + enforce all of them (range fields written directly; scale/zp routed + through ``OBSERVER_CLASS`` via a ``FixedQParamsFakeQuantize.with_args`` + partial). If the dtype differs, we only enforce ``QSCHEME`` — the + range/scale/zp values are dtype-specific and don't translate. + + For :data:`_always_affine_ops` (relu / relu6 / hardtanh(0,6)), only + ``QSCHEME`` (per_tensor_affine) is known; that's all we override. + + Every override that changes a value the user explicitly set is logged. + """ + fixed_spec = _fixed_q_params_ops.get(node.target) + override_qscheme = _op_intrinsic_qscheme(node) + assert override_qscheme is not None, "caller must gate on _op_intrinsic_qscheme" + + _override_field( + qspecs, output_slot, FieldName.QSCHEME, + override_qscheme, priority, node, user_spec.qscheme, "qscheme", + ) + + if fixed_spec is None: + # always-affine op — only qscheme is known. + return + + dtype_matches = user_spec.dtype == fixed_spec.dtype + if not dtype_matches: + logger.warning( + "Op-intrinsic override on %s (target=%s): user dtype %s doesn't " + "match intrinsic dtype %s. Enforcing qscheme only; skipping " + "range/scale/zero_point overrides because they're dtype-specific.", + node.name, node.target, user_spec.dtype, fixed_spec.dtype, + ) + return + + _override_field( + qspecs, output_slot, FieldName.QUANT_MIN, + fixed_spec.quant_min, priority, node, user_spec.quant_min, "quant_min", + ) + _override_field( + qspecs, output_slot, FieldName.QUANT_MAX, + fixed_spec.quant_max, priority, node, user_spec.quant_max, "quant_max", + ) + # scale/zp are encoded by swapping the observer class for a partial that + # pre-bakes the fixed parameters — a regular QuantizationSpec has no + # scale/zp fields. + fixed_observer_partial = FixedQParamsFakeQuantize.with_args( + scale=fixed_spec.scale, + zero_point=fixed_spec.zero_point, + quant_min=fixed_spec.quant_min, + quant_max=fixed_spec.quant_max, + dtype=fixed_spec.dtype, + qscheme=fixed_spec.qscheme, + ) + _override_field( + qspecs, output_slot, FieldName.OBSERVER_CLASS, + fixed_observer_partial, priority, node, + user_spec.observer_or_fake_quant_ctr, "observer_class", + ) + + +def _override_field( + qspecs: ProvisionalQSpecMap, + slot: NodeSlot, + field_name: FieldName, + intrinsic_value: Any, + priority: int, + node: fx.Node, + user_value: Any, + human_field_name: str, +) -> None: + """Write ``intrinsic_value`` at ``field_name`` on ``slot``, logging if + the user requested a different value. + + ``user_value`` is the value the user's spec carried (before the + override), and is compared to ``intrinsic_value`` to decide whether + the override is worth surfacing. Silent when they already agree or + the user didn't specify. + """ + qspec = _get_or_create(qspecs, slot) + qspec.fields[field_name] = FieldValue(value=intrinsic_value, priority=priority) + if user_value is not None and user_value != intrinsic_value: + logger.info( + "Op-intrinsic override on %s (target=%s): user %s=%r overridden " + "by intrinsic %s=%r.", + node.name, node.target, human_field_name, user_value, + human_field_name, intrinsic_value, + ) + + +def _op_intrinsic_qscheme(node: fx.Node) -> Any | None: + """Return the qscheme forced by op semantics, or ``None``.""" + if node.target in _fixed_q_params_ops: + return _fixed_q_params_ops[node.target].qscheme + if node.target in _always_affine_ops: + return torch.per_tensor_affine + if node.target in (torch.ops.aten.hardtanh.default, torch.ops.aten.hardtanh_.default): + if len(node.args) >= 3 and node.args[1] == 0 and node.args[2] == 6: + return torch.per_tensor_affine + return None + + +# --------------------------------------------------------------------------- +# Spec-population helpers. +# --------------------------------------------------------------------------- + + +def _populate_fields_from_spec( + qspecs: ProvisionalQSpecMap, + slot: NodeSlot, + spec: TorchAOQuantizationSpec, + priority: int, +) -> None: + """Copy every present field from a torchao :class:`QuantizationSpec` into ``qspecs``.""" + qspec = _get_or_create(qspecs, slot) + qspec.fields[FieldName.DTYPE] = FieldValue(value=spec.dtype, priority=priority) + qspec.fields[FieldName.OBSERVER_CLASS] = FieldValue( + value=spec.observer_or_fake_quant_ctr, priority=priority + ) + qspec.fields[FieldName.IS_DYNAMIC] = FieldValue(value=spec.is_dynamic, priority=priority) + if spec.quant_min is not None: + qspec.fields[FieldName.QUANT_MIN] = FieldValue(value=spec.quant_min, priority=priority) + if spec.quant_max is not None: + qspec.fields[FieldName.QUANT_MAX] = FieldValue(value=spec.quant_max, priority=priority) + if spec.qscheme is not None: + qspec.fields[FieldName.QSCHEME] = FieldValue(value=spec.qscheme, priority=priority) + if spec.ch_axis is not None: + qspec.fields[FieldName.CH_AXIS] = FieldValue(value=spec.ch_axis, priority=priority) + + +def _lookup_by_key(spec_map: dict[Any, Any], key: Any) -> Any: + """Look up ``key`` in an ``op_input/op_output/op_state`` map with ``*`` fallback.""" + if key in spec_map: + return spec_map[key] + if _ALL_TENSORS in spec_map: + return spec_map[_ALL_TENSORS] + return None diff --git a/src/coreai_opt/quantization/_graph/_qspec_constraint_generation.py b/src/coreai_opt/quantization/_graph/_qspec_constraint_generation.py new file mode 100644 index 0000000..5b7d5d3 --- /dev/null +++ b/src/coreai_opt/quantization/_graph/_qspec_constraint_generation.py @@ -0,0 +1,308 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Constraint generation for the qspec reconciliation drain loop. + +Every constraint the reconciler applies is emitted here. The pipeline +consults this module twice: + + * Once during seeding — every fx node in the graph is passed to + :func:`_generate_constraints_for_node` to build the initial + queue. + * Then on every re-enqueue — after a constraint mutates the state + of a slot, every fx node touching that slot's node is passed + back through :func:`_generate_constraints_for_node` so that a + newly-populated slot can trigger sharing constraints its neighbors + had previously suppressed. + +The three constraint sources dispatched from here are: + + * Adjacent-edge sharing (producer.OUTPUT ↔ consumer.INPUT[i]) — + :func:`_adjacent_edge_constraints`. + * Shared-observer patterns (concat, flatten, pooling) — + delegated to the pattern class via + :meth:`SharedObserverModulePattern.generate_qspec_sharing_constraints`. + * Shared-state (shared-weight) consumers — + :func:`_shared_state_constraints`. + +The context bundle :class:`_AnnotationContext` that every generator +takes is also defined here — it's the "inputs to constraint generation" +type — and :func:`_nodes_covered_by` lives here because the quantizer +uses it to enumerate a match's covered nodes when building that context. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch.fx as fx +from torch.fx.passes.utils.matcher_utils import InternalMatch +from torch.fx.passes.utils.source_matcher_utils import SourcePartition + +from coreai_opt._utils.fx_utils import is_coreai_compressed_state_node as _is_state_node +from coreai_opt.quantization.config import OpQuantizerConfig + +from ._annotation_pattern_registry import SharedObserverModulePattern +from ._annotation_utils import _get_call_function_node_from_partition +from ._qspec_constraints import Constraint, ShareObserverInstance +from ._qspec_types import NodeSlot, ProvisionalQSpecMap, SlotKind + + +# --------------------------------------------------------------------------- +# Context bundle (Phase-1/2 outputs, consumed by every generator). +# --------------------------------------------------------------------------- + + +@dataclass +class _AnnotationContext: + """Inputs to the reconciliation pipeline. + + Attributes: + winning_configs (dict[fx.Node, OpQuantizerConfig]): Highest-priority + config per covered node. + node_priorities (dict[fx.Node, int]): Position in the sorted-list + traversal. Lower = higher priority. Baked into every field + value the pipeline emits. + pattern_groups (dict[fx.Node, frozenset[fx.Node]]): For each + covered node, the set of fx nodes belonging to the same + winning pattern match. + shared_observer_nodes (dict[fx.Node, type[SharedObserverModulePattern]]): + Every shared-observer node in the model, mapped to the + pattern class that owns its qspec-sharing semantics. + :func:`_generate_constraints_for_node` dispatches through + this mapping. + """ + + winning_configs: dict[fx.Node, OpQuantizerConfig] + node_priorities: dict[fx.Node, int] + pattern_groups: dict[fx.Node, frozenset[fx.Node]] + shared_observer_nodes: dict[fx.Node, type[SharedObserverModulePattern]] + + +# --------------------------------------------------------------------------- +# Per-node constraint dispatch. +# --------------------------------------------------------------------------- + + +def _generate_constraints_for_node( + node: fx.Node, qspecs: ProvisionalQSpecMap, ctx: _AnnotationContext +) -> list[Constraint]: + """Return every constraint whose scope includes ``node``. + + Rules run against the current ``qspecs`` (dynamic peeking allowed) so + that when a slot on a downstream/upstream node has already been + decided, rules can inspect the decision and emit stronger constraints. + """ + constraints: list[Constraint] = [] + + constraints.extend(_adjacent_edge_constraints(node, qspecs, ctx)) + + pattern_class = ctx.shared_observer_nodes.get(node) + if pattern_class is not None: + constraints.extend( + pattern_class.generate_qspec_sharing_constraints(node, qspecs) + ) + + if _is_state_node(node) and len(node.users) >= 2: + constraints.extend(_shared_state_constraints(node, qspecs, ctx)) + + return constraints + + +# --------------------------------------------------------------------------- +# Adjacent-edge sharing. +# --------------------------------------------------------------------------- + + +def _adjacent_edge_constraints( + node: fx.Node, qspecs: ProvisionalQSpecMap, ctx: _AnnotationContext +) -> list[Constraint]: + """For every fx edge touching ``node`` whose endpoints are both + covered and non-internal, emit :class:`ShareObserverInstance` tying + the producer's OUTPUT slot to the consumer's INPUT slot when at + least one of the two has fields in ``qspecs``. + + Rationale mirrors the old code's :func:`_fill_input_qspec_map_for_input`: + if a producer's output is observed and its consumer would also want + to observe the same tensor, use one observer. + + Consequence — transitive grouping. Every constraint here is pairwise + (producer.OUTPUT, consumer.INPUT[i]), but + :meth:`ShareObserverInstance.apply` widens the group to include any + slot already sharing a :class:`ProvisionalQSpec` with a member. + Result: for a node ``A`` with fan-out to consumers ``B`` and ``C``, + once (A→B) and (A→C) constraints both apply, ``A.OUTPUT``, + ``B.INPUT[i]``, and ``C.INPUT[j]`` all end up in one observer group + — one fake_quant module at runtime. + + **Forecloses the requant-at-edge option.** Torchao's annotation + model natively supports having *different* concrete specs on + ``producer.output_qspec`` and ``consumer.input_qspec_map[producer]``. + At runtime that produces + ``quant(producer_spec) → dequant → quant(consumer_spec) → dequant`` + — the "intentional requant boundary" pattern. Typical use case: a + producer emits int8 activations while a specific consumer wants int4 + at that edge. By emitting :class:`ShareObserverInstance` on every + adjacent edge, this function closes off that option: divergent-spec + proposals lose to the reconciler's priority-based merge (with a + ``DTYPE_OVERRIDDEN`` relaxation), never becoming two separate + observers. + + If a future use case needs intentional requant boundaries, options + include: + + * A per-edge or per-config "don't share" opt-out that suppresses + the constraint for named edges. + * Detecting same-priority dtype conflicts and choosing "leave the + sides independent" instead of merging. + * Flipping adjacent-edge sharing from opt-out to opt-in. + + Internal edges (both endpoints in the same pattern group) are + skipped — pattern config declares specs for pattern boundary slots, + not for edges between the pattern's own constituent nodes. + """ + if node not in ctx.winning_configs: + # Uncovered nodes (placeholders, unmatched ops) never anchor an + # adjacent-edge constraint on either side; a covered neighbor + # that cares about this edge emits it from its own side. + return [] + + constraints: list[Constraint] = [] + pattern = ctx.pattern_groups.get(node, frozenset({node})) + + # Incoming edges: (producer, node, arg_index). + for arg_index, producer in enumerate(node.all_input_nodes): + if producer in pattern: + continue + if producer not in ctx.winning_configs: + continue + producer_output = NodeSlot(node=producer, kind=SlotKind.OUTPUT, arg_index=0) + consumer_input = NodeSlot(node=node, kind=SlotKind.INPUT, arg_index=arg_index) + if producer_output in qspecs or consumer_input in qspecs: + constraints.append( + ShareObserverInstance(frozenset({producer_output, consumer_input})) + ) + + # Outgoing edges: (node, consumer, arg_index). + for consumer in node.users: + if consumer in pattern: + continue + if consumer not in ctx.winning_configs: + continue + consumer_input = _find_input_slot_for_producer(consumer, node) + if consumer_input is None: + continue + producer_output = NodeSlot(node=node, kind=SlotKind.OUTPUT, arg_index=0) + if producer_output in qspecs or consumer_input in qspecs: + constraints.append( + ShareObserverInstance(frozenset({producer_output, consumer_input})) + ) + + return constraints + + +def _find_input_slot_for_producer( + consumer: fx.Node, producer: fx.Node +) -> NodeSlot | None: + """Return the consumer's INPUT slot that reads from ``producer``, or + ``None`` if no such slot exists (defensive; shouldn't happen for a + well-formed users-relation). + """ + for arg_index, actual_producer in enumerate(consumer.all_input_nodes): + if actual_producer is producer: + return NodeSlot(node=consumer, kind=SlotKind.INPUT, arg_index=arg_index) + return None + + +# --------------------------------------------------------------------------- +# Shared-state (shared-weight) sharing. +# --------------------------------------------------------------------------- + + +def _shared_state_constraints( + state_node: fx.Node, qspecs: ProvisionalQSpecMap, ctx: _AnnotationContext +) -> list[Constraint]: + """Emit one :class:`ShareObserverInstance` tying every covered + consumer's INPUT slot that reads ``state_node``. + + Weight tensors get quantized once at prepare time; all consumers of + the same underlying tensor see the same quantized version at runtime. + This constraint forces all covered consumers onto one qspec + one + observer instance. + + **Forecloses divergent-spec consumers of a shared weight.** Torchao's + annotation model natively supports two consumers of the same state + tensor having different ``input_qspec_map[state]`` entries — it + inserts one fake_quant per edge, each with its own spec. Both + consumers receive their own quantized-then-dequantized view of the + same underlying float tensor. + + **State-tensor-specific wrinkle.** Weight observer scale/zp is + *derived* from the tensor values at prepare time (deterministic), + not *learned* from calibration data at runtime. So + :class:`ShareObserverInstance` on shared weights bundles two things + that could be separated: + + 1. Same qspec across consumers (dtype / qscheme / range agreement). + 2. Same observer instance (compute scale/zp once vs. per-consumer). + + For weights (2) is a trivial optimization — two observers on the + same tensor compute identical values anyway. The strong constraint + is (1). If a future use case needs divergent-spec consumers of a + shared weight, loosening (1) is the right move — (2) falls out + naturally. + + **Downstream storage/export tradeoff.** If divergent specs are + allowed, ``convert_pt2e`` / export gets to choose the storage + strategy: store multiple quantized versions (Nx storage, no runtime + requant), or store the least aggressive version and requant to more + aggressive variants at each edge (1x storage, one requant per + divergent consumer). The reconciler doesn't need to pick — it + just needs to permit divergent specs. + """ + consumer_slots: list[NodeSlot] = [] + for consumer in state_node.users: + if consumer not in ctx.winning_configs: + continue + consumer_input = _find_input_slot_for_producer(consumer, state_node) + if consumer_input is None: + continue + consumer_slots.append(consumer_input) + + if len(consumer_slots) < 2: + return [] + # Only fire when at least one consumer wants observation — if none + # do, there's no observer to share. + if not any(slot in qspecs for slot in consumer_slots): + return [] + return [ShareObserverInstance(frozenset(consumer_slots))] + + +# --------------------------------------------------------------------------- +# Pattern coverage helper (used by quantizer.py when building the context). +# --------------------------------------------------------------------------- + + +def _nodes_covered_by(match_info: Any) -> list[fx.Node]: + """Return every fx node an annotator match annotates.""" + match = match_info.annotator_match + + if isinstance(match, InternalMatch): + return [ + node + for key, node in match.name_node_map.items() + if key in ("mod", "output", "bn") + ] + + if isinstance(match, tuple) and all( + isinstance(partition, SourcePartition) for partition in match + ): + return [_get_call_function_node_from_partition(partition) for partition in match] + + raise TypeError( + f"Unknown annotator match type: {type(match).__name__}. Update " + f"_nodes_covered_by when adding a new pattern family." + ) diff --git a/src/coreai_opt/quantization/_graph/_qspec_constraints.py b/src/coreai_opt/quantization/_graph/_qspec_constraints.py new file mode 100644 index 0000000..78a9c52 --- /dev/null +++ b/src/coreai_opt/quantization/_graph/_qspec_constraints.py @@ -0,0 +1,371 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Constraint types + per-field reconciliation policies. + +The two :class:`Constraint` subclasses (:class:`ShareFields`, +:class:`ShareObserverInstance`) are the vocabulary constraint generators +emit and the drain loop applies. Each per-field reconciliation rule +lives here as a small pure function, dispatched via +:data:`_FIELD_POLICY` in :func:`_reconcile_field`. +""" + +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch +from torchao.quantization.pt2e.fake_quantize import FixedQParamsFakeQuantize +from torchao.quantization.pt2e.observer import FixedQParamsObserver + +from ._qspec_types import ( + FieldName, + FieldValue, + NodeSlot, + ProvisionalQSpec, + ProvisionalQSpecMap, + ReconciliationError, +) + +logger = logging.getLogger(__name__) + + +_ALL_FIELDS: frozenset[FieldName] = frozenset(FieldName) + + +# --------------------------------------------------------------------------- +# Constraint ABC + implementations. +# --------------------------------------------------------------------------- + + +class Constraint(ABC): + """A relation over one or more slots. + + :meth:`apply` reads the current :class:`ProvisionalQSpecMap`, + mutates it to enforce the constraint, and returns the set of slots + whose state changed. Returning an empty set means the constraint was + already satisfied — that's the fixed-point signal the drain loop + uses to detect quiescence. + """ + + @property + @abstractmethod + def slots(self) -> frozenset[NodeSlot]: ... + + @abstractmethod + def apply(self, qspecs: ProvisionalQSpecMap) -> set[NodeSlot]: ... + + +@dataclass(frozen=True) +class ShareFields(Constraint): + """Every slot in ``slots`` must agree on each of ``fields``. + + Reconciliation is per-field (see :func:`_reconcile_field`). The + winning value is broadcast to every slot in the group; slots that + had no proposal for a field pick up the winner. Priority of the + reconciled value inherits the highest priority (lowest priority- + number) among the contributing proposals — a monotonicity that + guarantees the drain terminates. + """ + + _slots: frozenset[NodeSlot] + fields: frozenset[FieldName] + + @property + def slots(self) -> frozenset[NodeSlot]: + return self._slots + + def apply(self, qspecs: ProvisionalQSpecMap) -> set[NodeSlot]: + changed: set[NodeSlot] = set() + for field_name in self.fields: + reconciled = _reconcile_field(field_name, self._slots, qspecs) + if reconciled is None: + continue + for slot in self._slots: + qspec = _get_or_create(qspecs, slot) + current = qspec.fields.get(field_name) + if _field_value_stronger(reconciled, current): + qspec.fields[field_name] = reconciled + changed.add(slot) + return changed + + +@dataclass(frozen=True) +class ShareObserverInstance(Constraint): + """Every slot in ``slots`` must reference the same + :class:`ProvisionalQSpec` object (same observer instance at runtime). + + Applying: + + 1. Widen the group to include every slot already sharing a + :class:`ProvisionalQSpec` with any member of ``self.slots``. + 2. Reconcile every field across the widened group via the same + per-field policies :class:`ShareFields` uses. + 3. Re-point every widened slot's map entry at one shared + :class:`ProvisionalQSpec`. + + After apply, mutating any member's fields propagates to every + sharer — object identity IS the sharing relation. + """ + + _slots: frozenset[NodeSlot] + + @property + def slots(self) -> frozenset[NodeSlot]: + return self._slots + + def apply(self, qspecs: ProvisionalQSpecMap) -> set[NodeSlot]: + # Widen: pull in every slot that already shares a + # ProvisionalQSpec with any input slot. + target_ids = {id(qspecs[slot]) for slot in self._slots if slot in qspecs} + widened: set[NodeSlot] = set(self._slots) + for slot, qspec in qspecs.items(): + if id(qspec) in target_ids: + widened.add(slot) + + # Reconcile every field across the widened group. + reconciled_fields: dict[FieldName, FieldValue] = {} + for field_name in _ALL_FIELDS: + reconciled = _reconcile_field(field_name, frozenset(widened), qspecs) + if reconciled is not None: + reconciled_fields[field_name] = reconciled + + changed: set[NodeSlot] = set() + first_ref: ProvisionalQSpec | None = None + for slot in widened: + existing = qspecs.get(slot) + already_reconciled = existing is not None and existing.fields == reconciled_fields + if first_ref is None: + if already_reconciled: + first_ref = existing + else: + first_ref = ProvisionalQSpec(fields=dict(reconciled_fields)) + qspecs[slot] = first_ref + changed.add(slot) + else: + if qspecs.get(slot) is not first_ref: + qspecs[slot] = first_ref + changed.add(slot) + if not already_reconciled and first_ref.fields != reconciled_fields: + first_ref.fields = dict(reconciled_fields) + changed.update(widened) + if first_ref is not None and first_ref.fields != reconciled_fields: + first_ref.fields = dict(reconciled_fields) + changed.update(widened) + return changed + + +# --------------------------------------------------------------------------- +# Per-field reconciliation. +# --------------------------------------------------------------------------- + + +def _reconcile_field( + field_name: FieldName, slots: frozenset[NodeSlot], qspecs: ProvisionalQSpecMap +) -> FieldValue | None: + """Compute the reconciled :class:`FieldValue` for one field across ``slots``. + + Returns ``None`` when no slot has a proposal for the field (nothing + to do). Raises :class:`ReconciliationError` on infeasible fields per + the per-field policy. + """ + proposals = [ + qspecs[slot].fields[field_name] + for slot in slots + if slot in qspecs and field_name in qspecs[slot].fields + ] + if not proposals: + return None + policy = _FIELD_POLICY[field_name] + return policy(proposals) + + +def _priority_min(proposals: Sequence[FieldValue]) -> int: + """Highest priority (lowest priority-number) among the proposals.""" + return min(proposal.priority for proposal in proposals) + + +def _policy_priority_wins(proposals: Sequence[FieldValue]) -> FieldValue: + """Value from the highest-priority proposal wins. Ties: first encountered.""" + winner = min(proposals, key=lambda proposal: proposal.priority) + return FieldValue(value=winner.value, priority=_priority_min(proposals)) + + +def _policy_must_agree(proposals: Sequence[FieldValue]) -> FieldValue: + """All values must be equal, else :class:`ReconciliationError`.""" + values = {proposal.value for proposal in proposals} + if len(values) > 1: + raise ReconciliationError( + f"Incompatible values across slots: {sorted(str(value) for value in values)}. " + f"Proposals: {[(proposal.value, proposal.priority) for proposal in proposals]}" + ) + return FieldValue(value=next(iter(values)), priority=_priority_min(proposals)) + + +def _policy_qscheme_lattice(proposals: Sequence[FieldValue]) -> FieldValue: + """Lattice join: symmetric ⊂ affine; per_tensor ⊂ per_channel. Looser wins.""" + qschemes = [proposal.value for proposal in proposals if proposal.value is not None] + if not qschemes: + return FieldValue(value=None, priority=_priority_min(proposals)) + joined = _join_qscheme(set(qschemes)) + return FieldValue(value=joined, priority=_priority_min(proposals)) + + +def _policy_range_min(proposals: Sequence[FieldValue]) -> FieldValue: + """Union — take the smallest quant_min across proposals.""" + values = [proposal.value for proposal in proposals if proposal.value is not None] + if not values: + return FieldValue(value=None, priority=_priority_min(proposals)) + return FieldValue(value=min(values), priority=_priority_min(proposals)) + + +def _policy_range_max(proposals: Sequence[FieldValue]) -> FieldValue: + """Union — take the largest quant_max across proposals.""" + values = [proposal.value for proposal in proposals if proposal.value is not None] + if not values: + return FieldValue(value=None, priority=_priority_min(proposals)) + return FieldValue(value=max(values), priority=_priority_min(proposals)) + + +def _policy_observer_class_lattice(proposals: Sequence[FieldValue]) -> FieldValue: + """Observer-class lattice: ``fixed ⊂ learning``. + + Fixed observers (``FixedQParamsFakeQuantize`` / ``FixedQParamsObserver``, + optionally wrapped in a ``.with_args(...)`` partial) are a strict + specialization of learning observers: they clamp all statistics to a + baked-in scale / zero_point / range. If a fixed-observer proposal + ever shares a group with a learning-observer proposal, using the + fixed one would clamp the group's real observed values to the fixed + range — usually numerically wrong when the group contains slots + whose values fall outside the fixed range. + + Rule: + + * If any non-fixed proposal is present, prefer it + (priority-wins among non-fixed). Fixed proposals that lose + get a log entry noting the demotion — calibration behavior + shifts from "clamp to baked range" to "learn observed range". + * If every proposal is fixed, they must agree on their baked + parameters (a fixed observer's identity is its partial's + value); disagreement is a :class:`ReconciliationError`. + """ + non_fixed = [ + proposal for proposal in proposals if not _is_fixed_observer(proposal.value) + ] + fixed = [proposal for proposal in proposals if _is_fixed_observer(proposal.value)] + + if non_fixed: + winner = min(non_fixed, key=lambda proposal: proposal.priority) + for demoted in fixed: + logger.info( + "OBSERVER_CLASS reconciliation: fixed observer %r (priority %d) " + "demoted to learning observer %r (priority %d) — group contains " + "learning-observer members whose values may fall outside the " + "fixed observer's baked range.", + demoted.value, + demoted.priority, + winner.value, + winner.priority, + ) + return FieldValue(value=winner.value, priority=_priority_min(proposals)) + + # All fixed — must agree on the exact observer partial (baked + # scale / zp / range are part of the partial's identity). + return _policy_must_agree(proposals) + + +def _is_fixed_observer(observer_class: Any) -> bool: + """Return True if ``observer_class`` names a fixed-params observer. + + Accepts either a bare class or a ``.with_args(...)`` partial + (torchao's ``PartialWrapper``). Fixed observers include + :class:`FixedQParamsFakeQuantize` and :class:`FixedQParamsObserver` + and any partial whose underlying class is one of those. + """ + fixed_classes = (FixedQParamsFakeQuantize, FixedQParamsObserver) + if isinstance(observer_class, type) and issubclass(observer_class, fixed_classes): + return True + # Torchao's PartialWrapper exposes the underlying class via ``.p.func``. + underlying = getattr(getattr(observer_class, "p", None), "func", None) + if isinstance(underlying, type) and issubclass(underlying, fixed_classes): + return True + return False + + +_FIELD_POLICY: dict[FieldName, Any] = { + FieldName.DTYPE: _policy_priority_wins, + FieldName.QSCHEME: _policy_qscheme_lattice, + FieldName.CH_AXIS: _policy_must_agree, + FieldName.QUANT_MIN: _policy_range_min, + FieldName.QUANT_MAX: _policy_range_max, + FieldName.IS_DYNAMIC: _policy_must_agree, + FieldName.OBSERVER_CLASS: _policy_observer_class_lattice, +} + + +def _field_value_stronger(new: FieldValue, current: FieldValue | None) -> bool: + """Return True iff writing ``new`` would advance the state. + + "Advance" = value differs, or value equals but priority is stronger + (lower priority-number). Enables no-op detection for convergence. + """ + if current is None: + return True + if new.value != current.value: + return True + return new.priority < current.priority + + +# --------------------------------------------------------------------------- +# qscheme lattice helpers. +# --------------------------------------------------------------------------- + + +_QSCHEME_IS_SYMMETRIC: dict[Any, bool] = { + None: True, + torch.per_tensor_symmetric: True, + torch.per_tensor_affine: False, + torch.per_channel_symmetric: True, + torch.per_channel_affine: False, +} +_QSCHEME_IS_PER_CHANNEL: dict[Any, bool] = { + None: False, + torch.per_tensor_symmetric: False, + torch.per_tensor_affine: False, + torch.per_channel_symmetric: True, + torch.per_channel_affine: True, +} +PER_CHANNEL_SCHEMES = (torch.per_channel_symmetric, torch.per_channel_affine) + + +def _join_qscheme(qschemes: set[Any]) -> Any: + """Lattice join over ``qschemes``: looser wins in symmetry and granularity.""" + is_sym = all( + _QSCHEME_IS_SYMMETRIC.get(qscheme, False) + for qscheme in qschemes + if qscheme is not None + ) + is_per_channel = any( + _QSCHEME_IS_PER_CHANNEL.get(qscheme, False) for qscheme in qschemes + ) + if is_per_channel: + return torch.per_channel_symmetric if is_sym else torch.per_channel_affine + return torch.per_tensor_symmetric if is_sym else torch.per_tensor_affine + + +# --------------------------------------------------------------------------- +# Small helpers callers need. +# --------------------------------------------------------------------------- + + +def _get_or_create(qspecs: ProvisionalQSpecMap, slot: NodeSlot) -> ProvisionalQSpec: + """Return the :class:`ProvisionalQSpec` for ``slot``, creating one if absent.""" + if slot not in qspecs: + qspecs[slot] = ProvisionalQSpec() + return qspecs[slot] diff --git a/src/coreai_opt/quantization/_graph/_qspec_reconcile.py b/src/coreai_opt/quantization/_graph/_qspec_reconcile.py new file mode 100644 index 0000000..5a99b88 --- /dev/null +++ b/src/coreai_opt/quantization/_graph/_qspec_reconcile.py @@ -0,0 +1,92 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Constraint-queue reconciliation for graph-mode quantization annotation. + +Driver for the full pipeline. Types, per-field policies, constraint +generation, and resolution live in sibling modules: + + * ``_qspec_types`` — :class:`NodeSlot`, :class:`FieldName`, + :class:`ProvisionalQSpec`, :class:`ProvisionalQSpecMap`, etc. + * ``_qspec_constraints`` — :class:`Constraint` ABC, per-field + policies, :func:`_reconcile_field`. + * ``_provisional_qspec_generation`` — :func:`build_initial_state`. + * ``_qspec_constraint_generation`` — :class:`_AnnotationContext`, + :func:`_generate_constraints_for_node`, + :func:`_nodes_covered_by`. + * ``_qspec_resolution`` — :func:`resolve_qspecs` (per-node + annotation write). + +Pipeline: + + 1. Pattern matching (upstream, in ``_AnnotationHandler.annotate``). + 2. :func:`build_initial_state` populates per-slot proposals from + winning configs + op-intrinsic overrides. + 3. Seed the constraint queue by running + :func:`_generate_constraints_for_node` over every fx node. + 4. Drain: pop a constraint, apply it, re-enqueue constraints for + every node whose slots changed. Convergence guaranteed by + per-field priority monotonicity. + 5. :func:`resolve_qspecs` projects the final map onto per-node + :class:`QuantizationAnnotation` fields. +""" + +from __future__ import annotations + +from collections import deque + +import torch.fx as fx + +from ._provisional_qspec_generation import build_initial_provisional_qspecs +from ._qspec_constraint_generation import ( + _AnnotationContext, + _generate_constraints_for_node, + _nodes_covered_by, +) +from ._qspec_constraints import Constraint +from ._qspec_resolution import resolve_qspecs + +# Re-export so callers depend only on _qspec_reconcile. +__all__ = [ + "_AnnotationContext", + "_nodes_covered_by", + "annotate_via_reconciliation", +] + + +def annotate_via_reconciliation( + model: fx.GraphModule, ctx: _AnnotationContext +) -> fx.GraphModule: + """Annotate ``model`` in place using the constraint-queue reconciler.""" + qspecs = build_initial_provisional_qspecs( + model, + winning_configs=ctx.winning_configs, + node_priorities=ctx.node_priorities, + pattern_groups=ctx.pattern_groups, + ) + + queue: deque[Constraint] = deque() + for node in model.graph.nodes: + queue.extend(_generate_constraints_for_node(node, qspecs, ctx)) + + # Runaway-loop backstop; convergence is guaranteed by priority + # monotonicity but a bug in a policy could stall. + max_iters = 1000 * (1 + sum(1 for _ in model.graph.nodes)) * 8 + iters = 0 + while queue: + constraint = queue.popleft() + changed = constraint.apply(qspecs) + iters += 1 + if iters > max_iters: + raise RuntimeError( + f"Constraint drain did not converge after {iters} iterations — " + f"likely a bug in a reconciliation policy or generator." + ) + touched_nodes = {slot.node for slot in changed} + for touched in touched_nodes: + queue.extend(_generate_constraints_for_node(touched, qspecs, ctx)) + + resolve_qspecs(model, qspecs) + return model diff --git a/src/coreai_opt/quantization/_graph/_qspec_resolution.py b/src/coreai_opt/quantization/_graph/_qspec_resolution.py new file mode 100644 index 0000000..e107d5d --- /dev/null +++ b/src/coreai_opt/quantization/_graph/_qspec_resolution.py @@ -0,0 +1,306 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Phase 4 of the qspec reconciliation pipeline: resolve the final +:class:`ProvisionalQSpecMap` onto per-node +:class:`QuantizationAnnotation` fields on the fx graph. + +By the time this runs, the constraint queue has drained, every slot +has its reconciled fields, and slots that must share an observer at +runtime are already pointing at the same :class:`ProvisionalQSpec` +object. This module's job is to project that state onto torchao's +per-node data model: + + 1. Group slots by :class:`ProvisionalQSpec` object identity. + 2. For each group, pick an anchor slot (topologically-first; + INPUT preferred over OUTPUT within a node — see + :class:`SlotOrderKey` for why). + 3. Assign the anchor a concrete + :class:`torchao...QuantizationSpec`; the others a + :class:`torchao...SharedQuantizationSpec` pointing at the anchor. + 4. Bucket per node, backfill missing input slots with ``None`` + (a torchao interop concession), and write + :class:`QuantizationAnnotation` entries. +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass + +import torch.fx as fx +from torchao.quantization.pt2e.quantizer import ( + QuantizationAnnotation, + QuantizationSpec as TorchAOQuantizationSpec, + SharedQuantizationSpec as _SharedQuantizationSpec, +) + +# One of these is assigned to every observed slot: a concrete spec on +# the group's anchor slot, a SharedQuantizationSpec pointing at the +# anchor on every other slot in the group. +_SlotSpec = TorchAOQuantizationSpec | _SharedQuantizationSpec +from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY + +from ._qspec_types import ( + FieldName, + NodeSlot, + ProvisionalQSpec, + ProvisionalQSpecMap, + SlotKind, +) + + +# --------------------------------------------------------------------------- +# Entry point. +# --------------------------------------------------------------------------- + + +def resolve_qspecs(model: fx.GraphModule, qspecs: ProvisionalQSpecMap) -> None: + """Project reconciled per-slot state onto per-node annotations. + + Mutates ``model``'s node metadata in place. See module docstring for + the four-step pipeline. + """ + groups = _group_slots_by_qspec_identity(qspecs) + topo_index = _topo_index(model) + + slot_assignments: dict[NodeSlot, _SlotSpec] = {} + for group in groups: + _assign_group_specs(group, topo_index, slot_assignments) + + per_node_inputs, per_node_outputs = _bucket_per_node(slot_assignments) + _write_annotations(model, per_node_inputs, per_node_outputs) + + +# --------------------------------------------------------------------------- +# Grouping. +# --------------------------------------------------------------------------- + + +@dataclass +class _QSpecGroup: + """One equivalence class of slots — every slot points at the same + :class:`ProvisionalQSpec` object. + """ + + qspec: ProvisionalQSpec + slots: list[NodeSlot] + + +def _group_slots_by_qspec_identity(qspecs: ProvisionalQSpecMap) -> list[_QSpecGroup]: + """Bucket slots by ``id(ProvisionalQSpec)`` — same object = same group.""" + groups_by_id: dict[int, _QSpecGroup] = {} + for slot, qspec in qspecs.items(): + group = groups_by_id.get(id(qspec)) + if group is None: + group = _QSpecGroup(qspec=qspec, slots=[]) + groups_by_id[id(qspec)] = group + group.slots.append(slot) + return list(groups_by_id.values()) + + +# --------------------------------------------------------------------------- +# Anchor selection. +# --------------------------------------------------------------------------- + + +@dataclass(order=True, frozen=True) +class SlotOrderKey: + """Sort key for anchor selection within a shared-observer group. + + Dataclasses with ``order=True`` compare field-by-field in declaration + order, so the ordering below is exactly ``(topo_index, kind_order, + arg_index)`` lexicographically. + + Attributes: + topo_index (int): Position of the slot's node in + ``model.graph.nodes`` iteration order. Torchao processes + nodes in this order; the anchor must be topo-first so its + observer is registered in ``obs_or_fq_map`` before any + downstream SharedSpec references it. + kind_order (int): Within a single node, INPUT (0) sorts before + OUTPUT (1). Torchao's ``_get_edge_or_node_to_qspec`` + iterates ``input_qspec_map`` entries before ``output_qspec`` + for a given node; anchoring on the input side ensures the + observer is registered by the time the OUTPUT-side SharedSpec + on the same node resolves. + arg_index (int): For INPUT slots, distinguishes positional + inputs. Tiebreaker only; the anchor's position doesn't + affect torchao's semantics. + """ + + topo_index: int + kind_order: int + arg_index: int + + +def _topo_index(model: fx.GraphModule) -> dict[fx.Node, int]: + return {node: i for i, node in enumerate(model.graph.nodes)} + + +def _slot_order_key(slot: NodeSlot, topo_index: dict[fx.Node, int]) -> SlotOrderKey: + return SlotOrderKey( + topo_index=topo_index.get(slot.node, len(topo_index)), + kind_order=0 if slot.kind is SlotKind.INPUT else 1, + arg_index=slot.arg_index, + ) + + +def _pick_anchor(group: _QSpecGroup, topo_index: dict[fx.Node, int]) -> NodeSlot: + """Return the topologically-first slot in ``group``. See + :class:`SlotOrderKey` for why INPUT is preferred over OUTPUT within + a single node.""" + return min(group.slots, key=lambda slot: _slot_order_key(slot, topo_index)) + + +# --------------------------------------------------------------------------- +# Spec assignment. +# --------------------------------------------------------------------------- + + +def _assign_group_specs( + group: _QSpecGroup, + topo_index: dict[fx.Node, int], + slot_assignments: dict[NodeSlot, _SlotSpec], +) -> None: + """Fill ``slot_assignments`` with the spec each slot in ``group`` should carry. + + Singleton groups (or all-``None`` groups) just get the concrete spec + on the one slot. Multi-slot groups pick an anchor (concrete spec) + and give the rest a :class:`SharedQuantizationSpec` pointing back. + """ + concrete = _build_concrete_spec(group.qspec) + if concrete is None: + # Empty or None-only group — no annotation to emit. + return + if len(group.slots) == 1: + slot_assignments[group.slots[0]] = concrete + return + + anchor = _pick_anchor(group, topo_index) + shared_ref = _shared_spec_pointing_at(anchor) + for slot in group.slots: + slot_assignments[slot] = concrete if slot == anchor else shared_ref + + +def _shared_spec_pointing_at(anchor: NodeSlot) -> _SharedQuantizationSpec: + """Build a :class:`SharedQuantizationSpec` in the form matching where + torchao registers ``anchor``'s observer. + + Torchao's ``obs_or_fq_map`` uses different keys depending on which + annotation field placed the observer: + + * ``node.output_qspec = spec`` → registered under key ``node``. + Reference via **Node form** ``SharedQuantizationSpec(node)``. + * ``consumer.input_qspec_map[producer] = spec`` → registered under + key ``(producer, consumer)``. Reference via **edge form** + ``SharedQuantizationSpec((producer, consumer))``. + """ + if anchor.kind is SlotKind.OUTPUT: + return _SharedQuantizationSpec(anchor.node) + producer = anchor.node.all_input_nodes[anchor.arg_index] + return _SharedQuantizationSpec((producer, anchor.node)) + + +def _build_concrete_spec(qspec: ProvisionalQSpec) -> TorchAOQuantizationSpec | None: + """Assemble a :class:`TorchAOQuantizationSpec` from reconciled fields. + + Returns ``None`` if required fields (``DTYPE``, ``OBSERVER_CLASS``) are + missing — that ProvisionalQSpec represents a slot the pipeline chose + not to observe (e.g. a bias slot with ``op_state_spec`` unset). + """ + if FieldName.DTYPE not in qspec.fields or FieldName.OBSERVER_CLASS not in qspec.fields: + return None + return TorchAOQuantizationSpec( + dtype=qspec.fields[FieldName.DTYPE].value, + observer_or_fake_quant_ctr=qspec.fields[FieldName.OBSERVER_CLASS].value, + quant_min=( + qspec.fields[FieldName.QUANT_MIN].value + if FieldName.QUANT_MIN in qspec.fields + else None + ), + quant_max=( + qspec.fields[FieldName.QUANT_MAX].value + if FieldName.QUANT_MAX in qspec.fields + else None + ), + qscheme=( + qspec.fields[FieldName.QSCHEME].value + if FieldName.QSCHEME in qspec.fields + else None + ), + ch_axis=( + qspec.fields[FieldName.CH_AXIS].value + if FieldName.CH_AXIS in qspec.fields + else None + ), + is_dynamic=( + qspec.fields[FieldName.IS_DYNAMIC].value + if FieldName.IS_DYNAMIC in qspec.fields + else False + ), + ) + + +# --------------------------------------------------------------------------- +# Per-node bucketing + write. +# --------------------------------------------------------------------------- + + +def _bucket_per_node( + slot_assignments: dict[NodeSlot, _SlotSpec], +) -> tuple[dict[fx.Node, dict[fx.Node, _SlotSpec]], dict[fx.Node, _SlotSpec]]: + """Split per-slot spec assignments into per-node ``input_qspec_map`` + and per-node ``output_qspec`` buckets. + """ + per_node_inputs: dict[fx.Node, dict[fx.Node, _SlotSpec]] = defaultdict(dict) + per_node_outputs: dict[fx.Node, _SlotSpec] = {} + for slot, spec in slot_assignments.items(): + if slot.kind is SlotKind.OUTPUT: + per_node_outputs[slot.node] = spec + else: + producer = slot.node.all_input_nodes[slot.arg_index] + per_node_inputs[slot.node][producer] = spec + return per_node_inputs, per_node_outputs + + +def _write_annotations( + model: fx.GraphModule, + per_node_inputs: dict[fx.Node, dict[fx.Node, _SlotSpec]], + per_node_outputs: dict[fx.Node, _SlotSpec], +) -> None: + """Mutate each touched node's meta with a :class:`QuantizationAnnotation`. + + Nodes with any input-side annotation get their ``input_qspec_map`` + backfilled with ``None`` for every positional input that didn't get + an explicit spec — torchao's downstream passes (notably + ``qat_utils.py``'s ``_replace_target_node_with_quantization_annotation``) + look up slots by positional index and IndexError if entries are + missing. This is a torchao interop concession, not a reconciliation + decision. + """ + touched: set[fx.Node] = set(per_node_inputs) | set(per_node_outputs) + for node in touched: + annotation = node.meta.get(Q_ANNOTATION_KEY, QuantizationAnnotation()) + input_map = per_node_inputs.get(node) + if input_map: + annotation.input_qspec_map = _backfill_input_qspec_map(node, input_map) + if node in per_node_outputs: + annotation.output_qspec = per_node_outputs[node] + annotation._annotated = True + node.meta[Q_ANNOTATION_KEY] = annotation + + +def _backfill_input_qspec_map( + node: fx.Node, input_map: dict[fx.Node, _SlotSpec] +) -> dict[fx.Node, _SlotSpec | None]: + """Return an ``input_qspec_map`` with an entry for every + ``node.all_input_nodes`` position — explicit specs where set, + ``None`` otherwise. + """ + return { + producer: input_map.get(producer, None) + for producer in node.all_input_nodes + } diff --git a/src/coreai_opt/quantization/_graph/_qspec_types.py b/src/coreai_opt/quantization/_graph/_qspec_types.py new file mode 100644 index 0000000..36c3e44 --- /dev/null +++ b/src/coreai_opt/quantization/_graph/_qspec_types.py @@ -0,0 +1,118 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Shared types for the qspec reconciliation pipeline. + +These types are the vocabulary every phase of the pipeline speaks — +provisional-qspec generation, constraint generation + drain, and +resolution onto per-node annotations. Keeping them in a small, +dependency-free module lets each phase live in its own file without +importing sideways into another phase. + +Nothing here is behavior; only data + enums + one exception type. +""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import Any + +import torch.fx as fx + + +# --------------------------------------------------------------------------- +# Atomic units. +# --------------------------------------------------------------------------- + + +class SlotKind(enum.Enum): + """Which side of a node a slot lives on.""" + + INPUT = enum.auto() + OUTPUT = enum.auto() + + +@dataclass(frozen=True) +class NodeSlot: + """A single quantization decision point on a node. + + Each fx node contributes one ``OUTPUT`` slot (``arg_index=0``) and one + ``INPUT`` slot per unique input in ``node.all_input_nodes`` + (``arg_index`` tracks position). Slots are the reconciliation atom. + """ + + node: fx.Node + kind: SlotKind + arg_index: int + + +class FieldName(enum.Enum): + """Fields on a :class:`ProvisionalQSpec`. + + Each field has its own reconciliation policy — see + ``_qspec_constraints._FIELD_POLICY``. + """ + + DTYPE = enum.auto() + QSCHEME = enum.auto() # combined symmetry × granularity — one lattice-joined field + CH_AXIS = enum.auto() + QUANT_MIN = enum.auto() + QUANT_MAX = enum.auto() + IS_DYNAMIC = enum.auto() + OBSERVER_CLASS = enum.auto() + + +@dataclass(frozen=True) +class FieldValue: + """A value proposed for one field on one slot, plus its priority. + + Priority follows the sorted-list convention: **lower priority number = + higher priority** (top of the sort). Reconciliation only ever raises + priority-ness (lowers the number). + """ + + value: Any + priority: int + + +@dataclass +class ProvisionalQSpec: + """Mutable per-observer state. + + Multiple slots may reference one instance — that's how + ``ShareObserverInstance`` expresses sharing: two slots point at the + same Python object, so any field mutation is seen by both. Object + identity IS the sharing relation; there's no separate side-list. + """ + + fields: dict[FieldName, FieldValue] = field(default_factory=dict) + + +ProvisionalQSpecMap = dict[NodeSlot, ProvisionalQSpec] +"""Alias for the map every phase reads/writes. + +Two slots pointing at the same :class:`ProvisionalQSpec` object share an +observer at runtime. Two slots pointing at distinct objects reconcile +independently. +""" + + +# --------------------------------------------------------------------------- +# Cross-phase exception. +# --------------------------------------------------------------------------- + + +class ReconciliationError(RuntimeError): + """Raised when reconciliation hits an unresolvable state. + + Fatal cases: + + * ``IS_DYNAMIC`` disagreement across slots. + * Incompatible ``CH_AXIS`` values across per_channel slots. + * All-fixed ``OBSERVER_CLASS`` group whose fixed params disagree. + * Pattern-structure violations detected during provisional-qspec + generation (see ``_provisional_qspec_generation.py``). + """ diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py index 4c066e5..1c53a28 100644 --- a/src/coreai_opt/quantization/_graph/quantizer.py +++ b/src/coreai_opt/quantization/_graph/quantizer.py @@ -85,9 +85,7 @@ ) from ._annotation_utils import ( _get_input_qspec_map, - adjust_output_qspec_for_qscheme_and_propagate, annotate_module_level_specs as _annotate_module_level_specs, - is_node_annotated, ) from ._conv_bn_utils import ( fold_conv_bn_weights as _fold_conv_bn_weights, @@ -98,6 +96,11 @@ prepare_for_mil_export, prepare_for_mlir_export, ) +from ._qspec_reconcile import ( + _AnnotationContext, + _nodes_covered_by, + annotate_via_reconciliation, +) from ._utils import ( force_per_tensor_for_channel_altering_ops as _force_per_tensor_for_channel_altering_ops, get_source_module_name as _get_source_module_name, @@ -225,22 +228,30 @@ def _all_patterns(self) -> list[type]: """Globally-registered patterns plus per-instance ``extra_patterns``.""" return list(_AnnotationPatternRegistry.list_registry_values()) + self._extra_patterns - def _get_shared_observer_nodes(self, model: torch.fx.GraphModule) -> set[torch.fx.Node]: - """ - Return a set of all shared observer nodes in the model. + def _get_shared_observer_nodes( + self, model: torch.fx.GraphModule + ) -> dict[torch.fx.Node, type[_SharedObserverModulePattern]]: + """Return every shared-observer node in the model, mapped to the + pattern class that matched it. + + The pattern class is what owns the node's qspec-sharing semantics + (via :meth:`SharedObserverModulePattern.generate_qspec_sharing_constraints`). + Callers that only need the node set can take ``.keys()``. """ shared_observer_annotators = [ - a_class - for a_class in self._all_patterns() - if issubclass(a_class, _SharedObserverModulePattern) + annotator_class + for annotator_class in self._all_patterns() + if issubclass(annotator_class, _SharedObserverModulePattern) ] - shared_observer_nodes = set() + shared_observer_nodes: dict[ + torch.fx.Node, type[_SharedObserverModulePattern] + ] = {} for annotator in shared_observer_annotators: node_to_annotator_and_match_dict = annotator._match_all_patterns(model) - - # We only care about the nodes for node in node_to_annotator_and_match_dict: - shared_observer_nodes.add(node) + # First match wins if multiple patterns claim the same node — + # deterministic given registry iteration order. + shared_observer_nodes.setdefault(node, annotator) return shared_observer_nodes def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: @@ -260,40 +271,68 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: if not isinstance(model, torch.fx.GraphModule): raise TypeError("Model must be a torch.fx.GraphModule") - # Matching phase - apply all annotators to model + # Match all annotator patterns. node_to_annotator_match_info_dict = self._match_all_annotators(model) - # Sorting phase - sort all matches in the order to try annotating in + # Sort matches into priority order (config-level > pattern-length > topo). + # Used to pick each node's winning config when multiple matches touch it. sorted_nodes_with_annotation_match_info = self._sort_nodes_in_annotation_order( model, node_to_annotator_match_info_dict ) - # Annotation phase - go through sorted nodes with matches list to annotate - shared_observer_nodes = self._get_shared_observer_nodes(model) - # Build pass-invariant context once; shared by all annotator invocations. + # Winning config per node + priority + pattern group. Pattern-match + # dicts key on a single primary node but annotate every covered node + # (Conv+Sigmoid annotates both conv and sigmoid); expand the coverage + # so every covered node has (config, priority, pattern_group) for + # its slots. Priority = position in the sorted list (lower = higher + # priority); used by reconcile() to break dtype ties within a + # shared-observer group. + winning_configs: dict[torch.fx.Node, Any] = {} + node_priorities: dict[torch.fx.Node, int] = {} + pattern_groups: dict[torch.fx.Node, frozenset[torch.fx.Node]] = {} + for priority, (_primary_node, cfg, match_info) in enumerate( + sorted_nodes_with_annotation_match_info + ): + covered = frozenset(_nodes_covered_by(match_info)) + for covered_node in covered: + if covered_node in winning_configs: + continue # first (highest-priority) pattern wins + winning_configs[covered_node] = cfg + node_priorities[covered_node] = priority + pattern_groups[covered_node] = covered + + shared_observer_node_to_pattern = self._get_shared_observer_nodes(model) + # Build pass-invariant AnnotationContext once; still needed for the + # kv-cache-override pass below, which reuses the standard annotator + # helpers (``_get_input_qspec_map``) to compute its overrides. context = AnnotationContext( module_name_to_state_names_map=self._module_name_to_state_names_map, - shared_observer_nodes=shared_observer_nodes, ) - for node, config, annotator_match_info in sorted_nodes_with_annotation_match_info: - if is_node_annotated(node): - continue - annotation_config = AnnotationConfig.from_quantizer_config(config) - annotator_match_info.annotator_func( - annotator_match_info.annotator_match, - annotation_config, - context, - ) + # Reconciliation pipeline. Replaces the previous per-node annotator + # loop and the Phase-6 post-pass (fixed-range/always-affine qscheme + # overrides): both are now expressed as OP_INTRINSIC_* provisional + # qspecs that participate in per-component reconciliation. Any + # ordering surprises produced by the old interleaved traversal + # (see YOLOX cat-of-sigmoid) are eliminated by construction; when a + # shared-observer group has conflicting dtypes, reconciliation picks + # the highest-priority proposal. + ctx = _AnnotationContext( + winning_configs=winning_configs, + node_priorities=node_priorities, + pattern_groups=pattern_groups, + shared_observer_nodes=shared_observer_node_to_pattern, + ) + annotate_via_reconciliation(model, ctx) + + # Module-level input/output/state overrides. TODO: fold these into + # reconciliation as USER_CONFIG_* proposals so they participate in + # unification (per design). For now they still run as a post-pass and + # overwrite reconciled specs on module-boundary nodes. _annotate_module_level_specs( self._module_configs, self._module_name_to_state_names_map, model ) - # Post-annotation graph pass to go through the model and adjust any qspecs which - # are following always affine or fixed range nodes. - for node in model.graph.nodes: - adjust_output_qspec_for_qscheme_and_propagate(node, shared_observer_nodes) - # Force-apply each cache spec last so it wins over any module-scope # annotation that may have claimed the cache op via a wildcard. Cache # specs are global-only knobs (see KVCacheQuantConfig docstring). diff --git a/tests/quantization/test_annotation_pattern_registry.py b/tests/quantization/test_annotation_pattern_registry.py index 5ce28d6..b8393ec 100644 --- a/tests/quantization/test_annotation_pattern_registry.py +++ b/tests/quantization/test_annotation_pattern_registry.py @@ -3626,6 +3626,12 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]: patterns = [mock_pattern_1] return patterns + @classmethod + def generate_qspec_sharing_constraints(cls, node, qspecs): + # Stub — abstract on the base class; unused by this test, + # which only exercises the length-validation branch. + return [] + pattern = Pattern() with pytest.raises(RuntimeError): pattern.get_patterns() diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index f70e24d..8bc33a1 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -16,6 +16,10 @@ import torch.nn as nn from torch.export.dynamic_shapes import Dim from torch.ops import coreai +from torchao.quantization.pt2e.quantizer import ( + QuantizationSpec as TorchAOQuantizationSpec, + SharedQuantizationSpec, +) from coreai_opt import ExportBackend from coreai_opt._utils.metadata_utils import ( @@ -625,10 +629,6 @@ def match_single_pattern(cls, _model, _pattern): cls.match_single_pattern_call_count += 1 return {} - @classmethod - def get_annotator_func(cls): - pass - yield TestAnnotator del _AnnotationPatternRegistry.REGISTRY[pattern_name] @@ -666,6 +666,127 @@ def test_annotate_all_patterns_generate_patterns_called_once(self, basic_config, ) +class TestCatOfSigmoidReconciliation: + """End-to-end reconciliation on a YOLOX-shape cat-of-sigmoid model. + + Exercises the full annotation pipeline (pattern match → reconcile → + resolve) on the pattern that motivated the reconciler: a ``cat`` group + whose branches disagree on qscheme (one plain conv branch wants + symmetric; two sigmoid branches want affine via op-intrinsic override). + The reconciler must lattice-join to affine, pick a topo-first anchor, + and give every other slot a SharedQuantizationSpec pointing at it. + """ + + class _CatOfSigmoid(nn.Module): + """``cat([reg_conv(x), sigmoid(obj_conv(x)), sigmoid(cls_conv(x))])``.""" + + def __init__(self) -> None: + super().__init__() + self.reg = nn.Conv2d(3, 4, kernel_size=1) + self.obj = nn.Conv2d(3, 1, kernel_size=1) + self.cls = nn.Conv2d(3, 8, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + r = self.reg(x) + o = torch.sigmoid(self.obj(x)) + c = torch.sigmoid(self.cls(x)) + return torch.cat([r, o, c], dim=1) + + @staticmethod + def _by_target(model: torch.fx.GraphModule, needle: str) -> list[torch.fx.Node]: + return [ + n for n in model.graph.nodes + if n.op == "call_function" and needle in str(n.target) + ] + + @staticmethod + def _annotation(node: torch.fx.Node): + from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY + return node.meta.get(Q_ANNOTATION_KEY) + + def test_cat_of_sigmoid_shared_observer_topology( + self, weight_input_act_output_act_config + ): + """Every branch in the cat group should end up pointing at one anchor. + + Load-bearing invariants: + + - ``conv2d`` (the ``reg`` branch — topologically first covered node + in the cat group) carries the concrete ``QuantizationSpec``. + - Its qscheme is ``per_tensor_affine``, not symmetric — the two + sigmoid branches force affine via op-intrinsic override, and the + qscheme lattice-join takes the looser of the two. + - ``sigmoid``, ``sigmoid_1``, and ``cat`` all have + ``output_qspec = SharedQuantizationSpec(edge_or_node=conv2d)``. + - The cat's three input slots also share on ``conv2d``. + - ``conv2d_1`` / ``conv2d_2`` (the two convs feeding the sigmoids + inside the pattern group) have no ``output_qspec`` — the + internal-edge slot is intentionally left unannotated so torchao + doesn't insert a second observer between them and their sigmoid. + """ + model = self._CatOfSigmoid().eval() + example_inputs = (torch.randn(1, 3, 8, 8),) + + prepared = Quantizer(model, weight_input_act_output_act_config).prepare( + example_inputs + ) + + convs = self._by_target(prepared, "conv2d") + sigmoids = self._by_target(prepared, "sigmoid") + cats = self._by_target(prepared, "cat") + + assert [n.name for n in convs] == ["conv2d", "conv2d_1", "conv2d_2"] + assert [n.name for n in sigmoids] == ["sigmoid", "sigmoid_1"] + assert [n.name for n in cats] == ["cat"] + + anchor, internal_conv_1, internal_conv_2 = convs + sigmoid_a, sigmoid_b = sigmoids + (cat_node,) = cats + + # Anchor carries a concrete spec, lattice-joined to affine. + anchor_ann = self._annotation(anchor) + assert anchor_ann is not None and anchor_ann._annotated + assert isinstance(anchor_ann.output_qspec, TorchAOQuantizationSpec) + assert anchor_ann.output_qspec.qscheme == torch.per_tensor_affine + assert anchor_ann.output_qspec.dtype == torch.int8 + + # The two convs feeding the sigmoid pattern get no output_qspec: + # the sigmoid is annotated on the shared side and the intermediate + # edge is pattern-internal. + for internal in (internal_conv_1, internal_conv_2): + internal_ann = self._annotation(internal) + assert internal_ann is not None + assert internal_ann.output_qspec is None, ( + f"{internal.name} is a pattern-internal producer feeding " + f"sigmoid; its output slot must not be annotated." + ) + + # Every other slot in the group shares on the anchor. + for sharer in (sigmoid_a, sigmoid_b, cat_node): + sharer_ann = self._annotation(sharer) + assert sharer_ann is not None and sharer_ann._annotated + assert isinstance(sharer_ann.output_qspec, SharedQuantizationSpec) + assert sharer_ann.output_qspec.edge_or_node is anchor, ( + f"{sharer.name}.output_qspec should share on {anchor.name}, " + f"got {sharer_ann.output_qspec.edge_or_node}" + ) + + # cat's three inputs share on the anchor too. + cat_ann = self._annotation(cat_node) + cat_input_producers = list(cat_ann.input_qspec_map.keys()) + assert cat_input_producers == [anchor, sigmoid_a, sigmoid_b] + for producer, input_spec in cat_ann.input_qspec_map.items(): + assert isinstance(input_spec, SharedQuantizationSpec), ( + f"cat.input[{producer.name}] should be a SharedQuantizationSpec, " + f"got {type(input_spec).__name__}" + ) + assert input_spec.edge_or_node is anchor + + # Sanity: model still runs forward after annotation. + out = prepared(*example_inputs) + assert out.shape == (1, 4 + 1 + 8, 8, 8) + + class TestAPIOrdering: def test_multiple_prepare_calls( self, simple_conv_linear_model, basic_config, simple_model_input diff --git a/tests/quantization/test_qspec_reconcile.py b/tests/quantization/test_qspec_reconcile.py new file mode 100644 index 0000000..ec96afc --- /dev/null +++ b/tests/quantization/test_qspec_reconcile.py @@ -0,0 +1,389 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Unit tests for the constraint-queue qspec reconciliation pipeline. + +Tests here are pure — they construct :class:`ProvisionalQSpec` +:class:`ProvisionalQSpecMap` values and :class:`Constraint`s by hand and +exercise ``.apply()`` directly. No fx graph construction. End-to-end +pipeline behavior is covered by trace-driven checks against the +walkthrough toy. +""" + +from unittest.mock import Mock + +import pytest +import torch +from torchao.quantization.pt2e.observer import MinMaxObserver + +from coreai_opt.quantization._graph._qspec_constraints import ( + ShareFields, + ShareObserverInstance, + _reconcile_field, +) +from coreai_opt.quantization._graph._qspec_types import ( + FieldName, + FieldValue, + NodeSlot, + ProvisionalQSpec, + ProvisionalQSpecMap, + ReconciliationError, + SlotKind, +) + +# --------------------------------------------------------------------------- +# Test helpers. +# --------------------------------------------------------------------------- + + +def _slot(name: str = "s", kind: SlotKind = SlotKind.OUTPUT, arg_index: int = 0) -> NodeSlot: + """Opaque ``NodeSlot`` — reconciler never inspects the fx node.""" + return NodeSlot(node=Mock(name=name), kind=kind, arg_index=arg_index) + + +def _pspec(**fields: FieldValue) -> ProvisionalQSpec: + """Build a ProvisionalQSpec by keyword: DTYPE=FieldValue(int8, 0), ...""" + field_map = {FieldName[key]: value for key, value in fields.items()} + return ProvisionalQSpec(fields=field_map) + + +def _fv(value, priority: int = 0) -> FieldValue: + return FieldValue(value=value, priority=priority) + + +# --------------------------------------------------------------------------- +# _reconcile_field policies. +# --------------------------------------------------------------------------- + + +class TestReconcileField: + def test_dtype_priority_wins(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int4, priority=1)), + b: _pspec(DTYPE=_fv(torch.int8, priority=5)), + } + result = _reconcile_field(FieldName.DTYPE, frozenset({a, b}), state) + assert result.value == torch.int4 # priority 1 beats 5 + assert result.priority == 1 + + def test_dtype_tie_first_encountered(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int8, priority=3)), + b: _pspec(DTYPE=_fv(torch.int4, priority=3)), + } + result = _reconcile_field(FieldName.DTYPE, frozenset({a, b}), state) + # min() with equal keys returns the first — encounter order + assert result.value in (torch.int8, torch.int4) + + def test_qscheme_lattice_symmetric_and_affine(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(QSCHEME=_fv(torch.per_tensor_symmetric, priority=0)), + b: _pspec(QSCHEME=_fv(torch.per_tensor_affine, priority=5)), + } + result = _reconcile_field(FieldName.QSCHEME, frozenset({a, b}), state) + assert result.value == torch.per_tensor_affine # looser wins + assert result.priority == 0 # min of {0, 5} + + def test_qscheme_lattice_per_tensor_and_per_channel(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(QSCHEME=_fv(torch.per_tensor_symmetric, priority=0)), + b: _pspec(QSCHEME=_fv(torch.per_channel_symmetric, priority=0)), + } + result = _reconcile_field(FieldName.QSCHEME, frozenset({a, b}), state) + assert result.value == torch.per_channel_symmetric + + def test_range_min_union(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(QUANT_MIN=_fv(-127, priority=0)), + b: _pspec(QUANT_MIN=_fv(-128, priority=5)), + } + result = _reconcile_field(FieldName.QUANT_MIN, frozenset({a, b}), state) + assert result.value == -128 # smaller wins for min + assert result.priority == 0 + + def test_range_max_union(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(QUANT_MAX=_fv(127, priority=0)), + b: _pspec(QUANT_MAX=_fv(255, priority=5)), + } + result = _reconcile_field(FieldName.QUANT_MAX, frozenset({a, b}), state) + assert result.value == 255 # larger wins for max + assert result.priority == 0 + + def test_ch_axis_must_agree(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(CH_AXIS=_fv(0, priority=0)), + b: _pspec(CH_AXIS=_fv(1, priority=0)), + } + with pytest.raises(ReconciliationError): + _reconcile_field(FieldName.CH_AXIS, frozenset({a, b}), state) + + def test_is_dynamic_must_agree(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(IS_DYNAMIC=_fv(True, priority=0)), + b: _pspec(IS_DYNAMIC=_fv(False, priority=0)), + } + with pytest.raises(ReconciliationError): + _reconcile_field(FieldName.IS_DYNAMIC, frozenset({a, b}), state) + + def test_missing_field_returns_none(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = {a: _pspec(), b: _pspec()} + assert _reconcile_field(FieldName.DTYPE, frozenset({a, b}), state) is None + + +# --------------------------------------------------------------------------- +# ShareFields. +# --------------------------------------------------------------------------- + + +class TestShareFields: + def test_broadcasts_winner_to_all_slots(self) -> None: + a, b, c = _slot("a"), _slot("b"), _slot("c") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int4, priority=0)), + b: _pspec(DTYPE=_fv(torch.int8, priority=5)), + c: _pspec(), # no proposal + } + con = ShareFields(_slots=frozenset({a, b, c}), fields=frozenset({FieldName.DTYPE})) + changed = con.apply(state) + assert changed == {b, c} # a already had int4; b and c gain it + assert state[a].fields[FieldName.DTYPE].value == torch.int4 + assert state[b].fields[FieldName.DTYPE].value == torch.int4 + assert state[c].fields[FieldName.DTYPE].value == torch.int4 + # All at priority 0 after reconciliation. + assert state[a].fields[FieldName.DTYPE].priority == 0 + assert state[b].fields[FieldName.DTYPE].priority == 0 + + def test_noop_when_already_reconciled(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int8, priority=0)), + b: _pspec(DTYPE=_fv(torch.int8, priority=0)), + } + con = ShareFields(_slots=frozenset({a, b}), fields=frozenset({FieldName.DTYPE})) + assert con.apply(state) == set() + + def test_multiple_fields(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec( + DTYPE=_fv(torch.int4, priority=0), + QSCHEME=_fv(torch.per_tensor_symmetric, priority=0), + ), + b: _pspec( + DTYPE=_fv(torch.int8, priority=5), + QSCHEME=_fv(torch.per_tensor_affine, priority=5), + ), + } + con = ShareFields( + _slots=frozenset({a, b}), + fields=frozenset({FieldName.DTYPE, FieldName.QSCHEME}), + ) + con.apply(state) + # a's dtype wins (higher priority); qscheme resolves via lattice. + assert state[a].fields[FieldName.DTYPE].value == torch.int4 + assert state[b].fields[FieldName.DTYPE].value == torch.int4 + assert state[a].fields[FieldName.QSCHEME].value == torch.per_tensor_affine + assert state[b].fields[FieldName.QSCHEME].value == torch.per_tensor_affine + + +# --------------------------------------------------------------------------- +# ShareObserverInstance. +# --------------------------------------------------------------------------- + + +class TestShareObserverInstance: + def test_merges_two_slots_into_one_instance(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int8, priority=0)), + b: _pspec(DTYPE=_fv(torch.int8, priority=0)), + } + assert state[a] is not state[b] + con = ShareObserverInstance(_slots=frozenset({a, b})) + con.apply(state) + assert state[a] is state[b] # identity, not just value equality + + def test_field_mutation_after_share_propagates(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = {a: _pspec(), b: _pspec()} + ShareObserverInstance(_slots=frozenset({a, b})).apply(state) + state[a].fields[FieldName.DTYPE] = _fv(torch.int8, 0) + # b's ProvisionalQSpec is the same object → sees the new field. + assert state[b].fields[FieldName.DTYPE].value == torch.int8 + + def test_transitive_merge_pulls_in_prior_sharers(self) -> None: + a, b, c = _slot("a"), _slot("b"), _slot("c") + state: ProvisionalQSpecMap = {a: _pspec(), b: _pspec(), c: _pspec()} + # First merge a and b. + ShareObserverInstance(_slots=frozenset({a, b})).apply(state) + assert state[a] is state[b] + # Now merge a with c — b should be pulled in transitively. + ShareObserverInstance(_slots=frozenset({a, c})).apply(state) + assert state[a] is state[b] is state[c] + + def test_reconciles_fields_across_merged_group(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int4, priority=0)), + b: _pspec(DTYPE=_fv(torch.int8, priority=5)), + } + ShareObserverInstance(_slots=frozenset({a, b})).apply(state) + assert state[a].fields[FieldName.DTYPE].value == torch.int4 + assert state[b].fields[FieldName.DTYPE].value == torch.int4 + assert state[a] is state[b] + + def test_noop_when_already_shared_with_correct_fields(self) -> None: + a, b = _slot("a"), _slot("b") + shared = _pspec(DTYPE=_fv(torch.int8, priority=0)) + state: ProvisionalQSpecMap = {a: shared, b: shared} + changed = ShareObserverInstance(_slots=frozenset({a, b})).apply(state) + assert changed == set() + assert state[a] is shared + assert state[b] is shared + + +# --------------------------------------------------------------------------- +# Convergence / no-op behavior. +# --------------------------------------------------------------------------- + + +class TestConvergence: + def test_repeated_share_fields_stabilizes(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = { + a: _pspec(DTYPE=_fv(torch.int4, priority=1)), + b: _pspec(DTYPE=_fv(torch.int8, priority=5)), + } + con = ShareFields(_slots=frozenset({a, b}), fields=frozenset({FieldName.DTYPE})) + first = con.apply(state) + assert first # something changed + # Immediate re-apply is a no-op — proves the priority-inheritance + # rule prevents oscillation. + assert con.apply(state) == set() + + def test_repeated_share_observer_stabilizes(self) -> None: + a, b = _slot("a"), _slot("b") + state: ProvisionalQSpecMap = {a: _pspec(), b: _pspec()} + con = ShareObserverInstance(_slots=frozenset({a, b})) + first = con.apply(state) + assert first + assert con.apply(state) == set() + + +# --------------------------------------------------------------------------- +# YOLOX-shape scenario reconciliation. +# --------------------------------------------------------------------------- + + +def test_yolox_scenario_end_to_end_reconciliation() -> None: + """cat-of-sigmoid: two sigmoid-forced-affine outputs plus a wide-range + symmetric conv output. Group unifies to affine (lattice join), dtype + picks the highest-priority proposal.""" + conv_out = _slot("conv_out") + sig_a_out = _slot("sig_a_out") + sig_b_out = _slot("sig_b_out") + cat_out = _slot("cat_out") + + state: ProvisionalQSpecMap = { + conv_out: _pspec( + DTYPE=_fv(torch.int8, priority=5), + QSCHEME=_fv(torch.per_tensor_symmetric, priority=5), + QUANT_MIN=_fv(-128, priority=5), + QUANT_MAX=_fv(127, priority=5), + OBSERVER_CLASS=_fv(MinMaxObserver, priority=5), + IS_DYNAMIC=_fv(False, priority=5), + ), + sig_a_out: _pspec( + DTYPE=_fv(torch.int8, priority=5), + QSCHEME=_fv(torch.per_tensor_affine, priority=5), # sigmoid intrinsic + QUANT_MIN=_fv(-128, priority=5), + QUANT_MAX=_fv(127, priority=5), + OBSERVER_CLASS=_fv(MinMaxObserver, priority=5), + IS_DYNAMIC=_fv(False, priority=5), + ), + sig_b_out: _pspec( + DTYPE=_fv(torch.int8, priority=5), + QSCHEME=_fv(torch.per_tensor_affine, priority=5), + QUANT_MIN=_fv(-128, priority=5), + QUANT_MAX=_fv(127, priority=5), + OBSERVER_CLASS=_fv(MinMaxObserver, priority=5), + IS_DYNAMIC=_fv(False, priority=5), + ), + cat_out: _pspec( + DTYPE=_fv(torch.int8, priority=5), + QSCHEME=_fv(torch.per_tensor_symmetric, priority=5), + QUANT_MIN=_fv(-128, priority=5), + QUANT_MAX=_fv(127, priority=5), + OBSERVER_CLASS=_fv(MinMaxObserver, priority=5), + IS_DYNAMIC=_fv(False, priority=5), + ), + } + ShareObserverInstance(_slots=frozenset({conv_out, sig_a_out, sig_b_out, cat_out})).apply(state) + # All four share one ProvisionalQSpec now. + assert state[conv_out] is state[sig_a_out] is state[sig_b_out] is state[cat_out] + # Reconciled qscheme is affine (looser wins). + assert state[conv_out].fields[FieldName.QSCHEME].value == torch.per_tensor_affine + assert state[conv_out].fields[FieldName.DTYPE].value == torch.int8 + + +# --------------------------------------------------------------------------- +# Adjacent-edge sharing — the peephole case from precedence tests. +# --------------------------------------------------------------------------- + + +def test_adjacent_edge_share_resolves_dtype_conflict_by_priority() -> None: + """test_op_level_precedence-shape: linear1.OUTPUT wants int4 at lower + priority; linear2.INPUT wants int8 at higher priority. Adjacent-edge + sharing merges them; priority resolves dtype to int8.""" + linear1_out = _slot("linear1_out", kind=SlotKind.OUTPUT) + linear2_in = _slot("linear2_in", kind=SlotKind.INPUT) + state: ProvisionalQSpecMap = { + linear1_out: _pspec( + DTYPE=_fv(torch.int4, priority=5), + QSCHEME=_fv(torch.per_tensor_symmetric, priority=5), + OBSERVER_CLASS=_fv(MinMaxObserver, priority=5), + IS_DYNAMIC=_fv(False, priority=5), + ), + linear2_in: _pspec( + DTYPE=_fv(torch.int8, priority=0), + QSCHEME=_fv(torch.per_tensor_symmetric, priority=0), + OBSERVER_CLASS=_fv(MinMaxObserver, priority=0), + IS_DYNAMIC=_fv(False, priority=0), + ), + } + ShareObserverInstance(_slots=frozenset({linear1_out, linear2_in})).apply(state) + assert state[linear1_out] is state[linear2_in] + assert state[linear1_out].fields[FieldName.DTYPE].value == torch.int8 + + +# --------------------------------------------------------------------------- +# Concat axis-aware sharing — per-channel on concat axis stays independent. +# --------------------------------------------------------------------------- + + +def test_share_fields_dtype_only_leaves_scale_independent() -> None: + """When per-channel is along the concat axis, only DTYPE is shared — + scale/zp (via observer instance) stays per-input.""" + a_out = _slot("a_out") + b_out = _slot("b_out") + state: ProvisionalQSpecMap = { + a_out: _pspec(DTYPE=_fv(torch.int8, priority=0)), + b_out: _pspec(DTYPE=_fv(torch.int4, priority=5)), + } + ShareFields(_slots=frozenset({a_out, b_out}), fields=frozenset({FieldName.DTYPE})).apply(state) + # dtypes agree, but the ProvisionalQSpec objects are still distinct. + assert state[a_out] is not state[b_out] + assert state[a_out].fields[FieldName.DTYPE].value == torch.int8 + assert state[b_out].fields[FieldName.DTYPE].value == torch.int8