Adding Tests for CadenceFusedConvReluQuantizer (#16358)

Summary:
Pull Request resolved: #16358

A fused pattern is when the quantizer recognizes a sequence of operations and treats as a single unit for quantization purposes. So for example, for a Conv2D + ReLU fusion, rather than having something like this:
```
  input → [quantize] → conv2d → [dequantize] → [quantize] → relu → [dequantize] → output
```
a fused pattern quantizes them together like so:

```
  input → [quantize] → conv2d → relu → [dequantize] → output
```

We need to make a few changes in our framework to test this.

# Change 1: We allow graph builders to return a 3rd element for fused patterns
 For fused patterns like conv+relu, the quantization annotations are split across two nodes:
  - Output annotation is on the relu node (the final output of the fused pattern)
  - Input annotations are on the conv node (where the quantized inputs enter)

 The existing graph builders return (gm, target_node), which works for single-op patterns where both annotations are on the same node. For fused patterns, we need to know both nodes, so graph builders can now optionally return (gm, output_node,  input_source_node).

# Change 2: We check annotations on the correct nodes for fused patterns

The test previously assumed output_qspec and input_qspec_map were both on the same node. For fused patterns, they're on different nodes:
  - output_qspec is checked on the output node (relu)
  - input_qspec_map is checked on the input source node (conv)

This change is backwards-compatible: for non-fused patterns, both nodes are the same.

Reviewed By: hsharma35

Differential Revision: D89630759

Author: RahulC7

backends/cadence/aot/tests/test_quantizer_ops.py

@@ -43,15 +43,18 @@
 
 # Type alias for graph builder functions.
 # These functions take a test instance and return a graph module and the target op node.
+# For fused patterns (e.g., conv+relu), an optional third element specifies the node
+# whose args contain the quantized inputs (e.g., conv node for conv+relu fusion).
 GraphBuilderFn = Callable[
-    ["QuantizerAnnotationTest"], tuple[torch.fx.GraphModule, torch.fx.Node]
+    ["QuantizerAnnotationTest"],
+    tuple[torch.fx.GraphModule, torch.fx.Node]
+    | tuple[torch.fx.GraphModule, torch.fx.Node, torch.fx.Node],
 ]
 
 
 # Quantizers intentionally excluded from annotation testing.
 # These should be explicitly justified when added.
 EXCLUDED_FROM_ANNOTATION_TESTING: set[type[CadenceQuantizer]] = {
-    CadenceFusedConvReluQuantizer,  # TODO: T247438151 Add test coverage
     CadenceNopQuantizer,  # No-op quantizer, doesn't annotate anything
     CadenceW8A32MixedQuantizer,  # TODO: T247438158 Add test coverage
     CadenceRmsNormNopQuantizer,  # No-op quantizer, doesn't annotate anything, preserves rms_norm from decomposition
@@ -200,6 +203,16 @@
         # Use None to skip comparison for bias since it's a DerivedQuantizationSpec
         [None, qconfig_A8W8.input_activation, qconfig_A8W8.weight],
     ),
+    # CadenceFusedConvReluQuantizer test cases
+    (
+        "fused_conv2d_relu_A8W8sym",
+        lambda self: self._build_conv2d_relu_graph(),
+        CadenceFusedConvReluQuantizer(),
+        torch.ops.aten.relu.default,
+        qconfig_A8W8sym.output_activation,
+        # For fused conv2d+relu: [input_activation, weight] from conv2d node
+        [qconfig_A8W8sym.input_activation, qconfig_A8W8sym.weight],
+    ),
 ]
 
 # Derive the set of tested quantizer classes from the test cases.
@@ -441,6 +454,52 @@ def _build_addmm_graph(self) -> tuple[torch.fx.GraphModule, torch.fx.Node]:
         self.assertEqual(len(addmm_nodes), 1, "Should find exactly one addmm node")
         return gm, addmm_nodes[0]
 
