Fix quantized max_pool2d output observer for tuple-returning ops

Summary:
prepare_pt2e skips the output observer for max_pool2d_with_indices because its output is a tuple, not a single tensor. This caused the quantized_max_pool2d fusion to silently fail.

Fix: annotate getitem[0] (single tensor) instead of the tuple-returning op. Also bail out if indices are consumed or the graph structure is unexpected.

Differential Revision: D103436172

Author: abeakkas

backends/cadence/aot/quantizer/patterns.py

@@ -6,9 +6,10 @@
 
 # pyre-strict
 
+import operator
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import List, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 from executorch.backends.cadence.aot.quantizer.utils import get_bias_qparams
@@ -489,11 +490,32 @@ def partition_types(self) -> List[OpOverload]:
 
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule]
-    ) -> Tuple[PartitionAnchors, fx.Node]:
+    ) -> Optional[Tuple[PartitionAnchors, fx.Node]]:
         # pyre-fixme[29]: `Union[BoundMethod[typing.Callable(torch._C.TensorBase.__ge...
         max_pool_node = fused_partition[0].nodes[-1]
 
-        # Input and output share quantization parameters since max is order-preserving
+        # Find getitem[0] (values) node. Since max_pool2d_with_indices returns a tuple,
+        # the output observer must be placed on getitem[0] (a single tensor) rather than
+        # the tuple-returning op itself - otherwise prepare_pt2e silently skips the
+        # output observer.
+        getitem_0 = None
+        for user in max_pool_node.users:
+            if user.target is not operator.getitem:
+                # Unexpected consumer of tuple output - skip quantization
+                return None
+            if user.args[1] == 0:
+                if getitem_0 is not None:
+                    # Multiple getitem[0] nodes - unexpected graph structure
+                    return None
+                getitem_0 = user
+            elif user.args[1] == 1 and len(user.users) > 0:
+                # Indices are consumed downstream - can't quantize this op
+                return None
+
+        if getitem_0 is None:
+            # No getitem[0] found - values output is unused
+            return None
+
         return (
             PartitionAnchors(
                 inputs=[(max_pool_node, 0)],
@@ -505,7 +527,7 @@ def get_anchors(
                 ],
                 output=[
                     (
-                        max_pool_node,
+                        getitem_0,
                         SharedQuantizationSpec((max_pool_node.args[0], max_pool_node)),
                     )
                 ],