+    def _build_conv2d_relu_graph(
+        self,
+    ) -> tuple[torch.fx.GraphModule, torch.fx.Node, torch.fx.Node]:
+        """Build a graph with a conv2d followed by relu (fused pattern).
+
+        Returns:
+            A tuple of (graph_module, relu_node, conv_node).
+            The relu_node is the target node where the annotation is placed.
+            The conv_node is the input source node whose args contain the quantized inputs.
+        """
+        builder = GraphBuilder()
+        # Input shape: (batch, in_channels, height, width)
+        x = builder.placeholder("x", torch.randn(1, 3, 8, 8))
+        # Weight shape: (out_channels, in_channels, kernel_h, kernel_w)
+        weight = builder.placeholder("weight", torch.randn(6, 3, 3, 3))
+        conv2d = builder.call_operator(
+            op=torch.ops.aten.conv2d.default,
+            args=(x, weight),
+            meta=NodeMetadata(
+                {"source_fn_stack": [("conv2d", torch.ops.aten.conv2d.default)]}
+            ),
+        )
+        relu = builder.call_operator(
+            op=torch.ops.aten.relu.default,
+            args=(conv2d,),
+            meta=NodeMetadata(
+                {"source_fn_stack": [("relu", torch.ops.aten.relu.default)]}
+            ),
+        )
+        builder.output([relu])
+        gm = builder.get_graph_module()
+
+        relu_nodes = gm.graph.find_nodes(
+            op="call_function",
+            target=torch.ops.aten.relu.default,
+        )
+        self.assertEqual(len(relu_nodes), 1, "Should find exactly one relu node")
+
+        conv2d_nodes = gm.graph.find_nodes(
+            op="call_function",
+            target=torch.ops.aten.conv2d.default,
+        )
+        self.assertEqual(len(conv2d_nodes), 1, "Should find exactly one conv2d node")
+
+        return gm, relu_nodes[0], conv2d_nodes[0]
+
     @parameterized.expand(QUANTIZER_ANNOTATION_TEST_CASES)
     def test_quantizer_annotation(
         self,
@@ -452,28 +511,45 @@ def test_quantizer_annotation(
         expected_input_qspecs: list[QuantizationSpec | None],
     ) -> None:
         """Parameterized test for quantizer annotations."""
-        gm, op_node = graph_builder_fn(self)
+        result = graph_builder_fn(self)
+        # Handle both 2-element and 3-element returns from graph builders.
+        # For fused patterns, the 3rd element specifies the node whose args
+        # contain the quantized inputs (e.g., conv node for conv+relu fusion).
+        if len(result) == 3:
+            gm = result[0]
+            output_node = result[1]
+            input_source_node = result[2]
+        else:
+            gm = result[0]
+            output_node = result[1]
+            input_source_node = output_node
 
         quantizer.annotate(gm)
 
-        annotation: QuantizationAnnotation = op_node.meta[Q_ANNOTATION_KEY]
-        self.assertTrue(annotation._annotated)
-
-        # Verify output annotation
-        self.assertEqual(annotation.output_qspec, expected_output_qspec)
+        # Verify output annotation (always on the output node)
+        output_annotation: QuantizationAnnotation = output_node.meta[Q_ANNOTATION_KEY]
+        self.assertTrue(output_annotation._annotated)
+        self.assertEqual(output_annotation.output_qspec, expected_output_qspec)
 
-        # Verify input annotations
-        self.assertEqual(len(annotation.input_qspec_map), len(expected_input_qspecs))
-        for input_node, input_qspec in annotation.input_qspec_map.items():
-            # Find the index of this input node in the op's args
+        # Verify input annotations (on the input source node, which may differ for fused patterns)
+        input_annotation: QuantizationAnnotation = input_source_node.meta[
+            Q_ANNOTATION_KEY
+        ]
+        self.assertEqual(
+            len(input_annotation.input_qspec_map), len(expected_input_qspecs)
+        )
+        for input_node, input_qspec in input_annotation.input_qspec_map.items():
+            # Find the index of this input node in the input source node's args
             arg_index = None
-            for i, arg in enumerate(op_node.args):
+            args = input_source_node.args
+            assert isinstance(args, tuple)
+            for i, arg in enumerate(args):
                 if arg is input_node:
                     arg_index = i
                     break
             self.assertIsNotNone(
                 arg_index,
-                f"Input node {input_node} not found in op_node.args",
+                f"Input node {input_node} not found in input_source_node.args",
             )
             # Skip comparison if expected qspec is None (e.g., for DerivedQuantizationSpec)
             if expected_input_qspecs[arg_index] is not None: